diff --git a/.claude/agents/mcp-protocol-expert.md b/.claude/agents/mcp-protocol-expert.md index 825c67472b..be0cbb33e9 100644 --- a/.claude/agents/mcp-protocol-expert.md +++ b/.claude/agents/mcp-protocol-expert.md @@ -37,9 +37,10 @@ Defer to: oauth-expert (OAuth/OIDC), kubernetes-expert (K8s operator), toolhive- - For the full, current set of pages under that version, fetch `https://modelcontextprotocol.io/sitemap.xml` and filter for the resolved version string — this is self-correcting as pages are added, renamed, or removed, unlike a hand-maintained list. - Common pages to expect (verify against the sitemap, don't assume the exact path): architecture, basic, basic/authorization, basic/lifecycle, basic/transports, basic/utilities/{cancellation,ping,progress,tasks}, client/{elicitation,roots,sampling}, server, server/{tools,resources,prompts}, server/utilities/{completion,logging,pagination}, schema, changelog. - Also check `https://modelcontextprotocol.io/extensions/auth/overview` for MCP auth extensions, which live outside the versioned spec tree. -3. Cross-reference whichever source answered with ToolHive's implementation. -4. Provide guidance based on the latest spec. -5. Explicitly note any discrepancies between spec and implementation, and note which source (context7 vs. WebFetch) backed the answer if it's non-obvious. +3. **For schema-level questions (exact field names, required-vs-optional, discriminated unions, literal error-code values), prefer the raw TypeScript schema over the rendered `.../schema` page or context7 snippets.** The rendered/curated sources can summarize or truncate a ~3k-line file; the ground truth both are generated from is `https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema//schema.ts` (plain text, JSDoc-commented). Resolve `` the same way as above (don't hardcode a date) — list `https://github.com/modelcontextprotocol/modelcontextprotocol/tree/main/schema` if unsure which dated folder is current — then WebFetch that file directly with a narrow prompt for the type in question. +4. Cross-reference whichever source answered with ToolHive's implementation. +5. Provide guidance based on the latest spec. +6. Explicitly note any discrepancies between spec and implementation, and note which source (context7, WebFetch docs, or schema.ts) backed the answer if it's non-obvious. ## Your Expertise diff --git a/.claude/rules/go-style.md b/.claude/rules/go-style.md index 886aa651c2..8b86ac20fa 100644 --- a/.claude/rules/go-style.md +++ b/.claude/rules/go-style.md @@ -200,3 +200,18 @@ Before implementing any non-trivial functionality from scratch: 3. **Look for existing Go packages** — search for well-maintained OSS libraries that solve the problem before writing custom implementations. Implementing from scratch should be a last resort, justified by a specific gap no existing solution fills. + +## JSON-RPC Envelopes + +Never `json.Marshal` (or `json.NewEncoder().Encode`) a `jsonrpc2` type. +`jsonrpc2.Response` carries no json tags and `jsonrpc2.ID`'s only field is +unexported, so reflection emits Go-capitalized keys, drops the mandatory +`"jsonrpc":"2.0"` tag, and renders every id — valid or not — as `{}` (#5950). +Serialize through `jsonrpc2.EncodeMessage`; for HTTP error responses use +`pkg/mcp.EncodeJSONRPCError` / `pkg/mcp.WriteJSONRPCError`, which add +encode-before-header ordering and a conformant fallback body. + +When an envelope must be hand-built as a map (only where no `jsonrpc2.ID` is +in hand), omit the `"id"` key when the id is absent — never emit `"id":null`. +MCP's `RequestId` is `string | number`, so null is not representable, and the +reference TypeScript SDK client throws on it inside its transport (#6038). diff --git a/.github/workflows/api-compat.yml b/.github/workflows/api-compat.yml index f0c42ca464..6fcbdfe8d1 100644 --- a/.github/workflows/api-compat.yml +++ b/.github/workflows/api-compat.yml @@ -50,7 +50,7 @@ jobs: timeout-minutes: 10 steps: - name: Checkout PR HEAD - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index d2ce9b05f1..75ad13a273 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -46,7 +46,7 @@ jobs: actions: read # Required for Claude to read CI results on PRs steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: fetch-depth: 1 @@ -59,7 +59,7 @@ jobs: - name: Run Claude Code id: claude - uses: anthropics/claude-code-action@f1bd27ca5b54584506e40e17884d90bdaaa1a9b3 # v1 + uses: anthropics/claude-code-action@be7b93b1907a4abad570368f3c74b6fe3807510b # v1 with: anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} # Security: Restrict tools to prevent arbitrary code execution. diff --git a/.github/workflows/create-release-pr.yml b/.github/workflows/create-release-pr.yml index 2499bb16a4..be1f65bb42 100644 --- a/.github/workflows/create-release-pr.yml +++ b/.github/workflows/create-release-pr.yml @@ -38,7 +38,7 @@ jobs: private-key: ${{ secrets.RELEASE_APP_PRIVATE_KEY }} - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 with: diff --git a/.github/workflows/create-release-tag.yml b/.github/workflows/create-release-tag.yml index 5ca298fb6d..b9c4cbf5fc 100644 --- a/.github/workflows/create-release-tag.yml +++ b/.github/workflows/create-release-tag.yml @@ -28,7 +28,7 @@ jobs: private-key: ${{ secrets.RELEASE_APP_PRIVATE_KEY }} - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: fetch-depth: 0 diff --git a/.github/workflows/e2e-tests.yml b/.github/workflows/e2e-tests.yml index 6be21c0c40..4c398a42b1 100644 --- a/.github/workflows/e2e-tests.yml +++ b/.github/workflows/e2e-tests.yml @@ -11,7 +11,7 @@ jobs: name: Build ToolHive Binary runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -50,6 +50,7 @@ jobs: - title: core label_filter: core artifact: e2e-test-results-core + procs: 4 - title: mcp-run label_filter: mcp-run artifact: e2e-test-results-mcp-run @@ -57,8 +58,12 @@ jobs: label_filter: mcp-protocol artifact: e2e-test-results-mcp-protocol - title: proxy - label_filter: proxy + label_filter: "proxy && !isolation" artifact: e2e-test-results-proxy + - title: network-isolation + label_filter: isolation + artifact: e2e-test-results-network-isolation + test_timeout: 25m - title: middleware label_filter: 'middleware || stability' artifact: e2e-test-results-middleware @@ -86,7 +91,7 @@ jobs: label_filter: llm artifact: e2e-test-results-llm steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -99,7 +104,7 @@ jobs: go mod download - name: Install Ginkgo CLI run: | - go install github.com/onsi/ginkgo/v2/ginkgo@latest + go install github.com/onsi/ginkgo/v2/ginkgo@v2.28.1 - name: Download ToolHive binary uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7 @@ -121,22 +126,52 @@ jobs: - name: Pre-pull container images run: | - # Pre-pull images used by E2E tests so that workload creation - # does not pay the image-pull cost inside the 60s API timeout. - docker pull ghcr.io/stackloklabs/osv-mcp/server:0.1.3 & - docker pull ghcr.io/stackloklabs/gofetch/server:1.0.2 & - docker pull ghcr.io/stacklok/toolhive/egress-proxy:latest & + # Pre-pull images used by E2E tests so that workload startup does not + # pay the image-pull cost inside a spec's readiness timeout. + # + # Images that tests reach through a registry server name (`thv run + # fetch`) are resolved from the binary under test instead of being + # hardcoded here. Renovate bumps github.com/stacklok/toolhive-catalog + # daily but has no manager for these docker pull lines, so hardcoded + # tags drift out of sync with the catalog silently and the pre-pull + # stops covering what the tests actually run. + servers="fetch osv" + images="ghcr.io/stacklok/toolhive/egress-proxy:latest" # yardstick is only needed for the vmcp test suite if [ "${{ matrix.label_filter }}" = "vmcp" ]; then - docker pull ghcr.io/stackloklabs/yardstick/yardstick-server:1.1.1 & + images="$images ghcr.io/stackloklabs/yardstick/yardstick-server:1.2.0" + fi + # the time server and yardstick (dual-era barrier mode) are only + # used by the proxy (streamable-http) test suite + if [ "${{ matrix.label_filter }}" = "proxy && !isolation" ]; then + servers="$servers time" + images="$images ghcr.io/stackloklabs/yardstick/yardstick-server:1.2.0" fi - # the time server is only used by the proxy (streamable-http) test suite - if [ "${{ matrix.label_filter }}" = "proxy" ]; then - docker pull ghcr.io/stacklok/dockyard/uvx/mcp-server-time:2026.1.26 & + # Envoy is only used by the network-isolation test suite + if [ "${{ matrix.title }}" = "network-isolation" ]; then + images="$images envoyproxy/envoy-distroless:v1.32.3" fi + + for server in $servers; do + image="$(./bin/thv registry info "$server" --format json | jq -er '.image')" + echo "Resolved registry server '$server' -> $image" + images="$images $image" + done + + for image in $images; do + docker pull -q "$image" & + done wait + + # `wait` reports success even when a background pull failed, so check + # each image landed locally. Failing here is far cheaper to diagnose + # than a spec timing out later while it waits on a live pull. echo "Pre-pulled images:" - docker images --format '{{.Repository}}:{{.Tag}}' | grep -E 'osv-mcp|gofetch|egress-proxy|yardstick|mcp-server-time' + for image in $images; do + docker image inspect "$image" >/dev/null \ + || { echo "::error::failed to pre-pull $image"; exit 1; } + echo " $image" + done - name: Run E2E tests (${{ matrix.title }}) env: @@ -144,6 +179,7 @@ jobs: TOOLHIVE_EGRESS_IMAGE: ghcr.io/stacklok/toolhive/egress-proxy:latest TEST_TIMEOUT: ${{ matrix.test_timeout || '15m' }} LABEL_FILTER: ${{ matrix.label_filter }} + PROCS: ${{ matrix.procs || 1 }} run: ./test/e2e/run_tests.sh - name: Upload test results (${{ matrix.title }}) @@ -160,8 +196,12 @@ jobs: runs-on: ubuntu-8cores-32gb needs: build-binary timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + target: [proxy, vmcp] steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Node.js uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 @@ -185,12 +225,13 @@ jobs: - name: Run MCP conformance suite env: THV_BINARY: ${{ github.workspace }}/bin/thv + CONFORMANCE_TARGET: ${{ matrix.target }} run: ./test/conformance/run-conformance.sh - name: Upload conformance results if: always() uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6 with: - name: conformance-results + name: conformance-results-${{ matrix.target }} path: test/conformance/results retention-days: 7 diff --git a/.github/workflows/helm-charts-test.yml b/.github/workflows/helm-charts-test.yml index cb053044a8..2829eb9189 100644 --- a/.github/workflows/helm-charts-test.yml +++ b/.github/workflows/helm-charts-test.yml @@ -12,7 +12,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: fetch-depth: 0 diff --git a/.github/workflows/helm-publish.yml b/.github/workflows/helm-publish.yml index 083cfb8803..48701dca59 100644 --- a/.github/workflows/helm-publish.yml +++ b/.github/workflows/helm-publish.yml @@ -14,7 +14,7 @@ jobs: contents: read steps: - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Verify tag matches VERSION file run: | @@ -61,7 +61,7 @@ jobs: path: deploy/charts/operator-crds steps: - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Extract version id: version diff --git a/.github/workflows/image-build-and-publish.yml b/.github/workflows/image-build-and-publish.yml index 33524f0176..3e3e3948e4 100644 --- a/.github/workflows/image-build-and-publish.yml +++ b/.github/workflows/image-build-and-publish.yml @@ -17,7 +17,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -85,7 +85,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Compute version number id: version-string @@ -147,7 +147,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -227,7 +227,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -296,7 +296,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 diff --git a/.github/workflows/issue-triage.yml b/.github/workflows/issue-triage.yml index 4197f792ac..b1dfdf7449 100644 --- a/.github/workflows/issue-triage.yml +++ b/.github/workflows/issue-triage.yml @@ -14,12 +14,12 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: fetch-depth: 0 - name: Run Claude Code for Issue Triage - uses: anthropics/claude-code-action@558b1d6cab4085c7753fe402c10bef0fbb92ac7a # v1.0.165 + uses: anthropics/claude-code-action@be7b93b1907a4abad570368f3c74b6fe3807510b # v1.0.183 with: prompt: | You are an issue triage assistant. Analyze issue #${{ github.event.issue.number }} in ${{ github.repository }} and apply appropriate labels. diff --git a/.github/workflows/license-headers.yml b/.github/workflows/license-headers.yml index 6e081f06d7..ca0460c5c8 100644 --- a/.github/workflows/license-headers.yml +++ b/.github/workflows/license-headers.yml @@ -14,7 +14,7 @@ jobs: name: Check License Headers runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index bf42b7ad28..7970713744 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -12,7 +12,7 @@ jobs: name: Lint Go Code runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 diff --git a/.github/workflows/operator-ci.yml b/.github/workflows/operator-ci.yml index 2665c3f54b..013984a19e 100644 --- a/.github/workflows/operator-ci.yml +++ b/.github/workflows/operator-ci.yml @@ -12,7 +12,7 @@ jobs: name: Operator Tests runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -33,7 +33,7 @@ jobs: name: Operator Tests Integration runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -54,7 +54,7 @@ jobs: name: Build Operator runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -75,7 +75,7 @@ jobs: name: Generate CRDs runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -108,7 +108,7 @@ jobs: name: Generate CRD Docs runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -157,7 +157,7 @@ jobs: steps: - name: Checkout code - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Helm uses: azure/setup-helm@9bc31f4ebc9c6b171d7bfbaa5d006ae7abdb4310 # v5.0.1 diff --git a/.github/workflows/pr-size-justification-template.md b/.github/workflows/pr-size-justification-template.md deleted file mode 100644 index f027f4b295..0000000000 --- a/.github/workflows/pr-size-justification-template.md +++ /dev/null @@ -1,26 +0,0 @@ -## Large PR Detected - -This PR exceeds 1000 lines of changes and requires justification before it can be reviewed. - -### How to unblock this PR: - -Add a section to your PR description with the following format: - -```markdown -## Large PR Justification - -[Explain why this PR must be large, such as:] -- Generated code that cannot be split -- Large refactoring that must be atomic -- Multiple related changes that would break if separated -- Migration or data transformation -``` - -### Alternative: - -Consider splitting this PR into smaller, focused changes (< 1000 lines each) for easier review and reduced risk. - -See our [Contributing Guidelines](CONTRIBUTING_LINK) for more details. - ---- -*This review will be automatically dismissed once you add the justification section.* diff --git a/.github/workflows/pr-size-label-apply.yml b/.github/workflows/pr-size-label-apply.yml index 77c5ced21c..9c6285ff78 100644 --- a/.github/workflows/pr-size-label-apply.yml +++ b/.github/workflows/pr-size-label-apply.yml @@ -1,4 +1,4 @@ -name: PR Size Labeler - Apply and Enforce +name: PR Size Labeler - Apply on: workflow_run: @@ -75,151 +75,4 @@ jobs: repo: context.repo.repo, issue_number: prNumber, labels: [sizeLabel] - }); - - enforce-xl-justification: - name: Enforce XL PR Justification - runs-on: ubuntu-slim - if: github.event.workflow_run.conclusion == 'success' - steps: - - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 - - - name: Download artifact - uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7 - with: - name: pr-size-label - path: pr-size/ - github-token: ${{ secrets.GITHUB_TOKEN }} - run-id: ${{ github.event.workflow_run.id }} - - - name: Read PR number and check for XL justification - id: check - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8 - with: - script: | - const fs = require('fs'); - const prNumber = parseInt(fs.readFileSync('pr-size/pr-number.txt', 'utf8').trim()); - const sizeLabel = fs.readFileSync('pr-size/label.txt', 'utf8').trim().replace(/"/g, ''); - - console.log('PR Number:', prNumber); - console.log('Size Label:', sizeLabel); - - const pr = await github.rest.pulls.get({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: prNumber - }); - - const hasXLLabel = sizeLabel === 'size/XL'; - const prBody = pr.data.body || ''; - const hasJustification = /##\s*Large PR Justification/i.test(prBody); - - console.log('Has XL label:', hasXLLabel); - console.log('Has justification:', hasJustification); - - return { - prNumber: prNumber, - hasXLLabel: hasXLLabel, - hasJustification: hasJustification, - needsEnforcement: hasXLLabel && !hasJustification, - shouldDismiss: (hasXLLabel && hasJustification) || !hasXLLabel - }; - - - name: Request changes if no justification - if: fromJSON(steps.check.outputs.result).needsEnforcement - continue-on-error: true - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8 - env: - RESULT_JSON: ${{ steps.check.outputs.result }} - with: - script: | - const result = JSON.parse(process.env.RESULT_JSON); - const prNumber = result.prNumber; - - // Check if we already have a review requesting changes - const reviews = await github.rest.pulls.listReviews({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: prNumber - }); - - const botReview = reviews.data.find(review => - review.user.login === 'github-actions[bot]' && - review.state === 'CHANGES_REQUESTED' - ); - - if (botReview) { - console.log('Already requested changes in review:', botReview.id); - return; - } - - // Read the message template from file - const fs = require('fs'); - const template = fs.readFileSync('.github/workflows/pr-size-justification-template.md', 'utf8'); - const contributingLink = `https://github.com/${context.repo.owner}/${context.repo.repo}/blob/main/CONTRIBUTING.md#code-quality-expectations`; - const message = template.replace('CONTRIBUTING_LINK', contributingLink); - - // Request changes with explanation - await github.rest.pulls.createReview({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: prNumber, - event: 'REQUEST_CHANGES', - body: message - }); - - console.log('Created review requesting changes for PR #' + prNumber); - - - name: Dismiss review if justification added - if: fromJSON(steps.check.outputs.result).shouldDismiss - continue-on-error: true - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8 - env: - RESULT_JSON: ${{ steps.check.outputs.result }} - with: - script: | - const result = JSON.parse(process.env.RESULT_JSON); - const prNumber = result.prNumber; - - // Find our previous review requesting changes - const reviews = await github.rest.pulls.listReviews({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: prNumber - }); - - const botReview = reviews.data.find(review => - review.user.login === 'github-actions[bot]' && - review.state === 'CHANGES_REQUESTED' - ); - - if (botReview) { - const dismissMessage = result.hasXLLabel - ? 'Large PR justification has been provided. Thank you!' - : 'PR size has been reduced below the XL threshold. Thank you for splitting this up!'; - - await github.rest.pulls.dismissReview({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: prNumber, - review_id: botReview.id, - message: dismissMessage - }); - - console.log('Dismissed previous review:', botReview.id); - - // Add a comment confirming unblock - const commentBody = result.hasXLLabel - ? '✅ Large PR justification has been provided. The size review has been dismissed and this PR can now proceed with normal review.' - : '✅ PR size has been reduced below the XL threshold. The size review has been dismissed and this PR can now proceed with normal review. Thank you for splitting this up!'; - - await github.rest.issues.createComment({ - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: prNumber, - body: commentBody - }); - } else { - console.log('No previous blocking review found to dismiss'); - } \ No newline at end of file + }); \ No newline at end of file diff --git a/.github/workflows/release-notes.yml b/.github/workflows/release-notes.yml index 97d360f854..02771b2cd0 100644 --- a/.github/workflows/release-notes.yml +++ b/.github/workflows/release-notes.yml @@ -55,7 +55,7 @@ jobs: SLACK_WEBHOOK: ${{ secrets.SLACK_TOOLHIVE_RELEASE_WEBHOOK_URL }} steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: # Always check out the default branch, never the workflow_run head ref. # Checking out an event-supplied ref in a secrets-bearing workflow_run @@ -98,7 +98,7 @@ jobs: - name: Generate release notes with Claude id: claude - uses: anthropics/claude-code-action@f1bd27ca5b54584506e40e17884d90bdaaa1a9b3 # v1 + uses: anthropics/claude-code-action@be7b93b1907a4abad570368f3c74b6fe3807510b # v1 env: # gh CLI used by the skill (read-only PR/issue lookups) authenticates # via GH_TOKEN. Read access is sufficient — publishing happens in a diff --git a/.github/workflows/releaser.yml b/.github/workflows/releaser.yml index cb674b7966..705ced0bae 100644 --- a/.github/workflows/releaser.yml +++ b/.github/workflows/releaser.yml @@ -35,7 +35,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Verify tag matches VERSION file run: | @@ -74,7 +74,7 @@ jobs: tree-state: ${{ steps.ldflags.outputs.tree-state }} steps: - id: checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: fetch-depth: 0 - id: ldflags @@ -95,7 +95,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 with: fetch-depth: 0 diff --git a/.github/workflows/renovate-config-validation.yml b/.github/workflows/renovate-config-validation.yml index 0f3be4284e..94ffc02753 100644 --- a/.github/workflows/renovate-config-validation.yml +++ b/.github/workflows/renovate-config-validation.yml @@ -22,7 +22,7 @@ jobs: name: Validate Renovate Configuration runs-on: ubuntu-slim steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Verify configuration syntax run: | diff --git a/.github/workflows/security-scan.yml b/.github/workflows/security-scan.yml index b65e0bc4a4..2dfd880e9f 100644 --- a/.github/workflows/security-scan.yml +++ b/.github/workflows/security-scan.yml @@ -21,7 +21,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Run Grype vulnerability scanner id: grype-scan @@ -32,7 +32,7 @@ jobs: fail-build: false - name: Upload Grype scan results to GitHub Security tab - uses: github/codeql-action/upload-sarif@54f647b7e1bb85c95cddabcd46b0c578ec92bc1a # v4 + uses: github/codeql-action/upload-sarif@f205ea1c3313d32999d8d6a48b4f6530d4437b38 # v4 if: always() with: sarif_file: ${{ steps.grype-scan.outputs.sarif }} @@ -43,10 +43,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Run govulncheck - uses: golang/govulncheck-action@b625fbe08f3bccbe446d94fbf87fcc875a4f50ee # v1 + uses: golang/govulncheck-action@032d45514ae346b1db93c04b0c90b841c370344f # v1 with: go-version-input: 'stable' go-package: ./... diff --git a/.github/workflows/skills-build-and-publish.yml b/.github/workflows/skills-build-and-publish.yml index d8cf29f831..a08b75fc5f 100644 --- a/.github/workflows/skills-build-and-publish.yml +++ b/.github/workflows/skills-build-and-publish.yml @@ -42,7 +42,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 diff --git a/.github/workflows/spellcheck.yml b/.github/workflows/spellcheck.yml index b51c9b26f2..37a0bc0abf 100644 --- a/.github/workflows/spellcheck.yml +++ b/.github/workflows/spellcheck.yml @@ -12,7 +12,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout Code - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Codespell uses: codespell-project/actions-codespell@406322ec52dd7b488e48c1c4b82e2a8b3a1bf630 # v2 with: diff --git a/.github/workflows/test-e2e-lifecycle.yml b/.github/workflows/test-e2e-lifecycle.yml index 18c9925ad1..963043f5a7 100644 --- a/.github/workflows/test-e2e-lifecycle.yml +++ b/.github/workflows/test-e2e-lifecycle.yml @@ -8,6 +8,11 @@ on: - 'cmd/thv-operator/**' - 'pkg/**' - 'test/e2e/thv-operator/**' + # This job is the only one that runs the Ginkgo suite under + # test/e2e/thv-operator, and that suite sources every backend container + # image from test/e2e/images. A change there alters what these tests + # actually deploy, so it has to trigger them. + - 'test/e2e/images/**' - '.github/workflows/test-e2e-lifecycle.yml' permissions: @@ -19,7 +24,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 env: - YARDSTICK_IMAGE: ghcr.io/stackloklabs/yardstick/yardstick-server:1.1.1 + YARDSTICK_IMAGE: ghcr.io/stackloklabs/yardstick/yardstick-server:1.2.0 defaults: run: shell: bash @@ -34,7 +39,7 @@ jobs: steps: - name: Checkout code - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Disable containerd image store # Workaround for https://github.com/kubernetes-sigs/kind/issues/3795 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 884a7a9bd2..127f4b29ae 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -15,7 +15,7 @@ jobs: os: [ubuntu-latest] fail-fast: false steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - name: Set up Go uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 @@ -74,7 +74,7 @@ jobs: - name: Upload coverage to Coveralls if: startsWith(matrix.os, 'ubuntu') - uses: coverallsapp/github-action@5cbfd81b66ca5d10c19b062c04de0199c215fb6e # v2 + uses: coverallsapp/github-action@8d6379e14d29928660c4ba802d8e85393440b329 # v2 with: file: coverage/coverage.out fail-on-error: false diff --git a/.github/workflows/verify-docgen.yml b/.github/workflows/verify-docgen.yml index 1f28c9e7c6..b28d6a3e19 100644 --- a/.github/workflows/verify-docgen.yml +++ b/.github/workflows/verify-docgen.yml @@ -11,10 +11,13 @@ jobs: contents: read steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 with: go-version: 'stable' - name: Install swag + # Pinned to match go.mod and Taskfile.yml — see the comment on the + # swagger-install task for why this alone doesn't make the check + # fully deterministic. run: go install github.com/swaggo/swag/v2/cmd/swag@v2.0.0-rc5 - run: ./cmd/help/verify.sh diff --git a/.github/workflows/verify-gen.yml b/.github/workflows/verify-gen.yml index e92e6a11a5..cd952d3bf2 100644 --- a/.github/workflows/verify-gen.yml +++ b/.github/workflows/verify-gen.yml @@ -12,7 +12,7 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 - uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6 id: setup-go with: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5ae7bd7b00..23ebb8e983 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -33,7 +33,7 @@ report unacceptable behavior to If you think you have found a security vulnerability in ToolHive please DO NOT disclose it publicly until we've had a chance to fix it. Please don't report security vulnerabilities using GitHub issues; instead, please follow this -[process](https://github.com/stacklok/toolhive/blob/main/SECURITY.md) +[process](https://github.com/stacklok/toolhive/blob/main/SECURITY.md). ## How to contribute @@ -72,8 +72,8 @@ items first to build context. If you'd like to work on an existing issue: -1. Leave a comment saying "I'd like to work on this" -2. Wait for a team member to assign you before starting work +1. Leave a comment saying "I'd like to work on this". +2. Wait for a team member to assign you before starting work. This helps us avoid situations where multiple people work on the same thing. If you create an issue with the intent to implement it yourself, mention that diff --git a/FORK.md b/FORK.md index 4c0f626e2f..35ce75c383 100644 --- a/FORK.md +++ b/FORK.md @@ -4,126 +4,154 @@ fork patches we carry. Upstream releases are brought in with a merge commit (`Merge upstream vX.Y.Z into dp-stable`), never a rebase. -Merge base: upstream **v0.39.0** (`4ebb610f2`), merged in `75ba8d1d5` (PAIINFRA-233). +Merge base: upstream **v0.42.1** (`06940e81d`), merged in `3f02521c8` (PAIINFRA-209). -## Changes that live only in a merge commit - -Some fork behaviour was re-applied by hand while resolving an upstream merge, so it exists in no -standalone commit. Replaying the fork's commits onto a fresh upstream tag (instead of merging) would -silently lose it. Port these explicitly. +Goal: carry as few patches as possible. Each row below says what the patch does, why DocPlanner +needs it, and when it can go. Check the "Drop when" column on every upstream merge. -### `74e7a003a` — Merge upstream v0.23.1 +## Patch inventory -Upstream moved the vMCP server wiring from `cmd/vmcp/app/commands.go` to `pkg/vmcp/cli`. The fork's -changes to the old file were re-applied in the new package: +Commits are the fork commits that introduced the behaviour. Where a row says "merge", the behaviour +was re-applied by hand in a merge commit (see [Changes that live only in a merge +commit](#changes-that-live-only-in-a-merge-commit)), so replaying the commits alone loses it. -| Behaviour | Where it lives now | Original commit | -|---|---|---| -| Backend init timeout from `failureHandling.healthCheckTimeout`, falling back to `timeouts.default` | `pkg/vmcp/cli/serve_timeouts.go` (`backendInitTimeoutFromConfig`) | `5feb1b1c9` | -| Backend request timeouts (`timeouts.default` / `timeouts.perWorkload`) and the HTTP write timeout derived from them | `pkg/vmcp/cli/serve_timeouts.go` (`backendRequestTimeoutsFromConfig`, `serverWriteTimeoutFromConfig`) | `3533730ce` | -| Session TTL passed to the server | `pkg/vmcp/cli/serve.go` (`SessionTTL`) | `d104dfc9f` | +### Operator -Tests: `pkg/vmcp/cli/serve_timeouts_test.go`. - -The same merge also carried a groupRef transition shim (string-form `MCPServer.spec.groupRef`, -optional `VirtualMCPServer.spec.groupRef`). It was removed in `d0a48cbe3`; nothing to port. +| Change | Why | Commit(s) | Drop when | +|---|---|---|---| +| `GenerationChangedPredicate` on every primary watch; vMCP `MaxConcurrentReconciles=3` | Status-only updates started full reconciles, and one slow vMCP blocked the others | `2c68591f2` | Upstream filters status-only events | +| Follow-ups the predicate needs: MCPRegistry keeps reconciling after adding its finalizer, MCPRemoteProxy requeues and watches MCPGroup, EmbeddingServer stays Failed on an invalid PodTemplateSpec, MCPServer reconciles on annotation changes (restart annotation) | The predicate hides these changes | `c7660fbb4`, `1960086e7`, `b3b0da4b1`, `198b61950` | Together with the predicate | +| vMCP env compared with `equality.Semantic`, labels and annotations as a subset, status setters skip unchanged values | False diffs caused an endless vMCP reconcile loop | `8071da52d` | Upstream compares semantically | +| vMCP Service/Deployment AlreadyExists and Conflict errors requeue | Create/update races marked the vMCP as failed | `0532bd27f` | Upstream requeues on these errors | +| vMCP drift check also compares nodeSelector, tolerations and affinity; StatefulSet-era objects of a migrated MCPServer are deleted on reconcile | Scheduling changes did not roll the vMCP; old StatefulSets and Services stayed after the migration | `3a8e50139` | Upstream checks scheduling drift (cleanup part: once no cluster has StatefulSet-era objects) | +| Explicit `spec.config.backends` merged with the discovered backends (explicit wins on the same name) | Hubs declare remote backends there; upstream overwrites the list | `8e474c8b2` | Upstream merges static backends | +| MCPServer Service reconciled before the Deployment; `status.URL` follows port changes; `targetPort` drift detected | A port change left vMCP routing to a stale port | `5df636980` | Upstream updates the URL and targetPort | +| RunConfig checksum taken from the ConfigMap writer | A stale cache read meant config changes did not roll the pods | `e88959899` | Upstream stops reading the checksum back | +| Proxyrunner and vMCP pods take nodeSelector, tolerations and affinity from `spec.podTemplateSpec`; operator PodDisruptionBudget | Proxies must run on the backend's node pool; HA for the operator | `5b6448a0d` | Upstream schedules the proxyrunner from the template | +| Proxyrunner deletes a stuck old-revision pod when a StatefulSet rollout stalls (`pods: delete` in the roles) | Rollouts hung on a crash-looping old pod | `baa5927f2`, `9c8b62f61` | Upstream handles stalled rollouts | +| Operator ClusterRole and RoleBinding names from `operator.rbacName` | The prod and stg operators run in the same cluster | `808a8a693` | Upstream templates the names | +| `sessionStorage.usernameRef`: Redis ACL username for MCPServer, MCPRemoteProxy and vMCP | Our Dragonfly uses ACL users, not only a password | `8d09330f5`, merge v0.28.3 | Upstream adds a username field | +| MCPOIDCConfig `allowedClientIds`, optional `oidcConfigRef.audience`, `client_id` fallback in the token validator | Cognito `client_credentials` tokens have no `aud` | `6d9eee4c1`, `78034a265`, `106deb1b5` | Upstream validates by client ID | +| `MCPServer.spec.secrets` list key is `name` + `key` | One Secret supplies several keys | `b736c5bbb` | Upstream changes the list key | + +### Auth, auth server, authz, audit + +| Change | Why | Commit(s) | Drop when | +|---|---|---|---| +| 401 with `WWW-Authenticate` when `authz.primaryUpstreamProvider` has no token (empty provider names ignored) | Clients log in again on 401, not on 403 | `44b70f6a5`, `f387a57f2` | Upstream returns 401 for a missing upstream token | +| `upstreamInject` with `providerName: caller` forwards the already-validated caller bearer token to the backend | Lets a backend authenticate the hub caller with the same token, without the embedded auth server | `c333920ac` | Upstream supports forwarding the caller token | +| DCR and CIMD without `scope` default to `scopesSupported` (RFC 7591) | Clients that register without `scope` got no scopes | `16c2fba5e`, `7019ea20c` | Upstream defaults the scope (or `baselineClientScopes` gives the same result for our clients) | +| Upstream-token writes use WATCH/MULTI instead of Lua | Our Dragonfly does not run the Lua scripts upstream uses | `907d7040d`, merge v0.28.3 | Upstream stops using Lua there, or we leave Dragonfly | +| Audit user name falls back to the `username` claim | Cognito access tokens have no `name` or `email` | `a40493885` | Upstream reads `username` | +| Client JSON-RPC responses (replies to elicitation, sampling, roots) pass through authz and the validating webhook. Only a body with no `method` member and a `result` or `error` counts as a reply (`ParsedMCPRequest.IsClientResponse`) | Upstream answers them with 400, so server-to-client requests never complete | `f22ee7fa6`, v0.42.1 follow-up | Upstream parses responses | -### Merge upstream v0.28.3 +### Proxy runner and transport -Upstream moved the auth-server and session Redis clients onto toolhive-core's `redis` package and -renamed `address` to `addr`. Fork behaviour re-applied while resolving: +| Change | Why | Commit(s) | Drop when | +|---|---|---|---| +| `TOOLHIVE_PROXY_SESSION_TTL` sets the proxy session TTL | No CRD field exists; we set 24 h (docplanner-helm `mcpservers.yaml`) to match vMCP sessions | `3a0b285c6`, `d104dfc9f` | Upstream adds a TTL field | +| Proxy tuning env from `spec.env` mirrored into the runner process; on Kubernetes the proxy keeps running after repeated backend health failures (`pkg/runner/proxy_env.go`) | The env reached only the backend container, so the proxy kept a 2 h TTL and backends idle for 2 h answered 404 | `48c702ff2` | Upstream exposes these settings | +| Shallow `/health` on the transparent proxy | The proxyrunner health check pinged the backend, so a slow backend failed the proxy's own probe | `5d7663669` | Upstream separates liveness from backend health | +| The Host header stays the configured target host when `backend_url` pins a pod IP; any other Host is rewritten to the URL host, as upstream does | Backends with strict host validation rejected the pod IP | `ce851b471`, v0.42.1 follow-up | Upstream keeps the target Host | +| HTTP 400 "no valid session ID" or "expired session ID" counts as a lost session (re-initialize and replay); only the first 4 KiB of the body are read | TypeScript SDK servers and mcp-for-argocd answer 400, not 404 | `8abac892e`, `16ffb4874`, merge v0.42.1 | Upstream treats these as session loss | +| Streamable proxy answers a timeout or failure with a JSON-RPC error, not HTTP 504/500 | vMCP got a transport error instead of a call error | `cce5872bd` | Upstream returns JSON-RPC errors | +| `ValidatingCache.Remove` and `EvictIdle`; entries closed outside the lock | The vMCP session cache needs eviction and close | `036974570`, `51de52641` | Upstream adds eviction | -| Behaviour | Where it lives now | Original commit | -|---|---|---| -| Upstream-token writes use WATCH/MULTI instead of Lua (Dragonfly), with upstream's #5092 index-TTL rules | `pkg/authserver/storage/redis.go` (`queueUpstreamIndexTTL`) | `907d7040d` | -| Session Redis ACL username from `sessionStorage.usernameRef` passed to the toolhive-core client | `pkg/runner/runner.go`, `pkg/vmcp/server/server.go` (`buildSessionDataStorage`, re-ported in `bafe4d46a`) | `8d09330f5` | -| Redis credential env injected once, after `resourceOverrides` (upstream's password-only builder removed) | `cmd/thv-operator/controllers/mcpserver_controller.go` (`buildSessionRedisCredentialEnvVars`) | `8d09330f5` | -| `TOOLHIVE_PROXY_SESSION_TTL` is the fallback when RunConfig `session_ttl` is unset | `pkg/runner/runner.go`, proxy constructors | `d104dfc9f` | -| vMCP `--session-ttl` flag wins over the config-file `sessionTTL` | `pkg/vmcp/cli/serve_timeouts.go` (`sessionTTLFromConfig`) | `d104dfc9f` | +### Virtual MCP server -### Merge upstream v0.29.3 +| Change | Why | Commit(s) | Drop when | +|---|---|---|---| +| Backend init timeout from `failureHandling.healthCheckTimeout`, else `timeouts.default` (`serve_timeouts.go`) | The configured timeouts were not wired to backend initialization | `5feb1b1c9`, merge v0.23.1 | Upstream makes it configurable | +| `timeouts.default` / `timeouts.perWorkload` applied to each backend call (keyed by `WorkloadID`, 90 s fallback), on the Legacy and the Modern client; the HTTP write timeout covers the longest of them | Upstream uses a fixed 30 s per call | `3533730ce`, `4cf28e14c`, `703e8e849`, merge v0.42.1 | Upstream uses the configured timeouts (#6411, v0.45) | +| Session TTL from the config file, the `--session-ttl` flag winning | Session lifetime must match the 24 h of the proxies | `d104dfc9f` | Upstream wires `sessionTTL` | +| Session Redis ACL username and HTTP write timeout wired on the Serve path | See `usernameRef` above; long backend calls outlived the 30 s write timeout | `bafe4d46a` | With the rows above | +| Cached sessions closed on DELETE and after one session TTL idle on this replica (`FactoryConfig.IdleTTL`) | Upstream evicts only on the next Get, which never comes after a DELETE, so backend connections leaked | `036974570`, `51de52641`, merge v0.42.1 | Upstream closes on Terminate | +| Health probes between vMCPs neither create sessions nor list capabilities; the probe header never sets the `healthcontext` marker; a probe may only initialize, ping and list (else 403) | A hub probing peer vMCPs opened a session with backend connections per probe; the header must not switch outgoing auth to service credentials | `46dd7a465`, `37d40cf78`, `43623c120` | Upstream has a lightweight peer probe | +| Health checks of Kubernetes-managed MCPServer backends use the reconciled workload status; others keep the MCP probe | Each periodic check opened a full MCP session to every backend | `5d7663669` | Upstream reads workload status | +| No `structuredContent` synthesized for tools without an output schema, on errors too | Strict clients (Cursor) reject a result whose structured content does not match the schema | `a22a3fb96` | Upstream stops synthesizing it | +| Backend call errors wrapped with their class (`wrapBackendError`) on call, read and prompt | A backend timeout reached the client as a generic error, so the timeout chain could not be told apart | `cce5872bd` | Upstream classifies backend errors on every path | +| Capability refresh: a backend that crosses the advertise boundary (`health.ShouldAdvertise`) or a registry change triggers each local session's `list_changed` sink (15 s debounce, 60 s cooldown) | Upstream resyncs only when a connected backend sends `list_changed`; a backend that comes back or joins sends nothing | `c4984e59c`, merge v0.42.1 | Upstream reacts to health and registry changes | +| Partial capabilities: a query that misses a backend is marked partial and cached at most 5 s; a `tools/call` for a tool the session lacks tops the store up (own caller, not optimizer, add-only); `capabilityQueryTimeout` 10 s | Upstream caches a partial view 30 s and fills a store once, so calls failed with "tool not found" (PAIINFRA-343) | `d6d76710c`, `e2393c8d7`, `96a25efd3` | Upstream fixes #5786; replace the timeout with #6411 (v0.45) | +| A GET listen stream for a session bound to another identity gets 403 (`listenStreamGuardMiddleware`) | The SDK opens the stream for anyone with the session ID, on any pod, and routes the session's elicitation and sampling requests to it | merge v0.42.1 (replaces `ab1b5b667`, `fccda6842`, `854448185`) | Upstream checks the caller on GET | + +### CI, tests, dependencies + +| Change | Why | Commit(s) | Drop when | +|---|---|---|---| +| Jobs on `ubuntu-latest`; golangci-lint v2.14.0 with `only-new-issues` | The public repo gets no large runners; upstream's lint debt must not block fork PRs | `b387fe750` | Never (fork CI) | +| Lifecycle e2e with 4 ginkgo processes; optimizer e2e timeout 90 s; managed-backend health assertion skipped | Smaller runners; managed backends are checked from workload status | `22368126a`, `05373f634`, `560253e96` | With the runner and health rows | +| Module versions above upstream (x/net, x/crypto, x/sys, x/mod, x/text, x/tools, x/sync, x/term, grpc, zipkin) | govulncheck findings | `fd8d45c58`, `f90409530`, `b3929efa7` | Upstream reaches the same versions | -Upstream replaced HMAC token/subject session binding with an `(iss, sub)` identity binding -(#5306/#5378). The fork's subject binding (`e492cbfad`) is dropped; the binding is upstream's. -Fork behaviour re-applied on top (the vMCP session rows were dropped or re-ported in the v0.32.0 -merge, see below): +## Notes -| Behaviour | Where it lives now | Original commit | -|---|---|---| -| DCR without `scope` defaults to the server's supported scopes | `pkg/authserver/server/registration/dcr.go` | `16c2fba5e` | -| Digest-pinned image pulls go through the Docker daemon (upstream #5978, after v0.29.3) | `pkg/container/images/registry.go` | `09c117c56` | +**Capability refresh scope.** Only sessions registered on this pod (`refreshSinks`) are refreshed: +sessions a pod serves only through rehydration (registered on another pod) keep their stores until +the client lists again or reconnects. Resources and prompts resync add-only (upstream limitation). +Sinks of ended sessions are dropped at the next refresh. `startCapabilityRefresh` runs only from +`Server.Start`. -`VMCP_SESSION_HMAC_SECRET` is no longer used (upstream logs and ignores it). Sessions stored with the -old token-hash key are invalidated on read, so clients re-initialize once after the upgrade. +**Partial capabilities.** Upstream fills a session's tool store once per pod and caches a partial +aggregation for 30 s (#5786). Since #6216 the core also checks calls against the cached advertised +view, so the 5 s partial TTL matters on both paths. With a 3 s `capabilityQueryTimeout` about a +quarter of the production hub aggregations were partial on 2026-10-15. + +**Sessions across pods (v0.42.1).** mcpcompat registers a session (and runs `OnRegisterSession`) +only while it serves `initialize`, before the client gets the session ID. Any other pod rehydrates +the session without the hook, and go-sdk refuses a second `initialize`. So the binding can no longer +be rewritten by a GET or a repeated initialize (`TestCrossPodSessionRebind`), and a call can no +longer arrive before the binding exists. mcpcompat runs before-hooks for tools only, so a pod that +rehydrates a session serves no resources or prompts for it (upstream gap; `TestCrossPodRehydration` +covers tools only). Accepted for PAIINFRA-209: in Datadog APM from 2026-09-16 to 2026-10-16 the hub +backends got no `resources/read`, `resources/templates/list`, `prompts/list` or `prompts/get`; every +`resources/list` was the hubs' own capability discovery. -### `75ba8d1d5` — Merge upstream v0.39.0 +## Changes that live only in a merge commit -Fork behaviour kept while resolving: +Some fork behaviour was re-applied by hand while resolving an upstream merge, so it exists in no +standalone commit. Replaying the fork's commits onto a fresh upstream tag (instead of merging) would +silently lose it. Port these explicitly. -| Behaviour | Where it lives now | Original commit | -|---|---|---| -| The MCPServer Service is updated before the Deployment, so Deployment drift cannot starve a port fix; labels and annotations are now merged, as upstream #5731 does | `cmd/thv-operator/controllers/mcpserver_controller.go` (`Reconcile`) | `5df636980`, `e88959899` | -| Explicit `spec.config.backends` (hub remote backends) are merged with the discovered backends, an explicit backend winning on the same name; upstream's `excludedBackends` filters only the discovered ones | `cmd/thv-operator/controllers/virtualmcpserver_vmcpconfig.go` (`mergeStaticBackends`) | `8e474c8b2` | +### `74e7a003a` — Merge upstream v0.23.1 -`PodTemplateSpecBuilder.SchedulingSpec()` stays: MCPServer and the vMCP drift check still use it. +Upstream moved the vMCP server wiring from `cmd/vmcp/app/commands.go` to `pkg/vmcp/cli`. The fork's +changes to the old file were re-applied in the new package: backend init timeout +(`backendInitTimeoutFromConfig`, `5feb1b1c9`), backend request and write timeouts +(`backendRequestTimeoutsFromConfig`, `serverWriteTimeoutFromConfig`, `3533730ce`) and the session TTL +(`d104dfc9f`), all in `pkg/vmcp/cli/serve_timeouts.go`. -## Fork behaviour re-ported after the v0.32.0 merge +### Merge upstream v0.28.3 -In v0.32.0 every vMCP caller goes through `server.Serve`: tool, resource and prompt calls go -through the core with a new backend client per call, and any pod restores a session from Redis -(#5650). The merge commit `38edac4e3` resolved every conflict to upstream's behaviour; the fork -behaviour still needed came back in standalone commits on top of it (plan: -`docs/fork/v0.32-merge-plan.md`, PAIINFRA-227, DocPlanner/toolhive#21). Unlike the sections -above, a replay keeps these, as long as each commit is replayed. +Upstream moved the auth-server and session Redis clients onto toolhive-core's `redis` package. +Re-applied: WATCH/MULTI upstream-token writes (`queueUpstreamIndexTTL`, `907d7040d`), the session +Redis ACL username (`buildSessionDataStorage`, `8d09330f5`), Redis credential env injected once after +`resourceOverrides` (`buildSessionRedisCredentialEnvVars`), `TOOLHIVE_PROXY_SESSION_TTL` as the +RunConfig fallback and `--session-ttl` winning over `sessionTTL` (`d104dfc9f`). -| Behaviour | Where it lives now | Commit | Original commit | -|---|---|---|---| -| Session Redis ACL username, `WithSessionIdleTTL`, HTTP write timeout and backend init timeout wired on the Serve path; mcp-go's idle sweep clears only this pod's SDK state (`sessionIDManagerResolver`) and `Stop` stops it | `pkg/vmcp/server/server.go` (`buildSessionDataStorage`, `Handler`), `pkg/vmcp/cli/serve.go`, `pkg/vmcp/server/derive.go` | `bafe4d46a`, `7f57d1f6f` | `8d09330f5`, `d104dfc9f`, `3533730ce`, `5feb1b1c9` | -| Session registration runs on a detached context (request values kept, 45 s timeout) | `pkg/vmcp/server/server.go` (`handleSessionRegistrationImpl`) | `0d62ea671` | `596f6ccf8` | -| Session-cardinality gauges (transport sessions, SDK sessions, in-flight registrations) | `pkg/vmcp/server/session_telemetry.go`, `client_sessions.go` | `0d62ea671`, `f387a57f2` | `560c1d023` | -| New sessions connect only to backends eligible for exposure | `pkg/vmcp/server/sessionmanager/session_manager.go` (`FactoryConfig.HealthStatusProvider`) | `ee875852d` | `currentEligibleBackends` | -| `timeouts.default` / `timeouts.perWorkload` applied to each backend call (keyed by `WorkloadID`, 90 s fallback); the HTTP write timeout covers the longest of them, fallback included | `pkg/vmcp/client/client.go` (`WithRequestTimeouts`, `DefaultBackendRequestTimeout`), `pkg/vmcp/cli/serve_timeouts.go` (`serverWriteTimeoutFromConfig`) | `4cf28e14c`, `703e8e849` | `3533730ce` | -| `Terminate` closes the node-local cached session; idle sweep closes cached sessions after the session TTL; `Close` runs outside the cache lock | `pkg/vmcp/server/sessionmanager/session_manager.go`, `pkg/cache/validating_cache.go` (`Remove`, `EvictIdle`) | `036974570`, `51de52641` | `596f6ccf8`, `c9ab3ccbc` | -| A call that arrives before the registration wrote the identity binding waits (one store read per step, 25 ms doubling to 500 ms, 10 s) instead of failing | `pkg/vmcp/server/serve_handlers.go` (`waitForRegisteredSession`, `waitWhileRegistrationPending`) | `c94881448`, `7d04b4957` | `d6a23fcda` | -| Health probes between vMCPs neither create sessions nor list capabilities. The probe header sets a package-private context key, never the `healthcontext` marker (which switches outgoing auth to the vMCP's service credentials), and a probe may only initialize, ping and list (else 403) | `pkg/vmcp/server/peer_probe.go` (`peerProbeMiddleware`) | `37d40cf78`, `43623c120` | `46dd7a465` | -| Resources filled on a pod that did not register the session | `pkg/vmcp/server/server.go` (`lazyInjectSessionResources`), `serve.go` hooks | `536c603b4` | `d6a23fcda` | -| Backend health or registry change clears the sessions' tool and resource stores and sends `list_changed` (15 s debounce, 60 s cooldown) | `pkg/vmcp/server/serve_refresh.go`, `refresh_scheduler.go` | `c4984e59c` | `560c1d023`, `35680e28c`, `e755493fe` | -| GET re-registration of a bound session attaches a matching caller and rejects any other; it never rewrites the binding | `pkg/vmcp/server/server.go` (`attachToBoundSession`) | `ab1b5b667` | — (new, plan D5) | -| A GET registration never creates a session: it waits for the creating pod's binding, then attaches or rejects | `pkg/vmcp/server/server.go` (`attachToBoundSession`) | `fccda6842` | — (new) | -| Every GET is checked before mcp-go sees it: 400 without `Mcp-Session-Id`, 404 for a missing, terminated or never-bound session, 403 for a different caller (mcp-go answers 200 whatever the hook does) | `pkg/vmcp/server/server.go` (`listenStreamGuardMiddleware`, `checkListenStreamCaller`) | `854448185` | — (new) | -| A capability query that misses a backend is marked partial (`AggregationMetadata.Partial`) and cached for at most 5 s (`partialCapabilityTTL`); a `tools/call` for a tool the session store lacks tops the store up from a fresh query (own caller only, not in optimizer mode, add-only) | `pkg/vmcp/aggregator/caching_aggregator.go`, `default_aggregator.go`, `pkg/vmcp/server/server.go` (`lazyInjectSessionTools`) | `d6d76710c` | — (new, PAIINFRA-343) | -| Missing credentials for `authz.primaryUpstreamProvider` return 401 with `WWW-Authenticate` (empty provider names ignored) | `pkg/auth/token.go` (`WithRequiredUpstreamProviders`), `pkg/vmcp/auth/factory/incoming.go` | `44b70f6a5`, `f387a57f2` | `c2cb472a2` | - -Tests: `pkg/vmcp/server/serve_*_test.go`, `session_registration_context_test.go`, -`session_telemetry_test.go`, `sessionmanager/session_cache_close_test.go`, -`pkg/vmcp/client/request_timeout_test.go`, `pkg/vmcp/cli/serve_timeouts_test.go`, -`pkg/cache/validating_cache_test.go`, `pkg/auth/token_test.go`. - -**Capability refresh scope.** The refresh (`c4984e59c`) clears only the sessions in -`activeClientSessions`: sessions registered on this pod by initialize or a GET listen stream. A -session that a pod only serves through POSTs (registered on another pod, tools filled by lazy -injection) keeps that pod's tool and resource stores until mcp-go's idle sweep clears them, so it -can see a stale tool set there. `startCapabilityRefresh` runs only from `Server.Start`; an embedder -that serves `Handler` without `Start` gets no refresh. +### `75ba8d1d5` — Merge upstream v0.39.0 -**Partial capabilities.** Upstream fills a session's tool store once per pod and caches a partial -aggregation for 30 s (#5786), so a backend that misses one query (here the fork's -`capabilityQueryTimeout`, or a 502 during a rollout) leaves calls failing with `tool not found`. -`d6d76710c` (and the follow-up that bounds its cost) is the interim fix. A `tools/list` still never re-derives a filled store: a session -whose first list was partial gets the rest on the next backend refresh or reconnect. Drop the patch -once upstream lands the same fix (PAIINFRA-343). `capabilityQueryTimeout` is 10 s: with 3 s, about a -quarter of the production hub aggregations were partial on 2026-10-15, because the nested `global` -vMCP backend can take longer than 3 s to initialize on a cold pod. Upstream has no such timeout; -replace it with the configured backend timeouts (#6411) when the fork reaches v0.45. - -**GET guard.** mcp-go v0.55 registers a session again when a GET listen stream reaches a pod that -does not hold it, and upstream's `CreateSession` would re-bind it to the caller. Upstream stops -this only in v0.40.0, as a side effect of #5729 (mcp-go replaced by toolhive-core `mcpcompat`). -mcp-go also answers the GET with 200 whatever the hook does, so `854448185` checks the GET in an -HTTP middleware first; the hook check stays as a fallback. Drop `ab1b5b667`, `fccda6842` and the -middleware of `854448185` once the fork is on v0.40 or later, if `TestCrossPodSessionRebind`, -`TestCrossPodGETDuringRegistration` and `TestListenStreamGuard` -(`pkg/vmcp/server/serve_crosspod_rebind_test.go`) pass without them. +Kept: Service before Deployment with merged labels and annotations (`5df636980`, `e88959899`) and +`mergeStaticBackends` (`8e474c8b2`). `PodTemplateSpecBuilder.SchedulingSpec()` stays for MCPServer and +the vMCP drift check. + +### `3f02521c8` — Merge upstream v0.42.1 + +v0.40.0 replaced `mark3labs/mcp-go` with toolhive-core `mcpcompat` over the official go-sdk (#5729). +mcpcompat has no idle sweep, session-ID-manager resolver, unregister hook, resource hooks or +`SendNotificationToSpecificClient`, so the fork code built on them was dropped or rewritten. The +last two rows (client replies, Host header) are in the follow-up commit "Narrow the client-reply and +Host header patches", on top of the merge: + +| Behaviour | How it was re-applied | +|---|---| +| Capability refresh | `serve_refresh.go` triggers upstream's per-session `list_changed` sink (`buildListChangedSink`) stored in `Server.refreshSinks`, instead of clearing stores and notifying through mcp-go | +| Listen-stream guard | One HTTP middleware (`listenStreamGuardMiddleware`) using the stored binding: 403 for another identity, everything else left to the SDK | +| Per-backend timeouts | `requestTimeout(WorkloadID)` in upstream's `newStreamableHTTPClient` and `buildModernHTTPClient` instead of the fixed 30 s | +| No synthesized `structuredContent` | `toolResultFromMCP` takes `hasOutputSchema` | +| Peer probe | `peerProbeMiddleware` placed inside upstream's auth middleware (upstream now wraps auth in audit) | +| Idle close of cached sessions | `FactoryConfig.IdleTTL` kept next to upstream's `BackendHealth` gate | +| Partial-capabilities top-up | The call hook passes the tool name; the owner check reuses `enforceSessionBinding` | +| Client JSON-RPC responses | Narrowed: a reply needs no `method` member and a `result` or `error` (a webhook patch that removes `method` must fail closed, upstream test); one `IsClientResponse` check in authz and the validating webhook | +| Host header | `RoundTrip` keeps only the configured target host and rewrites any other value (upstream's host-injection regression test) | ## Dropped patches @@ -132,19 +160,29 @@ They stay in history; don't re-apply them on a replay. | Commit | Patch | Why it was dropped | |---|---|---| -| `a75f5027b` | Preserve OIDC config ref resource URL | Upstream has the same `ResourceURL` field and fallback since v0.23.1. No net diff. | -| `66be6ebe9` | Accept `best_effort` partial failure mode | Upstream accepts `best_effort` since v0.22.0 (#4865). The fork's extra `bestEffort` alias was removed; no DocPlanner config uses it. | -| `87be189f6` | Restore embedding config resolution | Upstream's `populateOptimizerEmbeddingService` resolves `embeddingServerRef` (with `Validate()` defaulting the optimizer). Upstream's version is used. | -| `690846dd8` | `InlineAuthzConfig.primaryUpstreamProvider` on `VirtualMCPServer` | Upstream added it in v0.27.1 and moved it to `spec.authServerConfig.primaryUpstreamProvider`; the inline location is still read (deprecated). Dropped in the v0.28.3 merge. | -| `4cca58224` | Prefer the caller's live identity on backend requests | Upstream #5335 (v0.28.0) does the same in `identityRoundTripper`. Dropped in the v0.28.3 merge. | -| `e492cbfad` | Bind sessions to the `sub` claim instead of the bearer token | Upstream #5378 (v0.29.x) binds to `(iss, sub)`. Dropped in the v0.29.3 merge. | -| `ee63a49a8`, `fd01fce88`, `f690bcd97`, `6ab79d0b9` | Session-owner advertise URL (`THV_VMCP_SESSION_OWNER_URL`, `POD_IP`), forwarding to the owner pod, claim of orphaned sessions, handover, dial fail-fast | v0.32.0 has no session owner: any pod restores a session from Redis (#5650). Dropped in the v0.32.0 merge. | -| `f7d033ff9`, `704a5c162`, `16ffb4874` (vMCP part), `c40943a83` | Reconnect expired backend sessions; live-session lookup and rebuild in the tool handlers | Each backend call uses a new client (`httpBackendClient`); the handlers were removed upstream (#5622). Dropped in the v0.32.0 merge. | -| `c9ab3ccbc` (stripe locks, `ReplaceSession`, `CompareAndSwap`), `560c1d023` (`CreatorIdentity`), `517f94007`, `e755493fe` (`HasBoundIdentity`) | Safe publication of refreshed sessions, rebuilt with the creator's identity | No session rebuild any more; the refresh clears and notifies instead (`c4984e59c`). Dropped in the v0.32.0 merge. | -| `517f94007` (`keepStoredBackendIDs`) | Stop token-less restores from killing bound sessions | The stored backend list no longer affects calls. Dropped in the v0.32.0 merge. | -| `b2dbce4b7` | Refresh on a backend `list_changed` notification | Covered by the health and registry refresh (`c4984e59c`). Dropped in the v0.32.0 merge. | -| `7a6c9a2f3` | Restore with the owner's live identity (`ownsStoredSession`) | Upstream #5650 restores with the request identity; `enforceSessionBinding` rejects other callers. Dropped in the v0.32.0 merge. | -| `c2cb472a2` (status map) | Upstream-token status map (`UpstreamCredentialStatus`, `UpstreamTokenStatuses`) | Upstream #5651 returns 401 when a refresh fails; the missing-token case is `44b70f6a5`. No Cedar policy read the status. Dropped in the v0.32.0 merge. | -| `8d09330f5` (alias) | Redis `address` alias of `addr` on `MCPExternalAuthConfig` (CRD, `EffectiveAddr`) and `LegacyAddress` in the auth-server RunConfig | All values and live objects use `addr` (PAIINFRA-219). Dropped in the v0.32.0 merge. | -| `5b6448a0d` (remote proxy part) | `MCPRemoteProxy.spec.podTemplateSpec`, applying only nodeSelector, tolerations and affinity | Upstream #5531 (v0.37.0) adds the same field with a full strategic-merge patch, a `PodTemplateValid` condition and drift detection by hash. No live MCPRemoteProxy exists. Dropped in the v0.39.0 merge. | -| `8e474c8b2` (env sort) | `sortEnvVarsByName` on the vMCP env, plus sorted workload and backend loops | Upstream already sorts each discovered and inline secret env group; the rest of the list has a fixed order. Dropping it changes the env order once, inside the rollout the new image causes anyway. Dropped in the v0.39.0 merge. | \ No newline at end of file +| `a75f5027b` | Preserve OIDC config ref resource URL | Upstream has the same field and fallback since v0.23.1. | +| `66be6ebe9` | Accept `best_effort` partial failure mode | Upstream accepts it since v0.22.0 (#4865). | +| `87be189f6` | Restore embedding config resolution | Upstream's `populateOptimizerEmbeddingService` resolves it. | +| `690846dd8` | `InlineAuthzConfig.primaryUpstreamProvider` on `VirtualMCPServer` | Upstream added it in v0.27.1. Dropped in the v0.28.3 merge. | +| `4cca58224` | Prefer the caller's live identity on backend requests | Upstream #5335 (v0.28.0). Dropped in the v0.28.3 merge. | +| `e492cbfad` | Bind sessions to the `sub` claim | Upstream #5378 binds to `(iss, sub)`. Dropped in the v0.29.3 merge. | +| `ee63a49a8`, `fd01fce88`, `f690bcd97`, `6ab79d0b9` | Session owner URL, forwarding to the owner pod, claim, handover, dial fail-fast | v0.32.0 restores a session on any pod (#5650). Dropped in the v0.32.0 merge. | +| `f7d033ff9`, `704a5c162`, `16ffb4874` (vMCP part), `c40943a83` | Reconnect expired backend sessions; live-session rebuild in the tool handlers | Each backend call uses a new client (#5622). Dropped in the v0.32.0 merge. | +| `c9ab3ccbc`, `560c1d023` (`CreatorIdentity`), `517f94007`, `e755493fe` | Safe publication of refreshed sessions | No session rebuild any more. Dropped in the v0.32.0 merge. | +| `b2dbce4b7` | Refresh on a backend `list_changed` | Upstream #5965/#5971 do it. Dropped in the v0.32.0 merge. | +| `7a6c9a2f3` | Restore with the owner's live identity | Upstream #5650. Dropped in the v0.32.0 merge. | +| `8d09330f5` (alias) | Redis `address` alias of `addr` | All values use `addr` (PAIINFRA-219). Dropped in the v0.32.0 merge. | +| `5b6448a0d` (remote proxy part) | `MCPRemoteProxy.spec.podTemplateSpec` | Upstream #5531 (v0.37.0). Dropped in the v0.39.0 merge. | +| `8e474c8b2` (env sort) | Sorted vMCP env, workload and backend loops | Upstream sorts each env group. Dropped in the v0.39.0 merge. | +| `09c117c56` | Digest-pinned image pulls through the Docker daemon | Upstream #5978 does the same. Dropped in the v0.42.1 merge. | +| `c2cb472a2` (what was left), `f387a57f2` (authz part) | Typed upstream-auth error in Cedar, 401 classification in authz middleware and response filter, list filtering failing on an authorizer error | `44b70f6a5` already answers 401 before authz runs, and upstream #5651 answers 401 on a failed refresh. Upstream now denies with a fixed 403 (#5841, #5850, #6066). Dropped in the v0.42.1 merge. | +| `497da99ae` | Default the sole allowed audience on every grant; require `resource` with several audiences | Upstream defaults the sole audience for `authorization_code`, and fosite carries it into refresh grants. Dropped in the v0.42.1 merge. | +| `9db565784`, `08d687ac6` | Structured token-endpoint error logs with refresh-failure classification | Upstream logs the error; `oauth_debug` could carry request fragments. Dropped in the v0.42.1 merge. | +| `8d09330f5` (validation) | Go validation of the auth-server Redis storage block | CEL rules and `controllerutil` cover it (only the `sentinelTls` with `addr` check is gone). Dropped in the v0.42.1 merge. | +| `ee875852d`, `afc71f865` | `health/exposure.go`: new sessions connect only to backends eligible for exposure | Upstream #6162 (`BackendHealth`, `health.ShouldAdvertise`); upstream also tries `Unknown` backends. Dropped in the v0.42.1 merge. | +| `ab1b5b667`, `fccda6842`, `854448185` | GET re-registration guard: attach or reject in the registration hook, wait for the creating pod, HTTP checks | mcpcompat never runs the hook outside `initialize`, so a GET cannot re-bind (`TestCrossPodSessionRebind`). The caller check lives on in `listenStreamGuardMiddleware`. Dropped in the v0.42.1 merge. | +| `bafe4d46a`, `7f57d1f6f` (SDK sweep) | `sessionIDManagerResolver`, `WithSessionIdleTTL`, `Shutdown` of the SDK sweeper | mcpcompat has no idle sweep. Dropped in the v0.42.1 merge. | +| `c94881448`, `7d04b4957` | Calls wait for the registration to write the binding | Registration finishes inside `initialize`. Dropped in the v0.42.1 merge. | +| `0d62ea671` (context), `560c1d023` / `0d62ea671` / `f387a57f2` (gauges) | Registration on a detached context; session-cardinality gauges | The hook runs before the initialize response; mcpcompat has no unregister hook to count sessions. Dropped in the v0.42.1 merge. | +| `536c603b4` | Resources filled on a pod that did not register the session | mcpcompat has no resource hooks (upstream gap, see Notes). Dropped in the v0.42.1 merge. | +| `525718fee`, `3d2af3d90`, `085637aa5`, `b387fe750` (swag pin) | e2e image pins, swag pin | Upstream pins them. Dropped in the v0.42.1 merge. | diff --git a/Taskfile.yml b/Taskfile.yml index 0cfc9735f8..62ce8127aa 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -13,11 +13,17 @@ tasks: - rm -rf docs/cli/* - go run cmd/help/main.go --dir docs/cli - swag init -g pkg/api/server.go --v3.1 -o docs/server --parseDependencyLevel 1 + # swag can non-deterministically double enum arrays for types from + # external modules (see cmd/help/dedupe-enums); normalize the output. + - go run ./cmd/help/dedupe-enums docs/server/docs.go docs/server/swagger.json docs/server/swagger.yaml - task: helm-docs swagger-install: desc: Install the swag tool for OpenAPI/Swagger generation cmds: + # Pinned to the version go.mod already requires, so @latest can't + # silently drift onto a newer/older swag release. + # Keep in sync with .github/workflows/verify-docgen.yml. - go install github.com/swaggo/swag/v2/cmd/swag@v2.0.0-rc5 helm-docs: @@ -220,6 +226,16 @@ tasks: deps: [build] env: THV_BINARY: "{{.PWD}}/bin/thv" + CONFORMANCE_TARGET: '{{.CONFORMANCE_TARGET | default "proxy"}}' + cmds: + - ./test/conformance/run-conformance.sh + + conformance-vmcp: + desc: Run the MCP conformance suite against a `thv vmcp serve` deployment (requires Docker + Node) + deps: [build] + env: + THV_BINARY: "{{.PWD}}/bin/thv" + CONFORMANCE_TARGET: vmcp cmds: - ./test/conformance/run-conformance.sh diff --git a/VERSION b/VERSION index 4ef2eb086f..339fd080af 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.39.0 +0.42.1 diff --git a/cmd/help/dedupe-enums/main.go b/cmd/help/dedupe-enums/main.go new file mode 100644 index 0000000000..655f7dddaa --- /dev/null +++ b/cmd/help/dedupe-enums/main.go @@ -0,0 +1,102 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Command dedupe-enums works around a swag v2.0.0-rc5 nondeterminism bug: +// loadExternalPackage (in swag's packages.go) reparses an external module's +// package without memoizing by import path, which can non-deterministically +// double the "enum"/"x-enum-varnames" arrays it emits for types defined +// outside this module (model.ArgumentType, model.Format, model.Icon from +// github.com/modelcontextprotocol/registry). When an array is exactly two +// identical halves back to back, this collapses it to one. +// +// Run after `swag init`, on each generated docs/server file: +// +// go run ./cmd/help/dedupe-enums docs/server/docs.go docs/server/swagger.json docs/server/swagger.yaml +package main + +import ( + "fmt" + "os" + "regexp" + "strings" +) + +// jsonArray only matches quoted-string enum values (swag emits every enum in +// docs/server today as strings); a doubled numeric/integer external enum +// would pass through undeduped. yamlArray matches any scalar, so it has no +// such gap. +var ( + jsonArray = regexp.MustCompile(`(?m)^([ \t]*"(?:enum|x-enum-varnames)": \[\n)((?:[ \t]*".*?",?\n)+?)([ \t]*\])`) + yamlArray = regexp.MustCompile(`(?m)^([ \t]*(?:enum|x-enum-varnames):\n)((?:[ \t]*- .+\n)+)`) +) + +func main() { + for _, path := range os.Args[1:] { + data, err := os.ReadFile(path) // #nosec G304 -- path is a CLI argument we control (Taskfile/verify.sh) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + + fixed := jsonArray.ReplaceAllStringFunc(string(data), dedupeJSONBlock) + fixed = yamlArray.ReplaceAllStringFunc(fixed, dedupeYAMLBlock) + + if fixed == string(data) { + continue + } + if err := os.WriteFile(path, []byte(fixed), 0o600); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + fmt.Println("deduped", path) + } +} + +func dedupeJSONBlock(match string) string { + g := jsonArray.FindStringSubmatch(match) + header, body, footer := g[1], g[2], g[3] + + lines := strings.Split(strings.TrimRight(body, "\n"), "\n") + trimmed := make([]string, len(lines)) + for i, l := range lines { + trimmed[i] = strings.TrimSuffix(l, ",") + } + + half, ok := firstHalfIfDoubled(trimmed) + if !ok { + return match + } + return header + strings.Join(half, ",\n") + "\n" + footer +} + +func dedupeYAMLBlock(match string) string { + g := yamlArray.FindStringSubmatch(match) + header, body := g[1], g[2] + + lines := strings.Split(strings.TrimRight(body, "\n"), "\n") + half, ok := firstHalfIfDoubled(lines) + if !ok { + return match + } + return header + strings.Join(half, "\n") + "\n" +} + +// firstHalfIfDoubled reports whether lines is exactly two identical halves +// back to back, returning the first half if so. It only catches this +// specific shape (contiguous, order-preserved duplication), not any 2x +// multiset — if swag's bug ever reordered values within the second copy, +// this wouldn't fire. That matches the upstream bug: it re-appends the same +// const list a second time, never reorders it. +func firstHalfIfDoubled(lines []string) ([]string, bool) { + n := len(lines) + if n == 0 || n%2 != 0 { + return nil, false + } + half := n / 2 + for i := 0; i < half; i++ { + if lines[i] != lines[half+i] { + return nil, false + } + } + return lines[:half], true +} diff --git a/cmd/help/dedupe-enums/main_test.go b/cmd/help/dedupe-enums/main_test.go new file mode 100644 index 0000000000..c2583e8e6e --- /dev/null +++ b/cmd/help/dedupe-enums/main_test.go @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "reflect" + "testing" +) + +func TestFirstHalfIfDoubled(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + lines []string + want []string + ok bool + }{ + { + name: "doubled", + lines: []string{"a", "b", "a", "b"}, + want: []string{"a", "b"}, + ok: true, + }, + { + name: "tripled is not a doubled shape", + lines: []string{"a", "b", "a", "b", "a", "b"}, + ok: false, + }, + { + name: "odd length", + lines: []string{"a", "b", "a"}, + ok: false, + }, + { + name: "single element", + lines: []string{"a"}, + ok: false, + }, + { + name: "empty", + lines: []string{}, + ok: false, + }, + { + name: "already deduped, idempotent", + lines: []string{"a", "b"}, + ok: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, ok := firstHalfIfDoubled(tt.lines) + if ok != tt.ok { + t.Fatalf("ok = %v, want %v", ok, tt.ok) + } + if ok && !reflect.DeepEqual(got, tt.want) { + t.Fatalf("got %v, want %v", got, tt.want) + } + }) + } +} diff --git a/cmd/help/verify.sh b/cmd/help/verify.sh index a40d4a99e2..4ed27ca59b 100755 --- a/cmd/help/verify.sh +++ b/cmd/help/verify.sh @@ -10,6 +10,9 @@ diff -Naur -I "^ date:" "$tmpdir" docs/cli/ api_tmpdir=$(mktemp -d) mkdir -p "$api_tmpdir/server" swag init -g pkg/api/server.go --v3.1 -o "$api_tmpdir/server" --parseDependencyLevel 1 +# swag can non-deterministically double enum arrays for types from external +# modules (see cmd/help/dedupe-enums); normalize before diffing. +go run ./cmd/help/dedupe-enums "$api_tmpdir/server/docs.go" "$api_tmpdir/server/swagger.json" "$api_tmpdir/server/swagger.yaml" # Exclude README.md from diff as it's manually maintained diff -Naur --exclude="README.md" "$api_tmpdir/server" docs/server/ diff --git a/cmd/thv-operator/api/v1alpha1/types.go b/cmd/thv-operator/api/v1alpha1/types.go index ae1258ac9b..f3d16e6e6f 100644 --- a/cmd/thv-operator/api/v1alpha1/types.go +++ b/cmd/thv-operator/api/v1alpha1/types.go @@ -47,7 +47,6 @@ type EmbeddingServerList struct { //+kubebuilder:resource:shortName=extauth;mcpextauth,categories=toolhive //+kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type` //+kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -//+kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` //+kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPExternalAuthConfig is the deprecated v1alpha1 version of the MCPExternalAuthConfig resource. @@ -105,7 +104,6 @@ type MCPGroupList struct { //+kubebuilder:resource:shortName=mcpoidc,categories=toolhive //+kubebuilder:printcolumn:name="Source",type=string,JSONPath=`.spec.type` //+kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -//+kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` //+kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPOIDCConfig is the deprecated v1alpha1 version of the MCPOIDCConfig resource. @@ -134,7 +132,6 @@ type MCPOIDCConfigList struct { //+kubebuilder:resource:shortName=authzcfg,categories=toolhive //+kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type` //+kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -//+kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` //+kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPAuthzConfig is the deprecated v1alpha1 version of the MCPAuthzConfig resource. @@ -318,7 +315,6 @@ type MCPTelemetryConfigList struct { //+kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true //+kubebuilder:resource:shortName=mwc,categories=toolhive //+kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -//+kubebuilder:printcolumn:name="References",type=string,JSONPath=`.status.referencingWorkloads` //+kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPWebhookConfig is the Schema for the mcpwebhookconfigs API. @@ -346,7 +342,6 @@ type MCPWebhookConfigList struct { //+kubebuilder:subresource:status //+kubebuilder:resource:shortName=tc;toolconfig,categories=toolhive //+kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -//+kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` //+kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPToolConfig is the deprecated v1alpha1 version of the MCPToolConfig resource. diff --git a/cmd/thv-operator/api/v1beta1/mcpauthzconfig_types.go b/cmd/thv-operator/api/v1beta1/mcpauthzconfig_types.go index 22e8820476..0a4360c7bb 100644 --- a/cmd/thv-operator/api/v1beta1/mcpauthzconfig_types.go +++ b/cmd/thv-operator/api/v1beta1/mcpauthzconfig_types.go @@ -58,21 +58,6 @@ type MCPAuthzConfigStatus struct { // ConfigHash is a hash of the current configuration for change detection // +optional ConfigHash string `json:"configHash,omitempty"` - - // ReferenceCount is the number of workloads referencing this config. - // +optional - ReferenceCount int32 `json:"referenceCount,omitempty"` - - // ReferencingWorkloads is a list of workload resources that reference this MCPAuthzConfig. - // Each entry identifies the workload by kind and name. The map key is the - // (kind, name) pair so two workloads of different kinds that share a name - // (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - // entries rather than colliding under merge-patch semantics. - // +listType=map - // +listMapKey=kind - // +listMapKey=name - // +optional - ReferencingWorkloads []WorkloadReference `json:"referencingWorkloads,omitempty"` } // +kubebuilder:object:root=true @@ -82,7 +67,6 @@ type MCPAuthzConfigStatus struct { // +kubebuilder:resource:shortName=authzcfg,categories=toolhive // +kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type` // +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -// +kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` // +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPAuthzConfig is the Schema for the mcpauthzconfigs API. diff --git a/cmd/thv-operator/api/v1beta1/mcpexternalauthconfig_types.go b/cmd/thv-operator/api/v1beta1/mcpexternalauthconfig_types.go index 63182d5cb8..2c3fd097e7 100644 --- a/cmd/thv-operator/api/v1beta1/mcpexternalauthconfig_types.go +++ b/cmd/thv-operator/api/v1beta1/mcpexternalauthconfig_types.go @@ -1422,21 +1422,6 @@ type MCPExternalAuthConfigStatus struct { // ConfigHash is a hash of the current configuration for change detection // +optional ConfigHash string `json:"configHash,omitempty"` - - // ReferenceCount is the number of workloads referencing this config. - // +optional - ReferenceCount int32 `json:"referenceCount,omitempty"` - - // ReferencingWorkloads is a list of workload resources that reference this MCPExternalAuthConfig. - // Each entry identifies the workload by kind and name. The map key is the - // (kind, name) pair so two workloads of different kinds that share a name - // (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - // entries rather than colliding under merge-patch semantics. - // +listType=map - // +listMapKey=kind - // +listMapKey=name - // +optional - ReferencingWorkloads []WorkloadReference `json:"referencingWorkloads,omitempty"` } // +kubebuilder:object:root=true @@ -1446,7 +1431,6 @@ type MCPExternalAuthConfigStatus struct { // +kubebuilder:resource:shortName=extauth;mcpextauth,categories=toolhive // +kubebuilder:printcolumn:name="Type",type=string,JSONPath=`.spec.type` // +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -// +kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` // +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPExternalAuthConfig is the Schema for the mcpexternalauthconfigs API. diff --git a/cmd/thv-operator/api/v1beta1/mcpoidcconfig_types.go b/cmd/thv-operator/api/v1beta1/mcpoidcconfig_types.go index 3839059a3e..90f58476d7 100644 --- a/cmd/thv-operator/api/v1beta1/mcpoidcconfig_types.go +++ b/cmd/thv-operator/api/v1beta1/mcpoidcconfig_types.go @@ -146,7 +146,7 @@ type InlineOIDCSharedConfig struct { // +optional ProtectedResourceAllowPrivateIP bool `json:"protectedResourceAllowPrivateIP"` - // InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing. + // InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing. // WARNING: This is insecure and should NEVER be used in production. // +kubebuilder:default=false // +optional @@ -189,21 +189,6 @@ type MCPOIDCConfigStatus struct { // ConfigHash is a hash of the current configuration for change detection // +optional ConfigHash string `json:"configHash,omitempty"` - - // ReferenceCount is the number of workloads referencing this config. - // +optional - ReferenceCount int32 `json:"referenceCount,omitempty"` - - // ReferencingWorkloads is a list of workload resources that reference this MCPOIDCConfig. - // Each entry identifies the workload by kind and name. The map key is the - // (kind, name) pair so two workloads of different kinds that share a name - // (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - // entries rather than colliding under merge-patch semantics. - // +listType=map - // +listMapKey=kind - // +listMapKey=name - // +optional - ReferencingWorkloads []WorkloadReference `json:"referencingWorkloads,omitempty"` } // +kubebuilder:object:root=true @@ -213,7 +198,6 @@ type MCPOIDCConfigStatus struct { // +kubebuilder:resource:shortName=mcpoidc,categories=toolhive // +kubebuilder:printcolumn:name="Source",type=string,JSONPath=`.spec.type` // +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -// +kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` // +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPOIDCConfig is the Schema for the mcpoidcconfigs API. diff --git a/cmd/thv-operator/api/v1beta1/mcptelemetryconfig_types.go b/cmd/thv-operator/api/v1beta1/mcptelemetryconfig_types.go index c35eb73de2..6286583184 100644 --- a/cmd/thv-operator/api/v1beta1/mcptelemetryconfig_types.go +++ b/cmd/thv-operator/api/v1beta1/mcptelemetryconfig_types.go @@ -128,12 +128,6 @@ type MCPTelemetryConfigStatus struct { // ConfigHash is a hash of the current configuration for change detection // +optional ConfigHash string `json:"configHash,omitempty"` - - // ReferencingWorkloads lists workloads that reference this MCPTelemetryConfig - // +listType=map - // +listMapKey=name - // +optional - ReferencingWorkloads []WorkloadReference `json:"referencingWorkloads,omitempty"` } // +kubebuilder:object:root=true diff --git a/cmd/thv-operator/api/v1beta1/mcpwebhookconfig_types.go b/cmd/thv-operator/api/v1beta1/mcpwebhookconfig_types.go index 579e327a27..88bb8c8013 100644 --- a/cmd/thv-operator/api/v1beta1/mcpwebhookconfig_types.go +++ b/cmd/thv-operator/api/v1beta1/mcpwebhookconfig_types.go @@ -103,11 +103,4 @@ type MCPWebhookConfigStatus struct { // ConfigHash is a hash of the spec, used for detecting changes // +optional ConfigHash string `json:"configHash,omitempty"` - - // ReferencingWorkloads is a list of workload resources that reference this MCPWebhookConfig. - // Each entry identifies the workload by kind and name. - // +listType=map - // +listMapKey=name - // +optional - ReferencingWorkloads []WorkloadReference `json:"referencingWorkloads,omitempty"` } diff --git a/cmd/thv-operator/api/v1beta1/toolconfig_types.go b/cmd/thv-operator/api/v1beta1/toolconfig_types.go index b28b37030a..ac709375f4 100644 --- a/cmd/thv-operator/api/v1beta1/toolconfig_types.go +++ b/cmd/thv-operator/api/v1beta1/toolconfig_types.go @@ -96,17 +96,6 @@ type MCPToolConfigStatus struct { // ConfigHash is a hash of the current configuration for change detection // +optional ConfigHash string `json:"configHash,omitempty"` - - // ReferenceCount is the number of workloads referencing this config. - // +optional - ReferenceCount int32 `json:"referenceCount,omitempty"` - - // ReferencingWorkloads is a list of workload resources that reference this MCPToolConfig. - // Each entry identifies the workload by kind and name. - // +listType=map - // +listMapKey=name - // +optional - ReferencingWorkloads []WorkloadReference `json:"referencingWorkloads,omitempty"` } // +kubebuilder:object:root=true @@ -115,7 +104,6 @@ type MCPToolConfigStatus struct { // +kubebuilder:metadata:labels=toolhive.stacklok.dev/auto-migrate-storage-version=true // +kubebuilder:resource:shortName=tc;toolconfig,categories=toolhive // +kubebuilder:printcolumn:name="Valid",type=string,JSONPath=`.status.conditions[?(@.type=='Valid')].status` -// +kubebuilder:printcolumn:name="References",type=integer,JSONPath=`.status.referenceCount` // +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` // MCPToolConfig is the Schema for the mcptoolconfigs API. diff --git a/cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go b/cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go index 856788d2cf..594e388d9e 100644 --- a/cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go +++ b/cmd/thv-operator/api/v1beta1/zz_generated.deepcopy.go @@ -804,11 +804,6 @@ func (in *MCPAuthzConfigStatus) DeepCopyInto(out *MCPAuthzConfigStatus) { (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ReferencingWorkloads != nil { - in, out := &in.ReferencingWorkloads, &out.ReferencingWorkloads - *out = make([]WorkloadReference, len(*in)) - copy(*out, *in) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPAuthzConfigStatus. @@ -945,11 +940,6 @@ func (in *MCPExternalAuthConfigStatus) DeepCopyInto(out *MCPExternalAuthConfigSt (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ReferencingWorkloads != nil { - in, out := &in.ReferencingWorkloads, &out.ReferencingWorkloads - *out = make([]WorkloadReference, len(*in)) - copy(*out, *in) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPExternalAuthConfigStatus. @@ -1202,11 +1192,6 @@ func (in *MCPOIDCConfigStatus) DeepCopyInto(out *MCPOIDCConfigStatus) { (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ReferencingWorkloads != nil { - in, out := &in.ReferencingWorkloads, &out.ReferencingWorkloads - *out = make([]WorkloadReference, len(*in)) - copy(*out, *in) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPOIDCConfigStatus. @@ -1948,11 +1933,6 @@ func (in *MCPTelemetryConfigStatus) DeepCopyInto(out *MCPTelemetryConfigStatus) (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ReferencingWorkloads != nil { - in, out := &in.ReferencingWorkloads, &out.ReferencingWorkloads - *out = make([]WorkloadReference, len(*in)) - copy(*out, *in) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPTelemetryConfigStatus. @@ -2110,11 +2090,6 @@ func (in *MCPToolConfigStatus) DeepCopyInto(out *MCPToolConfigStatus) { (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ReferencingWorkloads != nil { - in, out := &in.ReferencingWorkloads, &out.ReferencingWorkloads - *out = make([]WorkloadReference, len(*in)) - copy(*out, *in) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPToolConfigStatus. @@ -2166,11 +2141,6 @@ func (in *MCPWebhookConfigStatus) DeepCopyInto(out *MCPWebhookConfigStatus) { (*in)[i].DeepCopyInto(&(*out)[i]) } } - if in.ReferencingWorkloads != nil { - in, out := &in.ReferencingWorkloads, &out.ReferencingWorkloads - *out = make([]WorkloadReference, len(*in)) - copy(*out, *in) - } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new MCPWebhookConfigStatus. diff --git a/cmd/thv-operator/app/app.go b/cmd/thv-operator/app/app.go index 72b704120e..2aad8b882c 100644 --- a/cmd/thv-operator/app/app.go +++ b/cmd/thv-operator/app/app.go @@ -49,11 +49,11 @@ var ( setupLog = log.Log.WithName("setup") ) -// envEnableStorageVersionMigrator is the opt-in for the StorageVersionMigrator -// controller. The controller defaults to OFF in this release so the change can -// ship safely without functional impact. Set to "true" (or "1", "t") to enable. -// A follow-up release will flip the default to true alongside the helm chart -// surface and user docs. +// envEnableStorageVersionMigrator gates the StorageVersionMigrator controller. +// The binary itself defaults to OFF when the var is unset; the operator helm +// chart sets it to "true" by default, so chart-based installs run the migrator +// unless the operator explicitly opts out. Set to "true" (or "1", "t") to +// enable, "false" to disable. const envEnableStorageVersionMigrator = "TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR" func init() { @@ -197,10 +197,10 @@ func setupStorageVersionMigrator(mgr ctrl.Manager) error { } // isStorageVersionMigratorEnabled reports whether the StorageVersionMigrator -// controller should be registered. Defaults to false in this release — admins -// must explicitly opt in via TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR=true. -// An unparsable value returns an error so startup fails loudly rather than -// silently disabling the feature an admin asked to turn on. +// controller should be registered. Defaults to false when +// TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR is unset; the operator helm chart +// sets it to "true" by default. An unparsable value returns an error so startup +// fails loudly rather than silently disabling the feature an admin asked to turn on. func isStorageVersionMigratorEnabled() (bool, error) { value, found := os.LookupEnv(envEnableStorageVersionMigrator) if !found { diff --git a/cmd/thv-operator/controllers/config_controller_events_test.go b/cmd/thv-operator/controllers/config_controller_events_test.go index 56f8df8fdc..f1cc31a0f2 100644 --- a/cmd/thv-operator/controllers/config_controller_events_test.go +++ b/cmd/thv-operator/controllers/config_controller_events_test.go @@ -237,7 +237,7 @@ func TestMCPExternalAuthConfigReconciler_EmitsDeletionBlockedEvent(t *testing.T) } // TestMCPExternalAuthConfigReconciler_EmitsRecoveryOnSteadyStatePath covers the -// recovery transition that flows through updateReferencingWorkloads rather than +// recovery transition that flows through updateSteadyStateStatus rather than // handleConfigHashChange: a config whose stored hash already matches the current // spec (no hash change) but whose Valid condition is still False from an earlier // failure. Reconcile must emit a single Normal ConfigValid event. diff --git a/cmd/thv-operator/controllers/mcpauthzconfig_controller.go b/cmd/thv-operator/controllers/mcpauthzconfig_controller.go index 5b4685c331..9e68d03971 100644 --- a/cmd/thv-operator/controllers/mcpauthzconfig_controller.go +++ b/cmd/thv-operator/controllers/mcpauthzconfig_controller.go @@ -14,16 +14,11 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" - "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/tools/events" ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" - "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/predicate" - "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil" @@ -41,7 +36,8 @@ const ( // // This controller manages the lifecycle of MCPAuthzConfig resources: validation // via the authorizer factory registry, config hash computation, finalizer management, -// reference tracking, and deletion protection when workloads reference this config. +// and deletion protection when workloads reference this config. The set of referencing +// workloads is recomputed on demand during deletion rather than tracked in status. type MCPAuthzConfigReconciler struct { client.Client Scheme *runtime.Scheme @@ -128,18 +124,6 @@ func (r *MCPAuthzConfigReconciler) Reconcile(ctx context.Context, req ctrl.Reque wasInvalid := conditionStatusIs(authzConfig.Status.Conditions, mcpv1beta1.ConditionTypeAuthzConfigValid, metav1.ConditionFalse) - // Refresh the referencing workloads list. If the lookup fails we must - // requeue with backoff rather than continuing with a stale slice — a - // silent swallow would let ReferencingWorkloads / ReferenceCount drift - // permanently out of sync with the cluster on a transient apiserver - // hiccup. - referencingWorkloads, err := r.findReferencingWorkloads(ctx, authzConfig) - if err != nil { - logger.Error(err, "Failed to find referencing workloads") - return ctrl.Result{}, err - } - newRefCount := workloadReferenceCount(referencingWorkloads) - // Calculate the hash of the current configuration. // The spec is canonicalized first so that two semantically-equal configs // that differ only in whitespace or JSON key order produce the same hash — @@ -156,18 +140,15 @@ func (r *MCPAuthzConfigReconciler) Reconcile(ctx context.Context, req ctrl.Reque } // Single status patch covering the steady-state success path: Valid=True, - // fresh hash + observedGeneration, fresh references, and DeletionBlocked - // cleared (we only reach this path when DeletionTimestamp is zero). - // MutateAndPatchStatus short-circuits on an empty diff, so a no-op - // reconcile still produces no wire call (SteadyStateNoOp behaviour - // preserved). + // fresh hash + observedGeneration, and DeletionBlocked cleared (we only + // reach this path when DeletionTimestamp is zero). MutateAndPatchStatus + // short-circuits on an empty diff, so a no-op reconcile still produces no + // wire call (SteadyStateNoOp behaviour preserved). if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, authzConfig, func(c *mcpv1beta1.MCPAuthzConfig) { setValidTrueCondition(c) meta.RemoveStatusCondition(&c.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) c.Status.ConfigHash = configHash c.Status.ObservedGeneration = c.Generation - c.Status.ReferencingWorkloads = referencingWorkloads - c.Status.ReferenceCount = newRefCount }); err != nil { logger.Error(err, "Failed to update MCPAuthzConfig status") return ctrl.Result{}, err @@ -279,8 +260,6 @@ func (r *MCPAuthzConfigReconciler) handleDeletion( Message: fmt.Sprintf("Cannot delete: referenced by workloads: %v", referencingWorkloads), ObservedGeneration: c.Generation, }) - c.Status.ReferencingWorkloads = referencingWorkloads - c.Status.ReferenceCount = workloadReferenceCount(referencingWorkloads) }) if updateErr != nil { logger.Error(updateErr, "Failed to update status during deletion block") @@ -394,8 +373,6 @@ func (r *MCPAuthzConfigReconciler) findReferencingWorkloads( } // SetupWithManager sets up the controller with the Manager. -// Watches MCPServer, VirtualMCPServer, and MCPRemoteProxy changes to maintain -// accurate ReferencingWorkloads status. func (r *MCPAuthzConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { // Field indexes backing findReferencingWorkloads: each lets the controller // query only the workloads referencing a given config rather than listing @@ -416,75 +393,7 @@ func (r *MCPAuthzConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { return fmt.Errorf("failed to set up VirtualMCPServer authzConfigRef index: %w", err) } - // GenerationChangedPredicate also suppresses the workload-watch resync; the self-heal - // backstop for a stale ReferencingWorkloads entry (e.g. a workload deleted while the - // operator was down) is this config's own For() resync, which re-runs Reconcile and - // rebuilds ReferencingWorkloads from the index. return ctrl.NewControllerManagedBy(mgr). For(&mcpv1beta1.MCPAuthzConfig{}). - Watches(&mcpv1beta1.MCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPServerToAuthzConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{})). - Watches(&mcpv1beta1.VirtualMCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapVirtualMCPServerToAuthzConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{})). - Watches(&mcpv1beta1.MCPRemoteProxy{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPRemoteProxyToAuthzConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{})). Complete(r) } - -// mapMCPServerToAuthzConfig maps an MCPServer to the MCPAuthzConfig it currently references. -// EnqueueRequestsFromMapFunc invokes this on both the old and new object on update (and on the -// object for create/delete), so a ref change or deletion automatically enqueues both the -// previously- and newly-referenced config; the previously-referenced config then prunes the -// stale entry on reconcile. No manual stale-reference scan needed. -func (*MCPAuthzConfigReconciler) mapMCPServerToAuthzConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - server, ok := obj.(*mcpv1beta1.MCPServer) - if !ok || server.Spec.AuthzConfigRef == nil || server.Spec.AuthzConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: server.Spec.AuthzConfigRef.Name, - Namespace: server.Namespace, - }}} -} - -// mapVirtualMCPServerToAuthzConfig maps a VirtualMCPServer to the MCPAuthzConfig it currently -// references via spec.incomingAuth. EnqueueRequestsFromMapFunc invokes this on both the old and -// new object on update (and on the object for create/delete), so a ref change or deletion -// automatically enqueues both the previously- and newly-referenced config; the previously- -// referenced config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -func (*MCPAuthzConfigReconciler) mapVirtualMCPServerToAuthzConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - vmcp, ok := obj.(*mcpv1beta1.VirtualMCPServer) - if !ok || vmcp.Spec.IncomingAuth == nil || - vmcp.Spec.IncomingAuth.AuthzConfigRef == nil || vmcp.Spec.IncomingAuth.AuthzConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: vmcp.Spec.IncomingAuth.AuthzConfigRef.Name, - Namespace: vmcp.Namespace, - }}} -} - -// mapMCPRemoteProxyToAuthzConfig maps an MCPRemoteProxy to the MCPAuthzConfig it currently -// references. EnqueueRequestsFromMapFunc invokes this on both the old and new object on update -// (and on the object for create/delete), so a ref change or deletion automatically enqueues both -// the previously- and newly-referenced config; the previously-referenced config then prunes the -// stale entry on reconcile. No manual stale-reference scan needed. -func (*MCPAuthzConfigReconciler) mapMCPRemoteProxyToAuthzConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - proxy, ok := obj.(*mcpv1beta1.MCPRemoteProxy) - if !ok || proxy.Spec.AuthzConfigRef == nil || proxy.Spec.AuthzConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: proxy.Spec.AuthzConfigRef.Name, - Namespace: proxy.Namespace, - }}} -} diff --git a/cmd/thv-operator/controllers/mcpauthzconfig_controller_test.go b/cmd/thv-operator/controllers/mcpauthzconfig_controller_test.go index d377efb0bc..681e26c7b6 100644 --- a/cmd/thv-operator/controllers/mcpauthzconfig_controller_test.go +++ b/cmd/thv-operator/controllers/mcpauthzconfig_controller_test.go @@ -9,7 +9,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" @@ -488,7 +487,7 @@ func TestMCPAuthzConfigReconciler_ConfigChangeTriggersHashUpdate(t *testing.T) { assert.Equal(t, int64(2), finalConfig.Status.ObservedGeneration, "ObservedGeneration should be updated") } -func TestMCPAuthzConfigReconciler_HashAndRefsLandInOneReconcile(t *testing.T) { +func TestMCPAuthzConfigReconciler_HashAndGenerationLandInOneReconcile(t *testing.T) { t.Parallel() ctx := t.Context() @@ -512,14 +511,9 @@ func TestMCPAuthzConfigReconciler_HashAndRefsLandInOneReconcile(t *testing.T) { var after mcpv1beta1.MCPAuthzConfig require.NoError(t, fakeClient.Get(ctx, req.NamespacedName, &after)) - // F4 regression: hash, references, AND ObservedGeneration must all land - // in the same success-path reconcile. The previous shape returned early - // on hashChanged and left ReferenceCount at zero until the next event. + // The hash and ObservedGeneration must both land in the same success-path + // reconcile rather than the controller returning early on hashChanged. assert.NotEmpty(t, after.Status.ConfigHash, "ConfigHash should be set") - assert.Equal(t, int32(1), after.Status.ReferenceCount, - "ReferenceCount must match the referencing workload list in the same reconcile that wrote the hash") - require.Len(t, after.Status.ReferencingWorkloads, 1) - assert.Equal(t, "ref-server", after.Status.ReferencingWorkloads[0].Name) assert.Equal(t, int64(1), after.Status.ObservedGeneration, "ObservedGeneration must land in the same patch as the hash") } @@ -706,99 +700,6 @@ func TestMCPAuthzConfigReconciler_findReferencingWorkloads(t *testing.T) { } } -// TestMCPAuthzConfigReconciler_watchHandlers verifies that the workload watch map -// functions enqueue exactly the config the workload currently references (or nothing -// when the workload has no ref). The previously-referenced config is enqueued by -// EnqueueRequestsFromMapFunc, which runs the map function on both the old and new -// object on update — no manual stale-reference scan in the handler. -func TestMCPAuthzConfigReconciler_watchHandlers(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - obj client.Object - expected map[string]struct{} - }{ - { - name: "MCPServer with ref enqueues the current config", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithAuthzConfigRef("current-config"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "VirtualMCPServer with ref enqueues the current config", - obj: v1beta1test.NewVirtualMCPServer("vmcp", "default", - v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ - Type: "anonymous", - AuthzConfigRef: &mcpv1beta1.MCPAuthzConfigReference{Name: "current-config"}, - }), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPRemoteProxy with ref enqueues the current config", - obj: v1beta1test.NewMCPRemoteProxy("proxy", "default", - v1beta1test.WithRemoteProxyURL("https://example.com"), - v1beta1test.WithRemoteProxyAuthzConfigRef("current-config"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPServer without ref enqueues nothing", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - ), - expected: map[string]struct{}{}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := t.Context() - r, _ := newAuthzTestReconciler(t, tt.obj) - - requests := func() []reconcile.Request { - switch tt.obj.(type) { - case *mcpv1beta1.MCPServer: - return r.mapMCPServerToAuthzConfig(ctx, tt.obj) - case *mcpv1beta1.VirtualMCPServer: - return r.mapVirtualMCPServerToAuthzConfig(ctx, tt.obj) - case *mcpv1beta1.MCPRemoteProxy: - return r.mapMCPRemoteProxyToAuthzConfig(ctx, tt.obj) - default: - t.Fatalf("unexpected object type %T", tt.obj) - return nil - } - }() - - got := make(map[string]struct{}, len(requests)) - for _, req := range requests { - assert.Equal(t, "default", req.Namespace) - got[req.Name] = struct{}{} - } - assert.Equal(t, tt.expected, got) - }) - } -} - -// TestMCPAuthzConfigReconciler_watchHandlersWrongType verifies the map functions -// gracefully ignore objects of an unexpected type. -func TestMCPAuthzConfigReconciler_watchHandlersWrongType(t *testing.T) { - t.Parallel() - - ctx := t.Context() - r, _ := newAuthzTestReconciler(t) - - wrong := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{Name: "cm", Namespace: "default"}} - assert.Nil(t, r.mapMCPServerToAuthzConfig(ctx, wrong)) - assert.Nil(t, r.mapVirtualMCPServerToAuthzConfig(ctx, wrong)) - assert.Nil(t, r.mapMCPRemoteProxyToAuthzConfig(ctx, wrong)) -} - // TestMCPAuthzConfigReconciler_DeletionWithoutFinalizer verifies that handleDeletion // is a no-op when the config never had the finalizer (the object is passed directly // rather than created in the fake client, which rejects a deletionTimestamp without a diff --git a/cmd/thv-operator/controllers/mcpexternalauthconfig_controller.go b/cmd/thv-operator/controllers/mcpexternalauthconfig_controller.go index 35e2c95307..b6655b7db5 100644 --- a/cmd/thv-operator/controllers/mcpexternalauthconfig_controller.go +++ b/cmd/thv-operator/controllers/mcpexternalauthconfig_controller.go @@ -17,13 +17,9 @@ import ( "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/tools/events" ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" - "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/predicate" - "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil" @@ -158,9 +154,8 @@ func (r *MCPExternalAuthConfigReconciler) Reconcile(ctx context.Context, req ctr } // Steady-state success path: ensure Valid=True and the IdentitySynthesized - // advisory are set, and refresh the referencing-workloads list, in a single - // status patch. - return r.updateReferencingWorkloads(ctx, externalAuthConfig) + // advisory are set in a single status patch. + return r.updateSteadyStateStatus(ctx, externalAuthConfig) } // setValidTrueAndSynthesized stamps ConditionTypeValid=True and refreshes the @@ -332,36 +327,20 @@ func (r *MCPExternalAuthConfigReconciler) handleConfigHashChange( "oldHash", externalAuthConfig.Status.ConfigHash, "newHash", configHash) - // Find all MCPServers that reference this MCPExternalAuthConfig - referencingServers, err := r.findReferencingMCPServers(ctx, externalAuthConfig) - if err != nil { - logger.Error(err, "Failed to find referencing MCPServers") - return ctrl.Result{}, fmt.Errorf("failed to find referencing MCPServers: %w", err) - } - - // Build the list of referencing workloads - refs := make([]mcpv1beta1.WorkloadReference, 0, len(referencingServers)) - for _, server := range referencingServers { - refs = append(refs, mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPServer, Name: server.Name}) - } - ctrlutil.SortWorkloadRefs(refs) - // Capture the recovery transition before the patch mutates conditions in // place, so a single Normal event fires on the False->True transition. wasInvalid := conditionStatusIs(externalAuthConfig.Status.Conditions, mcpv1beta1.ConditionTypeValid, metav1.ConditionFalse) // Single status patch covering the hash-change success path: the new hash - // and generation, the refreshed reference list, and the Valid=True / - // IdentitySynthesized conditions. All mutations happen inside the closure so - // the pre-mutate snapshot stays clean (a MutateAndPatchStatus prerequisite). + // and generation, and the Valid=True / IdentitySynthesized conditions. All + // mutations happen inside the closure so the pre-mutate snapshot stays clean + // (a MutateAndPatchStatus prerequisite). if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, externalAuthConfig, func(c *mcpv1beta1.MCPExternalAuthConfig) { r.setValidTrueAndSynthesized(c) c.Status.ConfigHash = configHash c.Status.ObservedGeneration = c.Generation - c.Status.ReferencingWorkloads = refs - c.Status.ReferenceCount = workloadReferenceCount(refs) }); err != nil { logger.Error(err, "Failed to update MCPExternalAuthConfig status") return ctrl.Result{}, err @@ -404,8 +383,6 @@ func (r *MCPExternalAuthConfigReconciler) handleDeletion( Message: fmt.Sprintf("Cannot delete: referenced by workloads: %v", referencingWorkloads), ObservedGeneration: c.Generation, }) - c.Status.ReferencingWorkloads = referencingWorkloads - c.Status.ReferenceCount = workloadReferenceCount(referencingWorkloads) }) if updateErr != nil { logger.Error(updateErr, "Failed to update status during deletion block") @@ -573,7 +550,6 @@ func (r *MCPExternalAuthConfigReconciler) findReferencingWorkloads( } // SetupWithManager sets up the controller with the Manager. -// Watches MCPServer and MCPRemoteProxy changes to maintain accurate ReferencingWorkloads status. func (r *MCPExternalAuthConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { // Field indexes backing findReferencingMCPServers / findReferencingMCPRemoteProxies. // Each is a combined index covering both spec.externalAuthConfigRef and @@ -592,135 +568,21 @@ func (r *MCPExternalAuthConfigReconciler) SetupWithManager(mgr ctrl.Manager) err return fmt.Errorf("failed to set up MCPRemoteProxy externalAuthConfigRef index: %w", err) } - // GenerationChangedPredicate also suppresses the workload-watch resync; the self-heal - // backstop for a stale ReferencingWorkloads entry (e.g. a workload deleted while the - // operator was down) is this config's own For() resync, which re-runs Reconcile and - // rebuilds ReferencingWorkloads from the index. return ctrl.NewControllerManagedBy(mgr). For(&mcpv1beta1.MCPExternalAuthConfig{}). - Watches( - &mcpv1beta1.MCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPServerToExternalAuthConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). - Watches( - &mcpv1beta1.MCPRemoteProxy{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPRemoteProxyToExternalAuthConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). Complete(r) } -// mapMCPServerToExternalAuthConfig maps an MCPServer to the MCPExternalAuthConfig(s) -// it currently references via spec.externalAuthConfigRef and/or spec.authServerRef. -// EnqueueRequestsFromMapFunc invokes this on both the old and new object on update -// (and on the object for create/delete), so a ref change or deletion automatically -// enqueues both the previously- and newly-referenced config; the previously-referenced -// config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -// -// A server may name the same config via both fields, so the two referenced names are -// deduplicated against a seen set before being turned into requests. -func (*MCPExternalAuthConfigReconciler) mapMCPServerToExternalAuthConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - server, ok := obj.(*mcpv1beta1.MCPServer) - if !ok { - return nil - } - - seen := make(map[types.NamespacedName]struct{}) - var requests []reconcile.Request - - // Enqueue the currently-referenced MCPExternalAuthConfig via externalAuthConfigRef (if any) - if server.Spec.ExternalAuthConfigRef != nil && server.Spec.ExternalAuthConfigRef.Name != "" { - nn := types.NamespacedName{ - Name: server.Spec.ExternalAuthConfigRef.Name, - Namespace: server.Namespace, - } - seen[nn] = struct{}{} - requests = append(requests, reconcile.Request{NamespacedName: nn}) - } - - // Enqueue the MCPExternalAuthConfig referenced via authServerRef (if any), deduped - if server.Spec.AuthServerRef != nil && - server.Spec.AuthServerRef.Kind == authServerRefKindMCPExternalAuthConfig && - server.Spec.AuthServerRef.Name != "" { - nn := types.NamespacedName{ - Name: server.Spec.AuthServerRef.Name, - Namespace: server.Namespace, - } - if _, already := seen[nn]; !already { - seen[nn] = struct{}{} - requests = append(requests, reconcile.Request{NamespacedName: nn}) - } - } - - return requests -} - -// mapMCPRemoteProxyToExternalAuthConfig maps an MCPRemoteProxy to the MCPExternalAuthConfig(s) -// it currently references via spec.externalAuthConfigRef and/or spec.authServerRef. -// EnqueueRequestsFromMapFunc invokes this on both the old and new object on update -// (and on the object for create/delete), so a ref change or deletion automatically -// enqueues both the previously- and newly-referenced config; the previously-referenced -// config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -// -// A proxy may name the same config via both fields, so the two referenced names are -// deduplicated against a seen set before being turned into requests. -func (*MCPExternalAuthConfigReconciler) mapMCPRemoteProxyToExternalAuthConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - proxy, ok := obj.(*mcpv1beta1.MCPRemoteProxy) - if !ok { - return nil - } - - seen := make(map[types.NamespacedName]struct{}) - var requests []reconcile.Request - - // Enqueue the currently-referenced MCPExternalAuthConfig via externalAuthConfigRef (if any) - if proxy.Spec.ExternalAuthConfigRef != nil && proxy.Spec.ExternalAuthConfigRef.Name != "" { - nn := types.NamespacedName{ - Name: proxy.Spec.ExternalAuthConfigRef.Name, - Namespace: proxy.Namespace, - } - seen[nn] = struct{}{} - requests = append(requests, reconcile.Request{NamespacedName: nn}) - } - - // Enqueue the MCPExternalAuthConfig referenced via authServerRef (if any), deduped - if proxy.Spec.AuthServerRef != nil && - proxy.Spec.AuthServerRef.Kind == authServerRefKindMCPExternalAuthConfig && - proxy.Spec.AuthServerRef.Name != "" { - nn := types.NamespacedName{ - Name: proxy.Spec.AuthServerRef.Name, - Namespace: proxy.Namespace, - } - if _, already := seen[nn]; !already { - seen[nn] = struct{}{} - requests = append(requests, reconcile.Request{NamespacedName: nn}) - } - } - - return requests -} - -// updateReferencingWorkloads writes the steady-state success status in a single -// patch: it ensures Valid=True and the IdentitySynthesized advisory are set and -// refreshes the referencing-workloads list. MutateAndPatchStatus short-circuits -// on an empty diff, so a no-op reconcile skips the wire call. -func (r *MCPExternalAuthConfigReconciler) updateReferencingWorkloads( +// updateSteadyStateStatus writes the steady-state success status in a single +// patch: it ensures Valid=True and the IdentitySynthesized advisory are set. +// MutateAndPatchStatus short-circuits on an empty diff, so a no-op reconcile +// skips the wire call. +func (r *MCPExternalAuthConfigReconciler) updateSteadyStateStatus( ctx context.Context, externalAuthConfig *mcpv1beta1.MCPExternalAuthConfig, ) (ctrl.Result, error) { logger := log.FromContext(ctx) - refs, err := r.findReferencingWorkloads(ctx, externalAuthConfig) - if err != nil { - logger.Error(err, "Failed to find referencing workloads") - return ctrl.Result{}, fmt.Errorf("failed to find referencing workloads: %w", err) - } - // Capture the recovery transition before the patch mutates conditions in // place, so a single Normal event fires on the False->True transition. wasInvalid := conditionStatusIs(externalAuthConfig.Status.Conditions, @@ -729,11 +591,6 @@ func (r *MCPExternalAuthConfigReconciler) updateReferencingWorkloads( if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, externalAuthConfig, func(c *mcpv1beta1.MCPExternalAuthConfig) { r.setValidTrueAndSynthesized(c) - if !ctrlutil.WorkloadRefsEqual(c.Status.ReferencingWorkloads, refs) || - c.Status.ReferenceCount != workloadReferenceCount(refs) { - c.Status.ReferencingWorkloads = refs - c.Status.ReferenceCount = workloadReferenceCount(refs) - } }); err != nil { logger.Error(err, "Failed to update MCPExternalAuthConfig status") return ctrl.Result{}, err diff --git a/cmd/thv-operator/controllers/mcpexternalauthconfig_controller_test.go b/cmd/thv-operator/controllers/mcpexternalauthconfig_controller_test.go index 07db859bee..f99417b762 100644 --- a/cmd/thv-operator/controllers/mcpexternalauthconfig_controller_test.go +++ b/cmd/thv-operator/controllers/mcpexternalauthconfig_controller_test.go @@ -238,13 +238,6 @@ func TestMCPExternalAuthConfigReconciler_Reconcile(t *testing.T) { assert.NotEmpty(t, updatedConfig.Status.ConfigHash, "MCPExternalAuthConfig status should have config hash") } - - // Check referencing workloads in status - if tt.existingMCPServer != nil { - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: tt.existingMCPServer.Name}, - "Status should contain referencing MCPServer as WorkloadReference") - } }) } } @@ -594,144 +587,6 @@ func TestMCPExternalAuthConfigReconciler_ConfigChangeTriggersReconciliation(t *t assert.False(t, found, "MCPExternalAuthConfig controller should not annotate MCPServers") } -func TestMCPExternalAuthConfigReconciler_ReferencingWorkloadsUpdatedWithoutHashChange(t *testing.T) { - t.Parallel() - - ctx := t.Context() - - externalAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-config", - Namespace: "default", - Generation: 1, - }, - Spec: mcpv1beta1.MCPExternalAuthConfigSpec{ - Type: mcpv1beta1.ExternalAuthTypeTokenExchange, - TokenExchange: &mcpv1beta1.TokenExchangeConfig{ - TokenURL: "https://oauth.example.com/token", - ClientID: "test-client", - ClientSecretRef: &mcpv1beta1.SecretKeyRef{ - Name: "test-secret", - Key: "client-secret", - }, - Audience: "backend-service", - }, - }, - } - - r, fakeClient := newTestMCPExternalAuthConfigReconciler(t, externalAuthConfig) - - req := reconcile.Request{ - NamespacedName: types.NamespacedName{ - Name: externalAuthConfig.Name, - Namespace: externalAuthConfig.Namespace, - }, - } - - // First reconciliation - add finalizer - result, err := r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0)) - - // Second reconciliation - sets hash, no servers yet - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - var updatedConfig mcpv1beta1.MCPExternalAuthConfig - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.NotEmpty(t, updatedConfig.Status.ConfigHash) - assert.Empty(t, updatedConfig.Status.ReferencingWorkloads, "No workloads should be referencing yet") - assert.EqualValues(t, 0, updatedConfig.Status.ReferenceCount) - - // Now add an MCPServer that references this config (without changing the config spec) - mcpServer := v1beta1test.NewMCPServer("new-server", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithExternalAuthConfigRef("test-config"), - ) - require.NoError(t, fakeClient.Create(ctx, mcpServer)) - - // Reconcile again - hash hasn't changed, but referencing servers should be updated - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "new-server"}, - "ReferencingWorkloads should be updated even without hash change") - assert.EqualValues(t, 1, updatedConfig.Status.ReferenceCount) -} - -func TestMCPExternalAuthConfigReconciler_ReferencingWorkloadsRemovedOnServerDeletion(t *testing.T) { - t.Parallel() - - ctx := t.Context() - - externalAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-config", - Namespace: "default", - Generation: 1, - }, - Spec: mcpv1beta1.MCPExternalAuthConfigSpec{ - Type: mcpv1beta1.ExternalAuthTypeTokenExchange, - TokenExchange: &mcpv1beta1.TokenExchangeConfig{ - TokenURL: "https://oauth.example.com/token", - ClientID: "test-client", - ClientSecretRef: &mcpv1beta1.SecretKeyRef{ - Name: "test-secret", - Key: "client-secret", - }, - Audience: "backend-service", - }, - }, - } - - mcpServer := v1beta1test.NewMCPServer("server-to-delete", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithExternalAuthConfigRef("test-config"), - ) - - r, fakeClient := newTestMCPExternalAuthConfigReconciler(t, externalAuthConfig, mcpServer) - - req := reconcile.Request{ - NamespacedName: types.NamespacedName{ - Name: externalAuthConfig.Name, - Namespace: externalAuthConfig.Namespace, - }, - } - - // Add finalizer - result, err := r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0)) - - // Set hash and referencing servers - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - var updatedConfig mcpv1beta1.MCPExternalAuthConfig - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "server-to-delete"}) - assert.EqualValues(t, 1, updatedConfig.Status.ReferenceCount) - - // Delete the MCPServer - require.NoError(t, fakeClient.Delete(ctx, mcpServer)) - - // Reconcile again - referencing servers should be empty now - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Empty(t, updatedConfig.Status.ReferencingWorkloads, - "ReferencingWorkloads should be empty after server deletion") - assert.EqualValues(t, 0, updatedConfig.Status.ReferenceCount) -} - func TestMCPExternalAuthConfigReconciler_findReferencingWorkloads_authServerRef(t *testing.T) { t.Parallel() @@ -1561,116 +1416,3 @@ func TestMCPExternalAuthConfigReconciler_ReconcileKeepsExistingForeignCondition( require.NotNil(t, own, "controller-owned Valid condition must land") assert.Equal(t, metav1.ConditionTrue, own.Status) } - -// TestMCPExternalAuthConfigReconciler_watchHandlers verifies that the workload -// watch map functions enqueue exactly the config(s) the workload currently -// references via either externalAuthConfigRef or a MCPExternalAuthConfig-kind -// authServerRef, deduplicating when both name the same config. The -// previously-referenced config is enqueued by EnqueueRequestsFromMapFunc, which -// runs the map function on both the old and new object on update — no manual -// stale-reference scan in the handler. -func TestMCPExternalAuthConfigReconciler_watchHandlers(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - obj client.Object - expected map[string]struct{} - }{ - { - name: "MCPServer externalAuthConfigRef enqueues the current config", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithExternalAuthConfigRef("cfg-a"), - ), - expected: map[string]struct{}{"cfg-a": {}}, - }, - { - name: "MCPServer authServerRef of MCPExternalAuthConfig kind enqueues it", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithAuthServerRef(authServerRefKindMCPExternalAuthConfig, "cfg-b"), - ), - expected: map[string]struct{}{"cfg-b": {}}, - }, - { - name: "MCPServer with both refs to different configs enqueues both", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithExternalAuthConfigRef("cfg-a"), - v1beta1test.WithAuthServerRef(authServerRefKindMCPExternalAuthConfig, "cfg-b"), - ), - expected: map[string]struct{}{"cfg-a": {}, "cfg-b": {}}, - }, - { - name: "MCPServer with both refs to the same config enqueues it once", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithExternalAuthConfigRef("cfg-a"), - v1beta1test.WithAuthServerRef(authServerRefKindMCPExternalAuthConfig, "cfg-a"), - ), - expected: map[string]struct{}{"cfg-a": {}}, - }, - { - name: "MCPServer authServerRef of another kind is ignored", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithAuthServerRef("MCPOIDCConfig", "cfg-x"), - ), - expected: map[string]struct{}{}, - }, - { - name: "MCPServer without refs enqueues nothing", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - ), - expected: map[string]struct{}{}, - }, - { - name: "MCPRemoteProxy with both refs to different configs enqueues both", - obj: v1beta1test.NewMCPRemoteProxy("proxy", "default", - v1beta1test.WithRemoteProxyURL("https://example.com"), - v1beta1test.WithRemoteProxyExternalAuthConfigRef("cfg-a"), - v1beta1test.WithRemoteProxyAuthServerRef(authServerRefKindMCPExternalAuthConfig, "cfg-b"), - ), - expected: map[string]struct{}{"cfg-a": {}, "cfg-b": {}}, - }, - { - name: "MCPRemoteProxy with both refs to the same config enqueues it once", - obj: v1beta1test.NewMCPRemoteProxy("proxy", "default", - v1beta1test.WithRemoteProxyURL("https://example.com"), - v1beta1test.WithRemoteProxyExternalAuthConfigRef("cfg-a"), - v1beta1test.WithRemoteProxyAuthServerRef(authServerRefKindMCPExternalAuthConfig, "cfg-a"), - ), - expected: map[string]struct{}{"cfg-a": {}}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := t.Context() - r, _ := newTestMCPExternalAuthConfigReconciler(t) - - requests := func() []reconcile.Request { - switch tt.obj.(type) { - case *mcpv1beta1.MCPServer: - return r.mapMCPServerToExternalAuthConfig(ctx, tt.obj) - case *mcpv1beta1.MCPRemoteProxy: - return r.mapMCPRemoteProxyToExternalAuthConfig(ctx, tt.obj) - default: - t.Fatalf("unexpected object type %T", tt.obj) - return nil - } - }() - - got := make(map[string]struct{}, len(requests)) - for _, req := range requests { - assert.Equal(t, "default", req.Namespace) - got[req.Name] = struct{}{} - } - assert.Equal(t, tt.expected, got) - }) - } -} diff --git a/cmd/thv-operator/controllers/mcpoidcconfig_controller.go b/cmd/thv-operator/controllers/mcpoidcconfig_controller.go index f8f377d5fc..3dc65cb9a5 100644 --- a/cmd/thv-operator/controllers/mcpoidcconfig_controller.go +++ b/cmd/thv-operator/controllers/mcpoidcconfig_controller.go @@ -13,19 +13,15 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" - "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/tools/events" ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" - "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/predicate" - "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil" + "github.com/stacklok/toolhive/cmd/thv-operator/pkg/validation" ) const ( @@ -39,8 +35,9 @@ const ( // MCPOIDCConfigReconciler reconciles a MCPOIDCConfig object. // // This controller manages the lifecycle of MCPOIDCConfig resources: validation, -// config hash computation, finalizer management, reference tracking, and -// deletion protection when MCPServer resources reference this config. +// config hash computation, finalizer management, and deletion protection when +// MCPServer resources reference this config. The set of referencing workloads is +// recomputed on demand during deletion rather than tracked in status. type MCPOIDCConfigReconciler struct { client.Client Scheme *runtime.Scheme @@ -91,7 +88,7 @@ func (r *MCPOIDCConfigReconciler) Reconcile(ctx context.Context, req ctrl.Reques } // Validate spec configuration early - if err := oidcConfig.Validate(); err != nil { + if err := validateOIDCConfigSpec(oidcConfig); err != nil { logger.Error(err, "MCPOIDCConfig spec validation failed") return r.handleValidationFailure(ctx, oidcConfig, err) } @@ -124,27 +121,12 @@ func (r *MCPOIDCConfigReconciler) Reconcile(ctx context.Context, req ctrl.Reques return ctrl.Result{}, nil } - // Refresh ReferencingWorkloads list. On error, fall through with the lookup - // result skipped: the status patch below is best-effort and still ensures - // the Valid=True condition is set even when the reference refresh fails. - referencingWorkloads, findErr := r.findReferencingWorkloads(ctx, oidcConfig) - if findErr != nil { - logger.Error(findErr, "Failed to find referencing workloads") - } - // Single status patch covering the steady-state success path: ensure the - // Valid=True condition is set, and refresh the references list when the - // lookup succeeded and the list changed. MutateAndPatchStatus short-circuits - // on an empty diff so the no-op case still skips the wire call - // (SteadyStateNoOp behaviour is preserved). + // Valid=True condition is set. MutateAndPatchStatus short-circuits on an + // empty diff so the no-op case still skips the wire call (SteadyStateNoOp + // behaviour is preserved). if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, oidcConfig, func(c *mcpv1beta1.MCPOIDCConfig) { setOIDCConfigValidTrueCondition(c) - if findErr == nil && - (!ctrlutil.WorkloadRefsEqual(c.Status.ReferencingWorkloads, referencingWorkloads) || - c.Status.ReferenceCount != workloadReferenceCount(referencingWorkloads)) { - c.Status.ReferencingWorkloads = referencingWorkloads - c.Status.ReferenceCount = workloadReferenceCount(referencingWorkloads) - } }); err != nil { logger.Error(err, "Failed to update MCPOIDCConfig status") return ctrl.Result{}, err @@ -154,6 +136,26 @@ func (r *MCPOIDCConfigReconciler) Reconcile(ctx context.Context, req ctrl.Reques return ctrl.Result{}, nil } +// validateOIDCConfigSpec validates the MCPOIDCConfig spec: type/config +// consistency via Validate(), then issuer and JWKS URL checks for inline +// configs. The URL checks run here rather than inside Validate() because the +// validation package imports the v1beta1 API package (for example +// validation.ValidateCABundleSource), so calling it from the types package +// would create an import cycle. This mirrors the MCPRemoteProxy controller, +// which applies validation.ValidateRemoteURL in its validateSpec step. +func validateOIDCConfigSpec(oidcConfig *mcpv1beta1.MCPOIDCConfig) error { + if err := oidcConfig.Validate(); err != nil { + return err + } + if oidcConfig.Spec.Type != mcpv1beta1.MCPOIDCConfigTypeInline || oidcConfig.Spec.Inline == nil { + return nil + } + if err := validation.ValidateOIDCIssuerURL(oidcConfig.Spec.Inline.Issuer, oidcConfig.Spec.Inline.InsecureAllowHTTP); err != nil { + return err + } + return validation.ValidateJWKSURL(oidcConfig.Spec.Inline.JWKSURL, oidcConfig.Spec.Inline.InsecureAllowHTTP) +} + // handleValidationFailure records the Valid=False condition for a spec that // failed validation and emits a one-shot Warning event on the transition into // the invalid state. It never requeues — the user must fix the spec. @@ -243,8 +245,6 @@ func (r *MCPOIDCConfigReconciler) handleDeletion( Message: fmt.Sprintf("Cannot delete: referenced by workloads: %v", referencingWorkloads), ObservedGeneration: c.Generation, }) - c.Status.ReferencingWorkloads = referencingWorkloads - c.Status.ReferenceCount = workloadReferenceCount(referencingWorkloads) }) if updateErr != nil { logger.Error(updateErr, "Failed to update status during deletion block") @@ -358,7 +358,6 @@ func (r *MCPOIDCConfigReconciler) findReferencingWorkloads( } // SetupWithManager sets up the controller with the Manager. -// Watches MCPServer, VirtualMCPServer, and MCPRemoteProxy changes to maintain accurate ReferencingWorkloads status. func (r *MCPOIDCConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { // Field indexes backing findReferencingWorkloads: each lets the controller // query only the workloads referencing a given config rather than listing @@ -379,87 +378,7 @@ func (r *MCPOIDCConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { return fmt.Errorf("failed to set up MCPRemoteProxy oidcConfigRef index: %w", err) } - // Watch MCPServer changes to update ReferencingWorkloads on referenced MCPOIDCConfigs. - // The map function only returns the MCPOIDCConfig the server currently references. - // EnqueueRequestsFromMapFunc runs it on both the old and new object on update (and - // on the object for create/delete), so removing or changing the ref enqueues both - // the previously- and newly-referenced config — the previously-referenced config - // then reconciles and prunes the stale entry. No manual stale-reference scan needed. - // - // GenerationChangedPredicate also suppresses the workload-watch resync; the self-heal - // backstop for a stale ReferencingWorkloads entry (e.g. a workload deleted while the - // operator was down) is this config's own For() resync, which re-runs Reconcile and - // rebuilds ReferencingWorkloads from the index. return ctrl.NewControllerManagedBy(mgr). For(&mcpv1beta1.MCPOIDCConfig{}). - Watches(&mcpv1beta1.MCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPServerToOIDCConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{})). - Watches( - &mcpv1beta1.VirtualMCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapVirtualMCPServerToOIDCConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). - Watches( - &mcpv1beta1.MCPRemoteProxy{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPRemoteProxyToOIDCConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). Complete(r) } - -// mapMCPServerToOIDCConfig maps an MCPServer to the MCPOIDCConfig it currently -// references. EnqueueRequestsFromMapFunc invokes this on both the old and new object on -// update (and on the object for create/delete), so a ref change or deletion automatically -// enqueues both the previously- and newly-referenced config; the previously-referenced -// config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -func (*MCPOIDCConfigReconciler) mapMCPServerToOIDCConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - server, ok := obj.(*mcpv1beta1.MCPServer) - if !ok || server.Spec.OIDCConfigRef == nil || server.Spec.OIDCConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: server.Spec.OIDCConfigRef.Name, - Namespace: server.Namespace, - }}} -} - -// mapVirtualMCPServerToOIDCConfig maps a VirtualMCPServer to the MCPOIDCConfig it -// currently references via spec.incomingAuth. EnqueueRequestsFromMapFunc invokes this -// on both the old and new object on update (and on the object for create/delete), so a -// ref change or deletion automatically enqueues both the previously- and newly-referenced -// config; the previously-referenced config then prunes the stale entry on reconcile. No -// manual stale-reference scan needed. -func (*MCPOIDCConfigReconciler) mapVirtualMCPServerToOIDCConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - vmcp, ok := obj.(*mcpv1beta1.VirtualMCPServer) - if !ok || vmcp.Spec.IncomingAuth == nil || - vmcp.Spec.IncomingAuth.OIDCConfigRef == nil || vmcp.Spec.IncomingAuth.OIDCConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: vmcp.Spec.IncomingAuth.OIDCConfigRef.Name, - Namespace: vmcp.Namespace, - }}} -} - -// mapMCPRemoteProxyToOIDCConfig maps an MCPRemoteProxy to the MCPOIDCConfig it currently -// references. EnqueueRequestsFromMapFunc invokes this on both the old and new object on -// update (and on the object for create/delete), so a ref change or deletion automatically -// enqueues both the previously- and newly-referenced config; the previously-referenced -// config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -func (*MCPOIDCConfigReconciler) mapMCPRemoteProxyToOIDCConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - proxy, ok := obj.(*mcpv1beta1.MCPRemoteProxy) - if !ok || proxy.Spec.OIDCConfigRef == nil || proxy.Spec.OIDCConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: proxy.Spec.OIDCConfigRef.Name, - Namespace: proxy.Namespace, - }}} -} diff --git a/cmd/thv-operator/controllers/mcpoidcconfig_controller_test.go b/cmd/thv-operator/controllers/mcpoidcconfig_controller_test.go index 598d49b535..be7b5c0d33 100644 --- a/cmd/thv-operator/controllers/mcpoidcconfig_controller_test.go +++ b/cmd/thv-operator/controllers/mcpoidcconfig_controller_test.go @@ -4,8 +4,6 @@ package controllers import ( - "context" - "sync/atomic" "testing" "time" @@ -16,11 +14,9 @@ import ( "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" - "sigs.k8s.io/controller-runtime/pkg/client/interceptor" "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" - "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" "github.com/stacklok/toolhive/cmd/thv-operator/internal/testutil" ) @@ -517,57 +513,209 @@ func TestMCPOIDCConfig_Validate(t *testing.T) { } } -func TestMCPOIDCConfigReconciler_ReferenceCountUpdatedWithWorkloads(t *testing.T) { +func TestValidateOIDCConfigSpec(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + config *mcpv1beta1.MCPOIDCConfig + expectError bool + }{ + { + name: "valid inline config with HTTPS issuer and JWKS URL", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "https://accounts.google.com", + JWKSURL: "https://www.googleapis.com/oauth2/v3/certs", + ClientID: "test-client", + }, + }, + }, + expectError: false, + }, + { + name: "valid inline config with empty JWKS URL (auto-discovery)", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "https://accounts.google.com", + ClientID: "test-client", + }, + }, + }, + expectError: false, + }, + { + name: "valid inline config with HTTP issuer when insecureAllowHTTP is set", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "http://issuer.dev.local", + ClientID: "test-client", + InsecureAllowHTTP: true, + }, + }, + }, + expectError: false, + }, + { + name: "invalid inline config with HTTP issuer without insecureAllowHTTP", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "http://issuer.example.com", + ClientID: "test-client", + }, + }, + }, + expectError: true, + }, + { + name: "invalid inline config with malformed issuer", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "not-a-url", + ClientID: "test-client", + }, + }, + }, + expectError: true, + }, + { + name: "invalid inline config with issuer using unsupported scheme", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "ftp://issuer.example.com", + ClientID: "test-client", + }, + }, + }, + expectError: true, + }, + { + name: "invalid inline config with non-HTTPS JWKS URL", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "https://accounts.google.com", + JWKSURL: "http://www.googleapis.com/oauth2/v3/certs", + ClientID: "test-client", + }, + }, + }, + expectError: true, + }, + { + name: "valid inline config with HTTP JWKS URL and insecureAllowHTTP", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + Inline: &mcpv1beta1.InlineOIDCSharedConfig{ + Issuer: "http://keycloak:8080/realms/toolhive", + JWKSURL: "http://keycloak:8080/realms/toolhive/protocol/openid-connect/certs", + ClientID: "test-client", + InsecureAllowHTTP: true, + }, + }, + }, + expectError: false, + }, + { + name: "kubernetesServiceAccount config skips URL validation", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeKubernetesServiceAccount, + KubernetesServiceAccount: &mcpv1beta1.KubernetesServiceAccountOIDCConfig{ + ServiceAccount: "test-sa", + // Would fail the URL checks if they applied to this type + Issuer: "not-a-url", + }, + }, + }, + expectError: false, + }, + { + name: "type consistency failure is still reported", + config: &mcpv1beta1.MCPOIDCConfig{ + Spec: mcpv1beta1.MCPOIDCConfigSpec{ + Type: mcpv1beta1.MCPOIDCConfigTypeInline, + // Missing Inline config + }, + }, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := validateOIDCConfigSpec(tt.config) + + if tt.expectError { + assert.Error(t, err) + } else { + assert.NoError(t, err) + } + }) + } +} + +func TestMCPOIDCConfigReconciler_URLValidationFailureSetsCondition(t *testing.T) { t.Parallel() ctx := t.Context() + // Invalid config: HTTP issuer without insecureAllowHTTP oidcConfig := &mcpv1beta1.MCPOIDCConfig{ ObjectMeta: metav1.ObjectMeta{ - Name: "test-config", + Name: "insecure-issuer-config", Namespace: "default", + Finalizers: []string{OIDCConfigFinalizerName}, Generation: 1, }, Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "https://issuer.example.com", + Issuer: "http://issuer.example.com", ClientID: "test-client", }, }, } - mcpServer := v1beta1test.NewMCPServer("server-to-track", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithOIDCConfigRef("test-config", ""), - ) - - reconciler, fakeClient := newTestMCPOIDCConfigReconciler(t, oidcConfig, mcpServer) - req := reconcile.Request{NamespacedName: types.NamespacedName{Name: oidcConfig.Name, Namespace: oidcConfig.Namespace}} - result, err := reconciler.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0)) + r, fakeClient := newTestMCPOIDCConfigReconciler(t, oidcConfig) - _, err = reconciler.Reconcile(ctx, req) - require.NoError(t, err) + req := reconcile.Request{ + NamespacedName: types.NamespacedName{ + Name: oidcConfig.Name, + Namespace: oidcConfig.Namespace, + }, + } - _, err = reconciler.Reconcile(ctx, req) + // Reconcile should not return error (validation failures are not requeued) + _, err := r.Reconcile(ctx, req) require.NoError(t, err) var updatedConfig mcpv1beta1.MCPOIDCConfig - require.NoError(t, fakeClient.Get(ctx, req.NamespacedName, &updatedConfig)) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPServer, Name: "server-to-track"}) - assert.EqualValues(t, 1, updatedConfig.Status.ReferenceCount) - - require.NoError(t, fakeClient.Delete(ctx, mcpServer)) - - _, err = reconciler.Reconcile(ctx, req) + err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) require.NoError(t, err) - require.NoError(t, fakeClient.Get(ctx, req.NamespacedName, &updatedConfig)) - assert.Empty(t, updatedConfig.Status.ReferencingWorkloads) - assert.EqualValues(t, 0, updatedConfig.Status.ReferenceCount) + cond := meta.FindStatusCondition(updatedConfig.Status.Conditions, mcpv1beta1.ConditionTypeOIDCConfigValid) + require.NotNil(t, cond, "Should have a Valid condition") + assert.Equal(t, metav1.ConditionFalse, cond.Status, "Valid condition should be False") + assert.Equal(t, mcpv1beta1.ConditionReasonOIDCConfigInvalid, cond.Reason) + assert.Contains(t, cond.Message, "http://issuer.example.com", + "Message should surface the offending issuer URL") } // TestMCPOIDCConfigReconciler_ReconcileKeepsExistingForeignCondition verifies @@ -635,195 +783,3 @@ func TestMCPOIDCConfigReconciler_ReconcileKeepsExistingForeignCondition(t *testi require.NotNil(t, own, "controller-owned Valid condition must land") assert.Equal(t, metav1.ConditionTrue, own.Status) } - -// TestMCPOIDCConfigReconciler_ConcurrentForeignConditionSurvivesMergePatch -// proves the property the MutateAndPatchStatus migration actually buys over a -// full r.Status().Update: a condition written by a disjoint owner that lands -// between the reconciler's Get and its status patch survives, because the -// reconciler sends a JSON merge-patch carrying only the fields it changed -// (here referencingWorkloads/referenceCount) rather than sending a full PUT of -// its stale view of the whole Status.Conditions array. -// -// The merge-patch-vs-PUT behaviour lives entirely in the shared -// ctrlutil.MutateAndPatchStatus helper that all three config controllers use, -// so this single end-to-end proof guards the mechanism for the OIDC, -// ExternalAuth, and Authz reconcilers alike. -// -// A WithInterceptorFuncs Get hook simulates the concurrent writer: once armed, -// it injects a foreign condition into the backing store immediately after the -// reconciler reads the object. The reconcile that follows changes only the -// reference list — the spec, and hence the generation and the Valid -// condition's ObservedGeneration, is unchanged — so the merge-patch body omits -// the conditions array and the foreign entry is preserved. A regression to -// r.Status().Update would PUT conditions=[Valid] and erase it. -func TestMCPOIDCConfigReconciler_ConcurrentForeignConditionSurvivesMergePatch(t *testing.T) { - t.Parallel() - - ctx := t.Context() - scheme := testutil.NewScheme(t) - - oidcConfig := &mcpv1beta1.MCPOIDCConfig{ - ObjectMeta: metav1.ObjectMeta{Name: "concurrent-config", Namespace: "default", Generation: 1}, - Spec: mcpv1beta1.MCPOIDCConfigSpec{ - Type: mcpv1beta1.MCPOIDCConfigTypeInline, - Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "https://accounts.google.com", - ClientID: "test-client", - }, - }, - } - key := client.ObjectKeyFromObject(oidcConfig) - - foreign := metav1.Condition{ - Type: "ForeignControllerSays", - Status: metav1.ConditionTrue, - Reason: "ExternallySet", - Message: "written by a concurrent owner between Get and Patch", - LastTransitionTime: metav1.Now(), - } - - var armed atomic.Bool - inject := interceptor.Funcs{ - Get: func(ctx context.Context, cl client.WithWatch, k client.ObjectKey, obj client.Object, opts ...client.GetOption) error { - if err := cl.Get(ctx, k, obj, opts...); err != nil { - return err - } - // Simulate a disjoint owner writing a condition into the store - // right after the reconciler's read of the config object. cl is the - // inner (non-intercepted) client, so this does not recurse. - if k == key && armed.CompareAndSwap(true, false) { - cur := &mcpv1beta1.MCPOIDCConfig{} - if err := cl.Get(ctx, k, cur); err != nil { - return err - } - meta.SetStatusCondition(&cur.Status.Conditions, foreign) - if err := cl.Status().Update(ctx, cur); err != nil { - return err - } - } - return nil - }, - } - - fakeClient := withOIDCConfigRefIndexes(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(oidcConfig). - WithStatusSubresource(&mcpv1beta1.MCPOIDCConfig{}). - WithInterceptorFuncs(inject). - Build() - r := &MCPOIDCConfigReconciler{Client: fakeClient, Scheme: scheme} - req := reconcile.Request{NamespacedName: key} - - // Reach steady state: finalizer added, then Valid=True + ConfigHash set. - _, err := r.Reconcile(ctx, req) - require.NoError(t, err) - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - // Introduce a referencing workload so the next reconcile patches only - // referencingWorkloads/referenceCount — not the conditions array, and - // without bumping the config's generation. - server := v1beta1test.NewMCPServer("referencing-server", "default", - v1beta1test.WithImage("img"), - v1beta1test.WithOIDCConfigRef(oidcConfig.Name, "aud"), - ) - require.NoError(t, fakeClient.Create(ctx, server)) - - // Arm the concurrent writer and run the reference-refresh reconcile. - armed.Store(true) - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - require.False(t, armed.Load(), "interceptor must have fired on the config Get") - - var after mcpv1beta1.MCPOIDCConfig - require.NoError(t, fakeClient.Get(ctx, key, &after)) - - foreignAfter := meta.FindStatusCondition(after.Status.Conditions, "ForeignControllerSays") - require.NotNil(t, foreignAfter, - "a concurrently-written foreign condition must survive the reconciler's merge-patch; a full PUT would erase it") - assert.Equal(t, metav1.ConditionTrue, foreignAfter.Status) - - // The reconciler's own work still landed in the same patch cycle. - own := meta.FindStatusCondition(after.Status.Conditions, mcpv1beta1.ConditionTypeOIDCConfigValid) - require.NotNil(t, own, "controller-owned Valid condition must remain") - assert.Equal(t, metav1.ConditionTrue, own.Status) - assert.EqualValues(t, 1, after.Status.ReferenceCount, "reference refresh must have been applied") -} - -// TestMCPOIDCConfigReconciler_watchHandlers verifies that the workload watch map -// functions enqueue exactly the config the workload currently references (or nothing -// when the workload has no ref). The previously-referenced config is enqueued by -// EnqueueRequestsFromMapFunc, which runs the map function on both the old and new -// object on update — no manual stale-reference scan in the handler. -func TestMCPOIDCConfigReconciler_watchHandlers(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - obj client.Object - expected map[string]struct{} - }{ - { - name: "MCPServer with ref enqueues the current config", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithOIDCConfigRef("current-config", "aud"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "VirtualMCPServer with ref enqueues the current config", - obj: v1beta1test.NewVirtualMCPServer("vmcp", "default", - v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ - Type: "oidc", - OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{Name: "current-config"}, - }), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPRemoteProxy with ref enqueues the current config", - obj: v1beta1test.NewMCPRemoteProxy("proxy", "default", - v1beta1test.WithRemoteProxyURL("https://example.com"), - v1beta1test.WithRemoteProxyOIDCConfigRef("current-config", "aud"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPServer without ref enqueues nothing", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - ), - expected: map[string]struct{}{}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := t.Context() - r, _ := newTestMCPOIDCConfigReconciler(t, tt.obj) - - requests := func() []reconcile.Request { - switch tt.obj.(type) { - case *mcpv1beta1.MCPServer: - return r.mapMCPServerToOIDCConfig(ctx, tt.obj) - case *mcpv1beta1.VirtualMCPServer: - return r.mapVirtualMCPServerToOIDCConfig(ctx, tt.obj) - case *mcpv1beta1.MCPRemoteProxy: - return r.mapMCPRemoteProxyToOIDCConfig(ctx, tt.obj) - default: - t.Fatalf("unexpected object type %T", tt.obj) - return nil - } - }() - - got := make(map[string]struct{}, len(requests)) - for _, req := range requests { - assert.Equal(t, "default", req.Namespace) - got[req.Name] = struct{}{} - } - assert.Equal(t, tt.expected, got) - }) - } -} diff --git a/cmd/thv-operator/controllers/mcpremoteproxy_controller.go b/cmd/thv-operator/controllers/mcpremoteproxy_controller.go index b8953dd6f4..2353a415ba 100644 --- a/cmd/thv-operator/controllers/mcpremoteproxy_controller.go +++ b/cmd/thv-operator/controllers/mcpremoteproxy_controller.go @@ -1030,13 +1030,11 @@ func (r *MCPRemoteProxyReconciler) handleOIDCConfig(ctx context.Context, proxy * return err } - // ReferencingWorkloads on the MCPOIDCConfig is maintained solely by the - // MCPOIDCConfig controller, which watches MCPServer/VirtualMCPServer/ - // MCPRemoteProxy and recomputes the full list (additions and removals). The - // MCPRemoteProxy controller must not write the config's status: a full - // r.Status().Update here would clobber conditions the config controller - // owns, and the previous append-only write never removed stale entries. - // See #5511. + // The MCPRemoteProxy controller must not write the MCPOIDCConfig's status: + // that status is owned by the MCPOIDCConfig controller, and a full + // r.Status().Update here would clobber conditions it owns. The config + // controller no longer tracks referencing workloads in its status; deletion + // protection recomputes referrers on demand in its finalizer. See #5511. // Detect whether the condition is transitioning to True (e.g. recovering from // a transient error). Without this check the status update is skipped when the @@ -1958,8 +1956,8 @@ func (r *MCPRemoteProxyReconciler) mapMCPGroupToMCPRemoteProxy( // handleAuthzConfig validates the referenced MCPAuthzConfig, tracks its hash on // the MCPRemoteProxy status, and sets the AuthzConfigRefValidated condition. When // the ref is cleared it removes both the hash and the condition so a stale "valid" -// signal does not linger. ReferencingWorkloads on the MCPAuthzConfig is owned by -// the MCPAuthzConfig controller (#5511); this controller never writes it. +// signal does not linger. The MCPAuthzConfig's status is owned by the +// MCPAuthzConfig controller (#5511); this controller never writes it. func (r *MCPRemoteProxyReconciler) handleAuthzConfig(ctx context.Context, proxy *mcpv1beta1.MCPRemoteProxy) error { if proxy.Spec.AuthzConfigRef == nil { // No MCPAuthzConfig referenced: clear any stored hash and remove the diff --git a/cmd/thv-operator/controllers/mcpserver_controller.go b/cmd/thv-operator/controllers/mcpserver_controller.go index 86fc12ce70..5c7fc1b167 100644 --- a/cmd/thv-operator/controllers/mcpserver_controller.go +++ b/cmd/thv-operator/controllers/mcpserver_controller.go @@ -2483,13 +2483,11 @@ func (r *MCPServerReconciler) handleOIDCConfig(ctx context.Context, m *mcpv1beta return err } - // ReferencingWorkloads on the MCPOIDCConfig is maintained solely by the - // MCPOIDCConfig controller, which watches MCPServer/VirtualMCPServer/ - // MCPRemoteProxy and recomputes the full list (additions and removals). The - // MCPServer controller must not write the config's status: a full - // r.Status().Update here would clobber conditions the config controller - // owns, and the previous append-only write never removed stale entries. - // See #5511. + // The MCPServer controller must not write the MCPOIDCConfig's status: that + // status is owned by the MCPOIDCConfig controller, and a full + // r.Status().Update here would clobber conditions it owns. The config + // controller no longer tracks referencing workloads in its status; deletion + // protection recomputes referrers on demand in its finalizer. See #5511. // Detect whether the condition is transitioning to True (e.g. recovering from // a transient error). Without this check the status update is skipped when the @@ -2579,8 +2577,8 @@ func setOIDCConfigRefCondition(m *mcpv1beta1.MCPServer, status metav1.ConditionS // handleAuthzConfig validates the referenced MCPAuthzConfig, tracks its hash on // the MCPServer status, and sets the AuthzConfigRefValidated condition. When the // ref is cleared it removes both the hash and the condition so a stale "valid" -// signal does not linger. ReferencingWorkloads on the MCPAuthzConfig is owned by -// the MCPAuthzConfig controller (#5511); this controller never writes it. +// signal does not linger. The MCPAuthzConfig's status is owned by the +// MCPAuthzConfig controller (#5511); this controller never writes it. // // Revocation semantics (fail-stale, not fail-open): if a previously-valid ref // later becomes invalid or missing, this returns an error and Reconcile stops diff --git a/cmd/thv-operator/controllers/mcpserver_oidcconfig_test.go b/cmd/thv-operator/controllers/mcpserver_oidcconfig_test.go index 8f360d5ed4..fa0cc4fca1 100644 --- a/cmd/thv-operator/controllers/mcpserver_oidcconfig_test.go +++ b/cmd/thv-operator/controllers/mcpserver_oidcconfig_test.go @@ -294,15 +294,13 @@ func TestMCPOIDCConfigReconciler_handleDeletion_BlocksWhenReferenced(t *testing. assert.Contains(t, cfg.Finalizers, OIDCConfigFinalizerName, "finalizer must remain") // The DeletionBlocked condition is written through MutateAndPatchStatus; - // re-fetch to confirm it (and the referencing-workload bookkeeping) was - // persisted rather than only mutated in memory. + // re-fetch to confirm it was persisted rather than only mutated in memory. var after mcpv1beta1.MCPOIDCConfig require.NoError(t, fc.Get(ctx, client.ObjectKeyFromObject(cfg), &after)) blocked := meta.FindStatusCondition(after.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) require.NotNil(t, blocked, "DeletionBlocked condition must be set while referenced") assert.Equal(t, metav1.ConditionTrue, blocked.Status) assert.Equal(t, "ReferencedByWorkloads", blocked.Reason) - assert.EqualValues(t, 1, after.Status.ReferenceCount, "referencing workload must be recorded") } func TestMCPOIDCConfigReconciler_handleDeletion_AllowsWhenNotReferenced(t *testing.T) { @@ -374,20 +372,20 @@ func conditionStatusPtr(s metav1.ConditionStatus) *metav1.ConditionStatus { // TestMCPServerReconciler_handleOIDCConfig_DoesNotWriteConfigStatus guards the // trust boundary consolidated in #5511: the MCPServer controller may read the // referenced MCPOIDCConfig but must never write its status. The MCPOIDCConfig -// controller is the sole owner of that status (conditions and -// ReferencingWorkloads). This unit test is the actual enforcement of that -// boundary — RBAC does not enforce it, because the operator runs as a single -// ServiceAccount whose aggregated role still grants mcpoidcconfigs/status write -// for the MCPOIDCConfig controller's own use. A future reintroduction of a -// cross-controller status write would flip the ResourceVersion and fail here. +// controller is the sole owner of that status. This unit test is the actual +// enforcement of that boundary — RBAC does not enforce it, because the operator +// runs as a single ServiceAccount whose aggregated role still grants +// mcpoidcconfigs/status write for the MCPOIDCConfig controller's own use. A +// future reintroduction of a cross-controller status write would flip the +// ResourceVersion and fail here. func TestMCPServerReconciler_handleOIDCConfig_DoesNotWriteConfigStatus(t *testing.T) { t.Parallel() ctx := t.Context() scheme := testutil.NewScheme(t) - // Seed a config whose status carries a condition and a ReferencingWorkloads - // entry owned by the MCPOIDCConfig controller — neither must be touched. + // Seed a config whose status carries conditions owned by the MCPOIDCConfig + // controller — none must be touched. oidcConfig := &mcpv1beta1.MCPOIDCConfig{ ObjectMeta: metav1.ObjectMeta{Name: "cfg", Namespace: "default"}, Spec: mcpv1beta1.MCPOIDCConfigSpec{ @@ -406,8 +404,6 @@ func TestMCPServerReconciler_handleOIDCConfig_DoesNotWriteConfigStatus(t *testin Reason: "ExternallySet", LastTransitionTime: metav1.Now(), }, }, - ReferencingWorkloads: []mcpv1beta1.WorkloadReference{{Kind: "MCPServer", Name: "someone-else"}}, - ReferenceCount: 1, }, } mcpServer := &mcpv1beta1.MCPServer{ @@ -435,8 +431,6 @@ func TestMCPServerReconciler_handleOIDCConfig_DoesNotWriteConfigStatus(t *testin assert.Equal(t, before.ResourceVersion, after.ResourceVersion, "MCPServer reconcile must not write the MCPOIDCConfig — its status is owned by the MCPOIDCConfig controller") - assert.Equal(t, before.Status.ReferencingWorkloads, after.Status.ReferencingWorkloads, - "MCPServer reconcile must not touch the config's ReferencingWorkloads") assert.NotNil(t, meta.FindStatusCondition(after.Status.Conditions, "ForeignControllerSays"), "config-owned conditions must remain untouched") } diff --git a/cmd/thv-operator/controllers/mcptelemetryconfig_controller.go b/cmd/thv-operator/controllers/mcptelemetryconfig_controller.go index 02cb84f6af..2d579b4582 100644 --- a/cmd/thv-operator/controllers/mcptelemetryconfig_controller.go +++ b/cmd/thv-operator/controllers/mcptelemetryconfig_controller.go @@ -12,15 +12,10 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" - "k8s.io/apimachinery/pkg/types" ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" - "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/predicate" - "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil" @@ -37,7 +32,8 @@ const ( // MCPTelemetryConfigReconciler reconciles a MCPTelemetryConfig object. // // This controller manages the lifecycle of MCPTelemetryConfig resources: validation, -// config hash computation, finalizer management, reference tracking, and deletion protection. +// config hash computation, finalizer management, and deletion protection. The set of +// referencing workloads is recomputed on demand during deletion rather than tracked in status. type MCPTelemetryConfigReconciler struct { client.Client Scheme *runtime.Scheme @@ -109,17 +105,9 @@ func (r *MCPTelemetryConfigReconciler) Reconcile(ctx context.Context, req ctrl.R // Calculate the hash of the current configuration configHash := r.calculateConfigHash(telemetryConfig.Spec) - // Track referencing workloads - referencingWorkloads, err := r.findReferencingWorkloads(ctx, telemetryConfig) - if err != nil { - logger.Error(err, "Failed to find referencing workloads") - return ctrl.Result{}, err - } - // Check what changed hashChanged := telemetryConfig.Status.ConfigHash != configHash - refsChanged := !ctrlutil.WorkloadRefsEqual(telemetryConfig.Status.ReferencingWorkloads, referencingWorkloads) - needsUpdate := hashChanged || refsChanged || conditionChanged + needsUpdate := hashChanged || conditionChanged if hashChanged { logger.Info("MCPTelemetryConfig configuration changed", @@ -130,7 +118,6 @@ func (r *MCPTelemetryConfigReconciler) Reconcile(ctx context.Context, req ctrl.R if needsUpdate { telemetryConfig.Status.ConfigHash = configHash telemetryConfig.Status.ObservedGeneration = telemetryConfig.Generation - telemetryConfig.Status.ReferencingWorkloads = referencingWorkloads if err := r.Status().Update(ctx, telemetryConfig); err != nil { logger.Error(err, "Failed to update MCPTelemetryConfig status") @@ -178,7 +165,6 @@ func indexVirtualMCPServerByTelemetryConfigRef(obj client.Object) []string { } // SetupWithManager sets up the controller with the Manager. -// Watches MCPServer changes to maintain accurate ReferencingWorkloads status. func (r *MCPTelemetryConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { // Field indexes backing findReferencingWorkloads: each lets the controller // query only the workloads referencing a given config rather than listing @@ -199,91 +185,11 @@ func (r *MCPTelemetryConfigReconciler) SetupWithManager(mgr ctrl.Manager) error return fmt.Errorf("failed to set up VirtualMCPServer telemetryConfigRef index: %w", err) } - // Watch MCPServer changes to update ReferencingWorkloads on referenced MCPTelemetryConfigs. - // The map function only returns the MCPTelemetryConfig the server currently references. - // EnqueueRequestsFromMapFunc runs it on both the old and new object on update (and - // on the object for create/delete), so removing or changing the ref enqueues both - // the previously- and newly-referenced config — the previously-referenced config - // then reconciles and prunes the stale entry. No manual stale-reference scan needed. - // - // GenerationChangedPredicate also suppresses the workload-watch resync; the self-heal - // backstop for a stale ReferencingWorkloads entry (e.g. a workload deleted while the - // operator was down) is this config's own For() resync, which re-runs Reconcile and - // rebuilds ReferencingWorkloads from the index. return ctrl.NewControllerManagedBy(mgr). For(&mcpv1beta1.MCPTelemetryConfig{}). - Watches(&mcpv1beta1.MCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPServerToTelemetryConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{})). - Watches( - &mcpv1beta1.MCPRemoteProxy{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPRemoteProxyToTelemetryConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). - Watches( - &mcpv1beta1.VirtualMCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapVirtualMCPServerToTelemetryConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). Complete(r) } -// mapMCPServerToTelemetryConfig maps an MCPServer to the MCPTelemetryConfig it currently -// references. EnqueueRequestsFromMapFunc invokes this on both the old and new object on -// update (and on the object for create/delete), so a ref change or deletion automatically -// enqueues both the previously- and newly-referenced config; the previously-referenced -// config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -func (*MCPTelemetryConfigReconciler) mapMCPServerToTelemetryConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - server, ok := obj.(*mcpv1beta1.MCPServer) - if !ok || server.Spec.TelemetryConfigRef == nil || server.Spec.TelemetryConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: server.Spec.TelemetryConfigRef.Name, - Namespace: server.Namespace, - }}} -} - -// mapMCPRemoteProxyToTelemetryConfig maps an MCPRemoteProxy to the MCPTelemetryConfig it -// currently references. EnqueueRequestsFromMapFunc invokes this on both the old and new -// object on update (and on the object for create/delete), so a ref change or deletion -// automatically enqueues both the previously- and newly-referenced config; the -// previously-referenced config then prunes the stale entry on reconcile. No manual -// stale-reference scan needed. -func (*MCPTelemetryConfigReconciler) mapMCPRemoteProxyToTelemetryConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - proxy, ok := obj.(*mcpv1beta1.MCPRemoteProxy) - if !ok || proxy.Spec.TelemetryConfigRef == nil || proxy.Spec.TelemetryConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: proxy.Spec.TelemetryConfigRef.Name, - Namespace: proxy.Namespace, - }}} -} - -// mapVirtualMCPServerToTelemetryConfig maps a VirtualMCPServer to the MCPTelemetryConfig -// it currently references. EnqueueRequestsFromMapFunc invokes this on both the old and new -// object on update (and on the object for create/delete), so a ref change or deletion -// automatically enqueues both the previously- and newly-referenced config; the -// previously-referenced config then prunes the stale entry on reconcile. No manual -// stale-reference scan needed. -func (*MCPTelemetryConfigReconciler) mapVirtualMCPServerToTelemetryConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - vmcp, ok := obj.(*mcpv1beta1.VirtualMCPServer) - if !ok || vmcp.Spec.TelemetryConfigRef == nil || vmcp.Spec.TelemetryConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: vmcp.Spec.TelemetryConfigRef.Name, - Namespace: vmcp.Namespace, - }}} -} - // calculateConfigHash calculates a hash of the MCPTelemetryConfig spec using Kubernetes utilities func (*MCPTelemetryConfigReconciler) calculateConfigHash(spec mcpv1beta1.MCPTelemetryConfigSpec) string { return ctrlutil.CalculateConfigHash(spec) diff --git a/cmd/thv-operator/controllers/mcptelemetryconfig_controller_test.go b/cmd/thv-operator/controllers/mcptelemetryconfig_controller_test.go index 78cd099aae..c77440105a 100644 --- a/cmd/thv-operator/controllers/mcptelemetryconfig_controller_test.go +++ b/cmd/thv-operator/controllers/mcptelemetryconfig_controller_test.go @@ -11,7 +11,6 @@ import ( "github.com/stretchr/testify/require" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" - "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" "sigs.k8s.io/controller-runtime/pkg/reconcile" @@ -601,73 +600,6 @@ func TestMCPTelemetryConfigReconciler_ConditionOnlyUpdate(t *testing.T) { assert.True(t, foundValid, "Should have Valid=True condition after condition-only update") } -func TestMCPTelemetryConfigReconciler_ReferenceTracking(t *testing.T) { - t.Parallel() - - ctx := t.Context() - - scheme := testutil.NewScheme(t) - - telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "shared-config", - Namespace: "default", - Generation: 1, - }, - Spec: newTelemetrySpec("https://otel-collector:4317", true, false), - } - - // Two MCPServers reference this config, one does not - server1 := v1beta1test.NewMCPServer("server-a", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithTelemetryConfigRef("shared-config"), - ) - server2 := v1beta1test.NewMCPServer("server-b", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithTelemetryConfigRef("shared-config"), - ) - server3 := v1beta1test.NewMCPServer("server-c", "default", - v1beta1test.WithImage("test-image"), - // No TelemetryConfigRef - ) - - fakeClient := withTelemetryConfigRefIndexes(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(telemetryConfig, server1, server2, server3). - WithStatusSubresource(&mcpv1beta1.MCPTelemetryConfig{}). - Build() - - r := &MCPTelemetryConfigReconciler{ - Client: fakeClient, - Scheme: scheme, - } - - req := reconcile.Request{ - NamespacedName: types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, - } - - // First reconcile: add finalizer - result, err := r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0), "Should requeue after adding finalizer") - - // Second reconcile: set hash, condition, and referencing servers - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - var updated mcpv1beta1.MCPTelemetryConfig - err = fakeClient.Get(ctx, req.NamespacedName, &updated) - require.NoError(t, err) - - // ReferencingWorkloads should list server-a and server-b (sorted), but not server-c - assert.Equal(t, []mcpv1beta1.WorkloadReference{ - {Kind: "MCPServer", Name: "server-a"}, - {Kind: "MCPServer", Name: "server-b"}, - }, updated.Status.ReferencingWorkloads) -} - func TestMCPTelemetryConfigReconciler_handleDeletion_BlocksWhenReferenced(t *testing.T) { t.Parallel() @@ -801,84 +733,3 @@ func newTelemetrySpec(endpoint string, tracing, metrics bool) mcpv1beta1.MCPTele }, } } - -// TestMCPTelemetryConfigReconciler_watchHandlers verifies that the workload watch map -// functions enqueue exactly the config the workload currently references (or nothing -// when the workload has no ref). The previously-referenced config is enqueued by -// EnqueueRequestsFromMapFunc, which runs the map function on both the old and new -// object on update — no manual stale-reference scan in the handler. -func TestMCPTelemetryConfigReconciler_watchHandlers(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - obj client.Object - expected map[string]struct{} - }{ - { - name: "MCPServer with ref enqueues the current config", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithTelemetryConfigRef("current-config"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPRemoteProxy with ref enqueues the current config", - obj: v1beta1test.NewMCPRemoteProxy("proxy", "default", - v1beta1test.WithRemoteProxyURL("https://example.com"), - v1beta1test.WithRemoteProxyTelemetryConfigRef("current-config"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "VirtualMCPServer with ref enqueues the current config", - obj: v1beta1test.NewVirtualMCPServer("vmcp", "default", - v1beta1test.WithVMCPTelemetryConfigRef("current-config"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPServer without ref enqueues nothing", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - ), - expected: map[string]struct{}{}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := t.Context() - scheme := testutil.NewScheme(t) - fakeClient := withTelemetryConfigRefIndexes(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(tt.obj). - WithStatusSubresource(&mcpv1beta1.MCPTelemetryConfig{}). - Build() - r := &MCPTelemetryConfigReconciler{Client: fakeClient, Scheme: scheme} - - requests := func() []reconcile.Request { - switch tt.obj.(type) { - case *mcpv1beta1.MCPServer: - return r.mapMCPServerToTelemetryConfig(ctx, tt.obj) - case *mcpv1beta1.MCPRemoteProxy: - return r.mapMCPRemoteProxyToTelemetryConfig(ctx, tt.obj) - case *mcpv1beta1.VirtualMCPServer: - return r.mapVirtualMCPServerToTelemetryConfig(ctx, tt.obj) - default: - t.Fatalf("unexpected object type %T", tt.obj) - return nil - } - }() - - got := make(map[string]struct{}, len(requests)) - for _, req := range requests { - assert.Equal(t, "default", req.Namespace) - got[req.Name] = struct{}{} - } - assert.Equal(t, tt.expected, got) - }) - } -} diff --git a/cmd/thv-operator/controllers/mcpwebhookconfig_controller.go b/cmd/thv-operator/controllers/mcpwebhookconfig_controller.go index ad97b7796d..3d424f1128 100644 --- a/cmd/thv-operator/controllers/mcpwebhookconfig_controller.go +++ b/cmd/thv-operator/controllers/mcpwebhookconfig_controller.go @@ -13,13 +13,9 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" - "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/predicate" - "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1alpha1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1alpha1" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" @@ -119,7 +115,7 @@ func (r *MCPWebhookConfigReconciler) Reconcile(ctx context.Context, req ctrl.Req } } - return r.updateReferencingWorkloads(ctx, webhookConfig) + return ctrl.Result{}, nil } // calculateConfigHash calculates a hash of the MCPWebhookConfig spec @@ -138,18 +134,9 @@ func (r *MCPWebhookConfigReconciler) handleConfigHashChange( "oldHash", webhookConfig.Status.ConfigHash, "newHash", configHash) - referencingServers, err := r.findReferencingMCPServers(ctx, webhookConfig) - if err != nil { - logger.Error(err, "Failed to find referencing MCPServers") - return ctrl.Result{}, fmt.Errorf("failed to find referencing MCPServers: %w", err) - } - - refs := workloadRefsFromMCPServers(referencingServers) - if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, webhookConfig, func(cfg *mcpv1alpha1.MCPWebhookConfig) { cfg.Status.ConfigHash = configHash cfg.Status.ObservedGeneration = cfg.Generation - cfg.Status.ReferencingWorkloads = refs meta.SetStatusCondition(&cfg.Status.Conditions, metav1.Condition{ Type: mcpv1beta1.ConditionTypeValid, Status: metav1.ConditionTrue, @@ -193,7 +180,6 @@ func (r *MCPWebhookConfigReconciler) handleDeletion( Message: fmt.Sprintf("Cannot delete: referenced by workloads: %v", refs), ObservedGeneration: cfg.Generation, }) - cfg.Status.ReferencingWorkloads = refs }); err != nil { logger.Error(err, "Failed to update MCPWebhookConfig status during deletion") return ctrl.Result{}, err @@ -206,7 +192,6 @@ func (r *MCPWebhookConfigReconciler) handleDeletion( if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, webhookConfig, func(cfg *mcpv1alpha1.MCPWebhookConfig) { meta.RemoveStatusCondition(&cfg.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) - cfg.Status.ReferencingWorkloads = nil }); err != nil { logger.Error(err, "Failed to clear MCPWebhookConfig deletion block status") return ctrl.Result{}, err @@ -263,33 +248,6 @@ func (r *MCPWebhookConfigReconciler) findReferencingMCPServers( return serverList.Items, nil } -// updateReferencingWorkloads updates the list of workloads referencing this config -func (r *MCPWebhookConfigReconciler) updateReferencingWorkloads( - ctx context.Context, - webhookConfig *mcpv1alpha1.MCPWebhookConfig, -) (ctrl.Result, error) { - logger := log.FromContext(ctx) - - referencingServers, err := r.findReferencingMCPServers(ctx, webhookConfig) - if err != nil { - logger.Error(err, "Failed to find referencing MCPServers") - return ctrl.Result{}, err - } - - refs := workloadRefsFromMCPServers(referencingServers) - if !ctrlutil.WorkloadRefsEqual(webhookConfig.Status.ReferencingWorkloads, refs) { - if err := ctrlutil.MutateAndPatchStatus(ctx, r.Client, webhookConfig, - func(cfg *mcpv1alpha1.MCPWebhookConfig) { - cfg.Status.ReferencingWorkloads = refs - }); err != nil { - logger.Error(err, "Failed to update referencing workloads list") - return ctrl.Result{}, err - } - } - - return ctrl.Result{}, nil -} - func workloadRefsFromMCPServers(servers []mcpv1beta1.MCPServer) []mcpv1beta1.WorkloadReference { refs := make([]mcpv1beta1.WorkloadReference, 0, len(servers)) for _, server := range servers { @@ -315,34 +273,7 @@ func (r *MCPWebhookConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { return fmt.Errorf("failed to set up MCPServer webhookConfigRef index: %w", err) } - // GenerationChangedPredicate also suppresses the workload-watch resync; the self-heal - // backstop for a stale ReferencingWorkloads entry (e.g. a workload deleted while the - // operator was down) is this config's own For() resync, which re-runs Reconcile and - // rebuilds ReferencingWorkloads from the index. return ctrl.NewControllerManagedBy(mgr). For(&mcpv1alpha1.MCPWebhookConfig{}). - Watches( - &mcpv1beta1.MCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPServerToWebhookConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{}), - ). Complete(r) } - -// mapMCPServerToWebhookConfig maps an MCPServer to the MCPWebhookConfig it currently references. -// EnqueueRequestsFromMapFunc invokes this on both the old and new object on update (and on the -// object for create/delete), so a ref change or deletion automatically enqueues both the -// previously- and newly-referenced config; the previously-referenced config then prunes the stale -// entry on reconcile. No manual stale-reference scan needed. -func (*MCPWebhookConfigReconciler) mapMCPServerToWebhookConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - server, ok := obj.(*mcpv1beta1.MCPServer) - if !ok || server.Spec.WebhookConfigRef == nil || server.Spec.WebhookConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: client.ObjectKey{ - Name: server.Spec.WebhookConfigRef.Name, - Namespace: server.Namespace, - }}} -} diff --git a/cmd/thv-operator/controllers/mcpwebhookconfig_controller_test.go b/cmd/thv-operator/controllers/mcpwebhookconfig_controller_test.go index 92474777a2..24531d84f1 100644 --- a/cmd/thv-operator/controllers/mcpwebhookconfig_controller_test.go +++ b/cmd/thv-operator/controllers/mcpwebhookconfig_controller_test.go @@ -4,21 +4,17 @@ package controllers import ( - "context" - "errors" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" - "sigs.k8s.io/controller-runtime/pkg/client/interceptor" "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1alpha1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1alpha1" @@ -131,11 +127,6 @@ func TestMCPWebhookConfigReconciler_Reconcile(t *testing.T) { assert.NotEmpty(t, updatedConfig.Status.ConfigHash) } if tt.existingMCPServer != nil { - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, mcpv1beta1.WorkloadReference{ - Kind: mcpv1beta1.WorkloadKindMCPServer, - Name: tt.existingMCPServer.Name, - }) - var updatedServer mcpv1beta1.MCPServer require.NoError(t, fakeClient.Get(ctx, types.NamespacedName{ Name: tt.existingMCPServer.Name, @@ -252,7 +243,6 @@ func TestMCPWebhookConfigReconciler_handleDeletion(t *testing.T) { require.NotNil(t, cond) assert.Equal(t, metav1.ConditionTrue, cond.Status) assert.Equal(t, "ReferencedByWorkloads", cond.Reason) - assert.NotEmpty(t, blockedConfig.Status.ReferencingWorkloads) // Delete server and try again require.NoError(t, fakeClient.Delete(ctx, mcpServer)) @@ -268,7 +258,6 @@ func TestMCPWebhookConfigReconciler_handleDeletion(t *testing.T) { require.NoError(t, err) assert.Empty(t, deletedConfig.Finalizers) assert.Nil(t, meta.FindStatusCondition(deletedConfig.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked)) - assert.Empty(t, deletedConfig.Status.ReferencingWorkloads) } func TestMCPWebhookConfigReconciler_handleDeletionClearsBlockedCondition(t *testing.T) { @@ -283,10 +272,6 @@ func TestMCPWebhookConfigReconciler_handleDeletionClearsBlockedCondition(t *test Finalizers: []string{WebhookConfigFinalizerName}, }, Status: mcpv1beta1.MCPWebhookConfigStatus{ - ReferencingWorkloads: []mcpv1beta1.WorkloadReference{{ - Kind: mcpv1beta1.WorkloadKindMCPServer, - Name: "server1", - }}, Conditions: []metav1.Condition{{ Type: mcpv1beta1.ConditionTypeDeletionBlocked, Status: metav1.ConditionTrue, @@ -311,142 +296,6 @@ func TestMCPWebhookConfigReconciler_handleDeletionClearsBlockedCondition(t *test require.NoError(t, fakeClient.Get(ctx, client.ObjectKeyFromObject(webhookConfig), clearedConfig)) assert.Empty(t, clearedConfig.Finalizers) assert.Nil(t, meta.FindStatusCondition(clearedConfig.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked)) - assert.Empty(t, clearedConfig.Status.ReferencingWorkloads) -} - -// TestMCPWebhookConfigReconciler_mapMCPServerToWebhookConfig verifies the map function -// enqueues exactly the config the MCPServer currently references (or nothing when it has -// no ref). The previously-referenced config is enqueued by EnqueueRequestsFromMapFunc, -// which runs the map function on both the old and new object on update — no manual -// stale-reference scan in the handler. -func TestMCPWebhookConfigReconciler_mapMCPServerToWebhookConfig(t *testing.T) { - t.Parallel() - - ctx := t.Context() - scheme := testutil.NewScheme(t) - - server := &mcpv1beta1.MCPServer{ - ObjectMeta: metav1.ObjectMeta{Name: "server", Namespace: "default"}, - Spec: mcpv1beta1.MCPServerSpec{ - WebhookConfigRef: &mcpv1beta1.WebhookConfigRef{Name: "current-config"}, - }, - } - - t.Run("current WebhookConfigRef returns one request", func(t *testing.T) { - t.Parallel() - - serverCopy := server.DeepCopy() - fakeClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(serverCopy). - Build() - r := &MCPWebhookConfigReconciler{Client: fakeClient, Scheme: scheme} - - require.ElementsMatch(t, []reconcile.Request{{ - NamespacedName: types.NamespacedName{Name: "current-config", Namespace: "default"}, - }}, r.mapMCPServerToWebhookConfig(ctx, serverCopy)) - }) - - t.Run("no WebhookConfigRef returns nil", func(t *testing.T) { - t.Parallel() - - serverWithoutRef := server.DeepCopy() - serverWithoutRef.Spec.WebhookConfigRef = nil - fakeClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(serverWithoutRef). - Build() - r := &MCPWebhookConfigReconciler{Client: fakeClient, Scheme: scheme} - - assert.Nil(t, r.mapMCPServerToWebhookConfig(ctx, serverWithoutRef)) - }) - - t.Run("wrong object type returns nil", func(t *testing.T) { - t.Parallel() - - fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build() - r := &MCPWebhookConfigReconciler{Client: fakeClient, Scheme: scheme} - pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "pod", Namespace: "default"}} - - assert.Nil(t, r.mapMCPServerToWebhookConfig(ctx, pod)) - }) -} - -func TestMCPWebhookConfigReconciler_updateReferencingWorkloads(t *testing.T) { - t.Parallel() - - ctx := t.Context() - scheme := testutil.NewScheme(t) - - webhookConfig := &mcpv1alpha1.MCPWebhookConfig{ - ObjectMeta: metav1.ObjectMeta{Name: "webhook-config", Namespace: "default"}, - } - referencingServer := v1beta1test.NewMCPServer("server", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithWebhookConfigRef(webhookConfig.Name), - ) - otherServer := v1beta1test.NewMCPServer("other-server", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithWebhookConfigRef("other-config"), - ) - fakeClient := withWebhookConfigRefIndex(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(webhookConfig, referencingServer, otherServer). - WithStatusSubresource(&mcpv1alpha1.MCPWebhookConfig{}). - Build() - r := &MCPWebhookConfigReconciler{Client: fakeClient, Scheme: scheme} - - result, err := r.updateReferencingWorkloads(ctx, webhookConfig) - require.NoError(t, err) - assert.Equal(t, reconcile.Result{}, result) - - var updated mcpv1alpha1.MCPWebhookConfig - require.NoError(t, fakeClient.Get(ctx, client.ObjectKeyFromObject(webhookConfig), &updated)) - assert.Equal(t, []mcpv1beta1.WorkloadReference{{ - Kind: mcpv1beta1.WorkloadKindMCPServer, - Name: referencingServer.Name, - }}, updated.Status.ReferencingWorkloads) - - result, err = r.updateReferencingWorkloads(ctx, &updated) - require.NoError(t, err) - assert.Equal(t, reconcile.Result{}, result) - - var unchanged mcpv1alpha1.MCPWebhookConfig - require.NoError(t, fakeClient.Get(ctx, client.ObjectKeyFromObject(webhookConfig), &unchanged)) - assert.Equal(t, updated.Status.ReferencingWorkloads, unchanged.Status.ReferencingWorkloads) -} - -func TestMCPWebhookConfigReconciler_updateReferencingWorkloadsListError(t *testing.T) { - t.Parallel() - - ctx := t.Context() - scheme := testutil.NewScheme(t) - - webhookConfig := &mcpv1alpha1.MCPWebhookConfig{ - ObjectMeta: metav1.ObjectMeta{Name: "webhook-config", Namespace: "default"}, - } - listErr := errors.New("simulated MCPServer list failure") - fakeClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(webhookConfig). - WithInterceptorFuncs(interceptor.Funcs{ - List: func( - ctx context.Context, - c client.WithWatch, - list client.ObjectList, - opts ...client.ListOption, - ) error { - if _, ok := list.(*mcpv1beta1.MCPServerList); ok { - return listErr - } - return c.List(ctx, list, opts...) - }, - }). - Build() - r := &MCPWebhookConfigReconciler{Client: fakeClient, Scheme: scheme} - - result, err := r.updateReferencingWorkloads(ctx, webhookConfig) - require.ErrorIs(t, err, listErr) - assert.Equal(t, reconcile.Result{}, result) } func findWebhookConfigCondition(conditions []metav1.Condition, conditionType string) *metav1.Condition { diff --git a/cmd/thv-operator/controllers/toolconfig_controller.go b/cmd/thv-operator/controllers/toolconfig_controller.go index 51f961030b..a69e58b5fc 100644 --- a/cmd/thv-operator/controllers/toolconfig_controller.go +++ b/cmd/thv-operator/controllers/toolconfig_controller.go @@ -12,15 +12,10 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" - "k8s.io/apimachinery/pkg/types" ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" - "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/predicate" - "sigs.k8s.io/controller-runtime/pkg/reconcile" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil" @@ -99,17 +94,6 @@ func (r *ToolConfigReconciler) Reconcile(ctx context.Context, req ctrl.Request) return r.handleConfigHashChange(ctx, toolConfig, configHash) } - // Refresh ReferencingWorkloads list - referencingWorkloads, err := r.findReferencingWorkloads(ctx, toolConfig) - if err != nil { - logger.Error(err, "Failed to find referencing workloads") - } else if !ctrlutil.WorkloadRefsEqual(toolConfig.Status.ReferencingWorkloads, referencingWorkloads) || - toolConfig.Status.ReferenceCount != workloadReferenceCount(referencingWorkloads) { - toolConfig.Status.ReferencingWorkloads = referencingWorkloads - toolConfig.Status.ReferenceCount = workloadReferenceCount(referencingWorkloads) - conditionChanged = true - } - // Update condition if it changed (even without hash change) if conditionChanged { if err := r.Status().Update(ctx, toolConfig); err != nil { @@ -130,22 +114,10 @@ func (r *ToolConfigReconciler) handleConfigHashChange( logger := log.FromContext(ctx) logger.Info("MCPToolConfig configuration changed", "oldHash", toolConfig.Status.ConfigHash, "newHash", configHash) - // Find all MCPServers that reference this MCPToolConfig - refs, err := r.findReferencingWorkloads(ctx, toolConfig) - if err != nil { - logger.Error(err, "Failed to find referencing MCPServers") - // Don't persist the new hash on error; returning the error requeues so - // the next reconcile recomputes referencing workloads. - return ctrl.Result{}, fmt.Errorf("failed to find referencing MCPServers: %w", err) - } - - // Update the status with the new hash only after successful server lookup + // Record the new observed hash and generation. toolConfig.Status.ConfigHash = configHash toolConfig.Status.ObservedGeneration = toolConfig.Generation - toolConfig.Status.ReferencingWorkloads = refs - toolConfig.Status.ReferenceCount = workloadReferenceCount(refs) - // Update the MCPToolConfig status if err := r.Status().Update(ctx, toolConfig); err != nil { logger.Error(err, "Failed to update MCPToolConfig status") @@ -184,8 +156,6 @@ func (r *ToolConfigReconciler) handleDeletion(ctx context.Context, toolConfig *m Message: fmt.Sprintf("Cannot delete: referenced by workloads: %v", referencingWorkloads), ObservedGeneration: toolConfig.Generation, }) - toolConfig.Status.ReferencingWorkloads = referencingWorkloads - toolConfig.Status.ReferenceCount = workloadReferenceCount(referencingWorkloads) if updateErr := r.Status().Update(ctx, toolConfig); updateErr != nil { logger.Error(updateErr, "Failed to update status during deletion block") } @@ -246,7 +216,6 @@ func (r *ToolConfigReconciler) findReferencingWorkloads( } // SetupWithManager sets up the controller with the Manager. -// Watches MCPServer changes to maintain accurate ReferencingWorkloads status. func (r *ToolConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { // Field index backing findReferencingWorkloads: lets the controller query only // the MCPServers referencing a given config rather than listing every MCPServer @@ -257,39 +226,7 @@ func (r *ToolConfigReconciler) SetupWithManager(mgr ctrl.Manager) error { return fmt.Errorf("failed to set up MCPServer toolConfigRef index: %w", err) } - // Watch MCPServer changes to update ReferencingWorkloads on referenced MCPToolConfigs. - // The map function only returns the MCPToolConfig the server currently references. - // EnqueueRequestsFromMapFunc runs it on both the old and new object on update (and - // on the object for create/delete), so removing or changing the ref enqueues both - // the previously- and newly-referenced config — the previously-referenced config - // then reconciles and prunes the stale entry. No manual stale-reference scan needed. - // - // GenerationChangedPredicate also suppresses the workload-watch resync; the self-heal - // backstop for a stale ReferencingWorkloads entry (e.g. a workload deleted while the - // operator was down) is this config's own For() resync, which re-runs Reconcile and - // rebuilds ReferencingWorkloads from the index. return ctrl.NewControllerManagedBy(mgr). For(&mcpv1beta1.MCPToolConfig{}). - Watches(&mcpv1beta1.MCPServer{}, - handler.EnqueueRequestsFromMapFunc(r.mapMCPServerToToolConfig), - builder.WithPredicates(predicate.GenerationChangedPredicate{})). Complete(r) } - -// mapMCPServerToToolConfig maps an MCPServer to the MCPToolConfig it currently -// references. EnqueueRequestsFromMapFunc invokes this on both the old and new object on -// update (and on the object for create/delete), so a ref change or deletion automatically -// enqueues both the previously- and newly-referenced config; the previously-referenced -// config then prunes the stale entry on reconcile. No manual stale-reference scan needed. -func (*ToolConfigReconciler) mapMCPServerToToolConfig( - _ context.Context, obj client.Object, -) []reconcile.Request { - server, ok := obj.(*mcpv1beta1.MCPServer) - if !ok || server.Spec.ToolConfigRef == nil || server.Spec.ToolConfigRef.Name == "" { - return nil - } - return []reconcile.Request{{NamespacedName: types.NamespacedName{ - Name: server.Spec.ToolConfigRef.Name, - Namespace: server.Namespace, - }}} -} diff --git a/cmd/thv-operator/controllers/toolconfig_controller_edge_cases_test.go b/cmd/thv-operator/controllers/toolconfig_controller_edge_cases_test.go index 91c05e32ff..9a21498345 100644 --- a/cmd/thv-operator/controllers/toolconfig_controller_edge_cases_test.go +++ b/cmd/thv-operator/controllers/toolconfig_controller_edge_cases_test.go @@ -112,8 +112,6 @@ func TestToolConfigReconciler_EdgeCases(t *testing.T) { err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) require.NoError(t, err) assert.NotEmpty(t, updatedConfig.Status.ConfigHash) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "test-server"}) }) t.Run("reconcile with changed spec", func(t *testing.T) { @@ -176,24 +174,30 @@ func TestToolConfigReconciler_EdgeCases(t *testing.T) { func TestToolConfigReconciler_ErrorScenarios(t *testing.T) { t.Parallel() - t.Run("error updating status", func(t *testing.T) { + t.Run("error listing workloads during deletion", func(t *testing.T) { t.Parallel() ctx := t.Context() scheme := testutil.NewScheme(t) + // Object under deletion: it carries the finalizer and a deletion + // timestamp, so Reconcile routes into handleDeletion, which recomputes + // the referencing workloads on demand (the only remaining caller of the + // referrer lookup now that the status list is no longer stored). toolConfig := &mcpv1beta1.MCPToolConfig{ ObjectMeta: metav1.ObjectMeta{ - Name: "test-config", - Namespace: "default", - Finalizers: []string{ToolConfigFinalizerName}, + Name: "test-config", + Namespace: "default", + Finalizers: []string{ToolConfigFinalizerName}, + DeletionTimestamp: &metav1.Time{Time: time.Unix(0, 0)}, }, Spec: mcpv1beta1.MCPToolConfigSpec{ ToolsFilter: []string{"tool1"}, }, } - // Create a fake client that returns an error when listing MCPServers + // Create a fake client that returns an error when listing MCPServers, + // so the deletion-time referrer lookup fails. fakeClient := &errorClient{ Client: fake.NewClientBuilder(). WithScheme(scheme). @@ -216,8 +220,8 @@ func TestToolConfigReconciler_ErrorScenarios(t *testing.T) { } result, err := r.Reconcile(ctx, req) - assert.Error(t, err) - assert.Contains(t, err.Error(), "failed to find referencing MCPServers") + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to list MCPServers by toolConfigRef") assert.Equal(t, time.Duration(0), result.RequeueAfter) }) } @@ -238,89 +242,6 @@ func (c *errorClient) List(ctx context.Context, list client.ObjectList, opts ... func TestToolConfigReconciler_ComplexScenarios(t *testing.T) { t.Parallel() - t.Run("multiple MCPServers referencing same MCPToolConfig", func(t *testing.T) { - t.Parallel() - ctx := t.Context() - - scheme := testutil.NewScheme(t) - - toolConfig := &mcpv1beta1.MCPToolConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "shared-config", - Namespace: "default", - }, - Spec: mcpv1beta1.MCPToolConfigSpec{ - ToolsFilter: []string{"tool1", "tool2", "tool3"}, - ToolsOverride: map[string]mcpv1beta1.ToolOverride{ - "tool1": { - Name: "custom-tool1", - Description: "Customized tool 1", - }, - }, - }, - } - - // Create multiple MCPServers referencing the same MCPToolConfig - mcpServers := []*mcpv1beta1.MCPServer{ - v1beta1test.NewMCPServer("server1", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithToolConfigRef("shared-config"), - ), - v1beta1test.NewMCPServer("server2", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithToolConfigRef("shared-config"), - ), - v1beta1test.NewMCPServer("server3", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithToolConfigRef("shared-config"), - ), - } - - objs := []client.Object{toolConfig} - for _, server := range mcpServers { - objs = append(objs, server) - } - - fakeClient := withToolConfigRefIndex(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(objs...). - WithStatusSubresource(&mcpv1beta1.MCPToolConfig{}). - Build() - - r := &ToolConfigReconciler{ - Client: fakeClient, - Scheme: scheme, - } - - req := reconcile.Request{ - NamespacedName: types.NamespacedName{ - Name: toolConfig.Name, - Namespace: toolConfig.Namespace, - }, - } - - // First reconciliation adds finalizer - result, err := r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0)) - - // Second reconciliation updates status - result, err = r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Equal(t, time.Duration(0), result.RequeueAfter) - - // Verify all servers are listed in status - var updatedConfig mcpv1beta1.MCPToolConfig - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Len(t, updatedConfig.Status.ReferencingWorkloads, 3) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "server1"}) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "server2"}) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "server3"}) - }) - t.Run("empty MCPToolConfig spec", func(t *testing.T) { t.Parallel() ctx := t.Context() diff --git a/cmd/thv-operator/controllers/toolconfig_controller_test.go b/cmd/thv-operator/controllers/toolconfig_controller_test.go index 7eb043ec04..0b38393f5a 100644 --- a/cmd/thv-operator/controllers/toolconfig_controller_test.go +++ b/cmd/thv-operator/controllers/toolconfig_controller_test.go @@ -210,13 +210,6 @@ func TestToolConfigReconciler_Reconcile(t *testing.T) { "MCPToolConfig status should have config hash") } - // Check referencing workloads in status - if tt.existingMCPServer != nil { - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: tt.existingMCPServer.Name}, - "Status should contain referencing MCPServer as WorkloadReference") - } - // Check Valid condition is set after successful reconciliation cond := k8smeta.FindStatusCondition(updatedConfig.Status.Conditions, mcpv1beta1.ConditionToolConfigValid) require.NotNil(t, cond, "Valid condition must be set after successful reconciliation") @@ -276,149 +269,6 @@ func TestToolConfigReconciler_findReferencingWorkloads(t *testing.T) { assert.NotContains(t, refs, mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "server3"}) } -func TestToolConfigReconciler_ReferencingWorkloadsUpdatedWithoutHashChange(t *testing.T) { - t.Parallel() - - ctx := t.Context() - - scheme := testutil.NewScheme(t) - - toolConfig := &mcpv1beta1.MCPToolConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-config", - Namespace: "default", - }, - Spec: mcpv1beta1.MCPToolConfigSpec{ - ToolsFilter: []string{"tool1"}, - }, - } - - fakeClient := withToolConfigRefIndex(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(toolConfig). - WithStatusSubresource(&mcpv1beta1.MCPToolConfig{}). - Build() - - r := &ToolConfigReconciler{ - Client: fakeClient, - Scheme: scheme, - } - - req := reconcile.Request{ - NamespacedName: types.NamespacedName{ - Name: toolConfig.Name, - Namespace: toolConfig.Namespace, - }, - } - - // First reconciliation - add finalizer - result, err := r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0)) - - // Second reconciliation - sets hash, no servers yet - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - var updatedConfig mcpv1beta1.MCPToolConfig - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.NotEmpty(t, updatedConfig.Status.ConfigHash) - assert.Empty(t, updatedConfig.Status.ReferencingWorkloads, "No servers should be referencing yet") - assert.EqualValues(t, 0, updatedConfig.Status.ReferenceCount) - - // Verify Valid condition is set after initial reconciliation - cond := k8smeta.FindStatusCondition(updatedConfig.Status.Conditions, mcpv1beta1.ConditionToolConfigValid) - require.NotNil(t, cond, "Valid condition must be set after reconciliation") - assert.Equal(t, metav1.ConditionTrue, cond.Status) - - // Add an MCPServer that references this config (without changing the config spec) - mcpServer := v1beta1test.NewMCPServer("new-server", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithToolConfigRef("test-config"), - ) - require.NoError(t, fakeClient.Create(ctx, mcpServer)) - - // Reconcile again - hash hasn't changed, but referencing servers should be updated - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "new-server"}, - "ReferencingWorkloads should be updated even without hash change") - assert.EqualValues(t, 1, updatedConfig.Status.ReferenceCount) -} - -func TestToolConfigReconciler_ReferencingWorkloadsRemovedOnServerDeletion(t *testing.T) { - t.Parallel() - - ctx := t.Context() - - scheme := testutil.NewScheme(t) - - toolConfig := &mcpv1beta1.MCPToolConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-config", - Namespace: "default", - }, - Spec: mcpv1beta1.MCPToolConfigSpec{ - ToolsFilter: []string{"tool1"}, - }, - } - - mcpServer := v1beta1test.NewMCPServer("server-to-delete", "default", - v1beta1test.WithImage("test-image"), - v1beta1test.WithToolConfigRef("test-config"), - ) - - fakeClient := withToolConfigRefIndex(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(toolConfig, mcpServer). - WithStatusSubresource(&mcpv1beta1.MCPToolConfig{}). - Build() - - r := &ToolConfigReconciler{ - Client: fakeClient, - Scheme: scheme, - } - - req := reconcile.Request{ - NamespacedName: types.NamespacedName{ - Name: toolConfig.Name, - Namespace: toolConfig.Namespace, - }, - } - - // Add finalizer - result, err := r.Reconcile(ctx, req) - require.NoError(t, err) - assert.Greater(t, result.RequeueAfter, time.Duration(0)) - - // Set hash and referencing servers - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - var updatedConfig mcpv1beta1.MCPToolConfig - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Contains(t, updatedConfig.Status.ReferencingWorkloads, - mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: "server-to-delete"}) - assert.EqualValues(t, 1, updatedConfig.Status.ReferenceCount) - - // Delete the MCPServer - require.NoError(t, fakeClient.Delete(ctx, mcpServer)) - - // Reconcile again - referencing servers should be empty now - _, err = r.Reconcile(ctx, req) - require.NoError(t, err) - - err = fakeClient.Get(ctx, req.NamespacedName, &updatedConfig) - require.NoError(t, err) - assert.Empty(t, updatedConfig.Status.ReferencingWorkloads, - "ReferencingWorkloads should be empty after server deletion") - assert.EqualValues(t, 0, updatedConfig.Status.ReferenceCount) -} - func TestToolConfigReconciler_ValidConditionObservedGeneration(t *testing.T) { t.Parallel() @@ -497,57 +347,3 @@ func TestToolConfigReconciler_ValidConditionObservedGeneration(t *testing.T) { assert.Equal(t, int64(2), cond.ObservedGeneration, "ObservedGeneration should be updated to match new Generation") } - -// TestMCPToolConfigReconciler_watchHandlers verifies that the MCPServer watch map -// function enqueues exactly the config the server currently references (or nothing -// when the server has no ref). The previously-referenced config is enqueued by -// EnqueueRequestsFromMapFunc, which runs the map function on both the old and new -// object on update — no manual stale-reference scan in the handler. -func TestMCPToolConfigReconciler_watchHandlers(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - obj client.Object - expected map[string]struct{} - }{ - { - name: "MCPServer with ref enqueues the current config", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - v1beta1test.WithToolConfigRef("current-config"), - ), - expected: map[string]struct{}{"current-config": {}}, - }, - { - name: "MCPServer without ref enqueues nothing", - obj: v1beta1test.NewMCPServer("srv", "default", - v1beta1test.WithImage("example/mcp:latest"), - ), - expected: map[string]struct{}{}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := t.Context() - scheme := testutil.NewScheme(t) - fakeClient := withToolConfigRefIndex(fake.NewClientBuilder().WithScheme(scheme)). - WithObjects(tt.obj). - WithStatusSubresource(&mcpv1beta1.MCPToolConfig{}). - Build() - r := &ToolConfigReconciler{Client: fakeClient, Scheme: scheme} - - requests := r.mapMCPServerToToolConfig(ctx, tt.obj) - - got := make(map[string]struct{}, len(requests)) - for _, req := range requests { - assert.Equal(t, "default", req.Namespace) - got[req.Name] = struct{}{} - } - assert.Equal(t, tt.expected, got) - }) - } -} diff --git a/cmd/thv-operator/controllers/virtualmcpserver_controller.go b/cmd/thv-operator/controllers/virtualmcpserver_controller.go index d18ba781a9..87c2812b97 100644 --- a/cmd/thv-operator/controllers/virtualmcpserver_controller.go +++ b/cmd/thv-operator/controllers/virtualmcpserver_controller.go @@ -1810,10 +1810,12 @@ func (r *VirtualMCPServerReconciler) podTemplateMetadataNeedsUpdate( return false } -// podTemplateSpecNeedsUpdate checks if the user-provided PodTemplateSpec has changed. -// Instead of comparing full rendered templates (which always differ due to Kubernetes-defaulted -// fields like terminationGracePeriodSeconds, dnsPolicy, etc.), this compares a SHA256 hash of -// the raw PodTemplateSpec input stored as a deployment annotation. +// podTemplateSpecNeedsUpdate checks if the user-provided PodTemplateSpec has changed, by +// comparing a SHA256 hash of the raw input against the stored annotation rather than the +// full rendered template (which always differs due to Kubernetes-defaulted fields). +// Symmetric by design (#5818): an absent PodTemplateSpec expects an empty hash, so a stale +// annotation left over after the field is cleared is treated as drift and gets pruned by +// mergeDeploymentAnnotations on the next write, instead of being flagged forever. func (*VirtualMCPServerReconciler) podTemplateSpecNeedsUpdate( ctx context.Context, deployment *appsv1.Deployment, @@ -1824,21 +1826,14 @@ func (*VirtualMCPServerReconciler) podTemplateSpecNeedsUpdate( return true } - // If no PodTemplateSpec is provided, update is only needed if one was previously applied - if vmcp.Spec.PodTemplateSpec == nil || vmcp.Spec.PodTemplateSpec.Raw == nil { - _, hadPrevious := deployment.Annotations[podTemplateSpecHashAnnotation] - return hadPrevious - } - - // Compare hash of the raw PodTemplateSpec input against the stored annotation. - // Avoids comparing full rendered templates which always differ due to - // Kubernetes-defaulted fields (terminationGracePeriodSeconds, dnsPolicy, etc.). - // Uses HashRawJSON to ensure deterministic hashing regardless of JSON field ordering. - expectedHash, err := checksum.HashRawJSON(vmcp.Spec.PodTemplateSpec.Raw) - if err != nil { - // If we can't hash, assume update is needed - log.FromContext(ctx).Error(err, "Failed to hash PodTemplateSpec, assuming update needed") - return true + expectedHash := "" + if vmcp.Spec.PodTemplateSpec != nil && len(vmcp.Spec.PodTemplateSpec.Raw) > 0 { + hash, err := checksum.HashRawJSON(vmcp.Spec.PodTemplateSpec.Raw) + if err != nil { + log.FromContext(ctx).Error(err, "Failed to hash PodTemplateSpec, assuming update needed") + return true + } + expectedHash = hash } if deployment.Annotations[podTemplateSpecHashAnnotation] != expectedHash { return true @@ -1866,18 +1861,16 @@ func (*VirtualMCPServerReconciler) podTemplateSpecNeedsUpdate( return false } -// mergeDeploymentAnnotations merges desired Deployment annotations onto the -// live ones via ctrlutil.MergeAnnotations, then prunes imagePullRefsHashAnnotation -// from the result if desired no longer wants it. MergeAnnotations preserves any -// live key absent from desired so externally-managed annotations survive, but -// imagePullRefsHashAnnotation is operator-owned: when the desired imagePullSecrets -// list is empty, desired carries no annotation for it, and a stale value left by -// a prior reconcile must be pruned explicitly or it survives forever and -// imagePullSecretsNeedsUpdate flags drift on every reconcile (#5817). +// mergeDeploymentAnnotations merges desired annotations onto the live ones via +// ctrlutil.MergeAnnotations, then prunes the operator-owned hash annotations +// (imagePullRefsHashAnnotation, podTemplateSpecHashAnnotation) that desired no longer wants — +// MergeAnnotations otherwise preserves them forever once their source field goes empty (#5817, #5818). func mergeDeploymentAnnotations(desired, live map[string]string) map[string]string { merged := ctrlutil.MergeAnnotations(desired, live) - if _, wantHash := desired[imagePullRefsHashAnnotation]; !wantHash { - delete(merged, imagePullRefsHashAnnotation) + for _, key := range []string{imagePullRefsHashAnnotation, podTemplateSpecHashAnnotation} { + if _, want := desired[key]; !want { + delete(merged, key) + } } return merged } @@ -3538,8 +3531,8 @@ func (r *VirtualMCPServerReconciler) handleConfigRefs( // (spec.incomingAuth.authzConfigRef), tracks its hash on the VirtualMCPServer // status, and sets the AuthzConfigRefValidated condition. When the ref is // cleared it removes both the hash and the condition so a stale "valid" signal -// does not linger. ReferencingWorkloads on the MCPAuthzConfig is owned by the -// MCPAuthzConfig controller (#5511); this controller never writes it. +// does not linger. The MCPAuthzConfig's status is owned by the MCPAuthzConfig +// controller (#5511); this controller never writes it. // // Revocation semantics (fail-stale, not fail-open): if a previously-valid ref // later becomes invalid or missing, this returns an error and Reconcile stops @@ -3724,13 +3717,11 @@ func (r *VirtualMCPServerReconciler) handleOIDCConfig( return fmt.Errorf("%s", msg) } - // ReferencingWorkloads on the MCPOIDCConfig is maintained solely by the - // MCPOIDCConfig controller, which watches MCPServer/VirtualMCPServer/ - // MCPRemoteProxy and recomputes the full list (additions and removals). The - // VirtualMCPServer controller must not write the config's status: a full - // r.Status().Update here would clobber conditions the config controller - // owns, and the previous append-only write never removed stale entries. - // See #5511. + // The VirtualMCPServer controller must not write the MCPOIDCConfig's status: + // that status is owned by the MCPOIDCConfig controller, and a full + // r.Status().Update here would clobber conditions it owns. The config + // controller no longer tracks referencing workloads in its status; deletion + // protection recomputes referrers on demand in its finalizer. See #5511. // Set valid condition statusManager.SetCondition( diff --git a/cmd/thv-operator/controllers/virtualmcpserver_controller_test.go b/cmd/thv-operator/controllers/virtualmcpserver_controller_test.go index 3f2eb0fee1..4f08951923 100644 --- a/cmd/thv-operator/controllers/virtualmcpserver_controller_test.go +++ b/cmd/thv-operator/controllers/virtualmcpserver_controller_test.go @@ -29,6 +29,7 @@ import ( "k8s.io/apimachinery/pkg/types" "k8s.io/apimachinery/pkg/util/intstr" "k8s.io/client-go/tools/events" + "k8s.io/utils/ptr" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" @@ -588,80 +589,6 @@ func TestVirtualMCPServerEnsureRBACResources_CustomServiceAccount(t *testing.T) assert.Error(t, err, "RoleBinding should not be created when custom ServiceAccount is provided") } -// TestVirtualMCPServerEnsureDeployment tests Deployment creation -func TestVirtualMCPServerEnsureDeployment(t *testing.T) { - t.Parallel() - - vmcp := v1beta1test.NewVirtualMCPServer(testVmcpName, "default", - v1beta1test.WithVMCPGroupRef(testGroupName), - ) - - // Create MCPGroup that the VirtualMCPServer references - mcpGroup := &mcpv1beta1.MCPGroup{ - ObjectMeta: metav1.ObjectMeta{ - Name: testGroupName, - Namespace: "default", - }, - Status: mcpv1beta1.MCPGroupStatus{ - Phase: mcpv1beta1.MCPGroupPhaseReady, - }, - } - - // Create ConfigMap with checksum - configMap := &corev1.ConfigMap{ - ObjectMeta: metav1.ObjectMeta{ - Name: vmcpConfigMapName(vmcp.Name), - Namespace: "default", - Annotations: map[string]string{ - "toolhive.stacklok.dev/content-checksum": "test-checksum-123", - }, - }, - Data: map[string]string{ - "config.yaml": "{}", - }, - } - - scheme := testutil.NewScheme(t) - - fakeClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(vmcp, mcpGroup, configMap). - Build() - - r := &VirtualMCPServerReconciler{ - Client: fakeClient, - Scheme: scheme, - PlatformDetector: ctrlutil.NewSharedPlatformDetector(), - } - - result, err := r.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) - require.NoError(t, err) - assert.Equal(t, ctrl.Result{}, result) - - // Verify Deployment was created - deployment := &appsv1.Deployment{} - err = fakeClient.Get(context.Background(), types.NamespacedName{ - Name: vmcp.Name, - Namespace: vmcp.Namespace, - }, deployment) - require.NoError(t, err) - assert.Equal(t, vmcp.Name, deployment.Name) - // spec.replicas is nil — nil-passthrough for HPA compatibility - assert.Nil(t, deployment.Spec.Replicas) - - // Verify container configuration - require.Len(t, deployment.Spec.Template.Spec.Containers, 1) - container := deployment.Spec.Template.Spec.Containers[0] - assert.Equal(t, "vmcp", container.Name) - assert.NotEmpty(t, container.Image) - assert.Contains(t, container.Args, "serve") - assert.Contains(t, container.Args, "--config=/etc/vmcp-config/config.yaml") - - // Verify checksum annotation is set using standard annotation key - assert.Equal(t, "test-checksum-123", - deployment.Spec.Template.Annotations[checksum.RunConfigChecksumAnnotation]) -} - // TestVirtualMCPServerEnsureService tests Service creation func TestVirtualMCPServerEnsureService(t *testing.T) { t.Parallel() @@ -2036,6 +1963,70 @@ func TestVirtualMCPServerDeploymentNeedsUpdate(t *testing.T) { }, expectedUpdate: true, }, + { + name: "stale podTemplateSpec hash annotation", + deployment: &appsv1.Deployment{ + ObjectMeta: metav1.ObjectMeta{ + Labels: labelsForVirtualMCPServer(vmcp.Name), + Annotations: map[string]string{podTemplateSpecHashAnnotation: "stale-hash"}, + }, + Spec: appsv1.DeploymentSpec{ + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: expectedLabels, + Annotations: expectedAnnotations, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{ + { + Name: "vmcp", + Image: getVmcpImage(), + Ports: []corev1.ContainerPort{ + {ContainerPort: 4483}, + }, + Args: reconciler.buildContainerArgsForVmcp(vmcp), + Env: mustBuildEnvVarsForVmcp(reconciler, vmcp), + }, + }, + ServiceAccountName: vmcpServiceAccountName(vmcp.Name), + }, + }, + }, + }, + expectedUpdate: true, + }, + { + name: "stale imagePullSecrets hash annotation", + deployment: &appsv1.Deployment{ + ObjectMeta: metav1.ObjectMeta{ + Labels: labelsForVirtualMCPServer(vmcp.Name), + Annotations: map[string]string{imagePullRefsHashAnnotation: "stale-hash"}, + }, + Spec: appsv1.DeploymentSpec{ + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: expectedLabels, + Annotations: expectedAnnotations, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{ + { + Name: "vmcp", + Image: getVmcpImage(), + Ports: []corev1.ContainerPort{ + {ContainerPort: 4483}, + }, + Args: reconciler.buildContainerArgsForVmcp(vmcp), + Env: mustBuildEnvVarsForVmcp(reconciler, vmcp), + }, + }, + ServiceAccountName: vmcpServiceAccountName(vmcp.Name), + }, + }, + }, + }, + expectedUpdate: true, + }, { name: "no changes - no update needed", deployment: &appsv1.Deployment{ @@ -2080,6 +2071,56 @@ func TestVirtualMCPServerDeploymentNeedsUpdate(t *testing.T) { } } +// Direct unit tests for podTemplateSpecNeedsUpdate live in +// virtualmcpserver_podtemplatespec_reconcile_test.go, and for +// imagePullSecretsNeedsUpdate in virtualmcpserver_deployment_test.go / +// virtualmcpserver_default_imagepullsecrets_test.go — no need to duplicate here. + +func TestMergeDeploymentAnnotations(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + desired map[string]string + live map[string]string + expected map[string]string + }{ + { + name: "prunes stale imagePullRefsHashAnnotation when desired no longer wants it", + desired: map[string]string{}, + live: map[string]string{imagePullRefsHashAnnotation: "stale-hash"}, + expected: map[string]string{}, + }, + { + name: "prunes stale podTemplateSpecHashAnnotation when desired no longer wants it", + desired: map[string]string{}, + live: map[string]string{podTemplateSpecHashAnnotation: "stale-hash"}, + expected: map[string]string{}, + }, + { + name: "keeps hash annotations desired still wants", + desired: map[string]string{imagePullRefsHashAnnotation: "new-hash", podTemplateSpecHashAnnotation: "new-hash"}, + live: map[string]string{imagePullRefsHashAnnotation: "old-hash", podTemplateSpecHashAnnotation: "old-hash"}, + expected: map[string]string{imagePullRefsHashAnnotation: "new-hash", podTemplateSpecHashAnnotation: "new-hash"}, + }, + { + name: "preserves externally-managed annotations absent from desired", + desired: map[string]string{}, + live: map[string]string{"external.io/managed-by": "someone-else"}, + expected: map[string]string{"external.io/managed-by": "someone-else"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + merged := mergeDeploymentAnnotations(tt.desired, tt.live) + assert.Equal(t, tt.expected, merged) + }) + } +} + func TestVirtualMCPServerReconcile_HappyPath(t *testing.T) { t.Parallel() @@ -2380,6 +2421,19 @@ func TestVirtualMCPServerEnsureDeployment_CreateDeployment(t *testing.T) { }, deployment) assert.NoError(t, err) assert.Equal(t, vmcp.Name, deployment.Name) + // spec.replicas is nil — nil-passthrough for HPA compatibility + assert.Nil(t, deployment.Spec.Replicas) + + require.Len(t, deployment.Spec.Template.Spec.Containers, 1) + container := deployment.Spec.Template.Spec.Containers[0] + assert.Equal(t, "vmcp", container.Name) + assert.NotEmpty(t, container.Image) + assert.Contains(t, container.Args, "serve") + assert.Contains(t, container.Args, "--config=/etc/vmcp-config/config.yaml") + + // Verify checksum annotation is set using standard annotation key + assert.Equal(t, "test-checksum", + deployment.Spec.Template.Annotations[checksum.RunConfigChecksumAnnotation]) } func TestVirtualMCPServerEnsureDeployment_UpdateDeployment(t *testing.T) { @@ -2534,115 +2588,129 @@ func TestVirtualMCPServerEnsureDeployment_NoUpdateNeeded(t *testing.T) { assert.Equal(t, ctrl.Result{}, result) } -// TestVirtualMCPServerEnsureDeployment_RemovesStaleImagePullSecretsHash is a -// regression test for #5817: a stale imagePullRefsHashAnnotation left over -// from a prior reconcile (when imagePullSecrets was non-empty) must be -// removed once the desired list is empty, and reconciling again from that -// cleaned-up state must be a no-op rather than looping forever. -func TestVirtualMCPServerEnsureDeployment_RemovesStaleImagePullSecretsHash(t *testing.T) { +// TestVirtualMCPServerEnsureDeployment_RemovesStaleHashAnnotation is a regression test +// for #5817/#5818: a stale operator-owned hash annotation left over from a prior +// reconcile (when the corresponding field was non-empty) must be removed once that +// field is cleared, and reconciling again from the cleaned-up state must be a no-op +// rather than looping forever. +func TestVirtualMCPServerEnsureDeployment_RemovesStaleHashAnnotation(t *testing.T) { t.Parallel() - scheme := testutil.NewScheme(t) + tests := []struct { + name string + staleHashAnno string + }{ + {name: "imagePullSecrets", staleHashAnno: imagePullRefsHashAnnotation}, + {name: "podTemplateSpec", staleHashAnno: podTemplateSpecHashAnnotation}, + } - vmcp := v1beta1test.NewVirtualMCPServer(testVmcpName, "default", - v1beta1test.WithVMCPGroupRef(testGroupName), - ) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() - configMap := &corev1.ConfigMap{ - ObjectMeta: metav1.ObjectMeta{ - Name: vmcpConfigMapName(vmcp.Name), - Namespace: "default", - Annotations: map[string]string{ - checksum.ContentChecksumAnnotation: "test-checksum", - }, - }, - Data: map[string]string{ - "config.yaml": "test-config", - }, - } + scheme := testutil.NewScheme(t) - reconciler := &VirtualMCPServerReconciler{ - Client: fake.NewClientBuilder().WithScheme(scheme).Build(), - Scheme: scheme, - } + vmcp := v1beta1test.NewVirtualMCPServer(testVmcpName, "default", + v1beta1test.WithVMCPGroupRef(testGroupName), + ) - expectedLabels, expectedAnnotations := reconciler.buildPodTemplateMetadata( - labelsForVirtualMCPServer(vmcp.Name), vmcp, "test-checksum", - ) + configMap := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: vmcpConfigMapName(vmcp.Name), + Namespace: "default", + Annotations: map[string]string{ + checksum.ContentChecksumAnnotation: "test-checksum", + }, + }, + Data: map[string]string{ + "config.yaml": "test-config", + }, + } - // Deployment otherwise matches the desired state exactly, except it - // carries a stale imagePullRefsHashAnnotation from before imagePullSecrets - // was cleared on the VirtualMCPServer. - staleDeployment := &appsv1.Deployment{ - ObjectMeta: metav1.ObjectMeta{ - Name: testVmcpName, - Namespace: "default", - Labels: labelsForVirtualMCPServer(vmcp.Name), - Annotations: map[string]string{ - imagePullRefsHashAnnotation: "stale-hash", - }, - }, - Spec: appsv1.DeploymentSpec{ - Selector: &metav1.LabelSelector{ - MatchLabels: labelsForVirtualMCPServer(vmcp.Name), - }, - Template: corev1.PodTemplateSpec{ + reconciler := &VirtualMCPServerReconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).Build(), + Scheme: scheme, + } + + expectedLabels, expectedAnnotations := reconciler.buildPodTemplateMetadata( + labelsForVirtualMCPServer(vmcp.Name), vmcp, "test-checksum", + ) + + // Deployment otherwise matches the desired state exactly, except it + // carries a stale hash annotation from before the corresponding + // VirtualMCPServer field was cleared. + staleDeployment := &appsv1.Deployment{ ObjectMeta: metav1.ObjectMeta{ - Labels: expectedLabels, - Annotations: expectedAnnotations, + Name: testVmcpName, + Namespace: "default", + Labels: labelsForVirtualMCPServer(vmcp.Name), + Annotations: map[string]string{ + tt.staleHashAnno: "stale-hash", + }, }, - Spec: corev1.PodSpec{ - Containers: []corev1.Container{ - { - Name: "vmcp", - Image: getVmcpImage(), - Ports: []corev1.ContainerPort{ - {ContainerPort: 4483}, + Spec: appsv1.DeploymentSpec{ + Selector: &metav1.LabelSelector{ + MatchLabels: labelsForVirtualMCPServer(vmcp.Name), + }, + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: expectedLabels, + Annotations: expectedAnnotations, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{ + { + Name: "vmcp", + Image: getVmcpImage(), + Ports: []corev1.ContainerPort{ + {ContainerPort: 4483}, + }, + Env: mustBuildEnvVarsForVmcp(reconciler, vmcp), + }, }, - Env: mustBuildEnvVarsForVmcp(reconciler, vmcp), + ServiceAccountName: vmcpServiceAccountName(vmcp.Name), }, }, - ServiceAccountName: vmcpServiceAccountName(vmcp.Name), }, - }, - }, - } + } - k8sClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(vmcp, configMap, staleDeployment). - Build() - reconciler.Client = k8sClient + k8sClient := fake.NewClientBuilder(). + WithScheme(scheme). + WithObjects(vmcp, configMap, staleDeployment). + Build() + reconciler.Client = k8sClient - // First reconcile cleans up the stale annotation. - result, err := reconciler.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) - require.NoError(t, err) - assert.Equal(t, ctrl.Result{}, result) + // First reconcile cleans up the stale annotation. + result, err := reconciler.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) + require.NoError(t, err) + assert.Equal(t, ctrl.Result{}, result) - deployment := &appsv1.Deployment{} - require.NoError(t, k8sClient.Get(context.Background(), types.NamespacedName{ - Name: vmcp.Name, - Namespace: vmcp.Namespace, - }, deployment)) - _, present := deployment.Annotations[imagePullRefsHashAnnotation] - assert.False(t, present, "stale imagePullRefsHashAnnotation must be removed once desired list is empty") + deployment := &appsv1.Deployment{} + require.NoError(t, k8sClient.Get(context.Background(), types.NamespacedName{ + Name: vmcp.Name, + Namespace: vmcp.Namespace, + }, deployment)) + _, present := deployment.Annotations[tt.staleHashAnno] + assert.False(t, present, "stale hash annotation must be removed once the corresponding field is unset") - resourceVersionAfterCleanup := deployment.ResourceVersion + resourceVersionAfterCleanup := deployment.ResourceVersion - // Second reconcile from the now-clean steady state must be a no-op. - // Without both fixes, this would keep re-detecting drift and updating - // forever — the hot-reconcile loop from #5817. - result, err = reconciler.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) - require.NoError(t, err) - assert.Equal(t, ctrl.Result{}, result) + // Second reconcile from the now-clean steady state must be a no-op. + // Without both fixes, this would keep re-detecting drift and updating + // forever — the hot-reconcile loop from #5817/#5818. + result, err = reconciler.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) + require.NoError(t, err) + assert.Equal(t, ctrl.Result{}, result) - deployment = &appsv1.Deployment{} - require.NoError(t, k8sClient.Get(context.Background(), types.NamespacedName{ - Name: vmcp.Name, - Namespace: vmcp.Namespace, - }, deployment)) - assert.Equal(t, resourceVersionAfterCleanup, deployment.ResourceVersion, - "reconciling from steady state must not write again") + deployment = &appsv1.Deployment{} + require.NoError(t, k8sClient.Get(context.Background(), types.NamespacedName{ + Name: vmcp.Name, + Namespace: vmcp.Namespace, + }, deployment)) + assert.Equal(t, resourceVersionAfterCleanup, deployment.ResourceVersion, + "reconciling from steady state must not write again") + }) + } } func TestVirtualMCPServerEnsureService_CreateService(t *testing.T) { @@ -2962,144 +3030,101 @@ func TestVirtualMCPServerValidateEmbeddingServerRef(t *testing.T) { // TestVirtualMCPServerEnsureDeployment_ReplicaSync_SpecDriven verifies that when // spec.replicas is set, ensureDeployment updates the Deployment to match. -func TestVirtualMCPServerEnsureDeployment_ReplicaSync_SpecDriven(t *testing.T) { +// TestVirtualMCPServerEnsureDeployment_ReplicaSync covers spec.replicas +// syncing to the live Deployment: a spec-driven value overrides whatever is +// live, while nil leaves the live value untouched (HPA-managed passthrough). +func TestVirtualMCPServerEnsureDeployment_ReplicaSync(t *testing.T) { t.Parallel() - specReplicas := int32(3) - vmcp := v1beta1test.NewVirtualMCPServer("vmcp-replica-sync", "default", - v1beta1test.WithVMCPGroupRef(testGroupName), - v1beta1test.WithVMCPReplicas(specReplicas), - ) - - mcpGroup := &mcpv1beta1.MCPGroup{ - ObjectMeta: metav1.ObjectMeta{Name: testGroupName, Namespace: "default"}, - Status: mcpv1beta1.MCPGroupStatus{Phase: mcpv1beta1.MCPGroupPhaseReady}, - } - - configMap := &corev1.ConfigMap{ - ObjectMeta: metav1.ObjectMeta{ - Name: vmcpConfigMapName(vmcp.Name), - Namespace: "default", - Annotations: map[string]string{ - checksum.ContentChecksumAnnotation: testChecksumValue, - }, - }, - Data: map[string]string{"config.yaml": "{}"}, - } - - // Existing deployment has 1 replica — simulates a pre-existing state - existingReplicas := int32(1) - existingDeployment := &appsv1.Deployment{ - ObjectMeta: metav1.ObjectMeta{ - Name: vmcp.Name, - Namespace: "default", - Labels: labelsForVirtualMCPServer(vmcp.Name), + tests := []struct { + name string + specReplicas *int32 + existingReplicas int32 + wantReplicas int32 + }{ + { + name: "spec-driven value overrides existing replica count", + specReplicas: ptr.To(int32(3)), + existingReplicas: 1, + wantReplicas: 3, }, - Spec: appsv1.DeploymentSpec{ - Replicas: &existingReplicas, - Selector: &metav1.LabelSelector{MatchLabels: labelsForVirtualMCPServer(vmcp.Name)}, - Template: corev1.PodTemplateSpec{ - ObjectMeta: metav1.ObjectMeta{Labels: labelsForVirtualMCPServer(vmcp.Name)}, - Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "vmcp", Image: "test:latest"}}}, - }, + { + name: "nil spec.replicas does not overwrite HPA-managed count", + specReplicas: nil, + existingReplicas: 5, + wantReplicas: 5, }, } - scheme := testutil.NewScheme(t) - - fakeClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(vmcp, mcpGroup, configMap, existingDeployment). - Build() - - r := &VirtualMCPServerReconciler{ - Client: fakeClient, - Scheme: scheme, - PlatformDetector: ctrlutil.NewSharedPlatformDetector(), - } - - result, err := r.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) - require.NoError(t, err) - assert.Equal(t, ctrl.Result{}, result) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() - updated := &appsv1.Deployment{} - err = fakeClient.Get(context.Background(), types.NamespacedName{ - Name: vmcp.Name, Namespace: vmcp.Namespace, - }, updated) - require.NoError(t, err) - require.NotNil(t, updated.Spec.Replicas) - assert.Equal(t, int32(3), *updated.Spec.Replicas) -} + opts := []v1beta1test.VirtualMCPServerOption{v1beta1test.WithVMCPGroupRef(testGroupName)} + if tt.specReplicas != nil { + opts = append(opts, v1beta1test.WithVMCPReplicas(*tt.specReplicas)) + } + vmcp := v1beta1test.NewVirtualMCPServer(testVmcpName, "default", opts...) -// TestVirtualMCPServerEnsureDeployment_ReplicaSync_NilPassthrough verifies that when -// spec.replicas is nil, ensureDeployment does not overwrite a live replica count (HPA-managed). -func TestVirtualMCPServerEnsureDeployment_ReplicaSync_NilPassthrough(t *testing.T) { - t.Parallel() + mcpGroup := &mcpv1beta1.MCPGroup{ + ObjectMeta: metav1.ObjectMeta{Name: testGroupName, Namespace: "default"}, + Status: mcpv1beta1.MCPGroupStatus{Phase: mcpv1beta1.MCPGroupPhaseReady}, + } - // Replicas left nil — HPA manages replicas. - vmcp := v1beta1test.NewVirtualMCPServer("vmcp-nil-passthrough", "default", - v1beta1test.WithVMCPGroupRef(testGroupName), - ) + configMap := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: vmcpConfigMapName(vmcp.Name), + Namespace: "default", + Annotations: map[string]string{ + checksum.ContentChecksumAnnotation: testChecksumValue, + }, + }, + Data: map[string]string{"config.yaml": "{}"}, + } - mcpGroup := &mcpv1beta1.MCPGroup{ - ObjectMeta: metav1.ObjectMeta{Name: testGroupName, Namespace: "default"}, - Status: mcpv1beta1.MCPGroupStatus{Phase: mcpv1beta1.MCPGroupPhaseReady}, - } + existingReplicas := tt.existingReplicas + existingDeployment := &appsv1.Deployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: vmcp.Name, + Namespace: "default", + Labels: labelsForVirtualMCPServer(vmcp.Name), + }, + Spec: appsv1.DeploymentSpec{ + Replicas: &existingReplicas, + Selector: &metav1.LabelSelector{MatchLabels: labelsForVirtualMCPServer(vmcp.Name)}, + Template: corev1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{Labels: labelsForVirtualMCPServer(vmcp.Name)}, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "vmcp", Image: "test:latest"}}}, + }, + }, + } - configMap := &corev1.ConfigMap{ - ObjectMeta: metav1.ObjectMeta{ - Name: vmcpConfigMapName(vmcp.Name), - Namespace: "default", - Annotations: map[string]string{ - checksum.ContentChecksumAnnotation: testChecksumValue, - }, - }, - Data: map[string]string{"config.yaml": "{}"}, - } + scheme := testutil.NewScheme(t) - // Existing deployment has 5 replicas — set by HPA - hpaReplicas := int32(5) - existingDeployment := &appsv1.Deployment{ - ObjectMeta: metav1.ObjectMeta{ - Name: vmcp.Name, - Namespace: "default", - Labels: labelsForVirtualMCPServer(vmcp.Name), - }, - Spec: appsv1.DeploymentSpec{ - Replicas: &hpaReplicas, - Selector: &metav1.LabelSelector{MatchLabels: labelsForVirtualMCPServer(vmcp.Name)}, - Template: corev1.PodTemplateSpec{ - ObjectMeta: metav1.ObjectMeta{Labels: labelsForVirtualMCPServer(vmcp.Name)}, - Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "vmcp", Image: "test:latest"}}}, - }, - }, - } + fakeClient := fake.NewClientBuilder(). + WithScheme(scheme). + WithObjects(vmcp, mcpGroup, configMap, existingDeployment). + Build() - scheme := testutil.NewScheme(t) + r := &VirtualMCPServerReconciler{ + Client: fakeClient, + Scheme: scheme, + PlatformDetector: ctrlutil.NewSharedPlatformDetector(), + } - fakeClient := fake.NewClientBuilder(). - WithScheme(scheme). - WithObjects(vmcp, mcpGroup, configMap, existingDeployment). - Build() + result, err := r.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) + require.NoError(t, err) + assert.Equal(t, ctrl.Result{}, result) - r := &VirtualMCPServerReconciler{ - Client: fakeClient, - Scheme: scheme, - PlatformDetector: ctrlutil.NewSharedPlatformDetector(), + updated := &appsv1.Deployment{} + err = fakeClient.Get(context.Background(), types.NamespacedName{ + Name: vmcp.Name, Namespace: vmcp.Namespace, + }, updated) + require.NoError(t, err) + require.NotNil(t, updated.Spec.Replicas) + assert.Equal(t, tt.wantReplicas, *updated.Spec.Replicas) + }) } - - result, err := r.ensureDeployment(context.Background(), vmcp, nil, []workloads.TypedWorkload{}) - require.NoError(t, err) - assert.Equal(t, ctrl.Result{}, result) - - updated := &appsv1.Deployment{} - err = fakeClient.Get(context.Background(), types.NamespacedName{ - Name: vmcp.Name, Namespace: vmcp.Namespace, - }, updated) - require.NoError(t, err) - // HPA-managed replica count must not be overwritten - require.NotNil(t, updated.Spec.Replicas) - assert.Equal(t, int32(5), *updated.Spec.Replicas) } // mustBuildEnvVarsForVmcp is a test helper that calls buildEnvVarsForVmcp and panics on error. diff --git a/cmd/thv-operator/controllers/workload_reference_count.go b/cmd/thv-operator/controllers/workload_reference_count.go deleted file mode 100644 index aeb36b6c26..0000000000 --- a/cmd/thv-operator/controllers/workload_reference_count.go +++ /dev/null @@ -1,22 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package controllers - -import ( - "math" - - mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" -) - -func workloadReferenceCount(refs []mcpv1beta1.WorkloadReference) int32 { - return workloadReferenceCountFromLen(len(refs)) -} - -func workloadReferenceCountFromLen(length int) int32 { - if length > math.MaxInt32 { - return math.MaxInt32 - } - - return int32(length) //nolint:gosec // guarded above against int32 overflow -} diff --git a/cmd/thv-operator/controllers/workload_reference_count_test.go b/cmd/thv-operator/controllers/workload_reference_count_test.go deleted file mode 100644 index 7afa8c489e..0000000000 --- a/cmd/thv-operator/controllers/workload_reference_count_test.go +++ /dev/null @@ -1,28 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package controllers - -import ( - "math" - "testing" - - "github.com/stretchr/testify/assert" - - mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" -) - -func TestWorkloadReferenceCount(t *testing.T) { - t.Parallel() - - assert.EqualValues(t, 0, workloadReferenceCount(nil)) - assert.EqualValues(t, 0, workloadReferenceCount([]mcpv1beta1.WorkloadReference{})) - - refs := []mcpv1beta1.WorkloadReference{ - {Kind: mcpv1beta1.WorkloadKindMCPServer, Name: "server-1"}, - {Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: "proxy-1"}, - } - - assert.EqualValues(t, 2, workloadReferenceCount(refs)) - assert.Equal(t, int32(math.MaxInt32), workloadReferenceCountFromLen(math.MaxInt32+1)) -} diff --git a/cmd/thv-operator/pkg/controllerutil/config.go b/cmd/thv-operator/pkg/controllerutil/config.go index ff8f127d41..e3736a7251 100644 --- a/cmd/thv-operator/pkg/controllerutil/config.go +++ b/cmd/thv-operator/pkg/controllerutil/config.go @@ -54,14 +54,6 @@ func SortWorkloadRefs(refs []mcpv1beta1.WorkloadReference) { slices.SortFunc(refs, CompareWorkloadRefs) } -// WorkloadRefsEqual reports whether two WorkloadReference slices contain the same entries. -// Both slices must already be sorted (use SortWorkloadRefs) for correct results. -func WorkloadRefsEqual(a, b []mcpv1beta1.WorkloadReference) bool { - return slices.EqualFunc(a, b, func(x, y mcpv1beta1.WorkloadReference) bool { - return x.Kind == y.Kind && x.Name == y.Name - }) -} - // GetToolConfigForMCPRemoteProxy fetches MCPToolConfig referenced by MCPRemoteProxy func GetToolConfigForMCPRemoteProxy( ctx context.Context, diff --git a/cmd/thv-operator/pkg/controllerutil/config_test.go b/cmd/thv-operator/pkg/controllerutil/config_test.go index 2a6bbb793d..ef06f64227 100644 --- a/cmd/thv-operator/pkg/controllerutil/config_test.go +++ b/cmd/thv-operator/pkg/controllerutil/config_test.go @@ -128,56 +128,6 @@ func TestSortWorkloadRefs(t *testing.T) { }) } -func TestWorkloadRefsEqual(t *testing.T) { - t.Parallel() - - t.Run("equal slices", func(t *testing.T) { - t.Parallel() - a := []mcpv1beta1.WorkloadReference{ - {Kind: "MCPServer", Name: "alpha"}, - {Kind: "MCPServer", Name: "beta"}, - } - b := []mcpv1beta1.WorkloadReference{ - {Kind: "MCPServer", Name: "alpha"}, - {Kind: "MCPServer", Name: "beta"}, - } - assert.True(t, WorkloadRefsEqual(a, b)) - }) - - t.Run("different order is not equal", func(t *testing.T) { - t.Parallel() - a := []mcpv1beta1.WorkloadReference{ - {Kind: "MCPServer", Name: "alpha"}, - {Kind: "MCPServer", Name: "beta"}, - } - b := []mcpv1beta1.WorkloadReference{ - {Kind: "MCPServer", Name: "beta"}, - {Kind: "MCPServer", Name: "alpha"}, - } - assert.False(t, WorkloadRefsEqual(a, b)) - }) - - t.Run("different lengths", func(t *testing.T) { - t.Parallel() - a := []mcpv1beta1.WorkloadReference{{Kind: "MCPServer", Name: "alpha"}} - b := []mcpv1beta1.WorkloadReference{ - {Kind: "MCPServer", Name: "alpha"}, - {Kind: "MCPServer", Name: "beta"}, - } - assert.False(t, WorkloadRefsEqual(a, b)) - }) - - t.Run("both nil", func(t *testing.T) { - t.Parallel() - assert.True(t, WorkloadRefsEqual(nil, nil)) - }) - - t.Run("nil vs empty", func(t *testing.T) { - t.Parallel() - assert.True(t, WorkloadRefsEqual(nil, []mcpv1beta1.WorkloadReference{})) - }) -} - func TestGetTelemetryConfigForMCPRemoteProxy(t *testing.T) { t.Parallel() diff --git a/cmd/thv-operator/pkg/registryapi/deployment.go b/cmd/thv-operator/pkg/registryapi/deployment.go index ca9ac8fbba..23715d6c5d 100644 --- a/cmd/thv-operator/pkg/registryapi/deployment.go +++ b/cmd/thv-operator/pkg/registryapi/deployment.go @@ -138,6 +138,12 @@ func (m *manager) upsertDeployment( for k, v := range deployment.Annotations { existing.Annotations[k] = v } + // Prune the hash annotation if desired no longer wants it, otherwise it + // survives forever and deploymentNeedsUpdate reports drift on every + // reconcile (same bug class as #5817/#5818). + if _, wantHash := deployment.Annotations[podTemplateSpecHashAnnotation]; !wantHash { + delete(existing.Annotations, podTemplateSpecHashAnnotation) + } // Ensure owner reference is set on the existing object if err := controllerutil.SetControllerReference(mcpRegistry, existing, m.scheme); err != nil { diff --git a/cmd/thv-operator/pkg/registryapi/deployment_test.go b/cmd/thv-operator/pkg/registryapi/deployment_test.go index 49ab2f5563..2b2900f46d 100644 --- a/cmd/thv-operator/pkg/registryapi/deployment_test.go +++ b/cmd/thv-operator/pkg/registryapi/deployment_test.go @@ -13,8 +13,11 @@ import ( corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" + "github.com/stacklok/toolhive/cmd/thv-operator/internal/testutil" ) func TestGetRegistryAPIImage(t *testing.T) { @@ -526,6 +529,73 @@ func TestDeploymentNeedsUpdate(t *testing.T) { } } +// TestUpsertDeployment_RemovesStalePodTemplateSpecHash is a regression test +// for the registry-api analog of #5817/#5818: a stale +// podTemplateSpecHashAnnotation left over from a prior reconcile (when +// spec.podTemplateSpec was set) must be removed once the field is cleared, +// and upserting again from that cleaned-up state must be a no-op rather than +// looping forever. +func TestUpsertDeployment_RemovesStalePodTemplateSpecHash(t *testing.T) { + t.Parallel() + + scheme := testutil.NewScheme(t) + + mcpRegistry := &mcpv1beta1.MCPRegistry{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-registry", + Namespace: "test-namespace", + }, + Spec: mcpv1beta1.MCPRegistrySpec{ + ConfigYAML: "sources:\n - name: k8s\n kubernetes: {}\n", + }, + } + + mgr := &manager{scheme: scheme} + + desired, err := mgr.buildRegistryAPIDeployment(context.Background(), mcpRegistry, "test-registry-registry-server-config") + require.NoError(t, err) + + // Existing deployment matches desired exactly, except it carries a stale + // podTemplateSpecHashAnnotation from before spec.podTemplateSpec was + // cleared on the MCPRegistry. + staleDeployment := desired.DeepCopy() + staleDeployment.Annotations = map[string]string{ + podTemplateSpecHashAnnotation: "stale-hash", + } + + mgr.client = fake.NewClientBuilder(). + WithScheme(scheme). + WithObjects(staleDeployment). + Build() + + // First upsert cleans up the stale annotation. + _, err = mgr.upsertDeployment(context.Background(), mcpRegistry, desired.DeepCopy()) + require.NoError(t, err) + + deployment := &appsv1.Deployment{} + require.NoError(t, mgr.client.Get(context.Background(), client.ObjectKey{ + Name: mcpRegistry.GetAPIResourceName(), + Namespace: mcpRegistry.Namespace, + }, deployment)) + _, present := deployment.Annotations[podTemplateSpecHashAnnotation] + assert.False(t, present, "stale podTemplateSpecHashAnnotation must be removed once PodTemplateSpec is unset") + + resourceVersionAfterCleanup := deployment.ResourceVersion + + // Second upsert from the now-clean steady state must be a no-op. Without + // the prune fix, this would keep re-detecting drift and updating forever. + _, err = mgr.upsertDeployment(context.Background(), mcpRegistry, desired.DeepCopy()) + require.NoError(t, err) + + deployment = &appsv1.Deployment{} + require.NoError(t, mgr.client.Get(context.Background(), client.ObjectKey{ + Name: mcpRegistry.GetAPIResourceName(), + Namespace: mcpRegistry.Namespace, + }, deployment)) + assert.Equal(t, resourceVersionAfterCleanup, deployment.ResourceVersion, + "upserting from steady state must not write again") +} + func TestBuildRegistryAPIDeployment_PodTemplateSpecHash(t *testing.T) { t.Parallel() diff --git a/cmd/thv-operator/pkg/validation/url_validation.go b/cmd/thv-operator/pkg/validation/url_validation.go index 6d69d2d9d7..cf060f993b 100644 --- a/cmd/thv-operator/pkg/validation/url_validation.go +++ b/cmd/thv-operator/pkg/validation/url_validation.go @@ -116,9 +116,13 @@ func validateHostNotInternal(host string) error { // ValidateJWKSURL validates that rawURL, if non-empty, is a well-formed HTTPS // URL with a non-empty host. JWKS endpoints serve key material and must use -// HTTPS. An empty rawURL is allowed because JWKS discovery can determine the -// endpoint automatically. -func ValidateJWKSURL(rawURL string) error { +// HTTPS. If allowInsecure is true, HTTP is permitted (for development/testing +// only), matching ValidateOIDCIssuerURL. An empty rawURL is allowed because +// JWKS discovery can determine the endpoint automatically. +// +// Schemes other than http and https are always rejected, regardless of +// allowInsecure. +func ValidateJWKSURL(rawURL string, allowInsecure bool) error { if rawURL == "" { return nil } @@ -128,8 +132,15 @@ func ValidateJWKSURL(rawURL string) error { return fmt.Errorf("JWKS URL is invalid: %w", err) } - if u.Scheme != schemeHTTPS { - return fmt.Errorf("JWKS URL must use HTTPS scheme, got %q", u.Scheme) + if u.Scheme == schemeHTTP && !allowInsecure { + return fmt.Errorf( + "JWKS URL %q uses HTTP scheme, which is insecure; "+ + "use HTTPS or set insecureAllowHTTP: true for development only", rawURL, + ) + } + + if u.Scheme != schemeHTTP && u.Scheme != schemeHTTPS { + return fmt.Errorf("JWKS URL %q has unsupported scheme %q; must be http or https", rawURL, u.Scheme) } if u.Host == "" { diff --git a/cmd/thv-operator/pkg/validation/url_validation_test.go b/cmd/thv-operator/pkg/validation/url_validation_test.go index ab72dfacdf..230920ddf1 100644 --- a/cmd/thv-operator/pkg/validation/url_validation_test.go +++ b/cmd/thv-operator/pkg/validation/url_validation_test.go @@ -236,10 +236,11 @@ func TestValidateJWKSURL(t *testing.T) { t.Parallel() tests := []struct { - name string - rawURL string - wantErr bool - errContains string + name string + rawURL string + allowInsecure bool + wantErr bool + errContains string }{ { name: "empty URL allowed", @@ -255,13 +256,26 @@ func TestValidateJWKSURL(t *testing.T) { name: "http rejected", rawURL: "http://jwks.example.com", wantErr: true, - errContains: "HTTPS", + errContains: "insecureAllowHTTP", + }, + { + name: "http allowed with insecureAllowHTTP", + rawURL: "http://jwks.example.com", + allowInsecure: true, + wantErr: false, }, { name: "unsupported scheme", rawURL: "ftp://jwks.example.com", wantErr: true, - errContains: "HTTPS", + errContains: "unsupported scheme", + }, + { + name: "unsupported scheme rejected even with insecureAllowHTTP", + rawURL: "ftp://jwks.example.com", + allowInsecure: true, + wantErr: true, + errContains: "unsupported scheme", }, { name: "missing host", @@ -275,7 +289,7 @@ func TestValidateJWKSURL(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - err := validation.ValidateJWKSURL(tt.rawURL) + err := validation.ValidateJWKSURL(tt.rawURL, tt.allowInsecure) if tt.wantErr { require.Error(t, err) diff --git a/cmd/thv-operator/pkg/vmcpconfig/converter.go b/cmd/thv-operator/pkg/vmcpconfig/converter.go index af8a50338c..3b02fe7f26 100644 --- a/cmd/thv-operator/pkg/vmcpconfig/converter.go +++ b/cmd/thv-operator/pkg/vmcpconfig/converter.go @@ -469,6 +469,7 @@ func mapResolvedOIDCToVmcpConfigFromRef( JwksAllowPrivateIP: resolved.JWKSAllowPrivateIP, InsecureAllowHTTP: resolved.InsecureAllowHTTP, Scopes: resolved.Scopes, + CABundlePath: resolved.ThvCABundlePath, } // MCPOIDCConfig inline type may have a client secret @@ -727,11 +728,20 @@ func (c *Converter) convertAggregation( ctx context.Context, vmcp *mcpv1beta1.VirtualMCPServer, ) (*vmcpconfig.AggregationConfig, error) { - // Start with a deep copy of the source config + // Field-by-field copy, NOT a deep copy: scalars are copied by value, slices + // (ConflictResolutionConfig.PriorityOrder, WorkloadToolConfig.Filter) are shared + // with the source, and only Overrides is deep-copied. Two consequences: + // - Treat the source's slices as read-only; mutating them here would mutate the + // CR's in-memory spec. + // - Every new AggregationConfig field must be added HERE explicitly. A field + // omitted from this literal is accepted by the CRD and then silently dropped + // before the vMCP process ever sees it — which for a visibility setting means + // failing OPEN on config users rely on to withhold tools. srcAgg := vmcp.Spec.Config.Aggregation agg := &vmcpconfig.AggregationConfig{ - ConflictResolution: srcAgg.ConflictResolution, - ExcludeAllTools: srcAgg.ExcludeAllTools, + ConflictResolution: srcAgg.ConflictResolution, + ExcludeAllTools: srcAgg.ExcludeAllTools, + DefaultToolVisibility: srcAgg.DefaultToolVisibility, } // Apply defaults for conflict resolution diff --git a/cmd/thv-operator/pkg/vmcpconfig/converter_test.go b/cmd/thv-operator/pkg/vmcpconfig/converter_test.go index 11a980d561..313fb01314 100644 --- a/cmd/thv-operator/pkg/vmcpconfig/converter_test.go +++ b/cmd/thv-operator/pkg/vmcpconfig/converter_test.go @@ -184,6 +184,52 @@ func TestConverter_OIDCResolution(t *testing.T) { assert.Equal(t, "VMCP_OIDC_CLIENT_SECRET", config.IncomingAuth.OIDC.ClientSecretEnv) }, }, + { + name: "inline resolved ThvCABundlePath maps to CABundlePath", + oidcConfigRef: &mcpv1beta1.MCPOIDCConfigReference{Name: oidcConfigName, Audience: "test-audience"}, + oidcConfig: newTestMCPOIDCConfig(mcpv1beta1.MCPOIDCConfigTypeInline), + mockReturn: &oidc.OIDCConfig{ + Issuer: "https://issuer.example.com", + ThvCABundlePath: "/config/certs/example-ca/ca.crt", + }, + validate: func(t *testing.T, config *vmcpconfig.Config, err error) { + t.Helper() + require.NoError(t, err) + require.NotNil(t, config.IncomingAuth.OIDC) + assert.Equal(t, "/config/certs/example-ca/ca.crt", config.IncomingAuth.OIDC.CABundlePath) + }, + }, + { + name: "k8s service account ThvCABundlePath maps to CABundlePath", + oidcConfigRef: &mcpv1beta1.MCPOIDCConfigReference{Name: oidcConfigName, Audience: "test-audience"}, + oidcConfig: newTestMCPOIDCConfig(mcpv1beta1.MCPOIDCConfigTypeKubernetesServiceAccount), + mockReturn: &oidc.OIDCConfig{ + Issuer: "https://kubernetes.default.svc", + ThvCABundlePath: "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt", + }, + validate: func(t *testing.T, config *vmcpconfig.Config, err error) { + t.Helper() + require.NoError(t, err) + require.NotNil(t, config.IncomingAuth.OIDC) + assert.Equal(t, + "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt", + config.IncomingAuth.OIDC.CABundlePath) + }, + }, + { + name: "empty ThvCABundlePath results in empty CABundlePath", + oidcConfigRef: &mcpv1beta1.MCPOIDCConfigReference{Name: oidcConfigName, Audience: "test-audience"}, + oidcConfig: newTestMCPOIDCConfig(mcpv1beta1.MCPOIDCConfigTypeInline), + mockReturn: &oidc.OIDCConfig{ + Issuer: "https://issuer.example.com", + }, + validate: func(t *testing.T, config *vmcpconfig.Config, err error) { + t.Helper() + require.NoError(t, err) + require.NotNil(t, config.IncomingAuth.OIDC) + assert.Empty(t, config.IncomingAuth.OIDC.CABundlePath) + }, + }, { name: "non-inline type does not set ClientSecretEnv", oidcConfigRef: &mcpv1beta1.MCPOIDCConfigReference{Name: oidcConfigName, Audience: "test-audience"}, @@ -1432,6 +1478,56 @@ func TestConvert_MCPToolConfigFailClosed(t *testing.T) { } } +// TestConvert_DefaultToolVisibilityPreserved guards against silently dropping the +// aggregation visibility setting. convertAggregation hand-copies fields rather +// than deep-copying, so omitting DefaultToolVisibility there would let the CRD accept +// `defaultToolVisibility: deny` while the rendered vMCP config falls back to +// advertise-everything — a security-relevant setting that looks applied and is +// not. An empty value must survive as empty (the aggregator treats it as allow), +// so pre-existing CRs keep today's behavior. +func TestConvert_DefaultToolVisibilityPreserved(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + set vmcpconfig.DefaultToolVisibility + want vmcpconfig.DefaultToolVisibility + }{ + {name: "deny is carried through to the rendered config", set: vmcpconfig.DefaultToolVisibilityDeny, want: vmcpconfig.DefaultToolVisibilityDeny}, + {name: "allow is carried through to the rendered config", set: vmcpconfig.DefaultToolVisibilityAllow, want: vmcpconfig.DefaultToolVisibilityAllow}, + {name: "unset stays unset (treated as allow downstream)", set: "", want: ""}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + vmcpServer := v1beta1test.NewVirtualMCPServer("test-vmcp", "default", + v1beta1test.WithVMCPGroupRef("test-group"), + v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{Type: "anonymous"}), + v1beta1test.WithVMCPConfig(vmcpconfig.Config{ + Aggregation: &vmcpconfig.AggregationConfig{ + DefaultToolVisibility: tt.set, + Tools: []*vmcpconfig.WorkloadToolConfig{{Workload: "backend1"}}, + }, + }), + ) + + ctx := log.IntoContext(context.Background(), logr.Discard()) + converter := newTestConverter(t, newNoOpMockResolver(t)) + converter.k8sClient = newTestK8sClient(t) + + got, _, err := converter.Convert(ctx, vmcpServer, nil) + require.NoError(t, err) + require.NotNil(t, got) + require.NotNil(t, got.Aggregation) + assert.Equal(t, tt.want, got.Aggregation.DefaultToolVisibility, + "defaultToolVisibility must survive CRD-to-config conversion") + }) + } +} + // TestConverter_InlineTelemetryIgnored verifies that the operator-side converter // ignores Config.Telemetry (the standalone CLI field) and only uses TelemetryConfigRef. func TestConverter_InlineTelemetryIgnored(t *testing.T) { diff --git a/cmd/thv-operator/test-integration/mcp-external-auth/mcpexternalauthconfig_controller_integration_test.go b/cmd/thv-operator/test-integration/mcp-external-auth/mcpexternalauthconfig_controller_integration_test.go index a0eda3135c..4c3e81d690 100644 --- a/cmd/thv-operator/test-integration/mcp-external-auth/mcpexternalauthconfig_controller_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-external-auth/mcpexternalauthconfig_controller_integration_test.go @@ -11,6 +11,8 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" @@ -242,27 +244,6 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { }, timeout, interval).Should(BeTrue()) }) - It("Should update MCPExternalAuthConfig status with referencing workload", func() { - // Wait for the auth config status to be updated with the referencing workload - Eventually(func() bool { - updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: authConfigName, - Namespace: namespace, - }, updatedAuthConfig) - if err != nil { - return false - } - // Check if the server is in the referencing workloads list - for _, ref := range updatedAuthConfig.Status.ReferencingWorkloads { - if ref.Kind == "MCPServer" && ref.Name == mcpServerName { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - It("Should create ConfigMap with token exchange configuration", func() { // Wait for ConfigMap to be created configMapName := mcpServerName + "-runconfig" @@ -323,7 +304,7 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { }) }) - Context("When a referencing MCPServer is deleted", Ordered, func() { + Context("When updating an MCPExternalAuthConfig", Ordered, func() { var ( namespace string authConfigName string @@ -332,13 +313,14 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { mcpServer *mcpv1beta1.MCPServer oauthSecret *corev1.Secret oauthSecretName string + originalHash string ) BeforeAll(func() { namespace = defaultNamespace - authConfigName = "test-external-auth-prune" - mcpServerName = "external-auth-prune-server" - oauthSecretName = "oauth-test-secret-prune" + authConfigName = "test-external-auth-update" + mcpServerName = "external-auth-update-server" + oauthSecretName = "oauth-test-secret-update" // Create namespace if it doesn't exist ns := &corev1.Namespace{ @@ -355,7 +337,7 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { Namespace: namespace, }, StringData: map[string]string{ - "client-secret": "test-secret-value-prune", + "client-secret": "original-secret", }, } Expect(k8sClient.Create(ctx, oauthSecret)).Should(Succeed()) @@ -370,12 +352,12 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { Type: "tokenExchange", TokenExchange: &mcpv1beta1.TokenExchangeConfig{ TokenURL: "https://oauth.example.com/token", - ClientID: "test-client-id-prune", + ClientID: "original-client-id", ClientSecretRef: &mcpv1beta1.SecretKeyRef{ Name: oauthSecretName, Key: "client-secret", }, - Audience: "mcp-backend-prune", + Audience: "original-audience", Scopes: []string{"read"}, }, }, @@ -389,7 +371,11 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { Name: authConfigName, Namespace: namespace, }, updatedAuthConfig) - return err == nil && updatedAuthConfig.Status.ConfigHash != "" + if err != nil { + return false + } + originalHash = updatedAuthConfig.Status.ConfigHash + return originalHash != "" }, timeout, interval).Should(BeTrue()) // Create MCPServer with external auth reference @@ -409,40 +395,45 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { } Expect(k8sClient.Create(ctx, mcpServer)).Should(Succeed()) - // Wait until the server is tracked in ReferencingWorkloads. + // Wait for the MCPServer to have the original hash Eventually(func() bool { - updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} + updatedMCPServer := &mcpv1beta1.MCPServer{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: authConfigName, + Name: mcpServerName, Namespace: namespace, - }, updatedAuthConfig) + }, updatedMCPServer) if err != nil { return false } - for _, ref := range updatedAuthConfig.Status.ReferencingWorkloads { - if ref.Kind == "MCPServer" && ref.Name == mcpServerName { - return true - } - } - return false + return updatedMCPServer.Status.ExternalAuthConfigHash == originalHash }, timeout, interval).Should(BeTrue()) }) AfterAll(func() { - // mcpServer is already deleted by the spec; ignore NotFound on cleanup. - _ = k8sClient.Delete(ctx, mcpServer) + // Clean up resources + Expect(k8sClient.Delete(ctx, mcpServer)).Should(Succeed()) Expect(k8sClient.Delete(ctx, authConfig)).Should(Succeed()) Expect(k8sClient.Delete(ctx, oauthSecret)).Should(Succeed()) }) - It("Should prune the workload from ReferencingWorkloads and decrement ReferenceCount", func() { - // Delete the referencing MCPServer. The controller-runtime watch - // enqueues the referenced MCPExternalAuthConfig (via the old-object - // map function), and the level-triggered reconcile must prune the - // now-stale reference. - Expect(k8sClient.Delete(ctx, mcpServer)).Should(Succeed()) + It("Should update config hash when auth config is modified", func() { + // Update the auth config + Eventually(func() error { + updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} + if err := k8sClient.Get(ctx, types.NamespacedName{ + Name: authConfigName, + Namespace: namespace, + }, updatedAuthConfig); err != nil { + return err + } - // Eventually the reference is pruned and the count drops to zero. + // Modify the audience + updatedAuthConfig.Spec.TokenExchange.Audience = "updated-audience" + return k8sClient.Update(ctx, updatedAuthConfig) + }, timeout, interval).Should(Succeed()) + + // Wait for the config hash to change + var newHash string Eventually(func() bool { updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ @@ -452,17 +443,37 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { if err != nil { return false } - for _, ref := range updatedAuthConfig.Status.ReferencingWorkloads { - if ref.Kind == "MCPServer" && ref.Name == mcpServerName { - return false - } + newHash = updatedAuthConfig.Status.ConfigHash + return newHash != "" && newHash != originalHash + }, timeout, interval).Should(BeTrue()) + + // Verify the new hash is different + Expect(newHash).NotTo(Equal(originalHash)) + }) + + It("Should trigger MCPServer reconciliation with updated hash", func() { + // Wait for the MCPServer to get the updated hash + Eventually(func() bool { + updatedMCPServer := &mcpv1beta1.MCPServer{} + err := k8sClient.Get(ctx, types.NamespacedName{ + Name: mcpServerName, + Namespace: namespace, + }, updatedMCPServer) + if err != nil { + return false } - return updatedAuthConfig.Status.ReferenceCount == 0 + // Check if the hash has been updated + return updatedMCPServer.Status.ExternalAuthConfigHash != originalHash }, timeout, interval).Should(BeTrue()) }) }) - Context("When updating an MCPExternalAuthConfig", Ordered, func() { + Context("When deleting an MCPExternalAuthConfig with active references", Ordered, func() { + // The block->unblock deletion transition is driven purely by the + // controller's periodic requeue (the workload watches were removed), so + // the unblock assertion below must tolerate the up-to-30s requeue delay. + const deletionRequeueTimeout = time.Second * 45 + var ( namespace string authConfigName string @@ -471,14 +482,13 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { mcpServer *mcpv1beta1.MCPServer oauthSecret *corev1.Secret oauthSecretName string - originalHash string ) BeforeAll(func() { namespace = defaultNamespace - authConfigName = "test-external-auth-update" - mcpServerName = "external-auth-update-server" - oauthSecretName = "oauth-test-secret-update" + authConfigName = "test-external-auth-deletion" + mcpServerName = "external-auth-deletion-server" + oauthSecretName = "oauth-test-secret-deletion" // Create namespace if it doesn't exist ns := &corev1.Namespace{ @@ -495,7 +505,7 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { Namespace: namespace, }, StringData: map[string]string{ - "client-secret": "original-secret", + "client-secret": "deletion-secret-value", }, } Expect(k8sClient.Create(ctx, oauthSecret)).Should(Succeed()) @@ -510,33 +520,20 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { Type: "tokenExchange", TokenExchange: &mcpv1beta1.TokenExchangeConfig{ TokenURL: "https://oauth.example.com/token", - ClientID: "original-client-id", + ClientID: "deletion-client-id", ClientSecretRef: &mcpv1beta1.SecretKeyRef{ Name: oauthSecretName, Key: "client-secret", }, - Audience: "original-audience", + Audience: "mcp-backend-deletion", Scopes: []string{"read"}, }, }, } Expect(k8sClient.Create(ctx, authConfig)).Should(Succeed()) - // Wait for the auth config to have a hash - Eventually(func() bool { - updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: authConfigName, - Namespace: namespace, - }, updatedAuthConfig) - if err != nil { - return false - } - originalHash = updatedAuthConfig.Status.ConfigHash - return originalHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create MCPServer with external auth reference + // Create MCPServer referencing the auth config so it is observably + // referenced before we attempt deletion. mcpServer = &mcpv1beta1.MCPServer{ ObjectMeta: metav1.ObjectMeta{ Name: mcpServerName, @@ -553,7 +550,9 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { } Expect(k8sClient.Create(ctx, mcpServer)).Should(Succeed()) - // Wait for the MCPServer to have the original hash + // Wait for the MCPServer to be wired to the config (the external auth + // config hash propagates to its status) so the reference is + // observable before deletion. Eventually(func() bool { updatedMCPServer := &mcpv1beta1.MCPServer{} err := k8sClient.Get(ctx, types.NamespacedName{ @@ -563,66 +562,87 @@ var _ = Describe("MCPExternalAuthConfig Controller Integration Tests", func() { if err != nil { return false } - return updatedMCPServer.Status.ExternalAuthConfigHash == originalHash + return updatedMCPServer.Status.ExternalAuthConfigHash != "" }, timeout, interval).Should(BeTrue()) }) AfterAll(func() { - // Clean up resources - Expect(k8sClient.Delete(ctx, mcpServer)).Should(Succeed()) - Expect(k8sClient.Delete(ctx, authConfig)).Should(Succeed()) - Expect(k8sClient.Delete(ctx, oauthSecret)).Should(Succeed()) + // Best-effort cleanup; the specs may have already removed these. + _ = k8sClient.Delete(ctx, mcpServer) + _ = k8sClient.Delete(ctx, oauthSecret) + + // Ensure the MCPExternalAuthConfig is fully gone before the suite + // reuses the shared namespace. + Eventually(func() bool { + updated := &mcpv1beta1.MCPExternalAuthConfig{} + err := k8sClient.Get(ctx, types.NamespacedName{ + Name: authConfigName, + Namespace: namespace, + }, updated) + return errors.IsNotFound(err) + }, deletionRequeueTimeout, interval).Should(BeTrue()) }) - It("Should update config hash when auth config is modified", func() { - // Update the auth config - Eventually(func() error { - updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} + It("Should have the finalizer once it is being tracked", func() { + Eventually(func() []string { + updated := &mcpv1beta1.MCPExternalAuthConfig{} if err := k8sClient.Get(ctx, types.NamespacedName{ Name: authConfigName, Namespace: namespace, - }, updatedAuthConfig); err != nil { - return err + }, updated); err != nil { + return nil } + return updated.Finalizers + }, timeout, interval).Should(ContainElement("mcpexternalauthconfig.toolhive.stacklok.dev/finalizer")) + }) - // Modify the audience - updatedAuthConfig.Spec.TokenExchange.Audience = "updated-audience" - return k8sClient.Update(ctx, updatedAuthConfig) - }, timeout, interval).Should(Succeed()) + It("Should not be deleted while referenced and should set DeletionBlocked condition", func() { + // Request deletion; the finalizer must keep the object alive. + Expect(k8sClient.Delete(ctx, authConfig)).Should(Succeed()) - // Wait for the config hash to change - var newHash string + // The object should still exist (finalizer blocks deletion) with a + // pending deletion timestamp and a DeletionBlocked=True condition. Eventually(func() bool { - updatedAuthConfig := &mcpv1beta1.MCPExternalAuthConfig{} + updated := &mcpv1beta1.MCPExternalAuthConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ Name: authConfigName, Namespace: namespace, - }, updatedAuthConfig) + }, updated) if err != nil { return false } - newHash = updatedAuthConfig.Status.ConfigHash - return newHash != "" && newHash != originalHash + if updated.DeletionTimestamp.IsZero() { + return false + } + cond := meta.FindStatusCondition(updated.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) + return cond != nil && cond.Status == metav1.ConditionTrue }, timeout, interval).Should(BeTrue()) - // Verify the new hash is different - Expect(newHash).NotTo(Equal(originalHash)) + // Confirm it stays blocked (does not vanish) while still referenced. + Consistently(func() bool { + updated := &mcpv1beta1.MCPExternalAuthConfig{} + err := k8sClient.Get(ctx, types.NamespacedName{ + Name: authConfigName, + Namespace: namespace, + }, updated) + return err == nil + }, time.Second*2, interval).Should(BeTrue()) }) - It("Should trigger MCPServer reconciliation with updated hash", func() { - // Wait for the MCPServer to get the updated hash + It("Should be deleted after the referencing MCPServer is removed", func() { + // Remove the reference. + Expect(k8sClient.Delete(ctx, mcpServer)).Should(Succeed()) + + // With the workload watches removed, the unblock relies solely on the + // controller's periodic (up-to-30s) requeue, so allow extra time. Eventually(func() bool { - updatedMCPServer := &mcpv1beta1.MCPServer{} + updated := &mcpv1beta1.MCPExternalAuthConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: mcpServerName, + Name: authConfigName, Namespace: namespace, - }, updatedMCPServer) - if err != nil { - return false - } - // Check if the hash has been updated - return updatedMCPServer.Status.ExternalAuthConfigHash != originalHash - }, timeout, interval).Should(BeTrue()) + }, updated) + return errors.IsNotFound(err) + }, deletionRequeueTimeout, interval).Should(BeTrue()) }) }) }) diff --git a/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpremoteproxy_integration_test.go b/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpremoteproxy_integration_test.go index 3e08a884d3..1b8fc8aa39 100644 --- a/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpremoteproxy_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpremoteproxy_integration_test.go @@ -145,25 +145,6 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te }, timeout, interval).Should(BeTrue()) }) - It("should track MCPRemoteProxy in MCPOIDCConfig ReferencingWorkloads", func() { - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: proxyName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) }) Context("When MCPRemoteProxy references non-existent MCPOIDCConfig (fail-closed on missing)", Ordered, func() { @@ -414,23 +395,18 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te proxy = newTestMCPRemoteProxy(proxyName, namespace, configName) Expect(k8sClient.Create(ctx, proxy)).Should(Succeed()) - // Wait for ReferencingWorkloads to be populated + // Wait for the proxy to be wired to the config (OIDCConfigHash populated) + // so the config is observably referenced before we attempt deletion. Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} + updated := &mcpv1beta1.MCPRemoteProxy{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, + Name: proxyName, Namespace: namespace, }, updated) if err != nil { return false } - expectedRef := mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: proxyName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false + return updated.Status.OIDCConfigHash != "" }, timeout, interval).Should(BeTrue()) // Attempt to delete the MCPOIDCConfig (should be blocked by finalizer) @@ -455,8 +431,9 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should not be deleted while referenced by MCPRemoteProxy", func() { - // The object should still exist because the finalizer blocks deletion + It("should not be deleted while referenced and should set DeletionBlocked condition", func() { + // The object should still exist (finalizer blocks deletion) with a + // pending deletion timestamp and a DeletionBlocked=True condition. Eventually(func() bool { updated := &mcpv1beta1.MCPOIDCConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ @@ -466,7 +443,11 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te if err != nil { return false } - return !updated.DeletionTimestamp.IsZero() + if updated.DeletionTimestamp.IsZero() { + return false + } + cond := meta.FindStatusCondition(updated.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) + return cond != nil && cond.Status == metav1.ConditionTrue }, timeout, interval).Should(BeTrue()) }) @@ -549,25 +530,6 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te proxy = newTestMCPRemoteProxy(proxyName, namespace, configName) Expect(k8sClient.Create(ctx, proxy)).Should(Succeed()) - // Wait for ReferencingWorkloads to contain the proxy - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: proxyName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - // Wait for the proxy OIDCConfigHash to be populated Eventually(func() bool { updated := &mcpv1beta1.MCPRemoteProxy{} @@ -588,7 +550,7 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should clean up ReferencingWorkloads and clear OIDCConfigHash after ref removal", func() { + It("should clear OIDCConfigHash and remove condition after ref removal", func() { // Remove the OIDCConfigRef from the MCPRemoteProxy updated := &mcpv1beta1.MCPRemoteProxy{} Expect(k8sClient.Get(ctx, types.NamespacedName{ @@ -600,25 +562,6 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te updated.Spec.OIDCConfigRef = nil Expect(k8sClient.Update(ctx, updated)).Should(Succeed()) - // MCPOIDCConfig should no longer list MCPRemoteProxy in ReferencingWorkloads - Eventually(func() bool { - cfg := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, cfg) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: proxyName} - for _, ref := range cfg.Status.ReferencingWorkloads { - if ref == expectedRef { - return false - } - } - return true - }, timeout, interval).Should(BeTrue()) - // MCPRemoteProxy OIDCConfigHash should be cleared and condition removed Eventually(func() bool { proxyUpdated := &mcpv1beta1.MCPRemoteProxy{} @@ -638,109 +581,4 @@ var _ = Describe("MCPOIDCConfig and MCPRemoteProxy Cross-Resource Integration Te }, timeout, interval).Should(BeTrue()) }) }) - - Context("When MCPRemoteProxy is deleted, should clean up ReferencingWorkloads", Ordered, func() { - var ( - namespace string - configName string - proxyName string - oidcConfig *mcpv1beta1.MCPOIDCConfig - proxy *mcpv1beta1.MCPRemoteProxy - ns *corev1.Namespace - ) - - BeforeAll(func() { - ns = &corev1.Namespace{ - ObjectMeta: metav1.ObjectMeta{ - GenerateName: "test-proxy-oidcref-cleanup-", - }, - } - Expect(k8sClient.Create(ctx, ns)).Should(Succeed()) - namespace = ns.Name - - configName = testOIDCConfigName - proxyName = testRemoteProxyName - - // Create MCPOIDCConfig - oidcConfig = &mcpv1beta1.MCPOIDCConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: configName, - Namespace: namespace, - }, - Spec: mcpv1beta1.MCPOIDCConfigSpec{ - Type: mcpv1beta1.MCPOIDCConfigTypeInline, - Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "https://accounts.google.com", - ClientID: "test-client", - }, - }, - } - Expect(k8sClient.Create(ctx, oidcConfig)).Should(Succeed()) - - // Wait for ready - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - return updated.Status.ConfigHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create MCPRemoteProxy with OIDCConfigRef - proxy = newTestMCPRemoteProxy(proxyName, namespace, configName) - Expect(k8sClient.Create(ctx, proxy)).Should(Succeed()) - - // Wait for ReferencingWorkloads to contain the proxy - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: proxyName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - - AfterAll(func() { - _ = k8sClient.Delete(ctx, oidcConfig) - Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) - }) - - It("should remove MCPRemoteProxy from ReferencingWorkloads after deletion", func() { - // Delete the MCPRemoteProxy - Expect(k8sClient.Delete(ctx, proxy)).Should(Succeed()) - - // Eventually the referencing workloads list should not contain the proxy - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: mcpv1beta1.WorkloadKindMCPRemoteProxy, Name: proxyName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return false - } - } - return true - }, timeout, interval).Should(BeTrue()) - }) - }) }) diff --git a/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpserver_integration_test.go b/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpserver_integration_test.go index b61cfa66ac..d8be7502c9 100644 --- a/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpserver_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_mcpserver_integration_test.go @@ -139,138 +139,6 @@ var _ = Describe("MCPOIDCConfig and MCPServer Cross-Resource Integration Tests", }, timeout, interval).Should(BeTrue()) }) - It("should track MCPServer in MCPOIDCConfig ReferencingWorkloads", func() { - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: serverName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - }) - - Context("When MCPServer is deleted, should clean up ReferencingWorkloads", Ordered, func() { - var ( - namespace string - configName string - serverName string - oidcConfig *mcpv1beta1.MCPOIDCConfig - mcpServer *mcpv1beta1.MCPServer - ns *corev1.Namespace - ) - - BeforeAll(func() { - // Create a unique namespace for this test context - ns = &corev1.Namespace{ - ObjectMeta: metav1.ObjectMeta{ - GenerateName: "test-oidcref-cleanup-", - }, - } - Expect(k8sClient.Create(ctx, ns)).Should(Succeed()) - namespace = ns.Name - - configName = testOIDCConfigName - serverName = testServerName - - // Create MCPOIDCConfig - oidcConfig = &mcpv1beta1.MCPOIDCConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: configName, - Namespace: namespace, - }, - Spec: mcpv1beta1.MCPOIDCConfigSpec{ - Type: mcpv1beta1.MCPOIDCConfigTypeInline, - Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "https://accounts.google.com", - ClientID: "test-client", - }, - }, - } - Expect(k8sClient.Create(ctx, oidcConfig)).Should(Succeed()) - - // Wait for ready - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - return updated.Status.ConfigHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create MCPServer with OIDCConfigRef - mcpServer = &mcpv1beta1.MCPServer{ - ObjectMeta: metav1.ObjectMeta{ - Name: serverName, - Namespace: namespace, - }, - Spec: mcpv1beta1.MCPServerSpec{ - Image: testServerImage, - OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{ - Name: configName, - Audience: "test-audience", - Scopes: []string{"openid"}, - }, - }, - } - Expect(k8sClient.Create(ctx, mcpServer)).Should(Succeed()) - - // Wait for ReferencingWorkloads to contain the server - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: serverName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - - AfterAll(func() { - _ = k8sClient.Delete(ctx, oidcConfig) - Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) - }) - - It("should remove server from ReferencingWorkloads after MCPServer deletion", func() { - // Delete the MCPServer - Expect(k8sClient.Delete(ctx, mcpServer)).Should(Succeed()) - - // Eventually the referencing workloads list should be empty - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - return len(updated.Status.ReferencingWorkloads) == 0 - }, timeout, interval).Should(BeTrue()) - }) }) Context("When deleting MCPOIDCConfig with active references", Ordered, func() { @@ -342,23 +210,18 @@ var _ = Describe("MCPOIDCConfig and MCPServer Cross-Resource Integration Tests", } Expect(k8sClient.Create(ctx, mcpServer)).Should(Succeed()) - // Wait for ReferencingWorkloads to be populated + // Wait for the server to be wired to the config (OIDCConfigHash populated) + // so the config is observably referenced before we attempt deletion. Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} + updated := &mcpv1beta1.MCPServer{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, + Name: serverName, Namespace: namespace, }, updated) if err != nil { return false } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: serverName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false + return updated.Status.OIDCConfigHash != "" }, timeout, interval).Should(BeTrue()) // Attempt to delete the MCPOIDCConfig (should be blocked by finalizer) @@ -383,8 +246,9 @@ var _ = Describe("MCPOIDCConfig and MCPServer Cross-Resource Integration Tests", Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should not be deleted while referenced", func() { - // The object should still exist because the finalizer blocks deletion + It("should not be deleted while referenced and should set DeletionBlocked condition", func() { + // The object should still exist (finalizer blocks deletion) with a + // pending deletion timestamp and a DeletionBlocked=True condition. Eventually(func() bool { updated := &mcpv1beta1.MCPOIDCConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ @@ -394,7 +258,11 @@ var _ = Describe("MCPOIDCConfig and MCPServer Cross-Resource Integration Tests", if err != nil { return false } - return !updated.DeletionTimestamp.IsZero() + if updated.DeletionTimestamp.IsZero() { + return false + } + cond := meta.FindStatusCondition(updated.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) + return cond != nil && cond.Status == metav1.ConditionTrue }, timeout, interval).Should(BeTrue()) }) diff --git a/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_virtualmcpserver_integration_test.go b/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_virtualmcpserver_integration_test.go index c598e9e5c4..ec1f7f83b9 100644 --- a/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_virtualmcpserver_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-oidc-config/mcpoidcconfig_virtualmcpserver_integration_test.go @@ -6,6 +6,7 @@ package controllers import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/meta" @@ -183,155 +184,6 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration "resource should be the explicit resourceUrl, not the internal service URL") }) - It("should track VirtualMCPServer in MCPOIDCConfig ReferencingWorkloads", func() { - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "VirtualMCPServer", Name: vmcpName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - }) - - Context("When VirtualMCPServer is deleted, should clean up ReferencingWorkloads", Ordered, func() { - var ( - namespace string - configName string - vmcpName string - groupName string - oidcConfig *mcpv1beta1.MCPOIDCConfig - vmcpServer *mcpv1beta1.VirtualMCPServer - mcpGroup *mcpv1beta1.MCPGroup - ns *corev1.Namespace - ) - - BeforeAll(func() { - // Create a unique namespace for this test context - ns = &corev1.Namespace{ - ObjectMeta: metav1.ObjectMeta{ - GenerateName: "test-vmcp-oidcref-cleanup-", - }, - } - Expect(k8sClient.Create(ctx, ns)).Should(Succeed()) - namespace = ns.Name - - configName = testOIDCConfigName - vmcpName = testVMCPServerName - groupName = testVMCPGroupName - - // Create MCPGroup (required by VirtualMCPServer) - mcpGroup = &mcpv1beta1.MCPGroup{ - ObjectMeta: metav1.ObjectMeta{ - Name: groupName, - Namespace: namespace, - }, - } - Expect(k8sClient.Create(ctx, mcpGroup)).Should(Succeed()) - - // Create MCPOIDCConfig - oidcConfig = &mcpv1beta1.MCPOIDCConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: configName, - Namespace: namespace, - }, - Spec: mcpv1beta1.MCPOIDCConfigSpec{ - Type: mcpv1beta1.MCPOIDCConfigTypeInline, - Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "https://accounts.google.com", - ClientID: "test-client", - }, - }, - } - Expect(k8sClient.Create(ctx, oidcConfig)).Should(Succeed()) - - // Wait for ready - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - return updated.Status.ConfigHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create VirtualMCPServer with OIDCConfigRef - vmcpServer = v1beta1test.NewVirtualMCPServer(vmcpName, namespace, - v1beta1test.WithVMCPGroupRef(groupName), - v1beta1test.WithVMCPConfig(vmcpconfig.Config{Group: groupName}), - v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ - Type: "oidc", - OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{ - Name: configName, - Audience: "test-vmcp-audience", - Scopes: []string{"openid"}, - }, - }), - ) - Expect(k8sClient.Create(ctx, vmcpServer)).Should(Succeed()) - - // Wait for ReferencingWorkloads to contain the VirtualMCPServer - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "VirtualMCPServer", Name: vmcpName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - - AfterAll(func() { - _ = k8sClient.Delete(ctx, oidcConfig) - _ = k8sClient.Delete(ctx, mcpGroup) - Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) - }) - - It("should remove VirtualMCPServer from ReferencingWorkloads after deletion", func() { - // Delete the VirtualMCPServer - Expect(k8sClient.Delete(ctx, vmcpServer)).Should(Succeed()) - - // Eventually the referencing workloads list should not contain the vmcp entry - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "VirtualMCPServer", Name: vmcpName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return false - } - } - return true - }, timeout, interval).Should(BeTrue()) - }) }) Context("When deleting MCPOIDCConfig with active VirtualMCPServer references", Ordered, func() { @@ -413,23 +265,18 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration ) Expect(k8sClient.Create(ctx, vmcpServer)).Should(Succeed()) - // Wait for ReferencingWorkloads to be populated + // Wait for the VirtualMCPServer to be wired to the config (OIDCConfigHash + // populated) so the config is observably referenced before deletion. Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} + updated := &mcpv1beta1.VirtualMCPServer{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, + Name: vmcpName, Namespace: namespace, }, updated) if err != nil { return false } - expectedRef := mcpv1beta1.WorkloadReference{Kind: "VirtualMCPServer", Name: vmcpName} - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == expectedRef { - return true - } - } - return false + return updated.Status.OIDCConfigHash != "" }, timeout, interval).Should(BeTrue()) // Attempt to delete the MCPOIDCConfig (should be blocked by finalizer) @@ -455,8 +302,9 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should not be deleted while referenced by VirtualMCPServer", func() { - // The object should still exist because the finalizer blocks deletion + It("should not be deleted while referenced and should set DeletionBlocked condition", func() { + // The object should still exist (finalizer blocks deletion) with a + // pending deletion timestamp and a DeletionBlocked=True condition. Eventually(func() bool { updated := &mcpv1beta1.MCPOIDCConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ @@ -466,7 +314,11 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration if err != nil { return false } - return !updated.DeletionTimestamp.IsZero() + if updated.DeletionTimestamp.IsZero() { + return false + } + cond := meta.FindStatusCondition(updated.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) + return cond != nil && cond.Status == metav1.ConditionTrue }, timeout, interval).Should(BeTrue()) }) @@ -560,36 +412,47 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration }) }) - Context("When both MCPServer and VirtualMCPServer reference same MCPOIDCConfig", Ordered, func() { + Context("When MCPOIDCConfig inline.caBundleRef is set", Ordered, func() { var ( - namespace string - configName string - serverName string - vmcpName string - groupName string - oidcConfig *mcpv1beta1.MCPOIDCConfig - mcpServer *mcpv1beta1.MCPServer - vmcpServer *mcpv1beta1.VirtualMCPServer - mcpGroup *mcpv1beta1.MCPGroup - ns *corev1.Namespace + namespace string + configName string + vmcpName string + groupName string + caCMName string + caCMKey string + caConfigMap *corev1.ConfigMap + oidcConfig *mcpv1beta1.MCPOIDCConfig + vmcpServer *mcpv1beta1.VirtualMCPServer + mcpGroup *mcpv1beta1.MCPGroup + ns *corev1.Namespace ) BeforeAll(func() { - // Create a unique namespace for this test context ns = &corev1.Namespace{ - ObjectMeta: metav1.ObjectMeta{ - GenerateName: "test-vmcp-oidcref-both-", - }, + ObjectMeta: metav1.ObjectMeta{GenerateName: "test-vmcp-oidc-cabundle-"}, } Expect(k8sClient.Create(ctx, ns)).Should(Succeed()) namespace = ns.Name configName = testOIDCConfigName - serverName = testServerName vmcpName = testVMCPServerName groupName = testVMCPGroupName + caCMName = "vmcp-oidc-ca" + caCMKey = "ca.crt" + + // ConfigMap holding the CA bundle. Content is a placeholder — the operator + // only cares about mounting the ConfigMap at the right path. + caConfigMap = &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: caCMName, + Namespace: namespace, + }, + Data: map[string]string{ + caCMKey: "-----BEGIN CERTIFICATE-----\nplaceholder\n-----END CERTIFICATE-----\n", + }, + } + Expect(k8sClient.Create(ctx, caConfigMap)).Should(Succeed()) - // Create MCPGroup (required by VirtualMCPServer) mcpGroup = &mcpv1beta1.MCPGroup{ ObjectMeta: metav1.ObjectMeta{ Name: groupName, @@ -598,7 +461,6 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration } Expect(k8sClient.Create(ctx, mcpGroup)).Should(Succeed()) - // Create MCPOIDCConfig oidcConfig = &mcpv1beta1.MCPOIDCConfig{ ObjectMeta: metav1.ObjectMeta{ Name: configName, @@ -607,99 +469,123 @@ var _ = Describe("MCPOIDCConfig and VirtualMCPServer Cross-Resource Integration Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "https://accounts.google.com", + Issuer: "https://auth.example.internal/realms/demo", ClientID: "test-client", + CABundleRef: &mcpv1beta1.CABundleSource{ + ConfigMapRef: &corev1.ConfigMapKeySelector{ + LocalObjectReference: corev1.LocalObjectReference{Name: caCMName}, + Key: caCMKey, + }, + }, + JWKSAllowPrivateIP: true, }, }, } Expect(k8sClient.Create(ctx, oidcConfig)).Should(Succeed()) - // Wait for Valid condition and ConfigHash to be set Eventually(func() bool { updated := &mcpv1beta1.MCPOIDCConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ Name: configName, Namespace: namespace, }, updated) - if err != nil { - return false - } - if updated.Status.ConfigHash == "" { + if err != nil || updated.Status.ConfigHash == "" { return false } for _, cond := range updated.Status.Conditions { - if cond.Type == mcpv1beta1.ConditionTypeOIDCConfigValid && cond.Status == metav1.ConditionTrue { + if cond.Type == mcpv1beta1.ConditionTypeOIDCConfigValid && + cond.Status == metav1.ConditionTrue { return true } } return false }, timeout, interval).Should(BeTrue()) - // Create MCPServer with OIDCConfigRef - mcpServer = &mcpv1beta1.MCPServer{ + vmcpServer = &mcpv1beta1.VirtualMCPServer{ ObjectMeta: metav1.ObjectMeta{ - Name: serverName, + Name: vmcpName, Namespace: namespace, }, - Spec: mcpv1beta1.MCPServerSpec{ - Image: testServerImage, - OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{ - Name: configName, - Audience: "test-audience", - Scopes: []string{"openid"}, + Spec: mcpv1beta1.VirtualMCPServerSpec{ + GroupRef: &mcpv1beta1.MCPGroupRef{Name: groupName}, + Config: vmcpconfig.Config{Group: groupName}, + IncomingAuth: &mcpv1beta1.IncomingAuthConfig{ + Type: "oidc", + OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{ + Name: configName, + Audience: "test-vmcp-audience", + Scopes: []string{"openid"}, + }, }, }, } - Expect(k8sClient.Create(ctx, mcpServer)).Should(Succeed()) - - // Create VirtualMCPServer with OIDCConfigRef - vmcpServer = v1beta1test.NewVirtualMCPServer(vmcpName, namespace, - v1beta1test.WithVMCPGroupRef(groupName), - v1beta1test.WithVMCPConfig(vmcpconfig.Config{Group: groupName}), - v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ - Type: "oidc", - OIDCConfigRef: &mcpv1beta1.MCPOIDCConfigReference{ - Name: configName, - Audience: "test-vmcp-audience", - Scopes: []string{"openid"}, - }, - }), - ) Expect(k8sClient.Create(ctx, vmcpServer)).Should(Succeed()) }) AfterAll(func() { _ = k8sClient.Delete(ctx, vmcpServer) - _ = k8sClient.Delete(ctx, mcpServer) _ = k8sClient.Delete(ctx, oidcConfig) _ = k8sClient.Delete(ctx, mcpGroup) + _ = k8sClient.Delete(ctx, caConfigMap) Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should track both workloads in ReferencingWorkloads", func() { - Eventually(func() bool { - updated := &mcpv1beta1.MCPOIDCConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, + It("should render the mounted CA path into the vmcp ConfigMap's OIDC config", func() { + configMapName := vmcpName + "-vmcp-config" + expectedPath := "/config/certs/" + caCMName + "/" + caCMKey + + Eventually(func(g Gomega) { + cm := &corev1.ConfigMap{} + g.Expect(k8sClient.Get(ctx, types.NamespacedName{ + Name: configMapName, Namespace: namespace, - }, updated) - if err != nil { - return false - } - mcpServerRef := mcpv1beta1.WorkloadReference{Kind: "MCPServer", Name: serverName} - vmcpServerRef := mcpv1beta1.WorkloadReference{Kind: "VirtualMCPServer", Name: vmcpName} - hasMCPServer := false - hasVMCPServer := false - for _, ref := range updated.Status.ReferencingWorkloads { - if ref == mcpServerRef { - hasMCPServer = true + }, cm)).To(Succeed()) + g.Expect(cm.Data).To(HaveKey("config.yaml")) + + var config vmcpconfig.Config + g.Expect(yaml.Unmarshal([]byte(cm.Data["config.yaml"]), &config)).To(Succeed()) + g.Expect(config.IncomingAuth).NotTo(BeNil()) + g.Expect(config.IncomingAuth.OIDC).NotTo(BeNil()) + g.Expect(config.IncomingAuth.OIDC.CABundlePath).To(Equal(expectedPath), + "rendered vmcp config must contain the mounted CA path so the OIDC middleware can trust the issuer") + }, timeout, interval).Should(Succeed()) + }) + + It("should mount the CA ConfigMap as a read-only volume at /config/certs/", func() { + expectedMountPath := "/config/certs/" + caCMName + + Eventually(func(g Gomega) { + deployment := &appsv1.Deployment{} + g.Expect(k8sClient.Get(ctx, types.NamespacedName{ + Name: vmcpName, + Namespace: namespace, + }, deployment)).To(Succeed()) + + // The CA ConfigMap must be projected into a volume that sources from caCMName. + found := false + var volumeName string + for _, v := range deployment.Spec.Template.Spec.Volumes { + if v.ConfigMap != nil && v.ConfigMap.Name == caCMName { + found = true + volumeName = v.Name + break } - if ref == vmcpServerRef { - hasVMCPServer = true + } + g.Expect(found).To(BeTrue(), "Deployment must have a Volume sourcing from ConfigMap %q", caCMName) + + // The same volume must be mounted read-only at /config/certs/. + var mount *corev1.VolumeMount + for i := range deployment.Spec.Template.Spec.Containers[0].VolumeMounts { + m := &deployment.Spec.Template.Spec.Containers[0].VolumeMounts[i] + if m.Name == volumeName { + mount = m + break } } - return hasMCPServer && hasVMCPServer - }, timeout, interval).Should(BeTrue()) + g.Expect(mount).NotTo(BeNil(), "Deployment container must mount the CA volume") + g.Expect(mount.MountPath).To(Equal(expectedMountPath)) + g.Expect(mount.ReadOnly).To(BeTrue(), "CA bundle mount must be read-only") + }, timeout, interval).Should(Succeed()) }) }) }) diff --git a/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_authserverref_integration_test.go b/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_authserverref_integration_test.go index d38f291b5c..7b8d434dfc 100644 --- a/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_authserverref_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_authserverref_integration_test.go @@ -288,7 +288,8 @@ func newMCPOIDCConfig(name, namespace string) *mcpv1beta1.MCPOIDCConfig { Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "http://localhost:9090", + Issuer: "http://localhost:9090", + InsecureAllowHTTP: true, }, }, } diff --git a/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_cabundle_integration_test.go b/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_cabundle_integration_test.go index 5d8828d19d..c425d97636 100644 --- a/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_cabundle_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-remote-proxy/mcpremoteproxy_cabundle_integration_test.go @@ -25,7 +25,8 @@ func newMCPOIDCConfigWithCABundle(name, namespace, cmName, key string) *mcpv1bet Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "http://localhost:9090", + Issuer: "http://localhost:9090", + InsecureAllowHTTP: true, CABundleRef: &mcpv1beta1.CABundleSource{ ConfigMapRef: &corev1.ConfigMapKeySelector{ LocalObjectReference: corev1.LocalObjectReference{Name: cmName}, diff --git a/cmd/thv-operator/test-integration/mcp-server/mcpserver_authserverref_integration_test.go b/cmd/thv-operator/test-integration/mcp-server/mcpserver_authserverref_integration_test.go index 4fe4672ee1..7acb73208b 100644 --- a/cmd/thv-operator/test-integration/mcp-server/mcpserver_authserverref_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-server/mcpserver_authserverref_integration_test.go @@ -42,7 +42,8 @@ var _ = Describe("MCPServer AuthServerRef Integration Tests", func() { Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "http://localhost:9090", + Issuer: "http://localhost:9090", + InsecureAllowHTTP: true, }, }, } @@ -154,7 +155,8 @@ var _ = Describe("MCPServer AuthServerRef Integration Tests", func() { Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "http://localhost:9090", + Issuer: "http://localhost:9090", + InsecureAllowHTTP: true, }, }, } @@ -340,7 +342,8 @@ var _ = Describe("MCPServer AuthServerRef Integration Tests", func() { Spec: mcpv1beta1.MCPOIDCConfigSpec{ Type: mcpv1beta1.MCPOIDCConfigTypeInline, Inline: &mcpv1beta1.InlineOIDCSharedConfig{ - Issuer: "http://localhost:9090", + Issuer: "http://localhost:9090", + InsecureAllowHTTP: true, }, }, } diff --git a/cmd/thv-operator/test-integration/mcp-telemetry-config/mcptelemetryconfig_controller_integration_test.go b/cmd/thv-operator/test-integration/mcp-telemetry-config/mcptelemetryconfig_controller_integration_test.go index 1700aee0d4..08f0c8f945 100644 --- a/cmd/thv-operator/test-integration/mcp-telemetry-config/mcptelemetryconfig_controller_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-telemetry-config/mcptelemetryconfig_controller_integration_test.go @@ -172,147 +172,6 @@ var _ = Describe("MCPTelemetryConfig Controller", func() { }, timeout, interval).Should(BeTrue()) }) - It("should track referencing MCPServers in status", func() { - // Create a telemetry config - telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-ref-tracking", - Namespace: "default", - }, - } - telemetryConfig.Spec.OpenTelemetry = &mcpv1beta1.MCPTelemetryOTelConfig{ - Enabled: true, - Endpoint: testEndpoint, - Tracing: &mcpv1beta1.OpenTelemetryTracingConfig{Enabled: true}, - } - - Expect(k8sClient.Create(ctx, telemetryConfig)).To(Succeed()) - - // Wait for initial reconciliation (finalizer + hash) - Eventually(func() bool { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - return err == nil && fetched.Status.ConfigHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create an MCPServer that references this config - server := &mcpv1beta1.MCPServer{ - ObjectMeta: metav1.ObjectMeta{ - Name: "server-ref-tracking", - Namespace: "default", - }, - Spec: mcpv1beta1.MCPServerSpec{ - Image: "example/mcp-server:latest", - TelemetryConfigRef: &mcpv1beta1.MCPTelemetryConfigReference{ - Name: "test-ref-tracking", - }, - }, - } - Expect(k8sClient.Create(ctx, server)).To(Succeed()) - - // The MCPServer watch should trigger a reconciliation of the MCPTelemetryConfig. - // Verify ReferencingWorkloads is updated to include our server. - Eventually(func() []string { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - if err != nil { - return nil - } - names := make([]string, 0, len(fetched.Status.ReferencingWorkloads)) - for _, ref := range fetched.Status.ReferencingWorkloads { - names = append(names, ref.Name) - } - return names - }, timeout, interval).Should(ContainElement("server-ref-tracking")) - }) - - It("should prune referencing MCPServer from status after the MCPServer is deleted", func() { - // Create a telemetry config - telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-ref-prune", - Namespace: "default", - }, - } - telemetryConfig.Spec.OpenTelemetry = &mcpv1beta1.MCPTelemetryOTelConfig{ - Enabled: true, - Endpoint: testEndpoint, - Tracing: &mcpv1beta1.OpenTelemetryTracingConfig{Enabled: true}, - } - - Expect(k8sClient.Create(ctx, telemetryConfig)).To(Succeed()) - - // Wait for initial reconciliation (finalizer + hash) - Eventually(func() bool { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - return err == nil && fetched.Status.ConfigHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create an MCPServer that references this config - server := &mcpv1beta1.MCPServer{ - ObjectMeta: metav1.ObjectMeta{ - Name: "server-ref-prune", - Namespace: "default", - }, - Spec: mcpv1beta1.MCPServerSpec{ - Image: "example/mcp-server:latest", - TelemetryConfigRef: &mcpv1beta1.MCPTelemetryConfigReference{ - Name: "test-ref-prune", - }, - }, - } - Expect(k8sClient.Create(ctx, server)).To(Succeed()) - - // Wait until the server is tracked in ReferencingWorkloads. - Eventually(func() []string { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - if err != nil { - return nil - } - names := make([]string, 0, len(fetched.Status.ReferencingWorkloads)) - for _, ref := range fetched.Status.ReferencingWorkloads { - names = append(names, ref.Name) - } - return names - }, timeout, interval).Should(ContainElement("server-ref-prune")) - - // Delete the MCPServer. The controller-runtime watch enqueues the - // referenced MCPTelemetryConfig (via the old-object map function), and - // the level-triggered reconcile must prune the now-stale reference. - Expect(k8sClient.Delete(ctx, server)).To(Succeed()) - - // Eventually the reference is pruned from status. - Eventually(func() []string { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - if err != nil { - return nil - } - names := make([]string, 0, len(fetched.Status.ReferencingWorkloads)) - for _, ref := range fetched.Status.ReferencingWorkloads { - names = append(names, ref.Name) - } - return names - }, timeout, interval).ShouldNot(ContainElement("server-ref-prune")) - }) - It("should block deletion when MCPServers reference the config", func() { // Create a telemetry config telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{ @@ -362,25 +221,20 @@ var _ = Describe("MCPTelemetryConfig Controller", func() { } Expect(k8sClient.Create(ctx, server)).To(Succeed()) - // Wait for ReferencingWorkloads to be populated - Eventually(func() []string { - fetched := &mcpv1beta1.MCPTelemetryConfig{} + // Ensure the referencing MCPServer is observable before attempting deletion. + Eventually(func() bool { + fetched := &mcpv1beta1.MCPServer{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, + Name: server.Name, + Namespace: server.Namespace, }, fetched) - if err != nil { - return nil - } - names := make([]string, 0, len(fetched.Status.ReferencingWorkloads)) - for _, ref := range fetched.Status.ReferencingWorkloads { - names = append(names, ref.Name) - } - return names - }, timeout, interval).Should(ContainElement("server-deletion-blocker")) + return err == nil + }, timeout, interval).Should(BeTrue()) // Attempt to delete the config — the API call succeeds (sets DeletionTimestamp) - // but the finalizer blocks actual removal + // but the finalizer blocks actual removal because the config is still + // referenced. The controller recomputes referencing workloads on demand + // (via its field index) during the deletion reconcile. Expect(k8sClient.Delete(ctx, telemetryConfig)).To(Succeed()) // Verify the object still exists (finalizer prevents deletion) @@ -407,58 +261,6 @@ var _ = Describe("MCPTelemetryConfig Controller", func() { }, timeout, interval).Should(BeTrue(), "Config should be deleted after references are removed") }) - It("should track MCPRemoteProxy in ReferencingWorkloads", func() { - telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{ - ObjectMeta: metav1.ObjectMeta{ - Name: "test-proxy-ref-tracking", - Namespace: "default", - }, - } - telemetryConfig.Spec.OpenTelemetry = &mcpv1beta1.MCPTelemetryOTelConfig{ - Enabled: true, - Endpoint: testEndpoint, - Tracing: &mcpv1beta1.OpenTelemetryTracingConfig{Enabled: true}, - } - - Expect(k8sClient.Create(ctx, telemetryConfig)).To(Succeed()) - - // Wait for config to be ready - Eventually(func() bool { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - return err == nil && fetched.Status.ConfigHash != "" - }, timeout, interval).Should(BeTrue()) - - // Create an MCPRemoteProxy that references this config - proxy := v1beta1test.NewMCPRemoteProxy("proxy-ref-tracking", "default", - v1beta1test.WithRemoteProxyURL("https://example.com/mcp"), - v1beta1test.WithRemoteProxyPort(0), - v1beta1test.WithRemoteProxyTelemetryConfigRef("test-proxy-ref-tracking"), - ) - Expect(k8sClient.Create(ctx, proxy)).To(Succeed()) - - // The MCPRemoteProxy watch should trigger reconciliation of MCPTelemetryConfig. - // Verify ReferencingWorkloads includes the proxy. - Eventually(func() []string { - fetched := &mcpv1beta1.MCPTelemetryConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, - }, fetched) - if err != nil { - return nil - } - names := make([]string, 0, len(fetched.Status.ReferencingWorkloads)) - for _, ref := range fetched.Status.ReferencingWorkloads { - names = append(names, ref.Kind+"/"+ref.Name) - } - return names - }, timeout, interval).Should(ContainElement("MCPRemoteProxy/proxy-ref-tracking")) - }) - It("should block deletion when MCPRemoteProxy references the config", func() { telemetryConfig := &mcpv1beta1.MCPTelemetryConfig{ ObjectMeta: metav1.ObjectMeta{ @@ -500,24 +302,19 @@ var _ = Describe("MCPTelemetryConfig Controller", func() { ) Expect(k8sClient.Create(ctx, proxy)).To(Succeed()) - // Wait for ReferencingWorkloads to include the proxy - Eventually(func() []string { - fetched := &mcpv1beta1.MCPTelemetryConfig{} + // Ensure the referencing MCPRemoteProxy is observable before attempting deletion. + Eventually(func() bool { + fetched := &mcpv1beta1.MCPRemoteProxy{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: telemetryConfig.Name, - Namespace: telemetryConfig.Namespace, + Name: proxy.Name, + Namespace: proxy.Namespace, }, fetched) - if err != nil { - return nil - } - names := make([]string, 0, len(fetched.Status.ReferencingWorkloads)) - for _, ref := range fetched.Status.ReferencingWorkloads { - names = append(names, ref.Name) - } - return names - }, timeout, interval).Should(ContainElement("proxy-deletion-blocker")) + return err == nil + }, timeout, interval).Should(BeTrue()) - // Attempt to delete — finalizer blocks removal + // Attempt to delete — finalizer blocks removal because the config is still + // referenced. The controller recomputes referencing workloads on demand + // (via its field index) during the deletion reconcile. Expect(k8sClient.Delete(ctx, telemetryConfig)).To(Succeed()) // Verify object still exists diff --git a/cmd/thv-operator/test-integration/mcp-toolconfig/mcptoolconfig_controller_integration_test.go b/cmd/thv-operator/test-integration/mcp-toolconfig/mcptoolconfig_controller_integration_test.go index 54b03f486c..fe0b6fb9bb 100644 --- a/cmd/thv-operator/test-integration/mcp-toolconfig/mcptoolconfig_controller_integration_test.go +++ b/cmd/thv-operator/test-integration/mcp-toolconfig/mcptoolconfig_controller_integration_test.go @@ -305,25 +305,6 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should track referencing workloads in status", func() { - Eventually(func() bool { - updated := &mcpv1beta1.MCPToolConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - for _, ref := range updated.Status.ReferencingWorkloads { - if ref.Kind == "MCPServer" && ref.Name == mcpServerName { - return true - } - } - return false - }, timeout, interval).Should(BeTrue()) - }) - It("should propagate MCPToolConfig changes to referencing MCPServers", func() { var initialHash string Eventually(func() string { @@ -374,24 +355,6 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { return updated.Status.ToolConfigHash }, timeout, interval).Should(Equal(updatedConfigHash)) }) - - It("should remove server from status when MCPServer is deleted", func() { - // Delete the MCPServer - Expect(k8sClient.Delete(ctx, mcpServer)).Should(Succeed()) - - // Eventually the referencing workloads list should be empty - Eventually(func() bool { - updated := &mcpv1beta1.MCPToolConfig{} - err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated) - if err != nil { - return false - } - return len(updated.Status.ReferencingWorkloads) == 0 - }, timeout, interval).Should(BeTrue()) - }) }) Context("When deleting MCPToolConfig with active references", Ordered, func() { @@ -457,22 +420,18 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { } Expect(k8sClient.Create(ctx, mcpServer)).Should(Succeed()) - // Wait for ReferencingWorkloads to be populated + // Wait for the MCPServer to be wired to the config (ToolConfigHash + // populated) so the config is observably referenced before deletion. Eventually(func() bool { - updated := &mcpv1beta1.MCPToolConfig{} + updated := &mcpv1beta1.MCPServer{} err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, + Name: mcpServerName, Namespace: namespace, }, updated) if err != nil { return false } - for _, ref := range updated.Status.ReferencingWorkloads { - if ref.Kind == "MCPServer" && ref.Name == mcpServerName { - return true - } - } - return false + return updated.Status.ToolConfigHash != "" }, timeout, interval).Should(BeTrue()) // Attempt to delete the MCPToolConfig (should be blocked by finalizer) @@ -497,8 +456,9 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("should not be deleted while referenced", func() { - // The object should still exist because the finalizer blocks deletion + It("should not be deleted while referenced and should set DeletionBlocked condition", func() { + // The object should still exist (finalizer blocks deletion) with a + // pending deletion timestamp and a DeletionBlocked=True condition. Eventually(func() bool { updated := &mcpv1beta1.MCPToolConfig{} err := k8sClient.Get(ctx, types.NamespacedName{ @@ -508,7 +468,11 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { if err != nil { return false } - return !updated.DeletionTimestamp.IsZero() + if updated.DeletionTimestamp.IsZero() { + return false + } + cond := meta.FindStatusCondition(updated.Status.Conditions, mcpv1beta1.ConditionTypeDeletionBlocked) + return cond != nil && cond.Status == metav1.ConditionTrue }, timeout, interval).Should(BeTrue()) }) @@ -528,11 +492,11 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { }) }) - // These tests assert that an MCPToolConfig's ToolsFilter and ToolsOverride are - // not merely tracked via Status.ReferencingWorkloads, but actually propagate - // into the referencing MCPServer's rendered RunConfig ConfigMap. Both ride the - // same code path in createRunConfigFromMCPServer (WithToolsFilter / - // WithToolsOverride), so this single context covers filtering and renaming. + // These tests assert that an MCPToolConfig's ToolsFilter and ToolsOverride + // actually propagate into the referencing MCPServer's rendered RunConfig + // ConfigMap. Both ride the same code path in createRunConfigFromMCPServer + // (WithToolsFilter / WithToolsOverride), so this single context covers + // filtering and renaming. Context("When a referencing MCPServer renders its RunConfig", Ordered, func() { var ( namespace string @@ -707,25 +671,6 @@ var _ = Describe("MCPToolConfig Controller Integration Tests", func() { Expect(k8sClient.Delete(ctx, ns)).Should(Succeed()) }) - It("tracks every referencing MCPServer in status", func() { - Eventually(func() []string { - updated := &mcpv1beta1.MCPToolConfig{} - if err := k8sClient.Get(ctx, types.NamespacedName{ - Name: configName, - Namespace: namespace, - }, updated); err != nil { - return nil - } - var names []string - for _, ref := range updated.Status.ReferencingWorkloads { - if ref.Kind == "MCPServer" { - names = append(names, ref.Name) - } - } - return names - }, timeout, interval).Should(ConsistOf(serverNames)) - }) - It("propagates a config change to every referencing MCPServer", func() { // Capture the initial hash the servers converged on. initial := &mcpv1beta1.MCPToolConfig{} diff --git a/cmd/thv-operator/test-integration/testutil/envtest.go b/cmd/thv-operator/test-integration/testutil/envtest.go index bc2a6f5122..43ae4278e2 100644 --- a/cmd/thv-operator/test-integration/testutil/envtest.go +++ b/cmd/thv-operator/test-integration/testutil/envtest.go @@ -22,6 +22,7 @@ import ( "k8s.io/apimachinery/pkg/runtime" "k8s.io/client-go/kubernetes/scheme" "k8s.io/client-go/rest" + "k8s.io/utils/ptr" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/envtest" @@ -33,9 +34,19 @@ import ( mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" ) -// gracefulShutdownDelay gives a running manager a moment to stop cleanly after -// the suite context is cancelled, before the envtest control plane is torn down. -const gracefulShutdownDelay = 100 * time.Millisecond +const ( + // managerGracefulShutdownTimeout bounds how long the manager waits for its + // runnables (controllers and the informer cache) to return once the suite + // context is cancelled. It is deliberately well under envtest's 20s + // ControlPlaneStopTimeout so a wedged runnable surfaces as a manager + // shutdown timeout rather than as a kube-apiserver teardown failure. + managerGracefulShutdownTimeout = 10 * time.Second + + // managerStopTimeout bounds how long Stop waits for Manager.Start to return. + // It allows a little slack over managerGracefulShutdownTimeout so the + // manager's own bounded shutdown always wins the race. + managerStopTimeout = 15 * time.Second +) // SuiteOptions configures the envtest bootstrap performed by StartSuite. Add // fields only as concrete suites require them. @@ -64,6 +75,10 @@ type SuiteEnv struct { testEnv *envtest.Environment cancel context.CancelFunc + // managerDone is closed when the goroutine started by StartManager returns, + // i.e. once Manager.Start has fully shut down every runnable. It stays nil + // until StartManager is called. + managerDone chan struct{} } // StartSuite bootstraps a full operator integration suite: it starts an envtest @@ -104,6 +119,9 @@ func StartSuite(opts SuiteOptions) *SuiteEnv { BindAddress: "0", // Disable metrics server for tests to avoid port conflicts }, HealthProbeBindAddress: "0", // Disable health probe for tests + // Keep the manager's own shutdown bounded so Stop never waits on it + // longer than envtest is willing to wait for the control plane. + GracefulShutdownTimeout: ptr.To(managerGracefulShutdownTimeout), }) Expect(err).ToNot(HaveOccurred()) @@ -162,21 +180,40 @@ func (e *SuiteEnv) StartManager() { if e.Manager == nil { return } + e.managerDone = make(chan struct{}) go func() { defer GinkgoRecover() + // Closing before GinkgoRecover runs (deferred functions are LIFO) means + // Stop is released even if the manager fails. + defer close(e.managerDone) Expect(e.Manager.Start(e.Ctx)).To(Succeed(), "failed to run manager") }() } -// Stop cancels the suite context and tears down the envtest control plane. When -// a manager is running it first waits a short grace period so the manager can -// shut down cleanly before the API server goes away. +// Stop cancels the suite context and tears down the envtest control plane. +// +// When a manager is running it waits for Manager.Start to return before +// stopping the control plane. That wait is load-bearing, not politeness: the +// manager's informers hold long-running watch requests against kube-apiserver, +// and kube-apiserver's graceful shutdown blocks on draining them. Signalling +// the API server to stop while watches are still open makes it miss envtest's +// 20s ControlPlaneStopTimeout, which surfaces as a flaky +// "timeout waiting for process kube-apiserver to stop" failure in AfterSuite. +// The suites with fewer specs than Ginkgo procs (MCPTelemetryConfig, MCPGroup) +// hit it most often: their idle procs cancel the suite context moments after +// the informers established their watches. func (e *SuiteEnv) Stop() { By("tearing down the test environment") e.cancel() - if e.Manager != nil { - // Give the manager some time to shut down gracefully. - time.Sleep(gracefulShutdownDelay) + if e.managerDone != nil { + select { + case <-e.managerDone: + case <-time.After(managerStopTimeout): + // Fall through to the control plane teardown anyway: reporting the + // apiserver timeout on top of this would only obscure the real + // problem, which is a runnable that refused to stop. + AddReportEntry("manager did not shut down within " + managerStopTimeout.String()) + } } Expect(e.testEnv.Stop()).NotTo(HaveOccurred()) } diff --git a/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_compositetool_watch_test.go b/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_compositetool_watch_test.go index 33caa4682e..38e8bfd027 100644 --- a/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_compositetool_watch_test.go +++ b/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_compositetool_watch_test.go @@ -393,6 +393,39 @@ var _ = Describe("VirtualMCPServer CompositeToolDefinition Watch Integration Tes }) It("Should NOT trigger VirtualMCPServer reconciliation when unrelated composite tool definition is created", func() { + // BeforeAll only gates on ObservedGeneration > 0, so the controller + // may still be settling (the Ready condition flipping as discovery + // and config generation complete) when this spec starts. A baseline + // captured mid-settle fails the timestamp assertion below even + // though the unrelated create never triggered anything. Wait until + // the Ready condition has held the same transition time across + // several consecutive polls before taking the baseline. + var lastSeenReadyTime metav1.Time + stablePolls := 0 + Eventually(func() bool { + vmcpNow := &mcpv1beta1.VirtualMCPServer{} + if err := k8sClient.Get(ctx, types.NamespacedName{ + Name: vmcpName, + Namespace: namespace, + }, vmcpNow); err != nil { + return false + } + for _, cond := range vmcpNow.Status.Conditions { + if cond.Type == conditionReady { + if cond.LastTransitionTime.Equal(&lastSeenReadyTime) { + stablePolls++ + } else { + lastSeenReadyTime = cond.LastTransitionTime + stablePolls = 0 + } + // ~2s of stability at the 250ms poll interval: the + // observed flake transitioned ~2s after the baseline. + return stablePolls >= 8 + } + } + return false + }, timeout, interval).Should(BeTrue(), "Ready condition should settle before capturing the baseline") + // Get initial generation and observed generation initialVMCP := &mcpv1beta1.VirtualMCPServer{} Expect(k8sClient.Get(ctx, types.NamespacedName{ diff --git a/cmd/thv/app/ai_plugin.go b/cmd/thv/app/ai_plugin.go new file mode 100644 index 0000000000..45fc1c7820 --- /dev/null +++ b/cmd/thv/app/ai_plugin.go @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "github.com/spf13/cobra" +) + +var aiPluginCmd = &cobra.Command{ + Use: "ai-plugin", + Short: "Manage AI-tool plugins", + Long: `The ai-plugin command provides subcommands to manage plugins for AI tools +(e.g. Claude Code, Codex) — not plugins for ToolHive itself.`, +} diff --git a/cmd/thv/app/ai_plugin_build.go b/cmd/thv/app/ai_plugin_build.go new file mode 100644 index 0000000000..15d5270ddf --- /dev/null +++ b/cmd/thv/app/ai_plugin_build.go @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "fmt" + "path/filepath" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var aiPluginBuildTag string + +var aiPluginBuildCmd = &cobra.Command{ + Use: "build [path]", + Short: "Build an AI-tool plugin into a local OCI artifact", + Long: `Build a plugin from a local directory into an OCI artifact that can be pushed to a registry. + +On success, prints the OCI reference of the built artifact to stdout.`, + Args: cobra.ExactArgs(1), + ValidArgsFunction: func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + return nil, cobra.ShellCompDirectiveFilterDirs + }, + RunE: aiPluginBuildCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginBuildCmd) + + aiPluginBuildCmd.Flags().StringVarP(&aiPluginBuildTag, "tag", "t", "", "OCI tag for the built artifact") +} + +func aiPluginBuildCmdFunc(cmd *cobra.Command, args []string) error { + absPath, err := filepath.Abs(args[0]) + if err != nil { + return fmt.Errorf("failed to resolve path: %w", err) + } + + c := newAIPluginClient(cmd.Context()) + + result, err := c.Build(cmd.Context(), plugins.BuildOptions{ + Path: absPath, + Tag: aiPluginBuildTag, + }) + if err != nil { + return formatAIPluginError("build plugin", err) + } + + fmt.Println(result.Reference) + return nil +} diff --git a/cmd/thv/app/ai_plugin_builds.go b/cmd/thv/app/ai_plugin_builds.go new file mode 100644 index 0000000000..e9f6e108b1 --- /dev/null +++ b/cmd/thv/app/ai_plugin_builds.go @@ -0,0 +1,82 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "encoding/json" + "fmt" + "os" + "text/tabwriter" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var aiPluginBuildsFormat string + +var aiPluginBuildsCmd = &cobra.Command{ + Use: "builds", + Short: "List locally-built AI-tool plugin artifacts", + Long: `List all locally-built OCI plugin artifacts stored in the local OCI store.`, + PreRunE: chainPreRunE( + ValidateFormat(&aiPluginBuildsFormat), + ), + RunE: aiPluginBuildsCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginBuildsCmd) + + AddFormatFlag(aiPluginBuildsCmd, &aiPluginBuildsFormat) +} + +func aiPluginBuildsCmdFunc(cmd *cobra.Command, _ []string) error { + c := newAIPluginClient(cmd.Context()) + + builds, err := c.ListBuilds(cmd.Context()) + if err != nil { + return formatAIPluginError("list builds", err) + } + + switch aiPluginBuildsFormat { + case FormatJSON: + if builds == nil { + builds = []plugins.LocalBuild{} + } + data, err := json.MarshalIndent(builds, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal JSON: %w", err) + } + fmt.Println(string(data)) + default: + if len(builds) == 0 { + fmt.Println("No locally-built plugin artifacts found") + return nil + } + printAIPluginBuildsText(builds) + } + + return nil +} + +func printAIPluginBuildsText(builds []plugins.LocalBuild) { + w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', 0) + _, _ = fmt.Fprintln(w, "TAG\tDIGEST\tNAME\tVERSION") + + for _, b := range builds { + digest := b.Digest + if len(digest) > 19 { + digest = digest[:19] + "..." + } + _, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", + b.Tag, + digest, + b.Name, + b.Version, + ) + } + + _ = w.Flush() +} diff --git a/cmd/thv/app/ai_plugin_builds_remove.go b/cmd/thv/app/ai_plugin_builds_remove.go new file mode 100644 index 0000000000..bb18539313 --- /dev/null +++ b/cmd/thv/app/ai_plugin_builds_remove.go @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "github.com/spf13/cobra" +) + +var aiPluginBuildsRemoveCmd = &cobra.Command{ + Use: "remove ", + Short: "Remove a locally-built AI-tool plugin artifact", + Long: `Remove a locally-built OCI plugin artifact and its blobs from the local OCI store.`, + Args: cobra.ExactArgs(1), + RunE: aiPluginBuildsRemoveCmdFunc, +} + +func init() { + aiPluginBuildsCmd.AddCommand(aiPluginBuildsRemoveCmd) +} + +func aiPluginBuildsRemoveCmdFunc(cmd *cobra.Command, args []string) error { + c := newAIPluginClient(cmd.Context()) + if err := c.DeleteBuild(cmd.Context(), args[0]); err != nil { + return formatAIPluginError("remove build", err) + } + return nil +} diff --git a/cmd/thv/app/ai_plugin_helpers.go b/cmd/thv/app/ai_plugin_helpers.go new file mode 100644 index 0000000000..78f7b9677d --- /dev/null +++ b/cmd/thv/app/ai_plugin_helpers.go @@ -0,0 +1,64 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "context" + "errors" + "fmt" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" + pluginclient "github.com/stacklok/toolhive/pkg/plugins/client" +) + +// newAIPluginClient creates a new Plugins API HTTP client using default settings. +// The context is used for server discovery; it is not stored. +func newAIPluginClient(ctx context.Context) *pluginclient.Client { + return pluginclient.NewDefaultClient(ctx) +} + +// completeAIPluginNames provides shell completion for installed plugin names. +func completeAIPluginNames(cmd *cobra.Command, args []string, _ string) ([]string, cobra.ShellCompDirective) { + if len(args) > 0 { + return nil, cobra.ShellCompDirectiveNoFileComp + } + + c := newAIPluginClient(cmd.Context()) + installed, err := c.List(cmd.Context(), plugins.ListOptions{}) + if err != nil { + return nil, cobra.ShellCompDirectiveError + } + + names := make([]string, 0, len(installed)) + for _, p := range installed { + names = append(names, p.Metadata.Name) + } + return names, cobra.ShellCompDirectiveNoFileComp +} + +// formatAIPluginError wraps an error with contextual information, appending a +// hint that matches the actual failure — a timed-out request and an absent +// server need different advice. +func formatAIPluginError(action string, err error) error { + switch { + case errors.Is(err, pluginclient.ErrRequestTimeout): + return fmt.Errorf( + "failed to %s: %w\nHint: the server is running and was still working; "+ + "raise the limit with TOOLHIVE_API_TIMEOUT (e.g. TOOLHIVE_API_TIMEOUT=30m)", + action, err) + case errors.Is(err, pluginclient.ErrServerUnreachable): + return fmt.Errorf("failed to %s: %w\nHint: ensure 'thv serve' is running", action, err) + default: + return fmt.Errorf("failed to %s: %w", action, err) + } +} + +// validateAIPluginScope returns a PreRunE that validates the --scope flag. +func validateAIPluginScope(scopeVar *string) func(*cobra.Command, []string) error { + return func(_ *cobra.Command, _ []string) error { + return plugins.ValidateScope(plugins.Scope(*scopeVar)) + } +} diff --git a/cmd/thv/app/ai_plugin_info.go b/cmd/thv/app/ai_plugin_info.go new file mode 100644 index 0000000000..8d3c7c1317 --- /dev/null +++ b/cmd/thv/app/ai_plugin_info.go @@ -0,0 +1,153 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "encoding/json" + "fmt" + "maps" + "os" + "slices" + "sort" + "strings" + "text/tabwriter" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var ( + aiPluginInfoScope string + aiPluginInfoFormat string + aiPluginInfoProjectRoot string +) + +var aiPluginInfoCmd = &cobra.Command{ + Use: "info [plugin-name]", + Short: "Show details of an AI-tool plugin", + Long: `Display detailed information about a plugin, including metadata, version, and installation status.`, + Args: cobra.ExactArgs(1), + ValidArgsFunction: completeAIPluginNames, + PreRunE: chainPreRunE( + validateAIPluginScope(&aiPluginInfoScope), + ValidateFormat(&aiPluginInfoFormat), + ), + RunE: aiPluginInfoCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginInfoCmd) + + aiPluginInfoCmd.Flags().StringVar(&aiPluginInfoScope, "scope", "", "Filter by scope (user, project)") + AddFormatFlag(aiPluginInfoCmd, &aiPluginInfoFormat) + aiPluginInfoCmd.Flags().StringVar(&aiPluginInfoProjectRoot, "project-root", "", "Project root path for project-scoped plugins") +} + +func aiPluginInfoCmdFunc(cmd *cobra.Command, args []string) error { + c := newAIPluginClient(cmd.Context()) + + projectRoot, err := absProjectRoot(aiPluginInfoProjectRoot) + if err != nil { + return err + } + + info, err := c.Info(cmd.Context(), plugins.InfoOptions{ + Name: args[0], + Scope: plugins.Scope(aiPluginInfoScope), + ProjectRoot: projectRoot, + }) + if err != nil { + return formatAIPluginError("get plugin info", err) + } + + switch aiPluginInfoFormat { + case FormatJSON: + data, err := json.MarshalIndent(info, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal JSON: %w", err) + } + fmt.Println(string(data)) + default: + printAIPluginInfoText(info) + } + + return nil +} + +func printAIPluginInfoText(info *plugins.PluginInfo) { + w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0) + + _, _ = fmt.Fprintf(w, "Name:\t%s\n", info.Metadata.Name) + _, _ = fmt.Fprintf(w, "Version:\t%s\n", info.Metadata.Version) + _, _ = fmt.Fprintf(w, "Description:\t%s\n", info.Metadata.Description) + + if s := info.InstalledPlugin; s != nil { + _, _ = fmt.Fprintf(w, "Scope:\t%s\n", s.Scope) + _, _ = fmt.Fprintf(w, "Status:\t%s\n", s.Status) + _, _ = fmt.Fprintf(w, "Reference:\t%s\n", s.Reference) + _, _ = fmt.Fprintf(w, "Installed At:\t%s\n", s.InstalledAt.Format("2006-01-02 15:04:05")) + if len(s.Clients) > 0 { + _, _ = fmt.Fprintf(w, "Clients:\t%s\n", strings.Join(s.Clients, ", ")) + } + if len(s.Components) > 0 { + managed, declared := splitComponentInventory(s.Components) + if len(managed) > 0 { + _, _ = fmt.Fprintf(w, "Components:\t%s\n", formatComponentInventory(managed)) + } + if len(declared) > 0 { + _, _ = fmt.Fprintf(w, "Declared (not managed by ToolHive):\t%s\n", formatComponentInventory(declared)) + } + } + } + + if len(info.UnmaterializedComponents) > 0 { + _, _ = fmt.Fprintln(w, "\nNot activated on this client (declared by the plugin, but the client does not load them):") + for _, clientType := range slices.Sorted(maps.Keys(info.UnmaterializedComponents)) { + types := info.UnmaterializedComponents[clientType] + labels := make([]string, 0, len(types)) + for _, t := range types { + labels = append(labels, string(t)) + } + _, _ = fmt.Fprintf(w, " %s:\t%s\n", clientType, strings.Join(labels, ", ")) + } + } + if len(info.ProjectScopeDegradedClients) > 0 { + degraded := append([]string(nil), info.ProjectScopeDegradedClients...) + sort.Strings(degraded) + _, _ = fmt.Fprintf(w, "\nInstalled at user scope instead of project scope (client cannot scope per-project):\t%s\n", + strings.Join(degraded, ", ")) + } + + _ = w.Flush() +} + +// formatComponentInventory renders a ComponentInventory (map[string]int) as a +// sorted, space-separated "key=count" sequence for deterministic output. +func formatComponentInventory(inv plugins.ComponentInventory) string { + keys := slices.Sorted(maps.Keys(inv)) + parts := make([]string, 0, len(keys)) + for _, k := range keys { + parts = append(parts, fmt.Sprintf("%s=%d", k, inv[k])) + } + return strings.Join(parts, " ") +} + +// splitComponentInventory separates managed components (commands, agents, +// skills, hooks) from declared-but-unmanaged components (mcpServers, +// lspServers). MCP/LSP servers are declared in the plugin manifest but are +// not materialized or lifecycle-managed by ToolHive. +func splitComponentInventory(inv plugins.ComponentInventory) (managed, declared plugins.ComponentInventory) { + managed = plugins.ComponentInventory{} + declared = plugins.ComponentInventory{} + for k, v := range inv { + switch k { + case string(plugins.ComponentMCP), string(plugins.ComponentLSP): + declared[k] = v + default: + managed[k] = v + } + } + return managed, declared +} diff --git a/cmd/thv/app/ai_plugin_install.go b/cmd/thv/app/ai_plugin_install.go new file mode 100644 index 0000000000..3d4996b428 --- /dev/null +++ b/cmd/thv/app/ai_plugin_install.go @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var ( + aiPluginInstallScope string + aiPluginInstallClientsRaw string + aiPluginInstallForce bool + aiPluginInstallProjectRoot string + aiPluginInstallGroup string +) + +var aiPluginInstallCmd = &cobra.Command{ + Use: "install [plugin-name]", + Short: "Install an AI-tool plugin", + Long: `Install a plugin by name or OCI reference. +The plugin will be fetched from a remote registry and installed locally.`, + Args: cobra.ExactArgs(1), + PreRunE: chainPreRunE( + validateAIPluginScope(&aiPluginInstallScope), + validateProjectRootForScope(&aiPluginInstallScope, &aiPluginInstallProjectRoot), + validateGroupFlag(), + ), + RunE: aiPluginInstallCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginInstallCmd) + + aiPluginInstallCmd.Flags().StringVar(&aiPluginInstallClientsRaw, "clients", "", + `Comma-separated target client apps (e.g. claude-code,codex), or "all" for every available client`) + aiPluginInstallCmd.Flags().StringVar( + &aiPluginInstallScope, "scope", string(plugins.ScopeUser), "Installation scope (user, project)", + ) + aiPluginInstallCmd.Flags().BoolVar(&aiPluginInstallForce, "force", false, "Overwrite existing plugin directory") + aiPluginInstallCmd.Flags().StringVar( + &aiPluginInstallProjectRoot, "project-root", "", "Project root path for project-scoped installs", + ) + aiPluginInstallCmd.Flags().StringVar(&aiPluginInstallGroup, "group", "", "Group to add the plugin to after installation") +} + +func aiPluginInstallCmdFunc(cmd *cobra.Command, args []string) error { + c := newAIPluginClient(cmd.Context()) + + projectRoot, err := absProjectRoot(aiPluginInstallProjectRoot) + if err != nil { + return err + } + + _, err = c.Install(cmd.Context(), plugins.InstallOptions{ + Name: args[0], + Scope: plugins.Scope(aiPluginInstallScope), + Clients: parseSkillInstallClients(aiPluginInstallClientsRaw), + Force: aiPluginInstallForce, + ProjectRoot: projectRoot, + Group: aiPluginInstallGroup, + }) + if err != nil { + return formatAIPluginError("install plugin", err) + } + + return nil +} diff --git a/cmd/thv/app/ai_plugin_list.go b/cmd/thv/app/ai_plugin_list.go new file mode 100644 index 0000000000..62e45846cb --- /dev/null +++ b/cmd/thv/app/ai_plugin_list.go @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "encoding/json" + "fmt" + "os" + "strings" + "text/tabwriter" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var ( + aiPluginListScope string + aiPluginListClient string + aiPluginListFormat string + aiPluginListProjectRoot string + aiPluginListGroup string +) + +var aiPluginListCmd = &cobra.Command{ + Use: "list", + Aliases: []string{"ls"}, + Short: "List installed AI-tool plugins", + Long: `List all currently installed plugins and their status.`, + PreRunE: chainPreRunE( + validateAIPluginScope(&aiPluginListScope), + ValidateFormat(&aiPluginListFormat), + validateGroupFlag(), + ), + RunE: aiPluginListCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginListCmd) + + aiPluginListCmd.Flags().StringVar(&aiPluginListScope, "scope", "", "Filter by scope (user, project)") + aiPluginListCmd.Flags().StringVar(&aiPluginListClient, "client", "", "Filter by client application") + AddFormatFlag(aiPluginListCmd, &aiPluginListFormat) + AddGroupFlag(aiPluginListCmd, &aiPluginListGroup, false) + aiPluginListCmd.Flags().StringVar(&aiPluginListProjectRoot, "project-root", "", "Project root path for project-scoped plugins") +} + +func aiPluginListCmdFunc(cmd *cobra.Command, _ []string) error { + c := newAIPluginClient(cmd.Context()) + + projectRoot, err := absProjectRoot(aiPluginListProjectRoot) + if err != nil { + return err + } + + installed, err := c.List(cmd.Context(), plugins.ListOptions{ + Scope: plugins.Scope(aiPluginListScope), + ClientApp: aiPluginListClient, + ProjectRoot: projectRoot, + Group: aiPluginListGroup, + }) + if err != nil { + return formatAIPluginError("list plugins", err) + } + + switch aiPluginListFormat { + case FormatJSON: + if installed == nil { + installed = []plugins.InstalledPlugin{} + } + data, err := json.MarshalIndent(installed, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal JSON: %w", err) + } + fmt.Println(string(data)) + default: + if len(installed) == 0 { + if aiPluginListScope != "" || aiPluginListClient != "" { + fmt.Println("No plugins found matching filters") + } else { + fmt.Println("No plugins installed") + } + return nil + } + printAIPluginListText(installed) + } + + return nil +} + +func printAIPluginListText(installed []plugins.InstalledPlugin) { + w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', 0) + _, _ = fmt.Fprintln(w, "NAME\tVERSION\tSCOPE\tSTATUS\tCLIENTS\tREFERENCE") + + for _, p := range installed { + clients := strings.Join(p.Clients, ", ") + _, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", + p.Metadata.Name, + p.Metadata.Version, + p.Scope, + p.Status, + clients, + p.Reference, + ) + } + + _ = w.Flush() +} diff --git a/cmd/thv/app/ai_plugin_push.go b/cmd/thv/app/ai_plugin_push.go new file mode 100644 index 0000000000..0b6bb00598 --- /dev/null +++ b/cmd/thv/app/ai_plugin_push.go @@ -0,0 +1,35 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var aiPluginPushCmd = &cobra.Command{ + Use: "push [reference]", + Short: "Push a built AI-tool plugin to an OCI registry", + Long: `Push a previously built plugin artifact to a remote OCI registry.`, + Args: cobra.ExactArgs(1), + RunE: aiPluginPushCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginPushCmd) +} + +func aiPluginPushCmdFunc(cmd *cobra.Command, args []string) error { + c := newAIPluginClient(cmd.Context()) + + err := c.Push(cmd.Context(), plugins.PushOptions{ + Reference: args[0], + }) + if err != nil { + return formatAIPluginError("push plugin", err) + } + + return nil +} diff --git a/cmd/thv/app/ai_plugin_uninstall.go b/cmd/thv/app/ai_plugin_uninstall.go new file mode 100644 index 0000000000..b641e44ea0 --- /dev/null +++ b/cmd/thv/app/ai_plugin_uninstall.go @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/plugins" +) + +var ( + aiPluginUninstallScope string + aiPluginUninstallProjectRoot string +) + +var aiPluginUninstallCmd = &cobra.Command{ + Use: "uninstall [plugin-name]", + Short: "Uninstall an AI-tool plugin", + Long: `Remove a previously installed plugin by name.`, + Args: cobra.ExactArgs(1), + ValidArgsFunction: completeAIPluginNames, + PreRunE: chainPreRunE( + validateAIPluginScope(&aiPluginUninstallScope), + validateProjectRootForScope(&aiPluginUninstallScope, &aiPluginUninstallProjectRoot), + ), + RunE: aiPluginUninstallCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginUninstallCmd) + + aiPluginUninstallCmd.Flags().StringVar( + &aiPluginUninstallScope, "scope", string(plugins.ScopeUser), "Scope to uninstall from (user, project)", + ) + aiPluginUninstallCmd.Flags().StringVar( + &aiPluginUninstallProjectRoot, "project-root", "", "Project root path for project-scoped plugins", + ) +} + +func aiPluginUninstallCmdFunc(cmd *cobra.Command, args []string) error { + c := newAIPluginClient(cmd.Context()) + + projectRoot, err := absProjectRoot(aiPluginUninstallProjectRoot) + if err != nil { + return err + } + + err = c.Uninstall(cmd.Context(), plugins.UninstallOptions{ + Name: args[0], + Scope: plugins.Scope(aiPluginUninstallScope), + ProjectRoot: projectRoot, + }) + if err != nil { + return formatAIPluginError("uninstall plugin", err) + } + + return nil +} diff --git a/cmd/thv/app/ai_plugin_validate.go b/cmd/thv/app/ai_plugin_validate.go new file mode 100644 index 0000000000..8009486241 --- /dev/null +++ b/cmd/thv/app/ai_plugin_validate.go @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "encoding/json" + "fmt" + "path/filepath" + + "github.com/spf13/cobra" +) + +var aiPluginValidateFormat string + +var aiPluginValidateCmd = &cobra.Command{ + Use: "validate [path]", + Short: "Validate an AI-tool plugin directory", + Long: `Check that a plugin definition in the given directory is valid and well-formed.`, + Args: cobra.ExactArgs(1), + ValidArgsFunction: func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + return nil, cobra.ShellCompDirectiveFilterDirs + }, + PreRunE: ValidateFormat(&aiPluginValidateFormat), + RunE: aiPluginValidateCmdFunc, +} + +func init() { + aiPluginCmd.AddCommand(aiPluginValidateCmd) + + AddFormatFlag(aiPluginValidateCmd, &aiPluginValidateFormat) +} + +func aiPluginValidateCmdFunc(cmd *cobra.Command, args []string) error { + absPath, err := filepath.Abs(args[0]) + if err != nil { + return fmt.Errorf("failed to resolve path: %w", err) + } + + c := newAIPluginClient(cmd.Context()) + + result, err := c.Validate(cmd.Context(), absPath) + if err != nil { + return formatAIPluginError("validate plugin", err) + } + + switch aiPluginValidateFormat { + case FormatJSON: + data, err := json.MarshalIndent(result, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal JSON: %w", err) + } + fmt.Println(string(data)) + default: + for _, e := range result.Errors { + fmt.Printf("Error: %s\n", e) + } + for _, w := range result.Warnings { + fmt.Printf("Warning: %s\n", w) + } + } + + if !result.Valid { + return fmt.Errorf("plugin validation failed") + } + + return nil +} diff --git a/cmd/thv/app/build.go b/cmd/thv/app/build.go index e271a43dec..f2828792c2 100644 --- a/cmd/thv/app/build.go +++ b/cmd/thv/app/build.go @@ -11,6 +11,7 @@ import ( "github.com/spf13/cobra" "github.com/stacklok/toolhive/pkg/container/images" + "github.com/stacklok/toolhive/pkg/container/templates" "github.com/stacklok/toolhive/pkg/runner" ) @@ -59,9 +60,10 @@ var buildFlags BuildFlags // BuildFlags holds the configuration for building MCP server containers type BuildFlags struct { - Tag string - Output string - DryRun bool + Tag string + Output string + DryRun bool + BuildWith []string } func init() { @@ -77,6 +79,10 @@ func AddBuildFlags(cmd *cobra.Command, config *BuildFlags) { "(default builds an image instead of generating a Dockerfile)") cmd.Flags().BoolVar(&config.DryRun, "dry-run", false, "Generate Dockerfile without building (stdout output unless -o is set) "+ "(default false)") + cmd.Flags().StringArrayVar(&config.BuildWith, "build-with", []string{}, + "Build-time dependency constraint for protocol scheme builds, interpreted per package ecosystem "+ + "(uvx://: PEP 508 specifier passed to 'uv tool install --with', e.g. --build-with 'mcp<2'); "+ + "errors on ecosystems without constraint support (can be specified multiple times)") } func buildCmdFunc(cmd *cobra.Command, args []string) error { @@ -92,13 +98,22 @@ func buildCmdFunc(cmd *cobra.Command, args []string) error { buildArgs := parseCommandArguments(os.Args) slog.Debug(fmt.Sprintf("Build args: %v", buildArgs)) // #nosec G706 -- buildArgs are CLI arguments we control + // Build runtime config override from flags (if any) and validate it early. + var runtimeOverride *templates.RuntimeConfig + if len(buildFlags.BuildWith) > 0 { + runtimeOverride = &templates.RuntimeConfig{BuildWith: buildFlags.BuildWith} + if err := runtimeOverride.Validate(); err != nil { + return fmt.Errorf("invalid runtime configuration: %w", err) + } + } + // Create image manager (even for dry-run, we pass it but it won't be used) imageManager := images.NewImageManager(ctx) // If dry-run or output is specified, just generate the Dockerfile if buildFlags.DryRun || buildFlags.Output != "" { dockerfileContent, err := runner.BuildFromProtocolSchemeWithName( - ctx, imageManager, protocolScheme, "", buildFlags.Tag, buildArgs, nil, true) + ctx, imageManager, protocolScheme, "", buildFlags.Tag, buildArgs, runtimeOverride, true) if err != nil { return fmt.Errorf("failed to generate Dockerfile for %s: %w", protocolScheme, err) } @@ -121,7 +136,7 @@ func buildCmdFunc(cmd *cobra.Command, args []string) error { // Build the image using the new protocol handler with custom name imageName, err := runner.BuildFromProtocolSchemeWithName( - ctx, imageManager, protocolScheme, "", buildFlags.Tag, buildArgs, nil, false) + ctx, imageManager, protocolScheme, "", buildFlags.Tag, buildArgs, runtimeOverride, false) if err != nil { return fmt.Errorf("failed to build container for %s: %w", protocolScheme, err) } diff --git a/cmd/thv/app/commands.go b/cmd/thv/app/commands.go index b9b61fb8af..05ecd3593e 100644 --- a/cmd/thv/app/commands.go +++ b/cmd/thv/app/commands.go @@ -74,6 +74,7 @@ func NewRootCmd(enableUpdates bool) *cobra.Command { rootCmd.AddCommand(newLLMCommand()) rootCmd.AddCommand(groupCmd) rootCmd.AddCommand(skillCmd) + rootCmd.AddCommand(aiPluginCmd) rootCmd.AddCommand(statusCmd) rootCmd.AddCommand(tuiCmd) rootCmd.AddCommand(upgradeCmd) @@ -118,6 +119,7 @@ func IsInformationalCommand(args []string) bool { "mcp": true, "secret": true, "skill": true, + "ai-plugin": true, "vmcp": true, "llm": true, } diff --git a/cmd/thv/app/exitcode.go b/cmd/thv/app/exitcode.go new file mode 100644 index 0000000000..a8bfbbb236 --- /dev/null +++ b/cmd/thv/app/exitcode.go @@ -0,0 +1,61 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import "errors" + +// Exit codes for thv skill sync/upgrade, per RFC THV-0080's CI/scripting +// contract. 0 (success) and 1 (generic/unclassified error) are Go and +// cobra's own defaults and are not represented here. +const ( + // ExitCodeCheckFailure means sync --check or upgrade --fail-on-changes + // found drift/available changes: the project does not match its lock + // file, but nothing was installed, written, or removed. + ExitCodeCheckFailure = 2 + // ExitCodePartialFailure means some, but not all, of the targeted + // skills failed during sync or upgrade; check the reported outcomes for + // which ones. + ExitCodePartialFailure = 3 + // ExitCodePolicyRejection means the operation was refused by policy + // rather than attempted and failed: a non-interactive sync/upgrade + // declined the pre-install confirmation gate without --yes, or the + // ref-change guard blocked without --allow-ref-change. More policy + // gates (e.g. a signer-change guard) may map here in a future stack. + ExitCodePolicyRejection = 4 +) + +// exitCodeError pairs an error with the process exit code main() should use +// for it, so business logic in a RunE can request a specific exit code +// without main() needing to know the semantics of every command's errors. +type exitCodeError struct { + err error + code int +} + +// withExitCode wraps err so ExitCodeFromError reports code for it. Returns +// nil if err is nil, so callers can write `return withExitCode(err, ...)` +// unconditionally after an operation that may or may not have failed. +func withExitCode(err error, code int) error { + if err == nil { + return nil + } + return &exitCodeError{err: err, code: code} +} + +func (e *exitCodeError) Error() string { return e.err.Error() } +func (e *exitCodeError) Unwrap() error { return e.err } + +// ExitCodeFromError returns the process exit code for err: 0 for nil, the +// code carried by an exitCodeError (see withExitCode) if err wraps one, +// otherwise 1 — the generic failure code cobra callers already expect. +func ExitCodeFromError(err error) int { + if err == nil { + return 0 + } + var ece *exitCodeError + if errors.As(err, &ece) { + return ece.code + } + return 1 +} diff --git a/cmd/thv/app/exitcode_test.go b/cmd/thv/app/exitcode_test.go new file mode 100644 index 0000000000..71bb076f51 --- /dev/null +++ b/cmd/thv/app/exitcode_test.go @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestExitCodeFromError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want int + }{ + {name: "nil error", err: nil, want: 0}, + {name: "generic error", err: errors.New("boom"), want: 1}, + {name: "check failure", err: withExitCode(errors.New("drift"), ExitCodeCheckFailure), want: 2}, + {name: "partial failure", err: withExitCode(errors.New("partial"), ExitCodePartialFailure), want: 3}, + {name: "policy rejection", err: withExitCode(errors.New("refused"), ExitCodePolicyRejection), want: 4}, + { + name: "wrapped exit code error is still detected", + err: fmt.Errorf("context: %w", withExitCode(errors.New("drift"), ExitCodeCheckFailure)), + want: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, ExitCodeFromError(tt.err)) + }) + } +} + +func TestWithExitCodeNilIsNil(t *testing.T) { + t.Parallel() + assert.NoError(t, withExitCode(nil, ExitCodeCheckFailure)) +} + +func TestExitCodeErrorUnwraps(t *testing.T) { + t.Parallel() + inner := errors.New("boom") + err := withExitCode(inner, ExitCodePartialFailure) + assert.ErrorIs(t, err, inner) + assert.Equal(t, inner.Error(), err.Error()) +} diff --git a/cmd/thv/app/mcp.go b/cmd/thv/app/mcp.go index 876938f413..2606118ceb 100644 --- a/cmd/thv/app/mcp.go +++ b/cmd/thv/app/mcp.go @@ -13,9 +13,9 @@ import ( "text/tabwriter" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/spf13/cobra" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thclient "github.com/stacklok/toolhive/pkg/mcp/client" "github.com/stacklok/toolhive/pkg/workloads" ) diff --git a/cmd/thv/app/mcp_call.go b/cmd/thv/app/mcp_call.go index 2932124844..78deceb96b 100644 --- a/cmd/thv/app/mcp_call.go +++ b/cmd/thv/app/mcp_call.go @@ -12,9 +12,9 @@ import ( "os" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/spf13/cobra" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thclient "github.com/stacklok/toolhive/pkg/mcp/client" ) diff --git a/cmd/thv/app/mcp_call_test.go b/cmd/thv/app/mcp_call_test.go index 489db112dc..8f57e84454 100644 --- a/cmd/thv/app/mcp_call_test.go +++ b/cmd/thv/app/mcp_call_test.go @@ -9,9 +9,10 @@ import ( "strings" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestReadToolArgs(t *testing.T) { diff --git a/cmd/thv/app/project_root.go b/cmd/thv/app/project_root.go new file mode 100644 index 0000000000..7112fd262c --- /dev/null +++ b/cmd/thv/app/project_root.go @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "fmt" + "path/filepath" +) + +// absProjectRoot makes a user-supplied --project-root absolute, leaving an +// empty value empty — for most commands that means "not project-scoped", and +// substituting the working directory there would silently change the scope. +// +// The API server requires an absolute project root, so a relative value would +// otherwise travel all the way to the server and come back as a validation +// error naming the wire field rather than the flag the user typed. +// +// Shared by the skill and ai-plugin commands: both take a --project-root flag +// and both validate it through skills.ValidateProjectRoot server-side (the +// plugins package re-exports it), so they need the same normalization. +func absProjectRoot(explicit string) (string, error) { + if explicit == "" { + return "", nil + } + abs, err := filepath.Abs(explicit) + if err != nil { + return "", fmt.Errorf("resolving --project-root %q: %w", explicit, err) + } + return abs, nil +} diff --git a/cmd/thv/app/project_root_test.go b/cmd/thv/app/project_root_test.go new file mode 100644 index 0000000000..d058128569 --- /dev/null +++ b/cmd/thv/app/project_root_test.go @@ -0,0 +1,80 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +//nolint:paralleltest // uses t.Chdir, incompatible with t.Parallel +func TestAbsProjectRoot(t *testing.T) { + cwd := t.TempDir() + resolvedCwd, err := filepath.EvalSymlinks(cwd) + require.NoError(t, err) + t.Chdir(resolvedCwd) + + tests := []struct { + name string + explicit string + want string + }{ + { + name: "empty stays empty so the scope is not silently changed", + explicit: "", + want: "", + }, + { + name: "dot resolves to the working directory", + explicit: ".", + want: resolvedCwd, + }, + { + name: "relative child resolves against the working directory", + explicit: "child", + want: filepath.Join(resolvedCwd, "child"), + }, + { + name: "an absolute path is returned unchanged", + explicit: resolvedCwd, + want: resolvedCwd, + }, + { + name: "traversal is cleaned rather than rejected", + explicit: filepath.Join(resolvedCwd, "a", ".."), + want: resolvedCwd, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := absProjectRoot(tc.explicit) + require.NoError(t, err) + assert.Equal(t, tc.want, got) + }) + } +} + +// TestAbsProjectRootSharedBySkillAndPlugin documents why this helper is not +// in skill_helpers.go: thv ai-plugin takes the same --project-root flag and +// its server-side validation is literally the same function (pkg/plugins +// re-exports skills.ValidateProjectRoot), so both surfaces need identical +// normalization. A relative value from either must not reach the API. +// +//nolint:paralleltest // uses t.Chdir, incompatible with t.Parallel +func TestAbsProjectRootSharedBySkillAndPlugin(t *testing.T) { + dir := t.TempDir() + resolved, err := filepath.EvalSymlinks(dir) + require.NoError(t, err) + t.Chdir(resolved) + + for _, flagValue := range []string{".", "./", "sub/dir"} { + got, err := absProjectRoot(flagValue) + require.NoError(t, err) + assert.True(t, filepath.IsAbs(got), + "--project-root %q must be absolute before it reaches the API server", flagValue) + } +} diff --git a/cmd/thv/app/proxy.go b/cmd/thv/app/proxy.go index 46e5df577a..7a55fa0e01 100644 --- a/cmd/thv/app/proxy.go +++ b/cmd/thv/app/proxy.go @@ -391,15 +391,16 @@ func handleOutgoingAuthentication(ctx context.Context) (*discovery.OAuthFlowResu } flowConfig := &discovery.OAuthFlowConfig{ - ClientID: remoteAuthFlags.RemoteAuthClientID, - ClientSecret: clientSecret, - AuthorizeURL: remoteAuthFlags.RemoteAuthAuthorizeURL, - TokenURL: remoteAuthFlags.RemoteAuthTokenURL, - Scopes: remoteAuthFlags.RemoteAuthScopes, - CallbackPort: remoteAuthFlags.RemoteAuthCallbackPort, - Timeout: remoteAuthFlags.RemoteAuthTimeout, - SkipBrowser: remoteAuthFlags.RemoteAuthSkipBrowser, - ScopeParamName: remoteAuthFlags.RemoteAuthScopeParamName, + ClientID: remoteAuthFlags.RemoteAuthClientID, + ClientSecret: clientSecret, + AuthorizeURL: remoteAuthFlags.RemoteAuthAuthorizeURL, + TokenURL: remoteAuthFlags.RemoteAuthTokenURL, + Scopes: remoteAuthFlags.RemoteAuthScopes, + CallbackPort: remoteAuthFlags.RemoteAuthCallbackPort, + Timeout: remoteAuthFlags.RemoteAuthTimeout, + SkipBrowser: remoteAuthFlags.RemoteAuthSkipBrowser, + ScopeParamName: remoteAuthFlags.RemoteAuthScopeParamName, + AllowPrivateIPs: networking.TargetIsPrivate(ctx, proxyTargetURI), } result, err := discovery.PerformOAuthFlow(ctx, remoteAuthFlags.RemoteAuthIssuer, flowConfig) @@ -422,15 +423,16 @@ func handleOutgoingAuthentication(ctx context.Context) (*discovery.OAuthFlowResu // Perform OAuth flow with discovered configuration flowConfig := &discovery.OAuthFlowConfig{ - ClientID: remoteAuthFlags.RemoteAuthClientID, - ClientSecret: clientSecret, - AuthorizeURL: remoteAuthFlags.RemoteAuthAuthorizeURL, - TokenURL: remoteAuthFlags.RemoteAuthTokenURL, - Scopes: remoteAuthFlags.RemoteAuthScopes, - CallbackPort: remoteAuthFlags.RemoteAuthCallbackPort, - Timeout: remoteAuthFlags.RemoteAuthTimeout, - SkipBrowser: remoteAuthFlags.RemoteAuthSkipBrowser, - ScopeParamName: remoteAuthFlags.RemoteAuthScopeParamName, + ClientID: remoteAuthFlags.RemoteAuthClientID, + ClientSecret: clientSecret, + AuthorizeURL: remoteAuthFlags.RemoteAuthAuthorizeURL, + TokenURL: remoteAuthFlags.RemoteAuthTokenURL, + Scopes: remoteAuthFlags.RemoteAuthScopes, + CallbackPort: remoteAuthFlags.RemoteAuthCallbackPort, + Timeout: remoteAuthFlags.RemoteAuthTimeout, + SkipBrowser: remoteAuthFlags.RemoteAuthSkipBrowser, + ScopeParamName: remoteAuthFlags.RemoteAuthScopeParamName, + AllowPrivateIPs: networking.TargetIsPrivate(ctx, proxyTargetURI), } result, err := discovery.PerformOAuthFlow(ctx, authInfo.Realm, flowConfig) diff --git a/cmd/thv/app/run_flags.go b/cmd/thv/app/run_flags.go index e7f0eb3ff6..eb64ed1700 100644 --- a/cmd/thv/app/run_flags.go +++ b/cmd/thv/app/run_flags.go @@ -102,6 +102,10 @@ type RunFlags struct { // Proxy headers TrustProxyHeaders bool + // StrictProtocolValidation enables strict MCP-Protocol-Version validation + // on the streamable HTTP proxy. + StrictProtocolValidation bool + // Endpoint prefix for SSE endpoint URLs EndpointPrefix string @@ -150,6 +154,7 @@ type RunFlags struct { // Runtime configuration RuntimeImage string RuntimeAddPackages []string + BuildWith []string // WebhookConfigs is a list of paths to webhook configuration files. // Each file may define validating and/or mutating webhooks. @@ -218,6 +223,10 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) { "Override the default base image for protocol schemes (e.g., golang:1.24-alpine, node:20-alpine, python:3.11-slim)") cmd.Flags().StringArrayVar(&config.RuntimeAddPackages, "runtime-add-package", []string{}, "Add additional packages to install in the builder and runtime stages (can be repeated)") + cmd.Flags().StringArrayVar(&config.BuildWith, "build-with", []string{}, + "Build-time dependency constraint for protocol scheme builds, interpreted per package ecosystem "+ + "(uvx://: PEP 508 specifier passed to 'uv tool install --with', e.g. --build-with 'mcp<2'); "+ + "errors on ecosystems without constraint support (can be specified multiple times)") cmd.Flags().StringVar(&config.VerifyImage, "image-verification", retriever.VerifyImageWarn, fmt.Sprintf("Set image verification mode (%s, %s, %s)", retriever.VerifyImageWarn, retriever.VerifyImageEnabled, retriever.VerifyImageDisabled)) @@ -268,7 +277,8 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) { "Emit legacy attribute names alongside new OTEL semantic convention names (default true)") cmd.Flags().BoolVar(&config.IsolateNetwork, "isolate-network", true, - "Isolate the container network from the host. Use --isolate-network=false to opt out.") + "Isolate the container network from the host. Use --isolate-network=false to opt out. "+ + "Not enforced with --network host or --network none (isolation requires bridge networking).") cmd.Flags().BoolVar(&config.AllowDockerGateway, "allow-docker-gateway", false, "Allow outbound connections to Docker gateway addresses (host.docker.internal, gateway.docker.internal, 172.17.0.1). "+ "Only applies when --isolate-network is set. These are blocked by default even when insecure_allow_all is enabled. "+ @@ -277,6 +287,9 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) { cmd.Flags().BoolVar(&config.TrustProxyHeaders, "trust-proxy-headers", false, "Trust X-Forwarded-* headers from reverse proxies (X-Forwarded-Proto, X-Forwarded-Host, X-Forwarded-Port, X-Forwarded-Prefix) "+ "(default false)") + cmd.Flags().BoolVar(&config.StrictProtocolValidation, "strict-protocol-validation", false, + "Reject client requests whose MCP-Protocol-Version header is an unknown/unsupported MCP revision with HTTP 400 "+ + "(streamable-HTTP proxy only; an absent header is accepted). Off by default: any version is accepted.") cmd.Flags().BoolVar(&config.Stateless, "stateless", false, "Declare the server as stateless (POST-only, no SSE). "+ "Use for MCP servers implementing streamable-HTTP stateless mode.") @@ -285,7 +298,8 @@ func AddRunFlags(cmd *cobra.Command, config *RunFlags) { cmd.Flags().StringVar(&config.EndpointPrefix, "endpoint-prefix", "", "Path prefix to prepend to SSE endpoint URLs (e.g., /playwright)") cmd.Flags().StringVar(&config.Network, "network", "", - "Connect the container to a network (e.g., 'host' for host networking)") + "Connect the container to a network (e.g., 'host' for host networking). "+ + "Note: 'host' and 'none' cannot enforce network isolation, so isolation is dropped for those modes.") cmd.Flags().StringArrayVarP(&config.Labels, "label", "l", []string{}, "Set labels on the container (format: key=value)") cmd.Flags().BoolVarP(&config.Foreground, "foreground", "f", false, "Run in foreground mode (block until container exits) "+ "(default false)") @@ -361,10 +375,15 @@ func BuildRunnerConfig( return nil, err } + // Whether --isolate-network was explicitly passed (vs. defaulted). This + // controls whether an incompatible network mode fails fast or degrades. + isolateExplicit := cmd.Flags().Changed("isolate-network") + if runFlags.RemoteURL != "" { slog.Debug(fmt.Sprintf("Attempting to run remote MCP server: %s", runFlags.RemoteURL)) return buildRunnerConfig(ctx, runFlags, cmdArgs, debugMode, validatedHost, rt, runFlags.RemoteURL, nil, - nil, envVarValidator, oidcConfig, telemetryConfig, appConfig) + nil, envVarValidator, oidcConfig, telemetryConfig, appConfig, + runner.WithNetworkIsolationExplicit(isolateExplicit)) } // Resolve image from registry without pulling (fast registry lookup only). @@ -392,7 +411,8 @@ func BuildRunnerConfig( runConfig, err := buildRunnerConfig(ctx, runFlags, cmdArgs, debugMode, validatedHost, rt, imageURL, serverMetadata, envVars, envVarValidator, oidcConfig, telemetryConfig, appConfig, runner.WithRegistrySourceURLs(regAPIURL, regURL), - runner.WithRegistryServerName(regServerName)) + runner.WithRegistryServerName(regServerName), + runner.WithNetworkIsolationExplicit(isolateExplicit)) if err != nil { return nil, err } @@ -480,10 +500,11 @@ func handleImageResolution( // Validation here is intentionally duplicated with configureRuntimeOptions // so that invalid input is caught early before registry lookups. var runtimeOverride *templates.RuntimeConfig - if runFlags.RuntimeImage != "" || len(runFlags.RuntimeAddPackages) > 0 { + if runFlags.RuntimeImage != "" || len(runFlags.RuntimeAddPackages) > 0 || len(runFlags.BuildWith) > 0 { runtimeOverride = &templates.RuntimeConfig{ BuilderImage: runFlags.RuntimeImage, AdditionalPackages: runFlags.RuntimeAddPackages, + BuildWith: runFlags.BuildWith, } if err := runtimeOverride.Validate(); err != nil { return "", nil, fmt.Errorf("invalid runtime configuration: %w", err) @@ -615,13 +636,14 @@ func configureRemoteHeaderOptions(runFlags *RunFlags) ([]runner.RunConfigBuilder // It validates the configuration to prevent shell injection when values // are interpolated into Dockerfile templates. func configureRuntimeOptions(runFlags *RunFlags) ([]runner.RunConfigBuilderOption, error) { - if runFlags.RuntimeImage == "" && len(runFlags.RuntimeAddPackages) == 0 { + if runFlags.RuntimeImage == "" && len(runFlags.RuntimeAddPackages) == 0 && len(runFlags.BuildWith) == 0 { return nil, nil } runtimeConfig := &templates.RuntimeConfig{ BuilderImage: runFlags.RuntimeImage, AdditionalPackages: runFlags.RuntimeAddPackages, + BuildWith: runFlags.BuildWith, } if err := runtimeConfig.Validate(); err != nil { return nil, fmt.Errorf("invalid runtime configuration: %w", err) @@ -682,6 +704,7 @@ func buildRunnerConfig( runner.WithNetworkIsolation(runFlags.IsolateNetwork), runner.WithAllowDockerGateway(runFlags.AllowDockerGateway), runner.WithTrustProxyHeaders(runFlags.TrustProxyHeaders), + runner.WithStrictProtocolValidation(runFlags.StrictProtocolValidation), runner.WithStateless(runFlags.Stateless), runner.WithSessionTTL(runFlags.SessionTTL), runner.WithEndpointPrefix(runFlags.EndpointPrefix), diff --git a/cmd/thv/app/run_flags_test.go b/cmd/thv/app/run_flags_test.go index d93153148d..53511c4d67 100644 --- a/cmd/thv/app/run_flags_test.go +++ b/cmd/thv/app/run_flags_test.go @@ -18,6 +18,7 @@ import ( "github.com/stacklok/toolhive-core/logging" regtypes "github.com/stacklok/toolhive-core/registry/types" "github.com/stacklok/toolhive/pkg/config" + "github.com/stacklok/toolhive/pkg/runner" "github.com/stacklok/toolhive/pkg/webhook" ) @@ -809,3 +810,63 @@ validating: assert.Contains(t, err.Error(), "invalid webhook configuration") }) } + +// TestBuildRunnerConfig_NetworkIsolationExplicitWiring guards against a regression +// where the cmd layer drops WithNetworkIsolationExplicit: it reproduces the exact +// cmd.Flags().Changed("isolate-network") computation from BuildRunnerConfig and +// threads the result through buildRunnerConfig (the closest unit-testable entry; +// BuildRunnerConfig itself needs a live Docker runtime and registry). It confirms +// an explicit --isolate-network with --network host fails fast, while a defaulted +// isolation with --network host degrades to IsolateNetwork=false. See #5775. +func TestBuildRunnerConfig_NetworkIsolationExplicitWiring(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + network string + setIsolate string // "" means the flag is left defaulted (not changed) + expectError bool + expectIsolation bool + }{ + {"host with explicit isolate errors", "host", "true", true, false}, + {"host with defaulted isolate degrades", "host", "", false, false}, + {"bridge with explicit isolate keeps isolation", "", "true", false, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + runFlags := &RunFlags{} + cmd := &cobra.Command{} + AddRunFlags(cmd, runFlags) + + require.NoError(t, cmd.Flags().Set("permission-profile", "none")) + require.NoError(t, cmd.Flags().Set("transport", "stdio")) + if tt.network != "" { + require.NoError(t, cmd.Flags().Set("network", tt.network)) + } + if tt.setIsolate != "" { + require.NoError(t, cmd.Flags().Set("isolate-network", tt.setIsolate)) + } + + // Reproduce the exact explicitness computation from BuildRunnerConfig. + isolateExplicit := cmd.Flags().Changed("isolate-network") + + cfg, err := buildRunnerConfig( + t.Context(), runFlags, nil, false, "127.0.0.1", nil, "test:latest", nil, + map[string]string{}, &runner.DetachedEnvVarValidator{}, nil, nil, &config.Config{}, + runner.WithNetworkIsolationExplicit(isolateExplicit), + ) + + if tt.expectError { + require.Error(t, err) + assert.Contains(t, err.Error(), "network isolation cannot be enforced") + return + } + require.NoError(t, err) + require.NotNil(t, cfg) + assert.Equal(t, tt.expectIsolation, cfg.IsolateNetwork) + }) + } +} diff --git a/cmd/thv/app/skill_confirm.go b/cmd/thv/app/skill_confirm.go new file mode 100644 index 0000000000..f760beabf4 --- /dev/null +++ b/cmd/thv/app/skill_confirm.go @@ -0,0 +1,93 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "bufio" + "fmt" + "os" + "strings" + "unicode" + + "golang.org/x/term" + + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// requireConfirmation enforces RFC THV-0080's pre-install confirmation gate +// for sync/upgrade. With yes set, it returns immediately without prompting. +// On an interactive terminal, it prompts and reports whether the user +// confirmed. In a non-interactive context, it refuses outright with a +// policy-rejection exit code rather than silently proceeding: skill content +// is a set of AI-executed instructions, so unattended execution without an +// explicit --yes is not an acceptable default the way it might be for a +// lower-stakes operation. +// +// The prompt goes to stderr — stdout is reserved for the command's result +// (e.g. --format json output must stay machine-parseable) — and everything +// echoed into it is sanitized: a directory name carrying ANSI/OSC escapes +// must not be able to repaint the very prompt acting as the human gate. +func requireConfirmation(action string, yes bool) (confirmed bool, err error) { + if yes { + return true, nil + } + if !term.IsTerminal(int(os.Stdin.Fd())) { //nolint:gosec // uintptr fits int on all supported platforms + return false, withExitCode( + fmt.Errorf("%s requires confirmation; pass --yes to run non-interactively", sanitizeTerminal(action)), + ExitCodePolicyRejection, + ) + } + + fmt.Fprintf(os.Stderr, "%s? [y/N]: ", sanitizeTerminal(action)) + reader := bufio.NewReader(os.Stdin) + response, readErr := reader.ReadString('\n') + if readErr != nil { + return false, fmt.Errorf("failed to read user input: %w", readErr) + } + response = strings.TrimSpace(strings.ToLower(response)) + return response == "y" || response == "yes", nil +} + +// printLockEntriesSummary shows what the confirmation is actually about: +// the lock entries sync/upgrade will act on (name, source, pinned digest). +// A bare "proceed? [y/N]" tells the user nothing — a lock-file diff that +// swapped a digest or resolvedReference would sail past it unseen. Printed +// to stderr alongside the prompt; best-effort (an unreadable lock file is +// reported by the command itself, not here). +func printLockEntriesSummary(projectRoot string) { + root, err := lockfile.OpenRoot(projectRoot) + if err != nil { + return + } + lf, err := lockfile.Load(root) + if err != nil || len(lf.Skills) == 0 { + return + } + fmt.Fprintf(os.Stderr, "Lock file entries for %s:\n", sanitizeTerminal(projectRoot)) + for _, e := range lf.Skills { + fmt.Fprintf(os.Stderr, " %s %s %s\n", + sanitizeTerminal(e.Name), sanitizeTerminal(e.Source), sanitizeTerminal(shortDigest(e.Digest))) + } +} + +// shortDigest truncates a pin for display: enough hex to eyeball-compare, +// not enough to dominate the line. +func shortDigest(d string) string { + const keep = 19 // "sha256:" + 12 hex, or 19 chars of a commit hash + if len(d) <= keep { + return d + } + return d[:keep] + "…" +} + +// sanitizeTerminal strips non-graphic runes (control characters, ANSI/OSC +// escape introducers) from a string echoed to the terminal, keeping spaces. +func sanitizeTerminal(s string) string { + return strings.Map(func(r rune) rune { + if r == ' ' || unicode.IsGraphic(r) { + return r + } + return -1 + }, s) +} diff --git a/cmd/thv/app/skill_confirm_test.go b/cmd/thv/app/skill_confirm_test.go new file mode 100644 index 0000000000..e2bec9ae06 --- /dev/null +++ b/cmd/thv/app/skill_confirm_test.go @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRequireConfirmationYesSkipsPrompt(t *testing.T) { + t.Parallel() + confirmed, err := requireConfirmation("do the thing", true) + require.NoError(t, err) + assert.True(t, confirmed) +} + +// TestRequireConfirmationNonInteractiveWithoutYesRefuses exercises the +// non-interactive path: in this test process, os.Stdin is not a terminal +// (it's whatever `go test` wires up), so this reaches the policy-rejection +// branch without needing to fake TTY detection. +func TestRequireConfirmationNonInteractiveWithoutYesRefuses(t *testing.T) { + t.Parallel() + confirmed, err := requireConfirmation("do the thing", false) + require.Error(t, err) + assert.False(t, confirmed) + assert.Equal(t, ExitCodePolicyRejection, ExitCodeFromError(err)) + assert.Contains(t, err.Error(), "--yes") +} diff --git a/cmd/thv/app/skill_helpers.go b/cmd/thv/app/skill_helpers.go index 06f0bdc6f4..39371a769f 100644 --- a/cmd/thv/app/skill_helpers.go +++ b/cmd/thv/app/skill_helpers.go @@ -10,6 +10,7 @@ import ( "github.com/spf13/cobra" + tclient "github.com/stacklok/toolhive/pkg/client" "github.com/stacklok/toolhive/pkg/skills" skillclient "github.com/stacklok/toolhive/pkg/skills/client" ) @@ -39,13 +40,21 @@ func completeSkillNames(cmd *cobra.Command, args []string, _ string) ([]string, return names, cobra.ShellCompDirectiveNoFileComp } -// formatSkillError wraps an error with contextual information. If the -// underlying cause is ErrServerUnreachable it appends a helpful hint. +// formatSkillError wraps an error with contextual information, appending a +// hint that matches the actual failure — a timed-out request and an absent +// server need different advice. func formatSkillError(action string, err error) error { - if errors.Is(err, skillclient.ErrServerUnreachable) { + switch { + case errors.Is(err, skillclient.ErrRequestTimeout): + return fmt.Errorf( + "failed to %s: %w\nHint: the server is running and was still working; "+ + "raise the limit with TOOLHIVE_API_TIMEOUT (e.g. TOOLHIVE_API_TIMEOUT=30m)", + action, err) + case errors.Is(err, skillclient.ErrServerUnreachable): return fmt.Errorf("failed to %s: %w\nHint: ensure 'thv serve' is running", action, err) + default: + return fmt.Errorf("failed to %s: %w", action, err) } - return fmt.Errorf("failed to %s: %w", action, err) } // validateSkillScope returns a PreRunE that validates the --scope flag. @@ -65,3 +74,18 @@ func validateProjectRootForScope(scopeVar, projectRootVar *string) func(*cobra.C return nil } } + +// resolveProjectRoot returns explicit if set, otherwise auto-detects the +// project root by walking up from the current directory looking for .git — +// used by commands (sync, upgrade) that operate on "the project you're in" +// rather than requiring --project-root on every invocation. +func resolveProjectRoot(explicit string) (string, error) { + if explicit != "" { + return absProjectRoot(explicit) + } + root, err := tclient.DetectProjectRoot("") + if err != nil { + return "", fmt.Errorf("detecting project root: %w (use --project-root to specify it explicitly)", err) + } + return root, nil +} diff --git a/cmd/thv/app/skill_helpers_test.go b/cmd/thv/app/skill_helpers_test.go new file mode 100644 index 0000000000..b701a5fe41 --- /dev/null +++ b/cmd/thv/app/skill_helpers_test.go @@ -0,0 +1,29 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestResolveProjectRootAbsolutizesExplicit covers the sync/upgrade path, +// where an explicit value short-circuits auto-detection — it must still be +// made absolute, which is the regression behind #6211. +// +//nolint:paralleltest // uses t.Chdir, incompatible with t.Parallel +func TestResolveProjectRootAbsolutizesExplicit(t *testing.T) { + dir := t.TempDir() + resolved, err := filepath.EvalSymlinks(dir) + require.NoError(t, err) + t.Chdir(resolved) + + got, err := resolveProjectRoot(".") + require.NoError(t, err) + assert.Equal(t, resolved, got) + assert.True(t, filepath.IsAbs(got), "an explicit --project-root must reach the API server absolute") +} diff --git a/cmd/thv/app/skill_info.go b/cmd/thv/app/skill_info.go index 440429085d..0daa5c2dc5 100644 --- a/cmd/thv/app/skill_info.go +++ b/cmd/thv/app/skill_info.go @@ -45,10 +45,15 @@ func init() { func skillInfoCmdFunc(cmd *cobra.Command, args []string) error { c := newSkillClient(cmd.Context()) + projectRoot, err := absProjectRoot(skillInfoProjectRoot) + if err != nil { + return err + } + info, err := c.Info(cmd.Context(), skills.InfoOptions{ Name: args[0], Scope: skills.Scope(skillInfoScope), - ProjectRoot: skillInfoProjectRoot, + ProjectRoot: projectRoot, }) if err != nil { return formatSkillError("get skill info", err) diff --git a/cmd/thv/app/skill_install.go b/cmd/thv/app/skill_install.go index 57663ec3c2..f1d4c4f253 100644 --- a/cmd/thv/app/skill_install.go +++ b/cmd/thv/app/skill_install.go @@ -12,11 +12,12 @@ import ( ) var ( - skillInstallScope string - skillInstallClientsRaw string - skillInstallForce bool - skillInstallProjectRoot string - skillInstallGroup string + skillInstallScope string + skillInstallClientsRaw string + skillInstallForce bool + skillInstallProjectRoot string + skillInstallGroup string + skillInstallAllowUnsigned bool ) var skillInstallCmd = &cobra.Command{ @@ -42,18 +43,26 @@ func init() { skillInstallCmd.Flags().BoolVar(&skillInstallForce, "force", false, "Overwrite existing skill directory") skillInstallCmd.Flags().StringVar(&skillInstallProjectRoot, "project-root", "", "Project root path for project-scoped installs") skillInstallCmd.Flags().StringVar(&skillInstallGroup, "group", "", "Group to add the skill to after installation") + skillInstallCmd.Flags().BoolVar(&skillInstallAllowUnsigned, "allow-unsigned", false, + "Allow installing a project-scoped skill without a verified signature (recorded in the lock file)") } func skillInstallCmdFunc(cmd *cobra.Command, args []string) error { c := newSkillClient(cmd.Context()) - _, err := c.Install(cmd.Context(), skills.InstallOptions{ - Name: args[0], - Scope: skills.Scope(skillInstallScope), - Clients: parseSkillInstallClients(skillInstallClientsRaw), - Force: skillInstallForce, - ProjectRoot: skillInstallProjectRoot, - Group: skillInstallGroup, + projectRoot, err := absProjectRoot(skillInstallProjectRoot) + if err != nil { + return err + } + + _, err = c.Install(cmd.Context(), skills.InstallOptions{ + Name: args[0], + Scope: skills.Scope(skillInstallScope), + Clients: parseSkillInstallClients(skillInstallClientsRaw), + Force: skillInstallForce, + ProjectRoot: projectRoot, + Group: skillInstallGroup, + AllowUnsigned: skillInstallAllowUnsigned, }) if err != nil { return formatSkillError("install skill", err) diff --git a/cmd/thv/app/skill_list.go b/cmd/thv/app/skill_list.go index 0f749cc1ba..b226e4192c 100644 --- a/cmd/thv/app/skill_list.go +++ b/cmd/thv/app/skill_list.go @@ -49,10 +49,15 @@ func init() { func skillListCmdFunc(cmd *cobra.Command, _ []string) error { c := newSkillClient(cmd.Context()) + projectRoot, err := absProjectRoot(skillListProjectRoot) + if err != nil { + return err + } + installed, err := c.List(cmd.Context(), skills.ListOptions{ Scope: skills.Scope(skillListScope), ClientApp: skillListClient, - ProjectRoot: skillListProjectRoot, + ProjectRoot: projectRoot, Group: skillListGroup, }) if err != nil { diff --git a/cmd/thv/app/skill_sync.go b/cmd/thv/app/skill_sync.go new file mode 100644 index 0000000000..1c52f9faf8 --- /dev/null +++ b/cmd/thv/app/skill_sync.go @@ -0,0 +1,174 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "encoding/json" + "fmt" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/skills" +) + +var ( + skillSyncProjectRoot string + skillSyncClientsRaw string + skillSyncCheck bool + skillSyncAdopt bool + skillSyncPrune bool + skillSyncYes bool + skillSyncAllowUnsigned bool + skillSyncFormat string +) + +var skillSyncCmd = &cobra.Command{ + Use: "sync", + Short: "Restore project skills to match the lock file (experimental)", + Long: `Restore a project's installed skills to match toolhive.lock.yaml. + +Experimental: requires TOOLHIVE_SKILLS_LOCK_ENABLED=true on the ToolHive +server while the lock file feature rolls out. + +Missing or drifted skills are reinstalled at their pinned digest. Use +--check to report drift without installing anything (suitable for CI). +Use --adopt to record lock entries for existing unmanaged installs, and +--prune to remove installs no longer present in the lock file. + +Unless --check is set, sync prompts for confirmation before installing — +skill content is a set of AI-followed instructions. Pass --yes to skip the +prompt (required in non-interactive contexts such as CI).`, + PreRunE: chainPreRunE( + ValidateFormat(&skillSyncFormat), + ), + RunE: skillSyncCmdFunc, +} + +func init() { + skillCmd.AddCommand(skillSyncCmd) + + skillSyncCmd.Flags().StringVar(&skillSyncProjectRoot, "project-root", "", + "Project root path (default: auto-detected from the current directory)") + skillSyncCmd.Flags().StringVar(&skillSyncClientsRaw, "clients", "", + `Comma-separated target client apps (e.g. claude-code,opencode), or "all" for every available client`) + skillSyncCmd.Flags().BoolVar(&skillSyncCheck, "check", false, + "Report drift without installing, writing, or removing anything") + skillSyncCmd.Flags().BoolVar(&skillSyncAdopt, "adopt", false, + "Write lock entries for existing unmanaged project-scope installs") + skillSyncCmd.Flags().BoolVar(&skillSyncPrune, "prune", false, + "Remove installs no longer present in the lock file") + skillSyncCmd.Flags().BoolVar(&skillSyncYes, "yes", false, + "Skip the confirmation prompt (required when not running interactively)") + skillSyncCmd.Flags().BoolVar(&skillSyncAllowUnsigned, "allow-unsigned", false, + "Allow adopting skills whose signature state cannot be established (recorded as unsigned)") + AddFormatFlag(skillSyncCmd, &skillSyncFormat) +} + +func skillSyncCmdFunc(cmd *cobra.Command, _ []string) error { + projectRoot, err := resolveProjectRoot(skillSyncProjectRoot) + if err != nil { + return err + } + + if !skillSyncCheck { + if !skillSyncYes { + printLockEntriesSummary(projectRoot) + } + confirmed, confirmErr := requireConfirmation("Sync skills for "+projectRoot, skillSyncYes) + if confirmErr != nil { + return confirmErr + } + if !confirmed { + fmt.Println("Sync cancelled.") + return nil + } + } + + c := newSkillClient(cmd.Context()) + result, err := c.Sync(cmd.Context(), skills.SyncOptions{ + ProjectRoot: projectRoot, + Clients: parseSkillInstallClients(skillSyncClientsRaw), + Check: skillSyncCheck, + Adopt: skillSyncAdopt, + Prune: skillSyncPrune, + AllowUnsigned: skillSyncAllowUnsigned, + }) + if err != nil { + return formatSkillError("sync skills", err) + } + + if err := printSyncResult(result, skillSyncFormat); err != nil { + return err + } + return syncExitError(result, skillSyncCheck) +} + +// syncExitError maps a SyncResult to RFC THV-0080's exit-code contract. +// Partial failure takes precedence over check-mode drift: something +// actually going wrong is a stronger signal than the project merely not +// matching its lock file. Missing counts toward the check failure exactly +// like drift — on a fresh clone every locked skill is missing, and that is +// the canonical state the CI gate exists to catch. +func syncExitError(result *skills.SyncResult, check bool) error { + if len(result.Failed) > 0 { + return withExitCode(fmt.Errorf("sync failed for %d skill(s)", len(result.Failed)), ExitCodePartialFailure) + } + if outOfSync := len(result.Drifted) + len(result.Missing); check && outOfSync > 0 { + return withExitCode( + fmt.Errorf("%d skill(s) drifted from or are missing against the lock file", outOfSync), + ExitCodeCheckFailure, + ) + } + return nil +} + +func printSyncResult(result *skills.SyncResult, format string) error { + if format == FormatJSON { + data, err := json.MarshalIndent(result, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal JSON: %w", err) + } + fmt.Println(string(data)) + return nil + } + + printSkillNameGroup("Installed", result.Installed) + printSkillNameGroup("Drifted", result.Drifted) + printSkillNameGroup("Missing (not installed)", result.Missing) + printSkillNameGroup("Up to date", result.AlreadyCurrent) + printSkillNameGroup("Never managed (use --adopt to record)", result.NeverManaged) + printSkillNameGroup("Removed from lock (use --prune to remove)", result.RemovedFromLock) + printSkillNameGroup("Pruned", result.Pruned) + if len(result.Failed) > 0 { + fmt.Println("Failed:") + for _, f := range result.Failed { + fmt.Printf(" %s [%s]: %s\n", f.Name, f.Reason, f.Error) + } + } + if isSyncResultEmpty(result) { + fmt.Println("Nothing to sync — the project matches its lock file") + } + return nil +} + +func printSkillNameGroup(label string, names []string) { + if len(names) == 0 { + return + } + fmt.Printf("%s:\n", label) + for _, name := range names { + fmt.Printf(" %s\n", name) + } +} + +func isSyncResultEmpty(result *skills.SyncResult) bool { + return len(result.Installed) == 0 && + len(result.Drifted) == 0 && + len(result.Missing) == 0 && + len(result.AlreadyCurrent) == 0 && + len(result.NeverManaged) == 0 && + len(result.RemovedFromLock) == 0 && + len(result.Pruned) == 0 && + len(result.Failed) == 0 +} diff --git a/cmd/thv/app/skill_sync_test.go b/cmd/thv/app/skill_sync_test.go new file mode 100644 index 0000000000..b600d1c013 --- /dev/null +++ b/cmd/thv/app/skill_sync_test.go @@ -0,0 +1,112 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/skills" +) + +func TestSyncExitError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + result *skills.SyncResult + check bool + wantCode int + }{ + {name: "clean sync", result: &skills.SyncResult{AlreadyCurrent: []string{"a"}}, wantCode: 0}, + { + name: "check finds drift", + result: &skills.SyncResult{Drifted: []string{"a"}}, + check: true, + wantCode: ExitCodeCheckFailure, + }, + { + // The fresh-clone CI gate: a lock entry with no install record + // at all must fail --check exactly like drift does. + name: "check finds missing installs", + result: &skills.SyncResult{Missing: []string{"a"}}, + check: true, + wantCode: ExitCodeCheckFailure, + }, + { + name: "non-check drift is not a failure (already reinstalled)", + result: &skills.SyncResult{Drifted: []string{"a"}, Installed: []string{"a"}}, + check: false, + wantCode: 0, + }, + { + name: "non-check missing is not a failure (already installed)", + result: &skills.SyncResult{Missing: []string{"a"}, Installed: []string{"a"}}, + check: false, + wantCode: 0, + }, + { + name: "any failure is a partial failure", + result: &skills.SyncResult{Failed: []skills.SyncFailure{{Name: "a", Error: "boom"}}}, + wantCode: ExitCodePartialFailure, + }, + { + name: "failure takes precedence over check drift", + result: &skills.SyncResult{ + Drifted: []string{"a"}, + Failed: []skills.SyncFailure{{Name: "b", Error: "boom"}}, + }, + check: true, + wantCode: ExitCodePartialFailure, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := syncExitError(tt.result, tt.check) + assert.Equal(t, tt.wantCode, ExitCodeFromError(err)) + }) + } +} + +func TestIsSyncResultEmpty(t *testing.T) { + t.Parallel() + + assert.True(t, isSyncResultEmpty(&skills.SyncResult{})) + assert.False(t, isSyncResultEmpty(&skills.SyncResult{Installed: []string{"a"}})) + assert.False(t, isSyncResultEmpty(&skills.SyncResult{Failed: []skills.SyncFailure{{Name: "a"}}})) +} + +func TestPrintSyncResultJSON(t *testing.T) { + t.Parallel() + err := printSyncResult(&skills.SyncResult{Installed: []string{"my-skill"}}, FormatJSON) + require.NoError(t, err) +} + +func TestPrintSyncResultText(t *testing.T) { + t.Parallel() + + t.Run("every category populated", func(t *testing.T) { + t.Parallel() + err := printSyncResult(&skills.SyncResult{ + Installed: []string{"installed-skill"}, + Drifted: []string{"drifted-skill"}, + AlreadyCurrent: []string{"current-skill"}, + NeverManaged: []string{"unmanaged-skill"}, + RemovedFromLock: []string{"removed-skill"}, + Pruned: []string{"pruned-skill"}, + Failed: []skills.SyncFailure{{Name: "failed-skill", Reason: skills.FailureReasonUnknown, Error: "boom"}}, + }, FormatText) + require.NoError(t, err) + }) + + t.Run("empty result prints nothing-to-sync", func(t *testing.T) { + t.Parallel() + err := printSyncResult(&skills.SyncResult{}, FormatText) + require.NoError(t, err) + }) +} diff --git a/cmd/thv/app/skill_uninstall.go b/cmd/thv/app/skill_uninstall.go index 9f5ff6a755..119a648f10 100644 --- a/cmd/thv/app/skill_uninstall.go +++ b/cmd/thv/app/skill_uninstall.go @@ -41,10 +41,15 @@ func init() { func skillUninstallCmdFunc(cmd *cobra.Command, args []string) error { c := newSkillClient(cmd.Context()) - err := c.Uninstall(cmd.Context(), skills.UninstallOptions{ + projectRoot, err := absProjectRoot(skillUninstallProjectRoot) + if err != nil { + return err + } + + err = c.Uninstall(cmd.Context(), skills.UninstallOptions{ Name: args[0], Scope: skills.Scope(skillUninstallScope), - ProjectRoot: skillUninstallProjectRoot, + ProjectRoot: projectRoot, }) if err != nil { return formatSkillError("uninstall skill", err) diff --git a/cmd/thv/app/skill_upgrade.go b/cmd/thv/app/skill_upgrade.go new file mode 100644 index 0000000000..8a4364719f --- /dev/null +++ b/cmd/thv/app/skill_upgrade.go @@ -0,0 +1,218 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "encoding/json" + "fmt" + + "github.com/spf13/cobra" + + "github.com/stacklok/toolhive/pkg/skills" +) + +var ( + skillUpgradeProjectRoot string + skillUpgradeClientsRaw string + skillUpgradePreview bool + skillUpgradeFailOnChanges bool + skillUpgradeAllowRefChange bool + skillUpgradeAllowSignerChange bool + skillUpgradeYes bool + skillUpgradeFormat string +) + +var skillUpgradeCmd = &cobra.Command{ + Use: "upgrade [skill-name...]", + Short: "Upgrade project skills to newer pinned content (experimental)", + Long: `Re-resolve a project's lock entries and install newer content where available. + +Experimental: requires TOOLHIVE_SKILLS_LOCK_ENABLED=true on the ToolHive +server while the lock file feature rolls out. + +Skills pinned to an immutable reference (an OCI digest or a full git commit +hash) are reported not-upgradable — there is nothing newer to resolve to. +Use --preview to see what would change without persisting anything (OCI +sources are still fetched into the local artifact store to compare digests), +and --allow-ref-change to permit the artifact moving to a different +repository (a version bump within the same repository is not a change +this guard blocks). +--fail-on-changes evaluates the same plan and never installs: it is a CI +freshness gate. + +Unless --preview is set, upgrade prompts for confirmation before installing — +skill content is a set of AI-followed instructions. Pass --yes to skip the +prompt (required in non-interactive contexts such as CI).`, + PreRunE: chainPreRunE( + ValidateFormat(&skillUpgradeFormat), + ), + RunE: skillUpgradeCmdFunc, +} + +func init() { + skillCmd.AddCommand(skillUpgradeCmd) + + skillUpgradeCmd.Flags().StringVar(&skillUpgradeProjectRoot, "project-root", "", + "Project root path (default: auto-detected from the current directory)") + skillUpgradeCmd.Flags().StringVar(&skillUpgradeClientsRaw, "clients", "", + `Comma-separated target client apps (e.g. claude-code,opencode), or "all" for every available client`) + skillUpgradeCmd.Flags().BoolVar(&skillUpgradePreview, "preview", false, + "Report what would change without persisting anything (OCI sources are still fetched to compare digests)") + skillUpgradeCmd.Flags().BoolVar(&skillUpgradeFailOnChanges, "fail-on-changes", false, + "Report what would change without installing anything; a CI freshness gate") + skillUpgradeCmd.Flags().BoolVar(&skillUpgradeAllowSignerChange, "allow-signer-change", false, + "Permit upgrading to an artifact signed by a different identity; the new identity replaces the recorded one") + skillUpgradeCmd.Flags().BoolVar(&skillUpgradeAllowRefChange, "allow-ref-change", false, + "Permit the artifact to move to a different repository during upgrade") + skillUpgradeCmd.Flags().BoolVar(&skillUpgradeYes, "yes", false, + "Skip the confirmation prompt (required when not running interactively)") + AddFormatFlag(skillUpgradeCmd, &skillUpgradeFormat) +} + +func skillUpgradeCmdFunc(cmd *cobra.Command, args []string) error { + projectRoot, err := resolveProjectRoot(skillUpgradeProjectRoot) + if err != nil { + return err + } + + if !skillUpgradePreview && !skillUpgradeFailOnChanges { + if !skillUpgradeYes { + printLockEntriesSummary(projectRoot) + } + confirmed, confirmErr := requireConfirmation("Upgrade skills for "+projectRoot, skillUpgradeYes) + if confirmErr != nil { + return confirmErr + } + if !confirmed { + fmt.Println("Upgrade cancelled.") + return nil + } + } + + c := newSkillClient(cmd.Context()) + result, err := c.Upgrade(cmd.Context(), skills.UpgradeOptions{ + ProjectRoot: projectRoot, + Names: args, + Clients: parseSkillInstallClients(skillUpgradeClientsRaw), + Preview: skillUpgradePreview, + FailOnChanges: skillUpgradeFailOnChanges, + AllowRefChange: skillUpgradeAllowRefChange, + AllowSignerChange: skillUpgradeAllowSignerChange, + }) + if err != nil { + return formatSkillError("upgrade skills", err) + } + + if err := printUpgradeResult(result, skillUpgradeFormat, skillUpgradePreview || skillUpgradeFailOnChanges); err != nil { + return err + } + return upgradeExitError(result, skillUpgradePreview, skillUpgradeFailOnChanges) +} + +// upgradeExitError maps an UpgradeResult to RFC THV-0080's exit-code +// contract, entirely from the reported outcomes. Precedence: a failed +// outcome (exit 3) beats everything — a genuine failure must never be +// masked as "lock is stale" (exit 2) or a guard doing its job (exit 4). +// With failOnChanges, any would-change outcome is the CI freshness signal +// (exit 2). A ref-change block maps to a policy rejection (exit 4) only +// when the run wasn't a preview/gate evaluation — during those nothing was +// actually blocked, only reported. +// upgradeTally counts the exit-code-relevant outcome classes. +type upgradeTally struct { + failed, refBlocked, signerBlocked, wouldChange int +} + +func tallyUpgradeOutcomes(result *skills.UpgradeResult) upgradeTally { + var t upgradeTally + for _, o := range result.Outcomes { + switch o.Status { + case skills.UpgradeStatusFailed: + t.failed++ + case skills.UpgradeStatusRefChangeBlocked: + t.refBlocked++ + t.wouldChange++ + case skills.UpgradeStatusSignerChangeBlocked: + t.signerBlocked++ + t.wouldChange++ + case skills.UpgradeStatusUpgraded: + t.wouldChange++ + case skills.UpgradeStatusUpToDate, skills.UpgradeStatusNotUpgradable: + // No exit-code impact. + } + } + return t +} + +func upgradeExitError(result *skills.UpgradeResult, preview, failOnChanges bool) error { + tally := tallyUpgradeOutcomes(result) + failed, refBlocked, signerBlocked, wouldChange := + tally.failed, tally.refBlocked, tally.signerBlocked, tally.wouldChange + if failed > 0 { + return withExitCode(fmt.Errorf("upgrade failed for %d skill(s)", failed), ExitCodePartialFailure) + } + if failOnChanges && wouldChange > 0 { + return withExitCode( + fmt.Errorf("%d skill(s) would change; the lock file is stale", wouldChange), + ExitCodeCheckFailure, + ) + } + if !preview && !failOnChanges && signerBlocked > 0 { + return withExitCode( + fmt.Errorf("%d skill(s) blocked by a signer change; use --allow-signer-change", signerBlocked), + ExitCodePolicyRejection, + ) + } + if !preview && !failOnChanges && refBlocked > 0 { + return withExitCode( + fmt.Errorf("%d skill(s) blocked by a repository change; use --allow-ref-change", refBlocked), + ExitCodePolicyRejection, + ) + } + return nil +} + +// printUpgradeResult renders the outcomes. planOnly (a --preview or +// --fail-on-changes run) switches "upgraded" to "would upgrade": nothing +// was installed in those modes, and saying otherwise misreads the gate. +func printUpgradeResult(result *skills.UpgradeResult, format string, planOnly bool) error { + if format == FormatJSON { + data, err := json.MarshalIndent(result, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal JSON: %w", err) + } + fmt.Println(string(data)) + return nil + } + + if len(result.Outcomes) == 0 { + fmt.Println("No skills in the project's lock file") + return nil + } + upgradedVerb := "upgraded" + if planOnly { + upgradedVerb = "would upgrade" + } + for _, o := range result.Outcomes { + switch o.Status { + case skills.UpgradeStatusUpgraded: + fmt.Printf("%s: %s %s -> %s\n", o.Name, upgradedVerb, o.OldDigest, o.NewDigest) + case skills.UpgradeStatusUpToDate: + fmt.Printf("%s: up to date\n", o.Name) + case skills.UpgradeStatusNotUpgradable: + fmt.Printf("%s: not upgradable (pinned to an immutable reference)\n", o.Name) + case skills.UpgradeStatusRefChangeBlocked: + fmt.Printf("%s: repository change blocked (would move to %s; use --allow-ref-change)\n", + o.Name, o.NewResolvedReference) + case skills.UpgradeStatusSignerChangeBlocked: + newSigner := o.NewSignerIdentity + if newSigner == "" { + newSigner = "unsigned" + } + fmt.Printf("%s: signer change blocked (candidate is %s; use --allow-signer-change)\n", o.Name, newSigner) + case skills.UpgradeStatusFailed: + fmt.Printf("%s: failed [%s]: %s\n", o.Name, o.Reason, o.Error) + } + } + return nil +} diff --git a/cmd/thv/app/skill_upgrade_test.go b/cmd/thv/app/skill_upgrade_test.go new file mode 100644 index 0000000000..4d7aea0a4a --- /dev/null +++ b/cmd/thv/app/skill_upgrade_test.go @@ -0,0 +1,154 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package app + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/skills" +) + +func TestUpgradeExitError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + outcomes []skills.UpgradeOutcome + preview bool + failOnChanges bool + wantCode int + }{ + { + name: "all up to date", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusUpToDate}}, + wantCode: 0, + }, + { + name: "upgraded is not a failure", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusUpgraded}}, + wantCode: 0, + }, + { + name: "not upgradable is not a failure", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusNotUpgradable}}, + wantCode: 0, + }, + { + name: "ref change blocked is a policy rejection", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusRefChangeBlocked}}, + wantCode: ExitCodePolicyRejection, + }, + { + name: "signer change blocked is a policy rejection", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusSignerChangeBlocked}}, + wantCode: ExitCodePolicyRejection, + }, + { + name: "signer change blocked in preview is informational", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusSignerChangeBlocked}}, + preview: true, + wantCode: 0, + }, + { + name: "signer change blocked counts as would-change under fail-on-changes", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusSignerChangeBlocked}}, + failOnChanges: true, + wantCode: ExitCodeCheckFailure, + }, + { + name: "ref change blocked during preview is not a policy rejection", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusRefChangeBlocked}}, + preview: true, + wantCode: 0, + }, + { + name: "failed outcome is a partial failure", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusFailed}}, + wantCode: ExitCodePartialFailure, + }, + { + name: "failure takes precedence over ref-change-blocked", + outcomes: []skills.UpgradeOutcome{ + {Name: "a", Status: skills.UpgradeStatusRefChangeBlocked}, + {Name: "b", Status: skills.UpgradeStatusFailed}, + }, + wantCode: ExitCodePartialFailure, + }, + { + name: "failure still fails even during preview", + outcomes: []skills.UpgradeOutcome{ + {Name: "a", Status: skills.UpgradeStatusFailed}, + }, + preview: true, + wantCode: ExitCodePartialFailure, + }, + { + name: "fail-on-changes with pending upgrade is a check failure", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusUpgraded}}, + failOnChanges: true, + wantCode: ExitCodeCheckFailure, + }, + { + name: "fail-on-changes with a blocked ref change is a check failure", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusRefChangeBlocked}}, + failOnChanges: true, + wantCode: ExitCodeCheckFailure, + }, + { + name: "fail-on-changes with everything current exits clean", + outcomes: []skills.UpgradeOutcome{{Name: "a", Status: skills.UpgradeStatusUpToDate}}, + failOnChanges: true, + wantCode: 0, + }, + { + // The exit-code inversion from the panel review: a genuine + // resolution failure during the CI gate must exit 3, not be + // silently downgraded to "lock is stale" (exit 2). + name: "fail-on-changes with a genuine failure is a partial failure, not a check failure", + outcomes: []skills.UpgradeOutcome{ + {Name: "a", Status: skills.UpgradeStatusUpgraded}, + {Name: "b", Status: skills.UpgradeStatusFailed}, + }, + failOnChanges: true, + wantCode: ExitCodePartialFailure, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := upgradeExitError(&skills.UpgradeResult{Outcomes: tt.outcomes}, tt.preview, tt.failOnChanges) + assert.Equal(t, tt.wantCode, ExitCodeFromError(err)) + }) + } +} + +func TestPrintUpgradeResultJSON(t *testing.T) { + t.Parallel() + err := printUpgradeResult(&skills.UpgradeResult{ + Outcomes: []skills.UpgradeOutcome{{Name: "my-skill", Status: skills.UpgradeStatusUpgraded}}, + }, FormatJSON, false) + require.NoError(t, err) +} + +func TestPrintUpgradeResultTextNoOutcomes(t *testing.T) { + t.Parallel() + err := printUpgradeResult(&skills.UpgradeResult{}, FormatText, false) + require.NoError(t, err) +} + +func TestPrintUpgradeResultTextEveryStatus(t *testing.T) { + t.Parallel() + err := printUpgradeResult(&skills.UpgradeResult{Outcomes: []skills.UpgradeOutcome{ + {Name: "upgraded-skill", Status: skills.UpgradeStatusUpgraded, OldDigest: "old", NewDigest: "new"}, + {Name: "current-skill", Status: skills.UpgradeStatusUpToDate}, + {Name: "pinned-skill", Status: skills.UpgradeStatusNotUpgradable}, + {Name: "blocked-skill", Status: skills.UpgradeStatusRefChangeBlocked, NewResolvedReference: "new-ref"}, + {Name: "failed-skill", Status: skills.UpgradeStatusFailed, Reason: skills.FailureReasonUnknown, Error: "boom"}, + }}, FormatText, true) + require.NoError(t, err) +} diff --git a/cmd/thv/main.go b/cmd/thv/main.go index fb03b8282e..0f66538d2e 100644 --- a/cmd/thv/main.go +++ b/cmd/thv/main.go @@ -81,7 +81,7 @@ func main() { if err := cmd.ExecuteContext(ctx); err != nil { // Clean up any remaining lock files on error exit lockfile.CleanupAllLocks() - os.Exit(1) + os.Exit(app.ExitCodeFromError(err)) } // Clean up lock files on normal exit diff --git a/deploy/charts/operator-crds/Chart.yaml b/deploy/charts/operator-crds/Chart.yaml index bebdfdfd09..d56868a028 100644 --- a/deploy/charts/operator-crds/Chart.yaml +++ b/deploy/charts/operator-crds/Chart.yaml @@ -2,5 +2,5 @@ apiVersion: v2 name: toolhive-operator-crds description: A Helm chart for installing the ToolHive Operator CRDs into Kubernetes. type: application -version: 0.39.0 -appVersion: "v0.39.0" +version: 0.42.1 +appVersion: "v0.42.1" diff --git a/deploy/charts/operator-crds/README.md b/deploy/charts/operator-crds/README.md index daf71c89ac..3fea250bea 100644 --- a/deploy/charts/operator-crds/README.md +++ b/deploy/charts/operator-crds/README.md @@ -1,6 +1,6 @@ # ToolHive Operator CRDs Helm Chart -![Version: 0.39.0](https://img.shields.io/badge/Version-0.39.0-informational?style=flat-square) +![Version: 0.42.1](https://img.shields.io/badge/Version-0.42.1-informational?style=flat-square) ![Type: application](https://img.shields.io/badge/Type-application-informational?style=flat-square) A Helm chart for installing the ToolHive Operator CRDs into Kubernetes. diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpauthzconfigs.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpauthzconfigs.yaml index 518eb26970..eafdf71d39 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpauthzconfigs.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpauthzconfigs.yaml @@ -27,9 +27,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -155,43 +152,6 @@ spec: for this MCPAuthzConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPAuthzConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -205,9 +165,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -334,43 +291,6 @@ spec: for this MCPAuthzConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPAuthzConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml index d47545b4d6..c9b4245cc2 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml @@ -28,9 +28,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -1579,43 +1576,6 @@ spec: It corresponds to the MCPExternalAuthConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPExternalAuthConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -1629,9 +1589,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -3181,43 +3138,6 @@ spec: It corresponds to the MCPExternalAuthConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPExternalAuthConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpoidcconfigs.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpoidcconfigs.yaml index cbd795abb9..6dae7eaf23 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpoidcconfigs.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpoidcconfigs.yaml @@ -27,9 +27,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -127,7 +124,7 @@ spec: insecureAllowHTTP: default: false description: |- - InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing. + InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing. WARNING: This is insecure and should NEVER be used in production. type: boolean introspectionUrl: @@ -288,43 +285,6 @@ spec: for this MCPOIDCConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPOIDCConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -338,9 +298,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -439,7 +396,7 @@ spec: insecureAllowHTTP: default: false description: |- - InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing. + InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing. WARNING: This is insecure and should NEVER be used in production. type: boolean introspectionUrl: @@ -600,43 +557,6 @@ spec: for this MCPOIDCConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPOIDCConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptelemetryconfigs.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptelemetryconfigs.yaml index 2262801a88..a53f7b4926 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptelemetryconfigs.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptelemetryconfigs.yaml @@ -287,33 +287,6 @@ spec: for this MCPTelemetryConfig. format: int64 type: integer - referencingWorkloads: - description: ReferencingWorkloads lists workloads that reference this - MCPTelemetryConfig - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -588,33 +561,6 @@ spec: for this MCPTelemetryConfig. format: int64 type: integer - referencingWorkloads: - description: ReferencingWorkloads lists workloads that reference this - MCPTelemetryConfig - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptoolconfigs.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptoolconfigs.yaml index bdcfc3a3ca..681bd01ae1 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptoolconfigs.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcptoolconfigs.yaml @@ -25,9 +25,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -189,39 +186,6 @@ spec: It corresponds to the MCPToolConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPToolConfig. - Each entry identifies the workload by kind and name. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -232,9 +196,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -397,39 +358,6 @@ spec: It corresponds to the MCPToolConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPToolConfig. - Each entry identifies the workload by kind and name. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpwebhookconfigs.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpwebhookconfigs.yaml index f5ea07646a..b6a802683e 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpwebhookconfigs.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_mcpwebhookconfigs.yaml @@ -24,9 +24,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referencingWorkloads - name: References - type: string - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -348,34 +345,6 @@ spec: to the current status format: int64 type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPWebhookConfig. - Each entry identifies the workload by kind and name. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml index fdcad57a2f..608e365e00 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml @@ -68,6 +68,57 @@ spec: This embeds the CompositeToolConfig from pkg/vmcp/config to share the configuration model between CLI and operator usage. properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string @@ -467,6 +518,57 @@ spec: This embeds the CompositeToolConfig from pkg/vmcp/config to share the configuration model between CLI and operator usage. properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml index 9392de7830..cc4067207f 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml @@ -1017,16 +1017,43 @@ spec: prefixFormat: default: '{workload}_' description: |- - PrefixFormat defines the prefix format for the "prefix" strategy. + PrefixFormat defines the prefix format for the "prefix" tool strategy + and for backend-prefixed prompt names (the default for every prompt; + under the "priority" strategy, backends listed in priorityOrder keep + their own prompt names). Supports placeholders: {workload}, {workload}_, {workload}. type: string priorityOrder: - description: PriorityOrder defines the workload priority - order for the "priority" strategy. + description: |- + PriorityOrder defines the workload priority order for the "priority" strategy. + Listed workloads also keep their own prompt names (unlisted workloads' + prompts stay backend-prefixed). items: type: string type: array type: object + defaultToolVisibility: + description: |- + DefaultToolVisibility controls whether a backend with NO entry in Tools has its + tools advertised to MCP clients. + - allow (default): every tool from an unlisted backend is advertised, so + adding a workload to the group exposes it without further configuration. + - deny: an unlisted backend contributes no tools, so only backends named in + Tools are advertised. Use this when the set of exposed tools must be + enumerated deliberately rather than inherited from group membership. + + A backend that DOES have a Tools entry is unaffected by this setting: the + entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like + every other visibility setting here, this controls advertising only — hidden + tools remain in the routing table for composite tools (see the type doc). + + This gates TOOLS only. An unlisted backend's resources, resource templates, + and prompts are still advertised under deny; mergeResources/mergePrompts have + no equivalent check. + enum: + - allow + - deny + type: string excludeAllTools: description: |- ExcludeAllTools hides all backend tools from MCP clients when true. @@ -1072,24 +1099,34 @@ spec: Only specified fields are overridden; others pass through from the backend. properties: destructiveHint: - description: DestructiveHint overrides the - destructive hint annotation. + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean idempotentHint: - description: IdempotentHint overrides the - idempotent hint annotation. + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean openWorldHint: - description: OpenWorldHint overrides the open-world - hint annotation. + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean readOnlyHint: - description: ReadOnlyHint overrides the read-only - hint annotation. + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean title: - description: Title overrides the human-readable - title annotation. + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: string type: object description: @@ -1184,6 +1221,14 @@ spec: description: MaxDataSize limits the size of request/response data included in audit logs (in bytes). type: integer + maxDelegationDepth: + default: 10 + description: |- + MaxDelegationDepth caps how many nested RFC 8693 "act" entries are + recorded in an audit event's delegation chain. Deeper chains are + truncated (marked with truncated=true). Defaults to 10 when unset. + minimum: 1 + type: integer type: object backends: description: |- @@ -1309,6 +1354,57 @@ spec: CompositeToolConfig defines a composite tool workflow. This matches the YAML structure from the proposal (lines 173-255). properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string @@ -1642,6 +1738,16 @@ spec: audience: description: Audience is the required token audience. type: string + caBundlePath: + description: |- + CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle + used when the OIDC middleware performs HTTPS requests to the issuer + (OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle + at this path is added to the trust store used for verifying the issuer's + TLS certificate. Typically populated by the Kubernetes operator from + MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster + service-account CA when using Kubernetes service-account auth. + type: string clientId: description: ClientID is the OAuth client ID. type: string @@ -1816,6 +1922,10 @@ spec: When enabled, vMCP exposes only find_tool and call_tool operations to clients instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. + Enabling the optimizer currently pins this instance to MCP 2025-11-25: the + find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28) + clients are negotiated down to the Legacy revision (see + pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089). properties: embeddingHeaders: additionalProperties: @@ -3423,6 +3533,11 @@ spec: check format: date-time type: string + mcpRevision: + description: |- + MCPRevision is the backend's negotiated MCP protocol revision + ("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed. + type: string message: description: Message provides additional information about the backend status @@ -4477,16 +4592,43 @@ spec: prefixFormat: default: '{workload}_' description: |- - PrefixFormat defines the prefix format for the "prefix" strategy. + PrefixFormat defines the prefix format for the "prefix" tool strategy + and for backend-prefixed prompt names (the default for every prompt; + under the "priority" strategy, backends listed in priorityOrder keep + their own prompt names). Supports placeholders: {workload}, {workload}_, {workload}. type: string priorityOrder: - description: PriorityOrder defines the workload priority - order for the "priority" strategy. + description: |- + PriorityOrder defines the workload priority order for the "priority" strategy. + Listed workloads also keep their own prompt names (unlisted workloads' + prompts stay backend-prefixed). items: type: string type: array type: object + defaultToolVisibility: + description: |- + DefaultToolVisibility controls whether a backend with NO entry in Tools has its + tools advertised to MCP clients. + - allow (default): every tool from an unlisted backend is advertised, so + adding a workload to the group exposes it without further configuration. + - deny: an unlisted backend contributes no tools, so only backends named in + Tools are advertised. Use this when the set of exposed tools must be + enumerated deliberately rather than inherited from group membership. + + A backend that DOES have a Tools entry is unaffected by this setting: the + entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like + every other visibility setting here, this controls advertising only — hidden + tools remain in the routing table for composite tools (see the type doc). + + This gates TOOLS only. An unlisted backend's resources, resource templates, + and prompts are still advertised under deny; mergeResources/mergePrompts have + no equivalent check. + enum: + - allow + - deny + type: string excludeAllTools: description: |- ExcludeAllTools hides all backend tools from MCP clients when true. @@ -4532,24 +4674,34 @@ spec: Only specified fields are overridden; others pass through from the backend. properties: destructiveHint: - description: DestructiveHint overrides the - destructive hint annotation. + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean idempotentHint: - description: IdempotentHint overrides the - idempotent hint annotation. + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean openWorldHint: - description: OpenWorldHint overrides the open-world - hint annotation. + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean readOnlyHint: - description: ReadOnlyHint overrides the read-only - hint annotation. + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean title: - description: Title overrides the human-readable - title annotation. + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: string type: object description: @@ -4644,6 +4796,14 @@ spec: description: MaxDataSize limits the size of request/response data included in audit logs (in bytes). type: integer + maxDelegationDepth: + default: 10 + description: |- + MaxDelegationDepth caps how many nested RFC 8693 "act" entries are + recorded in an audit event's delegation chain. Deeper chains are + truncated (marked with truncated=true). Defaults to 10 when unset. + minimum: 1 + type: integer type: object backends: description: |- @@ -4769,6 +4929,57 @@ spec: CompositeToolConfig defines a composite tool workflow. This matches the YAML structure from the proposal (lines 173-255). properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string @@ -5102,6 +5313,16 @@ spec: audience: description: Audience is the required token audience. type: string + caBundlePath: + description: |- + CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle + used when the OIDC middleware performs HTTPS requests to the issuer + (OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle + at this path is added to the trust store used for verifying the issuer's + TLS certificate. Typically populated by the Kubernetes operator from + MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster + service-account CA when using Kubernetes service-account auth. + type: string clientId: description: ClientID is the OAuth client ID. type: string @@ -5276,6 +5497,10 @@ spec: When enabled, vMCP exposes only find_tool and call_tool operations to clients instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. + Enabling the optimizer currently pins this instance to MCP 2025-11-25: the + find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28) + clients are negotiated down to the Legacy revision (see + pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089). properties: embeddingHeaders: additionalProperties: @@ -6883,6 +7108,11 @@ spec: check format: date-time type: string + mcpRevision: + description: |- + MCPRevision is the backend's negotiated MCP protocol revision + ("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed. + type: string message: description: Message provides additional information about the backend status diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpauthzconfigs.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpauthzconfigs.yaml index 8a6d899874..ace9054aff 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpauthzconfigs.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpauthzconfigs.yaml @@ -30,9 +30,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -158,43 +155,6 @@ spec: for this MCPAuthzConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPAuthzConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -208,9 +168,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -337,43 +294,6 @@ spec: for this MCPAuthzConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPAuthzConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml index b73e66445b..c079660af4 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpexternalauthconfigs.yaml @@ -31,9 +31,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -1582,43 +1579,6 @@ spec: It corresponds to the MCPExternalAuthConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPExternalAuthConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -1632,9 +1592,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -3184,43 +3141,6 @@ spec: It corresponds to the MCPExternalAuthConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPExternalAuthConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpoidcconfigs.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpoidcconfigs.yaml index c46d3aa8f5..5ba2803faa 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpoidcconfigs.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpoidcconfigs.yaml @@ -30,9 +30,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -130,7 +127,7 @@ spec: insecureAllowHTTP: default: false description: |- - InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing. + InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing. WARNING: This is insecure and should NEVER be used in production. type: boolean introspectionUrl: @@ -291,43 +288,6 @@ spec: for this MCPOIDCConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPOIDCConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -341,9 +301,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -442,7 +399,7 @@ spec: insecureAllowHTTP: default: false description: |- - InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing. + InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing. WARNING: This is insecure and should NEVER be used in production. type: boolean introspectionUrl: @@ -603,43 +560,6 @@ spec: for this MCPOIDCConfig. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPOIDCConfig. - Each entry identifies the workload by kind and name. The map key is the - (kind, name) pair so two workloads of different kinds that share a name - (e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct - entries rather than colliding under merge-patch semantics. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - kind - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptelemetryconfigs.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptelemetryconfigs.yaml index 3377fe0408..d140312bcc 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptelemetryconfigs.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptelemetryconfigs.yaml @@ -290,33 +290,6 @@ spec: for this MCPTelemetryConfig. format: int64 type: integer - referencingWorkloads: - description: ReferencingWorkloads lists workloads that reference this - MCPTelemetryConfig - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -591,33 +564,6 @@ spec: for this MCPTelemetryConfig. format: int64 type: integer - referencingWorkloads: - description: ReferencingWorkloads lists workloads that reference this - MCPTelemetryConfig - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptoolconfigs.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptoolconfigs.yaml index d0b232ae59..de7120a4f2 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptoolconfigs.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcptoolconfigs.yaml @@ -28,9 +28,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -192,39 +189,6 @@ spec: It corresponds to the MCPToolConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPToolConfig. - Each entry identifies the workload by kind and name. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true @@ -235,9 +199,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referenceCount - name: References - type: integer - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -400,39 +361,6 @@ spec: It corresponds to the MCPToolConfig's generation, which is updated on mutation by the API Server. format: int64 type: integer - referenceCount: - description: ReferenceCount is the number of workloads referencing - this config. - format: int32 - type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPToolConfig. - Each entry identifies the workload by kind and name. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpwebhookconfigs.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpwebhookconfigs.yaml index 0640f21bc8..29c76d0901 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpwebhookconfigs.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_mcpwebhookconfigs.yaml @@ -27,9 +27,6 @@ spec: - jsonPath: .status.conditions[?(@.type=='Valid')].status name: Valid type: string - - jsonPath: .status.referencingWorkloads - name: References - type: string - jsonPath: .metadata.creationTimestamp name: Age type: date @@ -351,34 +348,6 @@ spec: to the current status format: int64 type: integer - referencingWorkloads: - description: |- - ReferencingWorkloads is a list of workload resources that reference this MCPWebhookConfig. - Each entry identifies the workload by kind and name. - items: - description: |- - WorkloadReference identifies a workload that references a shared configuration resource. - Namespace is implicit — cross-namespace references are not supported. - properties: - kind: - description: Kind is the type of workload resource - enum: - - MCPServer - - VirtualMCPServer - - MCPRemoteProxy - type: string - name: - description: Name is the name of the workload resource - minLength: 1 - type: string - required: - - kind - - name - type: object - type: array - x-kubernetes-list-map-keys: - - name - x-kubernetes-list-type: map type: object type: object served: true diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml index 8a7c7efc85..db2e16f117 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpcompositetooldefinitions.yaml @@ -71,6 +71,57 @@ spec: This embeds the CompositeToolConfig from pkg/vmcp/config to share the configuration model between CLI and operator usage. properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string @@ -470,6 +521,57 @@ spec: This embeds the CompositeToolConfig from pkg/vmcp/config to share the configuration model between CLI and operator usage. properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml index 3f29533245..cedb050902 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml @@ -1020,16 +1020,43 @@ spec: prefixFormat: default: '{workload}_' description: |- - PrefixFormat defines the prefix format for the "prefix" strategy. + PrefixFormat defines the prefix format for the "prefix" tool strategy + and for backend-prefixed prompt names (the default for every prompt; + under the "priority" strategy, backends listed in priorityOrder keep + their own prompt names). Supports placeholders: {workload}, {workload}_, {workload}. type: string priorityOrder: - description: PriorityOrder defines the workload priority - order for the "priority" strategy. + description: |- + PriorityOrder defines the workload priority order for the "priority" strategy. + Listed workloads also keep their own prompt names (unlisted workloads' + prompts stay backend-prefixed). items: type: string type: array type: object + defaultToolVisibility: + description: |- + DefaultToolVisibility controls whether a backend with NO entry in Tools has its + tools advertised to MCP clients. + - allow (default): every tool from an unlisted backend is advertised, so + adding a workload to the group exposes it without further configuration. + - deny: an unlisted backend contributes no tools, so only backends named in + Tools are advertised. Use this when the set of exposed tools must be + enumerated deliberately rather than inherited from group membership. + + A backend that DOES have a Tools entry is unaffected by this setting: the + entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like + every other visibility setting here, this controls advertising only — hidden + tools remain in the routing table for composite tools (see the type doc). + + This gates TOOLS only. An unlisted backend's resources, resource templates, + and prompts are still advertised under deny; mergeResources/mergePrompts have + no equivalent check. + enum: + - allow + - deny + type: string excludeAllTools: description: |- ExcludeAllTools hides all backend tools from MCP clients when true. @@ -1075,24 +1102,34 @@ spec: Only specified fields are overridden; others pass through from the backend. properties: destructiveHint: - description: DestructiveHint overrides the - destructive hint annotation. + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean idempotentHint: - description: IdempotentHint overrides the - idempotent hint annotation. + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean openWorldHint: - description: OpenWorldHint overrides the open-world - hint annotation. + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean readOnlyHint: - description: ReadOnlyHint overrides the read-only - hint annotation. + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean title: - description: Title overrides the human-readable - title annotation. + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: string type: object description: @@ -1187,6 +1224,14 @@ spec: description: MaxDataSize limits the size of request/response data included in audit logs (in bytes). type: integer + maxDelegationDepth: + default: 10 + description: |- + MaxDelegationDepth caps how many nested RFC 8693 "act" entries are + recorded in an audit event's delegation chain. Deeper chains are + truncated (marked with truncated=true). Defaults to 10 when unset. + minimum: 1 + type: integer type: object backends: description: |- @@ -1312,6 +1357,57 @@ spec: CompositeToolConfig defines a composite tool workflow. This matches the YAML structure from the proposal (lines 173-255). properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string @@ -1645,6 +1741,16 @@ spec: audience: description: Audience is the required token audience. type: string + caBundlePath: + description: |- + CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle + used when the OIDC middleware performs HTTPS requests to the issuer + (OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle + at this path is added to the trust store used for verifying the issuer's + TLS certificate. Typically populated by the Kubernetes operator from + MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster + service-account CA when using Kubernetes service-account auth. + type: string clientId: description: ClientID is the OAuth client ID. type: string @@ -1819,6 +1925,10 @@ spec: When enabled, vMCP exposes only find_tool and call_tool operations to clients instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. + Enabling the optimizer currently pins this instance to MCP 2025-11-25: the + find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28) + clients are negotiated down to the Legacy revision (see + pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089). properties: embeddingHeaders: additionalProperties: @@ -3426,6 +3536,11 @@ spec: check format: date-time type: string + mcpRevision: + description: |- + MCPRevision is the backend's negotiated MCP protocol revision + ("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed. + type: string message: description: Message provides additional information about the backend status @@ -4480,16 +4595,43 @@ spec: prefixFormat: default: '{workload}_' description: |- - PrefixFormat defines the prefix format for the "prefix" strategy. + PrefixFormat defines the prefix format for the "prefix" tool strategy + and for backend-prefixed prompt names (the default for every prompt; + under the "priority" strategy, backends listed in priorityOrder keep + their own prompt names). Supports placeholders: {workload}, {workload}_, {workload}. type: string priorityOrder: - description: PriorityOrder defines the workload priority - order for the "priority" strategy. + description: |- + PriorityOrder defines the workload priority order for the "priority" strategy. + Listed workloads also keep their own prompt names (unlisted workloads' + prompts stay backend-prefixed). items: type: string type: array type: object + defaultToolVisibility: + description: |- + DefaultToolVisibility controls whether a backend with NO entry in Tools has its + tools advertised to MCP clients. + - allow (default): every tool from an unlisted backend is advertised, so + adding a workload to the group exposes it without further configuration. + - deny: an unlisted backend contributes no tools, so only backends named in + Tools are advertised. Use this when the set of exposed tools must be + enumerated deliberately rather than inherited from group membership. + + A backend that DOES have a Tools entry is unaffected by this setting: the + entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like + every other visibility setting here, this controls advertising only — hidden + tools remain in the routing table for composite tools (see the type doc). + + This gates TOOLS only. An unlisted backend's resources, resource templates, + and prompts are still advertised under deny; mergeResources/mergePrompts have + no equivalent check. + enum: + - allow + - deny + type: string excludeAllTools: description: |- ExcludeAllTools hides all backend tools from MCP clients when true. @@ -4535,24 +4677,34 @@ spec: Only specified fields are overridden; others pass through from the backend. properties: destructiveHint: - description: DestructiveHint overrides the - destructive hint annotation. + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean idempotentHint: - description: IdempotentHint overrides the - idempotent hint annotation. + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean openWorldHint: - description: OpenWorldHint overrides the open-world - hint annotation. + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean readOnlyHint: - description: ReadOnlyHint overrides the read-only - hint annotation. + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: boolean title: - description: Title overrides the human-readable - title annotation. + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. type: string type: object description: @@ -4647,6 +4799,14 @@ spec: description: MaxDataSize limits the size of request/response data included in audit logs (in bytes). type: integer + maxDelegationDepth: + default: 10 + description: |- + MaxDelegationDepth caps how many nested RFC 8693 "act" entries are + recorded in an audit event's delegation chain. Deeper chains are + truncated (marked with truncated=true). Defaults to 10 when unset. + minimum: 1 + type: integer type: object backends: description: |- @@ -4772,6 +4932,57 @@ spec: CompositeToolConfig defines a composite tool workflow. This matches the YAML structure from the proposal (lines 173-255). properties: + annotations: + description: |- + Annotations declares MCP tool annotations for the composite tool. + + Annotation derivation runs at ADVERTISE TIME (when tools/list is served + and the backend tools are aggregated), not at CRD admission — thv vmcp + validate does NOT check annotation contradictions. The derived floor is + fail-closed: when the workflow has one or more tool steps the floor is + always non-nil, and any step whose annotations are nil/unknown taints the + floor conservatively (readOnly=false, destructive=true, openWorld=true). + A workflow with no tool steps (e.g. only elicitation) has no floor. + + When nil, annotations are derived from the annotations of the backend tools + referenced by the workflow's steps (e.g. readOnlyHint is true only when every + step tool is read-only). When set, the values are an explicit author + declaration merged over the derived floor — subject to a safety-floor + guardrail that drops the composite tool (with a warning naming the offending + step tools) if an explicit hint would make the tool look safer than its + steps allow. + properties: + destructiveHint: + description: |- + DestructiveHint sets the destructive hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + idempotentHint: + description: |- + IdempotentHint sets the idempotent hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + openWorldHint: + description: |- + OpenWorldHint sets the open-world hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + readOnlyHint: + description: |- + ReadOnlyHint sets the read-only hint annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: boolean + title: + description: |- + Title sets the human-readable title annotation. + For tool overrides this replaces the backend value; for composite tools + it is an explicit declaration merged over the derived safety floor. + type: string + type: object description: description: Description describes what the workflow does. type: string @@ -5105,6 +5316,16 @@ spec: audience: description: Audience is the required token audience. type: string + caBundlePath: + description: |- + CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle + used when the OIDC middleware performs HTTPS requests to the issuer + (OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle + at this path is added to the trust store used for verifying the issuer's + TLS certificate. Typically populated by the Kubernetes operator from + MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster + service-account CA when using Kubernetes service-account auth. + type: string clientId: description: ClientID is the OAuth client ID. type: string @@ -5279,6 +5500,10 @@ spec: When enabled, vMCP exposes only find_tool and call_tool operations to clients instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. + Enabling the optimizer currently pins this instance to MCP 2025-11-25: the + find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28) + clients are negotiated down to the Legacy revision (see + pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089). properties: embeddingHeaders: additionalProperties: @@ -6886,6 +7111,11 @@ spec: check format: date-time type: string + mcpRevision: + description: |- + MCPRevision is the backend's negotiated MCP protocol revision + ("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed. + type: string message: description: Message provides additional information about the backend status diff --git a/deploy/charts/operator/Chart.yaml b/deploy/charts/operator/Chart.yaml index dc00519ed1..c2a3bbdb80 100644 --- a/deploy/charts/operator/Chart.yaml +++ b/deploy/charts/operator/Chart.yaml @@ -2,5 +2,5 @@ apiVersion: v2 name: toolhive-operator description: A Helm chart for deploying the ToolHive Operator into Kubernetes. type: application -version: 0.39.0 -appVersion: "v0.39.0" +version: 0.42.1 +appVersion: "v0.42.1" diff --git a/deploy/charts/operator/README.md b/deploy/charts/operator/README.md index f129e60dc4..286b36198b 100644 --- a/deploy/charts/operator/README.md +++ b/deploy/charts/operator/README.md @@ -1,6 +1,6 @@ # ToolHive Operator Helm Chart -![Version: 0.39.0](https://img.shields.io/badge/Version-0.39.0-informational?style=flat-square) +![Version: 0.42.1](https://img.shields.io/badge/Version-0.42.1-informational?style=flat-square) ![Type: application](https://img.shields.io/badge/Type-application-informational?style=flat-square) A Helm chart for deploying the ToolHive Operator into Kubernetes. @@ -46,7 +46,7 @@ The command removes all the Kubernetes components associated with the chart and |-----|------|---------|-------------| | fullnameOverride | string | `"toolhive-operator"` | Provide a fully-qualified name override for resources | | nameOverride | string | `""` | Override the name of the chart | -| operator | object | `{"affinity":{},"autoscaling":{"enabled":false,"maxReplicas":100,"minReplicas":1,"targetCPUUtilizationPercentage":80},"containerSecurityContext":{"allowPrivilegeEscalation":false,"capabilities":{"drop":["ALL"]},"readOnlyRootFilesystem":true,"runAsNonRoot":true,"runAsUser":1000,"seccompProfile":{"type":"RuntimeDefault"}},"defaultImagePullSecrets":[],"defaultRedis":{"addr":"","existingSecret":"","existingSecretKey":""},"env":[],"features":{"experimental":false,"storageVersionMigrator":false},"gc":{"gogc":75,"gomemlimit":"110MiB"},"image":"ghcr.io/stacklok/toolhive/operator:v0.39.0","imagePullPolicy":"IfNotPresent","imagePullSecrets":[],"leaderElectionRole":{"binding":{"name":"toolhive-operator-leader-election-rolebinding"},"name":"toolhive-operator-leader-election-role","rules":[{"apiGroups":[""],"resources":["configmaps"],"verbs":["get","list","watch","create","update","patch","delete"]},{"apiGroups":["coordination.k8s.io"],"resources":["leases"],"verbs":["get","list","watch","create","update","patch","delete"]},{"apiGroups":["events.k8s.io"],"resources":["events"],"verbs":["create","patch"]}]},"livenessProbe":{"httpGet":{"path":"/healthz","port":"health"},"initialDelaySeconds":15,"periodSeconds":20},"nodeSelector":{},"pdb":{"enabled":false,"maxUnavailable":1},"podAnnotations":{},"podLabels":{},"podSecurityContext":{"runAsNonRoot":true},"ports":[{"containerPort":8080,"name":"metrics","protocol":"TCP"},{"containerPort":8081,"name":"health","protocol":"TCP"}],"proxyHost":"0.0.0.0","rbac":{"allowedNamespaces":[],"scope":"cluster"},"readinessProbe":{"httpGet":{"path":"/readyz","port":"health"},"initialDelaySeconds":5,"periodSeconds":10},"replicaCount":1,"resources":{"limits":{"cpu":"500m","memory":"128Mi"},"requests":{"cpu":"10m","memory":"64Mi"}},"serviceAccount":{"annotations":{},"automountServiceAccountToken":true,"create":true,"labels":{},"name":"toolhive-operator"},"tolerations":[],"toolhiveRunnerImage":"ghcr.io/stacklok/toolhive/proxyrunner:v0.39.0","vmcpImage":"ghcr.io/stacklok/toolhive/vmcp:v0.39.0","volumeMounts":[],"volumes":[]}` | All values for the operator deployment and associated resources | +| operator | object | `{"affinity":{},"autoscaling":{"enabled":false,"maxReplicas":100,"minReplicas":1,"targetCPUUtilizationPercentage":80},"containerSecurityContext":{"allowPrivilegeEscalation":false,"capabilities":{"drop":["ALL"]},"readOnlyRootFilesystem":true,"runAsNonRoot":true,"runAsUser":1000,"seccompProfile":{"type":"RuntimeDefault"}},"defaultImagePullSecrets":[],"defaultRedis":{"addr":"","existingSecret":"","existingSecretKey":""},"env":[],"features":{"experimental":false,"storageVersionMigrator":true},"gc":{"gogc":75,"gomemlimit":"110MiB"},"image":"ghcr.io/stacklok/toolhive/operator:v0.42.1","imageDiscovery":{"enabled":false,"resources":{}},"imagePullPolicy":"IfNotPresent","imagePullSecrets":[],"leaderElectionRole":{"binding":{"name":"toolhive-operator-leader-election-rolebinding"},"name":"toolhive-operator-leader-election-role","rules":[{"apiGroups":[""],"resources":["configmaps"],"verbs":["get","list","watch","create","update","patch","delete"]},{"apiGroups":["coordination.k8s.io"],"resources":["leases"],"verbs":["get","list","watch","create","update","patch","delete"]},{"apiGroups":["events.k8s.io"],"resources":["events"],"verbs":["create","patch"]}]},"livenessProbe":{"httpGet":{"path":"/healthz","port":"health"},"initialDelaySeconds":15,"periodSeconds":20},"nodeSelector":{},"pdb":{"enabled":false,"maxUnavailable":1},"podAnnotations":{},"podLabels":{},"podSecurityContext":{"runAsNonRoot":true},"ports":[{"containerPort":8080,"name":"metrics","protocol":"TCP"},{"containerPort":8081,"name":"health","protocol":"TCP"}],"proxyHost":"0.0.0.0","rbac":{"allowedNamespaces":[],"scope":"cluster"},"readinessProbe":{"httpGet":{"path":"/readyz","port":"health"},"initialDelaySeconds":5,"periodSeconds":10},"replicaCount":1,"resources":{"limits":{"cpu":"500m","memory":"128Mi"},"requests":{"cpu":"10m","memory":"64Mi"}},"serviceAccount":{"annotations":{},"automountServiceAccountToken":true,"create":true,"labels":{},"name":"toolhive-operator"},"tolerations":[],"toolhiveRunnerImage":"ghcr.io/stacklok/toolhive/proxyrunner:v0.42.1","vmcpImage":"ghcr.io/stacklok/toolhive/vmcp:v0.42.1","volumeMounts":[],"volumes":[]}` | All values for the operator deployment and associated resources | | operator.affinity | object | `{}` | Affinity settings for the operator pod | | operator.autoscaling | object | `{"enabled":false,"maxReplicas":100,"minReplicas":1,"targetCPUUtilizationPercentage":80}` | Configuration for horizontal pod autoscaling | | operator.autoscaling.enabled | bool | `false` | Enable autoscaling for the operator | @@ -61,11 +61,14 @@ The command removes all the Kubernetes components associated with the chart and | operator.defaultRedis.existingSecretKey | string | `""` | existingSecretKey is the key within existingSecret that holds the password. Empty means use global.redis.existingSecretKey or fall back to "redis-password". | | operator.env | list | `[]` | Environment variables to set in the operator container. Supported toolhive-specific variables include: - TOOLHIVE_SKIP_UPDATE_CHECK: set to "true" to disable the operator's periodic update check against the ToolHive update API. Also disables the usage-metrics collection that is gated on the same check. | | operator.features.experimental | bool | `false` | Enable experimental features | -| operator.features.storageVersionMigrator | bool | `false` | Enable the StorageVersionMigrator controller, which auto-cleans status.storedVersions on opted-in toolhive.stacklok.dev CRDs so a future release can drop deprecated versions (e.g. v1alpha1) without orphaning etcd objects in the cluster. Sets TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR in the operator deployment. | +| operator.features.storageVersionMigrator | bool | `true` | Enable the StorageVersionMigrator controller, which auto-cleans status.storedVersions on opted-in toolhive.stacklok.dev CRDs so a future release can drop deprecated versions (e.g. v1alpha1) without orphaning etcd objects in the cluster. Enabled by default; set to false to opt out and handle storage-version cleanup yourself. Sets TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR in the operator deployment. Requires `operator.rbac.scope=cluster` — the controller watches cluster-scoped CRDs and re-stores resources across all namespaces, so the chart rejects this being true when scope is namespace. | | operator.gc | object | `{"gogc":75,"gomemlimit":"110MiB"}` | Go memory limits and garbage collection percentage for the operator container | | operator.gc.gogc | int | `75` | Go garbage collection percentage for the operator container | | operator.gc.gomemlimit | string | `"110MiB"` | Go memory limits for the operator container | -| operator.image | string | `"ghcr.io/stacklok/toolhive/operator:v0.39.0"` | Container image for the operator | +| operator.image | string | `"ghcr.io/stacklok/toolhive/operator:v0.42.1"` | Container image for the operator | +| operator.imageDiscovery | object | `{"enabled":false,"resources":{}}` | Air-gap image-discovery aid. toolhiveRunnerImage, vmcpImage, and registryAPI.image are never rendered as chart `image:` keys — the operator only injects them into child workloads at runtime via env vars (TOOLHIVE_RUNNER_IMAGE / VMCP_IMAGE / TOOLHIVE_REGISTRY_API_IMAGE). Static manifest scanners used by air-gap tooling (e.g. a vendor portal's image list generator) discover images by reading `image:` keys out of `helm template` output, so they miss these three and an air-gapped customer can hit ImagePullBackOff on the first MCPServer / VirtualMCPServer / MCPRegistry create. Enabling this renders a `replicas: 0` Deployment (`-image-discovery`) whose pod template references all three images. A zero-replica Deployment never has Kubernetes schedule a pod from it, so the images are never pulled or run — this exists purely to put the refs in front of a static manifest scan. | +| operator.imageDiscovery.enabled | bool | `false` | Render the zero-replica image-discovery Deployment described above. | +| operator.imageDiscovery.resources | object | `{}` | Resource requests/limits for the image-discovery Deployment's containers. Admission control (a LimitRange, an OPA/Kyverno policy requiring explicit requests) validates these at CREATE time even though replicas: 0 means nothing ever actually runs, so this is exposed separately from operator.resources for clusters that need it set. | | operator.imagePullPolicy | string | `"IfNotPresent"` | Image pull policy for the operator container | | operator.imagePullSecrets | list | `[]` | List of image pull secrets to use | | operator.leaderElectionRole | object | `{"binding":{"name":"toolhive-operator-leader-election-rolebinding"},"name":"toolhive-operator-leader-election-role","rules":[{"apiGroups":[""],"resources":["configmaps"],"verbs":["get","list","watch","create","update","patch","delete"]},{"apiGroups":["coordination.k8s.io"],"resources":["leases"],"verbs":["get","list","watch","create","update","patch","delete"]},{"apiGroups":["events.k8s.io"],"resources":["events"],"verbs":["create","patch"]}]}` | Leader election role configuration | @@ -95,8 +98,8 @@ The command removes all the Kubernetes components associated with the chart and | operator.serviceAccount.labels | object | `{}` | Labels to add to the service account | | operator.serviceAccount.name | string | `"toolhive-operator"` | The name of the service account to use. If not set and create is true, a name is generated. | | operator.tolerations | list | `[]` | Tolerations for the operator pod | -| operator.toolhiveRunnerImage | string | `"ghcr.io/stacklok/toolhive/proxyrunner:v0.39.0"` | Image to use for Toolhive runners | -| operator.vmcpImage | string | `"ghcr.io/stacklok/toolhive/vmcp:v0.39.0"` | Image to use for Virtual MCP Server (vMCP) deployments | +| operator.toolhiveRunnerImage | string | `"ghcr.io/stacklok/toolhive/proxyrunner:v0.42.1"` | Image to use for Toolhive runners | +| operator.vmcpImage | string | `"ghcr.io/stacklok/toolhive/vmcp:v0.42.1"` | Image to use for Virtual MCP Server (vMCP) deployments | | operator.volumeMounts | list | `[]` | Additional volume mounts on the operator container | | operator.volumes | list | `[]` | Additional volumes to mount on the operator pod | | registryAPI | object | `{"image":"ghcr.io/stacklok/thv-registry-api:v1.3.0"}` | All values for the registry API deployment and associated resources | diff --git a/deploy/charts/operator/ci/imageDiscoveryEnabled-values.yaml b/deploy/charts/operator/ci/imageDiscoveryEnabled-values.yaml new file mode 100644 index 0000000000..894e66e4d3 --- /dev/null +++ b/deploy/charts/operator/ci/imageDiscoveryEnabled-values.yaml @@ -0,0 +1,6 @@ +operator: + image: ko.local/thv-operator:ci-test + toolhiveRunnerImage: ko.local/thv-proxyrunner:ci-test + vmcpImage: ko.local/vmcp:ci-test + imageDiscovery: + enabled: true diff --git a/deploy/charts/operator/templates/_helpers.tpl b/deploy/charts/operator/templates/_helpers.tpl index 193941d1c4..0147124539 100644 --- a/deploy/charts/operator/templates/_helpers.tpl +++ b/deploy/charts/operator/templates/_helpers.tpl @@ -79,4 +79,4 @@ Common labels for the toolhive resources {{- define "toolhive.labels" -}} app: toolhive app.kubernetes.io/name: toolhive -{{- end }} \ No newline at end of file +{{- end }} diff --git a/deploy/charts/operator/templates/image-discovery.yaml b/deploy/charts/operator/templates/image-discovery.yaml new file mode 100644 index 0000000000..a2ea587fab --- /dev/null +++ b/deploy/charts/operator/templates/image-discovery.yaml @@ -0,0 +1,64 @@ +{{- /* +See operator.imageDiscovery in values.yaml for why this exists. This +Deployment is always rendered at replicas: 0 — Kubernetes never schedules a +pod from it, so toolhiveRunnerImage / vmcpImage / registryAPI.image are never +pulled or run. Its only job is to put those refs behind an `image:` key so a +static manifest scan (used by offline/air-gapped image-mirroring tooling) +discovers them, since the operator otherwise only ever injects them into +child workloads via env vars at runtime. + +Despite replicas: 0, admission control (PSA `restricted`, Kyverno, OPA +Gatekeeper, a `LimitRange`) validates the PodTemplateSpec at CREATE time +regardless of replica count, so this mirrors deployment.yaml's security +context / resources / pull-secret handling rather than skipping it. +*/}} +{{- if .Values.operator.imageDiscovery.enabled }} +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ printf "%s-image-discovery" (include "toolhive-operator.fullname" .) | trunc 63 | trimSuffix "-" }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "toolhive-operator.labels" . | nindent 4 }} + app.kubernetes.io/component: image-discovery +spec: + replicas: 0 + selector: + matchLabels: + {{- include "toolhive-operator.selectorLabels" . | nindent 6 }} + app.kubernetes.io/component: image-discovery + template: + metadata: + labels: + {{- include "toolhive-operator.labels" . | nindent 8 }} + app.kubernetes.io/component: image-discovery + spec: + {{- with .Values.operator.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + securityContext: + {{- toYaml .Values.operator.podSecurityContext | nindent 8 }} + containers: + - name: toolhive-runner + securityContext: + {{- toYaml .Values.operator.containerSecurityContext | nindent 12 }} + image: "{{ .Values.operator.toolhiveRunnerImage }}" + imagePullPolicy: {{ .Values.operator.imagePullPolicy }} + resources: + {{- toYaml .Values.operator.imageDiscovery.resources | nindent 12 }} + - name: vmcp + securityContext: + {{- toYaml .Values.operator.containerSecurityContext | nindent 12 }} + image: "{{ .Values.operator.vmcpImage }}" + imagePullPolicy: {{ .Values.operator.imagePullPolicy }} + resources: + {{- toYaml .Values.operator.imageDiscovery.resources | nindent 12 }} + - name: registry-api + securityContext: + {{- toYaml .Values.operator.containerSecurityContext | nindent 12 }} + image: "{{ .Values.registryAPI.image }}" + imagePullPolicy: {{ .Values.operator.imagePullPolicy }} + resources: + {{- toYaml .Values.operator.imageDiscovery.resources | nindent 12 }} +{{- end }} diff --git a/deploy/charts/operator/templates/validate-features.yaml b/deploy/charts/operator/templates/validate-features.yaml new file mode 100644 index 0000000000..0850aac282 --- /dev/null +++ b/deploy/charts/operator/templates/validate-features.yaml @@ -0,0 +1,17 @@ +{{/* +Validate feature-flag / RBAC-scope combinations and fail the render early with +an actionable message rather than deploying a wedged operator. + +The StorageVersionMigrator controller watches cluster-scoped +CustomResourceDefinition objects and re-stores custom resources across every +namespace, so it requires cluster-scoped RBAC and a cluster-scoped manager +cache. A namespace-scoped operator (operator.rbac.scope=namespace) gets only +per-namespace RoleBindings and a namespace-restricted cache, so the controller +cannot sync its CRD informer and would prevent the manager from starting. We +fail loudly here instead of silently dropping the feature, because a silent +drop would let a namespace-scoped admin believe storedVersions are being kept +clean when they are not. +*/}} +{{- if and .Values.operator.features.storageVersionMigrator (ne .Values.operator.rbac.scope "cluster") -}} +{{- fail "operator.features.storageVersionMigrator requires operator.rbac.scope=cluster: the StorageVersionMigrator controller watches cluster-scoped CustomResourceDefinitions and re-stores resources across all namespaces, which a namespace-scoped operator cannot do. Set operator.features.storageVersionMigrator=false for namespace-scoped installs." -}} +{{- end -}} diff --git a/deploy/charts/operator/tests/default_install_test.yaml b/deploy/charts/operator/tests/default_install_test.yaml index adc4f2145e..440c0e9f4e 100644 --- a/deploy/charts/operator/tests/default_install_test.yaml +++ b/deploy/charts/operator/tests/default_install_test.yaml @@ -52,7 +52,7 @@ tests: template: deployment.yaml asserts: - contains: { path: 'spec.template.spec.containers[0].env', content: { name: ENABLE_EXPERIMENTAL_FEATURES, value: "false" } } - - contains: { path: 'spec.template.spec.containers[0].env', content: { name: TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR, value: "false" } } + - contains: { path: 'spec.template.spec.containers[0].env', content: { name: TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR, value: "true" } } - contains: { path: 'spec.template.spec.containers[0].env', content: { name: UNSTRUCTURED_LOGS, value: "false" } } - contains: { path: 'spec.template.spec.containers[0].env', content: { name: TOOLHIVE_USE_CONFIGMAP, value: "true" } } - contains: { path: 'spec.template.spec.containers[0].env', content: { name: GOGC, value: "75" } } diff --git a/deploy/charts/operator/tests/feature_flags_test.yaml b/deploy/charts/operator/tests/feature_flags_test.yaml index 6faaf9b521..11c1cbced7 100644 --- a/deploy/charts/operator/tests/feature_flags_test.yaml +++ b/deploy/charts/operator/tests/feature_flags_test.yaml @@ -1,15 +1,18 @@ suite: opt-in feature flags # The experimental and storage-version-migrator features are quoted-boolean env -# vars. A dropped quote or a renamed var silently disables the feature, so pin -# both the enabled and (in the baseline suite) disabled states. -set: - operator.features.experimental: true - operator.features.storageVersionMigrator: true +# vars. A dropped quote or a renamed var silently flips the feature, so pin both +# the enabled and disabled states. storageVersionMigrator defaults to true, so +# the disabled state is the one the baseline suite can't cover — pin it here. +# The migrator also requires cluster-scoped RBAC, so the chart rejects it in +# namespace scope; pin that validation too. templates: - deployment.yaml tests: - it: flips both feature-flag env vars to the quoted string "true" template: deployment.yaml + set: + operator.features.experimental: true + operator.features.storageVersionMigrator: true asserts: - contains: path: 'spec.template.spec.containers[0].env' @@ -17,3 +20,28 @@ tests: - contains: path: 'spec.template.spec.containers[0].env' content: { name: TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR, value: "true" } + + - it: opts out of the storage version migrator when set to false + template: deployment.yaml + set: + operator.features.storageVersionMigrator: false + asserts: + - contains: + path: 'spec.template.spec.containers[0].env' + content: { name: TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR, value: "false" } + + # Namespace-scope + migrator enabled is tested separately in + # validate_features_test.yaml (suite-level set is required for failedTemplate). + - it: allows namespace scope when the storage version migrator is disabled + template: deployment.yaml + set: + operator.rbac.scope: namespace + operator.rbac.allowedNamespaces: [toolhive-system] + operator.features.storageVersionMigrator: false + asserts: + - contains: + path: 'spec.template.spec.containers[0].env' + content: { name: TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR, value: "false" } + - contains: + path: 'spec.template.spec.containers[0].env' + content: { name: WATCH_NAMESPACE, value: "toolhive-system" } diff --git a/deploy/charts/operator/tests/image_discovery_test.yaml b/deploy/charts/operator/tests/image_discovery_test.yaml new file mode 100644 index 0000000000..261ec5e398 --- /dev/null +++ b/deploy/charts/operator/tests/image_discovery_test.yaml @@ -0,0 +1,129 @@ +suite: image discovery +# operator.imageDiscovery.enabled renders a zero-replica Deployment whose only +# purpose is to surface toolhiveRunnerImage / vmcpImage / registryAPI.image as +# `image:` keys for static air-gap manifest scanning. Pinning: the Deployment +# is absent by default, present with replicas:0 when enabled, carries the +# exact three image refs, and mirrors deployment.yaml's admission-relevant +# fields (security contexts, resources, pull secrets) — admission control +# validates the PodTemplateSpec at CREATE time regardless of replicas:0, so a +# missing field here would silently break enabling this on a hardened +# cluster. Also pins that enabling the flag doesn't disturb the real operator +# Deployment. +templates: + - image-discovery.yaml + - deployment.yaml +tests: + - it: renders nothing by default + template: image-discovery.yaml + asserts: + - hasDocuments: + count: 0 + + - it: renders a zero-replica Deployment referencing all three runtime-injected images + template: image-discovery.yaml + set: + operator.imageDiscovery.enabled: true + operator.toolhiveRunnerImage: ghcr.io/stacklok/toolhive/proxyrunner:v9.9.9 + operator.vmcpImage: ghcr.io/stacklok/toolhive/vmcp:v9.9.9 + registryAPI.image: ghcr.io/stacklok/thv-registry-api:v9.9.9 + asserts: + - hasDocuments: + count: 1 + - isKind: + of: Deployment + - isAPIVersion: + of: apps/v1 + - equal: + path: spec.replicas + value: 0 + - equal: + path: 'spec.template.spec.containers[0].name' + value: toolhive-runner + - equal: + path: 'spec.template.spec.containers[0].image' + value: ghcr.io/stacklok/toolhive/proxyrunner:v9.9.9 + - equal: + path: 'spec.template.spec.containers[1].name' + value: vmcp + - equal: + path: 'spec.template.spec.containers[1].image' + value: ghcr.io/stacklok/toolhive/vmcp:v9.9.9 + - equal: + path: 'spec.template.spec.containers[2].name' + value: registry-api + - equal: + path: 'spec.template.spec.containers[2].image' + value: ghcr.io/stacklok/thv-registry-api:v9.9.9 + + - it: selector does not collide with the operator Deployment's own selector + template: image-discovery.yaml + set: + operator.imageDiscovery.enabled: true + asserts: + - equal: + path: spec.selector.matchLabels["app.kubernetes.io/component"] + value: image-discovery + + - it: sets pod and container security contexts on every container so hardened-cluster admission (PSA/Kyverno/Gatekeeper) does not reject it + template: image-discovery.yaml + set: + operator.imageDiscovery.enabled: true + asserts: + - equal: + path: spec.template.spec.securityContext.runAsNonRoot + value: true + - equal: + path: 'spec.template.spec.containers[0].securityContext.allowPrivilegeEscalation' + value: false + - equal: + path: 'spec.template.spec.containers[1].securityContext.allowPrivilegeEscalation' + value: false + - equal: + path: 'spec.template.spec.containers[2].securityContext.allowPrivilegeEscalation' + value: false + + - it: sets resources and imagePullPolicy on every container so a LimitRange requiring explicit requests does not reject it + template: image-discovery.yaml + set: + operator.imageDiscovery.enabled: true + asserts: + - equal: + path: 'spec.template.spec.containers[0].resources' + value: {} + - equal: + path: 'spec.template.spec.containers[0].imagePullPolicy' + value: IfNotPresent + + - it: inherits operator.imagePullSecrets + template: image-discovery.yaml + set: + operator.imageDiscovery.enabled: true + operator.imagePullSecrets: + - name: regcred + asserts: + - contains: + path: spec.template.spec.imagePullSecrets + content: { name: regcred } + + - it: truncates the Deployment name at 63 characters for a long fullnameOverride + template: image-discovery.yaml + set: + operator.imageDiscovery.enabled: true + fullnameOverride: a-very-long-fullname-override-that-pushes-the-name-past-the-limit + asserts: + - matchRegex: + path: metadata.name + pattern: '^.{1,63}$' + + - it: still renders exactly one operator Deployment when imageDiscovery is enabled + template: deployment.yaml + set: + operator.imageDiscovery.enabled: true + asserts: + - hasDocuments: + count: 1 + - isKind: + of: Deployment + - equal: + path: spec.replicas + value: 1 diff --git a/deploy/charts/operator/tests/namespace_scope_test.yaml b/deploy/charts/operator/tests/namespace_scope_test.yaml index 85608a8791..acb401357b 100644 --- a/deploy/charts/operator/tests/namespace_scope_test.yaml +++ b/deploy/charts/operator/tests/namespace_scope_test.yaml @@ -10,6 +10,9 @@ release: set: operator.rbac.scope: namespace operator.rbac.allowedNamespaces: [team-a, team-b] + # storageVersionMigrator requires cluster scope; the chart rejects it under + # namespace scope (see feature_flags_test.yaml), so disable it here. + operator.features.storageVersionMigrator: false templates: - deployment.yaml - clusterrole/rolebinding.yaml diff --git a/deploy/charts/operator/tests/validate_features_test.yaml b/deploy/charts/operator/tests/validate_features_test.yaml new file mode 100644 index 0000000000..2c260de044 --- /dev/null +++ b/deploy/charts/operator/tests/validate_features_test.yaml @@ -0,0 +1,15 @@ +suite: validate-features — namespace scope rejects migrator +# helm-unittest requires failedTemplate values to be set at the suite level; +# test-level set is not applied when the template override differs from the +# suite templates list. +set: + operator.rbac.scope: namespace + operator.rbac.allowedNamespaces: [toolhive-system] + operator.features.storageVersionMigrator: true +templates: + - validate-features.yaml +tests: + - it: rejects the storage version migrator when scope is namespace + asserts: + - failedTemplate: + errorPattern: "operator.features.storageVersionMigrator requires operator.rbac.scope=cluster" diff --git a/deploy/charts/operator/values.yaml b/deploy/charts/operator/values.yaml index 6667e97704..13fe2b1c79 100644 --- a/deploy/charts/operator/values.yaml +++ b/deploy/charts/operator/values.yaml @@ -11,9 +11,13 @@ operator: # -- Enable the StorageVersionMigrator controller, which auto-cleans # status.storedVersions on opted-in toolhive.stacklok.dev CRDs so a # future release can drop deprecated versions (e.g. v1alpha1) without - # orphaning etcd objects in the cluster. Sets + # orphaning etcd objects in the cluster. Enabled by default; set to + # false to opt out and handle storage-version cleanup yourself. Sets # TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR in the operator deployment. - storageVersionMigrator: false + # Requires `operator.rbac.scope=cluster` — the controller watches + # cluster-scoped CRDs and re-stores resources across all namespaces, so + # the chart rejects this being true when scope is namespace. + storageVersionMigrator: true # -- Number of replicas for the operator deployment replicaCount: 1 @@ -59,15 +63,40 @@ operator: existingSecretKey: "" # -- Container image for the operator - image: ghcr.io/stacklok/toolhive/operator:v0.39.0 + image: ghcr.io/stacklok/toolhive/operator:v0.42.1 # -- Image pull policy for the operator container imagePullPolicy: IfNotPresent # -- Image to use for Toolhive runners - toolhiveRunnerImage: ghcr.io/stacklok/toolhive/proxyrunner:v0.39.0 + toolhiveRunnerImage: ghcr.io/stacklok/toolhive/proxyrunner:v0.42.1 # -- Image to use for Virtual MCP Server (vMCP) deployments - vmcpImage: ghcr.io/stacklok/toolhive/vmcp:v0.39.0 + vmcpImage: ghcr.io/stacklok/toolhive/vmcp:v0.42.1 + + # -- Air-gap image-discovery aid. toolhiveRunnerImage, vmcpImage, and + # registryAPI.image are never rendered as chart `image:` keys — the operator + # only injects them into child workloads at runtime via env vars + # (TOOLHIVE_RUNNER_IMAGE / VMCP_IMAGE / TOOLHIVE_REGISTRY_API_IMAGE). Static + # manifest scanners used by air-gap tooling (e.g. a vendor portal's image + # list generator) discover images by reading `image:` keys out of + # `helm template` output, so they miss these three and an air-gapped + # customer can hit ImagePullBackOff on the first MCPServer / + # VirtualMCPServer / MCPRegistry create. + # + # Enabling this renders a `replicas: 0` Deployment + # (`-image-discovery`) whose pod template references all + # three images. A zero-replica Deployment never has Kubernetes schedule a + # pod from it, so the images are never pulled or run — this exists purely + # to put the refs in front of a static manifest scan. + imageDiscovery: + # -- Render the zero-replica image-discovery Deployment described above. + enabled: false + # -- Resource requests/limits for the image-discovery Deployment's + # containers. Admission control (a LimitRange, an OPA/Kyverno policy + # requiring explicit requests) validates these at CREATE time even + # though replicas: 0 means nothing ever actually runs, so this is + # exposed separately from operator.resources for clusters that need it set. + resources: {} # -- Host for the proxy deployed by the operator proxyHost: 0.0.0.0 diff --git a/deploy/keycloak/mcpserver-with-auth.yaml b/deploy/keycloak/mcpserver-with-auth.yaml index 85ecf97a3a..f2bdc8110b 100644 --- a/deploy/keycloak/mcpserver-with-auth.yaml +++ b/deploy/keycloak/mcpserver-with-auth.yaml @@ -8,6 +8,8 @@ spec: inline: # Keycloak issuer URL for the toolhive realm issuer: http://keycloak:8080/realms/toolhive + # Dev Keycloak has no TLS, so the HTTP issuer and JWKS URLs need this opt-in. + insecureAllowHTTP: true # Explicit JWKS URL to avoid OIDC discovery issues jwksUrl: http://keycloak-dev-service.keycloak.svc.cluster.local:8080/realms/toolhive/protocol/openid-connect/certs # Optional: Allow private IP addresses for development diff --git a/docs/arch/02-core-concepts.md b/docs/arch/02-core-concepts.md index bfc7eb45bb..1215e059ef 100644 --- a/docs/arch/02-core-concepts.md +++ b/docs/arch/02-core-concepts.md @@ -659,6 +659,8 @@ See `pkg/audit/mcp_events.go` for complete list of event types. **Event data:** - Timestamp, source, outcome - Subjects (user, session) +- Delegation chain (`delegation`) — RFC 8693 `act` claim hops (issuer/subject + pairs), outermost first, when the caller used a delegated token - Target (endpoint, method, resource) - Request/response data (configurable) - Duration and metadata diff --git a/docs/arch/03-transport-architecture.md b/docs/arch/03-transport-architecture.md index 1b944cc836..75db441c78 100644 --- a/docs/arch/03-transport-architecture.md +++ b/docs/arch/03-transport-architecture.md @@ -211,10 +211,10 @@ All proxy types integrate with the middleware chain: ```mermaid graph LR - Client[Client Request] --> MW1[Middleware 1
Auth] - MW1 --> MW2[Middleware 2
Parser] - MW2 --> MW3[Middleware 3
Authz] - MW3 --> MW4[Middleware 4
Audit] + Client[Client Request] --> MW1[Middleware 1
Audit] + MW1 --> MW2[Middleware 2
Auth] + MW2 --> MW3[Middleware 3
Parser] + MW3 --> MW4[Middleware 4
Authz] MW4 --> Proxy[Proxy Handler] Proxy --> Container[MCP Server] @@ -227,7 +227,7 @@ graph LR **Implementation:** - `pkg/transport/types/transport.go` - `MiddlewareFunction` and `NamedMiddleware` types -- Middleware applied in reverse order (last registered = outermost) +- Middleware wraps the handler in reverse registration order, so the first registered entry is the outermost wrapper and runs first at request time - Each transport type accepts `[]NamedMiddleware` in constructor (each wraps a `MiddlewareFunction` with its name for logging) ## Remote MCP Server Proxying @@ -515,6 +515,194 @@ stdin, stdout, err := t.deployer.AttachToWorkload(ctx, t.containerName) - Ephemeral sessions for sessionless requests - DELETE `/mcp` to explicitly close session +### Server->Client Routing (Streamable HTTP) + +**Implementation**: `pkg/transport/proxy/streamable/streamable_proxy.go`, +`pkg/transport/proxy/streamable/dispatcher.go`, +`pkg/transport/proxy/streamable/dispatcher_streams.go`, +`pkg/transport/proxy/streamable/dispatcher_routing.go` + +The streamable HTTP proxy sits in front of a single shared backend process +(one container, one stdio pipe -- see `pkg/transport/stdio.go`) that may be +serving many concurrent MCP sessions. Every server->client message the +backend emits arrives on one shared channel with **no built-in session +attribution**, so the proxy must reconstruct "which session does this belong +to" itself before it can deliver anything, and must drop -- never guess or +broadcast -- whatever it cannot attribute. `dispatchResponses` +(`dispatcher.go`) is the single place every such message passes through; it +routes by concrete JSON-RPC type and method: + +- `*jsonrpc2.Response` -> the waiter-correlation path (`routeResponseToWaiter`, + unchanged): responses are matched back to the in-flight HTTP request that + sent them via a composite key (`compositeKey(sessID, idKey)`) baked into the + outgoing request's wire ID and echoed back by the backend. +- `*jsonrpc2.Request` with no valid ID (a **notification**) -> routed by + method (`routeNotification`), described in detail below. +- `*jsonrpc2.Request` with a valid ID (a server->client **request**, e.g. + `sampling/createMessage`, `elicitation/create`) -> `rejectServerRequestToBackend` + writes a JSON-RPC error (code `-32601`) back to the **backend** (not to any + client) via `SendMessageToDestination`, so the backend's own blocking call + unblocks instead of hanging until its own timeout. The shared-backend proxy + has no way to route a client's reply back to the correct originating backend + call across potentially many sessions, so it cannot forward this to any + client without either guessing (misdelivery) or broadcasting (cross-session + leakage). Per-session backends -- see the [Virtual MCP + architecture](10-virtual-mcp-architecture.md) -- are the correct place to + support server->client requests; tracked as a follow-up in #5744. + +#### Notification routing table + +| Method | Destination | Mechanism | +|---|---|---| +| `notifications/tools\|resources\|prompts/list_changed` | every connected session's standalone GET stream | `serverStreams.broadcast` -- GLOBAL, no session-specific payload | +| `notifications/progress` | the ONE originating request's POST SSE stream | `routeProgress` + `sessionRouter`'s progress-token table | +| `notifications/resources/updated` | only sessions subscribed to the notification's `uri` | `routeResourceUpdated` + `sessionRouter`'s subscription table | +| `notifications/message` (logging) | nobody -- **SECURE-DROP** | shared-backend log content is unattributable to any session | +| anything else | nobody | dropped, logged at Debug | + +**Progress correlation.** A client that wants progress for a request sets +`params._meta.progressToken` and sends the request with +`Accept: text/event-stream` (`handleSingleRequestSSE`). Before forwarding +upstream, the proxy mints a fresh, proxy-only token (`ptGlobal`, a UUID), +rewrites `_meta.progressToken` to it (`rewriteMetaProgressToken`, which copies +rather than mutates the caller's request), and records a `progressRoute` +(delivery channel + the client's *original* token) keyed by `ptGlobal` in +`sessionRouter`. `handleSingleRequestSSE`'s response loop then `select`s over +both that delivery channel and the final response's waiter channel, +interleaving every progress notification the backend sends as its own SSE +`data:` frame (with the original client token restored) and only returning +once the final response (or a context/shutdown signal) arrives. The route is +dropped as soon as the request completes (`defer`), so it can never be reused +by a later, unrelated notification. A request with **no** `Accept: +text/event-stream` (a plain JSON POST) gets no progress route at all -- its +token is forwarded unrewritten, and any progress the backend sends for it is +correctly dropped (no route is ever registered for it), since a non-streaming +client has no channel to receive interim progress on anyway. + +Because `ptGlobal` is unique per request regardless of what token value the +client chose, two different sessions asking for progress with the identical +client-visible token can never cross-deliver: each is keyed internally by its +own UUID. + +**Subscriptions.** `resources/subscribe` and `resources/unsubscribe` are +ref-counted per `uri` in `sessionRouter`'s subscription table, for **Legacy, +session-bearing requests only** (a request with no real, persisted session -- +Modern, or Legacy sessionless -- has nothing durable to track, so it always +forwards upstream unmodified). The interception happens in `handlePost` +*after* `applyMiddlewares`/authz has admitted the request, so an +authorization-denied subscribe is never recorded. A `resources/subscribe` for +a `uri` with no existing subscriber is forwarded upstream (`interceptSubscribe`, +via `doRequest`, under `uriLocks` for that `uri`), and is recorded in the +subscription table **only if that upstream call succeeds**: a backend +rejection, error, or timeout is returned to the client as-is and leaves no +trace in the subscription table, so a later session's subscribe for the same +`uri` is free to try again upstream rather than being dedup-served a +synthesized success for a subscription the backend never actually granted. A +`resources/subscribe` for a `uri` that already has a recorded subscriber is +answered with a locally synthesized success and never reaches the backend at +all, avoiding redundant upstream subscribe calls -- this dedup is safe +specifically because the entry it is deduping against is known to have +succeeded upstream. Symmetrically, `resources/unsubscribe` only reaches the +backend once the **last** subscriber for a `uri` leaves. `notifications/resources/updated` +is then delivered only to that `uri`'s current subscribers' standalone GET +streams (`serverStreams.deliverToMany`). + +**Trust-boundary note.** Because identical-`uri` subscriptions are ref-counted +down to a single upstream `resources/subscribe`, a backend that performs its +*own* per-resource access control inside its subscribe handler only ever sees +and evaluates the **first** subscriber for a given `uri` -- sessions 2..N are +served a synthesized success without the backend's handler ever running for +them. ToolHive's own per-session middleware/authz still runs in front of +*every* subscribe attempt regardless of dedup (see the authz-non-bypass note +above), so this is not a bypass of ToolHive's own controls. It is a +characteristic of the shared-backend proxy design worth calling out +explicitly: any authorization a backend implements *inside its own +subscribe handler*, keyed on the resource `uri`, is not re-invoked for the +2nd..Nth subscriber of that `uri`. Deployments relying on backend-side, +per-resource access control for `resources/subscribe` should account for this +when sharing one backend across sessions with different resource +entitlements. + +**Logging.** `logging/setLevel` is intercepted the same way: the client's +requested level is recorded per-session, and if it changes the **maximum** +verbosity requested across all sessions, a reconciled `logging/setLevel` is +sent upstream with that maximum (see `logLevelRank`, aligned with RFC 5424 +syslog severity numbers) so no session is under-served by another session's +less-verbose request silently overriding it. The client always gets an +immediate synthesized success. The resulting `notifications/message` the +backend then emits is dropped per the routing table above -- there is +currently no way to safely forward it to the session that actually wanted +that verbosity. + +**Lifecycle.** `sessionRouter`'s three tables (progress tokens, subscriptions, +log levels) each have their own mutex (no shared locking). Progress tokens are +request-scoped and dropped when their request completes. +Subscriptions/log-levels are session-scoped and are purged in two ways: (1) +explicitly, by `handleDelete` (`purgeSession`), and (2) periodically, by a +reaper goroutine (`reapRoutingState`, started alongside `dispatchResponses` in +`Start`) that ticks at `sessionTTL/2` (matching `session.Manager`'s own +cleanup cadence) and drops state for any session `session.Manager` no longer +reports as active. That liveness check is the session manager's ordinary +`Get`, which -- by design, to keep genuinely active sessions alive -- refreshes +the session's TTL on every call; using it here means a session that still owns +routing state can have its expiry delayed by up to one reap tick beyond actual +client inactivity. This is an accepted, documented tradeoff, not a leak: `Get` +on an already-deleted session simply returns false without refreshing +anything, so routing state can never outlive a session's true deletion. + +Because the GET handler (and every POST) is registered through the same +`applyMiddlewares` chain, tool-filter and other response-shaping middleware +see every frame written to any stream -- there is no separate, unfiltered code +path for server->client messages. + +**Behavior and configuration:** + +- Standalone SSE is **on by default**. Pass `WithStandaloneSSE(false)` to + restore the prior behavior (GET returns 405). +- Delivery is **best-effort and non-blocking** throughout: a stream (or + progress delivery channel) whose buffer is full has its message dropped + (logged) rather than blocking delivery to anything else. +- A notification dispatched while no GET stream is connected for its target + session is dropped. There is no replay buffer or pending-message queue + (unlike `httpsse`'s SSE proxy, which queues for reconnecting clients) -- a + client that wants server->client notifications must keep a GET stream open. +- Per MCP 2025-11-25, a server MUST NOT deliver the same notification to a + session more than once, so `serverStreamRegistry` allows **at most one + stream per session**: a second concurrent GET for the same session evicts + the first (its handler observes a closed per-stream `stop` signal and + returns). Each registered stream has two channels: a buffered `data` channel + that is **never closed** (only ever read from or sent to), and a `stop` + channel closed exactly once -- by an evicting `register` call or by proxy + `Stop`'s `closeAll` -- to signal the consumer to return. This avoids a + send-on-closed-channel panic if a concurrent broadcast/dispatch races + eviction/shutdown. + +**Relationship to Modern `subscriptions/listen`**: the 2026-07-28 (Modern) +revision (see `mcp.MCPVersionModern` and `pkg/mcp/revision.go`) replaces the +standalone GET with `subscriptions/listen`, a POST method whose response stream +stays open for server-pushed subscription events. + +vMCP now **serves** that method on its client edge +(`pkg/vmcp/server/modern_subscriptions.go`), and deliberately does **not** use +`serverStreamRegistry`. The reason is that the handler honors no subscriptions: +the honored set is intersected against the capabilities `server/discover` +advertises, every push flag there is false, so the acknowledged set is always +empty and the stream terminates immediately. With nothing to fan out, wiring in a +fan-out registry from another package would add an abstraction with zero +consumers. What the handler does implement is the wire contract — a long-lived +POST rather than a GET, keyed per-subscription-id (the listen request's own +JSON-RPC id) rather than per session ID since Modern has no sessions, an initial +`notifications/subscriptions/acknowledged`, and +`io.modelcontextprotocol/subscriptionId` tagging. + +`serverStreamRegistry` remains the sensible seam for the piece still missing: +actual **delivery** of notifications on a Modern listen stream (#5743). That work +needs per-notification-type AND per-URI opt-in filtering rather than blanket +fan-out, and per `schema/draft/schema.ts` must tag *every* notification with the +subscription id, not just the acknowledgement. `dispatcher_streams.go` is +therefore still intentionally kept transport-agnostic (no HTTP types), so that +when delivery lands it does not also require rewriting the fan-out primitives. + ## Error Handling ### Connection Failures @@ -594,6 +782,25 @@ For deployment behind reverse proxy, proxies respect X-Forwarded headers (Host, **Security**: Only enable if ToolHive is behind trusted reverse proxy. +### MCP-Protocol-Version Validation (Streamable HTTP) + +**Implementation**: `pkg/transport/proxy/streamable/streamable_proxy.go`, `pkg/transport/proxy/streamable/utils.go` + +By default, the streamable HTTP proxy is version-agnostic: it accepts any +`MCP-Protocol-Version` header value (including an absent header, per the +streamable HTTP spec's rule to assume `2025-03-26`). This is intentional -- +the proxy is transport-level and does not depend on a specific MCP revision, +so it avoids being pedantic and breaking on new protocol dates. + +**Opt-in strict mode**: `thv run --strict-protocol-validation` rejects a +request whose `MCP-Protocol-Version` header names an unknown/unsupported MCP +revision with HTTP 400. An absent header is still accepted in strict mode. +The set of recognized revisions is `supportedMCPVersions` in +`pkg/transport/proxy/streamable/utils.go`. This is wired through +`runner.WithStrictProtocolValidation` -> `RunConfig.StrictProtocolValidation` +-> `types.Config.StrictProtocolValidation` -> `StdioTransport.SetStrictProtocolValidation` +-> `streamable.WithStrictProtocolValidation`. + ### SSE Endpoint URL Rewriting **Problem**: When using path-based ingress routing that strips path prefixes: diff --git a/docs/arch/05-runconfig-and-permissions.md b/docs/arch/05-runconfig-and-permissions.md index 2ea393b280..0d3ab48578 100644 --- a/docs/arch/05-runconfig-and-permissions.md +++ b/docs/arch/05-runconfig-and-permissions.md @@ -386,7 +386,7 @@ type Profile struct { } ``` -`NetworkPermissions` also carries a `Mode` field (`"host"`, `"bridge"`, `"none"`) that selects the container runtime network mode; when empty, the runtime default is used. +`NetworkPermissions` also carries a `Mode` field (`"host"`, `"bridge"`, `"none"`) that selects the container runtime network mode; when empty, the runtime default is used. Network isolation is only enforceable in bridge mode; see [Network isolation requires bridge networking](#network-isolation-requires-bridge-networking). ### Filesystem Permissions @@ -534,6 +534,30 @@ When `isolate_network: true` in RunConfig: **Implementation**: `pkg/container/docker/squid.go`, `pkg/networking/` +##### Network isolation requires bridge networking + +Network isolation is a bridge-topology construct: it builds a per-workload +`internal=true` network plus DNS (dnsmasq) and egress/ingress (Squid) sidecars, +attaching the MCP container to the internal network only. This only works when +the network mode is bridge (`""`, `bridge`, or `default`). For non-bridge modes +the sidecars have no route out and the workload bypasses them entirely, so +isolation is dropped: + +| Network mode | `--isolate-network` | Behavior | +| --- | --- | --- | +| bridge/`""`/`default` | any | Isolation enforced as designed | +| `host` (or other non-bridge) | defaulted | Isolation dropped, `WARN` logged (dropping isolation reduces confinement) | +| `host` (or other non-bridge) | explicitly `true` | Build fails with an error naming both ways out | +| `none` | any | Isolation dropped silently (`DEBUG`), since `none` is already maximally confined | + +`host` and `none` are intentionally asymmetric: dropping isolation for a +host-networked workload is a real reduction in confinement (warn, or fail fast +when explicitly requested), whereas `none` already blocks all networking, so +isolation is merely redundant. This is enforced at three layers: config build +(`pkg/runner/config_builder.go`), config load (`pkg/runner/config.go` `ReadJSON` +self-heals legacy on-disk configs), and deploy (`pkg/container/docker/client.go` +as a safety net for restart paths). See issue #5775. + ### Privileged Mode **⚠️ Warning**: Privileged mode removes most security isolation! diff --git a/docs/arch/06-registry-system.md b/docs/arch/06-registry-system.md index ba8dd1699f..9432dc78b2 100644 --- a/docs/arch/06-registry-system.md +++ b/docs/arch/06-registry-system.md @@ -973,6 +973,7 @@ re-sync, modify `configYAML` (or restart the registry API pod). - [Operator Architecture](09-operator-architecture.md) - MCPRegistry CRD - [RunConfig and Permissions](05-runconfig-and-permissions.md) - How registry metadata feeds a server's RunConfig - [Skills System](12-skills-system.md) - Skills discovery and distribution via registry +- [Plugins System](14-plugins-system.md) - Plugins discovery and distribution via registry (analogue of skills) ### External Documentation - [ToolHive User Documentation](https://docs.stacklok.com/toolhive/) - User-facing guides diff --git a/docs/arch/09-operator-architecture.md b/docs/arch/09-operator-architecture.md index 525c41e457..01be9ec8d3 100644 --- a/docs/arch/09-operator-architecture.md +++ b/docs/arch/09-operator-architecture.md @@ -249,7 +249,7 @@ MCPOIDCConfig eliminates OIDC configuration duplication — define an identity p - `scopes` (optional) — Defaults to `["openid"]` - `resourceUrl` (optional) — Public URL for OAuth protected resource metadata (RFC 9728); defaults to internal service URL -**Status fields** include a `Ready` condition, `configHash` for change detection, and `referencingWorkloads` tracking which resources reference this config. Deletion is blocked while references exist (finalizer pattern). +**Status fields** include a `Ready` condition and `configHash` for change detection. Deletion is blocked while references exist (finalizer pattern): the controller recomputes the referencing workloads on demand at deletion time (a field-indexed lookup), rather than maintaining a stored list in status. **Referenced by**: MCPServer and MCPRemoteProxy (via `spec.oidcConfigRef`); VirtualMCPServer (via `spec.incomingAuth.oidcConfigRef`) @@ -270,7 +270,7 @@ MCPTelemetryConfig centralises telemetry infrastructure settings (collector endp - CEL validation prevents header name overlap between `headers` and `sensitiveHeaders` - Per-server `serviceName` override in the workload's `telemetryConfigRef` (since `service.name` must be unique per server) -**Status fields** include a `Ready` condition, `configHash` for change detection, and `referencingWorkloads` tracking. +**Status fields** include a `Ready` condition and `configHash` for change detection. Referencing workloads are recomputed on demand at deletion time (finalizer + field-indexed lookup), not stored in status. **Referenced by**: MCPServer, VirtualMCPServer, MCPRemoteProxy (via `telemetryConfigRef`) @@ -291,7 +291,7 @@ MCPWebhookConfig declares one or more `validating` and/or `mutating` webhooks (U - Optional `tlsConfig` with CA bundle, mTLS client cert, or `insecureSkipVerify` - Optional `hmacSecretRef` adds the `X-Toolhive-Signature` header to outgoing payloads -**Status fields** include a `Ready` condition, `configHash` for change detection, and `referencingWorkloads`. +**Status fields** include a `Ready` condition and `configHash` for change detection. Referencing workloads are recomputed on demand at deletion time (finalizer + field-indexed lookup), not stored in status. **Referenced by**: MCPServer (via `webhookConfigRef`) diff --git a/docs/arch/10-virtual-mcp-architecture.md b/docs/arch/10-virtual-mcp-architecture.md index cc03ec710c..4033f513c9 100644 --- a/docs/arch/10-virtual-mcp-architecture.md +++ b/docs/arch/10-virtual-mcp-architecture.md @@ -136,10 +136,106 @@ When backends expose tools with the same name, vMCP resolves the conflict using | **priority** | First backend in priority order wins, others hidden | | **manual** | Explicit mapping for each conflict | +The three strategies above cover **tools**. Resources, resource templates, and +prompts are resolved separately (`aggregator/capability_conflicts.go`), with +policies that deliberately differ by what the identity *is*: + +- **Prompt names are names**, like tool names: `prompts/get` is translated back + to the backend's own name via `BackendTarget.GetBackendCapabilityName`, so a + rename does not break invocation. By **default** every prompt is renamed to + its backend-prefixed form — `conflictResolutionConfig.prefixFormat` applied + to the backend ID (default `{workload}_`), the same formatting path the tool + prefix strategy uses. The **priority** strategy is the escape hatch for + clients that pin prompt names, exactly as it is for tools: backends listed + in `priorityOrder` keep their bare prompt names, a bare-name collision among + listed backends resolves to the highest-priority one, and unlisted backends + stay always-prefixed — deliberately stricter than tool priority resolution, + which lets a conflict-free unlisted tool keep its bare name. A `manual` + configuration changes tool resolution only. + + The invariant both modes preserve, scoped precisely: the advertised name of a + given (backendID, name) **pair** is a pure function of the aggregation config + and that pair, so it never shifts because an unrelated backend joined or left + the group. That stability matters beyond naming: authorization matches on the + advertised name (Cedar builds `Prompt::""` entities), so a + membership-dependent rename would detach `permit` and `forbid` policies from + the prompt they were written for — names move only on an explicit config + edit. + + The invariant does **not** promise that an advertised *string* keeps naming + the same prompt. Under `priority`, two listed backends can claim the same + bare name and the higher-ranked one takes it, so a + `permit(... resource == Prompt::"review")` written while `b1` owned `review` + silently begins authorizing `b2`'s **different** prompt once a higher-ranked + `b2` advertising that name deploys — no config edit involved. Cedar's + resource identity is the advertised name and nothing else (it carries no + backend attribute), so whoever wins a shared name inherits every policy + written for it. `forbid` still fails closed, because the priority loser is + **dropped rather than aliased**; it is `permit` that gets redirected. Keep + `priorityOrder` and name-scoped `permit` policies under review together. + + The loser being dropped instead of re-prefixed is deliberate: prefixing it so + "nothing is lost" would re-advertise that prompt under a name no policy + mentions, putting it beyond any `forbid` written on the bare name — the + fail-open this design exists to avoid. + + When two advertised names compose to the same string and no backend owns it — + backend `b1` prompt `x_y` vs backend `b1_x` prompt `y`, or a prefixed name + hitting a listed backend's literal name — the name is **ambiguous**, and + **every** colliding prompt is dropped with the collision logged at `ERROR`. + Aggregation does not fail: + erroring would take the group's entire aggregated view down (tools, + resources, templates, prompts, backend visibility) over one prompt name that + needs no conflict-resolution config to reach — just an unlucky combination of + operator-chosen workload names and backend-chosen prompt names. Keeping one + claimant is not an option either, since the survivor would inherit whatever + policy was written for the prompt it collided with. With the name advertised + by nobody, every `permit` and `forbid` on it is vacuous and the rest of the + group keeps serving. +- **Resource URIs and template strings are locators, not names.** The client + passes them back verbatim (`resources/read`, `resources/subscribe`, + completion refs), backends emit them in notifications and embedded resource + contents, and template-matched reads forward the client's concrete URI + untranslated — so vMCP never rewrites them. A URI (or template string) + advertised by several backends is instead advertised **once**: the backend + earliest in sorted-backend-ID order wins, later duplicates are dropped with a + warning. Reads are unchanged — the routing table keys by URI, so only one + backend was ever served per URI; the fix makes that pick deterministic and + the advertised list agree with it. Only the winning backend's `name`, + `description` and `mimeType` are advertised for that URI, so a duplicate that + differed in those fields loses them along with its entry. + +After aggregation, every capability identity (`Tool.Name`, `Resource.URI`, +`ResourceTemplate.URITemplate`, `Prompt.Name`) is unique in the aggregated view. +Resources, resource templates and prompts are processed in sorted-backend-ID +order, so their outcomes are stable across runs; tool resolution +(`prefix_resolver.go`, `manual_resolver.go`) still iterates backends in map +order, so a tool collision *after* prefixing or a manual override picks a +nondeterministic winner. + +The Modern list paginator still does **not** rely on that uniqueness: its cursor +names a position within a run of equal keys rather than assuming keys are +distinct, because the paginator is generic and a plain "resume after this key" +scan silently dropped an item whenever a duplicate landed on a page boundary +(see [Client-facing list pagination](#client-facing-list-pagination)). + ### Tool Filtering Beyond conflict resolution, vMCP can filter which tools are exposed through allow/deny lists, renaming, and description overrides. +By default a backend with no per-workload entry has all of its tools advertised, so +adding a workload to the group exposes it. `aggregation.defaultToolVisibility: deny` +inverts that, advertising only backends named in `aggregation.tools` — useful when the +exposed tool set should be enumerated deliberately rather than inherited from group +membership. A listed backend is opted in by its entry; its own `excludeAll`/`filter` +then decide which of its tools are advertised. + +All of these settings — `excludeAllTools`, `defaultToolVisibility`, per-workload +`excludeAll`, and `filter` — control **advertising only**. Every backend tool stays in +the routing table so composite tools can call hidden ones, and none of them affect +resources or prompts. Per-identity authorization is Cedar's job (see [Authorization +Enforcement](#authorization-enforcement-core-admission-seam--pre-dispatch-gate)). + **Implementation**: `pkg/vmcp/aggregator/` ## Composite Tools @@ -169,7 +265,589 @@ Steps can be of three types: - **elicitation**: Request user input via MCP elicitation protocol - **forEach**: Iterate over a collection from a previous step, executing an inner tool step per item with bounded parallelism -**Implementation**: `pkg/vmcp/composer/` +**Tool annotations**: Composite tools advertise MCP tool annotations computed at advertise time via a derive-then-merge ordering. First, a safety floor is derived from the annotations of the step tools (including `forEach` inner steps): `readOnlyHint` is the AND across steps, `destructiveHint` and `openWorldHint` are the OR across steps (unknown steps taint conservatively), and `idempotentHint` is never derived. Then any explicit `annotations` declared on the composite tool definition are merged over the floor, with explicitly set fields winning. An explicit hint may be more conservative than the floor, but if it would make the tool look safer than its steps allow (e.g. `readOnlyHint: true` when a step is not read-only), the composite tool is dropped from `tools/list` with a warning rather than advertised misleadingly. + +**Implementation**: `pkg/vmcp/composer/` (execution), `pkg/vmcp/internal/compositetools/` (advertised tool conversion and annotation derivation) + +## Backend MCP Revision Classification + +vMCP speaks both MCP revisions at once. The **client edge** is classified per +request (see [Transport Architecture](03-transport-architecture.md)); each +**backend edge** is classified independently, so a client on either revision can +reach backends on either revision. This section covers the backend edge, in +`pkg/vmcp/client`. + +Two names are used throughout the code: **Legacy** is 2025-11-25 (session-based, +`initialize` handshake, `Mcp-Session-Id`) and **Modern** is 2026-07-28 +(stateless, no handshake, per-request `_meta`). + +### The probe: Modern-first `server/discover` + +`probeRevision` decides a backend's revision once, then caches it. It is +Modern-first — it asks, rather than assuming Legacy: + +1. **SSE backends short-circuit to Legacy.** A `TransportType` of `sse` is the + 2024-11-05 two-endpoint transport (`GET /sse` + `POST /messages`), which has + no Modern endpoint at `BaseURL` to discover. This gate is *ToolHive routing*, + not a protocol rule — HTTP+SSE is Deprecated in the spec, not removed. +2. Otherwise vMCP POSTs `server/discover` (`discoverModernCapabilities`) and + reads `supportedVersions`. + +**`supportedVersions` is the authoritative signal — a clean `server/discover` +response is not, by itself, proof of Modern.** A go-sdk v1.7 shim answers +`server/discover` even when it is negotiating *down* to Legacy, so vMCP requires +`supportedVersions` to actually contain `2026-07-28`; if it does not, the probe +fails with `errModernNegotiatedDown` and the backend is classified Legacy. The +containment check is an **exact match**, deliberately stricter than go-sdk's +reference client (which accepts negotiated `>= 2026-07-28`), because vMCP's shim +speaks exactly one Modern wire shape. That exactness is a tripwire: adding a +newer Modern revision must widen this to a set/range check, or a newer-only +backend is misclassified Legacy. + +Probe outcomes: + +| Discover result | Classification | Why | +|---|---|---| +| Success, `supportedVersions` contains `2026-07-28` | **Modern** | Backend affirmatively serves the revision | +| A Modern-specific `-3202x` protocol error | **Modern** | The peer validated our Modern `_meta`, so it speaks the revision | +| A valid Modern envelope with `resultType != "complete"` | **Modern** | Only a Modern peer produces that envelope (`errModernInputRequired`) | +| Auth rejection (401/403/407) or transient (408/429/5xx, timeout) | **Inconclusive** | Says nothing about the revision; left unprobed so the next call re-probes | +| Anything else, including negotiated-down | **Legacy** | `errModernNegotiatedDown`, `errWrongEra`, and every other genuine not-Modern signal | + +A probe that fails inconclusively falls back to Legacy for that one call +**without caching it**, so a transient outage can never pin a backend to the +wrong revision. + +### The cache, and its asymmetric self-healing + +`dispatch` wraps every backend verb: it resolves the revision (probing on a +miss), runs the call, and on a revision-shaped failure (`isRevisionMismatch`) +re-probes and retries. The two directions self-heal by different mechanisms, +which is worth knowing when reading the code: + +- **Modern cached, backend is really Legacy** — `server/discover` no longer + advertises `2026-07-28`, producing `errModernNegotiatedDown`, which + `isRevisionMismatch` recognises on the Modern arm; the cache is reclassified. +- **Legacy cached, backend is really Modern** — corrected **in band** rather + than by a mismatch: `legacyInit` reads the genuinely-negotiated protocol + version off every `InitializeResult` and flips the cache when it comes back + `2026-07-28`. The SDK's Modern-first client negotiates Modern unaided even on + vMCP's Legacy code path, so the call still succeeds while the label corrects + itself. + +Reclassification never re-runs a request that may already have executed. An +`errLegacyResponseBody` (a lenient Legacy backend that answered a Modern-shaped +request) corrects the cache for *future* calls but returns the error, because the +backend may have performed the side effect. + +### vMCP owns the reserved `_meta` namespace on both hops + +vMCP is the client's MCP peer and the backends' MCP peer on two different hops, +so the reserved `io.modelcontextprotocol/*` `_meta` keys are vMCP's to own in +both directions. A single helper, `mcp.StripReservedMeta`, removes every key +under that prefix (except the end-to-end passthrough keys — `related-task`, +`model-immediate-response`) wherever a `_meta` map crosses vMCP: + +- **Request egress** — before any Legacy backend call (`pkg/vmcp/client`, + `pkg/vmcp/session/internal/backend`) and on the Modern request path, which + overlays vMCP's own authoritative values afterwards. This is not cosmetic: + go-sdk v1.7 rejects **any** `_meta.protocolVersion` on a stateful + streamable-HTTP server outright (HTTP 400), regardless of its value. +- **Response/request egress to the client** — Legacy strips inside + `conversion.ToMCPMeta` (the funnel every Legacy egress crosses); Modern strips + in `newModernResultMeta` and re-stamps its own `serverInfo` last; the + elicitation adapter (a server→client request) crosses the same chokepoint. + This stops a backend fabricating the client's own identity (`serverInfo`, + `protocolVersion`, `clientInfo`). + +Non-reserved caller/backend keys (progress tokens, W3C trace context) are +preserved throughout. + +### Observability + +Each backend's resolved revision is exposed as `mcpRevision` on the backend +status read-model (`vmcp.BackendStatus.MCPRevision`, fed from the health +monitor's `RecordRevision`), and surfaces on +`VirtualMCPServer.status.discoveredBackends[]`. It is empty until the backend has +been probed. + +### Limitation: elicitation and sampling are unavailable on Modern backends + +vMCP declares **empty** `clientCapabilities` on every Modern backend call +(`mcp.ModernRequestMeta`). That is honest — the Modern shim performs single-shot +dispatch and cannot drive multi-round retrieval — but it has a consequence worth +naming explicitly rather than leaving to be inferred: + +**Mid-call elicitation and sampling forwarding, which vMCP *does* implement for +Legacy backends (`forwardingClientOptions`), is structurally unavailable for +Modern backends.** A spec-compliant Modern backend that needs caller input +returns `-32021`, which surfaces as an opaque `errModernProtocolError`. The +2026-07-28 revision replaces server-initiated requests with client-polled +multi-round tool retrieval (MRTR), so the fix is shaped MRTR-first rather than by +extending the SSE standalone-stream model — see the epic (#5743) and the +mid-call forwarding section below for the Legacy behaviour this contrasts with. + +### Limitation: elicitation and sampling are unavailable to Modern clients + +The client edge mirrors the backend edge. The Modern dispatcher +(`pkg/vmcp/server`'s `dispatchModern`) is single-shot: every result it builds is +`resultType: "complete"`, and it never emits `"input_required"` — MRTR +(SEP-2322) is unimplemented on this edge too. When a backend tool issues a +mid-call server-initiated request during a **Modern** client's `tools/call`, +there is no client session to forward it to, so the call fails with an explicit +error naming the refused request (pinned by +`TestIntegration_Modern_RealBackend_ElicitingToolFailsCleanly`). This is a +deliberate honest-unsupported error, not a gap left by accident: + +**The `-32603` is a documented deviation, not the spec's answer — and the +spec's answer is unshippable today.** For a client that did NOT declare the +needed capability in its per-request `clientCapabilities`, SEP-2575 MUSTs a +`-32021` `MissingRequiredClientCapabilityError` at **HTTP 400**, with +explicitly execution-time language ("if processing a request requires a +capability…") — go-sdk's own doc comment on the error type tells handlers to +return it mid-execution, so the mid-call timing is not the problem. The +problem is the transport: go-sdk's streamable client treats any non-transient +4xx (its transient set is only 500/502/503/504/429) as a **connection** +failure — `checkResponse` → `fail()` → a one-shot, permanent session death — +so a conformant 400 would tear down the entire client session to punish one +call. And for a client that DID declare the capability, the 2026-07-28 +vocabulary has no conformant code at all: no "operation not supported", MRTR +is not a server-advertised capability, and SEP-2322 has no decline mechanism. +**#6061 (merged) implements exactly that two-path contract** in +`writeModernCallFailure`/`writeModernMissingCapability` (`pkg/vmcp/server`): +`-32021` with `data.requiredCapabilities` and a message naming both the +capability and the gateway limitation, served at **HTTP 200** for the +undeclared case — deviating from the mandated 400 for exactly the reason +above (tracked upstream as go-sdk#1117) — and an explicit `-32603` naming +SEP-2322 for the declared case, as a documented spec gap. The MRTR design +([16-vmcp-mrtr.md](16-vmcp-mrtr.md)) supersedes this contract slice by slice +where the backend is Modern; it is permanent for the Modern-client ↔ Legacy- +backend cell. + +**A clean error does not mean nothing happened.** The refusal reaches the +backend mid-call, so a real backend tool may have executed — including side +effects — up to the point it demanded input. A Modern client receiving this +error must not assume the call was side-effect-free. (The integration +fixture's tools elicit as their first action, so the tests cannot exhibit +this; production tools can.) + +- The 2026-07-28 revision **removed** server-initiated requests; go-sdk's + `ServerSession.assertServerInitiatedRequestAllowed` refuses + elicitation/sampling/roots purely by negotiated protocol version, so no + capability negotiation can restore the Legacy forwarding model for Modern + clients. +- A server that never returns `input_required` is fully SEP-2575-conformant: + the per-request `clientCapabilities` a client declares are an offer the + server may use, not an obligation. +- SEP-2577 deprecates sampling (and logging and roots) outright as of + 2026-07-28, with direct LLM-provider integration as the sanctioned + replacement — so elicitation is the only durable consumer a future MRTR + implementation would serve. + +Legacy clients keep the full mid-call forwarding behaviour unchanged; the +forwarding integration tests pin their downstream clients to Legacy explicitly +(`legacyPinningRoundTripper` in `pkg/vmcp/server`'s external test package) +because that surface exists only on a Legacy session. + +**Bridging was considered, costed, and rejected.** Serving MRTR to Modern +clients on top of a *Legacy* +backend would require parking the live, mid-flight backend call server-side +(the blocked goroutine and its open session cannot be serialized into the +opaque `requestState` the SEP designed for handler re-invocation) and keying +the resume on an unguessable token — per-round server state with TTL/eviction, +identity binding on a token that becomes a capability to resume someone else's +in-flight call, and replica affinity with no `Mcp-Session-Id` to route on. That +would reintroduce, in different clothes, the per-request server state the +2026-07-28 revision removed. The spec's own sanctioned path for genuinely stateful +`input_required` work is the **Tasks** extension (SEP-2663: `tools/call` +returns `resultType: "task"` with a `taskId`; the client polls `tasks/get` and +answers outstanding `inputRequests` via `inputResponses` on `tasks/update`; +note SEP-2663 supersedes SEP-1686 and removed the blocking `tasks/result` +method for the same reasons argued here) — if Modern-client elicitation over +Legacy backends is ever truly demanded, that is the machinery to reach for, +not parked `tools/call`. + +The coherent future MRTR shape for a re-aggregating gateway is +**Modern-client ↔ Modern-backend pass-through** — relay a Modern backend's +`inputRequests`/`requestState` to the client and the client's +`inputResponses` back, genuinely stateless at vMCP. The full design is +[16-vmcp-mrtr.md](16-vmcp-mrtr.md); its slice 1 (the egress half) has landed — +a Modern backend's `input_required` still classifies as +`errModernInputRequired`, but the typed `vmcp.InputRequiredError` now carries +the decoded round for the relay slices to consume. By the time Modern +backends exist to relay from, SEP-2577's deprecations make elicitation its +only durable consumer; see #5743 and #6059. + +Progress and log notifications toward Modern clients are a separate concern +from MRTR: they remain spec-legal as request-scoped notifications on the +POST-initiated SSE response stream (SEP-2260 requires messages on that stream +to relate to the originating request; `progressToken` is unchanged), which the +single-shot dispatcher does not produce today — a vMCP streaming-dispatch gap, +not a spec absence. + +## Served MCP Capabilities + +Beyond tools, vMCP aggregates and serves the full complement of MCP capabilities. Every served capability flows through the domain **core** (`pkg/vmcp/core`), so the same admission decision that filters `tools/list` also gates reads, gets, and completions. + +| Capability | Served? | Notes | +|------------|---------|-------| +| Tools (`tools/list`, `tools/call`) | Yes | Aggregated, conflict-resolved, admission-filtered; a backend's asynchronous `notifications/tools/list_changed` is propagated to already-registered sessions — see below | +| Resources (`resources/list`, `resources/read`) | Yes | Admission-filtered per identity; a backend's asynchronous `notifications/resources/list_changed` propagates ADDITIONS (not removals) to already-registered sessions — see below | +| Resource templates (`resources/templates/list`) | Yes | Templated reads route through the same `ReadResource` path; the router matches an expanded URI against the aggregated templates, and an exact template-string key routes to its backend; covered by the same `notifications/resources/list_changed` propagation as resources (MCP 2025-11-25 has no separate wire method for template changes) | +| Prompts (`prompts/list`, `prompts/get`) | Yes | Served per-session; a backend's asynchronous `notifications/prompts/list_changed` propagates ADDITIONS (not removals) to already-registered sessions — see below | +| Completions (`completion/complete`) | Yes | A `ref/prompt` routes via the prompts table; a `ref/resource` carries the URI-template string per the spec and routes via the resource-templates table (exact template-string key, with exact-resource and template-expansion fallbacks). Unroutable refs return an empty result (lenient completion); admission denial returns an error | +| Resource subscriptions (`resources/subscribe`, `resources/unsubscribe`) | Ack-level only | Legacy edge. vMCP accepts and records the subscription (after session-binding and resource-admission checks) but does **not** yet forward backend `notifications/resources/updated` — see the limitation below | +| Subscriptions (`subscriptions/listen`) | Ack-level only | Modern (2026-07-28) edge, the revision's only server→client push channel. Answered with an SSE stream carrying the mandatory `notifications/subscriptions/acknowledged` and then a terminating result, both tagged `io.modelcontextprotocol/subscriptionId` and keyed by the listen request's JSON-RPC id (Modern has no sessions). The honored set is intersected against the advertised capabilities, all of whose push flags are false, so it is always **empty** and the stream closes immediately rather than idling. Serving it is what lets a go-sdk v1.7 client complete `Connect` at all — see below | + +The four Modern list verbs (`tools/list`, `resources/list`, `resources/templates/list`, `prompts/list`) **paginate**, emitting `nextCursor` while items remain and capping a page at 1000 items to match the page size the SDK applies on the Legacy path. See [Client-facing list pagination](#client-facing-list-pagination). + +The completion handler is a single global handler installed via `WithCompletionHandler`, so it recovers the session from the SDK request context rather than a per-session closure. Setting it makes the shim auto-advertise the `completions` capability at initialize. + +### Served MCP revisions: the Modern capability gate + +vMCP serves two client-facing MCP revisions: **Legacy** (2025-11-25, the SDK +session path) always, and **Modern** (2026-07-28, `classifyingHandler → +dispatchModern`, stateless) **conditionally** — only when every enabled feature +of the instance is servable by the stateless dispatch path. The condition is +`modernDispatchBlockers` (`pkg/vmcp/server/modern_gate.go`), an explicit +enumeration that replaced the temporary `TOOLHIVE_VMCP_MODERN_STATELESS` +env-var kill-switch (#5959): instead of a global "don't serve Modern", the +instance serves Modern exactly when it can serve it correctly. + +Features that currently gate Modern off, and why: + +| Feature | Why the stateless path cannot serve it | +|---------|----------------------------------------| +| Optimizer (`find_tool`/`call_tool`) | The meta-tools are Serve-layer and **session-scoped** (`serve_optimizer.go`): each session builds an FTS5 index over its advertised set and swaps the two meta-tools in. The index is deliberately not in the stateless core, and `dispatchModern` serves `tools/*` straight from `core.ListTools`/`core.CallTool` — a Modern client would silently receive the raw aggregated tool set and `tools/call find_tool` would fail. Parity needs an identity- or instance-scoped index | + +"Cannot serve" means a Modern client would silently get different behavior than +the feature promises — not merely that the feature is session-flavored. +Redis-backed session sharing, for example, does **not** gate Modern: Legacy +clients keep their shared, reconstructible sessions while Modern clients are +sessionless by design and store nothing, a coexistence asserted end-to-end by +`test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_redis_test.go`. Rate +limiting does not gate Modern either: the limiter is a core decorator +(`pkg/vmcp/ratelimit`), so both eras meter the same `CallTool` seam, and the +Modern dispatcher preserves the limiter's coded error on the wire — a real +JSON-RPC error object, `429` with `data.retryAfterSeconds` +(`writeModernCodedError`, at HTTP 200 because go-sdk rejects a non-200 +response before decoding the body, so on an HTTP 429 the error object — and +its `retryAfterSeconds` — would be discarded unread) — where the Legacy SDK +seam can only smuggle the same code and data into an `IsError` tool result's +`structuredContent` (`conversion.ErrorToToolResult`). The code `429` mirrors +HTTP Too Many Requests (the same pattern as the `403` denial code) and sits +outside the JSON-RPC reserved range (`-32768`..`-32000`), making it +conformant under both MCP 2026-07-28 — which reserves `-32020`..`-32099` +exclusively for spec-defined codes — and 2025-11-25 and earlier, which +inherit plain JSON-RPC 2.0. It replaced `-32029`, which sat inside the +reserved band (#6101). + +"Does not gate Modern" is not "costs the same", though. The limiter wraps only +the `CallTool` seam, so the list verbs and `server/discover` are unmetered on +both eras — but Legacy aggregates once per session registration, while Modern +re-runs the full backend fan-out on every request with no cache, by design +(`core_vmcp.go`'s `aggregatedView`). A Modern client can therefore loop +unmetered, uncached fan-outs — reachable unauthenticated when incoming auth is +anonymous. Rate-limiting the list/discover verbs, or a short-TTL per-identity +capability cache, is deferred until profiling shows the per-request fan-out +cost matters (#5761 — the same deferral recorded in `dispatchModernDiscover`'s +doc comment). + +Wire behavior when the gate is closed: + +- **`server/discover` falls through to the SDK** instead of dispatching. go-sdk's + `filterSupportedVersions` keeps every version the transport's + `SupportsProtocolVersion` accepts, and the stateful transport excludes only + >= 2026-07-28 (those require `Stateless`), so the probe answer is the + transport-filtered list — everything except Modern: `[2025-11-25, + 2025-06-18, 2025-03-26, 2024-11-05]`. This matters because go-sdk v1.7+ + clients are Modern-first: `Connect` probes `server/discover` **before** + `initialize` and upgrades to whatever the server advertises. The + fall-through answer is what lands them on the Legacy handshake — where + sessions, and every gated feature, work. That answer over-advertises, + though: the `-32022` refusal below lists only 2025-11-25, and the refusal is + the accurate one — `mcpcompat`'s `handleInitialize` always responds with + `LATEST_PROTOCOL_VERSION` regardless of what the client requests, so the + three older revisions in the discover answer are not actually servable. The + mismatch is the SDK's discover answer to narrow, not the `-32022` list to + widen. +- **Every other well-formed Modern request** is refused with a conformant + 400 + `-32022 UnsupportedProtocolVersionError` whose data lists the Legacy + version, the shape a client negotiates down from. It is answered in + `classifyingHandler` rather than falling through, because go-sdk's stateful + rejection for Modern traffic is a plain-text 400 carrying no version list. +- **Legacy traffic is untouched** either way; the gate only ever affects + requests that classified Modern. + +The gate is derived from construction-time configuration, logged once at +startup ("MCP 2026-07-28 (Modern) dispatch disabled…"), and pinned by +`TestModernDispatchBlockers`, `TestClassifyingHandler_ModernCapabilityGate`, +and the full-handler pair in `modern_gate_integration_test.go`. Achieving +Modern parity for a feature means deleting its entry and updating those tests — +nothing else needs to change. + +### Subscription limitation (ack-level) + +vMCP advertises `resources.subscribe: true` and answers `resources/subscribe` / `resources/unsubscribe` at **ack level**: the request is accepted (enforcing session binding and validating the URI is an advertised, admitted resource), and go-sdk records the subscription. vMCP does **not** currently propagate backend `notifications/resources/updated` to the subscribed client — doing so requires persistent per-session backend connections, which is out of scope. Clients that subscribe will receive a success ack but no update stream yet. + +### Client-facing list pagination + +On the Legacy edge the SDK's session-scoped feature store splits list results into +pages. `dispatchModern` bypasses the SDK, so the Modern edge paginates itself +(`pkg/vmcp/server/modern_pagination.go`). + +Modern has no sessions, so a cursor may not denote server-held iteration state. +The draft pagination page makes cursors **opaque to clients** ("MUST treat cursors +as opaque tokens"), which is precisely what allows the server to encode position +*into* the token instead of remembering it. So the cursor is self-describing: +base64url over a small JSON payload naming the list kind, the last key delivered, +and how many items sharing that key have already been sent. + +Three properties worth knowing: + +- **Ordering.** Keyset paging needs a deterministic total order, and the + aggregator's fan-out order is not stable between calls, so Modern list results + are sorted by the item's key (`Tool.Name`, `Resource.URI`, + `ResourceTemplate.URITemplate`, `Prompt.Name`). Legacy ordering is unchanged. +- **Duplicate keys are tolerated, not assumed away.** The aggregator makes all + four keys unique (see [Conflict Resolution](#conflict-resolution)), but the + paginator is generic and does not depend on that caller invariant. The cursor + resumes *within* a run of equal keys, because a plain "resume after this key" + scan skipped every copy and permanently dropped items whose key collided at a + page boundary. +- **End of results omits `nextCursor` entirely.** The draft states that "an empty + string is a valid cursor and thus MUST NOT be treated as the end of results", so + emitting `""` would make a conformant client re-request and loop on page one. + +The cursor encodes a position in the **aggregated** ordering and never names a +backend, so adding or removing a backend cannot invalidate one. This is unrelated +to the aggregator's *upstream* cursor-following (#5851), which is vMCP acting as a +client walking a backend's pages. + +An invalid cursor — malformed, over-length, or minted for a different list verb — +is rejected with `-32602`, per the draft's error-handling rule. + +### Tools/resources/prompts list_changed propagation (#5748, #5969) + +Unlike the per-call backend client (`pkg/vmcp/client`), the **persistent** +per-session backend connection (`pkg/vmcp/session/internal/backend`) stays +open for the session's lifetime, so it can observe a backend's asynchronous +(out-of-band) notifications rather than only mid-call traffic. When a +non-nil `ListChangedSink` is supplied at connection time, the connector: + +- Enables `WithContinuousListening()` on streamable-HTTP backends (opens a + standalone GET stream) — gated **strictly** on a non-nil sink, because some + backends hang when this stream is opened against them. SSE backends need no + extra option: their whole session is already one continuous stream. +- Registers an `OnNotification` handler that dispatches + `notifications/tools/list_changed` (`kind=KindTools`), + `notifications/resources/list_changed` (`kind=KindResources` — also covers + resource templates, since MCP 2025-11-25 has no separate wire method for + template changes), and `notifications/prompts/list_changed` + (`kind=KindPrompts`) to the sink — `ChangeKind` is a typed constant rather + than a bare string so a typo is a compile error — and logs + `notifications/message` received out-of-call (log-only; no relay). + +This propagation mechanism only applies to a Legacy (2025-11-25) backend: the +per-session factory (`pkg/vmcp/session/factory.go`) skips the persistent +connection — and therefore the handshake, the sink registration, and the +standalone GET stream above — for any backend whose cached revision is known +Modern (2026-07-28); see [Backend MCP Revision Classification](#backend-mcp-revision-classification) +above for how that revision is resolved and cached. This is correct, not a gap in the skip: Modern removed +`initialize` and `Mcp-Session-Id`, so there is no Legacy-shaped persistent +connection to hold and no standalone GET stream a Modern backend could push +on. Modern's own server-push mechanism is `subscriptions/listen` (see +[Transport Architecture](03-transport-architecture.md)). vMCP **serves** that +method on its client edge (`pkg/vmcp/server/modern_subscriptions.go`), but only +at acknowledgement level: it computes the honored subscription set by +intersecting the client's request against the capabilities `server/discover` +advertises, and since every push-related flag there is deliberately false +(`newModernCapabilities`), the honored set is always empty and no notification is +ever pushed. So `list_changed` **delivery** remains Legacy-only, on both edges — +what the Modern client edge gained is a conformant, explicitly-empty answer +instead of a `-32601` that tore the client's connection down. Real Modern +delivery is tracked in #5743 and requires vMCP to start advertising a push +capability first. + +The sink is built once per session, at registration (`pkg/vmcp/server`'s +`buildListChangedSink`), closing over the SDK `ClientSession`, the session ID, +and the caller's identity **and per-request forwarded headers** captured **at +registration time** (a deliberate snapshot, not re-resolved per firing — see +the token-staleness note below). It is threaded through +`MultiSessionFactory.MakeSessionWithID` / `SessionManager.CreateSession` (a +single nilable `ListChangedSink` parameter) down to every backend connector +opened for that session. + +`buildListChangedSink` builds **one coalescing worker per `ChangeKind`** +(`KindTools`, `KindResources`, `KindPrompts`) rather than a single shared +worker: a tools notification resyncs only the session's tool store, without +forcing a redundant re-apply (and possible spurious downstream +`list_changed`) of its resources or prompts overlay, and vice versa. + +**The sink is non-blocking (runs on the backend receive-loop goroutine).** It +must not do real work inline: the mcpcompat client dispatches notifications +synchronously on its receive loop, so a blocking sink would stall that +backend's notification delivery and let a misbehaving backend amplify one +notification into unbounded work. The sink therefore only hands off to the +matching **per-(session, kind) coalescing worker** (`listChangedResyncWorker`): +`trigger()` sets a dirty flag and, at most, starts **one** worker goroutine — +never one goroutine per notification. At most one resync runs at a time per +(session, kind); notifications that arrive while a resync is in flight +collapse into a single follow-up run, so a notification storm is bounded to +O(1) concurrent work per (session, kind) — up to 3 concurrent resyncs per +session in the worst case, one per kind. Each worker goroutine exits when +idle, so an idle session holds no goroutine. + +Each worker's resync body (`runListChangedResync`), off the receive loop: + +1. **Liveness guard** — looks the session up via `GetMultiSession` and returns + immediately if it was terminated/expired, so a storm of notifications for a + dead session drives no work. +2. **Reconstructs the request context** — starts from a server-lifetime base + context (`Server.resyncBaseCtx`, cancelled on `Stop` so in-flight backend + sweeps never outlive the server) and layers on the captured identity + (`auth.WithIdentity`) and forwarded headers + (`headerforward.WithForwardedHeaders`). This is **security-critical**: the + capability cache key and the outbound backend authentication are derived + from the **context** (not from an explicit identity argument), so without + this the resync would enumerate backends *unauthenticated* — advertising + metadata the principal's own credentials would not surface, or wrongly + dropping credential-gated tools/resources/prompts, while replacing the + correctly-scoped registration-time set. +3. Calls `core.InvalidateCapabilityCache()`, which purges the **entire** + per-identity capability cache (`aggregator.CacheInvalidator`), covering all + capability kinds together (tools, resources, resource templates, prompts + share one cached `AggregatedCapabilities` per identity). This is coarse (not + scoped to the one backend that changed): the LRU key is a hash of identity + + forwarded headers + backend-set, so per-backend invalidation would need a + reverse index that is disproportionate here. The coalescing above already + bounds the purge to at most once per resync burst per kind, and a purge + only forces the next call per identity to re-sweep — an accepted de-opt / + follow-up. An aggregator that does not implement `CacheInvalidator` + WARN-logs instead of silently no-opping. +4. Re-derives and **replaces** — not merges — the session's advertised set for + the notification's kind, from the (now cold) cache: `KindTools` re-derives + via `serveSessionTools` and replaces the tool store + (`SessionWithTools.SetSessionTools`); `KindResources` re-derives via + `coreSessionResources`/`coreSessionResourceTemplates` and replaces BOTH the + resource store (`SessionWithResources.SetSessionResources`) and the + resource-template store + (`SessionWithResourceTemplates.SetSessionResourceTemplates`) — one + notification method covers both, per MCP 2025-11-25; `KindPrompts` + re-derives via `coreSessionPrompts` and replaces the prompt store + (`SessionWithPrompts.SetSessionPrompts`). Replacing (rather than merging) + means a tool the backend removed disappears rather than lingering from the + registration-time merge — but see the add-only caveat below for resources + and prompts. + +The go-sdk server auto-emits the corresponding `list_changed` notification +(`notifications/tools/list_changed`, `notifications/resources/list_changed`, +or `notifications/prompts/list_changed`) to the downstream client whenever the +matching `SetSession*` call changes something, now that +`WithToolCapabilities(true)`, `WithResourceCapabilities(true, true)`, and +`WithPromptCapabilities(true)` are all set (`pkg/vmcp/server/serve.go`). + +**Identity staleness (B1)**: each worker reuses the identity captured at +registration for its re-aggregation and admission view. If that identity's +upstream tokens are refreshed later (see #5323), the resync's +`core.ListTools`/`ListResources`/`ListPrompts` call authorizes/aggregates +against the (potentially stale) captured tokens, not the live per-request +ones. This only affects the *accuracy of the asynchronous resync's own view* +— every live call still authenticates via the fresh per-request identity +through `enforceSessionBinding`, so staleness here cannot grant a call that +would otherwise be denied. + +**Registration-time emission (R1)**: go-sdk's own `AddTool`/`RemoveTools` (and +the resource/prompt equivalents) debounce (10ms) and then broadcast the +corresponding `list_changed` notification to **every** session currently +connected to the shared `*gosdk.Server` — not only the session whose overlay +changed — and a session is eligible for that broadcast from the moment its +transport connects (`Server.bind`), before its own +`initialize`/`notifications/initialized` round-trip completes. go-sdk's own +source documents this as a known upstream gap ("potential spec violation... +when the feature list changes before the session ... is initialized"). This +means every session's registration-time per-session `AddTool`/resource/prompt +calls (`setSessionToolsDirect`/`setSessionResourcesDirect`/ +`setSessionResourceTemplatesDirect`/`setSessionPromptsDirect`) now cause a +broadcast to every other already-connected session too. There is no supported +hook to scope or suppress this without patching the vendored SDK, so it is +accepted as a benign nuisance for all three capability kinds: MCP +notifications are inherently a "you may want to refetch" hint, so an extra one +only costs an idempotent list round-trip — it never changes what any given +session's own list actually returns. + +**Scope**: additions AND removals propagate for tools. For resources +(including resource templates) and prompts, only **additions** propagate +today: toolhive-core's mcpcompat per-session sync for resources/resource +templates/prompts (`syncSessionResources`/`syncSessionResourceTemplates`/ +`syncSessionPrompts` in `mcpcompat/server/session.go`) is add-only — unlike +`syncSessionTools`, which also calls `RemoveTools` — so a backend-removed +resource/template/prompt stays registered on the session's go-sdk server +(listed until re-initialize) even though the overlay this PR replaces no +longer contains it. `resyncSessionResources`/`resyncSessionPrompts` still +REPLACE (not merge) their overlay, so once toolhive-core gains removal +reconciliation ([stacklok/toolhive-core#184](https://github.com/stacklok/toolhive-core/issues/184)), +removals start propagating with no change needed on the vMCP side. Advertising +`listChanged: true` stays honest in the meantime — notifications ARE emitted on +change (except that a pure removal which EMPTIES the overlay issues no `Add*` +and therefore triggers no downstream `list_changed`, since go-sdk notifies only +via `Add*`/`RemoveTools`); `listChanged` promises notification, not list +minimization — but until that follow-up lands, a refetch after a pure removal +returns a stale superset for resources/templates/prompts. Cross-pod session +restore (`RestoreSession`) does not thread a sink at all for any kind (no live +`ClientSession` to resync there). + +**Implementation**: `pkg/vmcp/session/internal/backend/mcp_session.go` +(`ChangeKind`/`KindTools`/`KindResources`/`KindPrompts`, `ListChangedSink`, +connector wiring), `pkg/vmcp/aggregator/aggregator.go` and +`caching_aggregator.go` (`CacheInvalidator`), `pkg/vmcp/core/core_vmcp.go` +(`InvalidateCapabilityCache`), `pkg/vmcp/server/serve_list_changed.go` +(`listChangedResyncWorker` coalescing, `buildListChangedSink`, +`runListChangedResync`, `resyncSessionTools`, `resyncSessionResources`, +`resyncSessionPrompts`) with `Server.resyncBaseCtx` cancelled on `Stop`. + +### Mid-call forwarding (elicitation / sampling / progress / logging) + +While a backend `tools/call` (or other request) is in flight, the backend may issue **server-initiated** requests and notifications back toward the client: elicitation, sampling, progress, and logging. vMCP forwards these mid-call in both directions through a per-call forwarder that bridges the backend connection to the originating client session, so a backend that needs user input (elicitation) or model completions (sampling), or that emits progress/log notifications, reaches the real client transparently. This is distinct from composite-tool elicitation (which the composer drives during a workflow); the mid-call forwarder handles the general request-scoped case for a single backend call. + +**Implementation**: `pkg/vmcp/forwarding.go`, `pkg/vmcp/client/forwarding.go`, `pkg/vmcp/server/serve_handlers.go` + +**Known limitation (Modern clients)**: everything in this section describes a +**Legacy (2025-11-25) client session**. For Modern (2026-07-28) clients there +is no session and no server-initiated request channel, so none of this +forwarding applies — see +[Limitation: elicitation and sampling are unavailable to Modern clients](#limitation-elicitation-and-sampling-are-unavailable-to-modern-clients) +for what a Modern caller gets instead. + +**Known limitation (logging level)**: forwarded backend logging is not yet filtered to the downstream client's requested level. On Legacy, vMCP requests debug-level logging from the backend (`logging/setLevel`) so it emits `notifications/message`, and every such notification is forwarded — the downstream client's own `logging/setLevel` preference is not applied to the relayed stream. The same is true on Modern (2026-07-28), where the RPC is removed and the level rides per-request in `_meta["io.modelcontextprotocol/logLevel"]`: vMCP strips that reserved per-hop key from the downstream request and overlays its own (`debug`, when forwarding is bound) on the backend hop, so a Modern client's per-request level preference is likewise not honored — the relay runs at debug either way. + +**Known limitation (advertised-but-no-stream elicitation fails fast)**: a client +that advertised the `elicitation` capability but holds **no open standalone SSE +stream** passes go-sdk's capability gate, yet the elicitation cannot be +delivered — under the shim's `JSONResponse` transport the go-sdk routes +server→client requests to the standalone stream, and a missing stream rejects +the write ("rejected by transport: stream not connected or already closed"). +The mid-call `tools/call` therefore fails fast with a tool error instead of +hanging to the deadline (pinned by +`TestForwarding_Elicitation_AdvertisedButNoStream_FastFails` in +`pkg/vmcp/server`). A cleaner pre-dispatch refusal awaits an upstream mcpcompat +accessor for stream presence (#5975). + +**Known limitation (cross-pod origination needs session affinity)**: a +server→client request can only be delivered by the replica currently holding +the client's standalone SSE stream. If the `tools/call` executes on replica A +but the client's GET stream is pinned to replica B, an elicitation or sampling +request originated from A cannot reach the client — the shim loads the go-sdk +session bound to *its* pod and has no cross-replica delivery channel for +request/response (only notifications rehydrate cross-replica). Multi-replica +deployments that rely on mid-call elicitation/sampling therefore need **session +affinity** at the load balancer pinning the standalone stream and the tool +calls to the same replica. The durable fix is the 2026-07-28 revision itself: +it replaces server-initiated requests with client-polled MRTR, which has no +stream-locality requirement — so this constraint is documented rather than +engineered around (#5975, #5743). + +**Known limitation (resource-template authorization)**: a resource template is advertised on the template-string entity (e.g. `file:///logs/{date}.txt`), but a concrete read is admission-checked on the **expanded** URI (e.g. `file:///logs/2025-01-01.txt`). Operators should therefore author resource authorization policies against concrete URI patterns, not the template string. ## Two-Boundary Authentication @@ -239,14 +917,17 @@ Middleware is applied by wrapping handlers, so execution order is outer-to-inner | 1 | Recovery | Always | Catches panics, returns HTTP 500 | | 2 | WriteTimeout | Always | Clears the server `WriteTimeout` for qualifying SSE connections | | 3 | Header Validation | Always | Rejects GETs without `Accept: text/event-stream` before they reach the MCP handler | -| 4 | Authentication (+ MCP parsing) | Optional | Validates incoming credentials (OIDC/local/anonymous); MCP parsing is composed inside so downstream layers see `ParsedMCPRequest` | -| 5 | Audit | Optional | Logs request events for compliance | +| 4 | Audit | Optional | Logs every request outcome, including 401s from the auth middleware it wraps; identity and parsed MCP data flow back via the `auth.IdentityHolder` / `mcp.ParsedRequestHolder` carriers | +| 5 | Authentication (+ MCP parsing) | Optional | Validates incoming credentials (OIDC/local/anonymous); MCP parsing is composed inside so downstream layers see `ParsedMCPRequest` | | 6 | Discovery | Always | Aggregates backend capabilities per session | | 7 | Annotation Enrichment | Optional | Injects tool annotations into context for annotation-aware authz (only when Authorization is configured) | | 8 | Authorization | Optional | Evaluates Cedar policies after discovery and annotation enrichment | | 9 | Backend Enrichment | Optional | Adds backend name to audit context (only when Audit is configured) | | 10 | MCP Parsing | Always | Second application is a no-op when auth already parsed; ensures telemetry can label metrics with `mcp_method` when auth is nil | | 11 | Telemetry | Optional | OpenTelemetry instrumentation | +| 12 | Pre-dispatch authorization gate | Optional | Innermost: runs inside the Streamable HTTP transport before session validation and SDK dispatch. Rejects a Cedar-denied `tools/call` / `resources/read` / `prompts/get` with HTTP 403 + JSON-RPC code 403, reusing the core admission decision. Installed only when Authorization is configured. See "Authorization Enforcement" below. | + +> On the New/Serve path, authorization is enforced by the **core admission seam**, not by rows 6–8 as standalone HTTP middleware; row 12 is the transport-level projection of that decision. See [Authorization Enforcement](#authorization-enforcement-core-admission-seam--pre-dispatch-gate). ### Discovery Middleware @@ -284,11 +965,65 @@ This enriches audit events with the backend name for better observability. The server wires them around discovery/annotation-enrichment so the effective execution order is: ``` -Authentication → MCP Parsing → Audit → Discovery → Annotation Enrichment → Authorization → Next Handler +Audit → Authentication → MCP Parsing → Discovery → Annotation Enrichment → Authorization → Next Handler ``` **Implementation**: `pkg/vmcp/server/server.go`, `pkg/vmcp/discovery/middleware.go`, `pkg/vmcp/auth/factory/` +### Authorization Enforcement (core admission seam + pre-dispatch gate) + +On the New/Serve path, authorization is enforced by the **core admission seam** +(`pkg/vmcp/core`), not by HTTP middleware. The seam applies one Cedar decision to both +the list side (`ListTools`/`ListResources`/`ListPrompts` filter the advertised set) and +the call side (`CallTool`/`ReadResource`/`GetPrompt` deny before dispatch), closing the +"list says yes / call says no" gap. + +Because the SDK maps a call-side deny to a tool result, a raw denied `tools/call` would +otherwise return **HTTP 200** (either the SDK's `-32602 "not found"` for a list-filtered +tool, or a `200 + IsError` tool result for an argument-gated deny). To make a denial a +first-class wire rejection, Serve installs a **pre-dispatch authorization gate** +(`pkg/vmcp/server/call_gate.go`) on the Streamable HTTP transport, but only when Cedar +policies are configured: + +- The gate re-runs the core admission decision for `tools/call`, `resources/read`, and + `prompts/get` via `core.CheckToolCall` / `CheckResourceRead` / `CheckPromptGet` — the + same helpers the call path uses, so a pre-check and the call can never drift. Non-gated + methods (e.g. `initialize`, `tools/list`) are admitted untouched. +- A denial is rejected as **HTTP 403 + JSON-RPC error code 403** (`pkg/mcp.JSONRPCCodeDenied`) + with a kind-only message (`"call denied by authorization policy"`, + `"read denied by authorization policy"`, `"prompt denied by authorization policy"`) — + identical to the single-server `thv run` authorization response. The message never + names the capability or reveals advertised-vs-nonexistent, so a denial is not an + **enumeration oracle**: a filtered tool, an argument-gated deny, and a nonexistent tool + under a default-deny policy all converge on the same 403. +- The gate runs **before session validation** (403-before-404): a denial is determinable + from the caller's own identity without session state. +- It sits **inside the audit middleware**, so a denied call is audited with outcome + `denied` (403 → `OutcomeDenied`) with no audit-layer changes. +- An authorizer error fails **closed** (treated as a denial); a non-authorization + (infrastructure) error admits, so the call path surfaces it through existing mapping — + the gate never converts a plumbing fault into a 403. +- **One decode per `tools/call`**: dispatch (`coreToolHandler`) prefers the transport + parse (`pkg/mcp`) the gate authorized on — via `gateParsedArgs`, keyed on matching + method + tool — so the gated decision, the enforced call-path decision, and the + forwarded backend arguments all derive from a single decoded map. Where no matching + parse exists (batch, embedders bypassing the transport, method/tool mismatch), dispatch + falls back to the SDK decode and makes a single decision on that single map, so no path + can produce an allow-then-deny split between gate and call. +- **Code mode carve-out**: `execute_tool_script` is not in the admission seam (the feature + flag is the grant, and each inner tool call the script makes is re-authorized by its real + name), so the codemode decorator's `CheckToolCall` admits it while delegating every other + name to the inner core. A backend that advertised a tool named `execute_tool_script` + would be silently shadowed by the virtual tool and skip its own Cedar admission, so the + decorator fails **loud** (`ErrReservedToolName`) on that collision — `ListTools`, + `LookupTool`, and the `CallTool` script-binding path all refuse to serve rather than mask it. + +The `Call*` methods keep their internal admission checks as defense-in-depth for other +embedders and misconfigured gates. + +**Implementation**: `pkg/vmcp/core/core_checks.go`, `pkg/vmcp/server/call_gate.go`, +`pkg/vmcp/server/serve_handlers.go`, `pkg/vmcp/codemode/decorator.go`, `pkg/mcp/errors.go` + ## Health Monitoring vMCP monitors backend health with configurable intervals. Health status (healthy, degraded, unhealthy, unauthenticated, unknown) affects routing decisions and is reported in VirtualMCPServer status. diff --git a/docs/arch/12-skills-system.md b/docs/arch/12-skills-system.md index 83e16b8409..9e6786e4ee 100644 --- a/docs/arch/12-skills-system.md +++ b/docs/arch/12-skills-system.md @@ -308,6 +308,7 @@ installed_skills table ├── client_apps (BLOB, JSONB-encoded []string) ├── status (installed | pending | failed) ├── installed_at (TEXT, ISO 8601) +├── managed (INTEGER, 0/1 — tracked in the project's toolhive.lock.yaml; see below) └── UNIQUE(entry_id, scope, project_root) skill_dependencies table @@ -322,7 +323,88 @@ oci_tags table (reserved; not currently populated) └── digest (TEXT NOT NULL — content digest) ``` -**Implementation:** `pkg/storage/sqlite/skill_store.go`, `pkg/storage/interfaces.go` (SkillStore), `pkg/storage/sqlite/migrations/001_create_entries_and_skills.sql` +**Implementation:** `pkg/storage/sqlite/skill_store.go`, `pkg/storage/interfaces.go` (SkillStore), `pkg/storage/sqlite/migrations/001_create_entries_and_skills.sql`, `003_add_managed_flag.sql` + +## Project Lock File + +RFC [THV-0080](https://github.com/stacklok/toolhive-rfcs/blob/main/rfcs/THV-0080-skills-lock-file.md) adds a project-level `toolhive.lock.yaml`, committed at the project root, that pins the exact content of every project-scoped skill install — the same guarantee `package-lock.json`, `Cargo.lock`, and `go.sum` provide elsewhere. Two teammates (or a CI runner) cloning the same repo restore identical skill content via `thv skill sync`, rather than whatever the source currently resolves to. + +**Trust model, stated plainly:** until the RFC's Sigstore signing/verification stack lands, the lock file provides *reproducibility and drift detection over a repository-editable file* — not verified integrity. A committed digest is an unauthenticated trust root: anyone who can change `toolhive.lock.yaml` in the repository controls what `sync` installs. Reviewing lock-file diffs (especially `digest` and `resolvedReference` changes) carries the same weight as reviewing the AI-executed skill content itself. Signature verification on consume is what upgrades this from drift detection to integrity, and is why the whole feature ships gated until that half exists. + +### Rollout + +The feature is gated behind the `TOOLHIVE_SKILLS_LOCK_ENABLED` environment variable (`skills.LockFileFeatureEnabled()`) while it lands across a stack of PRs, following the existing `TOOLHIVE_DEV` precedent for staged rollouts (`pkg/skills/gitresolver/reference.go`). With the flag unset, project-scope installs behave exactly as they did before this RFC — no lock file is written, no `toolhive.requires` materialization happens, `sync`/`upgrade` refuse with a clear "experimental" error. + +### Schema + +```yaml +version: 1 +skills: + - name: code-review + version: "1.0.0" + source: code-review # exactly what the user/registry resolver requested; never rewritten + resolvedReference: ghcr.io/org/code-review:1.0.0 + digest: sha256:9f2b1e... # the pin: OCI manifest digest or git commit hash + contentDigest: sha256:a1b2c3d4... # deterministic dirhash of the materialized files, for on-disk integrity + requiredBy: [parent-skill] # present only for transitively materialized toolhive.requires deps + explicit: true # false for a dependency that was never directly installed by name +``` + +Entries are sorted by name for stable diffs. `source` is never rewritten by `sync` or `upgrade` — only `upgrade` re-resolves it, and only to decide whether newer content exists. + +**Implementation:** `pkg/skills/lockfile/` (schema, load/save, atomic writes via `os.Root`, `contentDigest` dirhash algorithm) + +### Install and Uninstall Hooks + +For project-scope installs (with the feature enabled), `skillsvc.Install`'s single existing choke point (`installAndRegister` — every dispatch path, OCI or git, direct or registry-resolved, converges there) additionally: + +1. Computes `contentDigest` from the extracted files. +2. Materializes `toolhive.requires` dependencies recursively — reading `SKILL.md` back from disk (not from the resolver's own parse, so this works uniformly across OCI and git sources), with a `Visited` set guarding cycles and `skills.MaxDependencies` bounding the whole tree. +3. Upserts the lock entry, merging `requiredBy` when a dependency is shared by multiple parents. + +**A lock-write failure fails the entire install** (rolling back the DB record, matching the existing group-registration failure/rollback pattern) — RFC THV-0080 treats "installed but unpinned" as worse than "not installed." + +`Uninstall` mirrors this: for a lock-managed skill, it removes the skill's own lock entry and cascades to any dependency that consequently loses its last requiring parent (and is not itself `explicit`), via `Lockfile.RemoveParentFromRequiredBy` — itself cycle-safe. + +**Implementation:** `pkg/skills/skillsvc/lock.go`, `pkg/skills/skillsvc/install.go`, `pkg/skills/skillsvc/uninstall.go` + +### Sync + +`thv skill sync` reconciles installed skills against the lock file: + +| Lock file vs. installed state | Outcome | +|---|---| +| Digest and contentDigest both match | Reported as up to date | +| DB record exists but digest or contentDigest differs | Drifted — reinstalled at the pinned reference (`--check`: reported only, nothing written) | +| No DB record for a locked entry | Missing — installed fresh at the pinned reference | +| Installed, `managed`, but no lock entry | Removed from lock — reported (`--prune`: uninstalled) | +| Installed, not `managed`, no lock entry | Never managed — reported (`--adopt`: lock entry written from current state) | + +Reinstalling *at the pinned reference* (never re-resolving `source`) uses `buildPinnedReference` (`pkg/skills/skillsvc/pin.go`) to rewrite the entry's `resolvedReference` with its pinned digest substituted in (an OCI digest reference, or a git reference with the commit hash spliced in as the ref). A subtlety this surfaced: reinstalling at an *unchanged* digest — the normal case when repairing on-disk drift — would otherwise hit the install path's "same digest means content is already correct" fast path and silently skip re-extraction. An internal `InstallOptions.SyncRestore` flag bypasses that fast path specifically for this case. + +**Implementation:** `pkg/skills/skillsvc/sync.go`, `pkg/skills/skillsvc/pin.go` + +### Upgrade + +`thv skill upgrade [name...]` re-resolves each targeted entry's `source` (via the same git/OCI/registry dispatch order `Install` uses, stopping short of extraction — `resolveLatestState`) and installs newer content when the resolved digest changed. Entries pinned to an immutable reference (an OCI digest, or a git reference already pinned to a full commit hash — `isImmutableSource`) are reported not-upgradable. `--preview` reports what would change without persisting it (an OCI preview still pulls the artifact into the local store — there's no lighter "digest only" primitive — so it is not fully side-effect-free, matching the RFC's own caveat). `--allow-ref-change` permits the resolved reference itself changing; `--fail-on-changes` gives CI a freshness gate — it evaluates the same plan, never installs, and returns the full outcome set (exit codes are derived client-side from the outcomes, so a genuine resolution failure still surfaces as a partial failure rather than "lock is stale"). + +**Implementation:** `pkg/skills/skillsvc/upgrade.go` + +### CLI Confirmation and Exit Codes + +Because skill content is a set of AI-followed instructions, `sync` and `upgrade` gate real installs behind a confirmation prompt (skipped by `--check`/`--preview`/`--fail-on-changes`, which never write to the lock file or extracted skill directories — an OCI `--preview` still pulls the artifact to compare digests, but persists nothing). The prompt is printed to stderr together with a summary of the lock entries being acted on (name, source, short digest) so the human gate has something concrete to judge, and everything echoed into it is stripped of non-graphic characters so a hostile directory or entry name cannot repaint the prompt with terminal escapes. On a non-interactive terminal without `--yes`, the command refuses outright rather than silently proceeding. Until Sigstore verification lands, this prompt is a speed bump, not a security boundary — see the trust-model note above. + +Exit codes follow a CI-oriented contract distinct from the generic `1` used elsewhere in the CLI: + +| Code | Meaning | +|---|---| +| `0` | Success | +| `1` | Generic/unclassified error (cobra's default) | +| `2` | Check/freshness failure — `sync --check` found drifted or missing installs, or `upgrade --fail-on-changes` found available changes | +| `3` | Partial failure — some, but not all, targeted skills failed | +| `4` | Policy rejection — the confirmation gate declined a non-interactive run, or `--allow-ref-change` blocked a reference change | + +**Implementation:** `cmd/thv/app/exitcode.go`, `cmd/thv/app/skill_confirm.go`, `cmd/thv/app/skill_sync.go`, `cmd/thv/app/skill_upgrade.go` ## API @@ -342,6 +424,10 @@ oci_tags table (reserved; not currently populated) | `GET` | `/builds` | List local builds | | `DELETE` | `/builds/{tag}` | Delete a local build | | `GET` | `/content` | Get a skill's SKILL.md body and file listing for a reference | +| `POST` | `/sync` | Restore project skills to match the lock file ([Project Lock File](#project-lock-file), RFC THV-0080) | +| `POST` | `/upgrade` | Re-resolve project skills and install newer pinned content ([Project Lock File](#project-lock-file), RFC THV-0080) | + +`/sync` and `/upgrade` are served only when the configured `SkillService` also implements `skills.SkillLockService` (as `skillsvc.New`'s does) — otherwise they return `501 Not Implemented`. **Implementation:** `pkg/api/v1/skills.go` @@ -366,7 +452,9 @@ thv skill ├── build [path] Build skill to OCI artifact ├── push [reference] Push built skill to registry ├── builds List locally-built OCI artifacts -└── builds remove [tag] Delete a locally-built artifact +├── builds remove [tag] Delete a locally-built artifact +├── sync Restore project skills to match the lock file +└── upgrade [name...] Re-resolve project skills and install newer pinned content ``` **Implementation:** `cmd/thv/app/skill*.go` @@ -455,6 +543,12 @@ ToolHive owns the installation lifecycle, scoping model, CLI/API interfaces, and | HTTP client | `pkg/skills/client/` | | CLI commands | `cmd/thv/app/skill*.go` | | Group integration | `pkg/groups/skills.go` | +| Lock file schema | `pkg/skills/lockfile/` | +| Lock file rollout gate | `pkg/skills/feature_gate.go` | +| Install/uninstall lock hooks | `pkg/skills/skillsvc/lock.go` | +| Sync | `pkg/skills/skillsvc/sync.go`, `pkg/skills/skillsvc/pin.go` | +| Upgrade | `pkg/skills/skillsvc/upgrade.go` | +| CLI exit codes / confirmation | `cmd/thv/app/exitcode.go`, `cmd/thv/app/skill_confirm.go` | ## Related Documentation diff --git a/docs/arch/14-plugins-system.md b/docs/arch/14-plugins-system.md index 82c3d9a1ae..5f5dcf5a74 100644 --- a/docs/arch/14-plugins-system.md +++ b/docs/arch/14-plugins-system.md @@ -46,7 +46,9 @@ different component subsets. + SQLite storage + migration 002. - **Phase 3** (this doc): install/list/info/uninstall + MaterializationAdapter (Claude Code + Codex) + groups integration. -- **Phase 4** (planned, #5528): REST API + `thv plugin` CLI + app wiring. +- **Phase 4** (shipped, #5528): REST API + `thv ai-plugin` CLI + app wiring. +- **Phase 5b** (shipped, #5529): registry-name install wired via `lazyPluginLookup`; + plain-name resolution through the registry lookup chain. ## Core Concepts @@ -86,7 +88,8 @@ reporting which component types it deliberately dropped. OCI reference (`ghcr.io/org/plugin:v1`), git reference (`git://host/owner/repo[@ref][#subdir]`), or plain name (registry lookup — -seam declared, wired in a later phase). +wired via `lazyPluginLookup`; `thv ai-plugin install ` +resolves through the registry's `SearchPlugins` → first hit's OCI package). ### 2. Building @@ -267,6 +270,8 @@ The OCI plugin layer (`ociplugins.Store`, `PluginPackager`, `RegistryClient`, | `pkg/plugins/pluginsvc/list.go` / `info.go` / `uninstall.go` | lifecycle methods | | `pkg/plugins/adapters/claudecode.go` | Claude Code adapter (FS + settings.json mutation) | | `pkg/plugins/adapters/codex.go` | Codex adapter (FS + TOML) | +| `pkg/registry/api/plugins_client.go` | MCP v0.1 registry plugin search client | +| `pkg/api/v1/registry_v01_plugins.go` | HTTP handler for plugin search endpoint | | `pkg/storage/sqlite/plugin_store.go` | SQLite store | | `pkg/groups/plugins.go` | group membership | | `pkg/client/plugins.go` | client metadata + path resolution | diff --git a/docs/arch/15-envoy-network-proxy.md b/docs/arch/15-envoy-network-proxy.md new file mode 100644 index 0000000000..0117d216c3 --- /dev/null +++ b/docs/arch/15-envoy-network-proxy.md @@ -0,0 +1,271 @@ +# Envoy Network Proxy + +## Status + +Experimental — selected with `TOOLHIVE_NETWORK_PROXY=envoy`. Squid remains the +default. + +## Problem + +When network isolation is enabled (`--isolate-network`), ToolHive currently +starts **three** auxiliary containers per workload: + +| Container | Role | +|-----------|------| +| `-egress` | Squid forward proxy — routes outbound traffic through an allowlist | +| `-ingress` | Squid reverse proxy — receives traffic from the proxy runner | +| `-dns` | dnsmasq — provides DNS to the internal network | + +Three containers means three image pulls, three startup sequences, three sets of +resources, and three things that can fail or restart. The Squid egress and ingress +containers are logically a single gateway — splitting them into two processes is +an implementation artifact rather than a deliberate design. + +## Solution + +Replace the two Squid containers with a **single Envoy container** that handles +both egress and ingress as separate listeners inside the same process. The DNS +container (dnsmasq) is unchanged. + +``` +Before: -egress (Squid) + -ingress (Squid) + -dns +After: -egress (Envoy, two listeners) + -dns +``` + +This reduces auxiliary container count from 3 → 2, simplifies the startup +sequence, and uses a single bootstrap configuration file to describe the entire +proxy behaviour. + +## Why Envoy + +### Consolidation + +Envoy's `HttpConnectionManager` supports multiple listeners in a single process. +The egress forward proxy (`:3128`) and ingress reverse proxy share the same Envoy +instance, the same access logs, and the same lifecycle. + +### L7 authority-based enforcement + +Squid matches destination hostnames via `dstdomain` ACLs and destination IPs via +`dst` ACLs. Envoy, running as a forward proxy, enforces the equivalent controls +at L7 by matching the HTTP `:authority` header. + +Envoy's RBAC filter is used for both allow and deny decisions: +- **Header match on `:authority`** — L7 match on the forward-proxy target for + both plain HTTP requests and HTTPS CONNECT tunnels. This is the authority the + client asks the proxy to reach, so a single match covers hostnames *and* IP + literals presented as the CONNECT target. + +Note that RBAC `destination_ip` is **not** usable for target filtering in a +forward proxy: `destination_ip` resolves to Envoy's own listener socket (the +address the client connected to), not the proxied target, so a CIDR match on it +is inert. All target enforcement therefore happens through `:authority` regexes, +which mirror Squid's `dstdomain`/`dst` denies. The `Internal: true` Docker +network remains the fail-closed backstop for non-cooperative traffic that ignores +the proxy entirely. + +### Proper dynamic forward proxy + +Envoy's `dynamic_forward_proxy` cluster performs per-request DNS resolution and +handles HTTP CONNECT tunnelling natively. HTTPS flows through a CONNECT tunnel +exactly as a client would expect, with Envoy acting as a transparent TCP relay +after the CONNECT handshake — no TLS inspection, no certificate pinning, no CA +changes. + +### Access logging + +Both listeners write structured access logs to stdout, visible via `docker logs`. +Squid logged differently for egress and ingress with no unified view. + +### Configuration as code + +Envoy reads a protobuf-JSON bootstrap file generated by ToolHive at workload +start. The configuration is typed Go structs serialised to JSON — unit-testable, +diffable, and reproducible. Squid required template-rendered text files. + +### Future extensibility + +Envoy's xDS API makes it possible to update listeners, clusters, and RBAC +policies at runtime without a container restart. This is not used today, but the +groundwork is there for dynamic policy updates. The transparent L3/L4 +interception path (Phase 2, not yet implemented) requires an `original_dst` +listener and iptables rules that Envoy handles cleanly. + +## What Envoy Does Not Do + +- **Decrypt TLS.** Like Squid, Envoy filters HTTPS on the CONNECT target hostname + and then relays the encrypted stream as-is. No certificate inspection, no + man-in-the-middle. +- **Block non-cooperative traffic.** A workload that opens a raw TCP connection + ignoring `HTTP_PROXY` is contained by the `Internal: true` Docker network + blackhole, not by Envoy. Envoy only sees traffic that goes through the proxy. +- **Replace dnsmasq.** DNS for the workload's internal network is still served by + the dnsmasq container. +- **Run in Kubernetes.** Network isolation is a local-Docker feature only; the + Kubernetes operator has a separate egress gateway path. + +## Architecture + +### Egress listener (`:3128` — forward proxy) + +``` +HTTP_PROXY / HTTPS_PROXY → Envoy :3128 + └── HCM (upgrade: CONNECT) + ├── [optional] RBAC DENY — docker gateway (:authority: IP literal + hostnames) + ├── RBAC ALLOW — outbound allowlist (or allow-all) + ├── dynamic_forward_proxy — per-request DNS + CONNECT tunnel + └── router +``` + +The RBAC filters are evaluated top-to-bottom. The gateway DENY filter is present +unless `--allow-docker-gateway` is set. It blocks the Docker gateway purely at +L7, by matching the `:authority` header (the forward-proxy target for both plain +HTTP and HTTPS CONNECT) against anchored, case-insensitive, port-tolerant +regexes for: +- The resolved Docker bridge gateway IP literal (e.g. `172.17.0.1`, with an + optional `:port`) +- The Docker-internal hostnames `host.docker.internal` and + `gateway.docker.internal` (with an optional `:port`, so HTTPS CONNECT + authorities such as `host.docker.internal:443` match too) + +There is deliberately no `destination_ip`/L3 CIDR component: as noted above, a +forward proxy's `destination_ip` is Envoy's own listener, so it would never match +the gateway. Matching the IP literal in `:authority` is what actually catches a +client that connects directly to the gateway address, giving parity with Squid's +combined `dst`/`dstdomain` denies. + +The ALLOW filter implements the permission profile's `Outbound` rules: +- `InsecureAllowAll: true` → single wildcard policy (`any: true`) +- `AllowHost: [...]` → per-host `:authority` regex match that mirrors Squid's + `dstdomain` semantics: `example.com` matches that host exactly, while + `.example.com` (leading dot, also accepted as `*.example.com`) matches the apex + **and** all subdomains. The generated pattern is anchored, case-insensitive, + and tolerates an optional `:port` so HTTPS CONNECT authorities + (`example.com:443`) match too. +- No outbound permissions configured → empty policy map → Envoy deny-all + +This preserves parity with the existing Squid backend so that permission +profiles written for Squid's `dstdomain` syntax behave identically under Envoy. + +When `--allow-docker-gateway` is set, ToolHive not only omits the gateway DENY +filter but also injects explicit ALLOW policies for the gateway (the IP literal +and both Docker-internal hostnames, matched the same way as the deny). This is +required because the ALLOW filter is default-deny: under a non-permissive profile +(one without `InsecureAllowAll`), simply dropping the deny would leave the +gateway unreachable, since nothing in the allowlist would match it. The explicit +ALLOW entries ensure the gateway is actually permitted rather than merely +un-denied. + +### Ingress listener (`0.0.0.0:` — reverse proxy) + +``` +Proxy runner → host:127.0.0.1: → Docker port binding → Envoy :port + └── HCM + ├── router + └── route → STRICT_DNS cluster → : +``` + +The ingress listener binds to `0.0.0.0` inside the container so Docker's port +forwarding (which targets the container's bridge IP, not its loopback) can reach +it. The host-side port binding restricts to `127.0.0.1`, so the ingress is only +reachable from the local machine. + +The upstream STRICT_DNS cluster resolves the MCP container's hostname inside the +internal Docker network and forwards HTTP traffic to the MCP server port. + +The admin interface is intentionally omitted from the bootstrap configuration. +Envoy does not start an admin server when the field is absent, eliminating the +attack surface entirely. + +### Bootstrap lifecycle + +1. ToolHive generates a protobuf-JSON bootstrap file in `os.TempDir()` at mode + `0600`. +2. The file is bind-mounted read-only into the Envoy container at + `/etc/envoy/envoy.json`. +3. Envoy reads it once at startup. + +The bootstrap file is **not** removed when the workload is torn down — there is +no cleanup hook for it today. The only removal path is a best-effort cleanup if +generation fails partway through writing. As a result, a bootstrap file persists +in the OS temp directory (at mode `0600`) for the lifetime of the host until the +OS reclaims the temp directory. This is a known limitation (see below). + +## Selection + +```bash +TOOLHIVE_NETWORK_PROXY=envoy thv run --isolate-network +``` + +`TOOLHIVE_NETWORK_PROXY` accepts: +- `""` or `"squid"` — Squid backend (default) +- `"envoy"` — Envoy backend + +An unknown value causes `NewClient` to fail at startup with a descriptive error. +The env var is intentionally not exposed as a CLI flag or CRD field while the +backend is experimental; chart surface and `RunConfig` wiring come later once the +backend is stable. + +## Comparison with Squid + +| Aspect | Squid (current default) | Envoy | +|--------|------------------------|-------| +| Containers per workload | 3 (egress + ingress + dns) | 2 (combined + dns) | +| Forward proxy | ✓ | ✓ (dynamic_forward_proxy) | +| Reverse proxy (ingress) | ✓ (separate container) | ✓ (second listener, same container) | +| HTTPS CONNECT tunnelling | ✓ | ✓ | +| TLS inspection | ✗ | ✗ | +| L7 hostname deny | ✓ (`dstdomain`) | ✓ (`:authority` header match) | +| Destination IP deny | ✓ (`dst` ACL, resolved IP) | ~ (IP literal in `:authority` only; resolved IP not enforced — see Known Limitations) | +| Wildcard host allowlist | ✓ (`.example.com` dot-prefix) | ✓ (same syntax, regex match incl. apex + port) | +| Port-based allowlist | ✓ (`AllowPort` ACL, AND-d with host) | ✓ (combined host+port regex; bare hostname allowed when port 80 is listed, matching Squid's plain-HTTP behaviour) | +| Per-request DNS resolution | Via Squid resolver | Via DFP cluster | +| Access logs | Per-container, text format | Unified stdout, structured | +| Config format | Text template | Typed Go structs → protobuf-JSON | +| Runtime config update | Restart required | xDS-capable (not yet used) | +| Upstream image | Stacklok-built | Upstream distroless (pinned tag+digest) | + +## Known Limitations + +- **CONNECT access log timing.** Envoy logs CONNECT tunnel entries when the + tunnel closes, not when it opens. With keep-alive HTTP clients the log entry + may be delayed by minutes. Egress access logs are visible in `docker logs` but + appear after the connection closes. +- **`AllowPort` with port 443 only will deny plain-HTTP traffic.** Port 80 is + treated as the HTTP default: when 80 is in `AllowPort`, a bare `:authority` + header (plain HTTP, which omits the default port) is permitted alongside the + explicit `host:80` form — matching Squid's `allowed_ports` behaviour. When + port 80 is NOT in `AllowPort` (e.g., `AllowPort:[443]` only), bare-hostname + authorities are denied since port 80 is not listed. This is fail-closed (over- + blocks rather than a bypass). +- **No transparent L3/L4.** Non-cooperative traffic (workloads that ignore + `HTTP_PROXY`) is contained by the `Internal: true` network, not Envoy. True + non-bypassable enforcement requires iptables TPROXY + an init container with + `CAP_NET_ADMIN` — this is Phase 2 and requires its own architecture doc. +- **Gateway deny matches the requested name, not the resolved IP.** The deny + filter matches the `:authority` string (the gateway IP literal and the two + Docker-internal hostnames). Because an HTTP RBAC filter cannot see the address + the `dynamic_forward_proxy` cluster ultimately resolves and dials, a hostname + that resolves to the gateway IP — via DNS rebinding, an attacker-controlled + record, or an allowlisted host repointed at the gateway — is not caught, whereas + Squid's `dst` ACL denies on the resolved IP. Enforcing on the resolved address + requires L3/L4 interception (Phase 2, #5905); until then this is a known + weakening of the gateway block versus the Squid backend. +- **IPv4-only DFP DNS.** The `dynamic_forward_proxy` cluster's DNS lookup is + hardcoded to `V4_ONLY`. A host that resolves only to IPv6 addresses is + unreachable through the Envoy egress path even if it is allowlisted. +- **Wildcard ingress virtual host.** The ingress listener uses a wildcard virtual + host (`domains: ["*"]`) rather than a specific host match. This is safe because + it is bounded by the host-side port binding to `127.0.0.1`: only the local + proxy runner can reach the ingress listener, so the wildcard is not exposed + beyond the local machine. +- **Bootstrap file not cleaned up.** The generated bootstrap file (mode `0600` in + the OS temp directory) is not removed on workload teardown; it is only removed + on a mid-write generation failure. Files accumulate until the OS reclaims the + temp directory. +- **Deny-all on empty profile (deliberate divergence).** A nil or empty + permission profile produces an empty ALLOW policy map, which under Envoy's + default-deny RBAC means deny-all. Squid, by contrast, defaults to allow-all + when no rules are configured. The Envoy behaviour is fail-closed (safer) but is + an intentional divergence from the default backend, so a profile that is + permissive-by-omission under Squid will block all egress under Envoy. diff --git a/docs/arch/16-vmcp-mrtr.md b/docs/arch/16-vmcp-mrtr.md new file mode 100644 index 0000000000..cd69567549 --- /dev/null +++ b/docs/arch/16-vmcp-mrtr.md @@ -0,0 +1,544 @@ +# vMCP Multi Round-Trip Requests (MRTR) + +Design for serving MCP 2026-07-28 Multi Round-Trip Requests (SEP-2322) through +the Virtual MCP Server — the Modern replacement for the server-initiated +elicitation/sampling forwarding that exists only on Legacy (2025-11-25) +sessions. This work is tracked by issue **#6059** ("Support MRTR pass-through +for Modern backends in vMCP"), which supersedes the earlier design request +#5759 (closed as not-planned before #6059 was filed); it is the "coherent +future MRTR shape" the client-edge limitation section of +[10-virtual-mcp-architecture.md](10-virtual-mcp-architecture.md) names. That +document carries the **landed rationale** for both current limitations — the +backend-edge one ("elicitation and sampling are unavailable on Modern +backends": empty `clientCapabilities`, `-32021` surfacing) and the client-edge +one (the honest `-32603`/`-32021` contract) — and this document defers to it +for that ground rather than restating it: what follows is the design that +removes the limitations, cell by cell. + +Decision framing follows RFC-0083 (`stacklok/toolhive-rfcs`, +`rfcs/THV-0083-stateless-vmcp.md`), which defers MRTR and fixes the +workflow-handle design in its D5; this document sequences D5's enforcement +machinery and specifies the MRTR flows that trigger it. + +**Spec sources.** Written against the 2026-07-28 revision as published in +`modelcontextprotocol/modelcontextprotocol@main` on 2026-07-27, one day before +the revision's finalization date, when it still lived at `schema/draft` +(`LATEST_PROTOCOL_VERSION = "2026-07-28"`): `schema/draft/schema.ts`, +`/specification/draft/basic/patterns/mrtr`, SEP-2322 (Status: Final), SEP-2577 +(Status: Final), and the draft changelog. Anything cited from those sources +should be re-verified against the final text once cut; the MRTR SEP and schema +types were already Final at time of writing. + +*Re-checked 2026-07-28 (the nominal finalization date): the revision is +still not cut — no `schema/2026-07-28` directory exists, +`/specification/2026-07-28/` pages 404, and `schema/draft/schema.ts` is +byte-identical to the copy this design was written against. The checklist +below therefore remains pending the final cut.* + +**Draft-only dependencies (re-verify against the final cut).** The schema +types and the SEPs are Final and low-risk. Five load-bearing points in this +design rest on the draft *spec pages'* text, which refines (or diverges from) +the Final SEP and could still move before the final cut: + +- **The three-method table.** SEP-2322 (Final) also allows + `InputRequiredResult` on `GetTaskPayloadRequest`; the draft page dropped it + when tasks moved to an extension. If the final restores a fourth method, + the relay's method allow-list grows — a table edit, not a design change. +- **Capability gating** (server requirement 7, "MUST NOT send an + `inputRequests` that the client has not declared") is not stated by + SEP-2322, but it is not merely page-local either: the core `_meta` section + of `/specification/draft/basic` carries the general form as a MUST — a + server "MUST NOT rely on capabilities the client has not declared" and MUST + answer a request that needs one with `MissingRequiredClientCapabilityError` + (`-32021`) listing the missing capabilities in `data.requiredCapabilities`. + The MRTR requirement is the pattern-specific instance. Capability mirroring + is designed around it; the risk is only that the final cut reshuffles which + page states it. +- **Server requirement 6** ("at least one of `inputRequests` or + `requestState`") is stated on the page *and* in `schema/draft/schema.ts`'s + `InputRequiredResult` doc comment ("At least one of `inputRequests` or + `requestState` MUST be present") — anchored in the Final-typed schema, so + the least likely of these to move. It governs what vMCP treats as a backend + protocol violation (the egress decode fails closed on it — see slice 1). +- **Unrecognized `resultType` severity.** SEP-2322 says the client "SHOULD + treat unrecognized values as invalid protocol responses"; the draft basic + page's ResultType section tightens it: any unrecognized value "MUST be + considered invalid". The egress decode implements the MUST (sentinel-only, + no extractable round); a final cut that reverts to SHOULD costs nothing, + since rejecting remains SHOULD-conformant. +- **The `requestState` security language** differs between the two: SEP-2322 + says user-specific state "MUST use some mechanism to cryptographically + bind" it to the user, while the page allows omitting integrity protection + "only when tampering can cause nothing worse than request failure". The + workflow-handle design (below) satisfies the stricter reading — state never + leaves the server and the handle is owner-checked — so it is robust to + either landing in the final. + +## Protocol recap (what MRTR actually is) + +The 2026-07-28 revision removes server-initiated requests: "Servers **MUST** +send server-to-client requests (such as `roots/list`, `sampling/createMessage`, +or `elicitation/create`) using the MRTR pattern. The previous pattern of +server-initiated requests is no longer supported." (spec, `basic/patterns/mrtr`). + +Instead, a server answers a request it cannot complete with a `Result` whose +required `resultType` is `"input_required"`: + +```typescript +// schema/draft/schema.ts +export type InputRequest = CreateMessageRequest | ListRootsRequest | ElicitRequest; +export interface InputRequests { [key: string]: InputRequest; } +export interface InputRequiredResult extends Result { + inputRequests?: InputRequests; + requestState?: string; // opaque to the client +} +export interface InputResponseRequestParams extends RequestParams { + inputResponses?: InputResponses; + requestState?: string; +} +``` + +The client fulfills the `inputRequests`, then **retries the original request** +(new JSON-RPC id — "The JSON-RPC `id` MUST be different between the initial +request and the retry") carrying `inputResponses` and the echoed +`requestState`. The rounds are fully independent requests; any server replica +can serve any round. Only `tools/call`, `resources/read`, and `prompts/get` +may return `input_required` ("Servers **MUST NOT** send `InputRequiredResult` +responses on any other client requests"). + +Two requirements shape everything below: + +- **Capability gating** (server requirement 7): "Servers **MUST NOT** send an + `inputRequests` that the client has not declared support for in its + capabilities." Capabilities are declared per request in + `_meta.io.modelcontextprotocol/clientCapabilities`. +- **`requestState` is attacker-controlled** (server requirements 4–5): the + minting server "MUST treat `requestState` as an attacker-controlled input", + must protect its integrity (HMAC/AEAD) when it influences authorization or + business logic, and SHOULD bind the authenticated principal, a TTL, and an + originating-request identifier into the protected payload. + +## The four bridge cells + +vMCP terminates the client connection and re-originates backend calls, so MRTR +has a different shape in each client×backend era combination. The non-MRTR +halves of these cells were built in #6006; this table is the MRTR overlay. + +| Downstream client | Backend | MRTR shape | +|---|---|---| +| Modern | Modern | **Stateless pass-through** (the core design, below) | +| Modern | Legacy | **No MRTR bridge — deliberate.** Parking a live Legacy backend call server-side to synthesize `input_required` rounds is per-round server state with token-capability and replica-affinity costs: the session the revision removed, in different clothes. Costed and rejected in the client-edge limitation section of doc 10; the sanctioned stateful path is the Tasks extension (SEP-2663). The call fails with the two-path contract #6061 (merged) implements in `writeModernCallFailure`: `-32021` + `data.requiredCapabilities` when the client did not declare the capability, an explicit `-32603` naming SEP-2322 when it did. That contract is **permanent for this cell** — MRTR slices supersede it only where the backend is Modern. | +| Legacy | Modern | **In-request bridge**: vMCP fulfills the backend's `inputRequests` through the existing Legacy forwarding seams and retries the backend call itself (below) | +| Legacy | Legacy | Today's server-initiated forwarding, unchanged | + +### Cell 1 — Modern client ↔ Modern backend: stateless pass-through + +```mermaid +sequenceDiagram + participant C as Modern client + participant V as vMCP + participant B as Modern backend + C->>V: tools/call (id: 1, _meta.clientCapabilities) + V->>B: tools/call (id: v1, mirrored clientCapabilities) + B-->>V: input_required (inputRequests, requestState) + V-->>C: input_required (relayed verbatim, id: 1) + note over C: fulfills inputRequests locally + C->>V: tools/call retry (id: 2, inputResponses, echoed requestState) + note over V: fresh request — any replica;
backend re-derived from routing table + V->>B: tools/call retry (id: v2, both fields verbatim) + B-->>V: complete (id: v2) + V-->>C: complete (id: 2) +``` + +For a plain (non-composite) `tools/call` / `resources/read` / `prompts/get`, +vMCP adds **no state of its own**: + +1. **Capability mirroring (egress).** On the backend call, vMCP declares in + `_meta.io.modelcontextprotocol/clientCapabilities` exactly the relayable + capabilities the downstream client declared on this request — never more + (vMCP must not solicit `inputRequests` it cannot deliver downstream), never + less (an empty declaration makes a compliant backend that needs input + return `-32021 MissingRequiredClientCapability` — mandated by the core + `_meta` section's MUST, not merely predicted — which is today's behavior: + `mcpparser.ModernRequestMeta` hardcodes `clientCapabilities: {}`). +2. **`input_required` relay (ingress ← egress).** When the backend returns + `resultType:"input_required"`, vMCP relays `inputRequests` (values opaque, + verbatim) and `requestState` (opaque, verbatim) to the client in its own + `input_required` envelope, under the client's original request id. +3. **Retry relay.** The client's retry arrives as a fresh request carrying + `inputResponses` + `requestState`. vMCP routes it exactly like the first + round (routing is deterministic by capability name through the routing + table) and forwards both fields verbatim to the backend. + +Round-trip state lives **only** in the backend's `requestState`, whose +integrity/user-binding obligations sit with the backend that minted it, per +the spec. Whether the end-to-end principal binding survives that division of +labor **depends on the egress auth strategy**: + +- Under the **user-derived** strategies (`token_exchange`, `xaa`, + `upstream_inject`, `aws_sts`), vMCP's outgoing auth presents the calling + principal's identity to the backend on every round, so the backend can bind + its `requestState` to the same principal each time (server requirement 5) + and reject another user's echo. The claim above holds. +- Under the **shared-credential** strategies — `header_injection` (a static, + config-mounted secret whose own docstring says it "does not depend on user + identity") and `unauthenticated` (a no-op) — the backend sees **one vMCP + identity for every downstream user**. Its principal binding then binds all + rounds to that shared identity: it cannot protect one downstream user from + another, because no signal distinguishing them ever reaches it. If + cross-user misuse of a relayed `requestState` matters for such a backend, + **vMCP must bind the round to the downstream principal itself, because no + other component can** — which turns the wrap-vs-verbatim decision below + from a preference into a configuration-dependent obligation for + shared-credential egress. + +vMCP MUST NOT interpret, rewrite, or append to the relayed `requestState`; +the moment a design change requires vMCP-owned data inside it, that data +moves to the workflow-handle machinery below (D5 protections) or an +equivalent server-side entry, not into ad-hoc fields. + +**Open decision vs #6059 (resolve at slice 2/3, flagged, not silently +diverged):** the tracking issue sketches step 2 as "wrap the backend's opaque +`requestState` with vMCP routing context (which backend it came from)". Three +arms, no winner picked here: + +1. **Verbatim relay** (this document's default): re-derive the backend on the + retry from the capability name through the routing table. vMCP never mints + state, so MRTR server requirements 4–5 (attacker-controlled input, + integrity protection, principal binding) stay entirely with the backend — + no vMCP key-management surface. Costs: the rerouted-between-rounds edge + case, which is client-recoverable (the new backend rejects foreign state + and re-elicits); and under shared-credential egress it leaves the + cross-downstream-user gap above unclosed, since the backend cannot close + it and vMCP added nothing that could. +2. **#6059's wrapper**: vMCP wraps the backend's `requestState` with routing + context. The moment vMCP injects its own data, it becomes a state-minting + server under requirements 4–5; the wrapper must carry D5-grade protections + (integrity, principal binding, TTL) — cryptographic keys and their + management, the surface arm 1 avoids. +3. **Server-side handle over a verbatim backend round**: relay verbatim *to + the backend*, but make the **downstream-visible** `requestState` an opaque + handle into a server-side entry `{backend, capability, + owner=binding.Format(iss, sub), backend requestState, TTL}` — the same D5 + handle machinery the composite slice below already builds, and **no keys + at all** (the state never leaves the server, so there is nothing to sign + or seal). This closes cross-principal replay, cross-backend A→B + substitution, and the shared-credential gap above in one move. Its cost is + exactly the durable store this design otherwise keeps off the pass-through + path: every `input_required` round writes an entry, and plain-tool relay + stops being stateless at vMCP. + +Arm 1 maximizes statelessness, arm 3 maximizes containment with composite- +slice machinery instead of cryptography, arm 2 buys arm 3's properties at the +price of a key-management surface and is dominated by it unless the store is +unavailable. If shared-credential egress must support MRTR with cross-user +protection, arm 1 alone cannot deliver it (see above); the resolution at +slice 2/3 must say which arm serves that configuration. + +Failure modes, all client-recoverable per the spec's error-handling guidance +(a server that got unusable `inputResponses` "SHOULD respond with a new +`InputRequiredResult` requesting the missing information again"): + +- Backend set changed between rounds and the tool re-routed: the new backend + fails `requestState` validation and re-elicits or errors — no vMCP handling. +- Backend sends an `inputRequest` whose capability the client (and therefore + vMCP's mirrored declaration) did not declare: a backend protocol violation. + vMCP fails the call with an explicit error naming the violation; it never + forwards a request the downstream client cannot fulfill (server + requirement 7 applies to vMCP as a server in its own right). +- Unbounded re-elicitation: vMCP relays rounds without counting (the client + owns its retry budget; go-sdk's client middleware caps at 10 rounds, 3 + consecutive load-shedding rounds). vMCP's per-request deadline bounds each + individual round. + +### Cell 3 — Legacy client ↔ Modern backend: in-request bridge + +The inverse direction bridges cleanly because the blocking machinery already +exists. When a Modern backend returns `input_required` during a **Legacy** +client's call, vMCP — inside the same in-flight downstream request, on one +pod — fulfills each `inputRequest` through the existing server-initiated +forwarding seams (`vmcp.ElicitationRequester`, `vmcp.SamplingRequester`, both +capability-gated on what the Legacy client advertised at `initialize`), then +retries the backend call with the collected `inputResponses` and echoed +`requestState`. Precedent: go-sdk v1.7's own `serverMultiRoundTripMiddleware` +performs exactly this bridge for its handlers ("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"). Round count is bounded (mirror go-sdk's 10/3 caps) and the loop is +single-request-scoped: no durable state, no handle, no replica concern — +the downstream Legacy session vMCP already holds provides the delivery +channel. + +`ListRootsRequest` `inputRequests` are refused explicitly on this cell: vMCP +has never had a roots forwarding seam, and roots is SEP-2577-deprecated — an +explicit error beats building a new seam for a feature in its removal window. + +## Composite (workflow) tools: where durable state actually starts + +The pass-through above is stateless because vMCP is a relay. Composite tools +invert that: **vMCP is the server** executing a multi-step workflow +(`pkg/vmcp/composer`), so when a round trip interrupts a workflow, the +suspended state is vMCP's own. Two triggers: + +- a workflow step's Modern backend returns `input_required`, or +- the workflow itself needs user input (the composer's elicitation handler, + today implemented only over a Legacy session's `ElicitationRequester`). + +For a **Legacy** downstream client, neither suspends anything: the existing +blocking seams answer mid-call on the live session, and composites keep +executing synchronously inside one `tools/call` (RFC-0083 Alternative 3 — +the pod-local `InMemoryWorkflowStateStore` remains correct). + +For a **Modern** downstream client there is no session to block on, so the +workflow must **suspend**: vMCP returns `input_required` to the client with +the workflow's outstanding `inputRequests`, and the retry — a fresh request, +possibly on another pod — must **resume** it. This, precisely, is the trigger +RFC-0083 D5 sequences the durable machinery on: + +- **Round state**: the suspended `WorkflowStatus` (accumulated step outputs, + position, pending input keys) persists in a `WorkflowStateStore` backed by + Redis, injected via a `core.Config` seam (per RFC-0083's API-changes list). + It is too large and too sensitive to round-trip through the client, so the + client-visible `requestState` is a **handle**, not serialized state. +- **Handle format** (RFC-0083 D5, followed verbatim): a server-minted, + ≥128-bit opaque random token; an owner field on `WorkflowStatus` holding + the authenticated `(iss, sub)` pair encoded exactly as the existing + identity binding does — `pkg/vmcp/session/binding`'s + `Format(iss, sub)`, which returns `iss + "\x00" + sub` (and an error for + empty components), plaintext, not hashed — so the two owner-binding + mechanisms stay consistent. +- **Owner check on resume**: the retry's authenticated identity must satisfy + `binding.Format(iss, sub) == status.Owner` before any state is loaded; + a mismatch is an audit-visible security signal (owner-mismatch resume), + fails closed, and reveals nothing about the handle's existence. +- **TTL at write** (bounding the replay window per MRTR server requirement 5) + and **redaction**: the handle never appears in audit or telemetry output. +- **With authenticated incoming auth**, this satisfies the spec's requirement + to "cryptographically bind the data to the original user": the state never + leaves the server, and the handle is an unguessable capability whose owner + is checked server-side — stronger than client-side AEAD, with no + key-management surface. +- **Under `incomingAuth: anonymous`** the owner check is vacuous: every + anonymous session stores the same `binding.UnauthenticatedSentinel`, and + `validateCallerBinding` accepts any anonymous caller + (`pkg/vmcp/session/internal/security`), so "the original user" is not a + concept the deployment has — the owner-mismatch audit signal can never + fire, and protection reduces to handle unguessability plus TTL. Slice 5 + **accepts this as a documented single-tenant caveat rather than refusing to + suspend**: anonymous auth already grants any caller the tools themselves, + so a stolen handle adds no *tool-invocation* privilege the trust model + doesn't already concede — the thief could have run the workflow outright. + It does add a *data-read* privilege the trust model does not concede: + resuming loads another caller's accumulated step outputs, produced from + arguments the thief could not necessarily have supplied, so that + intermediate data is not something they could have obtained on their own — + a real residual exposure, gated on handle secrecy alone, which the + redaction requirement above covers. Refusing would break composites for + exactly the single-operator deployments anonymous auth exists for. A + multi-user deployment fronting vMCP with anything but authenticated + incoming auth gets no cross-user workflow isolation, and slice 5's docs + must say so where `anonymous` is documented. + +**The distinction that sizes the infrastructure**: plain-tool pass-through +rounds and the Legacy-client bridge need *no* durable store; only +cross-request workflow suspend/resume does. The store, its injection seam, +the owner field, and redaction all land with the composite slice — nothing +earlier. + +## What replaces `ElicitationRequester` + +Nothing replaces it one-for-one; the model inverts. `ElicitationRequester` / +`SamplingRequester` are *push* seams — block an in-flight call while the +server asks the client. They remain the implementation of the two Legacy +cells (unchanged) and become the fulfillment mechanism inside the +Legacy-client bridge (cell 3). On Modern paths, the seam is a *value*, not a +call: `input_required` results flow back through the ordinary return path — + +- domain: `vmcp.InputRequiredResult` (`inputRequests` as opaque raw values + keyed by string, plus `requestState`), carried by a typed error from + `BackendClient` so every existing complete-result path is untouched; +- core: `core.CallTool`/`ReadResource`/`GetPrompt` propagate it (suspending a + workflow first when the caller is the composer); +- server: `dispatchModern` renders it as the `resultType:"input_required"` + envelope, and accepts `inputResponses`/`requestState` params on the retry. + +## Sampling: deprecated pass-through, no feature work + +Explicit answer to "does sampling deserve an MRTR path": **vMCP builds no +sampling feature; sampling flows through the generic relay only.** + +- The MRTR union still carries it: `InputRequest = CreateMessageRequest | + ListRootsRequest | ElicitRequest` (schema/draft, above). SEP-2577's + explicit union freeze ("The following union types reference deprecated + types but MUST NOT be modified during the deprecation period", + `seps/2577-deprecate-roots-sampling-and-logging.md`) enumerates only + `ClientNotification`, `ClientResult`, `ServerRequest`, and + `ServerNotification` — `InputRequest` is **not** on that list (it + postdates the SEP), so its shape rests on the schema as published plus the + SEP's blanket "features remain fully functional during the deprecation + window" guarantee, not on a named freeze. Same practical effect during the + window; re-check the union at each revision rather than assuming it frozen. +- SEP-2577 (Final) deprecates sampling — along with roots and logging, the + latter irrelevant to MRTR — as of 2026-07-28: "New implementations SHOULD + NOT add support for deprecated features unless needed for backward + compatibility with existing counterparts" (SEP-2577, Capability + negotiation section), with "integrate directly with LLM provider APIs" as + the sanctioned replacement. +- vMCP's position under that carve-out: the pass-through relay is + **type-agnostic** — vMCP forwards whatever `inputRequests` the backend sent + and the client declared capability for, without interpreting them, so + sampling transits for the existing counterparts that still declare it; that + is compatibility plumbing, not added support. The one sampling-specific + code path retained is the cell-3 bridge fulfilling a sampling + `inputRequest` through the *existing* `SamplingRequester` — again existing + counterparts, no new capability. No new sampling machinery of any kind is + built, and when the deprecation window closes, sampling disappears from + vMCP by clients/backends dropping it, with zero vMCP code to remove beyond + the bridge's fulfillment case. + +Roots: same union, same deprecation, but unlike sampling there is no existing +vMCP counterpart (no roots forwarding seam has ever existed) — so roots gets +relay-only transit on cell 1 and an explicit refusal on cell 3. + +## Gaps in the pinned SDK surface (upstream candidates for toolhive-core) + +Verified against go-sdk `v1.7.0-pre.3` (the transitive pin via +`toolhive-core v0.0.34`) and `toolhive-core/mcpcompat`: + +1. **mcpcompat drops MRTR fields in both directions.** + `mcpcompat/mcp.CallToolResult` carries only an unexported `resultType` + (populated by `UnmarshalJSON`, feeds `NeedsInput()`); it has **no** + `InputRequests`/`RequestState` fields, so a decoded `input_required` + result loses its payload. `CallToolParams` (`Name`, `Arguments`, `Meta`, + `Task`) cannot carry `inputResponses`/`requestState` on a retry. Upstream + ask: add both field pairs (go-sdk's own `mcp.CallToolResult` already + exports `InputRequests InputRequestMap` and `RequestState string`). +2. **No no-initialize client primitive** in mcpcompat's public API (tracked + as #6018) — the reason `pkg/vmcp/client`'s hand-rolled `modernCall` exists + at all, and therefore where MRTR egress decode/retry params must live + until #6018's upstream work lands. +3. **go-sdk's MRTR internals are unexported where vMCP needs them.** The + client retry loop exists (`clientMultiRoundTripMiddleware`, + `MultiRoundTripOptions.Disabled` for manual handling) but its pieces + (`fulfillInputRequests`, `setMultiRoundTripRetryParams`) are unexported; + server-side, `resultType` is settable only via unexported + `setResultType`/`setCompleteResultType` inside `ServerSession` dispatch — + exactly the dispatch the stateless Modern path bypasses — so the + hand-rolled `modern_envelope.go` must emit `input_required` itself. + Upstream ask (shared with #6018's step 1): exported stateless + client/server MRTR primitives. +4. **`ServerSession.assertServerInitiatedRequestAllowed` gates by negotiated + protocol version only** (`v1.7.0-pre.3` `mcp/server.go:1544-1553`: + `iparams.ProtocolVersion >= protocolVersion20260728`), never consulting + capability declarations. Harmless for this design — the seams it guards + are only used on Legacy sessions — but it forecloses any notion of + re-enabling server-initiated requests for a capability-declaring Modern + client. Accounted for, no action. +5. **`-32021` at HTTP 400 is unshippable** with current go-sdk clients (their + transient set is 500/502/503/504/429; any other 4xx is permanent session + death), which is why the undeclared-capability error lands at HTTP 200 as + a documented deviation — implemented by #6061 + (`writeModernMissingCapability`), tracked upstream as go-sdk#1117; MRTR + does not change it, it *removes the case* for capabilities the client + does declare against Modern backends. + +## Sequencing (collision-aware) + +Collision set at time of writing: #6050 (modern pagination/subscriptions — +rewrites `modern_dispatch.go`/`modern_envelope.go`) and #6051 (Legacy-pins +the forwarding tests, adds doc 10's client-edge limitation section) have +both **merged** (2026-07-28); #6033 (kill-switch removal — +`classification.go`, `server.go`, `serve.go` and test helpers) remains open. +Slices are ordered so each lands without touching #6033's files until it +merges: + +1. **Egress MRTR surface (safe now; implemented alongside this doc).** Domain + types `vmcp.InputRequiredResult` and `vmcp.InputRequiredError` (carrier and + `InputRequiredFromError` extractor live beside the domain type, so slice + 2's server-side rendering never imports the concrete HTTP client and mock + consumers can construct the error); `pkg/vmcp/client`'s Modern shim decodes + an `input_required` envelope into that typed, `errors.Is`-compatible error + instead of the opaque `errModernInputRequired`. Extraction is + **fail-closed**: a round is extractable only on the three methods that may + carry one, only for `resultType:"input_required"` exactly, and only when + the payload decodes strictly and satisfies server requirement 6 — every + violation keeps pure sentinel semantics, so a slice-2 consumer can never + relay a malformed or misplaced round. `RequestState` is a `*string`, since + client requirement 2 makes absent-vs-present-and-empty an observable + distinction the retry must preserve. No behavior change: capabilities are + still declared empty, so a compliant backend cannot yet send + `input_required`; the seam #6051 calls "the egress half" simply becomes + load-bearing. Files: `pkg/vmcp/mrtr.go`, `pkg/vmcp/client/modern.go` (one + hook), `pkg/vmcp/client/modern_mrtr.go` — none touched by #6033. +2. **Ingress envelope + retry params (after #6050/#6033).** `dispatchModern` + accepts `inputResponses`/`requestState`, threads them through + `core.CallTool`→`BackendClient`, and renders `input_required` envelopes; + parser vocabulary for the two params. Post-#6061, the rendering branch + hooks into `writeModernCallFailure` via `vmcp.InputRequiredFromError`, + alongside (never instead of) its authz-priority and capability-refusal + branches — the two cannot co-occur: the refusal recorder fires only on + the Legacy-backend forwarding seams, `input_required` only from a Modern + backend's envelope. +3. **Capability mirroring (activates cell 1; with slice 2, this is what + issue #6059 tracks).** `ModernRequestMeta` gains a capabilities argument; + dispatch threads the downstream declaration to the egress call (reusing + #6061's `modernClientDeclaredCapability` reading of `_meta`); the relay + goes live end-to-end. The `requestState` open decision above (verbatim vs + #6059's routing-context wrapper) must be resolved here at the latest. The `-32021`-shaped + undeclared-capability error contract landed ahead of this in #6061 + (merged 2026-07-28; `-32021` at HTTP 200 as the documented deviation, + upstream go-sdk#1117). When this slice lands, #6061's declared-capability + `-32603` message ("multi-round retrieval … which this server does not + implement") stops being true for Modern backends and needs rescoping to + the Legacy-backend cell it permanently serves. +4. **Cell-3 bridge.** The in-request fulfillment loop over + `ElicitationRequester`/`SamplingRequester` with go-sdk-mirrored round caps. +5. **Composite suspend/resume.** Durable `WorkflowStateStore` (Redis) behind + the `core.Config` seam, D5 handle + owner check + TTL + redaction, resume + dispatch, owner-mismatch audit signal. + +Slices 2–5 each update doc 10's client-edge limitation section (whose current +text truthfully says MRTR is unserved) and the forwarding-test dispositions +from #6051 as their claims become false. + +## Testing strategy + +- Slice 1: envelope decode (payload extraction, `errors.Is` back-compat with + the sentinel, fail-closed rejection of malformed/empty/misplaced rounds), + no-behavior-change pins. +- Slice 2/3: dispatcher round-trip — `input_required` envelope shape + (resultType, echoed id, relayed keys), retry threading, capability-gate + refusals (undeclared → no forward; backend violation → explicit error). +- Cell 1 end-to-end: a Modern fake backend demanding one elicitation round; + assert both rounds' independence (distinct ids, verbatim + `requestState` echo, no vMCP-side state between them — serve round 2 from + a *fresh server instance* to prove it). +- Cell 3: Legacy downstream client + `input_required` Modern backend; assert + fulfillment via the existing seams, round caps, and the roots refusal. +- Slice 5: owner-binding (B cannot resume A's workflow), TTL expiry, + Redis-unavailable fail-closed — the RFC-0083 test list, verbatim. + +## References + +- SEP-2322 (Final), SEP-2577 (Final), SEP-2663 (tasks redesign); + `schema/draft/schema.ts` and `/specification/draft/basic/patterns/mrtr` as + of 2026-07-27. +- RFC-0083 (`stacklok/toolhive-rfcs`) — D5, Alternative 3, implementation + plan's deferred list. +- [10-virtual-mcp-architecture.md](10-virtual-mcp-architecture.md) — the + client-edge limitation section (added by #6051) this design supersedes + step by step; forwarding seams; #6006 bridge cells. +- Issues: #5743 (epic); **#6059 (tracks the MRTR pass-through this document + designs — supersedes #5759, the earlier not-planned design request)**; + #6018 (shim retirement); #5959/#6033 (kill-switch); #6050 + (pagination/subscriptions). +- Adjacent Modern-path issues this design deliberately does NOT cover: + **#6058** — per-request SSE streaming for `notifications/progress`/ + `notifications/message`; distinct from MRTR (those are notifications, not + server-initiated requests — MRTR does not touch them), but slice 2's + `input_required` envelope may ride the same future stream as its final + message, per the spec's "MAY be sent … as the final message on an SSE + stream". **#6064** — `ping` removal and the required `resultType` on the + dispatch path; pure conformance, orthogonal to MRTR. **#6065** — + `subscriptions/listen` push delivery; a different channel with a fixed + four-type subscribable set, structurally disjoint from MRTR rounds. diff --git a/docs/arch/17-token-exchange-delegation.md b/docs/arch/17-token-exchange-delegation.md new file mode 100644 index 0000000000..b34d412d34 --- /dev/null +++ b/docs/arch/17-token-exchange-delegation.md @@ -0,0 +1,257 @@ +# External Subject-Token Exchange (Delegation) + +This document covers the embedded authorization server's RFC 8693 token-exchange +grant when the `subject_token` comes from an external, trusted OIDC issuer +(`trustedIssuers`). It is unrelated to the *outgoing* token exchange described in +docs [02](02-core-concepts.md), [04](04-secrets-management.md), +[05](05-runconfig-and-permissions.md), [09](09-operator-architecture.md), and +[10](10-virtual-mcp-architecture.md), where `MCPExternalAuthConfig` middleware +exchanges a client's token for an upstream IdP token. Here the server is the +token *issuer*, accepting an externally-minted token as proof of identity and +re-issuing a ToolHive-scoped delegated token. + +Normally the grant accepts only the server's own self-issued tokens as +`subject_token`. `trustedIssuers` extends that to tokens minted by an external +IdP (e.g. a corporate IdP), so a client already holding a token from that IdP +can exchange it for a ToolHive delegated token without a separate ToolHive +login. + +## Prerequisite: not yet usable end to end + +- Using this grant requires a confidential client holding the token-exchange + grant. No supported deployment path provisions one today: DCR always + registers public clients restricted to `authorization_code`/`refresh_token`, + CIMD clients are public-only, and `RunConfig` has no client-seeding field. +- Discovery metadata does not yet advertise the token-exchange grant or any + secret-based client auth method, so a metadata-driven client will conclude + the grant is unsupported and never attempt it. +- This applies to the self-issued subject-token path too, not only + `trustedIssuers`. +- Tracked in [#6082](https://github.com/stacklok/toolhive/issues/6082). + +## Trust model + +A trusted issuer plus a matching audience and a valid signature only proves +the token was minted for ToolHive to consume — it authorizes ToolHive as a +*resource*, not any particular *client* as a delegate. Accepting the token on +that basis alone would let any client presenting a validly-signed token +exchange it: a confused-deputy risk (CWE-863). The handler therefore requires +an explicit consent signal before treating the request as an authorized +delegation. + +## Consent signals + +Two functions cooperate here. `resolveAllowedActor` +(`pkg/authserver/server/tokenexchange/multi_issuer_validator.go`), called from +`validateExternalToken` during signature/claim validation, checks the +external-issuer allowlist and — when it matches — sets +`ValidatedClaims.ExternalActor`. `checkDelegationConsent` +(`pkg/authserver/server/tokenexchange/handler.go`) runs afterwards and decides +whether to grant the exchange, in this order: + +1. **`may_act` (RFC 8693 §4.4).** If the subject token carries a well-formed + `may_act` claim, it is authoritative: the allowlist step above is skipped + entirely (`validateExternalToken` never calls `resolveAllowedActor` when + `may_act` is present), and `checkDelegationConsent` enforces `may_act.sub` + against the authenticated ToolHive client. A malformed `may_act` is + rejected outright by `validateMayActShape`, not silently ignored or + fallen through to the allowlist. On the external path, `may_act.iss` is + mandatory, not merely constrained when present as on the self-issued path + — `validateMayActShape`'s `requireIss` parameter is `true` there. Without + this, an external issuer omitting `iss` could authorize any ToolHive + client by naming its ID in a bare `may_act.sub`, bypassing + `allowedActors`/`allowedDelegateClients` entirely — the one consent path + that does. +2. **`ExternalActor`.** Otherwise, consent was already established by + `resolveAllowedActor`: the claim named by that issuer's `actorClaim` + (default `azp`; `appid` for Microsoft Entra v1, `cid` for Okta) matched an + entry in that issuer's `allowedActors`. `allowedActors` holds client IDs in + the *external* IdP's namespace, not ToolHive client IDs — the likeliest + misconfiguration. An empty `allowedActors` accepts only `may_act`-bearing + tokens from that issuer. `checkDelegationConsent` then additionally checks + the issuer's `allowedDelegateClients` against the authenticated ToolHive + client — this field is required (see below), so the check always applies + on this path. +3. **`client_id` binding.** For a self-issued subject token (not part of the + external path above), the token's `client_id` must match the authenticated + client. +4. **No binding.** If none of the above apply, the token carries no + verifiable client binding and the exchange is rejected. + +## Accepted limitations + +1. **`allowedDelegateClients` binds an allowlisted actor to specific ToolHive + clients, and is mandatory, not opt-in.** `allowedActors` by itself + authorizes "this external client's tokens may be exchanged here," not + "…by this particular ToolHive client" — the external actor claim is never + compared against the authenticated ToolHive client ID. Without a separate + binding, every ToolHive confidential client holding the token-exchange + grant would be delegation-equivalent with respect to an allowlisted + external actor, so compromise of the weakest such client would be as good + as compromise of all of them. + + `TrustedIssuer.AllowedDelegateClients` (`allowed_delegate_clients` on a + hand-written `authserver.RunConfig`) closes this per issuer: a list of + ToolHive client IDs permitted to use that issuer's allowlisted actors, or + the wildcard `"*"` to explicitly permit any ToolHive client holding the + token-exchange grant. `validateTrustedIssuer` rejects the field when it is + empty or absent — permissiveness must be *declared* with the wildcard, not + obtained by leaving the field out — and rejects the wildcard combined with + specific client IDs, since silently ignoring the specific IDs alongside it + would be worse than rejecting the config outright. `checkDelegationConsent` + checks the authenticated client against this list on both consent paths: + `may_act` bypasses `allowedActors` (the external-issuer allowlist) but NOT + `allowedDelegateClients`, since the validator sets `AllowedDelegateClients` + for every external token regardless of which path authorized it (see + limitation 4 below). A self-issued `may_act` (no external issuer involved) + has no `allowedDelegateClients` equivalent and is unaffected — it remains + bound by `may_act.sub` alone. + + Nothing can reach this code path in production today: the token-exchange + grant requires a confidential client, and no supported deployment path + provisions one (see issue #6082). That makes the fail-closed default free + right now — it would be a breaking change once a client can actually reach + this grant. +2. **Subject namespace collisions are closed, not accepted.** A trusted + issuer is trusted to assert *any* subject this server accepts for + delegation, and downstream authorization decisions (Cedar) key on `sub` + alone — not `iss`. Rather than requiring operators to keep every issuer's + subject namespace disjoint, the delegated token's `sub` is qualified as + `#` for every trusted-issuer exchange + (`tokenexchange.delegatedSubject`), never the external `sub` copied + verbatim. ToolHive's own native subjects are UUIDs minted by + `UserResolver.ResolveUser`, which cannot contain `#`, so a qualified + external subject can never collide with one. Scope names are not + qualified this way and remain the operator's responsibility to keep + disjoint across issuers. +3. **Provenance is recorded for every external token.** The RFC 8693 §4.1 + `act` claim records who acted: `act.sub` is always the ToolHive client, + with the external issuer nested one level in — `ValidatedClaims.ExternalIssuer` + is set for every token validated by the external-issuer path, whether or + not it also carries `may_act`. The nested entry additionally carries `sub` + (the allowlisted actor claim) when the allowlist path resolved one; + a `may_act`-bearing external token yields `act = {sub: , + act: {iss: }}` — no client-namespace actor to report, but + the issuer is still recorded. Either way, Cedar authorizers key on `sub` + and do not read `act` — it is an audit trail, not an access control. (AWS + STS role mapping can read arbitrary claims including `act` via its CEL + matcher, so "authorizers" here means Cedar specifically, not every + consumer.) +4. **A `may_act`-emitting issuer bypasses the allowlist entirely.** Since it + takes priority and can name any ToolHive client, that claim must be drawn + from ToolHive's own client namespace and must not be influenceable by an + untrusted party. +5. **ID/access-token discrimination on the external path rests on two + partial layers, neither of which is exhaustive alone.** Nothing inspects + the JWT header's `typ` claim: RFC 8725 §3.12 recommends `typ: at+jwt` for + access tokens, but Entra v1/v2, Okta, and Google all emit a bare `JWT` + `typ` for access tokens too, so requiring `at+jwt` would reject those + providers' genuine access tokens. + + Layer one: `rejectIDTokenClaims` (`validator.go`) rejects a subject token + carrying `at_hash` or `c_hash` — OIDC Core §3.3.2.11/§3.3.2.10 define both + for ID tokens only. This is a sound *positive* detector but not + exhaustive: both claims are OPTIONAL even on a valid + authorization-code-flow ID token, so an ID token that omits them is not + caught here. + + Layer two: `TrustedIssuer.ExpectedAudience` is documented as MUST be a + resource/API identifier and MUST NOT be a client ID, since an ID token's + `aud` is the requesting client's own ID while an access token's `aud` + names the resource. This check is operator-dependent, not + self-enforcing — nothing validates that the operator actually set a + resource identifier rather than a client ID. + `NewMultiIssuerTokenValidator` emits a `slog.Warn` at startup + (`looksLikeResourceIdentifier`) when `expected_audience` has no URI + scheme (`https://`, `api://`), naming the issuer and the consequence, but + does not hard-reject it: Entra v1 legitimately uses a bare app-ID GUID + (no scheme) as an access token's `aud`, so rejecting every non-URI + audience would break that real, supported provider. + + A third discriminator — rejecting when a subject token's `aud` equals the + issuer's configured actor claim, since OIDC Core requires an ID token's + `aud` to equal its own `azp` — was evaluated and not shipped: an app whose + API and client share one registration (the same ambiguous + bare-GUID-audience configuration layer two warns about) legitimately + issues *access* tokens where `aud` and `azp`/`appid` are the same value + too, so this would reject genuine access tokens from that setup. + + Taken together: an operator who sets `expected_audience` to a client ID + (despite the startup warning), talking to a real IdP whose ID tokens omit + `at_hash`/`c_hash` (both legal per OIDC Core), has no remaining + discriminator here. The blast radius is bounded: an admitted ID token + carries no `scope`/`scp` claim, so `grantScopes` (`handler.go`) either + rejects the exchange with `invalid_scope` or grants a zero-scope + delegated token — it must also be signed by a configured trusted issuer + and satisfy one of the consent signals above. + +## Operational notes + +- **Audience.** The requested audience is bounded by the subject token's own + `aud` claim (`ensureAudienceSubsetOfSubject`); a request for an audience the + subject token doesn't carry fails with `invalid_target`. An external IdP's + `aud` is typically an app-ID GUID or `api://`. For an exchange to + succeed, the external IdP's API identifier must be set as the per-issuer + `expectedAudience` *and* appear in the server's own `allowedAudiences` (plus + the calling client's registered audiences) — `expectedAudience` alone is not + enough, since it only governs whether the subject token validates, not what + the delegated token is granted. A startup WARN fires when `expectedAudience` + isn't in `allowedAudiences`. +- **Scopes.** `grantScopes` rejects — it does not intersect — any requested + scope absent from the subject token's granted scopes, with `invalid_scope`. + Both the RFC 9068 `"scope"` string claim and the `"scp"` array claim are + read (`"scope"` wins if both are present) — `"scp"` is not an Entra + quirk, it is what fosite's own default JWT claims strategy writes for a + self-issued ToolHive access token, so a genuine subject token minted by + this server's own token endpoint carries scopes this way already. +- **Delegated token lifetime.** `min(subject token's remaining lifetime, + configured delegationLifespan)`. A short-lived external token silently + yields a short delegated token. +- **Delegation chain depth.** Re-exchanging an already-delegated token nests + `act` further; `maxDelegationDepth` (10) bounds it, past which the exchange + fails with `invalid_grant` ("delegation chain is too deep"). +- **Discovery redirects.** The per-issuer HTTP client only follows same-host + redirects (`networking.SameHostRedirectPolicy()`), so an issuer whose + `/.well-known/openid-configuration` redirects cross-host cannot be + onboarded via discovery — set `jwksUrl` explicitly to skip discovery. +- **JWKS caching.** A successful fetch is cached and refreshed by jwx's + `jwk.Cache` on a schedule driven by the response's `Cache-Control`/`Expires` + headers (no fixed TTL). The 30-second figure is `jwksFetchFailureBackoff`, + which gates only how often a persistently-failing issuer is retried before + its first successful fetch — so a just-corrected `jwksUrl` can keep failing + for up to 30s before the next retry; once a fetch has ever succeeded, this + backoff no longer applies. +- **Diagnostics.** A non-allowlisted actor and an untrusted issuer both + surface as the same generic `invalid_request` + ("The subject token is invalid or could not be verified."), per RFC 8693 + §2.2.2. The only discriminator is a `slog.Debug` log line, so diagnosing a + failed exchange needs debug logging enabled. The two exceptions are startup + `slog.Warn` lines — an issuer with empty `allowedActors`, and an + `expectedAudience` missing from `allowedAudiences` — the only non-DEBUG + diagnostics this feature emits. +- **HTTPS enforcement for a trusted issuer.** `validateTrustedIssuerURL` + (`pkg/authserver/config.go`) validates `issuer_url` at config time and, + unlike the server's own issuer, grants no localhost exemption: an + `http://localhost` trusted issuer is rejected unless that issuer's own + `insecure_allow_http` is set. There is no separate runtime scheme check for + `issuer_url` — only `jwks_url` gets one (`ValidateJWKSURL`, on every fetch; + shared verbatim between the runtime choke point and the config-time check + so the two can't drift). There is no operator (CRD) surface for this + feature yet — see [#6082](https://github.com/stacklok/toolhive/issues/6082) + — so today `trustedIssuers` is reachable only via a hand-written + `authserver.RunConfig`. +- **`allowPrivateIPs` requires `jwksUrl`.** Enforced at config time by + `validateTrustedIssuers` (`pkg/authserver/config.go`) for every RunConfig. + Without a hand-configured `jwksUrl`, OIDC discovery — a document fetched + from, and thus influenceable by, the external issuer itself — would choose + the private JWKS dial target, which is exactly what pinning it to + operator-supplied config prevents. +- **Misconfiguration surfaces as a pod crash**, not an operator condition — + check pod logs, not `kubectl describe`. + +## Implementation + +- `pkg/authserver/server/tokenexchange/handler.go` — `checkDelegationConsent`, `grantScopes`, `grantAndBoundAudiences` +- `pkg/authserver/server/tokenexchange/multi_issuer_validator.go` — `MultiIssuerTokenValidator`, `TrustedIssuer`, `resolveAllowedActor`, `ValidateJWKSURL` +- `pkg/authserver/server/tokenexchange/validator.go` — `assignClaim`, `buildValidatedClaims`, `validateMayActShape`, `scpToScopeString` +- `pkg/authserver/config.go` — `validateTrustedIssuerURL`, `validateJWKSEndpointURL`, `validateTrustedIssuers`, `warnTrustedIssuerAudiences` diff --git a/docs/arch/README.md b/docs/arch/README.md index e590eee6b2..13df8eeb99 100644 --- a/docs/arch/README.md +++ b/docs/arch/README.md @@ -120,6 +120,26 @@ Welcome to the ToolHive architecture documentation. This directory contains comp - Component inventory and per-client dropped-component warnings - Name/repo consistency check, extraction safety, TOML mutation under file lock +15. **[Envoy Network Proxy](15-envoy-network-proxy.md)** + - Experimental Envoy backend (`TOOLHIVE_NETWORK_PROXY=envoy`) replacing two Squid containers with one + - Egress forward proxy and ingress reverse proxy as separate listeners in a single process + - L7 `:authority`-based RBAC allow/deny and `dstdomain` parity (incl. Docker-gateway blocking) + - Dynamic forward proxy with per-request DNS and native HTTPS CONNECT tunnelling + - Squid-vs-Envoy comparison and known limitations (`AllowPort` gap, V4_ONLY DNS, deny-all on empty profile) + +16. **[vMCP Multi Round-Trip Requests (MRTR)](16-vmcp-mrtr.md)** + - MRTR (SEP-2322) design for the Modern (2026-07-28) elicitation/sampling path through vMCP + - The four client×backend bridge cells: stateless Modern↔Modern pass-through, in-request Legacy-client bridge + - Per-request capability mirroring and verbatim `requestState` relay + - Composite-workflow suspend/resume as the sole durable-state trigger (RFC-0083 D5 handles, owner binding, TTL) + - Sampling's deprecated-pass-through-only disposition (SEP-2577) and the go-sdk/mcpcompat gap list + +17. **[External Subject-Token Exchange (Delegation)](17-token-exchange-delegation.md)** + - `trustedIssuers`: accepting subject tokens from an external OIDC issuer during RFC 8693 token exchange + - Confused-deputy trust model and the `may_act` / allowlist consent signals + - Accepted limitations: client-set equivalence, disjoint subject namespaces, partial provenance + - Operational gotchas: audience/scope binding, discovery redirects, JWKS caching, diagnostics + ### Existing Documentation For middleware architecture, see: **[docs/middleware.md](../middleware.md)** diff --git a/docs/arch/vmcp-local.md b/docs/arch/vmcp-local.md index ca5568511b..8cc5b19b5c 100644 --- a/docs/arch/vmcp-local.md +++ b/docs/arch/vmcp-local.md @@ -103,6 +103,8 @@ thv vmcp serve --config vmcp.yaml | 2 | `--optimizer-embedding` | FTS5 + TEI semantic | Managed TEI container | `find_tool`, `call_tool` only | | 3 | `optimizer.embeddingService` in config YAML | FTS5 + external embedding service | User-managed | `find_tool`, `call_tool` only | +Any tier >= 1 keeps clients on MCP 2025-11-25 (Legacy): the `find_tool`/`call_tool` meta-tools are session-scoped, so Modern-capable (2026-07-28) clients are negotiated down to Legacy by the capability gate (`pkg/vmcp/server/modern_gate.go`; Modern parity is tracked in #6089). + Tier 2 (`--optimizer-embedding`) implies `--optimizer`. The TEI container is started automatically and stopped on server shutdown. **Implementation**: `pkg/vmcp/optimizer/optimizer.go`, `pkg/vmcp/cli/embedding_manager.go` diff --git a/docs/authz.md b/docs/authz.md index 9575f6ec55..c20f19e111 100644 --- a/docs/authz.md +++ b/docs/authz.md @@ -51,6 +51,18 @@ occurs: 4. If the request is authorized, it is passed to the next handler. Otherwise, a 403 Forbidden response is returned. +### Virtual MCP parity + +Virtual MCP (`vmcp`) enforces the same Cedar policies through its core admission +seam rather than HTTP middleware, but the client-facing behavior matches a +single-server `thv run`: a denied `tools/call`, `resources/read`, or `prompts/get` +is rejected with **HTTP 403** and a JSON-RPC error whose code is `403`, and the +denied request is audited with outcome `denied`. The denial message is kind-only +(for example, `call denied by authorization policy`) and never reveals the +capability name or whether it exists, so it cannot be used to enumerate the +server. See the [Virtual MCP architecture](arch/10-virtual-mcp-architecture.md#authorization-enforcement-core-admission-seam--pre-dispatch-gate) +for details. + ## Configure authorization To set up authorization for an MCP server managed by ToolHive, follow these diff --git a/docs/cli/thv.md b/docs/cli/thv.md index 74824dfad8..170d95c26f 100644 --- a/docs/cli/thv.md +++ b/docs/cli/thv.md @@ -35,6 +35,7 @@ thv [flags] ### SEE ALSO +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins * [thv build](thv_build.md) - Build a container for an MCP server without running it * [thv client](thv_client.md) - Manage MCP clients * [thv config](thv_config.md) - Manage application configuration diff --git a/docs/cli/thv_ai-plugin.md b/docs/cli/thv_ai-plugin.md new file mode 100644 index 0000000000..1df12af56f --- /dev/null +++ b/docs/cli/thv_ai-plugin.md @@ -0,0 +1,44 @@ +--- +title: thv ai-plugin +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin` +last_update: + author: autogenerated +slug: thv_ai-plugin +mdx: + format: md +--- + +## thv ai-plugin + +Manage AI-tool plugins + +### Synopsis + +The ai-plugin command provides subcommands to manage plugins for AI tools +(e.g. Claude Code, Codex) — not plugins for ToolHive itself. + +### Options + +``` + -h, --help help for ai-plugin +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv](thv.md) - ToolHive (thv) is a lightweight, secure, and fast manager for MCP servers +* [thv ai-plugin build](thv_ai-plugin_build.md) - Build an AI-tool plugin into a local OCI artifact +* [thv ai-plugin builds](thv_ai-plugin_builds.md) - List locally-built AI-tool plugin artifacts +* [thv ai-plugin info](thv_ai-plugin_info.md) - Show details of an AI-tool plugin +* [thv ai-plugin install](thv_ai-plugin_install.md) - Install an AI-tool plugin +* [thv ai-plugin list](thv_ai-plugin_list.md) - List installed AI-tool plugins +* [thv ai-plugin push](thv_ai-plugin_push.md) - Push a built AI-tool plugin to an OCI registry +* [thv ai-plugin uninstall](thv_ai-plugin_uninstall.md) - Uninstall an AI-tool plugin +* [thv ai-plugin validate](thv_ai-plugin_validate.md) - Validate an AI-tool plugin directory + diff --git a/docs/cli/thv_ai-plugin_build.md b/docs/cli/thv_ai-plugin_build.md new file mode 100644 index 0000000000..094991b1f1 --- /dev/null +++ b/docs/cli/thv_ai-plugin_build.md @@ -0,0 +1,42 @@ +--- +title: thv ai-plugin build +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin build` +last_update: + author: autogenerated +slug: thv_ai-plugin_build +mdx: + format: md +--- + +## thv ai-plugin build + +Build an AI-tool plugin into a local OCI artifact + +### Synopsis + +Build a plugin from a local directory into an OCI artifact that can be pushed to a registry. + +On success, prints the OCI reference of the built artifact to stdout. + +``` +thv ai-plugin build [path] [flags] +``` + +### Options + +``` + -h, --help help for build + -t, --tag string OCI tag for the built artifact +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_ai-plugin_builds.md b/docs/cli/thv_ai-plugin_builds.md new file mode 100644 index 0000000000..b26da4a41e --- /dev/null +++ b/docs/cli/thv_ai-plugin_builds.md @@ -0,0 +1,41 @@ +--- +title: thv ai-plugin builds +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin builds` +last_update: + author: autogenerated +slug: thv_ai-plugin_builds +mdx: + format: md +--- + +## thv ai-plugin builds + +List locally-built AI-tool plugin artifacts + +### Synopsis + +List all locally-built OCI plugin artifacts stored in the local OCI store. + +``` +thv ai-plugin builds [flags] +``` + +### Options + +``` + --format string Output format (json, text) (default "text") + -h, --help help for builds +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins +* [thv ai-plugin builds remove](thv_ai-plugin_builds_remove.md) - Remove a locally-built AI-tool plugin artifact + diff --git a/docs/cli/thv_ai-plugin_builds_remove.md b/docs/cli/thv_ai-plugin_builds_remove.md new file mode 100644 index 0000000000..d265c16d83 --- /dev/null +++ b/docs/cli/thv_ai-plugin_builds_remove.md @@ -0,0 +1,39 @@ +--- +title: thv ai-plugin builds remove +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin builds remove` +last_update: + author: autogenerated +slug: thv_ai-plugin_builds_remove +mdx: + format: md +--- + +## thv ai-plugin builds remove + +Remove a locally-built AI-tool plugin artifact + +### Synopsis + +Remove a locally-built OCI plugin artifact and its blobs from the local OCI store. + +``` +thv ai-plugin builds remove [flags] +``` + +### Options + +``` + -h, --help help for remove +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin builds](thv_ai-plugin_builds.md) - List locally-built AI-tool plugin artifacts + diff --git a/docs/cli/thv_ai-plugin_info.md b/docs/cli/thv_ai-plugin_info.md new file mode 100644 index 0000000000..107f371c5a --- /dev/null +++ b/docs/cli/thv_ai-plugin_info.md @@ -0,0 +1,42 @@ +--- +title: thv ai-plugin info +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin info` +last_update: + author: autogenerated +slug: thv_ai-plugin_info +mdx: + format: md +--- + +## thv ai-plugin info + +Show details of an AI-tool plugin + +### Synopsis + +Display detailed information about a plugin, including metadata, version, and installation status. + +``` +thv ai-plugin info [plugin-name] [flags] +``` + +### Options + +``` + --format string Output format (json, text) (default "text") + -h, --help help for info + --project-root string Project root path for project-scoped plugins + --scope string Filter by scope (user, project) +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_ai-plugin_install.md b/docs/cli/thv_ai-plugin_install.md new file mode 100644 index 0000000000..bb4c04e4b6 --- /dev/null +++ b/docs/cli/thv_ai-plugin_install.md @@ -0,0 +1,45 @@ +--- +title: thv ai-plugin install +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin install` +last_update: + author: autogenerated +slug: thv_ai-plugin_install +mdx: + format: md +--- + +## thv ai-plugin install + +Install an AI-tool plugin + +### Synopsis + +Install a plugin by name or OCI reference. +The plugin will be fetched from a remote registry and installed locally. + +``` +thv ai-plugin install [plugin-name] [flags] +``` + +### Options + +``` + --clients string Comma-separated target client apps (e.g. claude-code,codex), or "all" for every available client + --force Overwrite existing plugin directory + --group string Group to add the plugin to after installation + -h, --help help for install + --project-root string Project root path for project-scoped installs + --scope string Installation scope (user, project) (default "user") +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_ai-plugin_list.md b/docs/cli/thv_ai-plugin_list.md new file mode 100644 index 0000000000..8fa59900de --- /dev/null +++ b/docs/cli/thv_ai-plugin_list.md @@ -0,0 +1,44 @@ +--- +title: thv ai-plugin list +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin list` +last_update: + author: autogenerated +slug: thv_ai-plugin_list +mdx: + format: md +--- + +## thv ai-plugin list + +List installed AI-tool plugins + +### Synopsis + +List all currently installed plugins and their status. + +``` +thv ai-plugin list [flags] +``` + +### Options + +``` + --client string Filter by client application + --format string Output format (json, text) (default "text") + --group string Filter by group + -h, --help help for list + --project-root string Project root path for project-scoped plugins + --scope string Filter by scope (user, project) +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_ai-plugin_push.md b/docs/cli/thv_ai-plugin_push.md new file mode 100644 index 0000000000..56af0314f9 --- /dev/null +++ b/docs/cli/thv_ai-plugin_push.md @@ -0,0 +1,39 @@ +--- +title: thv ai-plugin push +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin push` +last_update: + author: autogenerated +slug: thv_ai-plugin_push +mdx: + format: md +--- + +## thv ai-plugin push + +Push a built AI-tool plugin to an OCI registry + +### Synopsis + +Push a previously built plugin artifact to a remote OCI registry. + +``` +thv ai-plugin push [reference] [flags] +``` + +### Options + +``` + -h, --help help for push +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_ai-plugin_uninstall.md b/docs/cli/thv_ai-plugin_uninstall.md new file mode 100644 index 0000000000..156267d1f9 --- /dev/null +++ b/docs/cli/thv_ai-plugin_uninstall.md @@ -0,0 +1,41 @@ +--- +title: thv ai-plugin uninstall +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin uninstall` +last_update: + author: autogenerated +slug: thv_ai-plugin_uninstall +mdx: + format: md +--- + +## thv ai-plugin uninstall + +Uninstall an AI-tool plugin + +### Synopsis + +Remove a previously installed plugin by name. + +``` +thv ai-plugin uninstall [plugin-name] [flags] +``` + +### Options + +``` + -h, --help help for uninstall + --project-root string Project root path for project-scoped plugins + --scope string Scope to uninstall from (user, project) (default "user") +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_ai-plugin_validate.md b/docs/cli/thv_ai-plugin_validate.md new file mode 100644 index 0000000000..6146f0767c --- /dev/null +++ b/docs/cli/thv_ai-plugin_validate.md @@ -0,0 +1,40 @@ +--- +title: thv ai-plugin validate +hide_title: true +description: Reference for ToolHive CLI command `thv ai-plugin validate` +last_update: + author: autogenerated +slug: thv_ai-plugin_validate +mdx: + format: md +--- + +## thv ai-plugin validate + +Validate an AI-tool plugin directory + +### Synopsis + +Check that a plugin definition in the given directory is valid and well-formed. + +``` +thv ai-plugin validate [path] [flags] +``` + +### Options + +``` + --format string Output format (json, text) (default "text") + -h, --help help for validate +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv ai-plugin](thv_ai-plugin.md) - Manage AI-tool plugins + diff --git a/docs/cli/thv_build.md b/docs/cli/thv_build.md index 90b23bf49b..141364a487 100644 --- a/docs/cli/thv_build.md +++ b/docs/cli/thv_build.md @@ -53,10 +53,11 @@ thv build [flags] PROTOCOL [-- ARGS...] ### Options ``` - --dry-run Generate Dockerfile without building (stdout output unless -o is set) (default false) - -h, --help help for build - -o, --output string Write the Dockerfile to the specified file instead of building (default builds an image instead of generating a Dockerfile) - -t, --tag string Name and optionally a tag in the 'name:tag' format for the built image (default generates a unique image name based on the package and transport type) + --build-with stringArray Build-time dependency constraint for protocol scheme builds, interpreted per package ecosystem (uvx://: PEP 508 specifier passed to 'uv tool install --with', e.g. --build-with 'mcp<2'); errors on ecosystems without constraint support (can be specified multiple times) + --dry-run Generate Dockerfile without building (stdout output unless -o is set) (default false) + -h, --help help for build + -o, --output string Write the Dockerfile to the specified file instead of building (default builds an image instead of generating a Dockerfile) + -t, --tag string Name and optionally a tag in the 'name:tag' format for the built image (default generates a unique image name based on the package and transport type) ``` ### Options inherited from parent commands diff --git a/docs/cli/thv_run.md b/docs/cli/thv_run.md index 8a7fb8ecc7..1c8638d4e8 100644 --- a/docs/cli/thv_run.md +++ b/docs/cli/thv_run.md @@ -115,6 +115,7 @@ thv run [flags] SERVER_OR_IMAGE_OR_PROTOCOL [-- ARGS...] --allowed-origins stringArray Exact-match allowlist for the HTTP Origin header (repeatable). Recommended when binding publicly; loopback binds derive a default allowlist automatically, non-loopback binds log a warning when no value is supplied. Example: https://my-mcp.example.com --audit-config string Path to the audit configuration file --authz-config string Path to the authorization configuration file + --build-with stringArray Build-time dependency constraint for protocol scheme builds, interpreted per package ecosystem (uvx://: PEP 508 specifier passed to 'uv tool install --with', e.g. --build-with 'mcp<2'); errors on ecosystems without constraint support (can be specified multiple times) --ca-cert string Path to a custom CA certificate file to use for container builds --enable-audit Enable audit logging with default configuration (default false) --endpoint-prefix string Path prefix to prepend to SSE endpoint URLs (e.g., /playwright) @@ -128,12 +129,12 @@ thv run [flags] SERVER_OR_IMAGE_OR_PROTOCOL [-- ARGS...] --host string Host for the HTTP proxy to listen on (IP or hostname) (default "127.0.0.1") --ignore-globally Load global ignore patterns from ~/.config/toolhive/thvignore (default true) --image-verification string Set image verification mode (warn, enabled, disabled) (default "warn") - --isolate-network Isolate the container network from the host. Use --isolate-network=false to opt out. (default true) + --isolate-network Isolate the container network from the host. Use --isolate-network=false to opt out. Not enforced with --network host or --network none (isolation requires bridge networking). (default true) --jwks-allow-private-ip Allow JWKS/OIDC endpoints on private IP addresses (use with caution) (default false) --jwks-auth-token-file string Path to file containing bearer token for authenticating JWKS/OIDC requests -l, --label stringArray Set labels on the container (format: key=value) --name string Name of the MCP server (default to auto-generated from image) - --network string Connect the container to a network (e.g., 'host' for host networking) + --network string Connect the container to a network (e.g., 'host' for host networking). Note: 'host' and 'none' cannot enforce network isolation, so isolation is dropped for those modes. --oidc-audience string Expected audience for the token --oidc-client-id string OIDC client ID --oidc-client-secret string OIDC client secret (optional, for introspection) @@ -181,6 +182,7 @@ thv run [flags] SERVER_OR_IMAGE_OR_PROTOCOL [-- ARGS...] --secret stringArray Specify a secret to be fetched from the secrets manager and set as an environment variable (format: NAME,target=TARGET) --session-ttl duration Session inactivity timeout (e.g., 30m, 2h); zero uses the default (2h) --stateless Declare the server as stateless (POST-only, no SSE). Use for MCP servers implementing streamable-HTTP stateless mode. + --strict-protocol-validation Reject client requests whose MCP-Protocol-Version header is an unknown/unsupported MCP revision with HTTP 400 (streamable-HTTP proxy only; an absent header is accepted). Off by default: any version is accepted. --target-host string Host to forward traffic to (only applicable to SSE or Streamable HTTP transport) (default "127.0.0.1") --target-port int Port for the container to expose (only applicable to SSE or Streamable HTTP transport) --thv-ca-bundle string Path to CA certificate bundle for ToolHive HTTP operations (JWKS, OIDC discovery, etc.) diff --git a/docs/cli/thv_skill.md b/docs/cli/thv_skill.md index 3a93a62f77..2717cddc63 100644 --- a/docs/cli/thv_skill.md +++ b/docs/cli/thv_skill.md @@ -38,6 +38,8 @@ The skill command provides subcommands to manage skills. * [thv skill install](thv_skill_install.md) - Install a skill * [thv skill list](thv_skill_list.md) - List installed skills * [thv skill push](thv_skill_push.md) - Push a built skill +* [thv skill sync](thv_skill_sync.md) - Restore project skills to match the lock file (experimental) * [thv skill uninstall](thv_skill_uninstall.md) - Uninstall a skill +* [thv skill upgrade](thv_skill_upgrade.md) - Upgrade project skills to newer pinned content (experimental) * [thv skill validate](thv_skill_validate.md) - Validate a skill definition diff --git a/docs/cli/thv_skill_install.md b/docs/cli/thv_skill_install.md index 84ca857c8b..268829aadc 100644 --- a/docs/cli/thv_skill_install.md +++ b/docs/cli/thv_skill_install.md @@ -25,6 +25,7 @@ thv skill install [skill-name] [flags] ### Options ``` + --allow-unsigned Allow installing a project-scoped skill without a verified signature (recorded in the lock file) --clients string Comma-separated target client apps (e.g. claude-code,opencode), or "all" for every available client --force Overwrite existing skill directory --group string Group to add the skill to after installation diff --git a/docs/cli/thv_skill_sync.md b/docs/cli/thv_skill_sync.md new file mode 100644 index 0000000000..8cf0f6beb1 --- /dev/null +++ b/docs/cli/thv_skill_sync.md @@ -0,0 +1,59 @@ +--- +title: thv skill sync +hide_title: true +description: Reference for ToolHive CLI command `thv skill sync` +last_update: + author: autogenerated +slug: thv_skill_sync +mdx: + format: md +--- + +## thv skill sync + +Restore project skills to match the lock file (experimental) + +### Synopsis + +Restore a project's installed skills to match toolhive.lock.yaml. + +Experimental: requires TOOLHIVE_SKILLS_LOCK_ENABLED=true on the ToolHive +server while the lock file feature rolls out. + +Missing or drifted skills are reinstalled at their pinned digest. Use +--check to report drift without installing anything (suitable for CI). +Use --adopt to record lock entries for existing unmanaged installs, and +--prune to remove installs no longer present in the lock file. + +Unless --check is set, sync prompts for confirmation before installing — +skill content is a set of AI-followed instructions. Pass --yes to skip the +prompt (required in non-interactive contexts such as CI). + +``` +thv skill sync [flags] +``` + +### Options + +``` + --adopt Write lock entries for existing unmanaged project-scope installs + --allow-unsigned Allow adopting skills whose signature state cannot be established (recorded as unsigned) + --check Report drift without installing, writing, or removing anything + --clients string Comma-separated target client apps (e.g. claude-code,opencode), or "all" for every available client + --format string Output format (json, text) (default "text") + -h, --help help for sync + --project-root string Project root path (default: auto-detected from the current directory) + --prune Remove installs no longer present in the lock file + --yes Skip the confirmation prompt (required when not running interactively) +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv skill](thv_skill.md) - Manage skills + diff --git a/docs/cli/thv_skill_upgrade.md b/docs/cli/thv_skill_upgrade.md new file mode 100644 index 0000000000..51059c3b51 --- /dev/null +++ b/docs/cli/thv_skill_upgrade.md @@ -0,0 +1,64 @@ +--- +title: thv skill upgrade +hide_title: true +description: Reference for ToolHive CLI command `thv skill upgrade` +last_update: + author: autogenerated +slug: thv_skill_upgrade +mdx: + format: md +--- + +## thv skill upgrade + +Upgrade project skills to newer pinned content (experimental) + +### Synopsis + +Re-resolve a project's lock entries and install newer content where available. + +Experimental: requires TOOLHIVE_SKILLS_LOCK_ENABLED=true on the ToolHive +server while the lock file feature rolls out. + +Skills pinned to an immutable reference (an OCI digest or a full git commit +hash) are reported not-upgradable — there is nothing newer to resolve to. +Use --preview to see what would change without persisting anything (OCI +sources are still fetched into the local artifact store to compare digests), +and --allow-ref-change to permit the artifact moving to a different +repository (a version bump within the same repository is not a change +this guard blocks). +--fail-on-changes evaluates the same plan and never installs: it is a CI +freshness gate. + +Unless --preview is set, upgrade prompts for confirmation before installing — +skill content is a set of AI-followed instructions. Pass --yes to skip the +prompt (required in non-interactive contexts such as CI). + +``` +thv skill upgrade [skill-name...] [flags] +``` + +### Options + +``` + --allow-ref-change Permit the artifact to move to a different repository during upgrade + --allow-signer-change Permit upgrading to an artifact signed by a different identity; the new identity replaces the recorded one + --clients string Comma-separated target client apps (e.g. claude-code,opencode), or "all" for every available client + --fail-on-changes Report what would change without installing anything; a CI freshness gate + --format string Output format (json, text) (default "text") + -h, --help help for upgrade + --preview Report what would change without persisting anything (OCI sources are still fetched to compare digests) + --project-root string Project root path (default: auto-detected from the current directory) + --yes Skip the confirmation prompt (required when not running interactively) +``` + +### Options inherited from parent commands + +``` + --debug Enable debug mode +``` + +### SEE ALSO + +* [thv skill](thv_skill.md) - Manage skills + diff --git a/docs/middleware.md b/docs/middleware.md index 19bb148cd1..93b4b6201b 100644 --- a/docs/middleware.md +++ b/docs/middleware.md @@ -10,15 +10,15 @@ This document primarily covers the middleware system for `thv` and `thv-proxyrun The middleware chain consists of the following components: -1. **Authentication Middleware**: Validates JWT tokens and extracts client identity -2. **Upstream Token Swap Middleware**: Exchanges ToolHive JWTs for upstream IdP tokens (automatic with embedded auth server) -3. **Token Exchange Middleware**: Exchanges JWT tokens for external service tokens via OAuth 2.0 Token Exchange (optional) -4. **MCP Parsing Middleware**: Parses JSON-RPC MCP requests and extracts structured data -5. **Tool Mapping Middleware**: Enables tool filtering and override capabilities through two complementary middleware components that process outgoing `tools/list` responses and incoming `tools/call` requests (optional) -6. **Usage Metrics Middleware**: Collects anonymous usage metrics for ToolHive development (optional) -7. **Telemetry Middleware**: Instruments requests with OpenTelemetry (optional) -8. **Authorization Middleware**: Evaluates Cedar policies to authorize requests (optional) -9. **Audit Middleware**: Logs request events for compliance and monitoring (optional) +1. **Audit Middleware**: Wraps the rest of the chain and logs every request outcome — including rejections from authentication, webhooks, and authorization (optional) +2. **Authentication Middleware**: Validates JWT tokens and extracts client identity +3. **Upstream Token Swap Middleware**: Exchanges ToolHive JWTs for upstream IdP tokens (automatic with embedded auth server) +4. **Token Exchange Middleware**: Exchanges JWT tokens for external service tokens via OAuth 2.0 Token Exchange (optional) +5. **MCP Parsing Middleware**: Parses JSON-RPC MCP requests and extracts structured data +6. **Tool Mapping Middleware**: Enables tool filtering and override capabilities through two complementary middleware components that process outgoing `tools/list` responses and incoming `tools/call` requests (optional) +7. **Usage Metrics Middleware**: Collects anonymous usage metrics for ToolHive development (optional) +8. **Telemetry Middleware**: Instruments requests with OpenTelemetry (optional) +9. **Authorization Middleware**: Evaluates Cedar policies to authorize requests (optional) 10. **Header Forward Middleware**: Injects custom headers into requests to remote MCP servers (optional) 11. **Recovery Middleware**: Catches panics and returns HTTP 500 errors (always present) @@ -31,19 +31,24 @@ Two webhook types are supported: 1. **Mutating webhooks**: Transform the parsed MCP request before later policy evaluation. 2. **Validating webhooks**: Approve or deny the request after mutation has completed. +A mutating webhook therefore steers what authorization (and audit, telemetry, and usage metrics) evaluates, including the feature/operation mapping used to pick a Cedar policy. This is intended: an operator who configures a mutating webhook — including a fail-open one — is deliberately giving it influence over policy input, not introducing an accident. + When configured together, the effective order is: -1. Authentication -2. Token exchange and related auth middleware, when configured -3. MCP parsing -4. Mutating webhooks -5. Validating webhooks -6. Telemetry, authorization, and audit middleware +1. Audit (wraps everything below, so webhook denials are audited) +2. Authentication +3. Token exchange and related auth middleware, when configured +4. MCP parsing +5. Mutating webhooks +6. Validating webhooks +7. Telemetry and authorization middleware Multiple webhook definitions of the same type run in configuration order. When multiple `--webhook-config` files are provided, later files override earlier webhook definitions with the same `name`. Configuration files may be written in YAML or JSON. Duration values such as `timeout` accept strings like `5s`, and omitted timeouts default to `10s`. +When the caller authenticated with an RFC 8693 delegated token, the request payload sent to webhook receivers includes a `delegation` field on the principal object, with the same shape as the `delegation` field documented under [Audit Middleware](#9-audit-middleware) below. + Example: ```bash @@ -55,16 +60,21 @@ Example config files: - [`docs/examples/webhooks.yaml`](examples/webhooks.yaml) - [`docs/examples/webhooks.json`](examples/webhooks.json) +### Known limitations + +- **Stale Modern headers.** A mutating webhook patches only the JSON body, so after it renames a tool the `Mcp-Method`/`Mcp-Name` headers forwarded to the backend still describe the original name. `ValidateHeaderConsistency` (`pkg/mcp/revision.go:512`) requires the header and body to agree and returns a `RequestHeaderMismatchError` (`CodeHeaderMismatch`) if they don't, but today that check is only wired into vMCP (`pkg/vmcp/server/classification.go:72`), which never runs the mutating webhook — so nothing in ToolHive catches this yet. A spec-conformant Modern (2026-07-28) backend will reject the mismatched request itself, so this fails closed rather than silently mis-authorizing. Practical guidance: a mutating webhook must not rename tools on the Modern path. +- **Controls outside the parser see the pre-mutation request.** The tool-call filter (`pkg/mcp/tool_filter.go`) reads the raw request body directly, and the rate limiter runs before the mutating webhook in the chain (see the ordering rules below) — both decide against the request as received, not as mutated. A mutating webhook can therefore rename a call into a tool that `--tools` filtering excluded, and the rate limiter debits the bucket for the requested tool rather than the executed one. This is a known gap, not a regression: the tool filter's position ahead of the MCP parser is deliberate (it needs the raw request), as already noted in the ordering rules below. + ## Architecture Diagram ```mermaid graph TD - A[Incoming MCP Request] --> R[Recovery Middleware] - R --> B[Authentication Middleware] + A[Incoming MCP Request] --> E[Audit Middleware] + E --> B[Authentication Middleware] B --> C[MCP Parsing Middleware] C --> D[Authorization Middleware] - D --> E[Audit Middleware] - E --> F[MCP Server Handler] + D --> R[Recovery Middleware] + R --> F[MCP Server Handler] R --> R1[Catch Panics] R1 --> R2[Log Stack Trace] @@ -101,26 +111,27 @@ graph TD ```mermaid sequenceDiagram participant Client - participant Recovery as Recovery + participant Audit as Audit participant Auth as Authentication participant Parser as MCP Parser participant Authz as Authorization - participant Audit as Audit + participant Recovery as Recovery participant Server as MCP Server - Client->>Recovery: HTTP Request - Note over Recovery: Wraps entire chain to catch panics + Client->>Audit: HTTP Request + Note over Audit: Injects identity / parsed-request holders,
logs the outcome after the inner chain returns + Note over Recovery: Innermost wrapper: catches panics
from the handler and inner middleware - Recovery->>Auth: HTTP Request with JWT + Audit->>Auth: HTTP Request with JWT Auth->>Auth: Validate JWT Token Auth->>Auth: Extract Claims - Note over Auth: Add claims to context + Note over Auth: Add identity to context and holder Auth->>Parser: Request + JWT Claims Parser->>Parser: Parse JSON-RPC Parser->>Parser: Extract MCP Method Parser->>Parser: Extract Resource ID & Arguments - Note over Parser: Add parsed data to context + Note over Parser: Add parsed data to context and holder Parser->>Authz: Request + Parsed MCP Data Authz->>Authz: Get Parsed Data from Context @@ -128,13 +139,18 @@ sequenceDiagram Authz->>Authz: Evaluate Policies alt Authorized - Authz->>Audit: Authorized Request - Audit->>Audit: Extract Event Data - Audit->>Audit: Log Audit Event - Audit->>Server: Process Request - Server->>Client: Response - else Unauthorized - Authz->>Client: 403 Forbidden + Authz->>Server: Process Request + Server->>Audit: Response + Audit->>Audit: Log Audit Event (outcome success) + Audit->>Client: Response + else Authentication fails + Auth->>Audit: 401 Unauthorized + Audit->>Audit: Log Audit Event (outcome denied) + Audit->>Client: 401 Unauthorized + else Denied by policy + Authz->>Audit: 403 Forbidden + Audit->>Audit: Log Audit Event (outcome denied) + Audit->>Client: 403 Forbidden else Panic Occurs Recovery->>Recovery: Log stack trace Recovery->>Client: 500 Internal Server Error @@ -389,6 +405,7 @@ thv config usage-metrics enable - Log structured audit events as JSON - Track request duration and outcome - Support file-based and stdout log destinations +- Wrap the rest of the chain (authentication, webhooks, authorization) so rejected requests are still recorded (outcome `denied`); the identity and parsed MCP data flow back from the inner middlewares via holder carriers **Event Types**: - `mcp_initialize` - Client initialization events @@ -443,6 +460,7 @@ thv run --transport sse --name my-server --audit-config audit.json my-image:late **Important Notes**: - `excludeEventTypes` takes precedence over `eventTypes` +- Requests rejected before the MCP parser runs (e.g. authentication failures) are typed `http_request`, not `mcp_*`. An `eventTypes` allowlist containing only `mcp_*` types will drop those rejection events — include `http_request` to keep them. - When `includeRequestData` or `includeResponseData` is enabled, **`maxDataSize` must be set** (non-zero) for data capture to work - Log files are created with restrictive permissions (0600) for security - Logs are written in newline-delimited JSON format for easy parsing @@ -472,6 +490,15 @@ Audit events are logged as structured JSON objects: "client_name": "my-mcp-client", "client_version": "1.0.0" }, + "delegation": { + "chain": [ + {"iss": "https://auth.example.com", "sub": "agent-client-1"}, + {"sub": "agent-client-2"} + ], + "truncated": false, + "omitted": 0, + "malformed": false + }, "target": { "endpoint": "/messages", "method": "POST", @@ -508,6 +535,35 @@ Audit events are logged as structured JSON objects: - `user`: User display name (from `name` claim, `preferred_username`, or `email`) - `client_name`: MCP client name (from JWT claims) - `client_version`: MCP client version (from JWT claims) +- `delegation`: RFC 8693 delegation chain (toolhive-core's canonical schema), + present only when the caller authenticated with a delegated token (i.e. the + JWT carries an `act` claim) + - `chain`: Delegation hops, outermost (most recent) first — `chain[0]` is + the direct delegate that presented the token. Always an array, never + null; empty when `malformed` is `true` and no hop could be parsed at all + - `iss`: The hop's issuer (from the `act` claim's `iss` member), when + present and a string. Omitted otherwise + - `sub`: Acting party identifier (from the `act` claim's `sub` member), + when present and a string. Omitted otherwise. A non-string `iss` or + `sub` leaves the field out and flags the chain `malformed` instead. + Any other `act` claim members are deliberately never serialized + (data minimization per RFC 8693 §6); per OpenID Connect Core §5.7 the + (`iss`, `sub`) pair is the stable actor identifier + - `truncated`: `true` when the chain exceeded the configured maximum depth + (`maxDelegationDepth` in the audit config, default 10) and inner hops + were dropped; the outermost hops — including `chain[0]`, the current + actor — are always kept. Always present + - `omitted`: Number of well-formed hops dropped due to the depth cap. + Always present; `0` when nothing was dropped + - `malformed`: `true` when the token's `act` claim violated RFC 8693 + conformance somewhere in the chain (non-object `act`, non-string `iss` + or `sub`). Always present. Hops already parsed before the violation are + still reported alongside it. A malformed value sitting immediately past + the depth cap flags `malformed` but is not a countable hop for `omitted` + - `malformedReason`: Low-cardinality label for the first conformance + violation encountered (`act_not_object`, `nested_act_not_object`, + `iss_not_string`, `sub_not_string`). Present only when `malformed` is + `true` - `target`: Information about the operation target - `endpoint`: HTTP endpoint path - `method`: HTTP method @@ -560,8 +616,7 @@ thv run --transport sse --name my-server --audit-config <(echo '{"component":"my - Prevent server crashes from unhandled panics **Behavior**: -- Always added as the outermost middleware wrapper (added last in chain, executes first) -- Catches any panic from the entire middleware chain and MCP handlers +- On the `vmcp` path it is applied explicitly as the outermost wrapper and catches panics from the entire chain. On the `thv`/`thv-proxyrunner` path it is appended last to the config slice, which makes it the innermost wrapper: it catches panics from the proxy handler, but not from middleware that wrap it - Logs error with stack trace using `logger.Errorf` - Returns generic "Internal Server Error" message (no sensitive details exposed) @@ -616,10 +671,14 @@ The middleware chain uses Go's `context.Context` to pass data between components ```mermaid graph LR - A[Request Context] --> B[+ JWT Claims] + A[Request Context] --> E[+ Audit holder carriers] + E --> B[+ JWT Claims] B --> C[+ Parsed MCP Data] C --> D[+ Authorization Result] - D --> E[+ Audit Metadata] + + subgraph "Audit" + E + end subgraph "Authentication" B @@ -632,10 +691,6 @@ graph LR subgraph "Authorization" D end - - subgraph "Audit" - E - end ``` ## Configuration @@ -659,10 +714,10 @@ thv run --transport sse --name my-server --audit-config audit.yaml my-image:late The middleware order is critical and enforced by the system: -1. **Authentication** - Must be first to establish client identity -2. **MCP Parsing** - Must come after authentication to access JWT context -3. **Authorization** - Must come after parsing to access structured MCP data -4. **Audit** - Must be last to capture the complete request lifecycle +1. **Audit** - Wraps the rest of the chain so every request outcome is logged, including authentication failures (401) and policy denials (403, outcome `denied`). The identity and parsed MCP data are published back to it by the inner middlewares via holder carriers. +2. **Authentication** - Establishes client identity +3. **MCP Parsing** - Must come after authentication to access JWT context +4. **Authorization** - Must come after parsing to access structured MCP data ## Error Handling @@ -993,27 +1048,36 @@ func CreateMiddleware(config *types.MiddlewareConfig, runner types.MiddlewareRun The middleware chain execution order is critical and controlled by the order in `PopulateMiddlewareConfigs()` in `pkg/runner/middleware.go`. -1. **Authentication Middleware** (always present) - Validates JWT tokens and extracts claims -2. **Upstream Token Swap Middleware** (if embedded auth server configured) - Swaps ToolHive JWT for upstream IdP token -3. **Token Exchange Middleware** (if enabled) - Exchanges JWT for external service tokens via OAuth 2.0 Token Exchange -4. **Tool Filter Middleware** (if enabled) - Filters available tools in list responses -5. **Tool Call Filter Middleware** (if enabled) - Filters tool call requests -6. **MCP Parser Middleware** (always present) - Parses JSON-RPC MCP requests -7. **Usage Metrics Middleware** (if enabled) - Tracks tool call counts -8. **Telemetry Middleware** (if enabled) - OpenTelemetry instrumentation -9. **Authorization Middleware** (if enabled) - Cedar policy evaluation -10. **Audit Middleware** (if enabled) - Request logging -11. **Header Forward Middleware** (if configured for remote servers) - Injects custom headers -12. **Recovery Middleware** (always present) - Catches panics (outermost wrapper) +1. **Audit Middleware** (if enabled) - Request logging; wraps everything below so every rejection is audited +2. **Authentication Middleware** (always present) - Validates JWT tokens and extracts claims +3. **Upstream Token Swap Middleware** (if embedded auth server configured) - Swaps ToolHive JWT for upstream IdP token +4. **Token Exchange Middleware** (if enabled) - Exchanges JWT for external service tokens via OAuth 2.0 Token Exchange +5. **Tool Filter Middleware** (if enabled) - Filters available tools in list responses +6. **Tool Call Filter Middleware** (if enabled) - Filters tool call requests +7. **MCP Parser Middleware** (always present) - Parses JSON-RPC MCP requests +8. **Rate Limit Middleware** (if configured) - Enforces per-identity/tool limits using the parsed request +9. **Mutating Webhook Middleware** (if configured) - Patches the parsed MCP request and republishes the parse +10. **Validating Webhook Middleware** (if configured) - Approves or denies the (possibly mutated) request +11. **Usage Metrics Middleware** (if enabled) - Tracks tool call counts +12. **Telemetry Middleware** (if enabled) - OpenTelemetry instrumentation +13. **Authorization Middleware** (if enabled) - Cedar policy evaluation +14. **Header Forward Middleware** (if configured for remote servers) - Injects custom headers +15. **Recovery Middleware** (always present) - Catches panics **Important Ordering Rules**: -- Authentication must come first to establish client identity +- Audit wraps the whole chain (directly inside the body-size limit): every request that passes the size cap produces an audit event no matter which middleware rejects it. It does not need to run inside auth or the parser — those publish the identity and parsed MCP data back to it via `auth.IdentityHolder` and `mcp.ParsedRequestHolder`. + - Tradeoff: with `includeRequestData` enabled (default off), audit buffers the request body (up to `maxDataSize`, itself capped by the body-size limit) *before* authentication, so unauthenticated requests that will 401 also pay the buffer cost. This is inherent to auditing rejected requests — the body must be read before the outcome is known — and is bounded per-request. +- Authentication must come before the other middlewares to establish client identity - Upstream Token Swap must come after Authentication (requires `tsid` claim) and before Token Exchange (so it can read the original JWT) - Token Exchange must come after Upstream Swap if both are used (can further transform the upstream IdP token) - Tool filters should come before MCP Parser to operate on raw requests - MCP Parser must come before Authorization (provides structured MCP data) +- Mutating webhooks must come before Validating webhooks and before Authorization, so policy evaluation sees the patched request +- Middleware that rewrites the request body must republish the parsed request via `mcp.RepublishParsedMCPRequest` and refresh `r.ContentLength` — `ParsingMiddleware` deliberately parses only once, so every later consumer (authorization, audit, telemetry, usage metrics) reads the cached parse rather than re-reading the body - Header Forward executes close to the backend handler (innermost position) -- Recovery is always last in config (so it executes first as outermost wrapper) +- Recovery is always last in config, making it the innermost wrapper (the chain wraps in reverse config order, so the first entry is the outermost and runs first) +- Body-size limit and Origin validation stay OUTSIDE audit: oversized bodies must be rejected before audit buffers request data, and origin validation is a pre-auth DNS-rebind guard. Their rejections (413/403) are the only ones not audited. +- The list above (steps 1-15) describes the operator/proxyrunner path (`PopulateMiddlewareConfigs` in `pkg/runner/middleware.go`). The CLI flag path (`WithMiddlewareFromFlags` in `pkg/runner/config_builder.go`) has no rate limiting at all, and orders Usage Metrics before the webhooks rather than after (see the comment at `pkg/runner/config_builder.go:700`) — both paths still run Mutating before Validating webhooks and both before Authorization. ### Custom Authorization Policies diff --git a/docs/operator/composite-tools-quick-reference.md b/docs/operator/composite-tools-quick-reference.md index 719026a293..640350ee00 100644 --- a/docs/operator/composite-tools-quick-reference.md +++ b/docs/operator/composite-tools-quick-reference.md @@ -22,6 +22,23 @@ spec: description: Description of the parameter required: false + annotations: # Optional: explicit MCP tool annotations + title: My Workflow # Derivation runs at advertise time, not CRD + # admission. The floor is fail-closed: with + # one or more tool steps it is always non-nil, + # and any step with nil/unknown annotations + # taints it to readOnly=false, + # destructive=true, openWorld=true. + readOnlyHint: false # ⚠ An explicit readOnlyHint: true here would + # CONTRADICT the floor whenever a step tool + # does not declare readOnlyHint: true (the + # common case), dropping the composite tool + # from tools/list at runtime with a warning. + # thv vmcp validate does NOT catch this — it + # surfaces only in the vMCP logs. Prefer + # leaving readOnlyHint unset (derive it) or + # declaring a value no safer than the floor. + steps: # Required: workflow steps - id: step1 type: tool # tool|elicitation|forEach @@ -284,6 +301,7 @@ steps: - ✅ forEach maxParallel: 50 (hard cap), defaults to DAG maxParallel (10) - ✅ forEach inner step must be type `tool` (no nested forEach or elicitation) - ✅ forEach `itemVar` cannot be `index` (reserved) +- ✅ Annotations must not contradict the step-derived safety floor (e.g. `readOnlyHint: true` when a step is not read-only) — a contradicting composite tool is **dropped from `tools/list` at runtime** with a warning in the vMCP logs (more-conservative hints, e.g. `readOnlyHint: false`, are allowed) **Note**: Max retry and max steps limits are currently enforced at runtime. Future work may add CRD-level validation (`+kubebuilder:validation:MaxItems=100`) and webhook validation to fail at submission time rather than execution time. diff --git a/docs/operator/crd-api.md b/docs/operator/crd-api.md index 52cfb9ae45..651753ae23 100644 --- a/docs/operator/crd-api.md +++ b/docs/operator/crd-api.md @@ -42,6 +42,7 @@ _Appears in:_ | `includeResponseData` _boolean_ | IncludeResponseData determines whether to include response data in audit logs. | false | Optional: \{\}
| | `detectApplicationErrors` _boolean_ | DetectApplicationErrors controls whether the audit middleware inspects
JSON-RPC response bodies for application-level errors when the HTTP
status code indicates success (2xx). When enabled, a small prefix of
the response body is buffered to detect JSON-RPC error fields,
independent of the IncludeResponseData setting. | true | Optional: \{\}
| | `maxDataSize` _integer_ | MaxDataSize limits the size of request/response data included in audit logs (in bytes). | 1024 | Optional: \{\}
| +| `maxDelegationDepth` _integer_ | MaxDelegationDepth caps how many nested RFC 8693 "act" entries are
recorded in an audit event's delegation chain. Deeper chains are
truncated (marked with truncated=true). Defaults to 10 when unset. | 10 | Minimum: 1
Optional: \{\}
| | `logFile` _string_ | LogFile specifies the file path for audit logs. If empty, logs to stdout. | | Optional: \{\}
| @@ -262,8 +263,8 @@ _Appears in:_ AggregationConfig defines tool aggregation, filtering, and conflict resolution strategies. Tool Visibility vs Routing: - - ExcludeAllTools, per-workload ExcludeAll, and Filter control which tools are - advertised to MCP clients (visible in tools/list responses). + - ExcludeAllTools, DefaultToolVisibility, per-workload ExcludeAll, and Filter control + which tools are advertised to MCP clients (visible in tools/list responses). - ALL backend tools remain available in the internal routing table, allowing composite tools to call hidden backend tools. - This enables curated experiences where raw backend tools are hidden from @@ -280,6 +281,7 @@ _Appears in:_ | `conflictResolutionConfig` _[vmcp.config.ConflictResolutionConfig](#vmcpconfigconflictresolutionconfig)_ | ConflictResolutionConfig provides configuration for the chosen strategy. | | Optional: \{\}
| | `tools` _[vmcp.config.WorkloadToolConfig](#vmcpconfigworkloadtoolconfig) array_ | Tools defines per-workload tool filtering and overrides. | | Optional: \{\}
| | `excludeAllTools` _boolean_ | ExcludeAllTools hides all backend tools from MCP clients when true.
Hidden tools are NOT advertised in tools/list responses, but they ARE
available in the routing table for composite tools to use.
This enables the use case where you want to hide raw backend tools from
direct client access while exposing curated composite tool workflows. | | Optional: \{\}
| +| `defaultToolVisibility` _[vmcp.config.DefaultToolVisibility](#vmcpconfigdefaulttoolvisibility)_ | DefaultToolVisibility controls whether a backend with NO entry in Tools has its
tools advertised to MCP clients.
- allow (default): every tool from an unlisted backend is advertised, so
adding a workload to the group exposes it without further configuration.
- deny: an unlisted backend contributes no tools, so only backends named in
Tools are advertised. Use this when the set of exposed tools must be
enumerated deliberately rather than inherited from group membership.
A backend that DOES have a Tools entry is unaffected by this setting: the
entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like
every other visibility setting here, this controls advertising only — hidden
tools remain in the routing table for composite tools (see the type doc).
This gates TOOLS only. An unlisted backend's resources, resource templates,
and prompts are still advertised under deny; mergeResources/mergePrompts have
no equivalent check. | | Enum: [allow deny]
Optional: \{\}
| #### vmcp.config.AuthzConfig @@ -365,6 +367,7 @@ _Appears in:_ | `timeout` _[vmcp.config.Duration](#vmcpconfigduration)_ | Timeout is the maximum workflow execution time. | | Pattern: `^([0-9]+(\.[0-9]+)?(ns\|us\|µs\|ms\|s\|m\|h))+$`
Type: string
| | `steps` _[vmcp.config.WorkflowStepConfig](#vmcpconfigworkflowstepconfig) array_ | Steps are the workflow steps to execute. | | | | `output` _[vmcp.config.OutputConfig](#vmcpconfigoutputconfig)_ | Output defines the structured output schema for this workflow.
If not specified, the workflow returns the last step's output (backward compatible). | | Optional: \{\}
| +| `annotations` _[vmcp.config.ToolAnnotationsOverride](#vmcpconfigtoolannotationsoverride)_ | Annotations declares MCP tool annotations for the composite tool.
Annotation derivation runs at ADVERTISE TIME (when tools/list is served
and the backend tools are aggregated), not at CRD admission — thv vmcp
validate does NOT check annotation contradictions. The derived floor is
fail-closed: when the workflow has one or more tool steps the floor is
always non-nil, and any step whose annotations are nil/unknown taints the
floor conservatively (readOnly=false, destructive=true, openWorld=true).
A workflow with no tool steps (e.g. only elicitation) has no floor.
When nil, annotations are derived from the annotations of the backend tools
referenced by the workflow's steps (e.g. readOnlyHint is true only when every
step tool is read-only). When set, the values are an explicit author
declaration merged over the derived floor — subject to a safety-floor
guardrail that drops the composite tool (with a warning naming the offending
step tools) if an explicit hint would make the tool look safer than its
steps allow. | | Optional: \{\}
| #### vmcp.config.CompositeToolRef @@ -414,7 +417,7 @@ _Appears in:_ | `metadata` _object (keys:string, values:string)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | | | `telemetry` _[pkg.telemetry.Config](#pkgtelemetryconfig)_ | Telemetry configures OpenTelemetry-based observability for the Virtual MCP server
including distributed tracing, OTLP metrics export, and Prometheus metrics endpoint.
Deprecated (Kubernetes operator only): When deploying via the operator, use
VirtualMCPServer.spec.telemetryConfigRef to reference a shared MCPTelemetryConfig
resource instead. This field remains valid for standalone (non-operator) deployments. | | Optional: \{\}
| | `audit` _[pkg.audit.Config](#pkgauditconfig)_ | Audit configures audit logging for the Virtual MCP server.
When present, audit logs include MCP protocol operations.
See audit.Config for available configuration options. | | Optional: \{\}
| -| `optimizer` _[vmcp.config.OptimizerConfig](#vmcpconfigoptimizerconfig)_ | Optimizer configures the MCP optimizer for context optimization on large toolsets.
When enabled, vMCP exposes only find_tool and call_tool operations to clients
instead of all backend tools directly. This reduces token usage by allowing
LLMs to discover relevant tools on demand rather than receiving all tool definitions. | | Optional: \{\}
| +| `optimizer` _[vmcp.config.OptimizerConfig](#vmcpconfigoptimizerconfig)_ | Optimizer configures the MCP optimizer for context optimization on large toolsets.
When enabled, vMCP exposes only find_tool and call_tool operations to clients
instead of all backend tools directly. This reduces token usage by allowing
LLMs to discover relevant tools on demand rather than receiving all tool definitions.
Enabling the optimizer currently pins this instance to MCP 2025-11-25: the
find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28)
clients are negotiated down to the Legacy revision (see
pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089). | | Optional: \{\}
| | `codeMode` _[vmcp.config.CodeModeConfig](#vmcpconfigcodemodeconfig)_ | CodeMode configures vMCP code mode: server-side execution of Starlark scripts that
orchestrate multiple backend tool calls in a single request via the execute_tool_script
virtual tool. When enabled, execute_tool_script is advertised alongside the backend
tools; a script's inner tool calls are authorized individually, so a script can only
reach tools the caller is already permitted to use. Disabled by default. | | Optional: \{\}
| | `sessionStorage` _[vmcp.config.SessionStorageConfig](#vmcpconfigsessionstorageconfig)_ | SessionStorage configures session storage for stateful horizontal scaling.
When provider is "redis", the operator injects Redis connection parameters
(address, db, keyPrefix) here. The Redis password is provided separately via
the THV_SESSION_REDIS_PASSWORD environment variable. | | Optional: \{\}
| | `sessionTTL` _[vmcp.config.Duration](#vmcpconfigduration)_ | SessionTTL configures the idle TTL for vMCP sessions.
When omitted, the server default is used. | | Pattern: `^([0-9]+(\.[0-9]+)?(ns\|us\|µs\|ms\|s\|m\|h))+$`
Type: string
Optional: \{\}
| @@ -435,8 +438,27 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `prefixFormat` _string_ | PrefixFormat defines the prefix format for the "prefix" strategy.
Supports placeholders: \{workload\}, \{workload\}_, \{workload\}. | \{workload\}_ | Optional: \{\}
| -| `priorityOrder` _string array_ | PriorityOrder defines the workload priority order for the "priority" strategy. | | Optional: \{\}
| +| `prefixFormat` _string_ | PrefixFormat defines the prefix format for the "prefix" tool strategy
and for backend-prefixed prompt names (the default for every prompt;
under the "priority" strategy, backends listed in priorityOrder keep
their own prompt names).
Supports placeholders: \{workload\}, \{workload\}_, \{workload\}. | \{workload\}_ | Optional: \{\}
| +| `priorityOrder` _string array_ | PriorityOrder defines the workload priority order for the "priority" strategy.
Listed workloads also keep their own prompt names (unlisted workloads'
prompts stay backend-prefixed). | | Optional: \{\}
| + + +#### vmcp.config.DefaultToolVisibility + +_Underlying type:_ _string_ + +DefaultToolVisibility names the advertising default applied to backends with no +per-workload Tools entry. The zero value is "" (unset), which behaves as +DefaultToolVisibilityAllow so existing configs keep today's behavior. + + + +_Appears in:_ +- [vmcp.config.AggregationConfig](#vmcpconfigaggregationconfig) + +| Field | Description | +| --- | --- | +| `allow` | DefaultToolVisibilityAllow advertises every tool from a backend with no Tools
entry. This is the default and matches pre-DefaultToolVisibility behavior.
| +| `deny` | DefaultToolVisibilityDeny advertises no tools from a backend with no Tools entry.
| @@ -557,6 +579,7 @@ _Appears in:_ | `protectedResourceAllowPrivateIp` _boolean_ | ProtectedResourceAllowPrivateIP allows protected resource endpoint on private IP addresses
Use with caution - only enable for trusted internal IDPs or testing | | | | `jwksAllowPrivateIp` _boolean_ | JwksAllowPrivateIP allows OIDC discovery and JWKS fetches to private IP addresses.
Enable when the embedded auth server runs on a loopback address and
the OIDC middleware needs to fetch its JWKS from that address.
Use with caution - only enable for trusted internal IDPs or testing. | | | | `insecureAllowHttp` _boolean_ | InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing
WARNING: This is insecure and should NEVER be used in production | | | +| `caBundlePath` _string_ | CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle
used when the OIDC middleware performs HTTPS requests to the issuer
(OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle
at this path is added to the trust store used for verifying the issuer's
TLS certificate. Typically populated by the Kubernetes operator from
MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster
service-account CA when using Kubernetes service-account auth. | | Optional: \{\}
| #### vmcp.config.OperationalConfig @@ -763,7 +786,9 @@ All fields use pointers so nil means "don't override" while zero values _Appears in:_ +- [vmcp.config.CompositeToolConfig](#vmcpconfigcompositetoolconfig) - [vmcp.config.ToolOverride](#vmcpconfigtooloverride) +- [api.v1beta1.VirtualMCPCompositeToolDefinitionSpec](#apiv1beta1virtualmcpcompositetooldefinitionspec) @@ -2308,7 +2333,7 @@ _Appears in:_ | `jwksAuthTokenPath` _string_ | JWKSAuthTokenPath is the path to file containing bearer token for JWKS/OIDC requests | | Optional: \{\}
| | `jwksAllowPrivateIP` _boolean_ | JWKSAllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses.
Note: at runtime, if either JWKSAllowPrivateIP or ProtectedResourceAllowPrivateIP
is true, private IPs are allowed for all OIDC HTTP requests (JWKS, discovery, introspection). | false | Optional: \{\}
| | `protectedResourceAllowPrivateIP` _boolean_ | ProtectedResourceAllowPrivateIP allows protected resource endpoint on private IP addresses.
Note: at runtime, if either ProtectedResourceAllowPrivateIP or JWKSAllowPrivateIP
is true, private IPs are allowed for all OIDC HTTP requests (JWKS, discovery, introspection). | false | Optional: \{\}
| -| `insecureAllowHTTP` _boolean_ | InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing.
WARNING: This is insecure and should NEVER be used in production. | false | Optional: \{\}
| +| `insecureAllowHTTP` _boolean_ | InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuer and JWKS URLs for development/testing.
WARNING: This is insecure and should NEVER be used in production. | false | Optional: \{\}
| #### api.v1beta1.KubernetesServiceAccountOIDCConfig @@ -2432,8 +2457,6 @@ _Appears in:_ | `conditions` _[Condition](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#condition-v1-meta) array_ | Conditions represent the latest available observations of the MCPAuthzConfig's state | | Optional: \{\}
| | `observedGeneration` _integer_ | ObservedGeneration is the most recent generation observed for this MCPAuthzConfig. | | Optional: \{\}
| | `configHash` _string_ | ConfigHash is a hash of the current configuration for change detection | | Optional: \{\}
| -| `referenceCount` _integer_ | ReferenceCount is the number of workloads referencing this config. | | Optional: \{\}
| -| `referencingWorkloads` _[api.v1beta1.WorkloadReference](#apiv1beta1workloadreference) array_ | ReferencingWorkloads is a list of workload resources that reference this MCPAuthzConfig.
Each entry identifies the workload by kind and name. The map key is the
(kind, name) pair so two workloads of different kinds that share a name
(e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct
entries rather than colliding under merge-patch semantics. | | Optional: \{\}
| #### api.v1beta1.MCPExternalAuthConfig @@ -2523,8 +2546,6 @@ _Appears in:_ | `conditions` _[Condition](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#condition-v1-meta) array_ | Conditions represent the latest available observations of the MCPExternalAuthConfig's state | | Optional: \{\}
| | `observedGeneration` _integer_ | ObservedGeneration is the most recent generation observed for this MCPExternalAuthConfig.
It corresponds to the MCPExternalAuthConfig's generation, which is updated on mutation by the API Server. | | Optional: \{\}
| | `configHash` _string_ | ConfigHash is a hash of the current configuration for change detection | | Optional: \{\}
| -| `referenceCount` _integer_ | ReferenceCount is the number of workloads referencing this config. | | Optional: \{\}
| -| `referencingWorkloads` _[api.v1beta1.WorkloadReference](#apiv1beta1workloadreference) array_ | ReferencingWorkloads is a list of workload resources that reference this MCPExternalAuthConfig.
Each entry identifies the workload by kind and name. The map key is the
(kind, name) pair so two workloads of different kinds that share a name
(e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct
entries rather than colliding under merge-patch semantics. | | Optional: \{\}
| #### api.v1beta1.MCPGroup @@ -2768,8 +2789,6 @@ _Appears in:_ | `conditions` _[Condition](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#condition-v1-meta) array_ | Conditions represent the latest available observations of the MCPOIDCConfig's state | | Optional: \{\}
| | `observedGeneration` _integer_ | ObservedGeneration is the most recent generation observed for this MCPOIDCConfig. | | Optional: \{\}
| | `configHash` _string_ | ConfigHash is a hash of the current configuration for change detection | | Optional: \{\}
| -| `referenceCount` _integer_ | ReferenceCount is the number of workloads referencing this config. | | Optional: \{\}
| -| `referencingWorkloads` _[api.v1beta1.WorkloadReference](#apiv1beta1workloadreference) array_ | ReferencingWorkloads is a list of workload resources that reference this MCPOIDCConfig.
Each entry identifies the workload by kind and name. The map key is the
(kind, name) pair so two workloads of different kinds that share a name
(e.g., an MCPServer "foo" and a VirtualMCPServer "foo") are distinct
entries rather than colliding under merge-patch semantics. | | Optional: \{\}
| #### api.v1beta1.MCPRegistry @@ -3350,7 +3369,6 @@ _Appears in:_ | `conditions` _[Condition](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#condition-v1-meta) array_ | Conditions represent the latest available observations of the MCPTelemetryConfig's state | | Optional: \{\}
| | `observedGeneration` _integer_ | ObservedGeneration is the most recent generation observed for this MCPTelemetryConfig. | | Optional: \{\}
| | `configHash` _string_ | ConfigHash is a hash of the current configuration for change detection | | Optional: \{\}
| -| `referencingWorkloads` _[api.v1beta1.WorkloadReference](#apiv1beta1workloadreference) array_ | ReferencingWorkloads lists workloads that reference this MCPTelemetryConfig | | Optional: \{\}
| #### api.v1beta1.MCPTelemetryOTelConfig @@ -3463,8 +3481,6 @@ _Appears in:_ | `conditions` _[Condition](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#condition-v1-meta) array_ | Conditions represent the latest available observations of the MCPToolConfig's state | | Optional: \{\}
| | `observedGeneration` _integer_ | ObservedGeneration is the most recent generation observed for this MCPToolConfig.
It corresponds to the MCPToolConfig's generation, which is updated on mutation by the API Server. | | Optional: \{\}
| | `configHash` _string_ | ConfigHash is a hash of the current configuration for change detection | | Optional: \{\}
| -| `referenceCount` _integer_ | ReferenceCount is the number of workloads referencing this config. | | Optional: \{\}
| -| `referencingWorkloads` _[api.v1beta1.WorkloadReference](#apiv1beta1workloadreference) array_ | ReferencingWorkloads is a list of workload resources that reference this MCPToolConfig.
Each entry identifies the workload by kind and name. | | Optional: \{\}
| #### api.v1beta1.MCPWebhookConfigSpec @@ -3498,7 +3514,6 @@ MCPWebhookConfigStatus defines the observed state of MCPWebhookConfig | `conditions` _[Condition](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#condition-v1-meta) array_ | Conditions represent the latest available observations | | Optional: \{\}
| | `observedGeneration` _integer_ | ObservedGeneration is the last observed generation corresponding to the current status | | Optional: \{\}
| | `configHash` _string_ | ConfigHash is a hash of the spec, used for detecting changes | | Optional: \{\}
| -| `referencingWorkloads` _[api.v1beta1.WorkloadReference](#apiv1beta1workloadreference) array_ | ReferencingWorkloads is a list of workload resources that reference this MCPWebhookConfig.
Each entry identifies the workload by kind and name. | | Optional: \{\}
| #### api.v1beta1.ModelCacheConfig @@ -4456,6 +4471,7 @@ _Appears in:_ | `timeout` _[vmcp.config.Duration](#vmcpconfigduration)_ | Timeout is the maximum workflow execution time. | | Pattern: `^([0-9]+(\.[0-9]+)?(ns\|us\|µs\|ms\|s\|m\|h))+$`
Type: string
| | `steps` _[vmcp.config.WorkflowStepConfig](#vmcpconfigworkflowstepconfig) array_ | Steps are the workflow steps to execute. | | | | `output` _[vmcp.config.OutputConfig](#vmcpconfigoutputconfig)_ | Output defines the structured output schema for this workflow.
If not specified, the workflow returns the last step's output (backward compatible). | | Optional: \{\}
| +| `annotations` _[vmcp.config.ToolAnnotationsOverride](#vmcpconfigtoolannotationsoverride)_ | Annotations declares MCP tool annotations for the composite tool.
Annotation derivation runs at ADVERTISE TIME (when tools/list is served
and the backend tools are aggregated), not at CRD admission — thv vmcp
validate does NOT check annotation contradictions. The derived floor is
fail-closed: when the workflow has one or more tool steps the floor is
always non-nil, and any step whose annotations are nil/unknown taints the
floor conservatively (readOnly=false, destructive=true, openWorld=true).
A workflow with no tool steps (e.g. only elicitation) has no floor.
When nil, annotations are derived from the annotations of the backend tools
referenced by the workflow's steps (e.g. readOnlyHint is true only when every
step tool is read-only). When set, the values are an explicit author
declaration merged over the derived floor — subject to a safety-floor
guardrail that drops the composite tool (with a warning naming the offending
step tools) if an explicit hint would make the tool look safer than its
steps allow. | | Optional: \{\}
| #### api.v1beta1.VirtualMCPCompositeToolDefinitionStatus @@ -4689,27 +4705,6 @@ _Appears in:_ | `insecureSkipVerify` _boolean_ | InsecureSkipVerify disables server certificate verification.
WARNING: This should only be used for development/testing and not in production environments. | | Optional: \{\}
| -#### api.v1beta1.WorkloadReference - - - -WorkloadReference identifies a workload that references a shared configuration resource. -Namespace is implicit — cross-namespace references are not supported. - - - -_Appears in:_ -- [api.v1beta1.MCPAuthzConfigStatus](#apiv1beta1mcpauthzconfigstatus) -- [api.v1beta1.MCPExternalAuthConfigStatus](#apiv1beta1mcpexternalauthconfigstatus) -- [api.v1beta1.MCPOIDCConfigStatus](#apiv1beta1mcpoidcconfigstatus) -- [api.v1beta1.MCPTelemetryConfigStatus](#apiv1beta1mcptelemetryconfigstatus) -- [api.v1beta1.MCPToolConfigStatus](#apiv1beta1mcptoolconfigstatus) -- [api.v1beta1.MCPWebhookConfigStatus](#apiv1beta1mcpwebhookconfigstatus) - -| Field | Description | Default | Validation | -| --- | --- | --- | --- | -| `kind` _string_ | Kind is the type of workload resource | | Enum: [MCPServer VirtualMCPServer MCPRemoteProxy]
Required: \{\}
| -| `name` _string_ | Name is the name of the workload resource | | MinLength: 1
Required: \{\}
| #### api.v1beta1.XAASpec @@ -4757,6 +4752,14 @@ _Appears in:_ + + + + + + + + @@ -4811,6 +4814,31 @@ This type is shared with the Kubernetes operator CRD (VirtualMCPServer.Status.Di | `circuitBreakerState` _string_ | CircuitBreakerState is the current circuit breaker state (closed, open, half-open).
Empty when circuit breaker is disabled or not configured. | | Enum: [closed open half-open]
Optional: \{\}
| | `circuitLastChanged` _[Time](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#time-v1-meta)_ | CircuitLastChanged is the timestamp when the circuit breaker state last changed.
Empty when circuit breaker is disabled or has never changed state. | | Optional: \{\}
| | `consecutiveFailures` _integer_ | ConsecutiveFailures is the current count of consecutive health check failures.
Resets to 0 when the backend becomes healthy again. | | Optional: \{\}
| +| `mcpRevision` _string_ | MCPRevision is the backend's negotiated MCP protocol revision
("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed. | | Optional: \{\}
| + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/operator/storage-version-migration.md b/docs/operator/storage-version-migration.md index 36e4cc6c84..794b0bfefc 100644 --- a/docs/operator/storage-version-migration.md +++ b/docs/operator/storage-version-migration.md @@ -2,7 +2,7 @@ The ToolHive operator ships a `StorageVersionMigrator` controller that keeps every ToolHive CRD's `status.storedVersions` list clean, so a future operator release can drop deprecated API versions (e.g. `v1alpha1`) without orphaning objects in etcd. -The controller is **opt-in** via the `operator.features.storageVersionMigrator` chart value (or the `TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR` environment variable if you inject it directly via `operator.env`). +The controller is **enabled by default** via the `operator.features.storageVersionMigrator` chart value (or the `TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR` environment variable if you inject it directly via `operator.env`). Set the value to `false` to opt out. ## Why this exists @@ -71,18 +71,31 @@ Then `task operator-manifests` to regenerate the CRD YAML. CI verifies the marke ## Enabling the controller -Set the Helm feature flag at install or upgrade time: +The controller is enabled by default — `operator.features.storageVersionMigrator` defaults to `true`, which sets `TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR=true` on the operator Deployment and registers the reconciler with the manager. No action is required to use it. + +To opt out, set the Helm feature flag to `false` at install or upgrade time: ```yaml operator: features: - storageVersionMigrator: true + storageVersionMigrator: false ``` -This sets `TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR=true` on the operator Deployment and registers the reconciler with the manager. - Once enabled, the controller is dormant on CRDs whose `storedVersions` already equals `[]` — most of the time, most CRDs. It only does meaningful work when a CRD's stored-versions list is dirty (typically right after a graduation release). +### Requires cluster-scoped RBAC + +The migrator only works when the operator runs cluster-scoped (`operator.rbac.scope=cluster`, the default). It watches cluster-scoped `CustomResourceDefinition` objects and re-stores custom resources across **all** namespaces — neither is possible for a namespace-scoped operator, which gets only per-namespace `RoleBindings` and a namespace-restricted manager cache. To prevent a silently wedged operator, the chart **fails the render** if `storageVersionMigrator: true` is combined with `operator.rbac.scope=namespace`: + +``` +operator.features.storageVersionMigrator requires operator.rbac.scope=cluster: the +StorageVersionMigrator controller watches cluster-scoped CustomResourceDefinitions and +re-stores resources across all namespaces, which a namespace-scoped operator cannot do. +Set operator.features.storageVersionMigrator=false for namespace-scoped installs. +``` + +Namespace-scoped installs must set `storageVersionMigrator: false` and handle storage-version cleanup by other means — e.g. running the standalone [`kube-storage-version-migrator`](https://github.com/kubernetes-sigs/kube-storage-version-migrator), which is a cluster-scoped component with its own identity, before any version-removal release. + ## Per-CRD emergency escape hatch Removing the label on a live cluster excludes that single CRD from migration immediately: @@ -98,10 +111,10 @@ Intended for incident response only. If you deploy the operator via GitOps (Argo The `StorageVersionMigrator` must have run against your cluster *before* an operator release that drops a deprecated CRD version ships. The typical sequence is: -1. **Release N**: both versions served, newer version is storage, `StorageVersionMigrator` available (opt-in via the chart flag). Operators that enable the migrator have their `storedVersions` trimmed during this deprecation window. +1. **Release N**: both versions served, newer version is storage, `StorageVersionMigrator` enabled by default. Operators running the migrator have their `storedVersions` trimmed during this deprecation window. 2. **Release N+1+**: the deprecated version is removed from `spec.versions`. Because every cluster that enabled the migrator already has clean `storedVersions`, the CRD update applies. -> **⚠ Skip-a-version upgrade trap.** If your cluster upgrades directly from a pre-migrator release to the version-removal release without ever running an intermediate release that runs the migrator, the helm upgrade will **fail** when the API server refuses to remove the deprecated version from `spec.versions`. To recover: deploy [kube-storage-version-migrator](https://github.com/kubernetes-sigs/kube-storage-version-migrator) once to clean `storedVersions`, then retry the upgrade. To avoid the trap entirely, install each release in sequence, and enable `storageVersionMigrator: true` at least one release before any version-removal release. +> **⚠ Skip-a-version upgrade trap.** If your cluster upgrades directly from a pre-migrator release to the version-removal release without ever running an intermediate release that runs the migrator, the helm upgrade will **fail** when the API server refuses to remove the deprecated version from `spec.versions`. To recover: deploy [kube-storage-version-migrator](https://github.com/kubernetes-sigs/kube-storage-version-migrator) once to clean `storedVersions`, then retry the upgrade. To avoid the trap entirely, install each release in sequence, and keep `storageVersionMigrator` enabled (the default) for at least one release before any version-removal release. ## Verification diff --git a/docs/operator/upgrade-guide/README.md b/docs/operator/upgrade-guide/README.md index 4b0b0e3911..a7fa7f6e0c 100644 --- a/docs/operator/upgrade-guide/README.md +++ b/docs/operator/upgrade-guide/README.md @@ -11,6 +11,14 @@ This guide is the canonical reproducible verification for the migrator. Companio - [`docs/operator/storage-version-migration.md`](../storage-version-migration.md) — reference docs for the controller itself (label contract, opt-in model, mechanism). - [Issue #4969](https://github.com/stacklok/toolhive/issues/4969) — the motivating problem. +> **⚠ Upgrade break for namespace-scoped installs.** `operator.features.storageVersionMigrator` now defaults to `true`. The chart hard-fails the render when the migrator is enabled with `operator.rbac.scope=namespace`, because the controller requires cluster-wide CRD access that a namespace-scoped operator does not have. If you run the operator namespace-scoped and never explicitly set this flag, **`helm upgrade` will fail** with: +> +> ``` +> operator.features.storageVersionMigrator requires operator.rbac.scope=cluster +> ``` +> +> **Before upgrading**, add `operator.features.storageVersionMigrator=false` to your Helm values (or `--set` flags) for any namespace-scoped install. Cluster-scoped installs (the default) are unaffected. + ## Prerequisites - `kind`, `kubectl`, `helm`, `ko`, `task` on PATH (`go install github.com/google/ko@latest` for ko) @@ -211,10 +219,10 @@ kubectl rollout status deployment -n toolhive-system --timeout=180s # Confirm flag is now true — read off the Deployment spec directly (see step 5 # for why a pod-based check races with the old pod's Terminating state). -# NOTE: the flag must be set explicitly. The chart default is -# storageVersionMigrator=false (the feature is opt-in), and `helm upgrade` does -# not reuse the previous release's --set values, so omitting it here renders the -# env var as false and the migrator never runs. +# NOTE: this is set explicitly here only because step 5 explicitly disabled it +# (and `helm upgrade` does not reuse the previous release's --set values). The +# chart default is storageVersionMigrator=true, so a normal install/upgrade that +# omits the flag entirely already runs the migrator. kubectl get deploy -n toolhive-system toolhive-operator \ -o jsonpath='{.spec.template.spec.containers[0].env[?(@.name=="TOOLHIVE_ENABLE_STORAGE_VERSION_MIGRATOR")].value}' # expected: true diff --git a/docs/operator/virtualmcpcompositetooldefinition-guide.md b/docs/operator/virtualmcpcompositetooldefinition-guide.md index 14e47a4823..6921f823e7 100644 --- a/docs/operator/virtualmcpcompositetooldefinition-guide.md +++ b/docs/operator/virtualmcpcompositetooldefinition-guide.md @@ -347,6 +347,77 @@ spec: - `abort`: Stop on first failure (default) - `continue`: Execute all steps regardless of failures +### Annotations + +Declare MCP tool annotations for the composite tool. Annotations are behavioral +hints advertised to MCP clients in `tools/list` responses. + +```yaml +spec: + name: user_report + description: Generate a read-only user report + + annotations: + title: User Report + readOnlyHint: true + + steps: + - id: get_user + tool: users.get +``` + +**Fields** (all optional; all hints are booleans): + +| Field | Type | Default | Semantics | +|-------------------|--------|-------------------|-----------| +| `title` | string | derived (none) | Human-readable title. Explicit value wins when set. | +| `readOnlyHint` | bool | derived from steps | The composite tool does not modify its environment. | +| `destructiveHint` | bool | derived from steps | The composite tool may perform destructive updates. | +| `idempotentHint` | bool | none | Repeated calls with the same arguments have no additional effect. Never derived — only advertised when explicitly declared. | +| `openWorldHint` | bool | derived from steps | The composite tool interacts with external entities. | + +**When `annotations` is omitted**, annotations are derived at advertise time +from the annotations of the backend tools referenced by the workflow's steps +(including `forEach` inner steps): + +- `readOnlyHint`: **AND** across steps — `true` only when every step tool + declares `readOnlyHint: true`; any step that is nil or `false` makes it `false`. +- `destructiveHint`: **OR** across steps — `true` when any step declares it + `true`, or when any step's annotations are unknown. +- `openWorldHint`: **OR** across steps — same rule as `destructiveHint`. +- `idempotentHint`: never derived (left unset). + +> **Client compatibility note:** existing composite tools that previously +> advertised no annotations now advertise the fail-closed floor +> (`destructiveHint: true`, `openWorldHint: true`, `readOnlyHint: false`) when +> they have tool steps. That matches the MCP spec defaults for omitted hints, so +> conformant clients are unaffected — but clients that branch on +> `annotations != nil` may start treating composites as destructive. + +**When `annotations` is set**, the explicit values are merged over the derived +floor: each explicitly set field wins; unset fields keep the derived value. + +**Safety floor and contradiction guardrail**: an explicit hint may be *more +conservative* than the derived floor (e.g. `readOnlyHint: false` when the floor +is `true`), but it may never make the tool look *safer* than its steps allow +(e.g. `readOnlyHint: true` when a step is not read-only, or +`destructiveHint: false` when a step is destructive). The floor is +**fail-closed**: when a workflow has one or more tool steps the floor is always +non-nil, and any step whose annotations are nil or unknown taints it +conservatively (`readOnly=false`, `destructive=true`, `openWorld=true`). In the +common case where backends declare no annotations (e.g. yardstick echo), the +floor is that conservative set, so an explicit `readOnlyHint: true` will +contradict it and be dropped. This guardrail runs at **runtime (advertise +time)**, not at admission: a contradicting composite tool is omitted from +`tools/list` and is also uncallable via `tools/call` (with a warning naming the +offending step tools in the vMCP logs), while the VirtualMCPServer stays Ready. + +> **`thv vmcp validate` does NOT check annotation contradictions.** A composite +> tool definition that passes validation can still be dropped at runtime if its +> explicit annotations contradict the derived safety floor — contradictions only +> surface in the vMCP logs at advertise time. It cannot run at admission because +> step-tool annotations are only known once backends are aggregated. + ### Template Syntax Use Go template syntax for dynamic values: diff --git a/docs/operator/virtualmcpserver-api.md b/docs/operator/virtualmcpserver-api.md index b9444cb9da..61f13cdd49 100644 --- a/docs/operator/virtualmcpserver-api.md +++ b/docs/operator/virtualmcpserver-api.md @@ -231,33 +231,79 @@ Defines tool aggregation and conflict resolution strategies. - `conflictResolutionConfig` (ConflictResolutionConfig, optional): Configuration for the chosen strategy - `tools` ([]WorkloadToolConfig, optional): Per-workload tool filtering and overrides - `excludeAllTools` (bool, optional): Excludes all tools from aggregation when true +- `defaultToolVisibility` (string, optional, default: "allow"): Whether a backend with no `tools` entry is advertised + - `allow`: every tool from an unlisted backend is advertised + - `deny`: an unlisted backend contributes no tools, so only backends named in `tools` are advertised **Example (prefix strategy)**: ```yaml spec: groupRef: name: my-services - aggregation: - conflictResolution: prefix - conflictResolutionConfig: - prefixFormat: "{workload}_" - tools: - - workload: github - filter: ["create_pr", "merge_pr"] - - workload: jira - toolConfigRef: - name: jira-tool-config + config: + aggregation: + conflictResolution: prefix + conflictResolutionConfig: + prefixFormat: "{workload}_" + tools: + - workload: github + filter: ["create_pr", "merge_pr"] + - workload: jira + toolConfigRef: + name: jira-tool-config +``` + +**Example (deny-by-default)**: + +By default, a workload in the group with no `tools` entry has all of its tools +advertised, so adding a workload to the group exposes it. Set +`defaultToolVisibility: deny` to invert that, so only the workloads you list +contribute tools: + +```yaml +spec: + groupRef: + name: my-services + config: + aggregation: + conflictResolution: prefix + defaultToolVisibility: deny + tools: + # Only these two workloads are advertised. Any other workload in + # my-services contributes no tools, including ones added later. + - workload: github + filter: ["get_issue", "list_prs"] + - workload: jira + filter: ["get_issue"] ``` +A workload that *is* listed is opted in by its entry: `excludeAll` and `filter` +on that entry then decide which of its tools are advertised. An entry with +neither advertises all of that workload's tools. + +`defaultToolVisibility: deny` cannot be combined with the `priority` conflict +strategy unless every `priorityOrder` entry also has a `tools` entry. An unlisted +backend can win a name conflict and then be withheld, which would hide the tool +from every workload offering it; the configuration is rejected rather than +silently dropping the tool. + +> **Note**: `defaultToolVisibility` controls **advertising** only, as do +> `excludeAllTools` and `excludeAll`. Tools hidden this way remain in the routing +> table so composite tools can still call them, and `defaultToolVisibility` does +> not affect resources or prompts, which are advertised regardless. For +> per-identity authorization, use Cedar policies (see +> [Authorization](../authz.md)). + **Example (priority strategy)**: ```yaml spec: groupRef: name: my-services - aggregation: - conflictResolution: priority - conflictResolutionConfig: - priorityOrder: ["github", "jira", "slack"] + config: + aggregation: + conflictResolution: priority + conflictResolutionConfig: + priorityOrder: ["github", "jira", "slack"] ``` **Example (manual strategy)**: @@ -265,23 +311,24 @@ spec: spec: groupRef: name: my-services - aggregation: - conflictResolution: manual - tools: - - workload: github - filter: ["create_pr", "merge_pr", "list_repos"] - overrides: - create_pr: - name: github_create_pr - description: "Create a pull request in GitHub" - - workload: jira - filter: ["create_issue", "update_issue"] - overrides: - create_issue: - name: jira_create_issue - description: "Create an issue in Jira" - # All tool name conflicts must be explicitly resolved via overrides - # Runtime validation ensures no unresolved conflicts exist + config: + aggregation: + conflictResolution: manual + tools: + - workload: github + filter: ["create_pr", "merge_pr", "list_repos"] + overrides: + create_pr: + name: github_create_pr + description: "Create a pull request in GitHub" + - workload: jira + filter: ["create_issue", "update_issue"] + overrides: + create_issue: + name: jira_create_issue + description: "Create an issue in Jira" + # All tool name conflicts must be explicitly resolved via overrides + # Runtime validation ensures no unresolved conflicts exist ``` #### WorkloadToolConfig diff --git a/docs/operator/virtualmcpserver-kubernetes-guide.md b/docs/operator/virtualmcpserver-kubernetes-guide.md index 0f43265c42..93e7a7b010 100644 --- a/docs/operator/virtualmcpserver-kubernetes-guide.md +++ b/docs/operator/virtualmcpserver-kubernetes-guide.md @@ -650,6 +650,55 @@ spec: - 'permit(principal, action, resource);' ``` +**Which token's claims a policy sees**: Cedar reads the primary upstream +provider's *access token*. Claims that token asserts always win. Because many +OIDC providers put profile claims in the `id_token` only and omit them from the +access token, `name` and `email` fall back to **that same provider's `id_token`** +(captured at login and stored alongside its access token) when the access token +does not carry them. Provenance is preserved either way: `principal has +claim_email` means "this upstream asserted an email". When a fallback happens the +proxy logs, once per 30s: + +``` +WARN upstream access token lacks profile claims; using the upstream ID token's + values for Cedar evaluation provider=okta claims="[email]" +``` + +That fallback uses the stored `id_token` even after it has expired, by design. It +is read as a record of what the upstream asserted when the user logged in, not +presented to anyone as a credential. Since a session can outlive an `id_token` by +days, expiring it out of policy evaluation would let the same user be permitted +early in a session and denied later, with no configuration change. + +The token the client presented — the one ToolHive's auth server issued — is +never a claim source here, even though it mirrors a `name` and `email`. In a +multi-upstream chain those mirrored values come from the **first** configured +upstream (the identity provider), which need not be the provider +`primaryUpstreamProvider` names, so using them could attribute one IdP's email to +another. + +An OAuth 2.0 upstream that was never asked for the `openid` scope has no stored +`id_token`, so nothing is available to fall back to and the claim stays absent. +The proxy says so once per 30s: + +``` +WARN no upstream ID token stored for provider; policies referencing profile + claims the access token omits will deny provider=okta +``` + +Every other claim is upstream-access-token-only. In particular, group, role and +scope claims are not substituted from anywhere, so a policy such as `principal in +THVGroup::"platform-eng"` only ever matches groups the upstream access token +asserts. The same holds for the principal: `sub` is never substituted, so a rule +keyed on `Client::""` keeps matching the upstream identity, and +an access token with no `sub` is rejected outright rather than having a principal +chosen for it. + +If a policy references a claim that neither the access token nor the `id_token` +carries, `principal has claim_x` is false and the policy denies — author domain +and group gates defensively with `has`, and confirm the claim is present in one of +those two tokens before gating on it. + > **Migration: `primaryUpstreamProvider` location** > > The field used to live under diff --git a/docs/server/docs.go b/docs/server/docs.go index 21808b4888..21d847c9e3 100644 --- a/docs/server/docs.go +++ b/docs/server/docs.go @@ -8,37 +8,152 @@ const docTemplate = `{ "schemes": {{ marshal .Schemes }}, "components": { "schemas": { - "github_com_stacklok_toolhive-core_registry_types.Registry": { - "description": "Full registry data", + "auth.TokenValidatorConfig": { + "description": "DEPRECATED: Middleware configuration.\nOIDCConfig contains OIDC configuration", "properties": { - "groups": { - "description": "Groups is a slice of group definitions containing related MCP servers", + "allowPrivateIP": { + "description": "AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses", + "type": "boolean" + }, + "allowedClientIDs": { + "description": "AllowedClientIDs is the set of OAuth client IDs accepted when validating\ntokens by their client_id claim instead of an audience claim.", "items": { - "$ref": "#/components/schemas/registry.Group" + "type": "string" }, - "type": "array", - "uniqueItems": false + "type": "array" }, - "last_updated": { - "description": "LastUpdated is the timestamp when the registry was last updated, in RFC3339 format", + "audience": { + "description": "Audience is the expected audience for the token", "type": "string" }, - "remote_servers": { - "additionalProperties": { - "$ref": "#/components/schemas/registry.RemoteServerMetadata" + "authTokenFile": { + "description": "AuthTokenFile is the path to file containing bearer token for authentication", + "type": "string" + }, + "cacertPath": { + "description": "CACertPath is the path to the CA certificate bundle for HTTPS requests", + "type": "string" + }, + "clientID": { + "description": "ClientID is the OIDC client ID", + "type": "string" + }, + "clientSecret": { + "description": "ClientSecret is the optional OIDC client secret for introspection", + "type": "string" + }, + "insecureAllowHTTP": { + "description": "InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing\nWARNING: This is insecure and should NEVER be used in production", + "type": "boolean" + }, + "introspectionURL": { + "description": "IntrospectionURL is the optional introspection endpoint for validating tokens", + "type": "string" + }, + "issuer": { + "description": "Issuer is the OIDC issuer URL (e.g., https://accounts.google.com)", + "type": "string" + }, + "jwksurl": { + "description": "JWKSURL is the URL to fetch the JWKS from", + "type": "string" + }, + "resourceURL": { + "description": "ResourceURL is the explicit resource URL for OAuth discovery (RFC 9728)", + "type": "string" + }, + "scopes": { + "description": "Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728)\nIf empty, defaults to [\"openid\"]", + "items": { + "type": "string" }, - "description": "RemoteServers is a map of server names to their corresponding remote server definitions\nThese are MCP servers accessed via HTTP/HTTPS using the thv proxy command", - "type": "object" + "type": "array" + } + }, + "type": "object" + }, + "core.Workload": { + "properties": { + "created_at": { + "description": "CreatedAt is the timestamp when the workload was created.", + "type": "string" }, - "servers": { + "group": { + "description": "Group is the name of the group this workload belongs to, if any.", + "type": "string" + }, + "labels": { "additionalProperties": { - "$ref": "#/components/schemas/registry.ImageMetadata" + "type": "string" }, - "description": "Servers is a map of server names to their corresponding server definitions", + "description": "Labels are the container labels (excluding standard ToolHive labels)", "type": "object" }, - "version": { - "description": "Version is the schema version of the registry", + "name": { + "description": "Name is the name of the workload.\nIt is used as a unique identifier.", + "type": "string" + }, + "package": { + "description": "Package specifies the Workload Package used to create this Workload.", + "type": "string" + }, + "port": { + "description": "Port is the port on which the workload is exposed.\nThis is embedded in the URL.", + "type": "integer" + }, + "proxy_mode": { + "description": "ProxyMode is the proxy mode that clients should use to connect.\nFor stdio transports, this will be the proxy mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this will be the same as TransportType.", + "type": "string" + }, + "remote": { + "description": "Remote indicates whether this is a remote workload (true) or a container workload (false).", + "type": "boolean" + }, + "started_at": { + "description": "StartedAt is when the container was last started (changes on restart)", + "type": "string" + }, + "status": { + "description": "Status is the current status of the workload.", + "enum": [ + "running", + "stopped", + "error", + "starting", + "stopping", + "unhealthy", + "removing", + "unknown", + "unauthenticated", + "auth_retrying", + "policy_stopped" + ], + "type": "string" + }, + "status_context": { + "description": "StatusContext provides additional context about the workload's status.\nThe exact meaning is determined by the status and the underlying runtime.", + "type": "string" + }, + "tools": { + "description": "ToolsFilter is the filter on tools applied to the workload.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "transport_type": { + "description": "TransportType is the type of transport used for this workload.", + "enum": [ + "stdio", + "sse", + "streamable-http", + "inspector" + ], + "type": "string" + }, + "url": { + "description": "URL is the URL of the workload exposed by the ToolHive proxy.", "type": "string" } }, @@ -110,70 +225,10 @@ const docTemplate = `{ "maxDataSize": { "description": "MaxDataSize limits the size of request/response data included in audit logs (in bytes).\n+kubebuilder:default=1024\n+optional", "type": "integer" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_auth.TokenValidatorConfig": { - "description": "DEPRECATED: Middleware configuration.\nOIDCConfig contains OIDC configuration", - "properties": { - "allowPrivateIP": { - "description": "AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses", - "type": "boolean" - }, - "allowedClientIDs": { - "description": "AllowedClientIDs is the set of OAuth client IDs accepted when validating\ntokens by their client_id claim instead of an audience claim.", - "items": { - "type": "string" - }, - "type": "array" - }, - "audience": { - "description": "Audience is the expected audience for the token", - "type": "string" - }, - "authTokenFile": { - "description": "AuthTokenFile is the path to file containing bearer token for authentication", - "type": "string" - }, - "cacertPath": { - "description": "CACertPath is the path to the CA certificate bundle for HTTPS requests", - "type": "string" - }, - "clientID": { - "description": "ClientID is the OIDC client ID", - "type": "string" - }, - "clientSecret": { - "description": "ClientSecret is the optional OIDC client secret for introspection", - "type": "string" - }, - "insecureAllowHTTP": { - "description": "InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing\nWARNING: This is insecure and should NEVER be used in production", - "type": "boolean" - }, - "introspectionURL": { - "description": "IntrospectionURL is the optional introspection endpoint for validating tokens", - "type": "string" - }, - "issuer": { - "description": "Issuer is the OIDC issuer URL (e.g., https://accounts.google.com)", - "type": "string" - }, - "jwksurl": { - "description": "JWKSURL is the URL to fetch the JWKS from", - "type": "string" - }, - "resourceURL": { - "description": "ResourceURL is the explicit resource URL for OAuth discovery (RFC 9728)", - "type": "string" }, - "scopes": { - "description": "Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728)\nIf empty, defaults to [\"openid\"]", - "items": { - "type": "string" - }, - "type": "array" + "maxDelegationDepth": { + "description": "MaxDelegationDepth caps how many nested RFC 8693 \"act\" entries are\nrecorded in an audit event's delegation chain. Deeper chains are\ntruncated (marked with truncated=true). Defaults to 10 when unset.\n+kubebuilder:validation:Minimum=1\n+kubebuilder:default=10\n+optional", + "type": "integer" } }, "type": "object" @@ -241,157 +296,64 @@ const docTemplate = `{ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_auth_remote.Config": { - "description": "RemoteAuthConfig contains OAuth configuration for remote MCP servers", + "github_com_stacklok_toolhive_pkg_auth_upstreamswap.Config": { + "description": "UpstreamSwapConfig contains configuration for upstream token swap middleware.\nWhen set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs\nfor upstream IdP tokens before forwarding requests to the MCP server.", "properties": { - "authorize_url": { + "custom_header_name": { + "description": "CustomHeaderName is the header name when HeaderStrategy is \"custom\".", "type": "string" }, - "bearer_token": { - "description": "Bearer token configuration (alternative to OAuth)", + "header_strategy": { + "description": "HeaderStrategy determines how to inject the token: \"replace\" (default) or \"custom\".", "type": "string" }, - "bearer_token_file": { + "provider_name": { + "description": "ProviderName identifies which upstream provider's tokens to retrieve for injection.\nThis is required and must match a configured upstream provider name.", "type": "string" - }, - "cached_cimd_client_id": { - "description": "CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD\nauthentication was used. Kept separate from CachedClientID (which holds\nDCR-issued IDs) so the two can have independent lifecycles — DCR credential\nrotation clears CachedClientID without touching the stable CIMD URL.\nRead by resolveClientCredentials to send the correct client_id on token refresh.", + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig": { + "description": "CIMD controls client_id metadata document support. When enabled, the\nembedded authorization server accepts HTTPS URLs as client_id values\nand resolves them via the CIMD protocol instead of requiring DCR.", + "properties": { + "cache_fallback_ttl": { + "description": "CacheFallbackTTL is the fixed TTL applied to every cached CIMD document.\nCache-Control header parsing is not yet implemented; all entries use this value.\nFormat: Go duration string (e.g. \"5m\", \"10m\", \"1h\").\nDefaults to 5 minutes when Enabled is true and this field is omitted.", + "example": "5m", "type": "string" }, - "cached_client_id": { - "description": "Cached DCR client credentials for persistence across restarts.\nThese are obtained during Dynamic Client Registration and needed to refresh tokens.\nClientID is stored as plain text since it's public information.", - "type": "string" + "cache_max_size": { + "description": "CacheMaxSize is the maximum number of CIMD documents held in the LRU cache.\nDefaults to 256 when Enabled is true and this field is zero.", + "type": "integer" }, - "cached_client_secret_ref": { + "enabled": { + "description": "Enabled activates CIMD client lookup when true.", + "type": "boolean" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_authserver.DCRUpstreamConfig": { + "description": "DCRConfig enables RFC 7591 Dynamic Client Registration against the\nupstream authorization server. When set, the client credentials are\nobtained at runtime rather than being pre-provisioned via ClientID /\nClientSecretFile / ClientSecretEnvVar, and ClientID must be left empty.\nMutually exclusive with ClientID.", + "properties": { + "discovery_url": { + "description": "DiscoveryURL is the exact RFC 8414 / OIDC Discovery document URL to\nfetch at runtime. The resolver issues a single GET against this URL\n(no well-known-path fallback) and reads registration_endpoint,\nauthorization_endpoint, token_endpoint,\ntoken_endpoint_auth_methods_supported, and scopes_supported from the\nresponse. Per RFC 8414 §3.3, the document's \"issuer\" field must\nexactly match the upstream issuer configured on the parent\nrun-config.\n\nUse this field when the upstream publishes discovery metadata at a\npath that differs from the issuer-derived well-known paths — for\nexample a multi-tenant IdP whose metadata lives at\nhttps://idp.example.com/tenants/acme/.well-known/openid-configuration.\n\nMutually exclusive with RegistrationEndpoint.", "type": "string" }, - "cached_refresh_token_ref": { - "description": "Cached OAuth token reference for persistence across restarts.\nThe refresh token is stored securely in the secret manager, and this field\ncontains the reference to retrieve it (e.g., \"OAUTH_REFRESH_TOKEN_workload\").\nThis enables session restoration without requiring a new browser-based login.", + "initial_access_token_env_var": { + "description": "InitialAccessTokenEnvVar is the name of an environment variable\ncontaining the RFC 7591 initial access token. Mutually exclusive with\nInitialAccessTokenFile.", "type": "string" }, - "cached_reg_token_ref": { - "description": "RegistrationAccessToken is used to update/delete the client registration.\nStored as a secret reference since it's sensitive.", + "initial_access_token_file": { + "description": "InitialAccessTokenFile is the path to a file containing the RFC 7591\ninitial access token presented to the registration endpoint. Mutually\nexclusive with InitialAccessTokenEnvVar. Both may be omitted for open\nregistration endpoints.", "type": "string" }, - "cached_secret_expiry": { - "description": "ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server).\nA zero value means the secret does not expire.", + "registration_endpoint": { + "description": "RegistrationEndpoint is the RFC 7591 registration endpoint URL used\ndirectly, bypassing discovery. Because no discovery is performed,\nserver-capability fields (token_endpoint_auth_methods_supported,\nscopes_supported) are unavailable on this code path; the caller is\nexpected to also supply AuthorizationEndpoint, TokenEndpoint, and an\nexplicit Scopes list on the parent OAuth2UpstreamRunConfig. Auth\nmethod falls back to the resolver's default (client_secret_basic).\n\nMutually exclusive with DiscoveryURL.", "type": "string" }, - "cached_token_expiry": { - "type": "string" - }, - "callback_port": { - "type": "integer" - }, - "client_id": { - "type": "string" - }, - "client_secret": { - "type": "string" - }, - "client_secret_file": { - "type": "string" - }, - "issuer": { - "description": "OAuth endpoint configuration (from registry)", - "type": "string" - }, - "oauth_params": { - "additionalProperties": { - "type": "string" - }, - "description": "OAuth parameters for server-specific customization", - "type": "object" - }, - "resource": { - "description": "Resource is the OAuth 2.0 resource indicator (RFC 8707).", - "type": "string" - }, - "scope_param_name": { - "description": "ScopeParamName overrides the query parameter name used to send scopes in the\nauthorization URL. When empty, the standard \"scope\" parameter is used.\nSome providers require a non-standard name (e.g., Slack uses \"user_scope\").", - "type": "string" - }, - "scopes": { - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false - }, - "skip_browser": { - "type": "boolean" - }, - "timeout": { - "example": "5m", - "type": "string" - }, - "token_url": { - "type": "string" - }, - "use_pkce": { - "type": "boolean" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_auth_upstreamswap.Config": { - "description": "UpstreamSwapConfig contains configuration for upstream token swap middleware.\nWhen set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs\nfor upstream IdP tokens before forwarding requests to the MCP server.", - "properties": { - "custom_header_name": { - "description": "CustomHeaderName is the header name when HeaderStrategy is \"custom\".", - "type": "string" - }, - "header_strategy": { - "description": "HeaderStrategy determines how to inject the token: \"replace\" (default) or \"custom\".", - "type": "string" - }, - "provider_name": { - "description": "ProviderName identifies which upstream provider's tokens to retrieve for injection.\nThis is required and must match a configured upstream provider name.", - "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig": { - "description": "CIMD controls client_id metadata document support. When enabled, the\nembedded authorization server accepts HTTPS URLs as client_id values\nand resolves them via the CIMD protocol instead of requiring DCR.", - "properties": { - "cache_fallback_ttl": { - "description": "CacheFallbackTTL is the fixed TTL applied to every cached CIMD document.\nCache-Control header parsing is not yet implemented; all entries use this value.\nFormat: Go duration string (e.g. \"5m\", \"10m\", \"1h\").\nDefaults to 5 minutes when Enabled is true and this field is omitted.", - "example": "5m", - "type": "string" - }, - "cache_max_size": { - "description": "CacheMaxSize is the maximum number of CIMD documents held in the LRU cache.\nDefaults to 256 when Enabled is true and this field is zero.", - "type": "integer" - }, - "enabled": { - "description": "Enabled activates CIMD client lookup when true.", - "type": "boolean" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver.DCRUpstreamConfig": { - "description": "DCRConfig enables RFC 7591 Dynamic Client Registration against the\nupstream authorization server. When set, the client credentials are\nobtained at runtime rather than being pre-provisioned via ClientID /\nClientSecretFile / ClientSecretEnvVar, and ClientID must be left empty.\nMutually exclusive with ClientID.", - "properties": { - "discovery_url": { - "description": "DiscoveryURL is the exact RFC 8414 / OIDC Discovery document URL to\nfetch at runtime. The resolver issues a single GET against this URL\n(no well-known-path fallback) and reads registration_endpoint,\nauthorization_endpoint, token_endpoint,\ntoken_endpoint_auth_methods_supported, and scopes_supported from the\nresponse. Per RFC 8414 §3.3, the document's \"issuer\" field must\nexactly match the upstream issuer configured on the parent\nrun-config.\n\nUse this field when the upstream publishes discovery metadata at a\npath that differs from the issuer-derived well-known paths — for\nexample a multi-tenant IdP whose metadata lives at\nhttps://idp.example.com/tenants/acme/.well-known/openid-configuration.\n\nMutually exclusive with RegistrationEndpoint.", - "type": "string" - }, - "initial_access_token_env_var": { - "description": "InitialAccessTokenEnvVar is the name of an environment variable\ncontaining the RFC 7591 initial access token. Mutually exclusive with\nInitialAccessTokenFile.", - "type": "string" - }, - "initial_access_token_file": { - "description": "InitialAccessTokenFile is the path to a file containing the RFC 7591\ninitial access token presented to the registration endpoint. Mutually\nexclusive with InitialAccessTokenEnvVar. Both may be omitted for open\nregistration endpoints.", - "type": "string" - }, - "registration_endpoint": { - "description": "RegistrationEndpoint is the RFC 7591 registration endpoint URL used\ndirectly, bypassing discovery. Because no discovery is performed,\nserver-capability fields (token_endpoint_auth_methods_supported,\nscopes_supported) are unavailable on this code path; the caller is\nexpected to also supply AuthorizationEndpoint, TokenEndpoint, and an\nexplicit Scopes list on the parent OAuth2UpstreamRunConfig. Auth\nmethod falls back to the resolver's default (client_secret_basic).\n\nMutually exclusive with DiscoveryURL.", - "type": "string" - }, - "software_id": { - "description": "SoftwareID is the RFC 7591 \"software_id\" registration metadata value,\nidentifying the client software independent of any particular\nregistration instance.", + "software_id": { + "description": "SoftwareID is the RFC 7591 \"software_id\" registration metadata value,\nidentifying the client software independent of any particular\nregistration instance.", "type": "string" }, "software_statement": { @@ -430,7 +392,7 @@ const docTemplate = `{ "type": "object" }, "allow_private_ips": { - "description": "AllowPrivateIPs permits the upstream provider's HTTP client to connect to\nprivate IP ranges (RFC-1918, link-local). Use only when the upstream is\nhosted inside the same cluster and has no public endpoint. HTTP-scheme\nrestrictions are unchanged — HTTPS is still required for non-localhost hosts.\nDefaults to false.", + "description": "AllowPrivateIPs permits the upstream provider's HTTP client to connect to\nprivate IP ranges (RFC-1918, link-local). When DCRConfig is set, this\nalso gates the DCR discovery and registration calls made on this\nupstream's behalf (see pkg/authserver/runner/dcr_adapter.go), so a\nsingle flag covers the whole upstream rather than needing a separate\nDCR-specific setting. Use only when the upstream is hosted inside the\nsame cluster and has no public endpoint. HTTP-scheme restrictions are\nunchanged — HTTPS is still required for non-localhost hosts. Defaults\nto false.", "type": "boolean" }, "authorization_endpoint": { @@ -566,6 +528,10 @@ const docTemplate = `{ "cimd": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig" }, + "delegation_token_lifespan": { + "description": "DelegationTokenLifespan is the maximum lifetime for delegated tokens issued\nvia RFC 8693 token exchange. Specified as a Go duration string (e.g., \"15m\").\nIf empty, defaults to 15 minutes.", + "type": "string" + }, "disable_upstream_token_injection": { "description": "DisableUpstreamTokenInjection prevents the upstream swap middleware from being added.\nWhen true, the embedded auth server handles OAuth flows for clients, but instead of\ninjecting upstream IdP tokens the proxy strips the client's credential headers\n(Authorization, Cookie, Proxy-Authorization) after the JWT is validated — the\nbackend receives an unauthenticated request. Incompatible with token exchange\nand AWS STS, which would re-add credentials after the strip.", "type": "boolean" @@ -602,11 +568,19 @@ const docTemplate = `{ "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.SigningKeyRunConfig" }, "storage": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RunConfig" + "$ref": "#/components/schemas/storage.RunConfig" }, "token_lifespans": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.TokenLifespanRunConfig" }, + "trusted_issuers": { + "description": "TrustedIssuers lists external OIDC issuers whose tokens are accepted as\nsubject tokens during RFC 8693 token exchange. Empty (the default) means\nonly self-issued subject tokens are accepted.\n\nPrerequisite: the token-exchange grant requires a confidential client,\nand no supported deployment path provisions one today (DCR and CIMD\nclients are both public-only, and there is no client-seeding field on\nthis RunConfig), so this grant is not yet usable end to end for\nself-issued or external subject tokens alike. Tracked in\nhttps://github.com/stacklok/toolhive/issues/6082.\n\nSee tokenexchange.TrustedIssuer for the per-issuer field reference, and\ndocs/arch/17-token-exchange-delegation.md for the trust model, consent\nsignals, and operator-facing constraints (audience/scope bounding,\nsubject namespace qualification, required client binding) that aren't\nvisible from the config shape alone.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_server_tokenexchange.TrustedIssuer" + }, + "type": "array", + "uniqueItems": false + }, "upstreams": { "description": "Upstreams configures connections to upstream Identity Providers.\nAt least one upstream is required - the server delegates authentication to these providers.\nMultiple upstreams are supported for sequential authorization chains.", "items": { @@ -764,111 +738,47 @@ const docTemplate = `{ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_authserver_storage.ACLUserRunConfig": { - "description": "ACLUserConfig contains ACL user authentication configuration.", - "properties": { - "password_env_var": { - "description": "PasswordEnvVar is the environment variable containing the Redis password.", - "type": "string" - }, - "username_env_var": { - "description": "UsernameEnvVar is the environment variable containing the Redis username.", - "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.RedisRunConfig": { - "description": "RedisConfig is the Redis-specific configuration when Type is \"redis\".", + "github_com_stacklok_toolhive_pkg_authserver_server_tokenexchange.TrustedIssuer": { "properties": { - "acl_user_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.ACLUserRunConfig" - }, - "addr": { - "description": "Addr is the Redis server address (host:port). Required for standalone and cluster modes.\nMutually exclusive with SentinelConfig.", - "type": "string" - }, - "auth_type": { - "description": "AuthType must be \"aclUser\" - only ACL user authentication is supported.", + "actor_claim": { + "description": "ActorClaim names the claim identifying the client that requested the\nsubject token from THIS EXTERNAL ISSUER (used by AllowedActors below).\nValues are in the external issuer's namespace, NOT ToolHive client\nIDs. Defaults to \"azp\"; use \"appid\" for Microsoft Entra v1, \"cid\" for\nOkta. The special value \"client_id\" reads ValidatedClaims.ClientID\ninstead of Extra (assignClaim routes it to that field) — it is still\nthe external token's client_id claim, not a ToolHive one.", "type": "string" }, - "cluster_mode": { - "description": "ClusterMode enables the Redis Cluster protocol. Requires Addr to be set.", + "allow_private_ips": { + "description": "AllowPrivateIPs permits OIDC discovery and JWKS fetches for THIS\nissuer to resolve to a private or loopback address. Use only when the\nissuer is hosted inside the same cluster and has no public endpoint.", "type": "boolean" }, - "dial_timeout": { - "description": "DialTimeout is the timeout for establishing connections (e.g., \"5s\").", - "type": "string" - }, - "key_prefix": { - "description": "KeyPrefix for multi-tenancy, typically \"thv:auth:{ns}:{name}:\".", - "type": "string" - }, - "read_timeout": { - "description": "ReadTimeout is the timeout for read operations (e.g., \"3s\").", - "type": "string" - }, - "sentinel_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.SentinelRunConfig" - }, - "sentinel_tls": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig" + "allowed_actors": { + "description": "AllowedActors is the allowlist of ActorClaim values authorized to\nexchange a subject token from this issuer when it carries no\n\"may_act\" claim; empty means only may_act-bearing tokens are\naccepted. By itself names no ToolHive client — see\nAllowedDelegateClients and docs/arch/17-token-exchange-delegation.md\n(\"Accepted limitations\" #1).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "tls": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig" + "allowed_delegate_clients": { + "description": "AllowedDelegateClients restricts which ToolHive client IDs may\nexchange a subject token from this issuer, for BOTH consent paths.\nRequired (validateTrustedIssuer rejects empty/absent); \"*\" permits\nany confidential client holding the grant. See\ndocs/arch/17-token-exchange-delegation.md (\"Accepted limitations\" #1).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "write_timeout": { - "description": "WriteTimeout is the timeout for write operations (e.g., \"3s\").", - "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig": { - "description": "SentinelTLS configures TLS for Sentinel connections. Only applies when SentinelConfig is set.", - "properties": { - "ca_cert_file": { - "description": "CACertFile is the path to a PEM-encoded CA certificate file.", + "expected_audience": { + "description": "ExpectedAudience is the expected \"aud\" claim value that must appear\nin the token's audience list (a resource/API identifier, not a\nclient ID — required, but not enforced; see looksLikeResourceIdentifier).\nSee docs/arch/17-token-exchange-delegation.md (\"ID/access-token\ndiscrimination\") for why and its limits.", "type": "string" }, - "insecure_skip_verify": { - "description": "InsecureSkipVerify skips certificate verification.", + "insecure_allow_http": { + "description": "InsecureAllowHTTP permits plain-HTTP OIDC discovery and JWKS fetches\nfor THIS issuer only. Development and testing only — never set in\nproduction. Does not relax the private-IP guard; see AllowPrivateIPs.\nDeliberately per-issuer: this server's own InsecureAllowHTTP must not\nsilently permit plaintext discovery for every trusted external issuer\ntoo — a network attacker who can intercept that traffic could\nsubstitute a JWKS and forge subject tokens for that issuer's\nnamespace.", "type": "boolean" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.RunConfig": { - "description": "Storage configures the storage backend for the auth server.\nIf nil, defaults to in-memory storage.", - "properties": { - "redis_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisRunConfig" }, - "type": { - "description": "Type specifies the storage backend type. Defaults to \"memory\".", + "issuer_url": { + "description": "IssuerURL is the expected \"iss\" claim value (exact match).", "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.SentinelRunConfig": { - "description": "SentinelConfig contains Sentinel-specific configuration.\nMutually exclusive with Addr.", - "properties": { - "db": { - "description": "DB is the Redis database number (default: 0).", - "type": "integer" }, - "master_name": { - "description": "MasterName is the name of the Redis Sentinel master.", + "jwks_url": { + "description": "JWKSURL is the URL to fetch the issuer's JSON Web Key Set from.\nIf empty, it is resolved via OIDC discovery at {IssuerURL}/.well-known/openid-configuration.", "type": "string" - }, - "sentinel_addrs": { - "description": "SentinelAddrs is the list of Sentinel addresses (host:port).", - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false } }, "type": "object" @@ -982,230 +892,326 @@ const docTemplate = `{ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus": { - "description": "Current status of the workload", - "enum": [ - "running", - "stopped", - "error", - "starting", - "stopping", - "unhealthy", - "removing", - "unknown", - "unauthenticated", - "auth_retrying", - "policy_stopped", - "running", - "stopped", - "error", - "starting", - "stopping", - "unhealthy", - "removing", - "unknown", - "unauthenticated", - "auth_retrying", - "policy_stopped", - "running", - "stopped", - "error", - "starting", - "stopping", - "unhealthy", - "removing", - "unknown", - "unauthenticated", - "auth_retrying", - "policy_stopped" - ], - "type": "string", - "x-enum-varnames": [ - "WorkloadStatusRunning", - "WorkloadStatusStopped", - "WorkloadStatusError", - "WorkloadStatusStarting", - "WorkloadStatusStopping", - "WorkloadStatusUnhealthy", - "WorkloadStatusRemoving", - "WorkloadStatusUnknown", - "WorkloadStatusUnauthenticated", - "WorkloadStatusAuthRetrying", - "WorkloadStatusPolicyStopped" - ] - }, - "github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig": { - "description": "RuntimeConfig allows overriding the default runtime configuration\nfor this specific workload (base images and packages)", + "github_com_stacklok_toolhive_pkg_groups.Group": { "properties": { - "additional_packages": { - "description": "AdditionalPackages lists extra packages to install in the builder and\nruntime stages.\nExamples for Alpine: [\"git\", \"make\", \"gcc\"]\nExamples for Debian: [\"git\", \"build-essential\"]", + "name": { + "type": "string" + }, + "plugins": { "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "builder_image": { - "description": "BuilderImage is the full image reference for the builder stage.\nAn empty string signals \"use the default for this transport type\" during config merging.\nExamples: \"golang:1.26-alpine\", \"node:24-alpine\", \"python:3.14-slim\"", - "type": "string" + "registered_clients": { + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "runtime_env": { - "additionalProperties": { + "skills": { + "items": { "type": "string" }, - "description": "RuntimeEnv contains environment variables to inject into the Dockerfile's\nfinal runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv),\nwhich only affects the builder stage, these variables are baked into the\nshipped image and are present in the running container's process\nenvironment at startup. Use this for values a packaged MCP server reads at\nprocess start (e.g. feature flags, cache backend selection), not for\nbuild-time package manager configuration.\nKeys must be uppercase with underscores, values are validated for safety.", - "type": "object" + "type": "array", + "uniqueItems": false } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_core.Workload": { + "github_com_stacklok_toolhive_pkg_plugins.BuildResult": { "properties": { - "created_at": { - "description": "CreatedAt is the timestamp when the workload was created.", - "type": "string" - }, - "group": { - "description": "Group is the name of the group this workload belongs to, if any.", + "reference": { + "description": "Reference is the OCI reference of the built skill artifact.", + "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.ComponentInventory": { + "additionalProperties": { + "type": "integer" + }, + "description": "Components is the inventory of component types declared by the plugin\n(e.g. {\"commands\": 3, \"skills\": 2}). Extracted from the OCI artifact.", + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.ComponentType": { + "enum": [ + "commands", + "agents", + "skills", + "hooks", + "mcpServers", + "lspServers" + ], + "type": "string", + "x-enum-varnames": [ + "ComponentCommands", + "ComponentAgents", + "ComponentSkills", + "ComponentHooks", + "ComponentMCP", + "ComponentLSP" + ] + }, + "github_com_stacklok_toolhive_pkg_plugins.Dependency": { + "properties": { + "digest": { + "description": "Digest is the OCI digest for upgrade detection.", "type": "string" - }, - "labels": { - "additionalProperties": { - "type": "string" - }, - "description": "Labels are the container labels (excluding standard ToolHive labels)", - "type": "object" }, "name": { - "description": "Name is the name of the workload.\nIt is used as a unique identifier.", + "description": "Name is the dependency name.", "type": "string" }, - "package": { - "description": "Package specifies the Workload Package used to create this Workload.", + "reference": { + "description": "Reference is the OCI reference for the dependency.", "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.InstallStatus": { + "description": "Status is the current installation status.", + "enum": [ + "installed", + "pending", + "failed" + ], + "type": "string", + "x-enum-varnames": [ + "InstallStatusInstalled", + "InstallStatusPending", + "InstallStatusFailed" + ] + }, + "github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin": { + "description": "InstalledPlugin contains the full installation record.", + "properties": { + "clients": { + "description": "Clients is the list of client identifiers the plugin is installed for.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "port": { - "description": "Port is the port on which the workload is exposed.\nThis is embedded in the URL.", - "type": "integer" + "components": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ComponentInventory" }, - "proxy_mode": { - "description": "ProxyMode is the proxy mode that clients should use to connect.\nFor stdio transports, this will be the proxy mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this will be the same as TransportType.", + "dependencies": { + "description": "Dependencies is the list of external plugin dependencies.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Dependency" + }, + "type": "array", + "uniqueItems": false + }, + "digest": { + "description": "Digest is the OCI digest (sha256:...) for upgrade detection.", "type": "string" }, - "remote": { - "description": "Remote indicates whether this is a remote workload (true) or a container workload (false).", - "type": "boolean" + "installed_at": { + "description": "InstalledAt is the timestamp when the plugin was installed.", + "type": "string" }, - "started_at": { - "description": "StartedAt is when the container was last started (changes on restart)", + "metadata": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginMetadata" + }, + "project_root": { + "description": "ProjectRoot is the project root path for project-scoped plugins. Empty for user-scoped.", + "type": "string" + }, + "reference": { + "description": "Reference is the full OCI reference (e.g. ghcr.io/org/plugin:v1).", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Scope" + }, + "signature": { + "description": "Signature is the optional signing signature for the plugin artifact.", "type": "string" }, "status": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstallStatus" }, - "status_context": { - "description": "StatusContext provides additional context about the workload's status.\nThe exact meaning is determined by the status and the underlying runtime.", + "tag": { + "description": "Tag is the OCI tag (e.g. v1.0.0).", + "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.LocalBuild": { + "properties": { + "description": { + "description": "Description is the skill description extracted from the artifact metadata, if available.", "type": "string" }, - "tools": { - "description": "ToolsFilter is the filter on tools applied to the workload.", - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false + "digest": { + "description": "Digest is the OCI digest of the artifact (sha256:...).", + "type": "string" }, - "transport_type": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.TransportType" + "name": { + "description": "Name is the skill name extracted from the artifact metadata, if available.", + "type": "string" }, - "url": { - "description": "URL is the URL of the workload exposed by the ToolHive proxy.", + "tag": { + "description": "Tag is the OCI tag or name used to reference the artifact.", + "type": "string" + }, + "version": { + "description": "Version is the skill version extracted from the artifact metadata, if available.", "type": "string" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_groups.Group": { + "github_com_stacklok_toolhive_pkg_plugins.PluginContent": { "properties": { - "name": { + "description": { + "description": "Description is the plugin description from the OCI config labels.", "type": "string" }, - "plugins": { + "files": { + "description": "Files is the list of all files in the artifact with their sizes.", "items": { - "type": "string" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginFileEntry" }, "type": "array", "uniqueItems": false }, - "registered_clients": { - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false + "license": { + "description": "License is the SPDX license identifier from the OCI config labels.", + "type": "string" }, - "skills": { - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false + "manifest": { + "description": "Manifest is the raw .claude-plugin/plugin.json body.", + "type": "string" + }, + "name": { + "description": "Name is the plugin name from the OCI config labels.", + "type": "string" + }, + "version": { + "description": "Version is the plugin version from the OCI config labels.", + "type": "string" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_ignore.Config": { - "description": "IgnoreConfig contains configuration for ignore processing", + "github_com_stacklok_toolhive_pkg_plugins.PluginFileEntry": { "properties": { - "loadGlobal": { - "description": "Whether to load global ignore patterns", - "type": "boolean" + "path": { + "description": "Path is the file path within the artifact.", + "type": "string" }, - "printOverlays": { - "description": "Whether to print resolved overlay paths for debugging", - "type": "boolean" + "size": { + "description": "Size is the uncompressed file size in bytes.", + "type": "integer" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_oauthproto_tokenexchange.Config": { - "description": "TokenExchangeConfig contains token exchange configuration for external authentication", + "github_com_stacklok_toolhive_pkg_plugins.PluginInfo": { "properties": { - "audience": { - "description": "Audience is the target audience for the exchanged token", - "type": "string" + "installed_plugin": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin" }, - "client_id": { - "description": "ClientID is the OAuth 2.0 client identifier", - "type": "string" + "metadata": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginMetadata" }, - "client_secret": { - "description": "ClientSecret is the OAuth 2.0 client secret", - "type": "string" + "project_scope_degraded_clients": { + "description": "ProjectScopeDegradedClients lists the client types for which a\nproject-scoped install degraded (the adapter could only materialize at\nuser scope — e.g. Codex always writes to the user-scoped config.toml).\nPopulated by Info; empty for user-scoped installs. Recomputed at read\ntime from the stored scope + each adapter's capability, mirroring the\nUnmaterializedComponents pattern (no persistence needed — the degradation\nis deterministic from scope + client type).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "external_token_header_name": { - "description": "ExternalTokenHeaderName is the name of the custom header to use when HeaderStrategy is \"custom\"", + "unmaterialized_components": { + "additionalProperties": { + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ComponentType" + }, + "type": "array" + }, + "description": "UnmaterializedComponents lists, per client type, the component types the\nplugin declares that the installed client adapter does NOT load. Populated\nby Info by diffing InstalledPlugin.Components against each installed\nclient adapter's SupportedComponents.", + "type": "object" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.PluginMetadata": { + "description": "Metadata contains the plugin's metadata.", + "properties": { + "author": { + "description": "Author is the plugin author or maintainer.", "type": "string" }, - "header_strategy": { - "description": "HeaderStrategy determines how to inject the token\nValid values: HeaderStrategyReplace (default), HeaderStrategyCustom", + "description": { + "description": "Description is a human-readable description of the plugin.", "type": "string" }, - "scopes": { - "description": "Scopes is the list of scopes to request for the exchanged token", + "keywords": { + "description": "Keywords is a list of keywords for categorization/search.", "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "subject_token_type": { - "description": "SubjectTokenType specifies the type of the subject token being exchanged.\nCommon values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT.\nIf empty, defaults to oauthproto.TokenTypeAccessToken.", + "license": { + "description": "License is the SPDX license identifier for the plugin.", "type": "string" }, - "token_url": { - "description": "TokenURL is the OAuth 2.0 token endpoint URL", + "name": { + "description": "Name is the unique name of the plugin (kebab-case).", "type": "string" + }, + "version": { + "description": "Version is the semantic version of the plugin.", + "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.Scope": { + "description": "Scope for the installation", + "enum": [ + "user", + "project" + ], + "type": "string", + "x-enum-varnames": [ + "ScopeUser", + "ScopeProject" + ] + }, + "github_com_stacklok_toolhive_pkg_plugins.ValidationResult": { + "properties": { + "errors": { + "description": "Errors is a list of validation errors, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "valid": { + "description": "Valid indicates whether the skill definition is valid.", + "type": "boolean" + }, + "warnings": { + "description": "Warnings is a list of non-blocking validation warnings, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false } }, "type": "object" @@ -1285,7 +1291,7 @@ const docTemplate = `{ "additional_middleware_configs": { "description": "AdditionalMiddlewareConfigs carries pre-built middleware configs injected by\nexternal-auth handlers (reached via *[]RunConfigBuilderOption) rather than\nderived from typed RunConfig fields. PopulateMiddlewareConfigs splices these\ninto the chain in the backend-egress group — after auth and before recovery —\ninstead of discarding them. Upstream carries these configs verbatim and never\ninspects their parameters; the middleware type identity (e.g. an enterprise\nauth type) is supplied by the caller via types.MiddlewareConfig.Type.\n\nEach entry's Type is expected to be a NEW egress middleware type (e.g. OBO),\nnot one already produced from a typed RunConfig field (auth, authz, audit,\ntokenExchange, awssts, …). Dispatch in the proxyrunner is purely by Type\nstring, so an injected Type that shadows a typed-field type would add a\nsecond instance of that middleware to the chain; the seam does not validate\nagainst this.", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig" + "$ref": "#/components/schemas/types.MiddlewareConfig" }, "type": "array", "uniqueItems": false @@ -1376,7 +1382,7 @@ const docTemplate = `{ "type": "string" }, "ignore_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_ignore.Config" + "$ref": "#/components/schemas/ignore.Config" }, "image": { "description": "Image is the Docker image to run", @@ -1401,7 +1407,7 @@ const docTemplate = `{ "middleware_configs": { "description": "MiddlewareConfigs contains the list of middleware to apply to the transport\nand the configuration for each middleware.", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig" + "$ref": "#/components/schemas/types.MiddlewareConfig" }, "type": "array", "uniqueItems": false @@ -1419,7 +1425,7 @@ const docTemplate = `{ "type": "string" }, "oidc_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_auth.TokenValidatorConfig" + "$ref": "#/components/schemas/auth.TokenValidatorConfig" }, "permission_profile_name_or_path": { "description": "PermissionProfileNameOrPath is the name or path of the permission profile", @@ -1430,7 +1436,12 @@ const docTemplate = `{ "type": "integer" }, "proxy_mode": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.ProxyMode" + "description": "ProxyMode is the effective HTTP protocol the proxy uses.\nFor stdio transports, this is the configured mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this matches the transport type.\nNote: \"sse\" is deprecated; use \"streamable-http\" instead.", + "enum": [ + "sse", + "streamable-http" + ], + "type": "string" }, "publish": { "description": "Publish lists ports to publish to the host in format \"hostPort:containerPort\"", @@ -1460,14 +1471,14 @@ const docTemplate = `{ "type": "string" }, "remote_auth_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_auth_remote.Config" + "$ref": "#/components/schemas/remote.Config" }, "remote_url": { "description": "RemoteURL is the URL of the remote MCP server (if running remotely)", "type": "string" }, "runtime_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig" + "$ref": "#/components/schemas/templates.RuntimeConfig" }, "scaling_config": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_runner.ScalingConfig" @@ -1493,6 +1504,10 @@ const docTemplate = `{ "description": "Stateless indicates the server only supports POST (no SSE/GET).\nWhen true, the proxy returns 405 for incoming GET requests and uses a\nPOST-based health check instead of the default GET probe.\nApplies to both remote URLs and local container workloads.", "type": "boolean" }, + "strict_protocol_validation": { + "description": "StrictProtocolValidation enables strict MCP-Protocol-Version validation\non the streamable HTTP proxy: a request whose header names an unknown\nMCP revision is rejected with HTTP 400. Default false accepts any\nversion string (an absent header is always accepted in either mode).", + "type": "boolean" + }, "target_host": { "description": "TargetHost is the host to forward traffic to (only applicable to SSE transport)", "type": "string" @@ -1502,14 +1517,14 @@ const docTemplate = `{ "type": "integer" }, "telemetry_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_telemetry.Config" + "$ref": "#/components/schemas/telemetry.Config" }, "thv_ca_bundle": { "description": "DEPRECATED: No longer appears to be used.\nThvCABundle is the path to the CA certificate bundle for ToolHive HTTP operations", "type": "string" }, "token_exchange_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_oauthproto_tokenexchange.Config" + "$ref": "#/components/schemas/tokenexchange.Config" }, "tools_filter": { "description": "DEPRECATED: Middleware configuration.\nToolsFilter is the list of tools to filter", @@ -1527,7 +1542,14 @@ const docTemplate = `{ "type": "object" }, "transport": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.TransportType" + "description": "Transport is the transport mode (stdio, sse, or streamable-http)", + "enum": [ + "stdio", + "sse", + "streamable-http", + "inspector" + ], + "type": "string" }, "trust_proxy_headers": { "description": "TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies", @@ -1599,18 +1621,6 @@ const docTemplate = `{ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_secrets.SecretParameter": { - "description": "Bearer token for authentication (alternative to OAuth)", - "properties": { - "name": { - "type": "string" - }, - "target": { - "type": "string" - } - }, - "type": "object" - }, "github_com_stacklok_toolhive_pkg_skills.BuildResult": { "properties": { "reference": { @@ -1637,6 +1647,30 @@ const docTemplate = `{ }, "type": "object" }, + "github_com_stacklok_toolhive_pkg_skills.FailureReason": { + "description": "Reason is a typed failure reason when Status is UpgradeStatusFailed.", + "enum": [ + "registry-unreachable", + "digest-missing", + "validation-rejected", + "lock-write-failed", + "signature-invalid", + "signer-mismatch", + "unsigned-rejected", + "unknown" + ], + "type": "string", + "x-enum-varnames": [ + "FailureReasonRegistryUnreachable", + "FailureReasonDigestMissing", + "FailureReasonValidationRejected", + "FailureReasonLockWriteFailed", + "FailureReasonSignatureInvalid", + "FailureReasonSignerMismatch", + "FailureReasonUnsignedRejected", + "FailureReasonUnknown" + ] + }, "github_com_stacklok_toolhive_pkg_skills.InstallStatus": { "description": "Status is the current installation status.", "enum": [ @@ -1678,6 +1712,10 @@ const docTemplate = `{ "description": "InstalledAt is the timestamp when the skill was installed.", "type": "string" }, + "managed": { + "description": "Managed indicates this install is tracked in the project's\ntoolhive.lock.yaml. Only ever true for project-scoped installs. No\nomitempty: false is an observable state (unmanaged), not an absence.", + "type": "boolean" + }, "metadata": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillMetadata" }, @@ -1826,149 +1864,183 @@ const docTemplate = `{ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_skills.ValidationResult": { + "github_com_stacklok_toolhive_pkg_skills.SyncFailure": { "properties": { - "errors": { - "description": "Errors is a list of validation errors, if any.", + "error": { + "description": "Error is a human-readable description of the failure.", + "type": "string" + }, + "name": { + "description": "Name is the skill name that failed.", + "type": "string" + }, + "reason": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.FailureReason" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_skills.SyncResult": { + "properties": { + "already_current": { + "description": "AlreadyCurrent lists skills that already matched the lock file.", "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "valid": { - "description": "Valid indicates whether the skill definition is valid.", - "type": "boolean" - }, - "warnings": { - "description": "Warnings is a list of non-blocking validation warnings, if any.", + "drifted": { + "description": "Drifted lists skills whose on-disk contentDigest differed from the lock\nfile. Normally these are reinstalled to match it; when Check is set,\nnothing is written and this field reports the drift only.", "items": { "type": "string" }, "type": "array", "uniqueItems": false - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_telemetry.Config": { - "description": "DEPRECATED: Middleware configuration.\nTelemetryConfig contains the OpenTelemetry configuration", - "properties": { - "caCertPath": { - "description": "CACertPath is the file path to a CA certificate bundle for the OTLP endpoint.\nWhen set, the OTLP exporters use this CA to verify the collector's TLS certificate\ninstead of relying solely on the system CA pool.\n+optional", - "type": "string" }, - "customAttributes": { - "additionalProperties": { + "failed": { + "description": "Failed lists skills that could not be synced, with the reason for each.\nDrift alone is never reported here — see Drifted.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SyncFailure" + }, + "type": "array", + "uniqueItems": false + }, + "installed": { + "description": "Installed lists skills that were installed or reinstalled to match the lock file.", + "items": { "type": "string" }, - "description": "CustomAttributes contains custom resource attributes to be added to all telemetry signals.\nThese are parsed from CLI flags (--otel-custom-attributes) or environment variables\n(OTEL_RESOURCE_ATTRIBUTES) as key=value pairs.\n+optional", - "type": "object" + "type": "array", + "uniqueItems": false }, - "enablePrometheusMetricsPath": { - "description": "EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint.\nThe metrics are served on the main transport port at /metrics.\nThis is separate from OTLP metrics which are sent to the Endpoint.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "missing": { + "description": "Missing lists lock entries with no corresponding install record at all\n— the fresh-clone state. Normally these are installed at their pinned\nreference; when Check is set, nothing is written and this field\nreports the gap only.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "endpoint": { - "description": "Endpoint is the OTLP endpoint URL\n+optional", - "type": "string" + "never_managed": { + "description": "NeverManaged lists project-scoped skills never recorded as lock-managed.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "environmentVariables": { - "description": "EnvironmentVariables is a list of environment variable names that should be\nincluded in telemetry spans as attributes. Only variables in this list will\nbe read from the host machine and included in spans for observability.\nExample: [\"NODE_ENV\", \"DEPLOYMENT_ENV\", \"SERVICE_VERSION\"]\n+optional", + "pruned": { + "description": "Pruned lists removed-from-lock skills that were uninstalled because Prune was set.", "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "headers": { - "additionalProperties": { + "removed_from_lock": { + "description": "RemovedFromLock lists previously managed skills absent from the lock file.", + "items": { "type": "string" }, - "description": "Headers contains authentication headers for the OTLP endpoint.\n+optional", - "type": "object" + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_skills.UpgradeOutcome": { + "properties": { + "error": { + "description": "Error is a human-readable description of the failure, set only when Status is UpgradeStatusFailed.", + "type": "string" }, - "insecure": { - "description": "Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "name": { + "description": "Name is the skill name.", + "type": "string" }, - "metricsEnabled": { - "description": "MetricsEnabled controls whether OTLP metrics are enabled.\nWhen false, OTLP metrics are not sent even if an endpoint is configured.\nThis is independent of EnablePrometheusMetricsPath.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "new_digest": { + "description": "NewDigest is the digest the source currently resolves to. Equal to\nOldDigest when Status is UpgradeStatusUpToDate.", + "type": "string" }, - "samplingRate": { - "description": "SamplingRate is the trace sampling rate (0.0-1.0) as a string.\nOnly used when TracingEnabled is true.\nExample: \"0.05\" for 5% sampling.\n+kubebuilder:default=\"0.05\"\n+optional", + "new_resolved_reference": { + "description": "NewResolvedReference is the new resolvedReference when it changed.", "type": "string" }, - "serviceName": { - "description": "ServiceName is the service name for telemetry.\nWhen omitted, defaults to the server name (e.g., VirtualMCPServer name).\n+optional", + "new_signer_identity": { + "description": "NewSignerIdentity is the candidate's signer identity when it differs\nfrom the recorded one (empty when the candidate is unsigned).", "type": "string" }, - "serviceVersion": { - "description": "ServiceVersion is the service version for telemetry.\nWhen omitted, defaults to the ToolHive version.\n+optional", + "old_digest": { + "description": "OldDigest is the digest pinned in the lock file before this operation.", "type": "string" }, - "tracingEnabled": { - "description": "TracingEnabled controls whether distributed tracing is enabled.\nWhen false, no tracer provider is created even if an endpoint is configured.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "reason": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.FailureReason" }, - "useLegacyAttributes": { - "description": "UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names\nare emitted alongside the new standard attribute names. When true, spans include both\nold and new attribute names for backward compatibility with existing dashboards.\nCurrently defaults to true; this will change to false in a future release.\n+kubebuilder:default=true\n+optional", - "type": "boolean" + "status": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeStatus" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig": { + "github_com_stacklok_toolhive_pkg_skills.UpgradeResult": { "properties": { - "parameters": { - "description": "Parameters is a JSON object containing the middleware parameters.\nIt is stored as a raw message to allow flexible parameter types.", - "type": "object" - }, - "type": { - "description": "Type is a string representing the middleware type.", - "type": "string" + "outcomes": { + "description": "Outcomes contains one entry per skill considered for upgrade.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeOutcome" + }, + "type": "array", + "uniqueItems": false } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_transport_types.ProxyMode": { - "description": "ProxyMode is the effective HTTP protocol the proxy uses.\nFor stdio transports, this is the configured mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this matches the transport type.\nNote: \"sse\" is deprecated; use \"streamable-http\" instead.", + "github_com_stacklok_toolhive_pkg_skills.UpgradeStatus": { + "description": "Status is the outcome of the upgrade attempt.", "enum": [ - "sse", - "streamable-http", - "sse", - "streamable-http" + "upgraded", + "up-to-date", + "not-upgradable", + "ref-change-blocked", + "signer-change-blocked", + "failed" ], "type": "string", "x-enum-varnames": [ - "ProxyModeSSE", - "ProxyModeStreamableHTTP" + "UpgradeStatusUpgraded", + "UpgradeStatusUpToDate", + "UpgradeStatusNotUpgradable", + "UpgradeStatusRefChangeBlocked", + "UpgradeStatusSignerChangeBlocked", + "UpgradeStatusFailed" ] }, - "github_com_stacklok_toolhive_pkg_transport_types.TransportType": { - "description": "Transport is the transport mode (stdio, sse, or streamable-http)", - "enum": [ - "stdio", - "sse", - "streamable-http", - "inspector", - "stdio", - "sse", - "streamable-http", - "inspector", - "stdio", - "sse", - "streamable-http", - "inspector" - ], - "type": "string", - "x-enum-varnames": [ - "TransportTypeStdio", - "TransportTypeSSE", - "TransportTypeStreamableHTTP", - "TransportTypeInspector" - ] + "github_com_stacklok_toolhive_pkg_skills.ValidationResult": { + "properties": { + "errors": { + "description": "Errors is a list of validation errors, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "valid": { + "description": "Valid indicates whether the skill definition is valid.", + "type": "boolean" + }, + "warnings": { + "description": "Warnings is a list of non-blocking validation warnings, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" }, "github_com_stacklok_toolhive_pkg_webhook.Config": { "properties": { @@ -2155,6 +2227,20 @@ const docTemplate = `{ "StatusUnknown" ] }, + "ignore.Config": { + "description": "IgnoreConfig contains configuration for ignore processing", + "properties": { + "loadGlobal": { + "description": "Whether to load global ignore patterns", + "type": "boolean" + }, + "printOverlays": { + "description": "Whether to print resolved overlay paths for debugging", + "type": "boolean" + } + }, + "type": "object" + }, "model.Argument": { "properties": { "choices": { @@ -2633,6 +2719,20 @@ const docTemplate = `{ }, "type": "object" }, + "pkg_api_v1.buildPluginRequest": { + "description": "Request to build a plugin from a local directory", + "properties": { + "path": { + "description": "Path to the plugin definition directory", + "type": "string" + }, + "tag": { + "description": "OCI tag for the built artifact", + "type": "string" + } + }, + "type": "object" + }, "pkg_api_v1.buildSkillRequest": { "description": "Request to build a skill from a local directory", "properties": { @@ -2825,12 +2925,12 @@ const docTemplate = `{ "type": "string" }, "runtime_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig" + "$ref": "#/components/schemas/templates.RuntimeConfig" }, "secrets": { "description": "Secret parameters to inject", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "type": "array", "uniqueItems": false @@ -2946,7 +3046,7 @@ const docTemplate = `{ "type": "string" }, "registry": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive-core_registry_types.Registry" + "$ref": "#/components/schemas/registry.Registry" }, "server_count": { "description": "Number of servers in the registry", @@ -3032,9 +3132,59 @@ const docTemplate = `{ }, "type": "object" }, + "pkg_api_v1.installPluginRequest": { + "description": "Request to install a plugin", + "properties": { + "clients": { + "description": "Clients lists target client identifiers (e.g., \"claude-code\"),\nor [\"all\"] to target every plugin-supporting client.\nOmitting this field installs to all available clients.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "force": { + "description": "Force allows overwriting unmanaged plugin directories", + "type": "boolean" + }, + "group": { + "description": "Group is the group name to add the plugin to after installation", + "type": "string" + }, + "name": { + "description": "Name or OCI reference of the plugin to install", + "type": "string" + }, + "project_root": { + "description": "ProjectRoot is the project root path for project-scoped installs", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Scope" + }, + "version": { + "description": "Version to install (empty means latest)", + "type": "string" + } + }, + "type": "object" + }, + "pkg_api_v1.installPluginResponse": { + "description": "Response after successfully installing a plugin", + "properties": { + "plugin": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin" + } + }, + "type": "object" + }, "pkg_api_v1.installSkillRequest": { "description": "Request to install a skill", "properties": { + "allow_unsigned": { + "description": "AllowUnsigned permits installing a project-scoped skill without a\nverified signature; the exception is recorded in the project's lock\nfile.", + "type": "boolean" + }, "clients": { "description": "Clients lists target client identifiers (e.g., \"claude-code\"),\nor [\"all\"] to target every skill-supporting client.\nOmitting this field installs to all available clients.", "items": { @@ -3170,6 +3320,51 @@ const docTemplate = `{ }, "type": "object" }, + "pkg_api_v1.pluginBuildListResponse": { + "description": "Response containing a list of locally-built OCI plugin artifacts", + "properties": { + "builds": { + "description": "List of locally-built OCI plugin artifacts", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.LocalBuild" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "pkg_api_v1.pluginListResponse": { + "description": "Response containing a list of installed plugins", + "properties": { + "plugins": { + "description": "List of installed plugins", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "pkg_api_v1.pluginsV01Response": { + "description": "Paginated list of plugins from the registry", + "properties": { + "metadata": { + "$ref": "#/components/schemas/pkg_api_v1.paginationV01Metadata" + }, + "plugins": { + "description": "Plugins is the list of plugins on the current page", + "items": { + "$ref": "#/components/schemas/registry.Plugin" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, "pkg_api_v1.providerCapabilitiesResponse": { "description": "Capabilities of the secrets provider", "properties": { @@ -3196,6 +3391,16 @@ const docTemplate = `{ }, "type": "object" }, + "pkg_api_v1.pushPluginRequest": { + "description": "Request to push a built plugin artifact", + "properties": { + "reference": { + "description": "OCI reference to push", + "type": "string" + } + }, + "type": "object" + }, "pkg_api_v1.pushSkillRequest": { "description": "Request to push a built skill artifact", "properties": { @@ -3282,7 +3487,7 @@ const docTemplate = `{ "type": "string" }, "bearer_token": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "callback_port": { "description": "Specific port for OAuth callback server", @@ -3293,7 +3498,7 @@ const docTemplate = `{ "type": "string" }, "client_secret": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "issuer": { "description": "OAuth/OIDC issuer URL (e.g., https://accounts.google.com)", @@ -3423,25 +3628,59 @@ const docTemplate = `{ }, "type": "object" }, - "pkg_api_v1.toolOverride": { - "description": "Tool override", + "pkg_api_v1.syncSkillsRequest": { + "description": "Request to restore a project's installed skills to match its lock file", "properties": { - "description": { - "description": "Description of the tool", - "type": "string" + "adopt": { + "description": "Adopt writes lock entries for existing unmanaged project-scope installs", + "type": "boolean" }, - "name": { - "description": "Name of the tool", - "type": "string" - } - }, - "type": "object" - }, - "pkg_api_v1.updateRequest": { - "description": "Request to update an existing workload (name cannot be changed)", - "properties": { - "allow_docker_gateway": { - "description": "Whether to permit outbound connections to Docker gateway addresses\n(host.docker.internal, gateway.docker.internal, 172.17.0.1). These are\nblocked by default in the egress proxy even when network isolation is on.\nOnly applicable to Docker deployments with network isolation enabled.", + "allow_unsigned": { + "description": "AllowUnsigned permits adopting skills whose signature state cannot be\nestablished, recording them as unsigned", + "type": "boolean" + }, + "check": { + "description": "Check verifies on-disk content against the lock file without installing or writing anything", + "type": "boolean" + }, + "clients": { + "description": "Clients lists target client identifiers. Empty means every\nskill-supporting client detected on this host.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "project_root": { + "description": "ProjectRoot is the project root path whose lock file should be synced", + "type": "string" + }, + "prune": { + "description": "Prune removes project-scoped skills installed but not present in the lock file", + "type": "boolean" + } + }, + "type": "object" + }, + "pkg_api_v1.toolOverride": { + "description": "Tool override", + "properties": { + "description": { + "description": "Description of the tool", + "type": "string" + }, + "name": { + "description": "Name of the tool", + "type": "string" + } + }, + "type": "object" + }, + "pkg_api_v1.updateRequest": { + "description": "Request to update an existing workload (name cannot be changed)", + "properties": { + "allow_docker_gateway": { + "description": "Whether to permit outbound connections to Docker gateway addresses\n(host.docker.internal, gateway.docker.internal, 172.17.0.1). These are\nblocked by default in the egress proxy even when network isolation is on.\nOnly applicable to Docker deployments with network isolation enabled.", "type": "boolean" }, "authz_config": { @@ -3507,12 +3746,12 @@ const docTemplate = `{ "type": "integer" }, "runtime_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig" + "$ref": "#/components/schemas/templates.RuntimeConfig" }, "secrets": { "description": "Secret parameters to inject", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "type": "array", "uniqueItems": false @@ -3627,6 +3866,58 @@ const docTemplate = `{ }, "type": "object" }, + "pkg_api_v1.upgradeSkillsRequest": { + "description": "Request to re-resolve a project's lock entries and install newer content", + "properties": { + "allow_ref_change": { + "description": "AllowRefChange permits resolvedReference changes during upgrade", + "type": "boolean" + }, + "allow_signer_change": { + "description": "AllowSignerChange permits upgrading to an artifact signed by a\ndifferent identity than the recorded one", + "type": "boolean" + }, + "clients": { + "description": "Clients lists target client identifiers. Empty means every\nskill-supporting client detected on this host.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "fail_on_changes": { + "description": "FailOnChanges exits with an error when any mutable source would upgrade", + "type": "boolean" + }, + "names": { + "description": "Names restricts the upgrade to specific skill names. Empty means every entry.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "preview": { + "description": "Preview reports what would change without installing (still fetches to compare digests)", + "type": "boolean" + }, + "project_root": { + "description": "ProjectRoot is the project root path whose lock file should be upgraded", + "type": "string" + } + }, + "type": "object" + }, + "pkg_api_v1.validatePluginRequest": { + "description": "Request to validate a plugin definition", + "properties": { + "path": { + "description": "Path to the plugin definition directory", + "type": "string" + } + }, + "type": "object" + }, "pkg_api_v1.validateSkillRequest": { "description": "Request to validate a skill definition", "properties": { @@ -3651,7 +3942,7 @@ const docTemplate = `{ "workloads": { "description": "List of container information for each workload", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_core.Workload" + "$ref": "#/components/schemas/core.Workload" }, "type": "array", "uniqueItems": false @@ -3663,7 +3954,21 @@ const docTemplate = `{ "description": "Response containing workload status information", "properties": { "status": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus" + "description": "Current status of the workload", + "enum": [ + "running", + "stopped", + "error", + "starting", + "stopping", + "unhealthy", + "removing", + "unknown", + "unauthenticated", + "auth_retrying", + "policy_stopped" + ], + "type": "string" } }, "type": "object" @@ -3957,6 +4262,68 @@ const docTemplate = `{ }, "type": "object" }, + "registry.Plugin": { + "properties": { + "_meta": { + "additionalProperties": {}, + "description": "Meta is an opaque payload with extended meta data details of the plugin.", + "type": "object" + }, + "description": { + "description": "Description is the description of the plugin.", + "type": "string" + }, + "icons": { + "description": "Icons is the list of icons for the plugin.", + "items": { + "$ref": "#/components/schemas/registry.SkillIcon" + }, + "type": "array", + "uniqueItems": false + }, + "license": { + "description": "License is the SPDX license identifier of the plugin.", + "type": "string" + }, + "metadata": { + "additionalProperties": {}, + "description": "Metadata is the official metadata of the plugin as reported in the\nplugin manifest file.", + "type": "object" + }, + "name": { + "description": "Name is the name of the plugin.\nThe format is that of identifiers, e.g. \"my-plugin\".", + "type": "string" + }, + "namespace": { + "description": "Namespace is the namespace of the plugin.\nThe format is reverse-DNS, e.g. \"io.github.user\".", + "type": "string" + }, + "packages": { + "description": "Packages is the list of packages for the plugin.", + "items": { + "$ref": "#/components/schemas/registry.SkillPackage" + }, + "type": "array", + "uniqueItems": false + }, + "repository": { + "$ref": "#/components/schemas/registry.SkillRepository" + }, + "status": { + "description": "Status is the status of the plugin.\nCan be one of \"active\", \"deprecated\", or \"archived\".", + "type": "string" + }, + "title": { + "description": "Title is the title of the plugin.\nThis is for human consumption, not an identifier.", + "type": "string" + }, + "version": { + "description": "Version is the version of the plugin.\nAny non-empty string is valid, but ideally it should be either a\nsemantic version or a commit hash.", + "type": "string" + } + }, + "type": "object" + }, "registry.Provenance": { "description": "Provenance contains verification and signing metadata", "properties": { @@ -3984,6 +4351,42 @@ const docTemplate = `{ }, "type": "object" }, + "registry.Registry": { + "description": "Full registry data", + "properties": { + "groups": { + "description": "Groups is a slice of group definitions containing related MCP servers", + "items": { + "$ref": "#/components/schemas/registry.Group" + }, + "type": "array", + "uniqueItems": false + }, + "last_updated": { + "description": "LastUpdated is the timestamp when the registry was last updated, in RFC3339 format", + "type": "string" + }, + "remote_servers": { + "additionalProperties": { + "$ref": "#/components/schemas/registry.RemoteServerMetadata" + }, + "description": "RemoteServers is a map of server names to their corresponding remote server definitions\nThese are MCP servers accessed via HTTP/HTTPS using the thv proxy command", + "type": "object" + }, + "servers": { + "additionalProperties": { + "$ref": "#/components/schemas/registry.ImageMetadata" + }, + "description": "Servers is a map of server names to their corresponding server definitions", + "type": "object" + }, + "version": { + "description": "Version is the schema version of the registry", + "type": "string" + } + }, + "type": "object" + }, "registry.RemoteServerMetadata": { "description": "Remote server details (if it's a remote server)", "properties": { @@ -4232,145 +4635,1063 @@ const docTemplate = `{ }, "type": "object" }, - "v0.ServerJSON": { + "remote.Config": { + "description": "RemoteAuthConfig contains OAuth configuration for remote MCP servers", "properties": { - "$schema": { - "example": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json", - "format": "uri", - "minLength": 1, + "authorize_url": { "type": "string" }, - "_meta": { - "$ref": "#/components/schemas/v0.ServerMeta" + "bearer_token": { + "description": "Bearer token configuration (alternative to OAuth)", + "type": "string" }, - "description": { - "example": "MCP server providing weather data and forecasts via OpenWeatherMap API", - "maxLength": 100, - "minLength": 1, + "bearer_token_file": { "type": "string" }, - "icons": { - "items": { - "$ref": "#/components/schemas/model.Icon" - }, - "type": "array", - "uniqueItems": false + "cached_cimd_client_id": { + "description": "CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD\nauthentication was used. Kept separate from CachedClientID (which holds\nDCR-issued IDs) so the two can have independent lifecycles — DCR credential\nrotation clears CachedClientID without touching the stable CIMD URL.\nRead by resolveClientCredentials to send the correct client_id on token refresh.", + "type": "string" }, - "name": { - "example": "io.github.user/weather", - "maxLength": 200, - "minLength": 3, - "pattern": "^[a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+$", + "cached_client_id": { + "description": "Cached DCR client credentials for persistence across restarts.\nThese are obtained during Dynamic Client Registration and needed to refresh tokens.\nClientID is stored as plain text since it's public information.", "type": "string" }, - "packages": { - "items": { - "$ref": "#/components/schemas/model.Package" - }, - "type": "array", - "uniqueItems": false + "cached_client_secret_ref": { + "type": "string" }, - "remotes": { - "items": { - "$ref": "#/components/schemas/model.Transport" - }, - "type": "array", - "uniqueItems": false + "cached_dcr_callback_port": { + "description": "CachedDCRCallbackPort is the callback port that was actually registered\nduring DCR. It may differ from CallbackPort when the requested port was\nunavailable and a fallback port was selected.", + "type": "integer" }, - "repository": { - "$ref": "#/components/schemas/model.Repository" + "cached_refresh_token_ref": { + "description": "Cached OAuth token reference for persistence across restarts.\nThe refresh token is stored securely in the secret manager, and this field\ncontains the reference to retrieve it (e.g., \"OAUTH_REFRESH_TOKEN_workload\").\nThis enables session restoration without requiring a new browser-based login.", + "type": "string" }, - "title": { - "example": "Weather API", - "maxLength": 100, - "minLength": 1, + "cached_reg_client_uri": { + "description": "CachedRegClientURI is the registration_client_uri from the DCR response.\nThis is the endpoint used for RFC 7592 client read/update/delete operations.\nStored as plain text since it is not sensitive.", "type": "string" }, - "version": { - "example": "1.0.2", - "maxLength": 255, - "minLength": 1, + "cached_reg_token_ref": { + "description": "CachedRegTokenRef is a secret manager reference to the registration_access_token\nreturned in the DCR response. Used for RFC 7592 client update operations.\nStored as a secret reference since it's sensitive.", "type": "string" }, - "websiteUrl": { - "example": "https://modelcontextprotocol.io/examples", - "format": "uri", + "cached_secret_expiry": { + "description": "ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server).\nA zero value means the secret does not expire.", "type": "string" - } - }, - "type": "object" - }, - "v0.ServerMeta": { - "properties": { - "io.modelcontextprotocol.registry/publisher-provided": { - "additionalProperties": {}, - "type": "object" - } - }, - "type": "object" - }, - "v1.Duration": { - "description": "RefillPeriod is the duration to fully refill the bucket from zero to maxTokens.\nThe effective refill rate is maxTokens / refillPeriod tokens per second.\nFormat: Go duration string (e.g., \"1m0s\", \"30s\", \"1h0m0s\").\n+kubebuilder:validation:Required", - "type": "object" - } - } - }, - "info": { - "description": "{{escape .Description}}", - "title": "{{.Title}}", - "version": "{{.Version}}" - }, - "externalDocs": { - "description": "", - "url": "" - }, - "paths": { + }, + "cached_token_auth_method": { + "description": "CachedTokenEndpointAuthMethod is the auth method used for the token endpoint\n(e.g., \"client_secret_basic\", \"none\"). Persisted for RFC 7592 updates.", + "type": "string" + }, + "cached_token_expiry": { + "type": "string" + }, + "callback_port": { + "type": "integer" + }, + "client_id": { + "type": "string" + }, + "client_secret": { + "type": "string" + }, + "client_secret_file": { + "type": "string" + }, + "issuer": { + "description": "OAuth endpoint configuration (from registry)", + "type": "string" + }, + "oauth_params": { + "additionalProperties": { + "type": "string" + }, + "description": "OAuth parameters for server-specific customization", + "type": "object" + }, + "resource": { + "description": "Resource is the OAuth 2.0 resource indicator (RFC 8707).", + "type": "string" + }, + "scope_param_name": { + "description": "ScopeParamName overrides the query parameter name used to send scopes in the\nauthorization URL. When empty, the standard \"scope\" parameter is used.\nSome providers require a non-standard name (e.g., Slack uses \"user_scope\").", + "type": "string" + }, + "scopes": { + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "skip_browser": { + "type": "boolean" + }, + "timeout": { + "example": "5m", + "type": "string" + }, + "token_url": { + "type": "string" + }, + "use_pkce": { + "type": "boolean" + } + }, + "type": "object" + }, + "secrets.SecretParameter": { + "description": "Bearer token for authentication (alternative to OAuth)", + "properties": { + "name": { + "type": "string" + }, + "target": { + "type": "string" + } + }, + "type": "object" + }, + "storage.ACLUserRunConfig": { + "description": "ACLUserConfig contains ACL user authentication configuration.", + "properties": { + "password_env_var": { + "description": "PasswordEnvVar is the environment variable containing the Redis password.", + "type": "string" + }, + "username_env_var": { + "description": "UsernameEnvVar is the environment variable containing the Redis username.", + "type": "string" + } + }, + "type": "object" + }, + "storage.RedisRunConfig": { + "description": "RedisConfig is the Redis-specific configuration when Type is \"redis\".", + "properties": { + "acl_user_config": { + "$ref": "#/components/schemas/storage.ACLUserRunConfig" + }, + "addr": { + "description": "Addr is the Redis server address (host:port). Required for standalone and cluster modes.\nMutually exclusive with SentinelConfig.", + "type": "string" + }, + "auth_type": { + "description": "AuthType must be \"aclUser\" - only ACL user authentication is supported.", + "type": "string" + }, + "cluster_mode": { + "description": "ClusterMode enables the Redis Cluster protocol. Requires Addr to be set.", + "type": "boolean" + }, + "dial_timeout": { + "description": "DialTimeout is the timeout for establishing connections (e.g., \"5s\").", + "type": "string" + }, + "key_prefix": { + "description": "KeyPrefix for multi-tenancy, typically \"thv:auth:{ns}:{name}:\".", + "type": "string" + }, + "read_timeout": { + "description": "ReadTimeout is the timeout for read operations (e.g., \"3s\").", + "type": "string" + }, + "sentinel_config": { + "$ref": "#/components/schemas/storage.SentinelRunConfig" + }, + "sentinel_tls": { + "$ref": "#/components/schemas/storage.RedisTLSRunConfig" + }, + "tls": { + "$ref": "#/components/schemas/storage.RedisTLSRunConfig" + }, + "write_timeout": { + "description": "WriteTimeout is the timeout for write operations (e.g., \"3s\").", + "type": "string" + } + }, + "type": "object" + }, + "storage.RedisTLSRunConfig": { + "description": "SentinelTLS configures TLS for Sentinel connections. Only applies when SentinelConfig is set.", + "properties": { + "ca_cert_file": { + "description": "CACertFile is the path to a PEM-encoded CA certificate file.", + "type": "string" + }, + "insecure_skip_verify": { + "description": "InsecureSkipVerify skips certificate verification.", + "type": "boolean" + } + }, + "type": "object" + }, + "storage.RunConfig": { + "description": "Storage configures the storage backend for the auth server.\nIf nil, defaults to in-memory storage.", + "properties": { + "redis_config": { + "$ref": "#/components/schemas/storage.RedisRunConfig" + }, + "type": { + "description": "Type specifies the storage backend type. Defaults to \"memory\".", + "type": "string" + } + }, + "type": "object" + }, + "storage.SentinelRunConfig": { + "description": "SentinelConfig contains Sentinel-specific configuration.\nMutually exclusive with Addr.", + "properties": { + "db": { + "description": "DB is the Redis database number (default: 0).", + "type": "integer" + }, + "master_name": { + "description": "MasterName is the name of the Redis Sentinel master.", + "type": "string" + }, + "sentinel_addrs": { + "description": "SentinelAddrs is the list of Sentinel addresses (host:port).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "telemetry.Config": { + "description": "DEPRECATED: Middleware configuration.\nTelemetryConfig contains the OpenTelemetry configuration", + "properties": { + "caCertPath": { + "description": "CACertPath is the file path to a CA certificate bundle for the OTLP endpoint.\nWhen set, the OTLP exporters use this CA to verify the collector's TLS certificate\ninstead of relying solely on the system CA pool.\n+optional", + "type": "string" + }, + "customAttributes": { + "additionalProperties": { + "type": "string" + }, + "description": "CustomAttributes contains custom resource attributes to be added to all telemetry signals.\nThese are parsed from CLI flags (--otel-custom-attributes) or environment variables\n(OTEL_RESOURCE_ATTRIBUTES) as key=value pairs.\n+optional", + "type": "object" + }, + "enablePrometheusMetricsPath": { + "description": "EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint.\nThe metrics are served on the main transport port at /metrics.\nThis is separate from OTLP metrics which are sent to the Endpoint.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "endpoint": { + "description": "Endpoint is the OTLP endpoint URL\n+optional", + "type": "string" + }, + "environmentVariables": { + "description": "EnvironmentVariables is a list of environment variable names that should be\nincluded in telemetry spans as attributes. Only variables in this list will\nbe read from the host machine and included in spans for observability.\nExample: [\"NODE_ENV\", \"DEPLOYMENT_ENV\", \"SERVICE_VERSION\"]\n+optional", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "headers": { + "additionalProperties": { + "type": "string" + }, + "description": "Headers contains authentication headers for the OTLP endpoint.\n+optional", + "type": "object" + }, + "insecure": { + "description": "Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "metricsEnabled": { + "description": "MetricsEnabled controls whether OTLP metrics are enabled.\nWhen false, OTLP metrics are not sent even if an endpoint is configured.\nThis is independent of EnablePrometheusMetricsPath.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "samplingRate": { + "description": "SamplingRate is the trace sampling rate (0.0-1.0) as a string.\nOnly used when TracingEnabled is true.\nExample: \"0.05\" for 5% sampling.\n+kubebuilder:default=\"0.05\"\n+optional", + "type": "string" + }, + "serviceName": { + "description": "ServiceName is the service name for telemetry.\nWhen omitted, defaults to the server name (e.g., VirtualMCPServer name).\n+optional", + "type": "string" + }, + "serviceVersion": { + "description": "ServiceVersion is the service version for telemetry.\nWhen omitted, defaults to the ToolHive version.\n+optional", + "type": "string" + }, + "tracingEnabled": { + "description": "TracingEnabled controls whether distributed tracing is enabled.\nWhen false, no tracer provider is created even if an endpoint is configured.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "useLegacyAttributes": { + "description": "UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names\nare emitted alongside the new standard attribute names. When true, spans include both\nold and new attribute names for backward compatibility with existing dashboards.\nCurrently defaults to true; this will change to false in a future release.\n+kubebuilder:default=true\n+optional", + "type": "boolean" + } + }, + "type": "object" + }, + "templates.RuntimeConfig": { + "description": "RuntimeConfig allows overriding the default runtime configuration\nfor this specific workload (base images and packages)", + "properties": { + "additional_packages": { + "description": "AdditionalPackages lists extra packages to install in the builder and\nruntime stages.\nExamples for Alpine: [\"git\", \"make\", \"gcc\"]\nExamples for Debian: [\"git\", \"build-essential\"]", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "build_with": { + "description": "BuildWith lists build-time dependency constraints, interpreted per\npackage ecosystem. For uvx:// builds these are PEP 508 requirement\nspecifiers passed to ` + "`" + `uv tool install --with` + "`" + `, used to constrain\ntransitive dependencies the package itself leaves unbounded\n(e.g. \"mcp\u003c2\"). Ecosystems without constraint support (npx://, go://)\nreject a non-empty BuildWith at build time.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "builder_image": { + "description": "BuilderImage is the full image reference for the builder stage.\nAn empty string signals \"use the default for this transport type\" during config merging.\nExamples: \"golang:1.26-alpine\", \"node:24-alpine\", \"python:3.14-slim\"", + "type": "string" + }, + "runtime_env": { + "additionalProperties": { + "type": "string" + }, + "description": "RuntimeEnv contains environment variables to inject into the Dockerfile's\nfinal runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv),\nwhich only affects the builder stage, these variables are baked into the\nshipped image and are present in the running container's process\nenvironment at startup. Use this for values a packaged MCP server reads at\nprocess start (e.g. feature flags, cache backend selection), not for\nbuild-time package manager configuration.\nKeys must be uppercase with underscores, values are validated for safety.", + "type": "object" + } + }, + "type": "object" + }, + "tokenexchange.Config": { + "description": "TokenExchangeConfig contains token exchange configuration for external authentication", + "properties": { + "audience": { + "description": "Audience is the target audience for the exchanged token", + "type": "string" + }, + "client_id": { + "description": "ClientID is the OAuth 2.0 client identifier", + "type": "string" + }, + "client_secret": { + "description": "ClientSecret is the OAuth 2.0 client secret", + "type": "string" + }, + "external_token_header_name": { + "description": "ExternalTokenHeaderName is the name of the custom header to use when HeaderStrategy is \"custom\"", + "type": "string" + }, + "header_strategy": { + "description": "HeaderStrategy determines how to inject the token\nValid values: HeaderStrategyReplace (default), HeaderStrategyCustom", + "type": "string" + }, + "scopes": { + "description": "Scopes is the list of scopes to request for the exchanged token", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "subject_token_type": { + "description": "SubjectTokenType specifies the type of the subject token being exchanged.\nCommon values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT.\nIf empty, defaults to oauthproto.TokenTypeAccessToken.", + "type": "string" + }, + "token_url": { + "description": "TokenURL is the OAuth 2.0 token endpoint URL", + "type": "string" + } + }, + "type": "object" + }, + "types.MiddlewareConfig": { + "properties": { + "parameters": { + "description": "Parameters is a JSON object containing the middleware parameters.\nIt is stored as a raw message to allow flexible parameter types.", + "type": "object" + }, + "type": { + "description": "Type is a string representing the middleware type.", + "type": "string" + } + }, + "type": "object" + }, + "v0.ServerJSON": { + "properties": { + "$schema": { + "example": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json", + "format": "uri", + "minLength": 1, + "type": "string" + }, + "_meta": { + "$ref": "#/components/schemas/v0.ServerMeta" + }, + "description": { + "example": "MCP server providing weather data and forecasts via OpenWeatherMap API", + "maxLength": 100, + "minLength": 1, + "type": "string" + }, + "icons": { + "items": { + "$ref": "#/components/schemas/model.Icon" + }, + "type": "array", + "uniqueItems": false + }, + "name": { + "example": "io.github.user/weather", + "maxLength": 200, + "minLength": 3, + "pattern": "^[a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+$", + "type": "string" + }, + "packages": { + "items": { + "$ref": "#/components/schemas/model.Package" + }, + "type": "array", + "uniqueItems": false + }, + "remotes": { + "items": { + "$ref": "#/components/schemas/model.Transport" + }, + "type": "array", + "uniqueItems": false + }, + "repository": { + "$ref": "#/components/schemas/model.Repository" + }, + "title": { + "example": "Weather API", + "maxLength": 100, + "minLength": 1, + "type": "string" + }, + "version": { + "example": "1.0.2", + "maxLength": 255, + "minLength": 1, + "type": "string" + }, + "websiteUrl": { + "example": "https://modelcontextprotocol.io/examples", + "format": "uri", + "type": "string" + } + }, + "type": "object" + }, + "v0.ServerMeta": { + "properties": { + "io.modelcontextprotocol.registry/publisher-provided": { + "additionalProperties": {}, + "type": "object" + } + }, + "type": "object" + }, + "v1.Duration": { + "description": "RefillPeriod is the duration to fully refill the bucket from zero to maxTokens.\nThe effective refill rate is maxTokens / refillPeriod tokens per second.\nFormat: Go duration string (e.g., \"1m0s\", \"30s\", \"1h0m0s\").\n+kubebuilder:validation:Required", + "type": "object" + } + } + }, + "info": { + "description": "{{escape .Description}}", + "title": "{{.Title}}", + "version": "{{.Version}}" + }, + "externalDocs": { + "description": "", + "url": "" + }, + "paths": { "/api/openapi.json": { "get": { - "description": "Returns the OpenAPI specification for the API", + "description": "Returns the OpenAPI specification for the API", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "type": "object" + } + } + }, + "description": "OpenAPI specification" + } + }, + "summary": "Get OpenAPI specification", + "tags": [ + "system" + ] + } + }, + "/api/v1beta/clients": { + "get": { + "description": "List all registered clients in ToolHive", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_client.RegisteredClient" + }, + "type": "array" + } + } + }, + "description": "OK" + } + }, + "summary": "List all clients", + "tags": [ + "clients" + ] + }, + "post": { + "description": "Register a new client with ToolHive", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.createClientRequest", + "summary": "client", + "description": "Client to register" + } + ] + } + } + }, + "description": "Client to register", + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" + } + } + }, + "description": "OK" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Register a new client", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/register": { + "post": { + "description": "Register multiple clients with ToolHive", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", + "summary": "clients", + "description": "Clients to register" + } + ] + } + } + }, + "description": "Clients to register", + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" + }, + "type": "array" + } + } + }, + "description": "OK" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Register multiple clients", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/unregister": { + "post": { + "description": "Unregister multiple clients from ToolHive", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", + "summary": "clients", + "description": "Clients to unregister" + } + ] + } + } + }, + "description": "Clients to unregister", + "required": true + }, + "responses": { + "204": { + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Unregister multiple clients", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/{name}": { + "delete": { + "description": "Unregister a client from ToolHive", + "parameters": [ + { + "description": "Client name to unregister", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Unregister a client", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/{name}/groups/{group}": { + "delete": { + "description": "Unregister a client from a specific group in ToolHive", + "parameters": [ + { + "description": "Client name to unregister", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Group name to remove client from", + "in": "path", + "name": "group", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Client or group not found" + } + }, + "summary": "Unregister a client from a specific group", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/discovery/clients": { + "get": { + "description": "List all clients compatible with ToolHive and their status.\nEach object includes supports_skills when ToolHive can install skills for that client.", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.clientStatusResponse" + } + } + }, + "description": "OK" + } + }, + "summary": "List all clients status", + "tags": [ + "discovery" + ] + } + }, + "/api/v1beta/groups": { + "get": { + "description": "Get a list of all groups", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.groupListResponse" + } + } + }, + "description": "OK" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "List all groups", + "tags": [ + "groups" + ] + }, + "post": { + "description": "Create a new group with the specified name", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.createGroupRequest", + "summary": "group", + "description": "Group creation request" + } + ] + } + } + }, + "description": "Group creation request", + "required": true + }, + "responses": { + "201": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.createGroupResponse" + } + } + }, + "description": "Created" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "409": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Conflict" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "Create a new group", + "tags": [ + "groups" + ] + } + }, + "/api/v1beta/groups/{name}": { + "delete": { + "description": "Delete a group by name.", + "parameters": [ + { + "description": "Group name", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Delete all workloads in the group (default: false, moves workloads to default group)", + "in": "query", + "name": "with-workloads", + "schema": { + "type": "boolean" + } + } + ], + "responses": { + "204": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "No Content" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "Delete a group", + "tags": [ + "groups" + ] + }, + "get": { + "description": "Get details of a specific group", + "parameters": [ + { + "description": "Group name", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + } + ], "responses": { "200": { "content": { "application/json": { "schema": { - "type": "object" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_groups.Group" } } }, - "description": "OpenAPI specification" + "description": "OK" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "Get OpenAPI specification", + "summary": "Get group details", "tags": [ - "system" + "groups" ] } }, - "/api/v1beta/clients": { + "/api/v1beta/plugins": { "get": { - "description": "List all registered clients in ToolHive", + "description": "Get a list of all installed plugins", + "parameters": [ + { + "description": "Filter by scope (user or project)", + "in": "query", + "name": "scope", + "schema": { + "enum": [ + "user", + "project" + ], + "type": "string" + } + }, + { + "description": "Filter by client app", + "in": "query", + "name": "client", + "schema": { + "type": "string" + } + }, + { + "description": "Filter by project root path", + "in": "query", + "name": "project_root", + "schema": { + "type": "string" + } + }, + { + "description": "Filter by group name", + "in": "query", + "name": "group", + "schema": { + "type": "string" + } + } + ], "responses": { "200": { "content": { "application/json": { "schema": { - "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_client.RegisteredClient" - }, - "type": "array" + "$ref": "#/components/schemas/pkg_api_v1.pluginListResponse" } } }, "description": "OK" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "List all clients", + "summary": "List all installed plugins", "tags": [ - "clients" + "plugins" ] }, "post": { - "description": "Register a new client with ToolHive", + "description": "Install a plugin from a remote source", "requestBody": { "content": { "application/json": { @@ -4380,29 +5701,107 @@ const docTemplate = `{ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.createClientRequest", - "summary": "client", - "description": "Client to register" + "$ref": "#/components/schemas/pkg_api_v1.installPluginRequest", + "summary": "request", + "description": "Install request" } ] } } }, - "description": "Client to register", - "required": true - }, - "responses": { - "200": { + "description": "Install request", + "required": true + }, + "responses": { + "201": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.installPluginResponse" + } + } + }, + "description": "Created", + "headers": { + "Location": { + "description": "URI of the installed plugin resource", + "schema": { + "type": "string" + } + } + } + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "401": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Unauthorized (registry refused credentials)" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found (artifact not present in registry)" + }, + "409": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Conflict" + }, + "429": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Too Many Requests (registry rate limit)" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + }, + "502": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" + "type": "string" } } }, - "description": "OK" + "description": "Bad Gateway (upstream registry failure)" }, - "400": { + "504": { "content": { "application/json": { "schema": { @@ -4410,18 +5809,18 @@ const docTemplate = `{ } } }, - "description": "Invalid request or unsupported client type" + "description": "Gateway Timeout (upstream pull timed out)" } }, - "summary": "Register a new client", + "summary": "Install a plugin", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/register": { + "/api/v1beta/plugins/build": { "post": { - "description": "Register multiple clients with ToolHive", + "description": "Build a plugin from a local directory", "requestBody": { "content": { "application/json": { @@ -4431,15 +5830,15 @@ const docTemplate = `{ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", - "summary": "clients", - "description": "Clients to register" + "$ref": "#/components/schemas/pkg_api_v1.buildPluginRequest", + "summary": "request", + "description": "Build request" } ] } } }, - "description": "Clients to register", + "description": "Build request", "required": true }, "responses": { @@ -4447,10 +5846,7 @@ const docTemplate = `{ "content": { "application/json": { "schema": { - "items": { - "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" - }, - "type": "array" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.BuildResult" } } }, @@ -4464,43 +5860,40 @@ const docTemplate = `{ } } }, - "description": "Invalid request or unsupported client type" + "description": "Bad Request" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "Register multiple clients", + "summary": "Build a plugin", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/unregister": { - "post": { - "description": "Unregister multiple clients from ToolHive", - "requestBody": { - "content": { - "application/json": { - "schema": { - "oneOf": [ - { - "type": "object" - }, - { - "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", - "summary": "clients", - "description": "Clients to unregister" - } - ] - } - } - }, - "description": "Clients to unregister", - "required": true - }, + "/api/v1beta/plugins/builds": { + "get": { + "description": "Get a list of all locally-built OCI plugin artifacts in the local store", "responses": { - "204": { - "description": "No Content" + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.pluginBuildListResponse" + } + } + }, + "description": "OK" }, - "400": { + "500": { "content": { "application/json": { "schema": { @@ -4508,23 +5901,23 @@ const docTemplate = `{ } } }, - "description": "Invalid request or unsupported client type" + "description": "Internal Server Error" } }, - "summary": "Unregister multiple clients", + "summary": "List locally-built plugin artifacts", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/{name}": { + "/api/v1beta/plugins/builds/{tag}": { "delete": { - "description": "Unregister a client from ToolHive", + "description": "Remove a locally-built OCI plugin artifact and its blobs from the local store", "parameters": [ { - "description": "Client name to unregister", + "description": "Artifact tag", "in": "path", - "name": "name", + "name": "tag", "required": true, "schema": { "type": "string" @@ -4533,9 +5926,16 @@ const docTemplate = `{ ], "responses": { "204": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, "description": "No Content" }, - "400": { + "404": { "content": { "application/json": { "schema": { @@ -4543,32 +5943,33 @@ const docTemplate = `{ } } }, - "description": "Invalid request or unsupported client type" + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "Unregister a client", + "summary": "Delete a locally-built plugin artifact", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/{name}/groups/{group}": { - "delete": { - "description": "Unregister a client from a specific group in ToolHive", + "/api/v1beta/plugins/content": { + "get": { + "description": "Retrieve the plugin.json body and file listing from an artifact\nwithout installing it. Accepts OCI refs, git refs, or local tags.", "parameters": [ { - "description": "Client name to unregister", - "in": "path", - "name": "name", - "required": true, - "schema": { - "type": "string" - } - }, - { - "description": "Group name to remove client from", - "in": "path", - "name": "group", + "description": "OCI reference or local build tag", + "in": "query", + "name": "ref", "required": true, "schema": { "type": "string" @@ -4576,8 +5977,15 @@ const docTemplate = `{ } ], "responses": { - "204": { - "description": "No Content" + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginContent" + } + } + }, + "description": "OK" }, "400": { "content": { @@ -4587,7 +5995,17 @@ const docTemplate = `{ } } }, - "description": "Invalid request or unsupported client type" + "description": "Bad Request" + }, + "401": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Unauthorized (registry refused credentials)" }, "404": { "content": { @@ -4597,49 +6015,108 @@ const docTemplate = `{ } } }, - "description": "Client or group not found" - } - }, - "summary": "Unregister a client from a specific group", - "tags": [ - "clients" - ] - } - }, - "/api/v1beta/discovery/clients": { - "get": { - "description": "List all clients compatible with ToolHive and their status.\nEach object includes supports_skills when ToolHive can install skills for that client.", - "responses": { - "200": { + "description": "Not Found (artifact not present in registry)" + }, + "429": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.clientStatusResponse" + "type": "string" } } }, - "description": "OK" + "description": "Too Many Requests (registry rate limit)" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + }, + "502": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Gateway (upstream registry or git resolver failure)" + }, + "504": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Gateway Timeout (upstream pull timed out)" } }, - "summary": "List all clients status", + "summary": "Get plugin content", "tags": [ - "discovery" + "plugins" ] } }, - "/api/v1beta/groups": { - "get": { - "description": "Get a list of all groups", + "/api/v1beta/plugins/push": { + "post": { + "description": "Push a built plugin artifact to a remote registry", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.pushPluginRequest", + "summary": "request", + "description": "Push request" + } + ] + } + } + }, + "description": "Push request", + "required": true + }, "responses": { - "200": { + "204": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.groupListResponse" + "type": "string" } } }, - "description": "OK" + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" }, "500": { "content": { @@ -4652,13 +6129,15 @@ const docTemplate = `{ "description": "Internal Server Error" } }, - "summary": "List all groups", + "summary": "Push a plugin", "tags": [ - "groups" + "plugins" ] - }, + } + }, + "/api/v1beta/plugins/validate": { "post": { - "description": "Create a new group with the specified name", + "description": "Validate a plugin definition", "requestBody": { "content": { "application/json": { @@ -4668,27 +6147,27 @@ const docTemplate = `{ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.createGroupRequest", - "summary": "group", - "description": "Group creation request" + "$ref": "#/components/schemas/pkg_api_v1.validatePluginRequest", + "summary": "request", + "description": "Validate request" } ] } } }, - "description": "Group creation request", + "description": "Validate request", "required": true }, "responses": { - "201": { + "200": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.createGroupResponse" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ValidationResult" } } }, - "description": "Created" + "description": "OK" }, "400": { "content": { @@ -4700,16 +6179,6 @@ const docTemplate = `{ }, "description": "Bad Request" }, - "409": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Conflict" - }, "500": { "content": { "application/json": { @@ -4721,18 +6190,18 @@ const docTemplate = `{ "description": "Internal Server Error" } }, - "summary": "Create a new group", + "summary": "Validate a plugin", "tags": [ - "groups" + "plugins" ] } }, - "/api/v1beta/groups/{name}": { + "/api/v1beta/plugins/{name}": { "delete": { - "description": "Delete a group by name.", + "description": "Remove an installed plugin", "parameters": [ { - "description": "Group name", + "description": "Plugin name", "in": "path", "name": "name", "required": true, @@ -4741,11 +6210,23 @@ const docTemplate = `{ } }, { - "description": "Delete all workloads in the group (default: false, moves workloads to default group)", + "description": "Scope to uninstall from (user or project)", "in": "query", - "name": "with-workloads", + "name": "scope", "schema": { - "type": "boolean" + "enum": [ + "user", + "project" + ], + "type": "string" + } + }, + { + "description": "Project root path for project-scoped plugins", + "in": "query", + "name": "project_root", + "schema": { + "type": "string" } } ], @@ -4760,6 +6241,16 @@ const docTemplate = `{ }, "description": "No Content" }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, "404": { "content": { "application/json": { @@ -4781,22 +6272,42 @@ const docTemplate = `{ "description": "Internal Server Error" } }, - "summary": "Delete a group", + "summary": "Uninstall a plugin", "tags": [ - "groups" + "plugins" ] }, "get": { - "description": "Get details of a specific group", + "description": "Get detailed information about a specific plugin", "parameters": [ { - "description": "Group name", + "description": "Plugin name", "in": "path", "name": "name", "required": true, "schema": { "type": "string" } + }, + { + "description": "Filter by scope (user or project)", + "in": "query", + "name": "scope", + "schema": { + "enum": [ + "user", + "project" + ], + "type": "string" + } + }, + { + "description": "Project root path for project-scoped plugins", + "in": "query", + "name": "project_root", + "schema": { + "type": "string" + } } ], "responses": { @@ -4804,12 +6315,22 @@ const docTemplate = `{ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_groups.Group" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginInfo" } } }, "description": "OK" }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, "404": { "content": { "application/json": { @@ -4831,9 +6352,9 @@ const docTemplate = `{ "description": "Internal Server Error" } }, - "summary": "Get group details", + "summary": "Get plugin details", "tags": [ - "groups" + "plugins" ] } }, @@ -6052,32 +7573,214 @@ const docTemplate = `{ "description": "Internal Server Error" } }, - "summary": "Delete a locally-built skill artifact", + "summary": "Delete a locally-built skill artifact", + "tags": [ + "skills" + ] + } + }, + "/api/v1beta/skills/content": { + "get": { + "description": "Retrieve the SKILL.md body and file listing from an artifact\nwithout installing it. Accepts OCI refs, git refs, or local tags.", + "parameters": [ + { + "description": "OCI reference or local build tag", + "in": "query", + "name": "ref", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillContent" + } + } + }, + "description": "OK" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "401": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Unauthorized (registry refused credentials)" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found (artifact not present in registry)" + }, + "429": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Too Many Requests (registry rate limit)" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + }, + "502": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Gateway (upstream registry or git resolver failure)" + }, + "504": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Gateway Timeout (upstream pull timed out)" + } + }, + "summary": "Get skill content", + "tags": [ + "skills" + ] + } + }, + "/api/v1beta/skills/push": { + "post": { + "description": "Push a built skill artifact to a remote registry", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.pushSkillRequest", + "summary": "request", + "description": "Push request" + } + ] + } + } + }, + "description": "Push request", + "required": true + }, + "responses": { + "204": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "Push a skill", "tags": [ "skills" ] } }, - "/api/v1beta/skills/content": { - "get": { - "description": "Retrieve the SKILL.md body and file listing from an artifact\nwithout installing it. Accepts OCI refs, git refs, or local tags.", - "parameters": [ - { - "description": "OCI reference or local build tag", - "in": "query", - "name": "ref", - "required": true, - "schema": { - "type": "string" + "/api/v1beta/skills/sync": { + "post": { + "description": "Restore a project's installed skills to match toolhive.lock.yaml", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.syncSkillsRequest", + "summary": "request", + "description": "Sync request" + } + ] + } } - } - ], + }, + "description": "Sync request", + "required": true + }, "responses": { "200": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillContent" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SyncResult" } } }, @@ -6093,27 +7796,7 @@ const docTemplate = `{ }, "description": "Bad Request" }, - "401": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Unauthorized (registry refused credentials)" - }, - "404": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Not Found (artifact not present in registry)" - }, - "429": { + "403": { "content": { "application/json": { "schema": { @@ -6121,7 +7804,7 @@ const docTemplate = `{ } } }, - "description": "Too Many Requests (registry rate limit)" + "description": "Forbidden (feature not enabled)" }, "500": { "content": { @@ -6133,17 +7816,7 @@ const docTemplate = `{ }, "description": "Internal Server Error" }, - "502": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Bad Gateway (upstream registry or git resolver failure)" - }, - "504": { + "501": { "content": { "application/json": { "schema": { @@ -6151,18 +7824,18 @@ const docTemplate = `{ } } }, - "description": "Gateway Timeout (upstream pull timed out)" + "description": "Not Implemented" } }, - "summary": "Get skill content", + "summary": "Sync project skills from the lock file", "tags": [ "skills" ] } }, - "/api/v1beta/skills/push": { + "/api/v1beta/skills/upgrade": { "post": { - "description": "Push a built skill artifact to a remote registry", + "description": "Re-resolve a project's lock entries and install newer content where available", "requestBody": { "content": { "application/json": { @@ -6172,27 +7845,27 @@ const docTemplate = `{ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.pushSkillRequest", + "$ref": "#/components/schemas/pkg_api_v1.upgradeSkillsRequest", "summary": "request", - "description": "Push request" + "description": "Upgrade request" } ] } } }, - "description": "Push request", + "description": "Upgrade request", "required": true }, "responses": { - "204": { + "200": { "content": { "application/json": { "schema": { - "type": "string" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeResult" } } }, - "description": "No Content" + "description": "OK" }, "400": { "content": { @@ -6204,6 +7877,16 @@ const docTemplate = `{ }, "description": "Bad Request" }, + "403": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Forbidden (feature not enabled)" + }, "404": { "content": { "application/json": { @@ -6212,7 +7895,7 @@ const docTemplate = `{ } } }, - "description": "Not Found" + "description": "Not Found (a requested name is not in the lock file)" }, "500": { "content": { @@ -6223,9 +7906,19 @@ const docTemplate = `{ } }, "description": "Internal Server Error" + }, + "501": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Implemented" } }, - "summary": "Push a skill", + "summary": "Upgrade project skills", "tags": [ "skills" ] @@ -7572,6 +9265,162 @@ const docTemplate = `{ ] } }, + "/registry/{registryName}/v0.1/x/dev.toolhive/plugins": { + "get": { + "description": "Get a paginated list of plugins from the registry. Supports optional full-text search and pagination.", + "parameters": [ + { + "description": "Registry name (currently ignored, uses the default provider)", + "in": "path", + "name": "registryName", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Search filter — matches against plugin name, namespace, and description", + "in": "query", + "name": "q", + "schema": { + "type": "string" + } + }, + { + "description": "Page number, 1-based (default: 1)", + "in": "query", + "name": "page", + "schema": { + "type": "integer" + } + }, + { + "description": "Items per page, max 200 (default: 50)", + "in": "query", + "name": "limit", + "schema": { + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.pluginsV01Response" + } + } + }, + "description": "OK" + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Internal server error" + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Registry authentication required or upstream registry unavailable" + } + }, + "summary": "List available registry plugins", + "tags": [ + "registry-plugins" + ] + } + }, + "/registry/{registryName}/v0.1/x/dev.toolhive/plugins/{namespace}/{pluginName}": { + "get": { + "description": "Retrieve a single plugin by its namespace and name from the registry.", + "parameters": [ + { + "description": "Registry name (currently ignored, uses the default provider)", + "in": "path", + "name": "registryName", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Plugin namespace in reverse-DNS format (e.g. io.github.stacklok)", + "in": "path", + "name": "namespace", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Plugin name", + "in": "path", + "name": "pluginName", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/registry.Plugin" + } + } + }, + "description": "OK" + }, + "404": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Plugin not found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Internal server error" + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Registry authentication required or upstream registry unavailable" + } + }, + "summary": "Get a registry plugin", + "tags": [ + "registry-plugins" + ] + } + }, "/registry/{registryName}/v0.1/x/dev.toolhive/skills": { "get": { "description": "Get a paginated list of skills from the registry. Supports optional full-text search and pagination.", diff --git a/docs/server/swagger.json b/docs/server/swagger.json index cc3cfa8f59..3f34fef57f 100644 --- a/docs/server/swagger.json +++ b/docs/server/swagger.json @@ -1,37 +1,152 @@ { "components": { "schemas": { - "github_com_stacklok_toolhive-core_registry_types.Registry": { - "description": "Full registry data", + "auth.TokenValidatorConfig": { + "description": "DEPRECATED: Middleware configuration.\nOIDCConfig contains OIDC configuration", "properties": { - "groups": { - "description": "Groups is a slice of group definitions containing related MCP servers", + "allowPrivateIP": { + "description": "AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses", + "type": "boolean" + }, + "allowedClientIDs": { + "description": "AllowedClientIDs is the set of OAuth client IDs accepted when validating\ntokens by their client_id claim instead of an audience claim.", "items": { - "$ref": "#/components/schemas/registry.Group" + "type": "string" }, - "type": "array", - "uniqueItems": false + "type": "array" }, - "last_updated": { - "description": "LastUpdated is the timestamp when the registry was last updated, in RFC3339 format", + "audience": { + "description": "Audience is the expected audience for the token", "type": "string" }, - "remote_servers": { - "additionalProperties": { - "$ref": "#/components/schemas/registry.RemoteServerMetadata" + "authTokenFile": { + "description": "AuthTokenFile is the path to file containing bearer token for authentication", + "type": "string" + }, + "cacertPath": { + "description": "CACertPath is the path to the CA certificate bundle for HTTPS requests", + "type": "string" + }, + "clientID": { + "description": "ClientID is the OIDC client ID", + "type": "string" + }, + "clientSecret": { + "description": "ClientSecret is the optional OIDC client secret for introspection", + "type": "string" + }, + "insecureAllowHTTP": { + "description": "InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing\nWARNING: This is insecure and should NEVER be used in production", + "type": "boolean" + }, + "introspectionURL": { + "description": "IntrospectionURL is the optional introspection endpoint for validating tokens", + "type": "string" + }, + "issuer": { + "description": "Issuer is the OIDC issuer URL (e.g., https://accounts.google.com)", + "type": "string" + }, + "jwksurl": { + "description": "JWKSURL is the URL to fetch the JWKS from", + "type": "string" + }, + "resourceURL": { + "description": "ResourceURL is the explicit resource URL for OAuth discovery (RFC 9728)", + "type": "string" + }, + "scopes": { + "description": "Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728)\nIf empty, defaults to [\"openid\"]", + "items": { + "type": "string" }, - "description": "RemoteServers is a map of server names to their corresponding remote server definitions\nThese are MCP servers accessed via HTTP/HTTPS using the thv proxy command", - "type": "object" + "type": "array" + } + }, + "type": "object" + }, + "core.Workload": { + "properties": { + "created_at": { + "description": "CreatedAt is the timestamp when the workload was created.", + "type": "string" }, - "servers": { + "group": { + "description": "Group is the name of the group this workload belongs to, if any.", + "type": "string" + }, + "labels": { "additionalProperties": { - "$ref": "#/components/schemas/registry.ImageMetadata" + "type": "string" }, - "description": "Servers is a map of server names to their corresponding server definitions", + "description": "Labels are the container labels (excluding standard ToolHive labels)", "type": "object" }, - "version": { - "description": "Version is the schema version of the registry", + "name": { + "description": "Name is the name of the workload.\nIt is used as a unique identifier.", + "type": "string" + }, + "package": { + "description": "Package specifies the Workload Package used to create this Workload.", + "type": "string" + }, + "port": { + "description": "Port is the port on which the workload is exposed.\nThis is embedded in the URL.", + "type": "integer" + }, + "proxy_mode": { + "description": "ProxyMode is the proxy mode that clients should use to connect.\nFor stdio transports, this will be the proxy mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this will be the same as TransportType.", + "type": "string" + }, + "remote": { + "description": "Remote indicates whether this is a remote workload (true) or a container workload (false).", + "type": "boolean" + }, + "started_at": { + "description": "StartedAt is when the container was last started (changes on restart)", + "type": "string" + }, + "status": { + "description": "Status is the current status of the workload.", + "enum": [ + "running", + "stopped", + "error", + "starting", + "stopping", + "unhealthy", + "removing", + "unknown", + "unauthenticated", + "auth_retrying", + "policy_stopped" + ], + "type": "string" + }, + "status_context": { + "description": "StatusContext provides additional context about the workload's status.\nThe exact meaning is determined by the status and the underlying runtime.", + "type": "string" + }, + "tools": { + "description": "ToolsFilter is the filter on tools applied to the workload.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "transport_type": { + "description": "TransportType is the type of transport used for this workload.", + "enum": [ + "stdio", + "sse", + "streamable-http", + "inspector" + ], + "type": "string" + }, + "url": { + "description": "URL is the URL of the workload exposed by the ToolHive proxy.", "type": "string" } }, @@ -103,70 +218,10 @@ "maxDataSize": { "description": "MaxDataSize limits the size of request/response data included in audit logs (in bytes).\n+kubebuilder:default=1024\n+optional", "type": "integer" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_auth.TokenValidatorConfig": { - "description": "DEPRECATED: Middleware configuration.\nOIDCConfig contains OIDC configuration", - "properties": { - "allowPrivateIP": { - "description": "AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses", - "type": "boolean" - }, - "allowedClientIDs": { - "description": "AllowedClientIDs is the set of OAuth client IDs accepted when validating\ntokens by their client_id claim instead of an audience claim.", - "items": { - "type": "string" - }, - "type": "array" - }, - "audience": { - "description": "Audience is the expected audience for the token", - "type": "string" - }, - "authTokenFile": { - "description": "AuthTokenFile is the path to file containing bearer token for authentication", - "type": "string" - }, - "cacertPath": { - "description": "CACertPath is the path to the CA certificate bundle for HTTPS requests", - "type": "string" - }, - "clientID": { - "description": "ClientID is the OIDC client ID", - "type": "string" - }, - "clientSecret": { - "description": "ClientSecret is the optional OIDC client secret for introspection", - "type": "string" - }, - "insecureAllowHTTP": { - "description": "InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing\nWARNING: This is insecure and should NEVER be used in production", - "type": "boolean" - }, - "introspectionURL": { - "description": "IntrospectionURL is the optional introspection endpoint for validating tokens", - "type": "string" - }, - "issuer": { - "description": "Issuer is the OIDC issuer URL (e.g., https://accounts.google.com)", - "type": "string" - }, - "jwksurl": { - "description": "JWKSURL is the URL to fetch the JWKS from", - "type": "string" - }, - "resourceURL": { - "description": "ResourceURL is the explicit resource URL for OAuth discovery (RFC 9728)", - "type": "string" }, - "scopes": { - "description": "Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728)\nIf empty, defaults to [\"openid\"]", - "items": { - "type": "string" - }, - "type": "array" + "maxDelegationDepth": { + "description": "MaxDelegationDepth caps how many nested RFC 8693 \"act\" entries are\nrecorded in an audit event's delegation chain. Deeper chains are\ntruncated (marked with truncated=true). Defaults to 10 when unset.\n+kubebuilder:validation:Minimum=1\n+kubebuilder:default=10\n+optional", + "type": "integer" } }, "type": "object" @@ -234,157 +289,64 @@ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_auth_remote.Config": { - "description": "RemoteAuthConfig contains OAuth configuration for remote MCP servers", + "github_com_stacklok_toolhive_pkg_auth_upstreamswap.Config": { + "description": "UpstreamSwapConfig contains configuration for upstream token swap middleware.\nWhen set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs\nfor upstream IdP tokens before forwarding requests to the MCP server.", "properties": { - "authorize_url": { + "custom_header_name": { + "description": "CustomHeaderName is the header name when HeaderStrategy is \"custom\".", "type": "string" }, - "bearer_token": { - "description": "Bearer token configuration (alternative to OAuth)", + "header_strategy": { + "description": "HeaderStrategy determines how to inject the token: \"replace\" (default) or \"custom\".", "type": "string" }, - "bearer_token_file": { + "provider_name": { + "description": "ProviderName identifies which upstream provider's tokens to retrieve for injection.\nThis is required and must match a configured upstream provider name.", "type": "string" - }, - "cached_cimd_client_id": { - "description": "CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD\nauthentication was used. Kept separate from CachedClientID (which holds\nDCR-issued IDs) so the two can have independent lifecycles — DCR credential\nrotation clears CachedClientID without touching the stable CIMD URL.\nRead by resolveClientCredentials to send the correct client_id on token refresh.", + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig": { + "description": "CIMD controls client_id metadata document support. When enabled, the\nembedded authorization server accepts HTTPS URLs as client_id values\nand resolves them via the CIMD protocol instead of requiring DCR.", + "properties": { + "cache_fallback_ttl": { + "description": "CacheFallbackTTL is the fixed TTL applied to every cached CIMD document.\nCache-Control header parsing is not yet implemented; all entries use this value.\nFormat: Go duration string (e.g. \"5m\", \"10m\", \"1h\").\nDefaults to 5 minutes when Enabled is true and this field is omitted.", + "example": "5m", "type": "string" }, - "cached_client_id": { - "description": "Cached DCR client credentials for persistence across restarts.\nThese are obtained during Dynamic Client Registration and needed to refresh tokens.\nClientID is stored as plain text since it's public information.", - "type": "string" + "cache_max_size": { + "description": "CacheMaxSize is the maximum number of CIMD documents held in the LRU cache.\nDefaults to 256 when Enabled is true and this field is zero.", + "type": "integer" }, - "cached_client_secret_ref": { + "enabled": { + "description": "Enabled activates CIMD client lookup when true.", + "type": "boolean" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_authserver.DCRUpstreamConfig": { + "description": "DCRConfig enables RFC 7591 Dynamic Client Registration against the\nupstream authorization server. When set, the client credentials are\nobtained at runtime rather than being pre-provisioned via ClientID /\nClientSecretFile / ClientSecretEnvVar, and ClientID must be left empty.\nMutually exclusive with ClientID.", + "properties": { + "discovery_url": { + "description": "DiscoveryURL is the exact RFC 8414 / OIDC Discovery document URL to\nfetch at runtime. The resolver issues a single GET against this URL\n(no well-known-path fallback) and reads registration_endpoint,\nauthorization_endpoint, token_endpoint,\ntoken_endpoint_auth_methods_supported, and scopes_supported from the\nresponse. Per RFC 8414 §3.3, the document's \"issuer\" field must\nexactly match the upstream issuer configured on the parent\nrun-config.\n\nUse this field when the upstream publishes discovery metadata at a\npath that differs from the issuer-derived well-known paths — for\nexample a multi-tenant IdP whose metadata lives at\nhttps://idp.example.com/tenants/acme/.well-known/openid-configuration.\n\nMutually exclusive with RegistrationEndpoint.", "type": "string" }, - "cached_refresh_token_ref": { - "description": "Cached OAuth token reference for persistence across restarts.\nThe refresh token is stored securely in the secret manager, and this field\ncontains the reference to retrieve it (e.g., \"OAUTH_REFRESH_TOKEN_workload\").\nThis enables session restoration without requiring a new browser-based login.", + "initial_access_token_env_var": { + "description": "InitialAccessTokenEnvVar is the name of an environment variable\ncontaining the RFC 7591 initial access token. Mutually exclusive with\nInitialAccessTokenFile.", "type": "string" }, - "cached_reg_token_ref": { - "description": "RegistrationAccessToken is used to update/delete the client registration.\nStored as a secret reference since it's sensitive.", + "initial_access_token_file": { + "description": "InitialAccessTokenFile is the path to a file containing the RFC 7591\ninitial access token presented to the registration endpoint. Mutually\nexclusive with InitialAccessTokenEnvVar. Both may be omitted for open\nregistration endpoints.", "type": "string" }, - "cached_secret_expiry": { - "description": "ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server).\nA zero value means the secret does not expire.", + "registration_endpoint": { + "description": "RegistrationEndpoint is the RFC 7591 registration endpoint URL used\ndirectly, bypassing discovery. Because no discovery is performed,\nserver-capability fields (token_endpoint_auth_methods_supported,\nscopes_supported) are unavailable on this code path; the caller is\nexpected to also supply AuthorizationEndpoint, TokenEndpoint, and an\nexplicit Scopes list on the parent OAuth2UpstreamRunConfig. Auth\nmethod falls back to the resolver's default (client_secret_basic).\n\nMutually exclusive with DiscoveryURL.", "type": "string" }, - "cached_token_expiry": { - "type": "string" - }, - "callback_port": { - "type": "integer" - }, - "client_id": { - "type": "string" - }, - "client_secret": { - "type": "string" - }, - "client_secret_file": { - "type": "string" - }, - "issuer": { - "description": "OAuth endpoint configuration (from registry)", - "type": "string" - }, - "oauth_params": { - "additionalProperties": { - "type": "string" - }, - "description": "OAuth parameters for server-specific customization", - "type": "object" - }, - "resource": { - "description": "Resource is the OAuth 2.0 resource indicator (RFC 8707).", - "type": "string" - }, - "scope_param_name": { - "description": "ScopeParamName overrides the query parameter name used to send scopes in the\nauthorization URL. When empty, the standard \"scope\" parameter is used.\nSome providers require a non-standard name (e.g., Slack uses \"user_scope\").", - "type": "string" - }, - "scopes": { - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false - }, - "skip_browser": { - "type": "boolean" - }, - "timeout": { - "example": "5m", - "type": "string" - }, - "token_url": { - "type": "string" - }, - "use_pkce": { - "type": "boolean" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_auth_upstreamswap.Config": { - "description": "UpstreamSwapConfig contains configuration for upstream token swap middleware.\nWhen set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs\nfor upstream IdP tokens before forwarding requests to the MCP server.", - "properties": { - "custom_header_name": { - "description": "CustomHeaderName is the header name when HeaderStrategy is \"custom\".", - "type": "string" - }, - "header_strategy": { - "description": "HeaderStrategy determines how to inject the token: \"replace\" (default) or \"custom\".", - "type": "string" - }, - "provider_name": { - "description": "ProviderName identifies which upstream provider's tokens to retrieve for injection.\nThis is required and must match a configured upstream provider name.", - "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig": { - "description": "CIMD controls client_id metadata document support. When enabled, the\nembedded authorization server accepts HTTPS URLs as client_id values\nand resolves them via the CIMD protocol instead of requiring DCR.", - "properties": { - "cache_fallback_ttl": { - "description": "CacheFallbackTTL is the fixed TTL applied to every cached CIMD document.\nCache-Control header parsing is not yet implemented; all entries use this value.\nFormat: Go duration string (e.g. \"5m\", \"10m\", \"1h\").\nDefaults to 5 minutes when Enabled is true and this field is omitted.", - "example": "5m", - "type": "string" - }, - "cache_max_size": { - "description": "CacheMaxSize is the maximum number of CIMD documents held in the LRU cache.\nDefaults to 256 when Enabled is true and this field is zero.", - "type": "integer" - }, - "enabled": { - "description": "Enabled activates CIMD client lookup when true.", - "type": "boolean" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver.DCRUpstreamConfig": { - "description": "DCRConfig enables RFC 7591 Dynamic Client Registration against the\nupstream authorization server. When set, the client credentials are\nobtained at runtime rather than being pre-provisioned via ClientID /\nClientSecretFile / ClientSecretEnvVar, and ClientID must be left empty.\nMutually exclusive with ClientID.", - "properties": { - "discovery_url": { - "description": "DiscoveryURL is the exact RFC 8414 / OIDC Discovery document URL to\nfetch at runtime. The resolver issues a single GET against this URL\n(no well-known-path fallback) and reads registration_endpoint,\nauthorization_endpoint, token_endpoint,\ntoken_endpoint_auth_methods_supported, and scopes_supported from the\nresponse. Per RFC 8414 §3.3, the document's \"issuer\" field must\nexactly match the upstream issuer configured on the parent\nrun-config.\n\nUse this field when the upstream publishes discovery metadata at a\npath that differs from the issuer-derived well-known paths — for\nexample a multi-tenant IdP whose metadata lives at\nhttps://idp.example.com/tenants/acme/.well-known/openid-configuration.\n\nMutually exclusive with RegistrationEndpoint.", - "type": "string" - }, - "initial_access_token_env_var": { - "description": "InitialAccessTokenEnvVar is the name of an environment variable\ncontaining the RFC 7591 initial access token. Mutually exclusive with\nInitialAccessTokenFile.", - "type": "string" - }, - "initial_access_token_file": { - "description": "InitialAccessTokenFile is the path to a file containing the RFC 7591\ninitial access token presented to the registration endpoint. Mutually\nexclusive with InitialAccessTokenEnvVar. Both may be omitted for open\nregistration endpoints.", - "type": "string" - }, - "registration_endpoint": { - "description": "RegistrationEndpoint is the RFC 7591 registration endpoint URL used\ndirectly, bypassing discovery. Because no discovery is performed,\nserver-capability fields (token_endpoint_auth_methods_supported,\nscopes_supported) are unavailable on this code path; the caller is\nexpected to also supply AuthorizationEndpoint, TokenEndpoint, and an\nexplicit Scopes list on the parent OAuth2UpstreamRunConfig. Auth\nmethod falls back to the resolver's default (client_secret_basic).\n\nMutually exclusive with DiscoveryURL.", - "type": "string" - }, - "software_id": { - "description": "SoftwareID is the RFC 7591 \"software_id\" registration metadata value,\nidentifying the client software independent of any particular\nregistration instance.", + "software_id": { + "description": "SoftwareID is the RFC 7591 \"software_id\" registration metadata value,\nidentifying the client software independent of any particular\nregistration instance.", "type": "string" }, "software_statement": { @@ -423,7 +385,7 @@ "type": "object" }, "allow_private_ips": { - "description": "AllowPrivateIPs permits the upstream provider's HTTP client to connect to\nprivate IP ranges (RFC-1918, link-local). Use only when the upstream is\nhosted inside the same cluster and has no public endpoint. HTTP-scheme\nrestrictions are unchanged — HTTPS is still required for non-localhost hosts.\nDefaults to false.", + "description": "AllowPrivateIPs permits the upstream provider's HTTP client to connect to\nprivate IP ranges (RFC-1918, link-local). When DCRConfig is set, this\nalso gates the DCR discovery and registration calls made on this\nupstream's behalf (see pkg/authserver/runner/dcr_adapter.go), so a\nsingle flag covers the whole upstream rather than needing a separate\nDCR-specific setting. Use only when the upstream is hosted inside the\nsame cluster and has no public endpoint. HTTP-scheme restrictions are\nunchanged — HTTPS is still required for non-localhost hosts. Defaults\nto false.", "type": "boolean" }, "authorization_endpoint": { @@ -559,6 +521,10 @@ "cimd": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig" }, + "delegation_token_lifespan": { + "description": "DelegationTokenLifespan is the maximum lifetime for delegated tokens issued\nvia RFC 8693 token exchange. Specified as a Go duration string (e.g., \"15m\").\nIf empty, defaults to 15 minutes.", + "type": "string" + }, "disable_upstream_token_injection": { "description": "DisableUpstreamTokenInjection prevents the upstream swap middleware from being added.\nWhen true, the embedded auth server handles OAuth flows for clients, but instead of\ninjecting upstream IdP tokens the proxy strips the client's credential headers\n(Authorization, Cookie, Proxy-Authorization) after the JWT is validated — the\nbackend receives an unauthenticated request. Incompatible with token exchange\nand AWS STS, which would re-add credentials after the strip.", "type": "boolean" @@ -595,11 +561,19 @@ "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.SigningKeyRunConfig" }, "storage": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RunConfig" + "$ref": "#/components/schemas/storage.RunConfig" }, "token_lifespans": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.TokenLifespanRunConfig" }, + "trusted_issuers": { + "description": "TrustedIssuers lists external OIDC issuers whose tokens are accepted as\nsubject tokens during RFC 8693 token exchange. Empty (the default) means\nonly self-issued subject tokens are accepted.\n\nPrerequisite: the token-exchange grant requires a confidential client,\nand no supported deployment path provisions one today (DCR and CIMD\nclients are both public-only, and there is no client-seeding field on\nthis RunConfig), so this grant is not yet usable end to end for\nself-issued or external subject tokens alike. Tracked in\nhttps://github.com/stacklok/toolhive/issues/6082.\n\nSee tokenexchange.TrustedIssuer for the per-issuer field reference, and\ndocs/arch/17-token-exchange-delegation.md for the trust model, consent\nsignals, and operator-facing constraints (audience/scope bounding,\nsubject namespace qualification, required client binding) that aren't\nvisible from the config shape alone.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_server_tokenexchange.TrustedIssuer" + }, + "type": "array", + "uniqueItems": false + }, "upstreams": { "description": "Upstreams configures connections to upstream Identity Providers.\nAt least one upstream is required - the server delegates authentication to these providers.\nMultiple upstreams are supported for sequential authorization chains.", "items": { @@ -757,111 +731,47 @@ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_authserver_storage.ACLUserRunConfig": { - "description": "ACLUserConfig contains ACL user authentication configuration.", - "properties": { - "password_env_var": { - "description": "PasswordEnvVar is the environment variable containing the Redis password.", - "type": "string" - }, - "username_env_var": { - "description": "UsernameEnvVar is the environment variable containing the Redis username.", - "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.RedisRunConfig": { - "description": "RedisConfig is the Redis-specific configuration when Type is \"redis\".", + "github_com_stacklok_toolhive_pkg_authserver_server_tokenexchange.TrustedIssuer": { "properties": { - "acl_user_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.ACLUserRunConfig" - }, - "addr": { - "description": "Addr is the Redis server address (host:port). Required for standalone and cluster modes.\nMutually exclusive with SentinelConfig.", - "type": "string" - }, - "auth_type": { - "description": "AuthType must be \"aclUser\" - only ACL user authentication is supported.", + "actor_claim": { + "description": "ActorClaim names the claim identifying the client that requested the\nsubject token from THIS EXTERNAL ISSUER (used by AllowedActors below).\nValues are in the external issuer's namespace, NOT ToolHive client\nIDs. Defaults to \"azp\"; use \"appid\" for Microsoft Entra v1, \"cid\" for\nOkta. The special value \"client_id\" reads ValidatedClaims.ClientID\ninstead of Extra (assignClaim routes it to that field) — it is still\nthe external token's client_id claim, not a ToolHive one.", "type": "string" }, - "cluster_mode": { - "description": "ClusterMode enables the Redis Cluster protocol. Requires Addr to be set.", + "allow_private_ips": { + "description": "AllowPrivateIPs permits OIDC discovery and JWKS fetches for THIS\nissuer to resolve to a private or loopback address. Use only when the\nissuer is hosted inside the same cluster and has no public endpoint.", "type": "boolean" }, - "dial_timeout": { - "description": "DialTimeout is the timeout for establishing connections (e.g., \"5s\").", - "type": "string" - }, - "key_prefix": { - "description": "KeyPrefix for multi-tenancy, typically \"thv:auth:{ns}:{name}:\".", - "type": "string" - }, - "read_timeout": { - "description": "ReadTimeout is the timeout for read operations (e.g., \"3s\").", - "type": "string" - }, - "sentinel_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.SentinelRunConfig" - }, - "sentinel_tls": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig" + "allowed_actors": { + "description": "AllowedActors is the allowlist of ActorClaim values authorized to\nexchange a subject token from this issuer when it carries no\n\"may_act\" claim; empty means only may_act-bearing tokens are\naccepted. By itself names no ToolHive client — see\nAllowedDelegateClients and docs/arch/17-token-exchange-delegation.md\n(\"Accepted limitations\" #1).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "tls": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig" + "allowed_delegate_clients": { + "description": "AllowedDelegateClients restricts which ToolHive client IDs may\nexchange a subject token from this issuer, for BOTH consent paths.\nRequired (validateTrustedIssuer rejects empty/absent); \"*\" permits\nany confidential client holding the grant. See\ndocs/arch/17-token-exchange-delegation.md (\"Accepted limitations\" #1).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "write_timeout": { - "description": "WriteTimeout is the timeout for write operations (e.g., \"3s\").", - "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig": { - "description": "SentinelTLS configures TLS for Sentinel connections. Only applies when SentinelConfig is set.", - "properties": { - "ca_cert_file": { - "description": "CACertFile is the path to a PEM-encoded CA certificate file.", + "expected_audience": { + "description": "ExpectedAudience is the expected \"aud\" claim value that must appear\nin the token's audience list (a resource/API identifier, not a\nclient ID — required, but not enforced; see looksLikeResourceIdentifier).\nSee docs/arch/17-token-exchange-delegation.md (\"ID/access-token\ndiscrimination\") for why and its limits.", "type": "string" }, - "insecure_skip_verify": { - "description": "InsecureSkipVerify skips certificate verification.", + "insecure_allow_http": { + "description": "InsecureAllowHTTP permits plain-HTTP OIDC discovery and JWKS fetches\nfor THIS issuer only. Development and testing only — never set in\nproduction. Does not relax the private-IP guard; see AllowPrivateIPs.\nDeliberately per-issuer: this server's own InsecureAllowHTTP must not\nsilently permit plaintext discovery for every trusted external issuer\ntoo — a network attacker who can intercept that traffic could\nsubstitute a JWKS and forge subject tokens for that issuer's\nnamespace.", "type": "boolean" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.RunConfig": { - "description": "Storage configures the storage backend for the auth server.\nIf nil, defaults to in-memory storage.", - "properties": { - "redis_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisRunConfig" }, - "type": { - "description": "Type specifies the storage backend type. Defaults to \"memory\".", + "issuer_url": { + "description": "IssuerURL is the expected \"iss\" claim value (exact match).", "type": "string" - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_authserver_storage.SentinelRunConfig": { - "description": "SentinelConfig contains Sentinel-specific configuration.\nMutually exclusive with Addr.", - "properties": { - "db": { - "description": "DB is the Redis database number (default: 0).", - "type": "integer" }, - "master_name": { - "description": "MasterName is the name of the Redis Sentinel master.", + "jwks_url": { + "description": "JWKSURL is the URL to fetch the issuer's JSON Web Key Set from.\nIf empty, it is resolved via OIDC discovery at {IssuerURL}/.well-known/openid-configuration.", "type": "string" - }, - "sentinel_addrs": { - "description": "SentinelAddrs is the list of Sentinel addresses (host:port).", - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false } }, "type": "object" @@ -975,230 +885,326 @@ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus": { - "description": "Current status of the workload", - "enum": [ - "running", - "stopped", - "error", - "starting", - "stopping", - "unhealthy", - "removing", - "unknown", - "unauthenticated", - "auth_retrying", - "policy_stopped", - "running", - "stopped", - "error", - "starting", - "stopping", - "unhealthy", - "removing", - "unknown", - "unauthenticated", - "auth_retrying", - "policy_stopped", - "running", - "stopped", - "error", - "starting", - "stopping", - "unhealthy", - "removing", - "unknown", - "unauthenticated", - "auth_retrying", - "policy_stopped" - ], - "type": "string", - "x-enum-varnames": [ - "WorkloadStatusRunning", - "WorkloadStatusStopped", - "WorkloadStatusError", - "WorkloadStatusStarting", - "WorkloadStatusStopping", - "WorkloadStatusUnhealthy", - "WorkloadStatusRemoving", - "WorkloadStatusUnknown", - "WorkloadStatusUnauthenticated", - "WorkloadStatusAuthRetrying", - "WorkloadStatusPolicyStopped" - ] - }, - "github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig": { - "description": "RuntimeConfig allows overriding the default runtime configuration\nfor this specific workload (base images and packages)", + "github_com_stacklok_toolhive_pkg_groups.Group": { "properties": { - "additional_packages": { - "description": "AdditionalPackages lists extra packages to install in the builder and\nruntime stages.\nExamples for Alpine: [\"git\", \"make\", \"gcc\"]\nExamples for Debian: [\"git\", \"build-essential\"]", + "name": { + "type": "string" + }, + "plugins": { "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "builder_image": { - "description": "BuilderImage is the full image reference for the builder stage.\nAn empty string signals \"use the default for this transport type\" during config merging.\nExamples: \"golang:1.26-alpine\", \"node:24-alpine\", \"python:3.14-slim\"", - "type": "string" + "registered_clients": { + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "runtime_env": { - "additionalProperties": { + "skills": { + "items": { "type": "string" }, - "description": "RuntimeEnv contains environment variables to inject into the Dockerfile's\nfinal runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv),\nwhich only affects the builder stage, these variables are baked into the\nshipped image and are present in the running container's process\nenvironment at startup. Use this for values a packaged MCP server reads at\nprocess start (e.g. feature flags, cache backend selection), not for\nbuild-time package manager configuration.\nKeys must be uppercase with underscores, values are validated for safety.", - "type": "object" + "type": "array", + "uniqueItems": false } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_core.Workload": { + "github_com_stacklok_toolhive_pkg_plugins.BuildResult": { "properties": { - "created_at": { - "description": "CreatedAt is the timestamp when the workload was created.", - "type": "string" - }, - "group": { - "description": "Group is the name of the group this workload belongs to, if any.", + "reference": { + "description": "Reference is the OCI reference of the built skill artifact.", + "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.ComponentInventory": { + "additionalProperties": { + "type": "integer" + }, + "description": "Components is the inventory of component types declared by the plugin\n(e.g. {\"commands\": 3, \"skills\": 2}). Extracted from the OCI artifact.", + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.ComponentType": { + "enum": [ + "commands", + "agents", + "skills", + "hooks", + "mcpServers", + "lspServers" + ], + "type": "string", + "x-enum-varnames": [ + "ComponentCommands", + "ComponentAgents", + "ComponentSkills", + "ComponentHooks", + "ComponentMCP", + "ComponentLSP" + ] + }, + "github_com_stacklok_toolhive_pkg_plugins.Dependency": { + "properties": { + "digest": { + "description": "Digest is the OCI digest for upgrade detection.", "type": "string" - }, - "labels": { - "additionalProperties": { - "type": "string" - }, - "description": "Labels are the container labels (excluding standard ToolHive labels)", - "type": "object" }, "name": { - "description": "Name is the name of the workload.\nIt is used as a unique identifier.", + "description": "Name is the dependency name.", "type": "string" }, - "package": { - "description": "Package specifies the Workload Package used to create this Workload.", + "reference": { + "description": "Reference is the OCI reference for the dependency.", "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.InstallStatus": { + "description": "Status is the current installation status.", + "enum": [ + "installed", + "pending", + "failed" + ], + "type": "string", + "x-enum-varnames": [ + "InstallStatusInstalled", + "InstallStatusPending", + "InstallStatusFailed" + ] + }, + "github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin": { + "description": "InstalledPlugin contains the full installation record.", + "properties": { + "clients": { + "description": "Clients is the list of client identifiers the plugin is installed for.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "port": { - "description": "Port is the port on which the workload is exposed.\nThis is embedded in the URL.", - "type": "integer" + "components": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ComponentInventory" }, - "proxy_mode": { - "description": "ProxyMode is the proxy mode that clients should use to connect.\nFor stdio transports, this will be the proxy mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this will be the same as TransportType.", + "dependencies": { + "description": "Dependencies is the list of external plugin dependencies.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Dependency" + }, + "type": "array", + "uniqueItems": false + }, + "digest": { + "description": "Digest is the OCI digest (sha256:...) for upgrade detection.", "type": "string" }, - "remote": { - "description": "Remote indicates whether this is a remote workload (true) or a container workload (false).", - "type": "boolean" + "installed_at": { + "description": "InstalledAt is the timestamp when the plugin was installed.", + "type": "string" }, - "started_at": { - "description": "StartedAt is when the container was last started (changes on restart)", + "metadata": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginMetadata" + }, + "project_root": { + "description": "ProjectRoot is the project root path for project-scoped plugins. Empty for user-scoped.", + "type": "string" + }, + "reference": { + "description": "Reference is the full OCI reference (e.g. ghcr.io/org/plugin:v1).", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Scope" + }, + "signature": { + "description": "Signature is the optional signing signature for the plugin artifact.", "type": "string" }, "status": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstallStatus" }, - "status_context": { - "description": "StatusContext provides additional context about the workload's status.\nThe exact meaning is determined by the status and the underlying runtime.", + "tag": { + "description": "Tag is the OCI tag (e.g. v1.0.0).", + "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.LocalBuild": { + "properties": { + "description": { + "description": "Description is the skill description extracted from the artifact metadata, if available.", "type": "string" }, - "tools": { - "description": "ToolsFilter is the filter on tools applied to the workload.", - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false + "digest": { + "description": "Digest is the OCI digest of the artifact (sha256:...).", + "type": "string" }, - "transport_type": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.TransportType" + "name": { + "description": "Name is the skill name extracted from the artifact metadata, if available.", + "type": "string" }, - "url": { - "description": "URL is the URL of the workload exposed by the ToolHive proxy.", + "tag": { + "description": "Tag is the OCI tag or name used to reference the artifact.", + "type": "string" + }, + "version": { + "description": "Version is the skill version extracted from the artifact metadata, if available.", "type": "string" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_groups.Group": { + "github_com_stacklok_toolhive_pkg_plugins.PluginContent": { "properties": { - "name": { + "description": { + "description": "Description is the plugin description from the OCI config labels.", "type": "string" }, - "plugins": { + "files": { + "description": "Files is the list of all files in the artifact with their sizes.", "items": { - "type": "string" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginFileEntry" }, "type": "array", "uniqueItems": false }, - "registered_clients": { - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false + "license": { + "description": "License is the SPDX license identifier from the OCI config labels.", + "type": "string" }, - "skills": { - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": false + "manifest": { + "description": "Manifest is the raw .claude-plugin/plugin.json body.", + "type": "string" + }, + "name": { + "description": "Name is the plugin name from the OCI config labels.", + "type": "string" + }, + "version": { + "description": "Version is the plugin version from the OCI config labels.", + "type": "string" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_ignore.Config": { - "description": "IgnoreConfig contains configuration for ignore processing", + "github_com_stacklok_toolhive_pkg_plugins.PluginFileEntry": { "properties": { - "loadGlobal": { - "description": "Whether to load global ignore patterns", - "type": "boolean" + "path": { + "description": "Path is the file path within the artifact.", + "type": "string" }, - "printOverlays": { - "description": "Whether to print resolved overlay paths for debugging", - "type": "boolean" + "size": { + "description": "Size is the uncompressed file size in bytes.", + "type": "integer" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_oauthproto_tokenexchange.Config": { - "description": "TokenExchangeConfig contains token exchange configuration for external authentication", + "github_com_stacklok_toolhive_pkg_plugins.PluginInfo": { "properties": { - "audience": { - "description": "Audience is the target audience for the exchanged token", - "type": "string" + "installed_plugin": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin" }, - "client_id": { - "description": "ClientID is the OAuth 2.0 client identifier", - "type": "string" + "metadata": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginMetadata" }, - "client_secret": { - "description": "ClientSecret is the OAuth 2.0 client secret", - "type": "string" + "project_scope_degraded_clients": { + "description": "ProjectScopeDegradedClients lists the client types for which a\nproject-scoped install degraded (the adapter could only materialize at\nuser scope — e.g. Codex always writes to the user-scoped config.toml).\nPopulated by Info; empty for user-scoped installs. Recomputed at read\ntime from the stored scope + each adapter's capability, mirroring the\nUnmaterializedComponents pattern (no persistence needed — the degradation\nis deterministic from scope + client type).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "external_token_header_name": { - "description": "ExternalTokenHeaderName is the name of the custom header to use when HeaderStrategy is \"custom\"", + "unmaterialized_components": { + "additionalProperties": { + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ComponentType" + }, + "type": "array" + }, + "description": "UnmaterializedComponents lists, per client type, the component types the\nplugin declares that the installed client adapter does NOT load. Populated\nby Info by diffing InstalledPlugin.Components against each installed\nclient adapter's SupportedComponents.", + "type": "object" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.PluginMetadata": { + "description": "Metadata contains the plugin's metadata.", + "properties": { + "author": { + "description": "Author is the plugin author or maintainer.", "type": "string" }, - "header_strategy": { - "description": "HeaderStrategy determines how to inject the token\nValid values: HeaderStrategyReplace (default), HeaderStrategyCustom", + "description": { + "description": "Description is a human-readable description of the plugin.", "type": "string" }, - "scopes": { - "description": "Scopes is the list of scopes to request for the exchanged token", + "keywords": { + "description": "Keywords is a list of keywords for categorization/search.", "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "subject_token_type": { - "description": "SubjectTokenType specifies the type of the subject token being exchanged.\nCommon values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT.\nIf empty, defaults to oauthproto.TokenTypeAccessToken.", + "license": { + "description": "License is the SPDX license identifier for the plugin.", "type": "string" }, - "token_url": { - "description": "TokenURL is the OAuth 2.0 token endpoint URL", + "name": { + "description": "Name is the unique name of the plugin (kebab-case).", "type": "string" + }, + "version": { + "description": "Version is the semantic version of the plugin.", + "type": "string" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_plugins.Scope": { + "description": "Scope for the installation", + "enum": [ + "user", + "project" + ], + "type": "string", + "x-enum-varnames": [ + "ScopeUser", + "ScopeProject" + ] + }, + "github_com_stacklok_toolhive_pkg_plugins.ValidationResult": { + "properties": { + "errors": { + "description": "Errors is a list of validation errors, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "valid": { + "description": "Valid indicates whether the skill definition is valid.", + "type": "boolean" + }, + "warnings": { + "description": "Warnings is a list of non-blocking validation warnings, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false } }, "type": "object" @@ -1278,7 +1284,7 @@ "additional_middleware_configs": { "description": "AdditionalMiddlewareConfigs carries pre-built middleware configs injected by\nexternal-auth handlers (reached via *[]RunConfigBuilderOption) rather than\nderived from typed RunConfig fields. PopulateMiddlewareConfigs splices these\ninto the chain in the backend-egress group — after auth and before recovery —\ninstead of discarding them. Upstream carries these configs verbatim and never\ninspects their parameters; the middleware type identity (e.g. an enterprise\nauth type) is supplied by the caller via types.MiddlewareConfig.Type.\n\nEach entry's Type is expected to be a NEW egress middleware type (e.g. OBO),\nnot one already produced from a typed RunConfig field (auth, authz, audit,\ntokenExchange, awssts, …). Dispatch in the proxyrunner is purely by Type\nstring, so an injected Type that shadows a typed-field type would add a\nsecond instance of that middleware to the chain; the seam does not validate\nagainst this.", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig" + "$ref": "#/components/schemas/types.MiddlewareConfig" }, "type": "array", "uniqueItems": false @@ -1369,7 +1375,7 @@ "type": "string" }, "ignore_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_ignore.Config" + "$ref": "#/components/schemas/ignore.Config" }, "image": { "description": "Image is the Docker image to run", @@ -1394,7 +1400,7 @@ "middleware_configs": { "description": "MiddlewareConfigs contains the list of middleware to apply to the transport\nand the configuration for each middleware.", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig" + "$ref": "#/components/schemas/types.MiddlewareConfig" }, "type": "array", "uniqueItems": false @@ -1412,7 +1418,7 @@ "type": "string" }, "oidc_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_auth.TokenValidatorConfig" + "$ref": "#/components/schemas/auth.TokenValidatorConfig" }, "permission_profile_name_or_path": { "description": "PermissionProfileNameOrPath is the name or path of the permission profile", @@ -1423,7 +1429,12 @@ "type": "integer" }, "proxy_mode": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.ProxyMode" + "description": "ProxyMode is the effective HTTP protocol the proxy uses.\nFor stdio transports, this is the configured mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this matches the transport type.\nNote: \"sse\" is deprecated; use \"streamable-http\" instead.", + "enum": [ + "sse", + "streamable-http" + ], + "type": "string" }, "publish": { "description": "Publish lists ports to publish to the host in format \"hostPort:containerPort\"", @@ -1453,14 +1464,14 @@ "type": "string" }, "remote_auth_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_auth_remote.Config" + "$ref": "#/components/schemas/remote.Config" }, "remote_url": { "description": "RemoteURL is the URL of the remote MCP server (if running remotely)", "type": "string" }, "runtime_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig" + "$ref": "#/components/schemas/templates.RuntimeConfig" }, "scaling_config": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_runner.ScalingConfig" @@ -1486,6 +1497,10 @@ "description": "Stateless indicates the server only supports POST (no SSE/GET).\nWhen true, the proxy returns 405 for incoming GET requests and uses a\nPOST-based health check instead of the default GET probe.\nApplies to both remote URLs and local container workloads.", "type": "boolean" }, + "strict_protocol_validation": { + "description": "StrictProtocolValidation enables strict MCP-Protocol-Version validation\non the streamable HTTP proxy: a request whose header names an unknown\nMCP revision is rejected with HTTP 400. Default false accepts any\nversion string (an absent header is always accepted in either mode).", + "type": "boolean" + }, "target_host": { "description": "TargetHost is the host to forward traffic to (only applicable to SSE transport)", "type": "string" @@ -1495,14 +1510,14 @@ "type": "integer" }, "telemetry_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_telemetry.Config" + "$ref": "#/components/schemas/telemetry.Config" }, "thv_ca_bundle": { "description": "DEPRECATED: No longer appears to be used.\nThvCABundle is the path to the CA certificate bundle for ToolHive HTTP operations", "type": "string" }, "token_exchange_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_oauthproto_tokenexchange.Config" + "$ref": "#/components/schemas/tokenexchange.Config" }, "tools_filter": { "description": "DEPRECATED: Middleware configuration.\nToolsFilter is the list of tools to filter", @@ -1520,7 +1535,14 @@ "type": "object" }, "transport": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.TransportType" + "description": "Transport is the transport mode (stdio, sse, or streamable-http)", + "enum": [ + "stdio", + "sse", + "streamable-http", + "inspector" + ], + "type": "string" }, "trust_proxy_headers": { "description": "TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies", @@ -1592,18 +1614,6 @@ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_secrets.SecretParameter": { - "description": "Bearer token for authentication (alternative to OAuth)", - "properties": { - "name": { - "type": "string" - }, - "target": { - "type": "string" - } - }, - "type": "object" - }, "github_com_stacklok_toolhive_pkg_skills.BuildResult": { "properties": { "reference": { @@ -1630,6 +1640,30 @@ }, "type": "object" }, + "github_com_stacklok_toolhive_pkg_skills.FailureReason": { + "description": "Reason is a typed failure reason when Status is UpgradeStatusFailed.", + "enum": [ + "registry-unreachable", + "digest-missing", + "validation-rejected", + "lock-write-failed", + "signature-invalid", + "signer-mismatch", + "unsigned-rejected", + "unknown" + ], + "type": "string", + "x-enum-varnames": [ + "FailureReasonRegistryUnreachable", + "FailureReasonDigestMissing", + "FailureReasonValidationRejected", + "FailureReasonLockWriteFailed", + "FailureReasonSignatureInvalid", + "FailureReasonSignerMismatch", + "FailureReasonUnsignedRejected", + "FailureReasonUnknown" + ] + }, "github_com_stacklok_toolhive_pkg_skills.InstallStatus": { "description": "Status is the current installation status.", "enum": [ @@ -1671,6 +1705,10 @@ "description": "InstalledAt is the timestamp when the skill was installed.", "type": "string" }, + "managed": { + "description": "Managed indicates this install is tracked in the project's\ntoolhive.lock.yaml. Only ever true for project-scoped installs. No\nomitempty: false is an observable state (unmanaged), not an absence.", + "type": "boolean" + }, "metadata": { "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillMetadata" }, @@ -1819,149 +1857,183 @@ }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_skills.ValidationResult": { + "github_com_stacklok_toolhive_pkg_skills.SyncFailure": { "properties": { - "errors": { - "description": "Errors is a list of validation errors, if any.", + "error": { + "description": "Error is a human-readable description of the failure.", + "type": "string" + }, + "name": { + "description": "Name is the skill name that failed.", + "type": "string" + }, + "reason": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.FailureReason" + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_skills.SyncResult": { + "properties": { + "already_current": { + "description": "AlreadyCurrent lists skills that already matched the lock file.", "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "valid": { - "description": "Valid indicates whether the skill definition is valid.", - "type": "boolean" - }, - "warnings": { - "description": "Warnings is a list of non-blocking validation warnings, if any.", + "drifted": { + "description": "Drifted lists skills whose on-disk contentDigest differed from the lock\nfile. Normally these are reinstalled to match it; when Check is set,\nnothing is written and this field reports the drift only.", "items": { "type": "string" }, "type": "array", "uniqueItems": false - } - }, - "type": "object" - }, - "github_com_stacklok_toolhive_pkg_telemetry.Config": { - "description": "DEPRECATED: Middleware configuration.\nTelemetryConfig contains the OpenTelemetry configuration", - "properties": { - "caCertPath": { - "description": "CACertPath is the file path to a CA certificate bundle for the OTLP endpoint.\nWhen set, the OTLP exporters use this CA to verify the collector's TLS certificate\ninstead of relying solely on the system CA pool.\n+optional", - "type": "string" }, - "customAttributes": { - "additionalProperties": { + "failed": { + "description": "Failed lists skills that could not be synced, with the reason for each.\nDrift alone is never reported here — see Drifted.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SyncFailure" + }, + "type": "array", + "uniqueItems": false + }, + "installed": { + "description": "Installed lists skills that were installed or reinstalled to match the lock file.", + "items": { "type": "string" }, - "description": "CustomAttributes contains custom resource attributes to be added to all telemetry signals.\nThese are parsed from CLI flags (--otel-custom-attributes) or environment variables\n(OTEL_RESOURCE_ATTRIBUTES) as key=value pairs.\n+optional", - "type": "object" + "type": "array", + "uniqueItems": false }, - "enablePrometheusMetricsPath": { - "description": "EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint.\nThe metrics are served on the main transport port at /metrics.\nThis is separate from OTLP metrics which are sent to the Endpoint.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "missing": { + "description": "Missing lists lock entries with no corresponding install record at all\n— the fresh-clone state. Normally these are installed at their pinned\nreference; when Check is set, nothing is written and this field\nreports the gap only.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "endpoint": { - "description": "Endpoint is the OTLP endpoint URL\n+optional", - "type": "string" + "never_managed": { + "description": "NeverManaged lists project-scoped skills never recorded as lock-managed.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false }, - "environmentVariables": { - "description": "EnvironmentVariables is a list of environment variable names that should be\nincluded in telemetry spans as attributes. Only variables in this list will\nbe read from the host machine and included in spans for observability.\nExample: [\"NODE_ENV\", \"DEPLOYMENT_ENV\", \"SERVICE_VERSION\"]\n+optional", + "pruned": { + "description": "Pruned lists removed-from-lock skills that were uninstalled because Prune was set.", "items": { "type": "string" }, "type": "array", "uniqueItems": false }, - "headers": { - "additionalProperties": { + "removed_from_lock": { + "description": "RemovedFromLock lists previously managed skills absent from the lock file.", + "items": { "type": "string" }, - "description": "Headers contains authentication headers for the OTLP endpoint.\n+optional", - "type": "object" + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "github_com_stacklok_toolhive_pkg_skills.UpgradeOutcome": { + "properties": { + "error": { + "description": "Error is a human-readable description of the failure, set only when Status is UpgradeStatusFailed.", + "type": "string" }, - "insecure": { - "description": "Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "name": { + "description": "Name is the skill name.", + "type": "string" }, - "metricsEnabled": { - "description": "MetricsEnabled controls whether OTLP metrics are enabled.\nWhen false, OTLP metrics are not sent even if an endpoint is configured.\nThis is independent of EnablePrometheusMetricsPath.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "new_digest": { + "description": "NewDigest is the digest the source currently resolves to. Equal to\nOldDigest when Status is UpgradeStatusUpToDate.", + "type": "string" }, - "samplingRate": { - "description": "SamplingRate is the trace sampling rate (0.0-1.0) as a string.\nOnly used when TracingEnabled is true.\nExample: \"0.05\" for 5% sampling.\n+kubebuilder:default=\"0.05\"\n+optional", + "new_resolved_reference": { + "description": "NewResolvedReference is the new resolvedReference when it changed.", "type": "string" }, - "serviceName": { - "description": "ServiceName is the service name for telemetry.\nWhen omitted, defaults to the server name (e.g., VirtualMCPServer name).\n+optional", + "new_signer_identity": { + "description": "NewSignerIdentity is the candidate's signer identity when it differs\nfrom the recorded one (empty when the candidate is unsigned).", "type": "string" }, - "serviceVersion": { - "description": "ServiceVersion is the service version for telemetry.\nWhen omitted, defaults to the ToolHive version.\n+optional", + "old_digest": { + "description": "OldDigest is the digest pinned in the lock file before this operation.", "type": "string" }, - "tracingEnabled": { - "description": "TracingEnabled controls whether distributed tracing is enabled.\nWhen false, no tracer provider is created even if an endpoint is configured.\n+kubebuilder:default=false\n+optional", - "type": "boolean" + "reason": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.FailureReason" }, - "useLegacyAttributes": { - "description": "UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names\nare emitted alongside the new standard attribute names. When true, spans include both\nold and new attribute names for backward compatibility with existing dashboards.\nCurrently defaults to true; this will change to false in a future release.\n+kubebuilder:default=true\n+optional", - "type": "boolean" + "status": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeStatus" } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig": { + "github_com_stacklok_toolhive_pkg_skills.UpgradeResult": { "properties": { - "parameters": { - "description": "Parameters is a JSON object containing the middleware parameters.\nIt is stored as a raw message to allow flexible parameter types.", - "type": "object" - }, - "type": { - "description": "Type is a string representing the middleware type.", - "type": "string" + "outcomes": { + "description": "Outcomes contains one entry per skill considered for upgrade.", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeOutcome" + }, + "type": "array", + "uniqueItems": false } }, "type": "object" }, - "github_com_stacklok_toolhive_pkg_transport_types.ProxyMode": { - "description": "ProxyMode is the effective HTTP protocol the proxy uses.\nFor stdio transports, this is the configured mode (sse or streamable-http).\nFor direct transports (sse/streamable-http), this matches the transport type.\nNote: \"sse\" is deprecated; use \"streamable-http\" instead.", + "github_com_stacklok_toolhive_pkg_skills.UpgradeStatus": { + "description": "Status is the outcome of the upgrade attempt.", "enum": [ - "sse", - "streamable-http", - "sse", - "streamable-http" + "upgraded", + "up-to-date", + "not-upgradable", + "ref-change-blocked", + "signer-change-blocked", + "failed" ], "type": "string", "x-enum-varnames": [ - "ProxyModeSSE", - "ProxyModeStreamableHTTP" + "UpgradeStatusUpgraded", + "UpgradeStatusUpToDate", + "UpgradeStatusNotUpgradable", + "UpgradeStatusRefChangeBlocked", + "UpgradeStatusSignerChangeBlocked", + "UpgradeStatusFailed" ] }, - "github_com_stacklok_toolhive_pkg_transport_types.TransportType": { - "description": "Transport is the transport mode (stdio, sse, or streamable-http)", - "enum": [ - "stdio", - "sse", - "streamable-http", - "inspector", - "stdio", - "sse", - "streamable-http", - "inspector", - "stdio", - "sse", - "streamable-http", - "inspector" - ], - "type": "string", - "x-enum-varnames": [ - "TransportTypeStdio", - "TransportTypeSSE", - "TransportTypeStreamableHTTP", - "TransportTypeInspector" - ] + "github_com_stacklok_toolhive_pkg_skills.ValidationResult": { + "properties": { + "errors": { + "description": "Errors is a list of validation errors, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "valid": { + "description": "Valid indicates whether the skill definition is valid.", + "type": "boolean" + }, + "warnings": { + "description": "Warnings is a list of non-blocking validation warnings, if any.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" }, "github_com_stacklok_toolhive_pkg_webhook.Config": { "properties": { @@ -2148,6 +2220,20 @@ "StatusUnknown" ] }, + "ignore.Config": { + "description": "IgnoreConfig contains configuration for ignore processing", + "properties": { + "loadGlobal": { + "description": "Whether to load global ignore patterns", + "type": "boolean" + }, + "printOverlays": { + "description": "Whether to print resolved overlay paths for debugging", + "type": "boolean" + } + }, + "type": "object" + }, "model.Argument": { "properties": { "choices": { @@ -2626,6 +2712,20 @@ }, "type": "object" }, + "pkg_api_v1.buildPluginRequest": { + "description": "Request to build a plugin from a local directory", + "properties": { + "path": { + "description": "Path to the plugin definition directory", + "type": "string" + }, + "tag": { + "description": "OCI tag for the built artifact", + "type": "string" + } + }, + "type": "object" + }, "pkg_api_v1.buildSkillRequest": { "description": "Request to build a skill from a local directory", "properties": { @@ -2818,12 +2918,12 @@ "type": "string" }, "runtime_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig" + "$ref": "#/components/schemas/templates.RuntimeConfig" }, "secrets": { "description": "Secret parameters to inject", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "type": "array", "uniqueItems": false @@ -2939,7 +3039,7 @@ "type": "string" }, "registry": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive-core_registry_types.Registry" + "$ref": "#/components/schemas/registry.Registry" }, "server_count": { "description": "Number of servers in the registry", @@ -3025,9 +3125,59 @@ }, "type": "object" }, + "pkg_api_v1.installPluginRequest": { + "description": "Request to install a plugin", + "properties": { + "clients": { + "description": "Clients lists target client identifiers (e.g., \"claude-code\"),\nor [\"all\"] to target every plugin-supporting client.\nOmitting this field installs to all available clients.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "force": { + "description": "Force allows overwriting unmanaged plugin directories", + "type": "boolean" + }, + "group": { + "description": "Group is the group name to add the plugin to after installation", + "type": "string" + }, + "name": { + "description": "Name or OCI reference of the plugin to install", + "type": "string" + }, + "project_root": { + "description": "ProjectRoot is the project root path for project-scoped installs", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Scope" + }, + "version": { + "description": "Version to install (empty means latest)", + "type": "string" + } + }, + "type": "object" + }, + "pkg_api_v1.installPluginResponse": { + "description": "Response after successfully installing a plugin", + "properties": { + "plugin": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin" + } + }, + "type": "object" + }, "pkg_api_v1.installSkillRequest": { "description": "Request to install a skill", "properties": { + "allow_unsigned": { + "description": "AllowUnsigned permits installing a project-scoped skill without a\nverified signature; the exception is recorded in the project's lock\nfile.", + "type": "boolean" + }, "clients": { "description": "Clients lists target client identifiers (e.g., \"claude-code\"),\nor [\"all\"] to target every skill-supporting client.\nOmitting this field installs to all available clients.", "items": { @@ -3163,6 +3313,51 @@ }, "type": "object" }, + "pkg_api_v1.pluginBuildListResponse": { + "description": "Response containing a list of locally-built OCI plugin artifacts", + "properties": { + "builds": { + "description": "List of locally-built OCI plugin artifacts", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.LocalBuild" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "pkg_api_v1.pluginListResponse": { + "description": "Response containing a list of installed plugins", + "properties": { + "plugins": { + "description": "List of installed plugins", + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "pkg_api_v1.pluginsV01Response": { + "description": "Paginated list of plugins from the registry", + "properties": { + "metadata": { + "$ref": "#/components/schemas/pkg_api_v1.paginationV01Metadata" + }, + "plugins": { + "description": "Plugins is the list of plugins on the current page", + "items": { + "$ref": "#/components/schemas/registry.Plugin" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, "pkg_api_v1.providerCapabilitiesResponse": { "description": "Capabilities of the secrets provider", "properties": { @@ -3189,6 +3384,16 @@ }, "type": "object" }, + "pkg_api_v1.pushPluginRequest": { + "description": "Request to push a built plugin artifact", + "properties": { + "reference": { + "description": "OCI reference to push", + "type": "string" + } + }, + "type": "object" + }, "pkg_api_v1.pushSkillRequest": { "description": "Request to push a built skill artifact", "properties": { @@ -3275,7 +3480,7 @@ "type": "string" }, "bearer_token": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "callback_port": { "description": "Specific port for OAuth callback server", @@ -3286,7 +3491,7 @@ "type": "string" }, "client_secret": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "issuer": { "description": "OAuth/OIDC issuer URL (e.g., https://accounts.google.com)", @@ -3416,25 +3621,59 @@ }, "type": "object" }, - "pkg_api_v1.toolOverride": { - "description": "Tool override", + "pkg_api_v1.syncSkillsRequest": { + "description": "Request to restore a project's installed skills to match its lock file", "properties": { - "description": { - "description": "Description of the tool", - "type": "string" + "adopt": { + "description": "Adopt writes lock entries for existing unmanaged project-scope installs", + "type": "boolean" }, - "name": { - "description": "Name of the tool", - "type": "string" - } - }, - "type": "object" - }, - "pkg_api_v1.updateRequest": { - "description": "Request to update an existing workload (name cannot be changed)", - "properties": { - "allow_docker_gateway": { - "description": "Whether to permit outbound connections to Docker gateway addresses\n(host.docker.internal, gateway.docker.internal, 172.17.0.1). These are\nblocked by default in the egress proxy even when network isolation is on.\nOnly applicable to Docker deployments with network isolation enabled.", + "allow_unsigned": { + "description": "AllowUnsigned permits adopting skills whose signature state cannot be\nestablished, recording them as unsigned", + "type": "boolean" + }, + "check": { + "description": "Check verifies on-disk content against the lock file without installing or writing anything", + "type": "boolean" + }, + "clients": { + "description": "Clients lists target client identifiers. Empty means every\nskill-supporting client detected on this host.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "project_root": { + "description": "ProjectRoot is the project root path whose lock file should be synced", + "type": "string" + }, + "prune": { + "description": "Prune removes project-scoped skills installed but not present in the lock file", + "type": "boolean" + } + }, + "type": "object" + }, + "pkg_api_v1.toolOverride": { + "description": "Tool override", + "properties": { + "description": { + "description": "Description of the tool", + "type": "string" + }, + "name": { + "description": "Name of the tool", + "type": "string" + } + }, + "type": "object" + }, + "pkg_api_v1.updateRequest": { + "description": "Request to update an existing workload (name cannot be changed)", + "properties": { + "allow_docker_gateway": { + "description": "Whether to permit outbound connections to Docker gateway addresses\n(host.docker.internal, gateway.docker.internal, 172.17.0.1). These are\nblocked by default in the egress proxy even when network isolation is on.\nOnly applicable to Docker deployments with network isolation enabled.", "type": "boolean" }, "authz_config": { @@ -3500,12 +3739,12 @@ "type": "integer" }, "runtime_config": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig" + "$ref": "#/components/schemas/templates.RuntimeConfig" }, "secrets": { "description": "Secret parameters to inject", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter" + "$ref": "#/components/schemas/secrets.SecretParameter" }, "type": "array", "uniqueItems": false @@ -3620,6 +3859,58 @@ }, "type": "object" }, + "pkg_api_v1.upgradeSkillsRequest": { + "description": "Request to re-resolve a project's lock entries and install newer content", + "properties": { + "allow_ref_change": { + "description": "AllowRefChange permits resolvedReference changes during upgrade", + "type": "boolean" + }, + "allow_signer_change": { + "description": "AllowSignerChange permits upgrading to an artifact signed by a\ndifferent identity than the recorded one", + "type": "boolean" + }, + "clients": { + "description": "Clients lists target client identifiers. Empty means every\nskill-supporting client detected on this host.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "fail_on_changes": { + "description": "FailOnChanges exits with an error when any mutable source would upgrade", + "type": "boolean" + }, + "names": { + "description": "Names restricts the upgrade to specific skill names. Empty means every entry.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "preview": { + "description": "Preview reports what would change without installing (still fetches to compare digests)", + "type": "boolean" + }, + "project_root": { + "description": "ProjectRoot is the project root path whose lock file should be upgraded", + "type": "string" + } + }, + "type": "object" + }, + "pkg_api_v1.validatePluginRequest": { + "description": "Request to validate a plugin definition", + "properties": { + "path": { + "description": "Path to the plugin definition directory", + "type": "string" + } + }, + "type": "object" + }, "pkg_api_v1.validateSkillRequest": { "description": "Request to validate a skill definition", "properties": { @@ -3644,7 +3935,7 @@ "workloads": { "description": "List of container information for each workload", "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_core.Workload" + "$ref": "#/components/schemas/core.Workload" }, "type": "array", "uniqueItems": false @@ -3656,7 +3947,21 @@ "description": "Response containing workload status information", "properties": { "status": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus" + "description": "Current status of the workload", + "enum": [ + "running", + "stopped", + "error", + "starting", + "stopping", + "unhealthy", + "removing", + "unknown", + "unauthenticated", + "auth_retrying", + "policy_stopped" + ], + "type": "string" } }, "type": "object" @@ -3950,6 +4255,68 @@ }, "type": "object" }, + "registry.Plugin": { + "properties": { + "_meta": { + "additionalProperties": {}, + "description": "Meta is an opaque payload with extended meta data details of the plugin.", + "type": "object" + }, + "description": { + "description": "Description is the description of the plugin.", + "type": "string" + }, + "icons": { + "description": "Icons is the list of icons for the plugin.", + "items": { + "$ref": "#/components/schemas/registry.SkillIcon" + }, + "type": "array", + "uniqueItems": false + }, + "license": { + "description": "License is the SPDX license identifier of the plugin.", + "type": "string" + }, + "metadata": { + "additionalProperties": {}, + "description": "Metadata is the official metadata of the plugin as reported in the\nplugin manifest file.", + "type": "object" + }, + "name": { + "description": "Name is the name of the plugin.\nThe format is that of identifiers, e.g. \"my-plugin\".", + "type": "string" + }, + "namespace": { + "description": "Namespace is the namespace of the plugin.\nThe format is reverse-DNS, e.g. \"io.github.user\".", + "type": "string" + }, + "packages": { + "description": "Packages is the list of packages for the plugin.", + "items": { + "$ref": "#/components/schemas/registry.SkillPackage" + }, + "type": "array", + "uniqueItems": false + }, + "repository": { + "$ref": "#/components/schemas/registry.SkillRepository" + }, + "status": { + "description": "Status is the status of the plugin.\nCan be one of \"active\", \"deprecated\", or \"archived\".", + "type": "string" + }, + "title": { + "description": "Title is the title of the plugin.\nThis is for human consumption, not an identifier.", + "type": "string" + }, + "version": { + "description": "Version is the version of the plugin.\nAny non-empty string is valid, but ideally it should be either a\nsemantic version or a commit hash.", + "type": "string" + } + }, + "type": "object" + }, "registry.Provenance": { "description": "Provenance contains verification and signing metadata", "properties": { @@ -3977,6 +4344,42 @@ }, "type": "object" }, + "registry.Registry": { + "description": "Full registry data", + "properties": { + "groups": { + "description": "Groups is a slice of group definitions containing related MCP servers", + "items": { + "$ref": "#/components/schemas/registry.Group" + }, + "type": "array", + "uniqueItems": false + }, + "last_updated": { + "description": "LastUpdated is the timestamp when the registry was last updated, in RFC3339 format", + "type": "string" + }, + "remote_servers": { + "additionalProperties": { + "$ref": "#/components/schemas/registry.RemoteServerMetadata" + }, + "description": "RemoteServers is a map of server names to their corresponding remote server definitions\nThese are MCP servers accessed via HTTP/HTTPS using the thv proxy command", + "type": "object" + }, + "servers": { + "additionalProperties": { + "$ref": "#/components/schemas/registry.ImageMetadata" + }, + "description": "Servers is a map of server names to their corresponding server definitions", + "type": "object" + }, + "version": { + "description": "Version is the schema version of the registry", + "type": "string" + } + }, + "type": "object" + }, "registry.RemoteServerMetadata": { "description": "Remote server details (if it's a remote server)", "properties": { @@ -4225,145 +4628,1063 @@ }, "type": "object" }, - "v0.ServerJSON": { + "remote.Config": { + "description": "RemoteAuthConfig contains OAuth configuration for remote MCP servers", "properties": { - "$schema": { - "example": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json", - "format": "uri", - "minLength": 1, + "authorize_url": { "type": "string" }, - "_meta": { - "$ref": "#/components/schemas/v0.ServerMeta" + "bearer_token": { + "description": "Bearer token configuration (alternative to OAuth)", + "type": "string" }, - "description": { - "example": "MCP server providing weather data and forecasts via OpenWeatherMap API", - "maxLength": 100, - "minLength": 1, + "bearer_token_file": { "type": "string" }, - "icons": { - "items": { - "$ref": "#/components/schemas/model.Icon" - }, - "type": "array", - "uniqueItems": false + "cached_cimd_client_id": { + "description": "CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD\nauthentication was used. Kept separate from CachedClientID (which holds\nDCR-issued IDs) so the two can have independent lifecycles — DCR credential\nrotation clears CachedClientID without touching the stable CIMD URL.\nRead by resolveClientCredentials to send the correct client_id on token refresh.", + "type": "string" }, - "name": { - "example": "io.github.user/weather", - "maxLength": 200, - "minLength": 3, - "pattern": "^[a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+$", + "cached_client_id": { + "description": "Cached DCR client credentials for persistence across restarts.\nThese are obtained during Dynamic Client Registration and needed to refresh tokens.\nClientID is stored as plain text since it's public information.", "type": "string" }, - "packages": { - "items": { - "$ref": "#/components/schemas/model.Package" - }, - "type": "array", - "uniqueItems": false + "cached_client_secret_ref": { + "type": "string" }, - "remotes": { - "items": { - "$ref": "#/components/schemas/model.Transport" - }, - "type": "array", - "uniqueItems": false + "cached_dcr_callback_port": { + "description": "CachedDCRCallbackPort is the callback port that was actually registered\nduring DCR. It may differ from CallbackPort when the requested port was\nunavailable and a fallback port was selected.", + "type": "integer" }, - "repository": { - "$ref": "#/components/schemas/model.Repository" + "cached_refresh_token_ref": { + "description": "Cached OAuth token reference for persistence across restarts.\nThe refresh token is stored securely in the secret manager, and this field\ncontains the reference to retrieve it (e.g., \"OAUTH_REFRESH_TOKEN_workload\").\nThis enables session restoration without requiring a new browser-based login.", + "type": "string" }, - "title": { - "example": "Weather API", - "maxLength": 100, - "minLength": 1, + "cached_reg_client_uri": { + "description": "CachedRegClientURI is the registration_client_uri from the DCR response.\nThis is the endpoint used for RFC 7592 client read/update/delete operations.\nStored as plain text since it is not sensitive.", "type": "string" }, - "version": { - "example": "1.0.2", - "maxLength": 255, - "minLength": 1, + "cached_reg_token_ref": { + "description": "CachedRegTokenRef is a secret manager reference to the registration_access_token\nreturned in the DCR response. Used for RFC 7592 client update operations.\nStored as a secret reference since it's sensitive.", "type": "string" }, - "websiteUrl": { - "example": "https://modelcontextprotocol.io/examples", - "format": "uri", + "cached_secret_expiry": { + "description": "ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server).\nA zero value means the secret does not expire.", "type": "string" - } - }, - "type": "object" - }, - "v0.ServerMeta": { - "properties": { - "io.modelcontextprotocol.registry/publisher-provided": { - "additionalProperties": {}, - "type": "object" - } - }, - "type": "object" - }, - "v1.Duration": { - "description": "RefillPeriod is the duration to fully refill the bucket from zero to maxTokens.\nThe effective refill rate is maxTokens / refillPeriod tokens per second.\nFormat: Go duration string (e.g., \"1m0s\", \"30s\", \"1h0m0s\").\n+kubebuilder:validation:Required", - "type": "object" - } - } - }, - "info": { - "description": "This is the ToolHive API server.", - "title": "ToolHive API", - "version": "1.0" - }, - "externalDocs": { - "description": "", - "url": "" - }, - "paths": { + }, + "cached_token_auth_method": { + "description": "CachedTokenEndpointAuthMethod is the auth method used for the token endpoint\n(e.g., \"client_secret_basic\", \"none\"). Persisted for RFC 7592 updates.", + "type": "string" + }, + "cached_token_expiry": { + "type": "string" + }, + "callback_port": { + "type": "integer" + }, + "client_id": { + "type": "string" + }, + "client_secret": { + "type": "string" + }, + "client_secret_file": { + "type": "string" + }, + "issuer": { + "description": "OAuth endpoint configuration (from registry)", + "type": "string" + }, + "oauth_params": { + "additionalProperties": { + "type": "string" + }, + "description": "OAuth parameters for server-specific customization", + "type": "object" + }, + "resource": { + "description": "Resource is the OAuth 2.0 resource indicator (RFC 8707).", + "type": "string" + }, + "scope_param_name": { + "description": "ScopeParamName overrides the query parameter name used to send scopes in the\nauthorization URL. When empty, the standard \"scope\" parameter is used.\nSome providers require a non-standard name (e.g., Slack uses \"user_scope\").", + "type": "string" + }, + "scopes": { + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "skip_browser": { + "type": "boolean" + }, + "timeout": { + "example": "5m", + "type": "string" + }, + "token_url": { + "type": "string" + }, + "use_pkce": { + "type": "boolean" + } + }, + "type": "object" + }, + "secrets.SecretParameter": { + "description": "Bearer token for authentication (alternative to OAuth)", + "properties": { + "name": { + "type": "string" + }, + "target": { + "type": "string" + } + }, + "type": "object" + }, + "storage.ACLUserRunConfig": { + "description": "ACLUserConfig contains ACL user authentication configuration.", + "properties": { + "password_env_var": { + "description": "PasswordEnvVar is the environment variable containing the Redis password.", + "type": "string" + }, + "username_env_var": { + "description": "UsernameEnvVar is the environment variable containing the Redis username.", + "type": "string" + } + }, + "type": "object" + }, + "storage.RedisRunConfig": { + "description": "RedisConfig is the Redis-specific configuration when Type is \"redis\".", + "properties": { + "acl_user_config": { + "$ref": "#/components/schemas/storage.ACLUserRunConfig" + }, + "addr": { + "description": "Addr is the Redis server address (host:port). Required for standalone and cluster modes.\nMutually exclusive with SentinelConfig.", + "type": "string" + }, + "auth_type": { + "description": "AuthType must be \"aclUser\" - only ACL user authentication is supported.", + "type": "string" + }, + "cluster_mode": { + "description": "ClusterMode enables the Redis Cluster protocol. Requires Addr to be set.", + "type": "boolean" + }, + "dial_timeout": { + "description": "DialTimeout is the timeout for establishing connections (e.g., \"5s\").", + "type": "string" + }, + "key_prefix": { + "description": "KeyPrefix for multi-tenancy, typically \"thv:auth:{ns}:{name}:\".", + "type": "string" + }, + "read_timeout": { + "description": "ReadTimeout is the timeout for read operations (e.g., \"3s\").", + "type": "string" + }, + "sentinel_config": { + "$ref": "#/components/schemas/storage.SentinelRunConfig" + }, + "sentinel_tls": { + "$ref": "#/components/schemas/storage.RedisTLSRunConfig" + }, + "tls": { + "$ref": "#/components/schemas/storage.RedisTLSRunConfig" + }, + "write_timeout": { + "description": "WriteTimeout is the timeout for write operations (e.g., \"3s\").", + "type": "string" + } + }, + "type": "object" + }, + "storage.RedisTLSRunConfig": { + "description": "SentinelTLS configures TLS for Sentinel connections. Only applies when SentinelConfig is set.", + "properties": { + "ca_cert_file": { + "description": "CACertFile is the path to a PEM-encoded CA certificate file.", + "type": "string" + }, + "insecure_skip_verify": { + "description": "InsecureSkipVerify skips certificate verification.", + "type": "boolean" + } + }, + "type": "object" + }, + "storage.RunConfig": { + "description": "Storage configures the storage backend for the auth server.\nIf nil, defaults to in-memory storage.", + "properties": { + "redis_config": { + "$ref": "#/components/schemas/storage.RedisRunConfig" + }, + "type": { + "description": "Type specifies the storage backend type. Defaults to \"memory\".", + "type": "string" + } + }, + "type": "object" + }, + "storage.SentinelRunConfig": { + "description": "SentinelConfig contains Sentinel-specific configuration.\nMutually exclusive with Addr.", + "properties": { + "db": { + "description": "DB is the Redis database number (default: 0).", + "type": "integer" + }, + "master_name": { + "description": "MasterName is the name of the Redis Sentinel master.", + "type": "string" + }, + "sentinel_addrs": { + "description": "SentinelAddrs is the list of Sentinel addresses (host:port).", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + } + }, + "type": "object" + }, + "telemetry.Config": { + "description": "DEPRECATED: Middleware configuration.\nTelemetryConfig contains the OpenTelemetry configuration", + "properties": { + "caCertPath": { + "description": "CACertPath is the file path to a CA certificate bundle for the OTLP endpoint.\nWhen set, the OTLP exporters use this CA to verify the collector's TLS certificate\ninstead of relying solely on the system CA pool.\n+optional", + "type": "string" + }, + "customAttributes": { + "additionalProperties": { + "type": "string" + }, + "description": "CustomAttributes contains custom resource attributes to be added to all telemetry signals.\nThese are parsed from CLI flags (--otel-custom-attributes) or environment variables\n(OTEL_RESOURCE_ATTRIBUTES) as key=value pairs.\n+optional", + "type": "object" + }, + "enablePrometheusMetricsPath": { + "description": "EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint.\nThe metrics are served on the main transport port at /metrics.\nThis is separate from OTLP metrics which are sent to the Endpoint.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "endpoint": { + "description": "Endpoint is the OTLP endpoint URL\n+optional", + "type": "string" + }, + "environmentVariables": { + "description": "EnvironmentVariables is a list of environment variable names that should be\nincluded in telemetry spans as attributes. Only variables in this list will\nbe read from the host machine and included in spans for observability.\nExample: [\"NODE_ENV\", \"DEPLOYMENT_ENV\", \"SERVICE_VERSION\"]\n+optional", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "headers": { + "additionalProperties": { + "type": "string" + }, + "description": "Headers contains authentication headers for the OTLP endpoint.\n+optional", + "type": "object" + }, + "insecure": { + "description": "Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "metricsEnabled": { + "description": "MetricsEnabled controls whether OTLP metrics are enabled.\nWhen false, OTLP metrics are not sent even if an endpoint is configured.\nThis is independent of EnablePrometheusMetricsPath.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "samplingRate": { + "description": "SamplingRate is the trace sampling rate (0.0-1.0) as a string.\nOnly used when TracingEnabled is true.\nExample: \"0.05\" for 5% sampling.\n+kubebuilder:default=\"0.05\"\n+optional", + "type": "string" + }, + "serviceName": { + "description": "ServiceName is the service name for telemetry.\nWhen omitted, defaults to the server name (e.g., VirtualMCPServer name).\n+optional", + "type": "string" + }, + "serviceVersion": { + "description": "ServiceVersion is the service version for telemetry.\nWhen omitted, defaults to the ToolHive version.\n+optional", + "type": "string" + }, + "tracingEnabled": { + "description": "TracingEnabled controls whether distributed tracing is enabled.\nWhen false, no tracer provider is created even if an endpoint is configured.\n+kubebuilder:default=false\n+optional", + "type": "boolean" + }, + "useLegacyAttributes": { + "description": "UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names\nare emitted alongside the new standard attribute names. When true, spans include both\nold and new attribute names for backward compatibility with existing dashboards.\nCurrently defaults to true; this will change to false in a future release.\n+kubebuilder:default=true\n+optional", + "type": "boolean" + } + }, + "type": "object" + }, + "templates.RuntimeConfig": { + "description": "RuntimeConfig allows overriding the default runtime configuration\nfor this specific workload (base images and packages)", + "properties": { + "additional_packages": { + "description": "AdditionalPackages lists extra packages to install in the builder and\nruntime stages.\nExamples for Alpine: [\"git\", \"make\", \"gcc\"]\nExamples for Debian: [\"git\", \"build-essential\"]", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "build_with": { + "description": "BuildWith lists build-time dependency constraints, interpreted per\npackage ecosystem. For uvx:// builds these are PEP 508 requirement\nspecifiers passed to `uv tool install --with`, used to constrain\ntransitive dependencies the package itself leaves unbounded\n(e.g. \"mcp\u003c2\"). Ecosystems without constraint support (npx://, go://)\nreject a non-empty BuildWith at build time.", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "builder_image": { + "description": "BuilderImage is the full image reference for the builder stage.\nAn empty string signals \"use the default for this transport type\" during config merging.\nExamples: \"golang:1.26-alpine\", \"node:24-alpine\", \"python:3.14-slim\"", + "type": "string" + }, + "runtime_env": { + "additionalProperties": { + "type": "string" + }, + "description": "RuntimeEnv contains environment variables to inject into the Dockerfile's\nfinal runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv),\nwhich only affects the builder stage, these variables are baked into the\nshipped image and are present in the running container's process\nenvironment at startup. Use this for values a packaged MCP server reads at\nprocess start (e.g. feature flags, cache backend selection), not for\nbuild-time package manager configuration.\nKeys must be uppercase with underscores, values are validated for safety.", + "type": "object" + } + }, + "type": "object" + }, + "tokenexchange.Config": { + "description": "TokenExchangeConfig contains token exchange configuration for external authentication", + "properties": { + "audience": { + "description": "Audience is the target audience for the exchanged token", + "type": "string" + }, + "client_id": { + "description": "ClientID is the OAuth 2.0 client identifier", + "type": "string" + }, + "client_secret": { + "description": "ClientSecret is the OAuth 2.0 client secret", + "type": "string" + }, + "external_token_header_name": { + "description": "ExternalTokenHeaderName is the name of the custom header to use when HeaderStrategy is \"custom\"", + "type": "string" + }, + "header_strategy": { + "description": "HeaderStrategy determines how to inject the token\nValid values: HeaderStrategyReplace (default), HeaderStrategyCustom", + "type": "string" + }, + "scopes": { + "description": "Scopes is the list of scopes to request for the exchanged token", + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": false + }, + "subject_token_type": { + "description": "SubjectTokenType specifies the type of the subject token being exchanged.\nCommon values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT.\nIf empty, defaults to oauthproto.TokenTypeAccessToken.", + "type": "string" + }, + "token_url": { + "description": "TokenURL is the OAuth 2.0 token endpoint URL", + "type": "string" + } + }, + "type": "object" + }, + "types.MiddlewareConfig": { + "properties": { + "parameters": { + "description": "Parameters is a JSON object containing the middleware parameters.\nIt is stored as a raw message to allow flexible parameter types.", + "type": "object" + }, + "type": { + "description": "Type is a string representing the middleware type.", + "type": "string" + } + }, + "type": "object" + }, + "v0.ServerJSON": { + "properties": { + "$schema": { + "example": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json", + "format": "uri", + "minLength": 1, + "type": "string" + }, + "_meta": { + "$ref": "#/components/schemas/v0.ServerMeta" + }, + "description": { + "example": "MCP server providing weather data and forecasts via OpenWeatherMap API", + "maxLength": 100, + "minLength": 1, + "type": "string" + }, + "icons": { + "items": { + "$ref": "#/components/schemas/model.Icon" + }, + "type": "array", + "uniqueItems": false + }, + "name": { + "example": "io.github.user/weather", + "maxLength": 200, + "minLength": 3, + "pattern": "^[a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+$", + "type": "string" + }, + "packages": { + "items": { + "$ref": "#/components/schemas/model.Package" + }, + "type": "array", + "uniqueItems": false + }, + "remotes": { + "items": { + "$ref": "#/components/schemas/model.Transport" + }, + "type": "array", + "uniqueItems": false + }, + "repository": { + "$ref": "#/components/schemas/model.Repository" + }, + "title": { + "example": "Weather API", + "maxLength": 100, + "minLength": 1, + "type": "string" + }, + "version": { + "example": "1.0.2", + "maxLength": 255, + "minLength": 1, + "type": "string" + }, + "websiteUrl": { + "example": "https://modelcontextprotocol.io/examples", + "format": "uri", + "type": "string" + } + }, + "type": "object" + }, + "v0.ServerMeta": { + "properties": { + "io.modelcontextprotocol.registry/publisher-provided": { + "additionalProperties": {}, + "type": "object" + } + }, + "type": "object" + }, + "v1.Duration": { + "description": "RefillPeriod is the duration to fully refill the bucket from zero to maxTokens.\nThe effective refill rate is maxTokens / refillPeriod tokens per second.\nFormat: Go duration string (e.g., \"1m0s\", \"30s\", \"1h0m0s\").\n+kubebuilder:validation:Required", + "type": "object" + } + } + }, + "info": { + "description": "This is the ToolHive API server.", + "title": "ToolHive API", + "version": "1.0" + }, + "externalDocs": { + "description": "", + "url": "" + }, + "paths": { "/api/openapi.json": { "get": { - "description": "Returns the OpenAPI specification for the API", + "description": "Returns the OpenAPI specification for the API", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "type": "object" + } + } + }, + "description": "OpenAPI specification" + } + }, + "summary": "Get OpenAPI specification", + "tags": [ + "system" + ] + } + }, + "/api/v1beta/clients": { + "get": { + "description": "List all registered clients in ToolHive", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_client.RegisteredClient" + }, + "type": "array" + } + } + }, + "description": "OK" + } + }, + "summary": "List all clients", + "tags": [ + "clients" + ] + }, + "post": { + "description": "Register a new client with ToolHive", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.createClientRequest", + "summary": "client", + "description": "Client to register" + } + ] + } + } + }, + "description": "Client to register", + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" + } + } + }, + "description": "OK" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Register a new client", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/register": { + "post": { + "description": "Register multiple clients with ToolHive", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", + "summary": "clients", + "description": "Clients to register" + } + ] + } + } + }, + "description": "Clients to register", + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" + }, + "type": "array" + } + } + }, + "description": "OK" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Register multiple clients", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/unregister": { + "post": { + "description": "Unregister multiple clients from ToolHive", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", + "summary": "clients", + "description": "Clients to unregister" + } + ] + } + } + }, + "description": "Clients to unregister", + "required": true + }, + "responses": { + "204": { + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Unregister multiple clients", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/{name}": { + "delete": { + "description": "Unregister a client from ToolHive", + "parameters": [ + { + "description": "Client name to unregister", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + } + }, + "summary": "Unregister a client", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/clients/{name}/groups/{group}": { + "delete": { + "description": "Unregister a client from a specific group in ToolHive", + "parameters": [ + { + "description": "Client name to unregister", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Group name to remove client from", + "in": "path", + "name": "group", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Invalid request or unsupported client type" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Client or group not found" + } + }, + "summary": "Unregister a client from a specific group", + "tags": [ + "clients" + ] + } + }, + "/api/v1beta/discovery/clients": { + "get": { + "description": "List all clients compatible with ToolHive and their status.\nEach object includes supports_skills when ToolHive can install skills for that client.", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.clientStatusResponse" + } + } + }, + "description": "OK" + } + }, + "summary": "List all clients status", + "tags": [ + "discovery" + ] + } + }, + "/api/v1beta/groups": { + "get": { + "description": "Get a list of all groups", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.groupListResponse" + } + } + }, + "description": "OK" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "List all groups", + "tags": [ + "groups" + ] + }, + "post": { + "description": "Create a new group with the specified name", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.createGroupRequest", + "summary": "group", + "description": "Group creation request" + } + ] + } + } + }, + "description": "Group creation request", + "required": true + }, + "responses": { + "201": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.createGroupResponse" + } + } + }, + "description": "Created" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "409": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Conflict" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "Create a new group", + "tags": [ + "groups" + ] + } + }, + "/api/v1beta/groups/{name}": { + "delete": { + "description": "Delete a group by name.", + "parameters": [ + { + "description": "Group name", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Delete all workloads in the group (default: false, moves workloads to default group)", + "in": "query", + "name": "with-workloads", + "schema": { + "type": "boolean" + } + } + ], + "responses": { + "204": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "No Content" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "Delete a group", + "tags": [ + "groups" + ] + }, + "get": { + "description": "Get details of a specific group", + "parameters": [ + { + "description": "Group name", + "in": "path", + "name": "name", + "required": true, + "schema": { + "type": "string" + } + } + ], "responses": { "200": { "content": { "application/json": { "schema": { - "type": "object" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_groups.Group" } } }, - "description": "OpenAPI specification" + "description": "OK" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "Get OpenAPI specification", + "summary": "Get group details", "tags": [ - "system" + "groups" ] } }, - "/api/v1beta/clients": { + "/api/v1beta/plugins": { "get": { - "description": "List all registered clients in ToolHive", + "description": "Get a list of all installed plugins", + "parameters": [ + { + "description": "Filter by scope (user or project)", + "in": "query", + "name": "scope", + "schema": { + "enum": [ + "user", + "project" + ], + "type": "string" + } + }, + { + "description": "Filter by client app", + "in": "query", + "name": "client", + "schema": { + "type": "string" + } + }, + { + "description": "Filter by project root path", + "in": "query", + "name": "project_root", + "schema": { + "type": "string" + } + }, + { + "description": "Filter by group name", + "in": "query", + "name": "group", + "schema": { + "type": "string" + } + } + ], "responses": { "200": { "content": { "application/json": { "schema": { - "items": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_client.RegisteredClient" - }, - "type": "array" + "$ref": "#/components/schemas/pkg_api_v1.pluginListResponse" } } }, "description": "OK" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "List all clients", + "summary": "List all installed plugins", "tags": [ - "clients" + "plugins" ] }, "post": { - "description": "Register a new client with ToolHive", + "description": "Install a plugin from a remote source", "requestBody": { "content": { "application/json": { @@ -4373,29 +5694,107 @@ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.createClientRequest", - "summary": "client", - "description": "Client to register" + "$ref": "#/components/schemas/pkg_api_v1.installPluginRequest", + "summary": "request", + "description": "Install request" } ] } } }, - "description": "Client to register", - "required": true - }, - "responses": { - "200": { + "description": "Install request", + "required": true + }, + "responses": { + "201": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.installPluginResponse" + } + } + }, + "description": "Created", + "headers": { + "Location": { + "description": "URI of the installed plugin resource", + "schema": { + "type": "string" + } + } + } + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "401": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Unauthorized (registry refused credentials)" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found (artifact not present in registry)" + }, + "409": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Conflict" + }, + "429": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Too Many Requests (registry rate limit)" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + }, + "502": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" + "type": "string" } } }, - "description": "OK" + "description": "Bad Gateway (upstream registry failure)" }, - "400": { + "504": { "content": { "application/json": { "schema": { @@ -4403,18 +5802,18 @@ } } }, - "description": "Invalid request or unsupported client type" + "description": "Gateway Timeout (upstream pull timed out)" } }, - "summary": "Register a new client", + "summary": "Install a plugin", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/register": { + "/api/v1beta/plugins/build": { "post": { - "description": "Register multiple clients with ToolHive", + "description": "Build a plugin from a local directory", "requestBody": { "content": { "application/json": { @@ -4424,15 +5823,15 @@ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", - "summary": "clients", - "description": "Clients to register" + "$ref": "#/components/schemas/pkg_api_v1.buildPluginRequest", + "summary": "request", + "description": "Build request" } ] } } }, - "description": "Clients to register", + "description": "Build request", "required": true }, "responses": { @@ -4440,10 +5839,7 @@ "content": { "application/json": { "schema": { - "items": { - "$ref": "#/components/schemas/pkg_api_v1.createClientResponse" - }, - "type": "array" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.BuildResult" } } }, @@ -4457,43 +5853,40 @@ } } }, - "description": "Invalid request or unsupported client type" + "description": "Bad Request" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "Register multiple clients", + "summary": "Build a plugin", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/unregister": { - "post": { - "description": "Unregister multiple clients from ToolHive", - "requestBody": { - "content": { - "application/json": { - "schema": { - "oneOf": [ - { - "type": "object" - }, - { - "$ref": "#/components/schemas/pkg_api_v1.bulkClientRequest", - "summary": "clients", - "description": "Clients to unregister" - } - ] - } - } - }, - "description": "Clients to unregister", - "required": true - }, + "/api/v1beta/plugins/builds": { + "get": { + "description": "Get a list of all locally-built OCI plugin artifacts in the local store", "responses": { - "204": { - "description": "No Content" + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.pluginBuildListResponse" + } + } + }, + "description": "OK" }, - "400": { + "500": { "content": { "application/json": { "schema": { @@ -4501,23 +5894,23 @@ } } }, - "description": "Invalid request or unsupported client type" + "description": "Internal Server Error" } }, - "summary": "Unregister multiple clients", + "summary": "List locally-built plugin artifacts", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/{name}": { + "/api/v1beta/plugins/builds/{tag}": { "delete": { - "description": "Unregister a client from ToolHive", + "description": "Remove a locally-built OCI plugin artifact and its blobs from the local store", "parameters": [ { - "description": "Client name to unregister", + "description": "Artifact tag", "in": "path", - "name": "name", + "name": "tag", "required": true, "schema": { "type": "string" @@ -4526,9 +5919,16 @@ ], "responses": { "204": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, "description": "No Content" }, - "400": { + "404": { "content": { "application/json": { "schema": { @@ -4536,32 +5936,33 @@ } } }, - "description": "Invalid request or unsupported client type" + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" } }, - "summary": "Unregister a client", + "summary": "Delete a locally-built plugin artifact", "tags": [ - "clients" + "plugins" ] } }, - "/api/v1beta/clients/{name}/groups/{group}": { - "delete": { - "description": "Unregister a client from a specific group in ToolHive", + "/api/v1beta/plugins/content": { + "get": { + "description": "Retrieve the plugin.json body and file listing from an artifact\nwithout installing it. Accepts OCI refs, git refs, or local tags.", "parameters": [ { - "description": "Client name to unregister", - "in": "path", - "name": "name", - "required": true, - "schema": { - "type": "string" - } - }, - { - "description": "Group name to remove client from", - "in": "path", - "name": "group", + "description": "OCI reference or local build tag", + "in": "query", + "name": "ref", "required": true, "schema": { "type": "string" @@ -4569,8 +5970,15 @@ } ], "responses": { - "204": { - "description": "No Content" + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginContent" + } + } + }, + "description": "OK" }, "400": { "content": { @@ -4580,7 +5988,17 @@ } } }, - "description": "Invalid request or unsupported client type" + "description": "Bad Request" + }, + "401": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Unauthorized (registry refused credentials)" }, "404": { "content": { @@ -4590,49 +6008,108 @@ } } }, - "description": "Client or group not found" - } - }, - "summary": "Unregister a client from a specific group", - "tags": [ - "clients" - ] - } - }, - "/api/v1beta/discovery/clients": { - "get": { - "description": "List all clients compatible with ToolHive and their status.\nEach object includes supports_skills when ToolHive can install skills for that client.", - "responses": { - "200": { + "description": "Not Found (artifact not present in registry)" + }, + "429": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.clientStatusResponse" + "type": "string" } } }, - "description": "OK" + "description": "Too Many Requests (registry rate limit)" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + }, + "502": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Gateway (upstream registry or git resolver failure)" + }, + "504": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Gateway Timeout (upstream pull timed out)" } }, - "summary": "List all clients status", + "summary": "Get plugin content", "tags": [ - "discovery" + "plugins" ] } }, - "/api/v1beta/groups": { - "get": { - "description": "Get a list of all groups", + "/api/v1beta/plugins/push": { + "post": { + "description": "Push a built plugin artifact to a remote registry", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.pushPluginRequest", + "summary": "request", + "description": "Push request" + } + ] + } + } + }, + "description": "Push request", + "required": true + }, "responses": { - "200": { + "204": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.groupListResponse" + "type": "string" } } }, - "description": "OK" + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" }, "500": { "content": { @@ -4645,13 +6122,15 @@ "description": "Internal Server Error" } }, - "summary": "List all groups", + "summary": "Push a plugin", "tags": [ - "groups" + "plugins" ] - }, + } + }, + "/api/v1beta/plugins/validate": { "post": { - "description": "Create a new group with the specified name", + "description": "Validate a plugin definition", "requestBody": { "content": { "application/json": { @@ -4661,27 +6140,27 @@ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.createGroupRequest", - "summary": "group", - "description": "Group creation request" + "$ref": "#/components/schemas/pkg_api_v1.validatePluginRequest", + "summary": "request", + "description": "Validate request" } ] } } }, - "description": "Group creation request", + "description": "Validate request", "required": true }, "responses": { - "201": { + "200": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/pkg_api_v1.createGroupResponse" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ValidationResult" } } }, - "description": "Created" + "description": "OK" }, "400": { "content": { @@ -4693,16 +6172,6 @@ }, "description": "Bad Request" }, - "409": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Conflict" - }, "500": { "content": { "application/json": { @@ -4714,18 +6183,18 @@ "description": "Internal Server Error" } }, - "summary": "Create a new group", + "summary": "Validate a plugin", "tags": [ - "groups" + "plugins" ] } }, - "/api/v1beta/groups/{name}": { + "/api/v1beta/plugins/{name}": { "delete": { - "description": "Delete a group by name.", + "description": "Remove an installed plugin", "parameters": [ { - "description": "Group name", + "description": "Plugin name", "in": "path", "name": "name", "required": true, @@ -4734,11 +6203,23 @@ } }, { - "description": "Delete all workloads in the group (default: false, moves workloads to default group)", + "description": "Scope to uninstall from (user or project)", "in": "query", - "name": "with-workloads", + "name": "scope", "schema": { - "type": "boolean" + "enum": [ + "user", + "project" + ], + "type": "string" + } + }, + { + "description": "Project root path for project-scoped plugins", + "in": "query", + "name": "project_root", + "schema": { + "type": "string" } } ], @@ -4753,6 +6234,16 @@ }, "description": "No Content" }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, "404": { "content": { "application/json": { @@ -4774,22 +6265,42 @@ "description": "Internal Server Error" } }, - "summary": "Delete a group", + "summary": "Uninstall a plugin", "tags": [ - "groups" + "plugins" ] }, "get": { - "description": "Get details of a specific group", + "description": "Get detailed information about a specific plugin", "parameters": [ { - "description": "Group name", + "description": "Plugin name", "in": "path", "name": "name", "required": true, "schema": { "type": "string" } + }, + { + "description": "Filter by scope (user or project)", + "in": "query", + "name": "scope", + "schema": { + "enum": [ + "user", + "project" + ], + "type": "string" + } + }, + { + "description": "Project root path for project-scoped plugins", + "in": "query", + "name": "project_root", + "schema": { + "type": "string" + } } ], "responses": { @@ -4797,12 +6308,22 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_groups.Group" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginInfo" } } }, "description": "OK" }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, "404": { "content": { "application/json": { @@ -4824,9 +6345,9 @@ "description": "Internal Server Error" } }, - "summary": "Get group details", + "summary": "Get plugin details", "tags": [ - "groups" + "plugins" ] } }, @@ -6045,32 +7566,214 @@ "description": "Internal Server Error" } }, - "summary": "Delete a locally-built skill artifact", + "summary": "Delete a locally-built skill artifact", + "tags": [ + "skills" + ] + } + }, + "/api/v1beta/skills/content": { + "get": { + "description": "Retrieve the SKILL.md body and file listing from an artifact\nwithout installing it. Accepts OCI refs, git refs, or local tags.", + "parameters": [ + { + "description": "OCI reference or local build tag", + "in": "query", + "name": "ref", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillContent" + } + } + }, + "description": "OK" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "401": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Unauthorized (registry refused credentials)" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found (artifact not present in registry)" + }, + "429": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Too Many Requests (registry rate limit)" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + }, + "502": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Gateway (upstream registry or git resolver failure)" + }, + "504": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Gateway Timeout (upstream pull timed out)" + } + }, + "summary": "Get skill content", + "tags": [ + "skills" + ] + } + }, + "/api/v1beta/skills/push": { + "post": { + "description": "Push a built skill artifact to a remote registry", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.pushSkillRequest", + "summary": "request", + "description": "Push request" + } + ] + } + } + }, + "description": "Push request", + "required": true + }, + "responses": { + "204": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "No Content" + }, + "400": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Bad Request" + }, + "404": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Internal Server Error" + } + }, + "summary": "Push a skill", "tags": [ "skills" ] } }, - "/api/v1beta/skills/content": { - "get": { - "description": "Retrieve the SKILL.md body and file listing from an artifact\nwithout installing it. Accepts OCI refs, git refs, or local tags.", - "parameters": [ - { - "description": "OCI reference or local build tag", - "in": "query", - "name": "ref", - "required": true, - "schema": { - "type": "string" + "/api/v1beta/skills/sync": { + "post": { + "description": "Restore a project's installed skills to match toolhive.lock.yaml", + "requestBody": { + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "type": "object" + }, + { + "$ref": "#/components/schemas/pkg_api_v1.syncSkillsRequest", + "summary": "request", + "description": "Sync request" + } + ] + } } - } - ], + }, + "description": "Sync request", + "required": true + }, "responses": { "200": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillContent" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SyncResult" } } }, @@ -6086,27 +7789,7 @@ }, "description": "Bad Request" }, - "401": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Unauthorized (registry refused credentials)" - }, - "404": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Not Found (artifact not present in registry)" - }, - "429": { + "403": { "content": { "application/json": { "schema": { @@ -6114,7 +7797,7 @@ } } }, - "description": "Too Many Requests (registry rate limit)" + "description": "Forbidden (feature not enabled)" }, "500": { "content": { @@ -6126,17 +7809,7 @@ }, "description": "Internal Server Error" }, - "502": { - "content": { - "application/json": { - "schema": { - "type": "string" - } - } - }, - "description": "Bad Gateway (upstream registry or git resolver failure)" - }, - "504": { + "501": { "content": { "application/json": { "schema": { @@ -6144,18 +7817,18 @@ } } }, - "description": "Gateway Timeout (upstream pull timed out)" + "description": "Not Implemented" } }, - "summary": "Get skill content", + "summary": "Sync project skills from the lock file", "tags": [ "skills" ] } }, - "/api/v1beta/skills/push": { + "/api/v1beta/skills/upgrade": { "post": { - "description": "Push a built skill artifact to a remote registry", + "description": "Re-resolve a project's lock entries and install newer content where available", "requestBody": { "content": { "application/json": { @@ -6165,27 +7838,27 @@ "type": "object" }, { - "$ref": "#/components/schemas/pkg_api_v1.pushSkillRequest", + "$ref": "#/components/schemas/pkg_api_v1.upgradeSkillsRequest", "summary": "request", - "description": "Push request" + "description": "Upgrade request" } ] } } }, - "description": "Push request", + "description": "Upgrade request", "required": true }, "responses": { - "204": { + "200": { "content": { "application/json": { "schema": { - "type": "string" + "$ref": "#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeResult" } } }, - "description": "No Content" + "description": "OK" }, "400": { "content": { @@ -6197,6 +7870,16 @@ }, "description": "Bad Request" }, + "403": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Forbidden (feature not enabled)" + }, "404": { "content": { "application/json": { @@ -6205,7 +7888,7 @@ } } }, - "description": "Not Found" + "description": "Not Found (a requested name is not in the lock file)" }, "500": { "content": { @@ -6216,9 +7899,19 @@ } }, "description": "Internal Server Error" + }, + "501": { + "content": { + "application/json": { + "schema": { + "type": "string" + } + } + }, + "description": "Not Implemented" } }, - "summary": "Push a skill", + "summary": "Upgrade project skills", "tags": [ "skills" ] @@ -7565,6 +9258,162 @@ ] } }, + "/registry/{registryName}/v0.1/x/dev.toolhive/plugins": { + "get": { + "description": "Get a paginated list of plugins from the registry. Supports optional full-text search and pagination.", + "parameters": [ + { + "description": "Registry name (currently ignored, uses the default provider)", + "in": "path", + "name": "registryName", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Search filter — matches against plugin name, namespace, and description", + "in": "query", + "name": "q", + "schema": { + "type": "string" + } + }, + { + "description": "Page number, 1-based (default: 1)", + "in": "query", + "name": "page", + "schema": { + "type": "integer" + } + }, + { + "description": "Items per page, max 200 (default: 50)", + "in": "query", + "name": "limit", + "schema": { + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.pluginsV01Response" + } + } + }, + "description": "OK" + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Internal server error" + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Registry authentication required or upstream registry unavailable" + } + }, + "summary": "List available registry plugins", + "tags": [ + "registry-plugins" + ] + } + }, + "/registry/{registryName}/v0.1/x/dev.toolhive/plugins/{namespace}/{pluginName}": { + "get": { + "description": "Retrieve a single plugin by its namespace and name from the registry.", + "parameters": [ + { + "description": "Registry name (currently ignored, uses the default provider)", + "in": "path", + "name": "registryName", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Plugin namespace in reverse-DNS format (e.g. io.github.stacklok)", + "in": "path", + "name": "namespace", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Plugin name", + "in": "path", + "name": "pluginName", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/registry.Plugin" + } + } + }, + "description": "OK" + }, + "404": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Plugin not found" + }, + "500": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Internal server error" + }, + "503": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/pkg_api_v1.registryErrorResponse" + } + } + }, + "description": "Registry authentication required or upstream registry unavailable" + } + }, + "summary": "Get a registry plugin", + "tags": [ + "registry-plugins" + ] + } + }, "/registry/{registryName}/v0.1/x/dev.toolhive/skills": { "get": { "description": "Get a paginated list of skills from the registry. Supports optional full-text search and pagination.", diff --git a/docs/server/swagger.yaml b/docs/server/swagger.yaml index 970dcf9db8..8734b14b0f 100644 --- a/docs/server/swagger.yaml +++ b/docs/server/swagger.yaml @@ -1,34 +1,143 @@ components: schemas: - github_com_stacklok_toolhive-core_registry_types.Registry: - description: Full registry data + auth.TokenValidatorConfig: + description: |- + DEPRECATED: Middleware configuration. + OIDCConfig contains OIDC configuration properties: - groups: - description: Groups is a slice of group definitions containing related MCP - servers + allowPrivateIP: + description: AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses + type: boolean + allowedClientIDs: + description: |- + AllowedClientIDs is the set of OAuth client IDs accepted when validating + tokens by their client_id claim instead of an audience claim. items: - $ref: '#/components/schemas/registry.Group' + type: string type: array - uniqueItems: false - last_updated: - description: LastUpdated is the timestamp when the registry was last updated, - in RFC3339 format + audience: + description: Audience is the expected audience for the token type: string - remote_servers: - additionalProperties: - $ref: '#/components/schemas/registry.RemoteServerMetadata' + authTokenFile: + description: AuthTokenFile is the path to file containing bearer token for + authentication + type: string + cacertPath: + description: CACertPath is the path to the CA certificate bundle for HTTPS + requests + type: string + clientID: + description: ClientID is the OIDC client ID + type: string + clientSecret: + description: ClientSecret is the optional OIDC client secret for introspection + type: string + insecureAllowHTTP: description: |- - RemoteServers is a map of server names to their corresponding remote server definitions - These are MCP servers accessed via HTTP/HTTPS using the thv proxy command - type: object - servers: + InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing + WARNING: This is insecure and should NEVER be used in production + type: boolean + introspectionURL: + description: IntrospectionURL is the optional introspection endpoint for + validating tokens + type: string + issuer: + description: Issuer is the OIDC issuer URL (e.g., https://accounts.google.com) + type: string + jwksurl: + description: JWKSURL is the URL to fetch the JWKS from + type: string + resourceURL: + description: ResourceURL is the explicit resource URL for OAuth discovery + (RFC 9728) + type: string + scopes: + description: |- + Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728) + If empty, defaults to ["openid"] + items: + type: string + type: array + type: object + core.Workload: + properties: + created_at: + description: CreatedAt is the timestamp when the workload was created. + type: string + group: + description: Group is the name of the group this workload belongs to, if + any. + type: string + labels: additionalProperties: - $ref: '#/components/schemas/registry.ImageMetadata' - description: Servers is a map of server names to their corresponding server - definitions + type: string + description: Labels are the container labels (excluding standard ToolHive + labels) type: object - version: - description: Version is the schema version of the registry + name: + description: |- + Name is the name of the workload. + It is used as a unique identifier. + type: string + package: + description: Package specifies the Workload Package used to create this + Workload. + type: string + port: + description: |- + Port is the port on which the workload is exposed. + This is embedded in the URL. + type: integer + proxy_mode: + description: |- + ProxyMode is the proxy mode that clients should use to connect. + For stdio transports, this will be the proxy mode (sse or streamable-http). + For direct transports (sse/streamable-http), this will be the same as TransportType. + type: string + remote: + description: Remote indicates whether this is a remote workload (true) or + a container workload (false). + type: boolean + started_at: + description: StartedAt is when the container was last started (changes on + restart) + type: string + status: + description: Status is the current status of the workload. + enum: + - running + - stopped + - error + - starting + - stopping + - unhealthy + - removing + - unknown + - unauthenticated + - auth_retrying + - policy_stopped + type: string + status_context: + description: |- + StatusContext provides additional context about the workload's status. + The exact meaning is determined by the status and the underlying runtime. + type: string + tools: + description: ToolsFilter is the filter on tools applied to the workload. + items: + type: string + type: array + uniqueItems: false + transport_type: + description: TransportType is the type of transport used for this workload. + enum: + - stdio + - sse + - streamable-http + - inspector + type: string + url: + description: URL is the URL of the workload exposed by the ToolHive proxy. type: string type: object github_com_stacklok_toolhive_cmd_thv-operator_api_v1beta1.RateLimitConfig: @@ -120,65 +229,15 @@ components: +kubebuilder:default=1024 +optional type: integer - type: object - github_com_stacklok_toolhive_pkg_auth.TokenValidatorConfig: - description: |- - DEPRECATED: Middleware configuration. - OIDCConfig contains OIDC configuration - properties: - allowPrivateIP: - description: AllowPrivateIP allows JWKS/OIDC endpoints on private IP addresses - type: boolean - allowedClientIDs: - description: |- - AllowedClientIDs is the set of OAuth client IDs accepted when validating - tokens by their client_id claim instead of an audience claim. - items: - type: string - type: array - audience: - description: Audience is the expected audience for the token - type: string - authTokenFile: - description: AuthTokenFile is the path to file containing bearer token for - authentication - type: string - cacertPath: - description: CACertPath is the path to the CA certificate bundle for HTTPS - requests - type: string - clientID: - description: ClientID is the OIDC client ID - type: string - clientSecret: - description: ClientSecret is the optional OIDC client secret for introspection - type: string - insecureAllowHTTP: - description: |- - InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing - WARNING: This is insecure and should NEVER be used in production - type: boolean - introspectionURL: - description: IntrospectionURL is the optional introspection endpoint for - validating tokens - type: string - issuer: - description: Issuer is the OIDC issuer URL (e.g., https://accounts.google.com) - type: string - jwksurl: - description: JWKSURL is the URL to fetch the JWKS from - type: string - resourceURL: - description: ResourceURL is the explicit resource URL for OAuth discovery - (RFC 9728) - type: string - scopes: + maxDelegationDepth: description: |- - Scopes is the list of OAuth scopes to advertise in the well-known endpoint (RFC 9728) - If empty, defaults to ["openid"] - items: - type: string - type: array + MaxDelegationDepth caps how many nested RFC 8693 "act" entries are + recorded in an audit event's delegation chain. Deeper chains are + truncated (marked with truncated=true). Defaults to 10 when unset. + +kubebuilder:validation:Minimum=1 + +kubebuilder:default=10 + +optional + type: integer type: object github_com_stacklok_toolhive_pkg_auth_awssts.Config: description: AWSStsConfig contains AWS STS token exchange configuration for @@ -248,126 +307,41 @@ components: description: RoleArn is the IAM role ARN to assume when this mapping matches. type: string type: object - github_com_stacklok_toolhive_pkg_auth_remote.Config: - description: RemoteAuthConfig contains OAuth configuration for remote MCP servers + github_com_stacklok_toolhive_pkg_auth_upstreamswap.Config: + description: |- + UpstreamSwapConfig contains configuration for upstream token swap middleware. + When set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs + for upstream IdP tokens before forwarding requests to the MCP server. properties: - authorize_url: - type: string - bearer_token: - description: Bearer token configuration (alternative to OAuth) - type: string - bearer_token_file: - type: string - cached_cimd_client_id: - description: |- - CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD - authentication was used. Kept separate from CachedClientID (which holds - DCR-issued IDs) so the two can have independent lifecycles — DCR credential - rotation clears CachedClientID without touching the stable CIMD URL. - Read by resolveClientCredentials to send the correct client_id on token refresh. - type: string - cached_client_id: - description: |- - Cached DCR client credentials for persistence across restarts. - These are obtained during Dynamic Client Registration and needed to refresh tokens. - ClientID is stored as plain text since it's public information. + custom_header_name: + description: CustomHeaderName is the header name when HeaderStrategy is + "custom". type: string - cached_client_secret_ref: + header_strategy: + description: 'HeaderStrategy determines how to inject the token: "replace" + (default) or "custom".' type: string - cached_refresh_token_ref: + provider_name: description: |- - Cached OAuth token reference for persistence across restarts. - The refresh token is stored securely in the secret manager, and this field - contains the reference to retrieve it (e.g., "OAUTH_REFRESH_TOKEN_workload"). - This enables session restoration without requiring a new browser-based login. + ProviderName identifies which upstream provider's tokens to retrieve for injection. + This is required and must match a configured upstream provider name. type: string - cached_reg_token_ref: + type: object + github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig: + description: |- + CIMD controls client_id metadata document support. When enabled, the + embedded authorization server accepts HTTPS URLs as client_id values + and resolves them via the CIMD protocol instead of requiring DCR. + properties: + cache_fallback_ttl: description: |- - RegistrationAccessToken is used to update/delete the client registration. - Stored as a secret reference since it's sensitive. + CacheFallbackTTL is the fixed TTL applied to every cached CIMD document. + Cache-Control header parsing is not yet implemented; all entries use this value. + Format: Go duration string (e.g. "5m", "10m", "1h"). + Defaults to 5 minutes when Enabled is true and this field is omitted. + example: 5m type: string - cached_secret_expiry: - description: |- - ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server). - A zero value means the secret does not expire. - type: string - cached_token_expiry: - type: string - callback_port: - type: integer - client_id: - type: string - client_secret: - type: string - client_secret_file: - type: string - issuer: - description: OAuth endpoint configuration (from registry) - type: string - oauth_params: - additionalProperties: - type: string - description: OAuth parameters for server-specific customization - type: object - resource: - description: Resource is the OAuth 2.0 resource indicator (RFC 8707). - type: string - scope_param_name: - description: |- - ScopeParamName overrides the query parameter name used to send scopes in the - authorization URL. When empty, the standard "scope" parameter is used. - Some providers require a non-standard name (e.g., Slack uses "user_scope"). - type: string - scopes: - items: - type: string - type: array - uniqueItems: false - skip_browser: - type: boolean - timeout: - example: 5m - type: string - token_url: - type: string - use_pkce: - type: boolean - type: object - github_com_stacklok_toolhive_pkg_auth_upstreamswap.Config: - description: |- - UpstreamSwapConfig contains configuration for upstream token swap middleware. - When set along with EmbeddedAuthServerConfig, this middleware exchanges ToolHive JWTs - for upstream IdP tokens before forwarding requests to the MCP server. - properties: - custom_header_name: - description: CustomHeaderName is the header name when HeaderStrategy is - "custom". - type: string - header_strategy: - description: 'HeaderStrategy determines how to inject the token: "replace" - (default) or "custom".' - type: string - provider_name: - description: |- - ProviderName identifies which upstream provider's tokens to retrieve for injection. - This is required and must match a configured upstream provider name. - type: string - type: object - github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig: - description: |- - CIMD controls client_id metadata document support. When enabled, the - embedded authorization server accepts HTTPS URLs as client_id values - and resolves them via the CIMD protocol instead of requiring DCR. - properties: - cache_fallback_ttl: - description: |- - CacheFallbackTTL is the fixed TTL applied to every cached CIMD document. - Cache-Control header parsing is not yet implemented; all entries use this value. - Format: Go duration string (e.g. "5m", "10m", "1h"). - Defaults to 5 minutes when Enabled is true and this field is omitted. - example: 5m - type: string - cache_max_size: + cache_max_size: description: |- CacheMaxSize is the maximum number of CIMD documents held in the LRU cache. Defaults to 256 when Enabled is true and this field is zero. @@ -476,10 +450,14 @@ components: allow_private_ips: description: |- AllowPrivateIPs permits the upstream provider's HTTP client to connect to - private IP ranges (RFC-1918, link-local). Use only when the upstream is - hosted inside the same cluster and has no public endpoint. HTTP-scheme - restrictions are unchanged — HTTPS is still required for non-localhost hosts. - Defaults to false. + private IP ranges (RFC-1918, link-local). When DCRConfig is set, this + also gates the DCR discovery and registration calls made on this + upstream's behalf (see pkg/authserver/runner/dcr_adapter.go), so a + single flag covers the whole upstream rather than needing a separate + DCR-specific setting. Use only when the upstream is hosted inside the + same cluster and has no public endpoint. HTTP-scheme restrictions are + unchanged — HTTPS is still required for non-localhost hosts. Defaults + to false. type: boolean authorization_endpoint: description: AuthorizationEndpoint is the URL for the OAuth authorization @@ -641,6 +619,12 @@ components: uniqueItems: false cimd: $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.CIMDRunConfig' + delegation_token_lifespan: + description: |- + DelegationTokenLifespan is the maximum lifetime for delegated tokens issued + via RFC 8693 token exchange. Specified as a Go duration string (e.g., "15m"). + If empty, defaults to 15 minutes. + type: string disable_upstream_token_injection: description: |- DisableUpstreamTokenInjection prevents the upstream swap middleware from being added. @@ -687,9 +671,31 @@ components: signing_key_config: $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.SigningKeyRunConfig' storage: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RunConfig' + $ref: '#/components/schemas/storage.RunConfig' token_lifespans: $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver.TokenLifespanRunConfig' + trusted_issuers: + description: |- + TrustedIssuers lists external OIDC issuers whose tokens are accepted as + subject tokens during RFC 8693 token exchange. Empty (the default) means + only self-issued subject tokens are accepted. + + Prerequisite: the token-exchange grant requires a confidential client, + and no supported deployment path provisions one today (DCR and CIMD + clients are both public-only, and there is no client-seeding field on + this RunConfig), so this grant is not yet usable end to end for + self-issued or external subject tokens alike. Tracked in + https://github.com/stacklok/toolhive/issues/6082. + + See tokenexchange.TrustedIssuer for the per-issuer field reference, and + docs/arch/17-token-exchange-delegation.md for the trust model, consent + signals, and operator-facing constraints (audience/scope bounding, + subject namespace qualification, required client binding) that aren't + visible from the config shape alone. + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_server_tokenexchange.TrustedIssuer' + type: array + uniqueItems: false upstreams: description: |- Upstreams configures connections to upstream Identity Providers. @@ -853,95 +859,74 @@ components: If not specified, defaults to GET. type: string type: object - github_com_stacklok_toolhive_pkg_authserver_storage.ACLUserRunConfig: - description: ACLUserConfig contains ACL user authentication configuration. - properties: - password_env_var: - description: PasswordEnvVar is the environment variable containing the Redis - password. - type: string - username_env_var: - description: UsernameEnvVar is the environment variable containing the Redis - username. - type: string - type: object - github_com_stacklok_toolhive_pkg_authserver_storage.RedisRunConfig: - description: RedisConfig is the Redis-specific configuration when Type is "redis". + github_com_stacklok_toolhive_pkg_authserver_server_tokenexchange.TrustedIssuer: properties: - acl_user_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.ACLUserRunConfig' - addr: + actor_claim: description: |- - Addr is the Redis server address (host:port). Required for standalone and cluster modes. - Mutually exclusive with SentinelConfig. - type: string - auth_type: - description: AuthType must be "aclUser" - only ACL user authentication is - supported. - type: string - cluster_mode: - description: ClusterMode enables the Redis Cluster protocol. Requires Addr - to be set. - type: boolean - dial_timeout: - description: DialTimeout is the timeout for establishing connections (e.g., - "5s"). - type: string - key_prefix: - description: KeyPrefix for multi-tenancy, typically "thv:auth:{ns}:{name}:". - type: string - read_timeout: - description: ReadTimeout is the timeout for read operations (e.g., "3s"). - type: string - sentinel_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.SentinelRunConfig' - sentinel_tls: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig' - tls: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig' - write_timeout: - description: WriteTimeout is the timeout for write operations (e.g., "3s"). - type: string - type: object - github_com_stacklok_toolhive_pkg_authserver_storage.RedisTLSRunConfig: - description: SentinelTLS configures TLS for Sentinel connections. Only applies - when SentinelConfig is set. - properties: - ca_cert_file: - description: CACertFile is the path to a PEM-encoded CA certificate file. + ActorClaim names the claim identifying the client that requested the + subject token from THIS EXTERNAL ISSUER (used by AllowedActors below). + Values are in the external issuer's namespace, NOT ToolHive client + IDs. Defaults to "azp"; use "appid" for Microsoft Entra v1, "cid" for + Okta. The special value "client_id" reads ValidatedClaims.ClientID + instead of Extra (assignClaim routes it to that field) — it is still + the external token's client_id claim, not a ToolHive one. type: string - insecure_skip_verify: - description: InsecureSkipVerify skips certificate verification. + allow_private_ips: + description: |- + AllowPrivateIPs permits OIDC discovery and JWKS fetches for THIS + issuer to resolve to a private or loopback address. Use only when the + issuer is hosted inside the same cluster and has no public endpoint. type: boolean - type: object - github_com_stacklok_toolhive_pkg_authserver_storage.RunConfig: - description: |- - Storage configures the storage backend for the auth server. - If nil, defaults to in-memory storage. - properties: - redis_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_authserver_storage.RedisRunConfig' - type: - description: Type specifies the storage backend type. Defaults to "memory". - type: string - type: object - github_com_stacklok_toolhive_pkg_authserver_storage.SentinelRunConfig: - description: |- - SentinelConfig contains Sentinel-specific configuration. - Mutually exclusive with Addr. - properties: - db: - description: 'DB is the Redis database number (default: 0).' - type: integer - master_name: - description: MasterName is the name of the Redis Sentinel master. - type: string - sentinel_addrs: - description: SentinelAddrs is the list of Sentinel addresses (host:port). + allowed_actors: + description: |- + AllowedActors is the allowlist of ActorClaim values authorized to + exchange a subject token from this issuer when it carries no + "may_act" claim; empty means only may_act-bearing tokens are + accepted. By itself names no ToolHive client — see + AllowedDelegateClients and docs/arch/17-token-exchange-delegation.md + ("Accepted limitations" #1). + items: + type: string + type: array + uniqueItems: false + allowed_delegate_clients: + description: |- + AllowedDelegateClients restricts which ToolHive client IDs may + exchange a subject token from this issuer, for BOTH consent paths. + Required (validateTrustedIssuer rejects empty/absent); "*" permits + any confidential client holding the grant. See + docs/arch/17-token-exchange-delegation.md ("Accepted limitations" #1). items: type: string type: array uniqueItems: false + expected_audience: + description: |- + ExpectedAudience is the expected "aud" claim value that must appear + in the token's audience list (a resource/API identifier, not a + client ID — required, but not enforced; see looksLikeResourceIdentifier). + See docs/arch/17-token-exchange-delegation.md ("ID/access-token + discrimination") for why and its limits. + type: string + insecure_allow_http: + description: |- + InsecureAllowHTTP permits plain-HTTP OIDC discovery and JWKS fetches + for THIS issuer only. Development and testing only — never set in + production. Does not relax the private-IP guard; see AllowPrivateIPs. + Deliberately per-issuer: this server's own InsecureAllowHTTP must not + silently permit plaintext discovery for every trusted external issuer + too — a network attacker who can intercept that traffic could + substitute a JWKS and forge subject tokens for that issuer's + namespace. + type: boolean + issuer_url: + description: IssuerURL is the expected "iss" claim value (exact match). + type: string + jwks_url: + description: |- + JWKSURL is the URL to fetch the issuer's JSON Web Key Set from. + If empty, it is resolved via OIDC discovery at {IssuerURL}/.well-known/openid-configuration. + type: string type: object github_com_stacklok_toolhive_pkg_authz.Config: description: |- @@ -1039,219 +1024,261 @@ components: name: $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_client.ClientApp' type: object - github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus: - description: Current status of the workload - enum: - - running - - stopped - - error - - starting - - stopping - - unhealthy - - removing - - unknown - - unauthenticated - - auth_retrying - - policy_stopped - - running - - stopped - - error - - starting - - stopping - - unhealthy - - removing - - unknown - - unauthenticated - - auth_retrying - - policy_stopped - - running - - stopped - - error - - starting - - stopping - - unhealthy - - removing - - unknown - - unauthenticated - - auth_retrying - - policy_stopped - type: string - x-enum-varnames: - - WorkloadStatusRunning - - WorkloadStatusStopped - - WorkloadStatusError - - WorkloadStatusStarting - - WorkloadStatusStopping - - WorkloadStatusUnhealthy - - WorkloadStatusRemoving - - WorkloadStatusUnknown - - WorkloadStatusUnauthenticated - - WorkloadStatusAuthRetrying - - WorkloadStatusPolicyStopped - github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig: - description: |- - RuntimeConfig allows overriding the default runtime configuration - for this specific workload (base images and packages) + github_com_stacklok_toolhive_pkg_groups.Group: properties: - additional_packages: - description: |- - AdditionalPackages lists extra packages to install in the builder and - runtime stages. - Examples for Alpine: ["git", "make", "gcc"] - Examples for Debian: ["git", "build-essential"] + name: + type: string + plugins: items: type: string type: array uniqueItems: false - builder_image: - description: |- - BuilderImage is the full image reference for the builder stage. - An empty string signals "use the default for this transport type" during config merging. - Examples: "golang:1.26-alpine", "node:24-alpine", "python:3.14-slim" - type: string - runtime_env: - additionalProperties: + registered_clients: + items: type: string - description: |- - RuntimeEnv contains environment variables to inject into the Dockerfile's - final runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv), - which only affects the builder stage, these variables are baked into the - shipped image and are present in the running container's process - environment at startup. Use this for values a packaged MCP server reads at - process start (e.g. feature flags, cache backend selection), not for - build-time package manager configuration. - Keys must be uppercase with underscores, values are validated for safety. - type: object + type: array + uniqueItems: false + skills: + items: + type: string + type: array + uniqueItems: false type: object - github_com_stacklok_toolhive_pkg_core.Workload: + github_com_stacklok_toolhive_pkg_plugins.BuildResult: properties: - created_at: - description: CreatedAt is the timestamp when the workload was created. + reference: + description: Reference is the OCI reference of the built skill artifact. type: string - group: - description: Group is the name of the group this workload belongs to, if - any. + type: object + github_com_stacklok_toolhive_pkg_plugins.ComponentInventory: + additionalProperties: + type: integer + description: |- + Components is the inventory of component types declared by the plugin + (e.g. {"commands": 3, "skills": 2}). Extracted from the OCI artifact. + type: object + github_com_stacklok_toolhive_pkg_plugins.ComponentType: + enum: + - commands + - agents + - skills + - hooks + - mcpServers + - lspServers + type: string + x-enum-varnames: + - ComponentCommands + - ComponentAgents + - ComponentSkills + - ComponentHooks + - ComponentMCP + - ComponentLSP + github_com_stacklok_toolhive_pkg_plugins.Dependency: + properties: + digest: + description: Digest is the OCI digest for upgrade detection. type: string - labels: - additionalProperties: - type: string - description: Labels are the container labels (excluding standard ToolHive - labels) - type: object name: - description: |- - Name is the name of the workload. - It is used as a unique identifier. + description: Name is the dependency name. type: string - package: - description: Package specifies the Workload Package used to create this - Workload. + reference: + description: Reference is the OCI reference for the dependency. type: string - port: - description: |- - Port is the port on which the workload is exposed. - This is embedded in the URL. - type: integer - proxy_mode: - description: |- - ProxyMode is the proxy mode that clients should use to connect. - For stdio transports, this will be the proxy mode (sse or streamable-http). - For direct transports (sse/streamable-http), this will be the same as TransportType. + type: object + github_com_stacklok_toolhive_pkg_plugins.InstallStatus: + description: Status is the current installation status. + enum: + - installed + - pending + - failed + type: string + x-enum-varnames: + - InstallStatusInstalled + - InstallStatusPending + - InstallStatusFailed + github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin: + description: InstalledPlugin contains the full installation record. + properties: + clients: + description: Clients is the list of client identifiers the plugin is installed + for. + items: + type: string + type: array + uniqueItems: false + components: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ComponentInventory' + dependencies: + description: Dependencies is the list of external plugin dependencies. + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Dependency' + type: array + uniqueItems: false + digest: + description: Digest is the OCI digest (sha256:...) for upgrade detection. type: string - remote: - description: Remote indicates whether this is a remote workload (true) or - a container workload (false). - type: boolean - started_at: - description: StartedAt is when the container was last started (changes on - restart) + installed_at: + description: InstalledAt is the timestamp when the plugin was installed. + type: string + metadata: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginMetadata' + project_root: + description: ProjectRoot is the project root path for project-scoped plugins. + Empty for user-scoped. + type: string + reference: + description: Reference is the full OCI reference (e.g. ghcr.io/org/plugin:v1). + type: string + scope: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Scope' + signature: + description: Signature is the optional signing signature for the plugin + artifact. type: string status: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus' - status_context: - description: |- - StatusContext provides additional context about the workload's status. - The exact meaning is determined by the status and the underlying runtime. + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstallStatus' + tag: + description: Tag is the OCI tag (e.g. v1.0.0). type: string - tools: - description: ToolsFilter is the filter on tools applied to the workload. + type: object + github_com_stacklok_toolhive_pkg_plugins.LocalBuild: + properties: + description: + description: Description is the skill description extracted from the artifact + metadata, if available. + type: string + digest: + description: Digest is the OCI digest of the artifact (sha256:...). + type: string + name: + description: Name is the skill name extracted from the artifact metadata, + if available. + type: string + tag: + description: Tag is the OCI tag or name used to reference the artifact. + type: string + version: + description: Version is the skill version extracted from the artifact metadata, + if available. + type: string + type: object + github_com_stacklok_toolhive_pkg_plugins.PluginContent: + properties: + description: + description: Description is the plugin description from the OCI config labels. + type: string + files: + description: Files is the list of all files in the artifact with their sizes. items: - type: string + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginFileEntry' type: array uniqueItems: false - transport_type: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.TransportType' - url: - description: URL is the URL of the workload exposed by the ToolHive proxy. + license: + description: License is the SPDX license identifier from the OCI config + labels. + type: string + manifest: + description: Manifest is the raw .claude-plugin/plugin.json body. + type: string + name: + description: Name is the plugin name from the OCI config labels. + type: string + version: + description: Version is the plugin version from the OCI config labels. type: string type: object - github_com_stacklok_toolhive_pkg_groups.Group: + github_com_stacklok_toolhive_pkg_plugins.PluginFileEntry: properties: - name: + path: + description: Path is the file path within the artifact. type: string - plugins: + size: + description: Size is the uncompressed file size in bytes. + type: integer + type: object + github_com_stacklok_toolhive_pkg_plugins.PluginInfo: + properties: + installed_plugin: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin' + metadata: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginMetadata' + project_scope_degraded_clients: + description: |- + ProjectScopeDegradedClients lists the client types for which a + project-scoped install degraded (the adapter could only materialize at + user scope — e.g. Codex always writes to the user-scoped config.toml). + Populated by Info; empty for user-scoped installs. Recomputed at read + time from the stored scope + each adapter's capability, mirroring the + UnmaterializedComponents pattern (no persistence needed — the degradation + is deterministic from scope + client type). items: type: string type: array uniqueItems: false - registered_clients: + unmaterialized_components: + additionalProperties: + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ComponentType' + type: array + description: |- + UnmaterializedComponents lists, per client type, the component types the + plugin declares that the installed client adapter does NOT load. Populated + by Info by diffing InstalledPlugin.Components against each installed + client adapter's SupportedComponents. + type: object + type: object + github_com_stacklok_toolhive_pkg_plugins.PluginMetadata: + description: Metadata contains the plugin's metadata. + properties: + author: + description: Author is the plugin author or maintainer. + type: string + description: + description: Description is a human-readable description of the plugin. + type: string + keywords: + description: Keywords is a list of keywords for categorization/search. items: type: string type: array uniqueItems: false - skills: + license: + description: License is the SPDX license identifier for the plugin. + type: string + name: + description: Name is the unique name of the plugin (kebab-case). + type: string + version: + description: Version is the semantic version of the plugin. + type: string + type: object + github_com_stacklok_toolhive_pkg_plugins.Scope: + description: Scope for the installation + enum: + - user + - project + type: string + x-enum-varnames: + - ScopeUser + - ScopeProject + github_com_stacklok_toolhive_pkg_plugins.ValidationResult: + properties: + errors: + description: Errors is a list of validation errors, if any. items: type: string type: array uniqueItems: false - type: object - github_com_stacklok_toolhive_pkg_ignore.Config: - description: IgnoreConfig contains configuration for ignore processing - properties: - loadGlobal: - description: Whether to load global ignore patterns - type: boolean - printOverlays: - description: Whether to print resolved overlay paths for debugging + valid: + description: Valid indicates whether the skill definition is valid. type: boolean - type: object - github_com_stacklok_toolhive_pkg_oauthproto_tokenexchange.Config: - description: TokenExchangeConfig contains token exchange configuration for external - authentication - properties: - audience: - description: Audience is the target audience for the exchanged token - type: string - client_id: - description: ClientID is the OAuth 2.0 client identifier - type: string - client_secret: - description: ClientSecret is the OAuth 2.0 client secret - type: string - external_token_header_name: - description: ExternalTokenHeaderName is the name of the custom header to - use when HeaderStrategy is "custom" - type: string - header_strategy: - description: |- - HeaderStrategy determines how to inject the token - Valid values: HeaderStrategyReplace (default), HeaderStrategyCustom - type: string - scopes: - description: Scopes is the list of scopes to request for the exchanged token + warnings: + description: Warnings is a list of non-blocking validation warnings, if + any. items: type: string type: array uniqueItems: false - subject_token_type: - description: |- - SubjectTokenType specifies the type of the subject token being exchanged. - Common values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT. - If empty, defaults to oauthproto.TokenTypeAccessToken. - type: string - token_url: - description: TokenURL is the OAuth 2.0 token endpoint URL - type: string type: object github_com_stacklok_toolhive_pkg_ratelimit_types.RateLimitBucket: description: |- @@ -1340,7 +1367,7 @@ components: second instance of that middleware to the chain; the seam does not validate against this. items: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig' + $ref: '#/components/schemas/types.MiddlewareConfig' type: array uniqueItems: false allow_docker_gateway: @@ -1427,7 +1454,7 @@ components: description: Host is the host for the HTTP proxy type: string ignore_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_ignore.Config' + $ref: '#/components/schemas/ignore.Config' image: description: Image is the Docker image to run type: string @@ -1457,7 +1484,7 @@ components: MiddlewareConfigs contains the list of middleware to apply to the transport and the configuration for each middleware. items: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig' + $ref: '#/components/schemas/types.MiddlewareConfig' type: array uniqueItems: false mutating_webhooks: @@ -1472,7 +1499,7 @@ components: description: Name is the name of the MCP server type: string oidc_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_auth.TokenValidatorConfig' + $ref: '#/components/schemas/auth.TokenValidatorConfig' permission_profile_name_or_path: description: PermissionProfileNameOrPath is the name or path of the permission profile @@ -1481,7 +1508,15 @@ components: description: Port is the port for the HTTP proxy to listen on (host port) type: integer proxy_mode: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.ProxyMode' + description: |- + ProxyMode is the effective HTTP protocol the proxy uses. + For stdio transports, this is the configured mode (sse or streamable-http). + For direct transports (sse/streamable-http), this matches the transport type. + Note: "sse" is deprecated; use "streamable-http" instead. + enum: + - sse + - streamable-http + type: string publish: description: Publish lists ports to publish to the host in format "hostPort:containerPort" items: @@ -1510,12 +1545,12 @@ components: Empty when the server was not discovered via registry lookup. type: string remote_auth_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_auth_remote.Config' + $ref: '#/components/schemas/remote.Config' remote_url: description: RemoteURL is the URL of the remote MCP server (if running remotely) type: string runtime_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig' + $ref: '#/components/schemas/templates.RuntimeConfig' scaling_config: $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_runner.ScalingConfig' schema_version: @@ -1546,6 +1581,13 @@ components: POST-based health check instead of the default GET probe. Applies to both remote URLs and local container workloads. type: boolean + strict_protocol_validation: + description: |- + StrictProtocolValidation enables strict MCP-Protocol-Version validation + on the streamable HTTP proxy: a request whose header names an unknown + MCP revision is rejected with HTTP 400. Default false accepts any + version string (an absent header is always accepted in either mode). + type: boolean target_host: description: TargetHost is the host to forward traffic to (only applicable to SSE transport) @@ -1555,14 +1597,14 @@ components: to SSE transport) type: integer telemetry_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_telemetry.Config' + $ref: '#/components/schemas/telemetry.Config' thv_ca_bundle: description: |- DEPRECATED: No longer appears to be used. ThvCABundle is the path to the CA certificate bundle for ToolHive HTTP operations type: string token_exchange_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_oauthproto_tokenexchange.Config' + $ref: '#/components/schemas/tokenexchange.Config' tools_filter: description: |- DEPRECATED: Middleware configuration. @@ -1579,7 +1621,13 @@ components: ToolsOverride is a map from an actual tool to its overridden name and/or description type: object transport: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_transport_types.TransportType' + description: Transport is the transport mode (stdio, sse, or streamable-http) + enum: + - stdio + - sse + - streamable-http + - inspector + type: string trust_proxy_headers: description: TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies @@ -1645,14 +1693,6 @@ components: description: Name is the redefined name of the tool type: string type: object - github_com_stacklok_toolhive_pkg_secrets.SecretParameter: - description: Bearer token for authentication (alternative to OAuth) - properties: - name: - type: string - target: - type: string - type: object github_com_stacklok_toolhive_pkg_skills.BuildResult: properties: reference: @@ -1671,6 +1711,27 @@ components: description: Reference is the OCI reference for the dependency. type: string type: object + github_com_stacklok_toolhive_pkg_skills.FailureReason: + description: Reason is a typed failure reason when Status is UpgradeStatusFailed. + enum: + - registry-unreachable + - digest-missing + - validation-rejected + - lock-write-failed + - signature-invalid + - signer-mismatch + - unsigned-rejected + - unknown + type: string + x-enum-varnames: + - FailureReasonRegistryUnreachable + - FailureReasonDigestMissing + - FailureReasonValidationRejected + - FailureReasonLockWriteFailed + - FailureReasonSignatureInvalid + - FailureReasonSignerMismatch + - FailureReasonUnsignedRejected + - FailureReasonUnknown github_com_stacklok_toolhive_pkg_skills.InstallStatus: description: Status is the current installation status. enum: @@ -1705,6 +1766,12 @@ components: installed_at: description: InstalledAt is the timestamp when the skill was installed. type: string + managed: + description: |- + Managed indicates this install is tracked in the project's + toolhive.lock.yaml. Only ever true for project-scoped installs. No + omitempty: false is an observable state (unmanaged), not an absence. + type: boolean metadata: $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SkillMetadata' project_root: @@ -1815,175 +1882,157 @@ components: description: Version is the semantic version of the skill. type: string type: object - github_com_stacklok_toolhive_pkg_skills.ValidationResult: + github_com_stacklok_toolhive_pkg_skills.SyncFailure: properties: - errors: - description: Errors is a list of validation errors, if any. + error: + description: Error is a human-readable description of the failure. + type: string + name: + description: Name is the skill name that failed. + type: string + reason: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.FailureReason' + type: object + github_com_stacklok_toolhive_pkg_skills.SyncResult: + properties: + already_current: + description: AlreadyCurrent lists skills that already matched the lock file. items: type: string type: array uniqueItems: false - valid: - description: Valid indicates whether the skill definition is valid. - type: boolean - warnings: - description: Warnings is a list of non-blocking validation warnings, if - any. + drifted: + description: |- + Drifted lists skills whose on-disk contentDigest differed from the lock + file. Normally these are reinstalled to match it; when Check is set, + nothing is written and this field reports the drift only. items: type: string type: array uniqueItems: false - type: object - github_com_stacklok_toolhive_pkg_telemetry.Config: - description: |- - DEPRECATED: Middleware configuration. - TelemetryConfig contains the OpenTelemetry configuration - properties: - caCertPath: + failed: description: |- - CACertPath is the file path to a CA certificate bundle for the OTLP endpoint. - When set, the OTLP exporters use this CA to verify the collector's TLS certificate - instead of relying solely on the system CA pool. - +optional - type: string - customAttributes: - additionalProperties: + Failed lists skills that could not be synced, with the reason for each. + Drift alone is never reported here — see Drifted. + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SyncFailure' + type: array + uniqueItems: false + installed: + description: Installed lists skills that were installed or reinstalled to + match the lock file. + items: type: string + type: array + uniqueItems: false + missing: description: |- - CustomAttributes contains custom resource attributes to be added to all telemetry signals. - These are parsed from CLI flags (--otel-custom-attributes) or environment variables - (OTEL_RESOURCE_ATTRIBUTES) as key=value pairs. - +optional - type: object - enablePrometheusMetricsPath: - description: |- - EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint. - The metrics are served on the main transport port at /metrics. - This is separate from OTLP metrics which are sent to the Endpoint. - +kubebuilder:default=false - +optional - type: boolean - endpoint: - description: |- - Endpoint is the OTLP endpoint URL - +optional - type: string - environmentVariables: - description: |- - EnvironmentVariables is a list of environment variable names that should be - included in telemetry spans as attributes. Only variables in this list will - be read from the host machine and included in spans for observability. - Example: ["NODE_ENV", "DEPLOYMENT_ENV", "SERVICE_VERSION"] - +optional + Missing lists lock entries with no corresponding install record at all + — the fresh-clone state. Normally these are installed at their pinned + reference; when Check is set, nothing is written and this field + reports the gap only. items: type: string type: array uniqueItems: false - headers: - additionalProperties: + never_managed: + description: NeverManaged lists project-scoped skills never recorded as + lock-managed. + items: type: string - description: |- - Headers contains authentication headers for the OTLP endpoint. - +optional - type: object - insecure: - description: |- - Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint. - +kubebuilder:default=false - +optional - type: boolean - metricsEnabled: - description: |- - MetricsEnabled controls whether OTLP metrics are enabled. - When false, OTLP metrics are not sent even if an endpoint is configured. - This is independent of EnablePrometheusMetricsPath. - +kubebuilder:default=false - +optional - type: boolean - samplingRate: - description: |- - SamplingRate is the trace sampling rate (0.0-1.0) as a string. - Only used when TracingEnabled is true. - Example: "0.05" for 5% sampling. - +kubebuilder:default="0.05" - +optional + type: array + uniqueItems: false + pruned: + description: Pruned lists removed-from-lock skills that were uninstalled + because Prune was set. + items: + type: string + type: array + uniqueItems: false + removed_from_lock: + description: RemovedFromLock lists previously managed skills absent from + the lock file. + items: + type: string + type: array + uniqueItems: false + type: object + github_com_stacklok_toolhive_pkg_skills.UpgradeOutcome: + properties: + error: + description: Error is a human-readable description of the failure, set only + when Status is UpgradeStatusFailed. type: string - serviceName: - description: |- - ServiceName is the service name for telemetry. - When omitted, defaults to the server name (e.g., VirtualMCPServer name). - +optional + name: + description: Name is the skill name. type: string - serviceVersion: + new_digest: description: |- - ServiceVersion is the service version for telemetry. - When omitted, defaults to the ToolHive version. - +optional + NewDigest is the digest the source currently resolves to. Equal to + OldDigest when Status is UpgradeStatusUpToDate. type: string - tracingEnabled: - description: |- - TracingEnabled controls whether distributed tracing is enabled. - When false, no tracer provider is created even if an endpoint is configured. - +kubebuilder:default=false - +optional - type: boolean - useLegacyAttributes: + new_resolved_reference: + description: NewResolvedReference is the new resolvedReference when it changed. + type: string + new_signer_identity: description: |- - UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names - are emitted alongside the new standard attribute names. When true, spans include both - old and new attribute names for backward compatibility with existing dashboards. - Currently defaults to true; this will change to false in a future release. - +kubebuilder:default=true - +optional - type: boolean + NewSignerIdentity is the candidate's signer identity when it differs + from the recorded one (empty when the candidate is unsigned). + type: string + old_digest: + description: OldDigest is the digest pinned in the lock file before this + operation. + type: string + reason: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.FailureReason' + status: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeStatus' type: object - github_com_stacklok_toolhive_pkg_transport_types.MiddlewareConfig: + github_com_stacklok_toolhive_pkg_skills.UpgradeResult: properties: - parameters: - description: |- - Parameters is a JSON object containing the middleware parameters. - It is stored as a raw message to allow flexible parameter types. - type: object - type: - description: Type is a string representing the middleware type. - type: string + outcomes: + description: Outcomes contains one entry per skill considered for upgrade. + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeOutcome' + type: array + uniqueItems: false type: object - github_com_stacklok_toolhive_pkg_transport_types.ProxyMode: - description: |- - ProxyMode is the effective HTTP protocol the proxy uses. - For stdio transports, this is the configured mode (sse or streamable-http). - For direct transports (sse/streamable-http), this matches the transport type. - Note: "sse" is deprecated; use "streamable-http" instead. - enum: - - sse - - streamable-http - - sse - - streamable-http - type: string - x-enum-varnames: - - ProxyModeSSE - - ProxyModeStreamableHTTP - github_com_stacklok_toolhive_pkg_transport_types.TransportType: - description: Transport is the transport mode (stdio, sse, or streamable-http) + github_com_stacklok_toolhive_pkg_skills.UpgradeStatus: + description: Status is the outcome of the upgrade attempt. enum: - - stdio - - sse - - streamable-http - - inspector - - stdio - - sse - - streamable-http - - inspector - - stdio - - sse - - streamable-http - - inspector + - upgraded + - up-to-date + - not-upgradable + - ref-change-blocked + - signer-change-blocked + - failed type: string x-enum-varnames: - - TransportTypeStdio - - TransportTypeSSE - - TransportTypeStreamableHTTP - - TransportTypeInspector + - UpgradeStatusUpgraded + - UpgradeStatusUpToDate + - UpgradeStatusNotUpgradable + - UpgradeStatusRefChangeBlocked + - UpgradeStatusSignerChangeBlocked + - UpgradeStatusFailed + github_com_stacklok_toolhive_pkg_skills.ValidationResult: + properties: + errors: + description: Errors is a list of validation errors, if any. + items: + type: string + type: array + uniqueItems: false + valid: + description: Valid indicates whether the skill definition is valid. + type: boolean + warnings: + description: Warnings is a list of non-blocking validation warnings, if + any. + items: + type: string + type: array + uniqueItems: false + type: object github_com_stacklok_toolhive_pkg_webhook.Config: properties: failure_policy: @@ -2145,6 +2194,16 @@ components: - StatusNotRegistrySourced - StatusServerNotFound - StatusUnknown + ignore.Config: + description: IgnoreConfig contains configuration for ignore processing + properties: + loadGlobal: + description: Whether to load global ignore patterns + type: boolean + printOverlays: + description: Whether to print resolved overlay paths for debugging + type: boolean + type: object model.Argument: properties: choices: @@ -2511,6 +2570,16 @@ components: type: array uniqueItems: false type: object + pkg_api_v1.buildPluginRequest: + description: Request to build a plugin from a local directory + properties: + path: + description: Path to the plugin definition directory + type: string + tag: + description: OCI tag for the built artifact + type: string + type: object pkg_api_v1.buildSkillRequest: description: Request to build a skill from a local directory properties: @@ -2659,11 +2728,11 @@ components: from (e.g. "default"). type: string runtime_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig' + $ref: '#/components/schemas/templates.RuntimeConfig' secrets: description: Secret parameters to inject items: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter' + $ref: '#/components/schemas/secrets.SecretParameter' type: array uniqueItems: false server: @@ -2755,7 +2824,7 @@ components: description: Name of the registry type: string registry: - $ref: '#/components/schemas/github_com_stacklok_toolhive-core_registry_types.Registry' + $ref: '#/components/schemas/registry.Registry' server_count: description: Number of servers in the registry type: integer @@ -2822,9 +2891,51 @@ components: Use AddHeadersFromSecret for sensitive data like API keys. type: object type: object + pkg_api_v1.installPluginRequest: + description: Request to install a plugin + properties: + clients: + description: |- + Clients lists target client identifiers (e.g., "claude-code"), + or ["all"] to target every plugin-supporting client. + Omitting this field installs to all available clients. + items: + type: string + type: array + uniqueItems: false + force: + description: Force allows overwriting unmanaged plugin directories + type: boolean + group: + description: Group is the group name to add the plugin to after installation + type: string + name: + description: Name or OCI reference of the plugin to install + type: string + project_root: + description: ProjectRoot is the project root path for project-scoped installs + type: string + scope: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.Scope' + version: + description: Version to install (empty means latest) + type: string + type: object + pkg_api_v1.installPluginResponse: + description: Response after successfully installing a plugin + properties: + plugin: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin' + type: object pkg_api_v1.installSkillRequest: description: Request to install a skill properties: + allow_unsigned: + description: |- + AllowUnsigned permits installing a project-scoped skill without a + verified signature; the exception is recorded in the project's lock + file. + type: boolean clients: description: |- Clients lists target client identifiers (e.g., "claude-code"), @@ -2925,6 +3036,38 @@ components: description: Total is the total number of items matching the query type: integer type: object + pkg_api_v1.pluginBuildListResponse: + description: Response containing a list of locally-built OCI plugin artifacts + properties: + builds: + description: List of locally-built OCI plugin artifacts + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.LocalBuild' + type: array + uniqueItems: false + type: object + pkg_api_v1.pluginListResponse: + description: Response containing a list of installed plugins + properties: + plugins: + description: List of installed plugins + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.InstalledPlugin' + type: array + uniqueItems: false + type: object + pkg_api_v1.pluginsV01Response: + description: Paginated list of plugins from the registry + properties: + metadata: + $ref: '#/components/schemas/pkg_api_v1.paginationV01Metadata' + plugins: + description: Plugins is the list of plugins on the current page + items: + $ref: '#/components/schemas/registry.Plugin' + type: array + uniqueItems: false + type: object pkg_api_v1.providerCapabilitiesResponse: description: Capabilities of the secrets provider properties: @@ -2944,6 +3087,13 @@ components: description: Whether the provider can write secrets type: boolean type: object + pkg_api_v1.pushPluginRequest: + description: Request to push a built plugin artifact + properties: + reference: + description: OCI reference to push + type: string + type: object pkg_api_v1.pushSkillRequest: description: Request to push a built skill artifact properties: @@ -3015,7 +3165,7 @@ components: non-OIDC OAuth) type: string bearer_token: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter' + $ref: '#/components/schemas/secrets.SecretParameter' callback_port: description: Specific port for OAuth callback server type: integer @@ -3023,7 +3173,7 @@ components: description: OAuth client ID for authentication type: string client_secret: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter' + $ref: '#/components/schemas/secrets.SecretParameter' issuer: description: OAuth/OIDC issuer URL (e.g., https://accounts.google.com) type: string @@ -3119,6 +3269,40 @@ components: type: array uniqueItems: false type: object + pkg_api_v1.syncSkillsRequest: + description: Request to restore a project's installed skills to match its lock + file + properties: + adopt: + description: Adopt writes lock entries for existing unmanaged project-scope + installs + type: boolean + allow_unsigned: + description: |- + AllowUnsigned permits adopting skills whose signature state cannot be + established, recording them as unsigned + type: boolean + check: + description: Check verifies on-disk content against the lock file without + installing or writing anything + type: boolean + clients: + description: |- + Clients lists target client identifiers. Empty means every + skill-supporting client detected on this host. + items: + type: string + type: array + uniqueItems: false + project_root: + description: ProjectRoot is the project root path whose lock file should + be synced + type: string + prune: + description: Prune removes project-scoped skills installed but not present + in the lock file + type: boolean + type: object pkg_api_v1.toolOverride: description: Tool override properties: @@ -3191,11 +3375,11 @@ components: description: Port for the HTTP proxy to listen on type: integer runtime_config: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_container_templates.RuntimeConfig' + $ref: '#/components/schemas/templates.RuntimeConfig' secrets: description: Secret parameters to inject items: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_secrets.SecretParameter' + $ref: '#/components/schemas/secrets.SecretParameter' type: array uniqueItems: false target_port: @@ -3286,6 +3470,53 @@ components: type: array uniqueItems: false type: object + pkg_api_v1.upgradeSkillsRequest: + description: Request to re-resolve a project's lock entries and install newer + content + properties: + allow_ref_change: + description: AllowRefChange permits resolvedReference changes during upgrade + type: boolean + allow_signer_change: + description: |- + AllowSignerChange permits upgrading to an artifact signed by a + different identity than the recorded one + type: boolean + clients: + description: |- + Clients lists target client identifiers. Empty means every + skill-supporting client detected on this host. + items: + type: string + type: array + uniqueItems: false + fail_on_changes: + description: FailOnChanges exits with an error when any mutable source would + upgrade + type: boolean + names: + description: Names restricts the upgrade to specific skill names. Empty + means every entry. + items: + type: string + type: array + uniqueItems: false + preview: + description: Preview reports what would change without installing (still + fetches to compare digests) + type: boolean + project_root: + description: ProjectRoot is the project root path whose lock file should + be upgraded + type: string + type: object + pkg_api_v1.validatePluginRequest: + description: Request to validate a plugin definition + properties: + path: + description: Path to the plugin definition directory + type: string + type: object pkg_api_v1.validateSkillRequest: description: Request to validate a skill definition properties: @@ -3304,7 +3535,7 @@ components: workloads: description: List of container information for each workload items: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_core.Workload' + $ref: '#/components/schemas/core.Workload' type: array uniqueItems: false type: object @@ -3312,7 +3543,20 @@ components: description: Response containing workload status information properties: status: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_container_runtime.WorkloadStatus' + description: Current status of the workload + enum: + - running + - stopped + - error + - starting + - stopping + - unhealthy + - removing + - unknown + - unauthenticated + - auth_retrying + - policy_stopped + type: string type: object registry.EnvVar: properties: @@ -3594,49 +3838,140 @@ components: Defaults to true for enhanced security type: boolean type: object - registry.Provenance: - description: Provenance contains verification and signing metadata - properties: - attestation: - $ref: '#/components/schemas/registry.VerifiedAttestation' - cert_issuer: - type: string - repository_ref: - type: string - repository_uri: - type: string - runner_environment: - type: string - signer_identity: - type: string - sigstore_url: - type: string - type: object - registry.RemoteServerMetadata: - description: Remote server details (if it's a remote server) + registry.Plugin: properties: - custom_metadata: + _meta: additionalProperties: {} - description: CustomMetadata allows for additional user-defined metadata + description: Meta is an opaque payload with extended meta data details of + the plugin. type: object description: - description: Description is a human-readable description of the server's - purpose and functionality + description: Description is the description of the plugin. type: string - env_vars: - description: |- - EnvVars defines environment variables that can be passed to configure the client - These might be needed for client-side configuration when connecting to the remote server + icons: + description: Icons is the list of icons for the plugin. items: - $ref: '#/components/schemas/registry.EnvVar' + $ref: '#/components/schemas/registry.SkillIcon' type: array uniqueItems: false - headers: + license: + description: License is the SPDX license identifier of the plugin. + type: string + metadata: + additionalProperties: {} description: |- - Headers defines HTTP headers that can be passed to the remote server for authentication - These are used with the thv proxy command's authentication features - items: - $ref: '#/components/schemas/registry.Header' + Metadata is the official metadata of the plugin as reported in the + plugin manifest file. + type: object + name: + description: |- + Name is the name of the plugin. + The format is that of identifiers, e.g. "my-plugin". + type: string + namespace: + description: |- + Namespace is the namespace of the plugin. + The format is reverse-DNS, e.g. "io.github.user". + type: string + packages: + description: Packages is the list of packages for the plugin. + items: + $ref: '#/components/schemas/registry.SkillPackage' + type: array + uniqueItems: false + repository: + $ref: '#/components/schemas/registry.SkillRepository' + status: + description: |- + Status is the status of the plugin. + Can be one of "active", "deprecated", or "archived". + type: string + title: + description: |- + Title is the title of the plugin. + This is for human consumption, not an identifier. + type: string + version: + description: |- + Version is the version of the plugin. + Any non-empty string is valid, but ideally it should be either a + semantic version or a commit hash. + type: string + type: object + registry.Provenance: + description: Provenance contains verification and signing metadata + properties: + attestation: + $ref: '#/components/schemas/registry.VerifiedAttestation' + cert_issuer: + type: string + repository_ref: + type: string + repository_uri: + type: string + runner_environment: + type: string + signer_identity: + type: string + sigstore_url: + type: string + type: object + registry.Registry: + description: Full registry data + properties: + groups: + description: Groups is a slice of group definitions containing related MCP + servers + items: + $ref: '#/components/schemas/registry.Group' + type: array + uniqueItems: false + last_updated: + description: LastUpdated is the timestamp when the registry was last updated, + in RFC3339 format + type: string + remote_servers: + additionalProperties: + $ref: '#/components/schemas/registry.RemoteServerMetadata' + description: |- + RemoteServers is a map of server names to their corresponding remote server definitions + These are MCP servers accessed via HTTP/HTTPS using the thv proxy command + type: object + servers: + additionalProperties: + $ref: '#/components/schemas/registry.ImageMetadata' + description: Servers is a map of server names to their corresponding server + definitions + type: object + version: + description: Version is the schema version of the registry + type: string + type: object + registry.RemoteServerMetadata: + description: Remote server details (if it's a remote server) + properties: + custom_metadata: + additionalProperties: {} + description: CustomMetadata allows for additional user-defined metadata + type: object + description: + description: Description is a human-readable description of the server's + purpose and functionality + type: string + env_vars: + description: |- + EnvVars defines environment variables that can be passed to configure the client + These might be needed for client-side configuration when connecting to the remote server + items: + $ref: '#/components/schemas/registry.EnvVar' + type: array + uniqueItems: false + headers: + description: |- + Headers defines HTTP headers that can be passed to the remote server for authentication + These are used with the thv proxy command's authentication features + items: + $ref: '#/components/schemas/registry.Header' type: array uniqueItems: false metadata: @@ -3833,343 +4168,1183 @@ components: predicate_type: type: string type: object - v0.ServerJSON: + remote.Config: + description: RemoteAuthConfig contains OAuth configuration for remote MCP servers properties: - $schema: - example: https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json - format: uri - minLength: 1 + authorize_url: type: string - _meta: - $ref: '#/components/schemas/v0.ServerMeta' - description: - example: MCP server providing weather data and forecasts via OpenWeatherMap - API - maxLength: 100 - minLength: 1 + bearer_token: + description: Bearer token configuration (alternative to OAuth) type: string - icons: - items: - $ref: '#/components/schemas/model.Icon' - type: array - uniqueItems: false - name: - example: io.github.user/weather - maxLength: 200 - minLength: 3 - pattern: ^[a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+$ + bearer_token_file: type: string - packages: - items: - $ref: '#/components/schemas/model.Package' - type: array - uniqueItems: false - remotes: + cached_cimd_client_id: + description: |- + CachedCIMDClientID stores the CIMD metadata URL used as client_id when CIMD + authentication was used. Kept separate from CachedClientID (which holds + DCR-issued IDs) so the two can have independent lifecycles — DCR credential + rotation clears CachedClientID without touching the stable CIMD URL. + Read by resolveClientCredentials to send the correct client_id on token refresh. + type: string + cached_client_id: + description: |- + Cached DCR client credentials for persistence across restarts. + These are obtained during Dynamic Client Registration and needed to refresh tokens. + ClientID is stored as plain text since it's public information. + type: string + cached_client_secret_ref: + type: string + cached_dcr_callback_port: + description: |- + CachedDCRCallbackPort is the callback port that was actually registered + during DCR. It may differ from CallbackPort when the requested port was + unavailable and a fallback port was selected. + type: integer + cached_refresh_token_ref: + description: |- + Cached OAuth token reference for persistence across restarts. + The refresh token is stored securely in the secret manager, and this field + contains the reference to retrieve it (e.g., "OAUTH_REFRESH_TOKEN_workload"). + This enables session restoration without requiring a new browser-based login. + type: string + cached_reg_client_uri: + description: |- + CachedRegClientURI is the registration_client_uri from the DCR response. + This is the endpoint used for RFC 7592 client read/update/delete operations. + Stored as plain text since it is not sensitive. + type: string + cached_reg_token_ref: + description: |- + CachedRegTokenRef is a secret manager reference to the registration_access_token + returned in the DCR response. Used for RFC 7592 client update operations. + Stored as a secret reference since it's sensitive. + type: string + cached_secret_expiry: + description: |- + ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server). + A zero value means the secret does not expire. + type: string + cached_token_auth_method: + description: |- + CachedTokenEndpointAuthMethod is the auth method used for the token endpoint + (e.g., "client_secret_basic", "none"). Persisted for RFC 7592 updates. + type: string + cached_token_expiry: + type: string + callback_port: + type: integer + client_id: + type: string + client_secret: + type: string + client_secret_file: + type: string + issuer: + description: OAuth endpoint configuration (from registry) + type: string + oauth_params: + additionalProperties: + type: string + description: OAuth parameters for server-specific customization + type: object + resource: + description: Resource is the OAuth 2.0 resource indicator (RFC 8707). + type: string + scope_param_name: + description: |- + ScopeParamName overrides the query parameter name used to send scopes in the + authorization URL. When empty, the standard "scope" parameter is used. + Some providers require a non-standard name (e.g., Slack uses "user_scope"). + type: string + scopes: items: - $ref: '#/components/schemas/model.Transport' + type: string type: array uniqueItems: false - repository: - $ref: '#/components/schemas/model.Repository' - title: - example: Weather API - maxLength: 100 - minLength: 1 - type: string - version: - example: 1.0.2 - maxLength: 255 - minLength: 1 + skip_browser: + type: boolean + timeout: + example: 5m type: string - websiteUrl: - example: https://modelcontextprotocol.io/examples - format: uri + token_url: type: string + use_pkce: + type: boolean type: object - v0.ServerMeta: + secrets.SecretParameter: + description: Bearer token for authentication (alternative to OAuth) properties: - io.modelcontextprotocol.registry/publisher-provided: - additionalProperties: {} - type: object + name: + type: string + target: + type: string type: object - v1.Duration: - description: |- - RefillPeriod is the duration to fully refill the bucket from zero to maxTokens. - The effective refill rate is maxTokens / refillPeriod tokens per second. - Format: Go duration string (e.g., "1m0s", "30s", "1h0m0s"). - +kubebuilder:validation:Required + storage.ACLUserRunConfig: + description: ACLUserConfig contains ACL user authentication configuration. + properties: + password_env_var: + description: PasswordEnvVar is the environment variable containing the Redis + password. + type: string + username_env_var: + description: UsernameEnvVar is the environment variable containing the Redis + username. + type: string type: object -externalDocs: - description: "" - url: "" -info: - description: This is the ToolHive API server. - title: ToolHive API - version: "1.0" + storage.RedisRunConfig: + description: RedisConfig is the Redis-specific configuration when Type is "redis". + properties: + acl_user_config: + $ref: '#/components/schemas/storage.ACLUserRunConfig' + addr: + description: |- + Addr is the Redis server address (host:port). Required for standalone and cluster modes. + Mutually exclusive with SentinelConfig. + type: string + auth_type: + description: AuthType must be "aclUser" - only ACL user authentication is + supported. + type: string + cluster_mode: + description: ClusterMode enables the Redis Cluster protocol. Requires Addr + to be set. + type: boolean + dial_timeout: + description: DialTimeout is the timeout for establishing connections (e.g., + "5s"). + type: string + key_prefix: + description: KeyPrefix for multi-tenancy, typically "thv:auth:{ns}:{name}:". + type: string + read_timeout: + description: ReadTimeout is the timeout for read operations (e.g., "3s"). + type: string + sentinel_config: + $ref: '#/components/schemas/storage.SentinelRunConfig' + sentinel_tls: + $ref: '#/components/schemas/storage.RedisTLSRunConfig' + tls: + $ref: '#/components/schemas/storage.RedisTLSRunConfig' + write_timeout: + description: WriteTimeout is the timeout for write operations (e.g., "3s"). + type: string + type: object + storage.RedisTLSRunConfig: + description: SentinelTLS configures TLS for Sentinel connections. Only applies + when SentinelConfig is set. + properties: + ca_cert_file: + description: CACertFile is the path to a PEM-encoded CA certificate file. + type: string + insecure_skip_verify: + description: InsecureSkipVerify skips certificate verification. + type: boolean + type: object + storage.RunConfig: + description: |- + Storage configures the storage backend for the auth server. + If nil, defaults to in-memory storage. + properties: + redis_config: + $ref: '#/components/schemas/storage.RedisRunConfig' + type: + description: Type specifies the storage backend type. Defaults to "memory". + type: string + type: object + storage.SentinelRunConfig: + description: |- + SentinelConfig contains Sentinel-specific configuration. + Mutually exclusive with Addr. + properties: + db: + description: 'DB is the Redis database number (default: 0).' + type: integer + master_name: + description: MasterName is the name of the Redis Sentinel master. + type: string + sentinel_addrs: + description: SentinelAddrs is the list of Sentinel addresses (host:port). + items: + type: string + type: array + uniqueItems: false + type: object + telemetry.Config: + description: |- + DEPRECATED: Middleware configuration. + TelemetryConfig contains the OpenTelemetry configuration + properties: + caCertPath: + description: |- + CACertPath is the file path to a CA certificate bundle for the OTLP endpoint. + When set, the OTLP exporters use this CA to verify the collector's TLS certificate + instead of relying solely on the system CA pool. + +optional + type: string + customAttributes: + additionalProperties: + type: string + description: |- + CustomAttributes contains custom resource attributes to be added to all telemetry signals. + These are parsed from CLI flags (--otel-custom-attributes) or environment variables + (OTEL_RESOURCE_ATTRIBUTES) as key=value pairs. + +optional + type: object + enablePrometheusMetricsPath: + description: |- + EnablePrometheusMetricsPath controls whether to expose Prometheus-style /metrics endpoint. + The metrics are served on the main transport port at /metrics. + This is separate from OTLP metrics which are sent to the Endpoint. + +kubebuilder:default=false + +optional + type: boolean + endpoint: + description: |- + Endpoint is the OTLP endpoint URL + +optional + type: string + environmentVariables: + description: |- + EnvironmentVariables is a list of environment variable names that should be + included in telemetry spans as attributes. Only variables in this list will + be read from the host machine and included in spans for observability. + Example: ["NODE_ENV", "DEPLOYMENT_ENV", "SERVICE_VERSION"] + +optional + items: + type: string + type: array + uniqueItems: false + headers: + additionalProperties: + type: string + description: |- + Headers contains authentication headers for the OTLP endpoint. + +optional + type: object + insecure: + description: |- + Insecure indicates whether to use HTTP instead of HTTPS for the OTLP endpoint. + +kubebuilder:default=false + +optional + type: boolean + metricsEnabled: + description: |- + MetricsEnabled controls whether OTLP metrics are enabled. + When false, OTLP metrics are not sent even if an endpoint is configured. + This is independent of EnablePrometheusMetricsPath. + +kubebuilder:default=false + +optional + type: boolean + samplingRate: + description: |- + SamplingRate is the trace sampling rate (0.0-1.0) as a string. + Only used when TracingEnabled is true. + Example: "0.05" for 5% sampling. + +kubebuilder:default="0.05" + +optional + type: string + serviceName: + description: |- + ServiceName is the service name for telemetry. + When omitted, defaults to the server name (e.g., VirtualMCPServer name). + +optional + type: string + serviceVersion: + description: |- + ServiceVersion is the service version for telemetry. + When omitted, defaults to the ToolHive version. + +optional + type: string + tracingEnabled: + description: |- + TracingEnabled controls whether distributed tracing is enabled. + When false, no tracer provider is created even if an endpoint is configured. + +kubebuilder:default=false + +optional + type: boolean + useLegacyAttributes: + description: |- + UseLegacyAttributes controls whether legacy (pre-MCP OTEL semconv) attribute names + are emitted alongside the new standard attribute names. When true, spans include both + old and new attribute names for backward compatibility with existing dashboards. + Currently defaults to true; this will change to false in a future release. + +kubebuilder:default=true + +optional + type: boolean + type: object + templates.RuntimeConfig: + description: |- + RuntimeConfig allows overriding the default runtime configuration + for this specific workload (base images and packages) + properties: + additional_packages: + description: |- + AdditionalPackages lists extra packages to install in the builder and + runtime stages. + Examples for Alpine: ["git", "make", "gcc"] + Examples for Debian: ["git", "build-essential"] + items: + type: string + type: array + uniqueItems: false + build_with: + description: |- + BuildWith lists build-time dependency constraints, interpreted per + package ecosystem. For uvx:// builds these are PEP 508 requirement + specifiers passed to `uv tool install --with`, used to constrain + transitive dependencies the package itself leaves unbounded + (e.g. "mcp<2"). Ecosystems without constraint support (npx://, go://) + reject a non-empty BuildWith at build time. + items: + type: string + type: array + uniqueItems: false + builder_image: + description: |- + BuilderImage is the full image reference for the builder stage. + An empty string signals "use the default for this transport type" during config merging. + Examples: "golang:1.26-alpine", "node:24-alpine", "python:3.14-slim" + type: string + runtime_env: + additionalProperties: + type: string + description: |- + RuntimeEnv contains environment variables to inject into the Dockerfile's + final runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv), + which only affects the builder stage, these variables are baked into the + shipped image and are present in the running container's process + environment at startup. Use this for values a packaged MCP server reads at + process start (e.g. feature flags, cache backend selection), not for + build-time package manager configuration. + Keys must be uppercase with underscores, values are validated for safety. + type: object + type: object + tokenexchange.Config: + description: TokenExchangeConfig contains token exchange configuration for external + authentication + properties: + audience: + description: Audience is the target audience for the exchanged token + type: string + client_id: + description: ClientID is the OAuth 2.0 client identifier + type: string + client_secret: + description: ClientSecret is the OAuth 2.0 client secret + type: string + external_token_header_name: + description: ExternalTokenHeaderName is the name of the custom header to + use when HeaderStrategy is "custom" + type: string + header_strategy: + description: |- + HeaderStrategy determines how to inject the token + Valid values: HeaderStrategyReplace (default), HeaderStrategyCustom + type: string + scopes: + description: Scopes is the list of scopes to request for the exchanged token + items: + type: string + type: array + uniqueItems: false + subject_token_type: + description: |- + SubjectTokenType specifies the type of the subject token being exchanged. + Common values: oauthproto.TokenTypeAccessToken (default), oauthproto.TokenTypeIDToken, oauthproto.TokenTypeJWT. + If empty, defaults to oauthproto.TokenTypeAccessToken. + type: string + token_url: + description: TokenURL is the OAuth 2.0 token endpoint URL + type: string + type: object + types.MiddlewareConfig: + properties: + parameters: + description: |- + Parameters is a JSON object containing the middleware parameters. + It is stored as a raw message to allow flexible parameter types. + type: object + type: + description: Type is a string representing the middleware type. + type: string + type: object + v0.ServerJSON: + properties: + $schema: + example: https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json + format: uri + minLength: 1 + type: string + _meta: + $ref: '#/components/schemas/v0.ServerMeta' + description: + example: MCP server providing weather data and forecasts via OpenWeatherMap + API + maxLength: 100 + minLength: 1 + type: string + icons: + items: + $ref: '#/components/schemas/model.Icon' + type: array + uniqueItems: false + name: + example: io.github.user/weather + maxLength: 200 + minLength: 3 + pattern: ^[a-zA-Z0-9.-]+/[a-zA-Z0-9._-]+$ + type: string + packages: + items: + $ref: '#/components/schemas/model.Package' + type: array + uniqueItems: false + remotes: + items: + $ref: '#/components/schemas/model.Transport' + type: array + uniqueItems: false + repository: + $ref: '#/components/schemas/model.Repository' + title: + example: Weather API + maxLength: 100 + minLength: 1 + type: string + version: + example: 1.0.2 + maxLength: 255 + minLength: 1 + type: string + websiteUrl: + example: https://modelcontextprotocol.io/examples + format: uri + type: string + type: object + v0.ServerMeta: + properties: + io.modelcontextprotocol.registry/publisher-provided: + additionalProperties: {} + type: object + type: object + v1.Duration: + description: |- + RefillPeriod is the duration to fully refill the bucket from zero to maxTokens. + The effective refill rate is maxTokens / refillPeriod tokens per second. + Format: Go duration string (e.g., "1m0s", "30s", "1h0m0s"). + +kubebuilder:validation:Required + type: object +externalDocs: + description: "" + url: "" +info: + description: This is the ToolHive API server. + title: ToolHive API + version: "1.0" openapi: 3.1.0 paths: /api/openapi.json: get: - description: Returns the OpenAPI specification for the API + description: Returns the OpenAPI specification for the API + responses: + "200": + content: + application/json: + schema: + type: object + description: OpenAPI specification + summary: Get OpenAPI specification + tags: + - system + /api/v1beta/clients: + get: + description: List all registered clients in ToolHive + responses: + "200": + content: + application/json: + schema: + items: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_client.RegisteredClient' + type: array + description: OK + summary: List all clients + tags: + - clients + post: + description: Register a new client with ToolHive + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.createClientRequest' + description: Client to register + summary: client + description: Client to register + required: true + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.createClientResponse' + description: OK + "400": + content: + application/json: + schema: + type: string + description: Invalid request or unsupported client type + summary: Register a new client + tags: + - clients + /api/v1beta/clients/{name}: + delete: + description: Unregister a client from ToolHive + parameters: + - description: Client name to unregister + in: path + name: name + required: true + schema: + type: string + responses: + "204": + description: No Content + "400": + content: + application/json: + schema: + type: string + description: Invalid request or unsupported client type + summary: Unregister a client + tags: + - clients + /api/v1beta/clients/{name}/groups/{group}: + delete: + description: Unregister a client from a specific group in ToolHive + parameters: + - description: Client name to unregister + in: path + name: name + required: true + schema: + type: string + - description: Group name to remove client from + in: path + name: group + required: true + schema: + type: string + responses: + "204": + description: No Content + "400": + content: + application/json: + schema: + type: string + description: Invalid request or unsupported client type + "404": + content: + application/json: + schema: + type: string + description: Client or group not found + summary: Unregister a client from a specific group + tags: + - clients + /api/v1beta/clients/register: + post: + description: Register multiple clients with ToolHive + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.bulkClientRequest' + description: Clients to register + summary: clients + description: Clients to register + required: true + responses: + "200": + content: + application/json: + schema: + items: + $ref: '#/components/schemas/pkg_api_v1.createClientResponse' + type: array + description: OK + "400": + content: + application/json: + schema: + type: string + description: Invalid request or unsupported client type + summary: Register multiple clients + tags: + - clients + /api/v1beta/clients/unregister: + post: + description: Unregister multiple clients from ToolHive + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.bulkClientRequest' + description: Clients to unregister + summary: clients + description: Clients to unregister + required: true + responses: + "204": + description: No Content + "400": + content: + application/json: + schema: + type: string + description: Invalid request or unsupported client type + summary: Unregister multiple clients + tags: + - clients + /api/v1beta/discovery/clients: + get: + description: |- + List all clients compatible with ToolHive and their status. + Each object includes supports_skills when ToolHive can install skills for that client. responses: "200": content: application/json: schema: - type: object - description: OpenAPI specification - summary: Get OpenAPI specification + $ref: '#/components/schemas/pkg_api_v1.clientStatusResponse' + description: OK + summary: List all clients status tags: - - system - /api/v1beta/clients: + - discovery + /api/v1beta/groups: get: - description: List all registered clients in ToolHive + description: Get a list of all groups + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.groupListResponse' + description: OK + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: List all groups + tags: + - groups + post: + description: Create a new group with the specified name + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.createGroupRequest' + description: Group creation request + summary: group + description: Group creation request + required: true + responses: + "201": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.createGroupResponse' + description: Created + "400": + content: + application/json: + schema: + type: string + description: Bad Request + "409": + content: + application/json: + schema: + type: string + description: Conflict + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: Create a new group + tags: + - groups + /api/v1beta/groups/{name}: + delete: + description: Delete a group by name. + parameters: + - description: Group name + in: path + name: name + required: true + schema: + type: string + - description: 'Delete all workloads in the group (default: false, moves workloads + to default group)' + in: query + name: with-workloads + schema: + type: boolean + responses: + "204": + content: + application/json: + schema: + type: string + description: No Content + "404": + content: + application/json: + schema: + type: string + description: Not Found + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: Delete a group + tags: + - groups + get: + description: Get details of a specific group + parameters: + - description: Group name + in: path + name: name + required: true + schema: + type: string + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_groups.Group' + description: OK + "404": + content: + application/json: + schema: + type: string + description: Not Found + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: Get group details + tags: + - groups + /api/v1beta/plugins: + get: + description: Get a list of all installed plugins + parameters: + - description: Filter by scope (user or project) + in: query + name: scope + schema: + enum: + - user + - project + type: string + - description: Filter by client app + in: query + name: client + schema: + type: string + - description: Filter by project root path + in: query + name: project_root + schema: + type: string + - description: Filter by group name + in: query + name: group + schema: + type: string responses: "200": content: application/json: schema: - items: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_client.RegisteredClient' - type: array - description: OK - summary: List all clients + $ref: '#/components/schemas/pkg_api_v1.pluginListResponse' + description: OK + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: List all installed plugins tags: - - clients + - plugins post: - description: Register a new client with ToolHive + description: Install a plugin from a remote source requestBody: content: application/json: schema: oneOf: - type: object - - $ref: '#/components/schemas/pkg_api_v1.createClientRequest' - description: Client to register - summary: client - description: Client to register + - $ref: '#/components/schemas/pkg_api_v1.installPluginRequest' + description: Install request + summary: request + description: Install request required: true responses: - "200": + "201": content: application/json: schema: - $ref: '#/components/schemas/pkg_api_v1.createClientResponse' - description: OK + $ref: '#/components/schemas/pkg_api_v1.installPluginResponse' + description: Created + headers: + Location: + description: URI of the installed plugin resource + schema: + type: string "400": content: application/json: schema: type: string - description: Invalid request or unsupported client type - summary: Register a new client + description: Bad Request + "401": + content: + application/json: + schema: + type: string + description: Unauthorized (registry refused credentials) + "404": + content: + application/json: + schema: + type: string + description: Not Found (artifact not present in registry) + "409": + content: + application/json: + schema: + type: string + description: Conflict + "429": + content: + application/json: + schema: + type: string + description: Too Many Requests (registry rate limit) + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + "502": + content: + application/json: + schema: + type: string + description: Bad Gateway (upstream registry failure) + "504": + content: + application/json: + schema: + type: string + description: Gateway Timeout (upstream pull timed out) + summary: Install a plugin tags: - - clients - /api/v1beta/clients/{name}: + - plugins + /api/v1beta/plugins/{name}: delete: - description: Unregister a client from ToolHive + description: Remove an installed plugin parameters: - - description: Client name to unregister + - description: Plugin name in: path name: name required: true schema: type: string + - description: Scope to uninstall from (user or project) + in: query + name: scope + schema: + enum: + - user + - project + type: string + - description: Project root path for project-scoped plugins + in: query + name: project_root + schema: + type: string responses: "204": + content: + application/json: + schema: + type: string description: No Content "400": content: application/json: schema: type: string - description: Invalid request or unsupported client type - summary: Unregister a client + description: Bad Request + "404": + content: + application/json: + schema: + type: string + description: Not Found + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: Uninstall a plugin tags: - - clients - /api/v1beta/clients/{name}/groups/{group}: - delete: - description: Unregister a client from a specific group in ToolHive + - plugins + get: + description: Get detailed information about a specific plugin parameters: - - description: Client name to unregister + - description: Plugin name in: path name: name required: true schema: type: string - - description: Group name to remove client from - in: path - name: group - required: true + - description: Filter by scope (user or project) + in: query + name: scope + schema: + enum: + - user + - project + type: string + - description: Project root path for project-scoped plugins + in: query + name: project_root schema: type: string responses: - "204": - description: No Content + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginInfo' + description: OK "400": content: application/json: schema: type: string - description: Invalid request or unsupported client type + description: Bad Request "404": content: application/json: schema: type: string - description: Client or group not found - summary: Unregister a client from a specific group + description: Not Found + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: Get plugin details tags: - - clients - /api/v1beta/clients/register: + - plugins + /api/v1beta/plugins/build: post: - description: Register multiple clients with ToolHive + description: Build a plugin from a local directory requestBody: content: application/json: schema: oneOf: - type: object - - $ref: '#/components/schemas/pkg_api_v1.bulkClientRequest' - description: Clients to register - summary: clients - description: Clients to register + - $ref: '#/components/schemas/pkg_api_v1.buildPluginRequest' + description: Build request + summary: request + description: Build request required: true responses: "200": content: application/json: schema: - items: - $ref: '#/components/schemas/pkg_api_v1.createClientResponse' - type: array + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.BuildResult' description: OK "400": content: application/json: schema: type: string - description: Invalid request or unsupported client type - summary: Register multiple clients - tags: - - clients - /api/v1beta/clients/unregister: - post: - description: Unregister multiple clients from ToolHive - requestBody: - content: - application/json: - schema: - oneOf: - - type: object - - $ref: '#/components/schemas/pkg_api_v1.bulkClientRequest' - description: Clients to unregister - summary: clients - description: Clients to unregister - required: true - responses: - "204": - description: No Content - "400": + description: Bad Request + "500": content: application/json: schema: type: string - description: Invalid request or unsupported client type - summary: Unregister multiple clients + description: Internal Server Error + summary: Build a plugin tags: - - clients - /api/v1beta/discovery/clients: + - plugins + /api/v1beta/plugins/builds: get: - description: |- - List all clients compatible with ToolHive and their status. - Each object includes supports_skills when ToolHive can install skills for that client. + description: Get a list of all locally-built OCI plugin artifacts in the local + store responses: "200": content: application/json: schema: - $ref: '#/components/schemas/pkg_api_v1.clientStatusResponse' + $ref: '#/components/schemas/pkg_api_v1.pluginBuildListResponse' description: OK - summary: List all clients status + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + summary: List locally-built plugin artifacts tags: - - discovery - /api/v1beta/groups: - get: - description: Get a list of all groups + - plugins + /api/v1beta/plugins/builds/{tag}: + delete: + description: Remove a locally-built OCI plugin artifact and its blobs from the + local store + parameters: + - description: Artifact tag + in: path + name: tag + required: true + schema: + type: string responses: - "200": + "204": content: application/json: schema: - $ref: '#/components/schemas/pkg_api_v1.groupListResponse' - description: OK + type: string + description: No Content + "404": + content: + application/json: + schema: + type: string + description: Not Found "500": content: application/json: schema: type: string description: Internal Server Error - summary: List all groups + summary: Delete a locally-built plugin artifact tags: - - groups - post: - description: Create a new group with the specified name - requestBody: - content: - application/json: - schema: - oneOf: - - type: object - - $ref: '#/components/schemas/pkg_api_v1.createGroupRequest' - description: Group creation request - summary: group - description: Group creation request + - plugins + /api/v1beta/plugins/content: + get: + description: |- + Retrieve the plugin.json body and file listing from an artifact + without installing it. Accepts OCI refs, git refs, or local tags. + parameters: + - description: OCI reference or local build tag + in: query + name: ref required: true + schema: + type: string responses: - "201": + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.PluginContent' + description: OK + "400": + content: + application/json: + schema: + type: string + description: Bad Request + "401": content: application/json: schema: - $ref: '#/components/schemas/pkg_api_v1.createGroupResponse' - description: Created - "400": + type: string + description: Unauthorized (registry refused credentials) + "404": content: application/json: schema: type: string - description: Bad Request - "409": + description: Not Found (artifact not present in registry) + "429": content: application/json: schema: type: string - description: Conflict + description: Too Many Requests (registry rate limit) "500": content: application/json: schema: type: string description: Internal Server Error - summary: Create a new group + "502": + content: + application/json: + schema: + type: string + description: Bad Gateway (upstream registry or git resolver failure) + "504": + content: + application/json: + schema: + type: string + description: Gateway Timeout (upstream pull timed out) + summary: Get plugin content tags: - - groups - /api/v1beta/groups/{name}: - delete: - description: Delete a group by name. - parameters: - - description: Group name - in: path - name: name + - plugins + /api/v1beta/plugins/push: + post: + description: Push a built plugin artifact to a remote registry + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.pushPluginRequest' + description: Push request + summary: request + description: Push request required: true - schema: - type: string - - description: 'Delete all workloads in the group (default: false, moves workloads - to default group)' - in: query - name: with-workloads - schema: - type: boolean responses: "204": content: @@ -4177,6 +5352,12 @@ paths: schema: type: string description: No Content + "400": + content: + application/json: + schema: + type: string + description: Bad Request "404": content: application/json: @@ -4189,40 +5370,45 @@ paths: schema: type: string description: Internal Server Error - summary: Delete a group + summary: Push a plugin tags: - - groups - get: - description: Get details of a specific group - parameters: - - description: Group name - in: path - name: name + - plugins + /api/v1beta/plugins/validate: + post: + description: Validate a plugin definition + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.validatePluginRequest' + description: Validate request + summary: request + description: Validate request required: true - schema: - type: string responses: "200": content: application/json: schema: - $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_groups.Group' + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_plugins.ValidationResult' description: OK - "404": + "400": content: application/json: schema: type: string - description: Not Found + description: Bad Request "500": content: application/json: schema: type: string description: Internal Server Error - summary: Get group details + summary: Validate a plugin tags: - - groups + - plugins /api/v1beta/registry: get: description: Get a list of the current registries @@ -5179,6 +6365,109 @@ paths: summary: Push a skill tags: - skills + /api/v1beta/skills/sync: + post: + description: Restore a project's installed skills to match toolhive.lock.yaml + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.syncSkillsRequest' + description: Sync request + summary: request + description: Sync request + required: true + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.SyncResult' + description: OK + "400": + content: + application/json: + schema: + type: string + description: Bad Request + "403": + content: + application/json: + schema: + type: string + description: Forbidden (feature not enabled) + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + "501": + content: + application/json: + schema: + type: string + description: Not Implemented + summary: Sync project skills from the lock file + tags: + - skills + /api/v1beta/skills/upgrade: + post: + description: Re-resolve a project's lock entries and install newer content where + available + requestBody: + content: + application/json: + schema: + oneOf: + - type: object + - $ref: '#/components/schemas/pkg_api_v1.upgradeSkillsRequest' + description: Upgrade request + summary: request + description: Upgrade request + required: true + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/github_com_stacklok_toolhive_pkg_skills.UpgradeResult' + description: OK + "400": + content: + application/json: + schema: + type: string + description: Bad Request + "403": + content: + application/json: + schema: + type: string + description: Forbidden (feature not enabled) + "404": + content: + application/json: + schema: + type: string + description: Not Found (a requested name is not in the lock file) + "500": + content: + application/json: + schema: + type: string + description: Internal Server Error + "501": + content: + application/json: + schema: + type: string + description: Not Implemented + summary: Upgrade project skills + tags: + - skills /api/v1beta/skills/validate: post: description: Validate a skill definition @@ -5916,6 +7205,104 @@ paths: summary: Get a registry server tags: - registry-servers + /registry/{registryName}/v0.1/x/dev.toolhive/plugins: + get: + description: Get a paginated list of plugins from the registry. Supports optional + full-text search and pagination. + parameters: + - description: Registry name (currently ignored, uses the default provider) + in: path + name: registryName + required: true + schema: + type: string + - description: Search filter — matches against plugin name, namespace, and description + in: query + name: q + schema: + type: string + - description: 'Page number, 1-based (default: 1)' + in: query + name: page + schema: + type: integer + - description: 'Items per page, max 200 (default: 50)' + in: query + name: limit + schema: + type: integer + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.pluginsV01Response' + description: OK + "500": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.registryErrorResponse' + description: Internal server error + "503": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.registryErrorResponse' + description: Registry authentication required or upstream registry unavailable + summary: List available registry plugins + tags: + - registry-plugins + /registry/{registryName}/v0.1/x/dev.toolhive/plugins/{namespace}/{pluginName}: + get: + description: Retrieve a single plugin by its namespace and name from the registry. + parameters: + - description: Registry name (currently ignored, uses the default provider) + in: path + name: registryName + required: true + schema: + type: string + - description: Plugin namespace in reverse-DNS format (e.g. io.github.stacklok) + in: path + name: namespace + required: true + schema: + type: string + - description: Plugin name + in: path + name: pluginName + required: true + schema: + type: string + responses: + "200": + content: + application/json: + schema: + $ref: '#/components/schemas/registry.Plugin' + description: OK + "404": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.registryErrorResponse' + description: Plugin not found + "500": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.registryErrorResponse' + description: Internal server error + "503": + content: + application/json: + schema: + $ref: '#/components/schemas/pkg_api_v1.registryErrorResponse' + description: Registry authentication required or upstream registry unavailable + summary: Get a registry plugin + tags: + - registry-plugins /registry/{registryName}/v0.1/x/dev.toolhive/skills: get: description: Get a paginated list of skills from the registry. Supports optional diff --git a/examples/operator/virtual-mcps/composite_tool_simple.yaml b/examples/operator/virtual-mcps/composite_tool_simple.yaml index 7614787dc2..5f0ced1a6b 100644 --- a/examples/operator/virtual-mcps/composite_tool_simple.yaml +++ b/examples/operator/virtual-mcps/composite_tool_simple.yaml @@ -100,6 +100,20 @@ spec: required: - url + # Explicit MCP tool annotations advertised to clients. When omitted, + # annotations are derived from the step tools at advertise time. The floor + # is fail-closed: any step whose annotations are nil/unknown (gofetch + # fetch and yardstick echo declare none) taints it conservatively to + # readOnly=false, destructive=true, openWorld=true. Declaring + # readOnlyHint: true here would CONTRADICT that floor and cause the tool + # to be DROPPED from tools/list at runtime with a warning — so we declare + # openWorldHint: true (honest: fetch makes HTTP calls) and let derivation + # fill the rest. This guardrail runs at advertise time, not at CRD + # admission; thv vmcp validate does not catch it. + annotations: + title: Fetch and Validate Data + openWorldHint: true + # Sequential workflow steps steps: # Step 1: Fetch data from URL diff --git a/examples/operator/virtual-mcps/vmcp_optimizer_quickstart.yaml b/examples/operator/virtual-mcps/vmcp_optimizer_quickstart.yaml index 35694a384c..fd3d6bfd3c 100644 --- a/examples/operator/virtual-mcps/vmcp_optimizer_quickstart.yaml +++ b/examples/operator/virtual-mcps/vmcp_optimizer_quickstart.yaml @@ -288,6 +288,9 @@ spec: --- # Step 2j: MCPServer backend - ida-pro-mcp (IDA Pro reverse engineering) +# The 1.4.0 tag is mutable and its latest rebuild resolves mcp Python SDK 2.0.0, +# which dropped mcp.server.fastmcp and makes the container crashloop. Pinned to +# the last build that resolved mcp 1.x. See test/e2e/images/images.go. apiVersion: toolhive.stacklok.dev/v1beta1 kind: MCPServer metadata: @@ -296,7 +299,7 @@ metadata: spec: groupRef: name: optimizer-services - image: ghcr.io/stacklok/dockyard/uvx/ida-pro-mcp:1.4.0 + image: ghcr.io/stacklok/dockyard/uvx/ida-pro-mcp:1.4.0@sha256:5a596d965fe05052d615a41152213f93ebc72bf46b07304718de948833d73c70 transport: stdio proxyPort: 8080 resources: diff --git a/examples/vmcp-config.yaml b/examples/vmcp-config.yaml index 7b02dea7b7..438968ad3f 100644 --- a/examples/vmcp-config.yaml +++ b/examples/vmcp-config.yaml @@ -157,6 +157,18 @@ operational: # required: ["pr_number"] # timeout: "30m" # +# # Optional: explicit MCP tool annotations advertised to clients. +# # When omitted, annotations are derived from the step tools at advertise +# # time (readOnlyHint = AND across steps, destructiveHint/openWorldHint = OR, +# # idempotentHint = never derived). Explicit values are merged over the +# # derived floor but must never contradict it (e.g. readOnlyHint: true when +# # a step is not read-only) — a contradicting composite tool is dropped +# # from tools/list with a warning at runtime. +# annotations: +# title: "Deploy and Notify" +# destructiveHint: true +# openWorldHint: true +# # steps: # - id: "merge" # tool: "github.merge_pr" diff --git a/go.mod b/go.mod index 58e389acef..2c838809f5 100644 --- a/go.mod +++ b/go.mod @@ -7,9 +7,9 @@ require ( github.com/1password/onepassword-sdk-go v0.3.1 github.com/alicebob/miniredis/v2 v2.38.0 github.com/atotto/clipboard v0.1.4 - github.com/aws/aws-sdk-go-v2 v1.42.1 - github.com/aws/aws-sdk-go-v2/config v1.32.30 - github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 + github.com/aws/aws-sdk-go-v2 v1.43.3 + github.com/aws/aws-sdk-go-v2/config v1.32.34 + github.com/aws/aws-sdk-go-v2/service/sts v1.45.3 github.com/cedar-policy/cedar-go v1.8.0 github.com/cenkalti/backoff/v5 v5.0.3 github.com/charmbracelet/bubbles v1.0.0 @@ -28,12 +28,11 @@ require ( github.com/gofrs/flock v0.13.0 github.com/google/cel-go v0.30.0 github.com/google/go-cmp v0.7.0 - github.com/google/go-containerregistry v0.21.7 + github.com/google/go-containerregistry v0.21.8 github.com/google/uuid v1.6.0 github.com/lestrrat-go/httprc/v3 v3.0.6 github.com/lestrrat-go/jwx/v3 v3.0.13 - github.com/mark3labs/mcp-go v0.55.1 - github.com/moby/moby/client v0.4.1 + github.com/moby/moby/client v0.5.1 github.com/modelcontextprotocol/registry v1.8.0 github.com/oauth2-proxy/mockoidc v0.0.0-20240214162133-caebfff84d25 github.com/olekukonko/tablewriter v1.1.4 @@ -43,30 +42,30 @@ require ( github.com/pelletier/go-toml/v2 v2.3.1 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c github.com/pressly/goose/v3 v3.27.1 - github.com/prometheus/client_golang v1.23.2 + github.com/prometheus/client_golang v1.24.1 // indirect github.com/prometheus/client_model v0.6.2 - github.com/prometheus/common v0.67.5 + github.com/prometheus/common v0.70.1 github.com/redis/go-redis/v9 v9.21.0 github.com/shirou/gopsutil/v4 v4.26.5 github.com/spf13/viper v1.21.0 - github.com/stacklok/toolhive-catalog v0.20260706.0 - github.com/stacklok/toolhive-core v0.0.28 + github.com/stacklok/toolhive-catalog v0.20260804.0 + github.com/stacklok/toolhive-core v0.0.38 github.com/stretchr/testify v1.11.1 github.com/swaggo/swag/v2 v2.0.0-rc5 - github.com/tailscale/hujson v0.0.0-20260302212456-ecc657c15afd + github.com/tailscale/hujson v0.0.0-20260727124030-b80ff77dac4f github.com/testcontainers/testcontainers-go v0.42.0 github.com/tidwall/gjson v1.18.0 github.com/xeipuuv/gojsonschema v1.2.0 github.com/zalando/go-keyring v0.2.8 - go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v1.45.0 - go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 - go.opentelemetry.io/otel/exporters/prometheus v0.66.0 + go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v1.45.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 // indirect + go.opentelemetry.io/otel/exporters/prometheus v0.67.0 // indirect go.opentelemetry.io/otel/sdk v1.45.0 go.opentelemetry.io/otel/sdk/metric v1.45.0 go.uber.org/mock v0.6.0 go.uber.org/zap v1.28.0 golang.ngrok.com/ngrok/v2 v2.1.4 - golang.org/x/exp/jsonrpc2 v0.0.0-20260611194520-c48552f49976 + golang.org/x/exp/jsonrpc2 v0.0.0-20260709172345-9ea1abe57597 golang.org/x/mod v0.41.0 golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.23.0 @@ -85,14 +84,24 @@ require github.com/getsentry/sentry-go/otel v0.44.1 require github.com/hashicorp/golang-lru/v2 v2.0.7 -require go.starlark.net v0.0.0-20260630144053-529d8e869a14 +require ( + github.com/github/smimesign v0.2.0 + go.starlark.net v0.0.0-20260630144053-529d8e869a14 +) require ( - github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.35 // indirect github.com/go-openapi/runtime/server-middleware v0.30.0 // indirect + github.com/go-openapi/swag/pools v0.27.3 // indirect + github.com/jedisct1/go-minisign v0.0.0-20241212093149-d2f9f49435c7 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/mitchellh/mapstructure v1.5.1-0.20231216201459-8508981c8b6c // indirect + github.com/modelcontextprotocol/go-sdk v1.7.0 // indirect github.com/oklog/ulid/v2 v2.1.1 // indirect - github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect + github.com/oleiade/reflections v1.1.0 // indirect + github.com/sassoftware/relic v7.2.1+incompatible // indirect + github.com/segmentio/encoding v0.5.4 // indirect + github.com/theupdateframework/go-tuf v0.7.0 // indirect github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect ) @@ -101,19 +110,19 @@ require ( github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c // indirect github.com/KyleBanks/depth v1.2.1 // indirect github.com/Masterminds/semver/v3 v3.4.0 // indirect - github.com/ProtonMail/go-crypto v1.1.6 // indirect + github.com/ProtonMail/go-crypto v1.4.1 // indirect github.com/antlr4-go/antlr/v4 v4.13.1 // indirect github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect - github.com/aws/aws-sdk-go-v2/credentials v1.19.29 // indirect - github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect - github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 // indirect - github.com/aws/aws-sdk-go-v2/service/signin v1.4.1 // indirect - github.com/aws/aws-sdk-go-v2/service/sso v1.32.1 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 // indirect - github.com/aws/smithy-go v1.27.3 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.19.33 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.34 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.34 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.34 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.15 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.34 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.5.3 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.33.3 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.3 // indirect + github.com/aws/smithy-go v1.27.6 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/blang/semver v3.5.1+incompatible // indirect @@ -135,14 +144,14 @@ require ( github.com/cyphar/filepath-securejoin v0.6.1 // indirect github.com/danieljoos/wincred v1.2.3 // indirect github.com/dgraph-io/ristretto v1.0.0 // indirect - github.com/digitorus/pkcs7 v0.0.0-20230818184609-3a137a874352 // indirect - github.com/digitorus/timestamp v0.0.0-20231217203849-220c5c2851b7 // indirect - github.com/docker/cli v29.6.0+incompatible // indirect - github.com/docker/docker-credential-helpers v0.9.3 // indirect + github.com/digitorus/pkcs7 v0.0.0-20250730155240-ffadbf3f398c // indirect + github.com/digitorus/timestamp v0.0.0-20250524132541-c45532741eea // indirect + github.com/docker/cli v29.6.2+incompatible // indirect + github.com/docker/docker-credential-helpers v0.9.5 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/dylibso/observe-sdk/go v0.0.0-20240819160327-2d926c5d788a // indirect github.com/ebitengine/purego v0.10.0 // indirect - github.com/emicklei/go-restful/v3 v3.12.2 // indirect + github.com/emicklei/go-restful/v3 v3.13.0 // indirect github.com/emirpasic/gods v1.18.1 // indirect github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/extism/go-sdk v1.7.0 // indirect @@ -152,28 +161,28 @@ require ( github.com/getsentry/sentry-go v0.47.0 github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 // indirect github.com/go-logr/zapr v1.3.0 // indirect - github.com/go-ole/go-ole v1.2.6 // indirect - github.com/go-openapi/analysis v0.25.2 // indirect + github.com/go-ole/go-ole v1.3.0 // indirect + github.com/go-openapi/analysis v0.25.5 // indirect github.com/go-openapi/errors v0.22.8 // indirect - github.com/go-openapi/jsonpointer v0.23.1 // indirect - github.com/go-openapi/jsonreference v0.21.6 // indirect - github.com/go-openapi/loads v0.24.0 // indirect - github.com/go-openapi/runtime v0.32.4 // indirect - github.com/go-openapi/spec v0.22.6 // indirect - github.com/go-openapi/strfmt v0.26.4 // indirect + github.com/go-openapi/jsonpointer v1.0.0 // indirect + github.com/go-openapi/jsonreference v1.0.0 // indirect + github.com/go-openapi/loads v0.25.0 // indirect + github.com/go-openapi/runtime v0.33.0 // indirect + github.com/go-openapi/spec v0.22.9 // indirect + github.com/go-openapi/strfmt v0.27.0 // indirect github.com/go-openapi/swag v0.26.1 // indirect - github.com/go-openapi/swag/cmdutils v0.26.1 // indirect - github.com/go-openapi/swag/conv v0.27.0 // indirect - github.com/go-openapi/swag/fileutils v0.26.1 // indirect + github.com/go-openapi/swag/cmdutils v0.27.0 // indirect + github.com/go-openapi/swag/conv v0.27.3 // indirect + github.com/go-openapi/swag/fileutils v0.27.3 // indirect github.com/go-openapi/swag/jsonname v0.26.1 // indirect - github.com/go-openapi/swag/jsonutils v0.26.1 // indirect - github.com/go-openapi/swag/loading v0.26.1 // indirect - github.com/go-openapi/swag/mangling v0.26.1 // indirect - github.com/go-openapi/swag/netutils v0.26.1 // indirect - github.com/go-openapi/swag/stringutils v0.26.1 // indirect - github.com/go-openapi/swag/typeutils v0.27.0 // indirect - github.com/go-openapi/swag/yamlutils v0.26.1 // indirect - github.com/go-openapi/validate v0.26.0 // indirect + github.com/go-openapi/swag/jsonutils v0.27.3 // indirect + github.com/go-openapi/swag/loading v0.27.3 // indirect + github.com/go-openapi/swag/mangling v0.27.3 // indirect + github.com/go-openapi/swag/netutils v0.27.0 // indirect + github.com/go-openapi/swag/stringutils v0.27.3 // indirect + github.com/go-openapi/swag/typeutils v0.27.3 // indirect + github.com/go-openapi/swag/yamlutils v0.27.3 // indirect + github.com/go-openapi/validate v0.26.1 // indirect github.com/go-task/slim-sprig/v3 v3.0.0 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect github.com/gobuffalo/pop/v6 v6.1.1 // indirect @@ -183,7 +192,7 @@ require ( github.com/golang/mock v1.7.0-rc.1 // indirect github.com/google/btree v1.1.3 // indirect github.com/google/certificate-transparency-go v1.3.3 // indirect - github.com/google/gnostic-models v0.7.0 // indirect + github.com/google/gnostic-models v0.7.1 // indirect github.com/google/jsonschema-go v0.4.3 // indirect github.com/google/pprof v0.0.0-20260402051712-545e8a4df936 // indirect github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect @@ -196,11 +205,11 @@ require ( github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect github.com/jpillora/backoff v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect - github.com/kevinburke/ssh_config v1.2.0 // indirect - github.com/klauspost/compress v1.18.7 // indirect + github.com/kevinburke/ssh_config v1.6.0 // indirect + github.com/klauspost/compress v1.19.1 // indirect github.com/lestrrat-go/option/v2 v2.0.0 // indirect github.com/lucasb-eyer/go-colorful v1.4.0 // indirect - github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/lufia/plan9stats v0.0.0-20250317134145-8bc96cf8fc35 // indirect github.com/magiconair/properties v1.8.10 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.21 // indirect @@ -209,7 +218,7 @@ require ( github.com/mattn/goveralls v0.0.12 // indirect github.com/mfridman/interpolate v0.0.2 // indirect github.com/moby/go-archive v0.2.0 // indirect - github.com/moby/moby/api v1.54.2 + github.com/moby/moby/api v1.55.0 github.com/moby/patternmatcher v0.6.1 // indirect github.com/moby/spdystream v0.5.1 // indirect github.com/moby/sys/sequential v0.6.0 // indirect @@ -232,29 +241,28 @@ require ( github.com/openzipkin/zipkin-go v0.4.3 // indirect github.com/ory/go-acc v0.2.9-0.20230103102148-6b1c9a70dbbe // indirect github.com/ory/go-convenience v0.1.0 // indirect - github.com/ory/x v0.0.665 // indirect + github.com/ory/x v0.0.665 github.com/pjbgf/sha1cd v0.6.0 // indirect github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect github.com/prometheus/otlptranslator v1.0.0 // indirect - github.com/prometheus/procfs v0.20.1 // indirect + github.com/prometheus/procfs v0.21.1 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/rivo/uniseg v0.4.7 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect - github.com/sagikazarmark/locafero v0.11.0 // indirect + github.com/sagikazarmark/locafero v0.12.0 // indirect github.com/seatgeek/logrus-gelf-formatter v0.0.0-20210414080842-5b05eb8ff761 // indirect github.com/secure-systems-lab/go-securesystemslib v0.11.0 // indirect github.com/sergi/go-diff v1.4.0 // indirect github.com/sethvargo/go-retry v0.3.0 // indirect github.com/shibumi/go-pathspec v1.3.0 // indirect - github.com/sigstore/protobuf-specs v0.5.1 // indirect + github.com/sigstore/protobuf-specs v0.5.1 github.com/sigstore/rekor v1.5.3 // indirect github.com/sigstore/rekor-tiles/v2 v2.3.0 // indirect - github.com/sigstore/sigstore v1.10.8 // indirect - github.com/sigstore/sigstore-go v1.2.2 // indirect - github.com/sigstore/timestamp-authority/v2 v2.1.2 // indirect + github.com/sigstore/sigstore v1.10.8 + github.com/sigstore/sigstore-go v1.3.0 + github.com/sigstore/timestamp-authority/v2 v2.1.3 // indirect github.com/sirupsen/logrus v1.9.4 // indirect - github.com/skeema/knownhosts v1.3.1 // indirect - github.com/sourcegraph/conc v0.3.1-0.20240121214520-5f936abd7ae8 // indirect + github.com/skeema/knownhosts v1.3.2 // indirect github.com/spf13/afero v1.15.0 // indirect github.com/spf13/cast v1.10.0 // indirect github.com/stretchr/objx v0.5.3 // indirect @@ -274,23 +282,23 @@ require ( github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect - github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 github.com/yuin/gopher-lua v1.1.1 // indirect github.com/yusufpapurcu/wmi v1.2.4 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0.46.1 // indirect - go.opentelemetry.io/contrib/propagators/b3 v1.21.0 // indirect - go.opentelemetry.io/contrib/propagators/jaeger v1.21.1 // indirect + go.opentelemetry.io/contrib/propagators/b3 v1.40.0 // indirect + go.opentelemetry.io/contrib/propagators/jaeger v1.40.0 // indirect go.opentelemetry.io/contrib/samplers/jaegerremote v0.15.1 // indirect go.opentelemetry.io/otel/exporters/jaeger v1.17.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 // indirect go.opentelemetry.io/otel/exporters/zipkin v1.45.0 // indirect - go.opentelemetry.io/proto/otlp v1.11.0 // indirect + go.opentelemetry.io/proto/otlp v1.11.0 go.uber.org/atomic v1.11.0 // indirect go.uber.org/multierr v1.11.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect golang.ngrok.com/muxado/v2 v2.0.1 // indirect - golang.org/x/exp/event v0.0.0-20260603202125-055de637280b // indirect + golang.org/x/exp/event v0.0.0-20260611194520-c48552f49976 // indirect golang.org/x/net v0.60.0 // indirect golang.org/x/text v0.42.0 // indirect golang.org/x/tools v0.49.0 // indirect @@ -299,13 +307,13 @@ require ( google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d // indirect google.golang.org/grpc v1.83.2 // indirect - google.golang.org/protobuf v1.36.11 // indirect + google.golang.org/protobuf v1.36.11 gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/warnings.v0 v0.1.2 // indirect k8s.io/apiextensions-apiserver v0.35.0 k8s.io/klog/v2 v2.140.0 // indirect - k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912 // indirect + k8s.io/kube-openapi v0.0.0-20251125145642-4e65d59e963e // indirect modernc.org/libc v1.72.1 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect diff --git a/go.sum b/go.sum index 309c5d958e..f8487279c1 100644 --- a/go.sum +++ b/go.sum @@ -49,8 +49,8 @@ github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lpr github.com/Microsoft/go-winio v0.5.2/go.mod h1:WpS1mjBmmwHBEWmogvA2mj8546UReBk4v8QkMxJ6pZY= github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= -github.com/ProtonMail/go-crypto v1.1.6 h1:ZcV+Ropw6Qn0AX9brlQLAUXfqLBc7Bl+f/DmNxpLfdw= -github.com/ProtonMail/go-crypto v1.1.6/go.mod h1:rA3QumHc/FZ8pAHreoekgiAbzpNsfQAosU5td4SnOrE= +github.com/ProtonMail/go-crypto v1.4.1 h1:9RfcZHqEQUvP8RzecWEUafnZVtEvrBVL9BiF67IQOfM= +github.com/ProtonMail/go-crypto v1.4.1/go.mod h1:e1OaTyu5SYVrO9gKOEhTc+5UcXtTUa+P3uLudwcgPqo= github.com/adrg/xdg v0.5.3 h1:xRnxJXne7+oWDatRhR1JLnvuccuIeCoBu2rtuLqQB78= github.com/adrg/xdg v0.5.3/go.mod h1:nlTsY+NNiCBGCK2tpm09vRqfVzrc2fLmXGpBLF0zlTQ= github.com/alicebob/miniredis/v2 v2.38.0 h1:nZAzCR+Lj+Vxk4ZXzm2NuKq2O33RXj1XxJ2e2uP9jiw= @@ -65,36 +65,36 @@ github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3d github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw= github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= -github.com/aws/aws-sdk-go-v2 v1.42.1 h1:9eOTgu1z/dVtYpNZ3/8/XbbaX0x/BqE3HUzAzs6K0ek= -github.com/aws/aws-sdk-go-v2 v1.42.1/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= -github.com/aws/aws-sdk-go-v2/config v1.32.30 h1:XwsEzpTJfQYJbFicz/QMLwAZdyeNVVoOEkbF7R3gPJk= -github.com/aws/aws-sdk-go-v2/config v1.32.30/go.mod h1:Ud32SuMc+/9BGxfpSVld7HrE2o05JwKmXY4M3jOQNZU= -github.com/aws/aws-sdk-go-v2/credentials v1.19.29 h1:WHZGssHH887cO0ox07SIQZsFx3MKD4ps6w0xUEmnKYQ= -github.com/aws/aws-sdk-go-v2/credentials v1.19.29/go.mod h1:Mhl0xR6zjguiuj00XRx2wMx22sAltk7oya39sT7fdg8= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 h1:/hi1JADLEW9YYryEz1w4GQu0EtP23pP553Cf9KgsDV4= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30/go.mod h1:/3AOgy4K17Dm4ucMZVC/MJkzy5kmfKUcINRHZyo0koQ= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 h1:xM/Is9cKMHa8Jj8zkvWhvrFkZsXJV9E+BB4g0HW0duQ= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30/go.mod h1:WueJeNDZvK1fMYEWJIkcivBfEzUkTpBhzlrUKKY8EuA= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 h1:jn46zC9LdsVR/ZpMIJqMqb8hHv31BlLx3ulVqNspUOk= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30/go.mod h1:1hTMsAgbdS/AtUi4bw8+gUuh1pceo+eXRLfpSuSQj3M= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 h1:3GUprIsfmGcC5SACIyB0e7E0BM1O1b3Erl5CePYIAeQ= -github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31/go.mod h1:7PuV1yl5e2xnUbm+RqvVg5i2iBM8EyijZNoI9wsOoOc= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur/BiHK6SKPjoBIXSE/hJ6g6JGRLuxQy1jGjlN4= -github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13/go.mod h1:ITg9em2KbJx1s0y4aqRX5OYWG6HBZ5TVR//OdpEZ2CQ= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30 h1:/Z5jmNrKsSD7EmDjzAPsm/3L9IuOkzaynklJZ1qX7S4= -github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.30/go.mod h1:lEzEZnOosE7zi8Z6royW1cFJTD9fpab4Ul1SBrllewk= +github.com/aws/aws-sdk-go-v2 v1.43.3 h1:XJIcfv8uDs2ukdQsoAC8/Ebu1ejxwzlayl2ZsiFns2A= +github.com/aws/aws-sdk-go-v2 v1.43.3/go.mod h1:70vwSy16txshwG+g55WkpgPKDIByzHI8ccBsOteo3bQ= +github.com/aws/aws-sdk-go-v2/config v1.32.34 h1:o+YAizrX562nEZXaB38uYTK8RvIsvW0uuRP+e5e0Pfk= +github.com/aws/aws-sdk-go-v2/config v1.32.34/go.mod h1:wc0zYRChOniiufvdWiRVf3jgXSgbkvaD683IHHHc2ZQ= +github.com/aws/aws-sdk-go-v2/credentials v1.19.33 h1:/e5V3EWfeDiW6cuRxHsC8gbwko4/vvVYPJR2afBKFFY= +github.com/aws/aws-sdk-go-v2/credentials v1.19.33/go.mod h1:ZxAmkcyOM9beY/WO9oxp2oVPXiP3rq5N1/p4NbenJdE= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.34 h1:1EsGke6rTD2CG3j2MMVB77n6Q+FlbQWYI/dFdLWBNtM= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.34/go.mod h1:5B1Z/QbaWzqoWRzYxZfmCbDDRcvUHcfAIQw/S+KfDmc= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.34 h1:vuIfjzoeqhQMGJyOBU3t0ZEjn2jrN8Bbg1N4CgjzM5Q= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.34/go.mod h1:hP28cN4CPJLZHirdQPrZR50JcLN4ApRJP2tzG8cRlhY= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.34 h1:9faHsnqxJ1vDvB4wMZy/ajIDyz5QhllQjjc72RJpXAw= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.34/go.mod h1:Yp6nIyejpa23nzlB/LhT63KTla9Jdi06nv/HH/OkAH8= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.35 h1:Oe8gMKJLO5awqpa5EhAGKVnBv1s+brdWVuxM2mDa7zA= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.35/go.mod h1:FZevcG9cOST/FWAAUhHIchjR9fXFXFRCWodOhx+PDLA= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.15 h1:JJLBQxwY+AFwuPAi5ivGc1ChnTdUt4cXMv7e76m2c/Y= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.15/go.mod h1:lQknBIe78MVL0cQOQDlag8KGflMbMEVFx9mB6O8ENvk= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.34 h1:sYg4qHWLqsjp15PzX7XCOHSOgKEGoZ5vQY43VvZ1pas= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.34/go.mod h1:N58SSz3roKf1HzW5qRaOiyk6MbDLTKgLPvlTfJ90iyI= github.com/aws/aws-sdk-go-v2/service/kms v1.52.0 h1:QNtg+Mtj1zmepk568+UKBD5DFfqh+ESTUUqQT27JkQc= github.com/aws/aws-sdk-go-v2/service/kms v1.52.0/go.mod h1:Y0+uxvxz6ib4KktRdK0V4X45Vcs/JyYoz8H71pO8xeI= -github.com/aws/aws-sdk-go-v2/service/signin v1.4.1 h1:V7ZZ300WPXGjvkyore5DGe0ljVPOxCXie/thWdtSBXE= -github.com/aws/aws-sdk-go-v2/service/signin v1.4.1/go.mod h1:mxC0nT/C8wMMS97DemZPzvUZxvIt+2Iq+eS3JdFZGgg= -github.com/aws/aws-sdk-go-v2/service/sso v1.32.1 h1:gYFYh4iLLcAOJRLNPY2aD2g9DIhKn4eof8UkIrr1rTk= -github.com/aws/aws-sdk-go-v2/service/sso v1.32.1/go.mod h1:u8af9Nqkmqnr96f7v9nHqzZT9XBwbXEkTiqT4ROuJSE= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1 h1:arjT9Cm3/WYbGmD5TUZHk4UQn4Lle1fUNZs5FC6CtF0= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.37.1/go.mod h1:DMPWJBjYs6+3+f/qhBFEFPPlQ6NlhWjai3dJNvipJ84= -github.com/aws/aws-sdk-go-v2/service/sts v1.44.1 h1:RvfHDg+xvAeZ+5741vUEjpOVtYSIm93W2zhx10Xtydw= -github.com/aws/aws-sdk-go-v2/service/sts v1.44.1/go.mod h1:9gdl4RrflIdpDb2TlXshWgR1F9TeCkvqDx77Vpr4Z/Q= -github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= -github.com/aws/smithy-go v1.27.3/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.3 h1:togAtAmgV5IGMnQDuBDJeM8z5Y5RN6G7xeOgphWz+Yc= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.3/go.mod h1:T7xKUUUvN7W3RW8UmMvKnD12xqh+Ux2gCPHPhnt64Dg= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.3 h1:YjH64OUytnWZBHUtM9GMyi4ZWBiSQdEJkZuPykOIe44= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.3/go.mod h1:5qoHcDZDTSJotoKk1bvVRPv1MXaL/NhfY9ng8D1g/ig= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.3 h1:A4o1di/XGaqtw6r3toSBrFX2U7mVSLqg7jo9wL4I+cU= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.3/go.mod h1:sKuKz2kHtrGVtFu34vbM3LWSA9CKD9YZUmm6e5PPqRA= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.3 h1:Fi7+DiKN1+QphlajvE6FqeZ8GRbnnRul7zTdUiRpbGc= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.3/go.mod h1:KCc3e27fHZUGtzpek7wZcp6dyCpGkJJo/+3PBujh/yU= +github.com/aws/smithy-go v1.27.6 h1:0zjT8jgK3jbrTT7JJ3EE6JsMhX8JTrZ+f1sEndYDXrA= +github.com/aws/smithy-go v1.27.6/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= @@ -112,6 +112,7 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= +github.com/certifi/gocertifi v0.0.0-20180118203423-deb3ae2ef261/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc= @@ -177,18 +178,16 @@ github.com/dgraph-io/ristretto v1.0.0/go.mod h1:jTi2FiYEhQ1NsMmA7DeBykizjOuY88Nh github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 h1:fAjc9m62+UWV/WAFKLNi6ZS0675eEUC9y3AlwSbQu1Y= github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/digitorus/pkcs7 v0.0.0-20230713084857-e76b763bdc49/go.mod h1:SKVExuS+vpu2l9IoOc0RwqE7NYnb0JlcFHFnEJkVDzc= -github.com/digitorus/pkcs7 v0.0.0-20230818184609-3a137a874352 h1:ge14PCmCvPjpMQMIAH7uKg0lrtNSOdpYsRXlwk3QbaE= -github.com/digitorus/pkcs7 v0.0.0-20230818184609-3a137a874352/go.mod h1:SKVExuS+vpu2l9IoOc0RwqE7NYnb0JlcFHFnEJkVDzc= -github.com/digitorus/timestamp v0.0.0-20231217203849-220c5c2851b7 h1:lxmTCgmHE1GUYL7P0MlNa00M67axePTq+9nBSGddR8I= -github.com/digitorus/timestamp v0.0.0-20231217203849-220c5c2851b7/go.mod h1:GvWntX9qiTlOud0WkQ6ewFm0LPy5JUR1Xo0Ngbd1w6Y= +github.com/digitorus/pkcs7 v0.0.0-20250730155240-ffadbf3f398c h1:g349iS+CtAvba7i0Ee9EP1TlTZ9w+UncBY6HSmsFZa0= +github.com/digitorus/pkcs7 v0.0.0-20250730155240-ffadbf3f398c/go.mod h1:mCGGmWkOQvEuLdIRfPIpXViBfpWto4AhwtJlAvo62SQ= +github.com/digitorus/timestamp v0.0.0-20250524132541-c45532741eea h1:ALRwvjsSP53QmnN3Bcj0NpR8SsFLnskny/EIMebAk1c= +github.com/digitorus/timestamp v0.0.0-20250524132541-c45532741eea/go.mod h1:GvWntX9qiTlOud0WkQ6ewFm0LPy5JUR1Xo0Ngbd1w6Y= github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= -github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI= -github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= -github.com/docker/cli v29.6.0+incompatible h1:nw9himxMMZ7eIeherJNlKQq+acnlzGgHd+4uf10QRSc= -github.com/docker/cli v29.6.0+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= -github.com/docker/docker-credential-helpers v0.9.3 h1:gAm/VtF9wgqJMoxzT3Gj5p4AqIjCBS4wrsOh9yRqcz8= -github.com/docker/docker-credential-helpers v0.9.3/go.mod h1:x+4Gbw9aGmChi3qTLZj8Dfn0TD20M/fuWy0E5+WDeCo= +github.com/docker/cli v29.6.2+incompatible h1:/bjePvcbbFTnRrMfWJBY7AjfICdsiLVgHn6LwTVOcqw= +github.com/docker/cli v29.6.2+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= +github.com/docker/docker-credential-helpers v0.9.5 h1:EFNN8DHvaiK8zVqFA2DT6BjXE0GzfLOZ38ggPTKePkY= +github.com/docker/docker-credential-helpers v0.9.5/go.mod h1:v1S+hepowrQXITkEfw6o4+BMbGot02wiKpzWhGUZK6c= github.com/docker/go-connections v0.7.0 h1:6SsRfJddP22WMrCkj19x9WKjEDTB+ahsdiGYf0mN39c= github.com/docker/go-connections v0.7.0/go.mod h1:no1qkHdjq7kLMGUXYAduOhYPSJxxvgWBh7ogVvptn3Q= github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= @@ -201,8 +200,8 @@ github.com/ebitengine/purego v0.10.0 h1:QIw4xfpWT6GWTzaW5XEKy3HXoqrJGx1ijYHzTF0/ github.com/ebitengine/purego v0.10.0/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= github.com/elazarl/goproxy v1.7.2 h1:Y2o6urb7Eule09PjlhQRGNsqRfPmYI3KKQLFpCAV3+o= github.com/elazarl/goproxy v1.7.2/go.mod h1:82vkLNir0ALaW14Rc399OTTjyNREgmdL2cVoIbS6XaE= -github.com/emicklei/go-restful/v3 v3.12.2 h1:DhwDP0vY3k8ZzE0RunuJy8GhNpPL6zqLkDf9B/a0/xU= -github.com/emicklei/go-restful/v3 v3.12.2/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= +github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes= +github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= @@ -229,6 +228,8 @@ github.com/getsentry/sentry-go v0.47.0 h1:AnSMSyrYA5qZCIN/2xpgAAwv63sVULV+vBq37a github.com/getsentry/sentry-go v0.47.0/go.mod h1:h+b4VHpKnK7aUXB5wc+KDnPgp9ZtfliRD4eV85FbiSA= github.com/getsentry/sentry-go/otel v0.44.1 h1:RV2zUHEvGHJmCCpMaJ52tZZAlcbMgvtasQn/g3CcKKc= github.com/getsentry/sentry-go/otel v0.44.1/go.mod h1:CfzTxocQJ6JX4SLFvnBrGULBAARFAd1fHmbJCTQlOP4= +github.com/github/smimesign v0.2.0 h1:Hho4YcX5N1I9XNqhq0fNx0Sts8MhLonHd+HRXVGNjvk= +github.com/github/smimesign v0.2.0/go.mod h1:iZiiwNT4HbtGRVqCQu7uJPEZCuEE5sfSSttcnePkDl4= github.com/gkampitakis/ciinfo v0.3.2 h1:JcuOPk8ZU7nZQjdUhctuhQofk7BGHuIy0c9Ez8BNhXs= github.com/gkampitakis/ciinfo v0.3.2/go.mod h1:1NIwaOcFChN4fa/B0hEBdAb6npDlFL8Bwx4dfRLRqAo= github.com/gkampitakis/go-diff v1.3.2 h1:Qyn0J9XJSDTgnsgHRdz9Zp24RaJeKMUHg2+PDZZdC4M= @@ -262,58 +263,63 @@ github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/go-logr/zapr v1.3.0 h1:XGdV8XW8zdwFiwOA2Dryh1gj2KRQyOOoNmBy4EplIcQ= github.com/go-logr/zapr v1.3.0/go.mod h1:YKepepNBd1u/oyhd/yQmtjVXmm9uML4IXUgMOwR8/Gg= -github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= -github.com/go-openapi/analysis v0.25.2 h1:I0vy4n3alz+DHTiN1PRhCb7QZxkK6g5YmswZKv2TKuw= -github.com/go-openapi/analysis v0.25.2/go.mod h1:Uhs1t/2XR10EnwONYILGEzw8gcfGIG5Xk5K2AxnhqDo= +github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE= +github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78= +github.com/go-openapi/analysis v0.25.5 h1:xPYEvTb90o1y0epuiOPAoG4QqahjP3cdp5xNlHeKJRI= +github.com/go-openapi/analysis v0.25.5/go.mod h1:d3UGtQC5uq5Kqqqis2VH09Km/v3vwsWrYkbp4gdm+Rc= github.com/go-openapi/errors v0.22.8 h1:oP7sW7TWc3wFFjrzzj0nI83H2qMBkNjNfSd+XRejk/I= github.com/go-openapi/errors v0.22.8/go.mod h1:BuUoHcYrU6E7V9gfj1I5wLQqgtIHnup/alXZ8KdgQ0w= -github.com/go-openapi/jsonpointer v0.23.1 h1:1HBACs7XIwR2RcmItfdSFlALhGbe6S92p0ry4d1GWg4= -github.com/go-openapi/jsonpointer v0.23.1/go.mod h1:iWRmZTrGn7XwYhtPt/fvdSFj1OfNBngqRT2UG3BxSqY= -github.com/go-openapi/jsonreference v0.21.6 h1:NZ5nGfnaM1n4I43Xjm1e5/M2GjOwQwndQz22uhxwD+Y= -github.com/go-openapi/jsonreference v0.21.6/go.mod h1:xzbgtQ3ZbWxvET3AxdzCJlJt6vkovbf+IfSPJjD0tUY= -github.com/go-openapi/loads v0.24.0 h1:4LLorXRPTzIN9V6ngMUZbAscsBOUBk3Oa8cClu/bFrQ= -github.com/go-openapi/loads v0.24.0/go.mod h1:xQMgX+hw5xRAhGrcDXxeMw78IFqUpIzhleu3HqPhyF4= -github.com/go-openapi/runtime v0.32.4 h1:8ElGj/3goG0itt0nBPP6Cm57ehcYyuHoI3O20nxgvkw= -github.com/go-openapi/runtime v0.32.4/go.mod h1:Bz6keOZw1NX4T6f+m42OoT1MBPDt6Re13dbccHyGH/4= +github.com/go-openapi/jsonpointer v1.0.0 h1:kR9tHqY0CtZaOPVFm622dPVNhrvYpwr4uCxgL3h1H8s= +github.com/go-openapi/jsonpointer v1.0.0/go.mod h1:Z3rw7dWu1p9IgitXCFamSlA5lmDiklEB6vkaxcNZW5Y= +github.com/go-openapi/jsonreference v1.0.0 h1:jlmTr6torcd1YgDQvSfNmRtKzYDO4FGBkrAdlAVWnpY= +github.com/go-openapi/jsonreference v1.0.0/go.mod h1:jtwdyGbJk0Xhe5Y+rwtglQP6Sb1WZST4rT32LWB+sv0= +github.com/go-openapi/loads v0.25.0 h1:74Bc2snfaVlsHzwdQj/3gsA9XJz3daXTJVs+4ZaK7jI= +github.com/go-openapi/loads v0.25.0/go.mod h1:JFBw4SIB9+PTIFHDfcXuSSy5h6aWzjtUCrPYyx3qWU8= +github.com/go-openapi/runtime v0.33.0 h1:Dd3Oj2ig+WH8ckK95l0Wn2V8a4bH/UqWPRZVT0vc8yU= +github.com/go-openapi/runtime v0.33.0/go.mod h1:+rsupH3+TFKqmFysqkmgBOTxpVJV8eV+j9myvvea2Xw= github.com/go-openapi/runtime/server-middleware v0.30.0 h1:8rPoJ/xv7JL8BsovaqboKETlpWBArVh8n+0L/GyePog= github.com/go-openapi/runtime/server-middleware v0.30.0/go.mod h1:OYNT/TxNvB/VK5oe4htM2jDTwlEXuejVJmu0DVZfAMs= -github.com/go-openapi/spec v0.22.6 h1:Tyy1pLaNCM8GBCFLoGYLonjJi6zykqyLCjXLc19ZPic= -github.com/go-openapi/spec v0.22.6/go.mod h1:HZvTHat+iH0PALQRWhrqIHtU/PEqxqd89fu0MxGlMeM= -github.com/go-openapi/strfmt v0.26.4 h1:yI6IAEfcWow459BD5UzFY430KUwXZwBHrYusPFkhWlc= -github.com/go-openapi/strfmt v0.26.4/go.mod h1:hNJi6nb5ETD6i7A1yRo03M9S6ZoTPPoWff1iUexmfUc= +github.com/go-openapi/spec v0.22.9 h1:/vKIFDcGKp0ktZWGbym/tJEWbk6/XOEmAVU0kqKMH+w= +github.com/go-openapi/spec v0.22.9/go.mod h1:b/mNUYIOQOyIiUzUzXEE8xzyZqf93KvM9hQGP91yfl0= +github.com/go-openapi/strfmt v0.27.0 h1:kbcTeaD9TXuXD0hhMXzuYa1sdTo6+dWGvwjW93E80IM= +github.com/go-openapi/strfmt v0.27.0/go.mod h1:s/qhDqfY72irigXUGJmtgid2Rm+3tnz3k8hZaRmvWYc= github.com/go-openapi/swag v0.26.1 h1:l5sVEyVpwj+DDYeZyo7wQI/Ebn/mKYIyGB/pFwAfGoQ= github.com/go-openapi/swag v0.26.1/go.mod h1:yNY38BbIVthxbkDtq1UHBCGasBqjakW3lCR6ANzdBEw= -github.com/go-openapi/swag/cmdutils v0.26.1 h1:f2iE1ijYaJ3nuu5PaEMx3zpEhzhZFgivCJObWEObLIQ= -github.com/go-openapi/swag/cmdutils v0.26.1/go.mod h1:Sm1MVFMkF6guJJ+pQqHnQA3N0j9qALV3NxzDSv6bETM= -github.com/go-openapi/swag/conv v0.27.0 h1:EKOH4feXrvdo8DbSsXSAqRT8fz1epEnS5O2IfXUOzE8= -github.com/go-openapi/swag/conv v0.27.0/go.mod h1:pfiv0uKQTbaGApk8Zs/lZV3uSjmSpa2FO1y183YngN8= -github.com/go-openapi/swag/fileutils v0.26.1 h1:K1XCM2CGhfNsc6YDt6v7Q5+1e59rftYWdcu/isZhvFw= -github.com/go-openapi/swag/fileutils v0.26.1/go.mod h1:mYUgxQAKX4ShS3qvvySx+/9yrlUnDhjiD1CalaQl8lQ= +github.com/go-openapi/swag/cmdutils v0.27.0 h1:aIKiqhB29AaP+7xm8/CPg3uOpeHx2SUp6TvMpu/a31Y= +github.com/go-openapi/swag/cmdutils v0.27.0/go.mod h1:Sm1MVFMkF6guJJ+pQqHnQA3N0j9qALV3NxzDSv6bETM= +github.com/go-openapi/swag/conv v0.27.3 h1:iqJFmGEjmX3AY0lSszABFqRVqOSt99XS0LzNIMJYuhU= +github.com/go-openapi/swag/conv v0.27.3/go.mod h1:nPRmN6jgNme99hpf+nM0auDZGALWIqlwhisKPK/bQhQ= +github.com/go-openapi/swag/fileutils v0.27.3 h1:3UVoZ2RLaIs1lt+2jcKzL8RM3Yk0rmsDE9FLA/HGxFE= +github.com/go-openapi/swag/fileutils v0.27.3/go.mod h1:VvJFZLTZS0AI854gEQz5tk7dBESdLjiNUMSZ/th2ry8= github.com/go-openapi/swag/jsonname v0.26.1 h1:VReupaV6WxlAsCn0e4DUfgV6bPmINnPpyJDLqSfNPcE= github.com/go-openapi/swag/jsonname v0.26.1/go.mod h1:OvdW6BoWoj33pTfi7x9vFrgmT+fk7aw0BRwvCE0YOuc= -github.com/go-openapi/swag/jsonutils v0.26.1 h1:2hdBfFkHg+7Wrz2VsCbeyR6hzkRDs7AztnMR2u84yOY= -github.com/go-openapi/swag/jsonutils v0.26.1/go.mod h1:U+RMJH3wa+6BRiphuRtIyI8fW9HPFqFQ4sHk2oRx0UQ= -github.com/go-openapi/swag/jsonutils/fixtures_test v0.26.1 h1:1CD7NiLLb/TXl3tOnFYU4b+mNfb5rtgHkaA+q7RMYYQ= -github.com/go-openapi/swag/jsonutils/fixtures_test v0.26.1/go.mod h1:ZWafc8nMdYzTE3uYY6W86f0n46+IF0g4uUyRhJw/kXc= -github.com/go-openapi/swag/loading v0.26.1 h1:E9K4wqXeROlhjFQ13K9zMz6ojFGXIggGe+ad1odrK9w= -github.com/go-openapi/swag/loading v0.26.1/go.mod h1:3qvRIlWzWdq1HvmldwmuJ2ohpcAryN6xVt2OTKd0/7E= -github.com/go-openapi/swag/mangling v0.26.1 h1:gpYI4WuPKFJJVjV5cDLGlDVJhFIxYjQc7yN5eEb4CqM= -github.com/go-openapi/swag/mangling v0.26.1/go.mod h1:POETDH01hqAdASXfw7ISEd9bCOE6xBHOt8NHmGZRmYM= -github.com/go-openapi/swag/netutils v0.26.1 h1:BNctoc39WTAUMxyAs355fExOPzMZtPbZ0ZZ1Am2FR5M= -github.com/go-openapi/swag/netutils v0.26.1/go.mod h1:y02vByhZhQPAVwOX+0KipXFZ/hUbk6G/Enhf5rGaOkQ= -github.com/go-openapi/swag/stringutils v0.26.1 h1:f88uYyTso7TnHrKM/bUBsQ5e2wKf37cpgo6pvbzd9yU= -github.com/go-openapi/swag/stringutils v0.26.1/go.mod h1:Sc6d3bU8fgk5AyZR8/8jEQ+Is/Ald+TD/IIggPN8UJk= -github.com/go-openapi/swag/typeutils v0.27.0 h1:aCf4MSGo8NLwZP8Q6t32DWLJSvl/WwNqgmEG+xJ6v2o= -github.com/go-openapi/swag/typeutils v0.27.0/go.mod h1:Srm0xFNRZ1Y+vCxJclo5qzx8aj+1pAKda/YfFPrG0dQ= -github.com/go-openapi/swag/yamlutils v0.26.1 h1:0TSLK+lXs9vfIhAWzBeI/lOzEnIoot6WTCO1aAeWFTk= -github.com/go-openapi/swag/yamlutils v0.26.1/go.mod h1:7W5b7PRX9MxwL7TjeG7H8HkyBGRsIDRObhyMWFgBI2M= -github.com/go-openapi/testify/enable/yaml/v2 v2.5.1 h1:q9NtHwK4qHF7yZziBPvZyv7zWAIk8ok88Gh2mR6Jpc8= -github.com/go-openapi/testify/enable/yaml/v2 v2.5.1/go.mod h1:JW0MXIotCYps/XsgJnG3a8Q7rE5xAiBwoOD5OfaIQBk= +github.com/go-openapi/swag/jsonutils v0.27.3 h1:1DEz+O82frtSMBcos/7XIn1GnpNTbsD4Bru4Dc/uhRc= +github.com/go-openapi/swag/jsonutils v0.27.3/go.mod h1:qiDCoQvzkMxrV3G8FLEdIU5L+EFYc0zcDOHWT3Yofvo= +github.com/go-openapi/swag/jsonutils/fixtures_test v0.27.3 h1:h/eT9kmGCDdFLJF29lOhzLtF0FmP1AX2MhLJWVebsb8= +github.com/go-openapi/swag/jsonutils/fixtures_test v0.27.3/go.mod h1:mofwUWx70wvskwESqRJ//k/9kURmCgyJl5m5Ppoh5kY= +github.com/go-openapi/swag/loading v0.27.3 h1:L9nQkEgzU7QgFQL+pLEMfGUKxeM4pWwGwbET9Z3weW0= +github.com/go-openapi/swag/loading v0.27.3/go.mod h1:rJ0NeaKsF4CVPnMGjPQl7JlSHzvD0bc2DKXLss1hiuE= +github.com/go-openapi/swag/mangling v0.27.3 h1:gRzzD1PAUoLTtGMgI3KpBmCSOlTuLTFWnviLxLcTnyg= +github.com/go-openapi/swag/mangling v0.27.3/go.mod h1:jtBE2+V+3pILxOR7Vgce+Cwp6A2PgZbvVqfNntbVs0w= +github.com/go-openapi/swag/netutils v0.27.0 h1:lEUG+hHvPvLggB3A8snFk0IRKNf9uC0YKc+7WYqvAF8= +github.com/go-openapi/swag/netutils v0.27.0/go.mod h1:J+WYyFMLtvtCGqa6jLv+YNUmIKI3ZRQRrvfNDMoQoEQ= +github.com/go-openapi/swag/pools v0.27.3 h1:gXjImP3F6/56wRRcFgEPld084Y6u2gs21ikPBt8NKBk= +github.com/go-openapi/swag/pools v0.27.3/go.mod h1:kVQefhSK5RWuRe7BXsL8htgBPAMpN7HDGpGEknqugeE= +github.com/go-openapi/swag/stringutils v0.27.3 h1:Ru28hnbAvN5wycALQYy8IobHvASq+FUFMlp1QzLM0JI= +github.com/go-openapi/swag/stringutils v0.27.3/go.mod h1:lzRN95CxXmA03XcDWHLOb6nOMcxCqR5rGY0lOgsfRoM= +github.com/go-openapi/swag/typeutils v0.27.3 h1:l6SSrx5eR5/WVwrGNzN6bQ9WqL04mrxNBl9YgQ3rcJ4= +github.com/go-openapi/swag/typeutils v0.27.3/go.mod h1:Srm0xFNRZ1Y+vCxJclo5qzx8aj+1pAKda/YfFPrG0dQ= +github.com/go-openapi/swag/yamlutils v0.27.3 h1:cRFCAoYtslYn9L9T0xWryHy1t7c1MACC+DMj3CLvwvs= +github.com/go-openapi/swag/yamlutils v0.27.3/go.mod h1:6JYBGj8sw/NawMllyZY+cTA8Mzk2etS3ZBASdcyPsiU= +github.com/go-openapi/testify/enable/yaml/v2 v2.6.0 h1:gGHwAJ0R/5jU8BEGDbfRNR3hL68dAVi84WuOApp29B0= +github.com/go-openapi/testify/enable/yaml/v2 v2.6.0/go.mod h1:tY+St1SGq4NFl0QIqdTY4aEdbChAHxhyB77XQi9iJCo= github.com/go-openapi/testify/v2 v2.6.0 h1:5PKH2HE7YJ/LuRPQGvSxBRlFXNQhSetBLlGAgUEu3ug= github.com/go-openapi/testify/v2 v2.6.0/go.mod h1:SgsVHtfooshd0tublTtJ50FPKhujf47YRqauXXOUxfw= -github.com/go-openapi/validate v0.26.0 h1:dxWzQ3F+vb1SajqUxHjwb5T4mTpSHmdrtv5Bi7+ZNhw= -github.com/go-openapi/validate v0.26.0/go.mod h1:b4o00uq7fJeJA+wWhVFCJpKTctzeFwzZImGGmHsl2JA= +github.com/go-openapi/validate v0.26.1 h1:pZSbvtRO8G2R2FpWTYRn3w8LrsNwbtaVhP2dWiBa0Us= +github.com/go-openapi/validate v0.26.1/go.mod h1:B8UMgXiQiwwQWIbmuROlwJZDPGlikPuh7iHV1vPX9Oo= +github.com/go-rod/rod v0.116.2 h1:A5t2Ky2A+5eD/ZJQr1EfsQSe5rms5Xof/qj296e+ZqA= +github.com/go-rod/rod v0.116.2/go.mod h1:H+CMO9SCNc2TJ2WfrG+pKhITz57uGNYU43qYHh438Mg= github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg= github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= @@ -369,14 +375,13 @@ github.com/google/cel-go v0.30.0 h1:ll54AkzKunWkBn9wSoiUXbFZXYZTkdJGNXTBXUoolGo= github.com/google/cel-go v0.30.0/go.mod h1:X0bD6iVNR8pkROSOoHVdgTkzmRcosof7WQqCD6wcMc8= github.com/google/certificate-transparency-go v1.3.3 h1:hq/rSxztSkXN2tx/3jQqF6Xc0O565UQPdHrOWvZwybo= github.com/google/certificate-transparency-go v1.3.3/go.mod h1:iR17ZgSaXRzSa5qvjFl8TnVD5h8ky2JMVio+dzoKMgA= -github.com/google/gnostic-models v0.7.0 h1:qwTtogB15McXDaNqTZdzPJRHvaVJlAl+HVQnLmJEJxo= -github.com/google/gnostic-models v0.7.0/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= -github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/gnostic-models v0.7.1 h1:SisTfuFKJSKM5CPZkffwi6coztzzeYUhc3v4yxLWH8c= +github.com/google/gnostic-models v0.7.1/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 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/go-containerregistry v0.21.7 h1:/vPFuVXDjtFREsVArW+0h1CIl5urnOhzei4X2DMW9IU= -github.com/google/go-containerregistry v0.21.7/go.mod h1:kjSbt7/zMsKLWfnHrIvKvhXHUw91jbe9DNjPPJ32gXE= +github.com/google/go-containerregistry v0.21.8 h1:Ig/zIsnztdCUNaiNNczE+MoP5xcyUMfvpvfOr1xyMLE= +github.com/google/go-containerregistry v0.21.8/go.mod h1:dP5XNKcL7kMFF/TB3LfvWmVhAcv7iqkHb3oDK8aauTo= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0= github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= @@ -481,8 +486,8 @@ github.com/jandelgado/gcov2lcov v1.0.5 h1:rkBt40h0CVK4oCb8Dps950gvfd1rYvQ8+cWa34 github.com/jandelgado/gcov2lcov v1.0.5/go.mod h1:NnSxK6TMlg1oGDBfGelGbjgorT5/L3cchlbtgFYZSss= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOlocH6Fxy8MmwDt+yVQYULKfN0RoTN8A= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo= -github.com/jedisct1/go-minisign v0.0.0-20230811132847-661be99b8267 h1:TMtDYDHKYY15rFihtRfck/bfFqNfvcabqvXAFQfAUpY= -github.com/jedisct1/go-minisign v0.0.0-20230811132847-661be99b8267/go.mod h1:h1nSAbGFqGVzn6Jyl1R/iCcBUHN4g+gW1u9CoBTrb9E= +github.com/jedisct1/go-minisign v0.0.0-20241212093149-d2f9f49435c7 h1:FWpSWRD8FbVkKQu8M1DM9jF5oXFLyE+XpisIYfdzbic= +github.com/jedisct1/go-minisign v0.0.0-20241212093149-d2f9f49435c7/go.mod h1:BMxO138bOokdgt4UaxZiEfypcSHX0t6SIFimVP1oRfk= github.com/jellydator/ttlcache/v3 v3.4.0 h1:YS4P125qQS0tNhtL6aeYkheEaB/m8HCqdMMP4mnWdTY= github.com/jellydator/ttlcache/v3 v3.4.0/go.mod h1:Hw9EgjymziQD3yGsQdf1FqFdpp7YjFMd4Srg5EJlgD4= github.com/jmoiron/sqlx v1.3.5/go.mod h1:nRVWtLre0KfCLJvgxzCsLVMogSvQ1zNJtpYr2Ccp0mQ= @@ -494,12 +499,12 @@ github.com/jpillora/backoff v1.0.0/go.mod h1:J/6gKK9jxlEcS3zixgDgUAsiuZ7yrSoa/FX github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8= -github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4= -github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM= +github.com/kevinburke/ssh_config v1.6.0 h1:J1FBfmuVosPHf5GRdltRLhPJtJpTlMdKTBjRgTaQBFY= +github.com/kevinburke/ssh_config v1.6.0/go.mod h1:q2RIzfka+BXARoNexmF9gkxEX7DmvbW9P4hIVx2Kg4M= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= -github.com/klauspost/compress v1.18.7 h1:aUyZsS4kH3QTKurYhAOwAHxllVPnOthb3vPfnF1Ehjw= -github.com/klauspost/compress v1.18.7/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= +github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= github.com/knadh/koanf/maps v0.1.1 h1:G5TjmUh2D7G2YWf5SQQqSiHRJEjaicvU0KpypqB3NIs= @@ -545,13 +550,11 @@ github.com/lib/pq v1.10.2/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/lib/pq v1.10.7/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= -github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= -github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/lufia/plan9stats v0.0.0-20250317134145-8bc96cf8fc35 h1:PpXWgLPs+Fqr325bN2FD2ISlRRztXibcX6e8f5FR5Dc= +github.com/lufia/plan9stats v0.0.0-20250317134145-8bc96cf8fc35/go.mod h1:autxFIvghDt3jPTLoqZ9OZ7s9qTGNAWmYCjVFWPX/zg= github.com/luna-duclos/instrumentedsql v1.1.3/go.mod h1:9J1njvFds+zN7y85EDhN9XNQLANWwZt2ULeIC8yMNYs= github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE= github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= -github.com/mark3labs/mcp-go v0.55.1 h1:GLYqNm9qdMGPhCtK4g1t1y1vhAPfayOBuaibDi4mrSA= -github.com/mark3labs/mcp-go v0.55.1/go.mod h1:+8WclSK1ZUweCP3hvktSji8n8ABG/95QaEkeVE/Uwas= github.com/maruel/natural v1.1.1 h1:Hja7XhhmvEFhcByqDoHz9QZbkWey+COd9xWfCfn1ioo= github.com/maruel/natural v1.1.1/go.mod h1:v+Rfd79xlw1AgVBjbO0BEQmptqb5HvL/k9GRHB7ZKEg= github.com/mattn/go-colorable v0.1.1/go.mod h1:FuOcm+DKB9mbwrcAfNl7/TZVBZ6rcnceauSikq3lYCQ= @@ -583,18 +586,18 @@ github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa1 github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s= github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= -github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= -github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= +github.com/mitchellh/mapstructure v1.5.1-0.20231216201459-8508981c8b6c h1:cqn374mizHuIWj+OSJCajGr/phAmuMug9qIX3l9CflE= +github.com/mitchellh/mapstructure v1.5.1-0.20231216201459-8508981c8b6c/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ= github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw= github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= github.com/moby/go-archive v0.2.0 h1:zg5QDUM2mi0JIM9fdQZWC7U8+2ZfixfTYoHL7rWUcP8= github.com/moby/go-archive v0.2.0/go.mod h1:mNeivT14o8xU+5q1YnNrkQVpK+dnNe/K6fHqnTg4qPU= -github.com/moby/moby/api v1.54.2 h1:wiat9QAhnDQjA7wk1kh/TqHz2I1uUA7M7t9SAl/JNXg= -github.com/moby/moby/api v1.54.2/go.mod h1:+RQ6wluLwtYaTd1WnPLykIDPekkuyD/ROWQClE83pzs= -github.com/moby/moby/client v0.4.1 h1:DMQgisVoMkmMs7fp3ROSdiBnoAu8+vo3GggFl06M/wY= -github.com/moby/moby/client v0.4.1/go.mod h1:z52C9O2POPOsnxZAy//WtKcQ32P+jT/NGeXu/7nfjGQ= +github.com/moby/moby/api v1.55.0 h1:2/sexvQyqIWS8pRSCFddBfpW2qE7vR7FCL+vN8pxwMc= +github.com/moby/moby/api v1.55.0/go.mod h1:+RQ6wluLwtYaTd1WnPLykIDPekkuyD/ROWQClE83pzs= +github.com/moby/moby/client v0.5.1 h1:tYNaJno4c0HXz12y5BiqEDy0rVTYkWzI26lGvnTMiJw= +github.com/moby/moby/client v0.5.1/go.mod h1:odLstlZ6uSnfvAgVxMpvgmb8SUdd+siH2T0GBuxVAlM= github.com/moby/patternmatcher v0.6.1 h1:qlhtafmr6kgMIJjKJMDmMWq7WLkKIo23hsrpR3x084U= github.com/moby/patternmatcher v0.6.1/go.mod h1:hDPoyOpDY7OrrMDLaYoY3hf52gNCR/YOUYxkhApJIxc= github.com/moby/spdystream v0.5.1 h1:9sNYeYZUcci9R6/w7KDaFWEWeV4LStVG78Mpyq/Zm/Y= @@ -607,6 +610,8 @@ github.com/moby/sys/userns v0.1.0 h1:tVLXkFOxVu9A64/yh59slHVv9ahO9UIev4JZusOLG/g github.com/moby/sys/userns v0.1.0/go.mod h1:IHUYgu/kao6N8YZlp9Cf444ySSvCmDlmzUcYfDHOl28= github.com/moby/term v0.5.2 h1:6qk3FJAFDs6i/q3W/pQ97SX192qKfZgGjCQqfCJkgzQ= github.com/moby/term v0.5.2/go.mod h1:d3djjFCrjnB+fl8NJux+EJzu0msscUP+f8it8hPkFLc= +github.com/modelcontextprotocol/go-sdk v1.7.0 h1:yqjY2dsbKAC0LSuWZVBMrHgiG8ukXv6NRo0JiALay44= +github.com/modelcontextprotocol/go-sdk v1.7.0/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= github.com/modelcontextprotocol/registry v1.8.0 h1:x/seX0ji4iqRUpSovmkBcGbxfRiZZ7dgPwBcpCJrSTM= github.com/modelcontextprotocol/registry v1.8.0/go.mod h1:G6AUpTpZSekQvcLl5griUjijEE8vedARE/TyCaHEFdo= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -637,8 +642,8 @@ github.com/oauth2-proxy/mockoidc v0.0.0-20240214162133-caebfff84d25 h1:9bCMuD3Tc github.com/oauth2-proxy/mockoidc v0.0.0-20240214162133-caebfff84d25/go.mod h1:eDjgYHYDJbPLBLsyZ6qRaugP0mX8vePOhZ5id1fdzJw= github.com/oklog/ulid/v2 v2.1.1 h1:suPZ4ARWLOJLegGFiZZ1dFAkqzhMjL3J1TzI+5wHz8s= github.com/oklog/ulid/v2 v2.1.1/go.mod h1:rcEKHmBBKfef9DhnvX7y1HZBYxjXb0cP5ExxNsTT1QQ= -github.com/oleiade/reflections v1.0.1 h1:D1XO3LVEYroYskEsoSiGItp9RUxG6jWnCVvrqH0HHQM= -github.com/oleiade/reflections v1.0.1/go.mod h1:rdFxbxq4QXVZWj0F+e9jqjDkc7dbp97vkRixKo2JR60= +github.com/oleiade/reflections v1.1.0 h1:D+I/UsXQB4esMathlt0kkZRJZdUDmhv5zGi/HOwYTWo= +github.com/oleiade/reflections v1.1.0/go.mod h1:mCxx0QseeVCHs5Um5HhJeCKVC7AwS8kO67tky4rdisA= github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6 h1:zrbMGy9YXpIeTnGj4EljqMiZsIcE09mmF8XsD5AYOJc= github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6/go.mod h1:rEKTHC9roVVicUIfZK7DYrdIoM0EOr8mK1Hj5s3JjH0= github.com/olekukonko/errors v1.2.0 h1:10Zcn4GeV59t/EGqJc8fUjtFT/FuUh5bTMzZ1XwmCRo= @@ -670,6 +675,7 @@ github.com/ory/jsonschema/v3 v3.0.8/go.mod h1:ZPzqjDkwd3QTnb2Z6PAS+OTvBE2x5i6m25 github.com/ory/x v0.0.665 h1:61vv0ObCDSX1vOQYbxBeqDiv4YiPmMT91lYxDaaKX08= github.com/ory/x v0.0.665/go.mod h1:7SCTki3N0De3ZpqlxhxU/94ZrOCfNEnXwVtd0xVt+L8= github.com/pborman/getopt v0.0.0-20170112200414-7148bc3a4c30/go.mod h1:85jBQOZwpVEaDAr341tbn15RS4fCAsIst0qp7i8ex1o= +github.com/pborman/getopt v0.0.0-20180811024354-2b5b3bfb099b/go.mod h1:85jBQOZwpVEaDAr341tbn15RS4fCAsIst0qp7i8ex1o= github.com/pelletier/go-toml/v2 v2.3.1 h1:MYEvvGnQjeNkRF1qUuGolNtNExTDwct51yp7olPtrEc= github.com/pelletier/go-toml/v2 v2.3.1/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= github.com/pingcap/errors v0.11.4 h1:lFuQV/oaUMGcD2tqt+01ROSmJs75VG1ToEOkZIZ4nE4= @@ -689,16 +695,16 @@ github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= github.com/pressly/goose/v3 v3.27.1 h1:6uEvcprBybDmW4hcz3gYujhARhye+GoWKhEWyzD5sh4= github.com/pressly/goose/v3 v3.27.1/go.mod h1:maruOxsPnIG2yHHyo8UqKWXYKFcH7Q76csUV7+7KYoM= -github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= -github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg= +github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU= +github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE= github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= -github.com/prometheus/common v0.67.5 h1:pIgK94WWlQt1WLwAC5j2ynLaBRDiinoAb86HZHTUGI4= -github.com/prometheus/common v0.67.5/go.mod h1:SjE/0MzDEEAyrdr5Gqc6G+sXI67maCxzaT3A2+HqjUw= +github.com/prometheus/common v0.70.1 h1:1HvjP4D5oL3t8RsPlwxA9onvvStjtIHYE5XuuwOi/PY= +github.com/prometheus/common v0.70.1/go.mod h1:VdFUQDMZK3VLkurFUVhia6uys/0suUp86TJz5qbJRhc= github.com/prometheus/otlptranslator v1.0.0 h1:s0LJW/iN9dkIH+EnhiD3BlkkP5QVIUVEoIwkU+A6qos= github.com/prometheus/otlptranslator v1.0.0/go.mod h1:vRYWnXvI6aWGpsdY/mOT/cbeVRBlPWtBNDb7kGR3uKM= -github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= -github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= +github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI= +github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= github.com/redis/go-redis/v9 v9.21.0 h1:FPBE4hhbAke+TLmcY3WkpbDffJEomdqPn3HYiqAtL9E= github.com/redis/go-redis/v9 v9.21.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= @@ -716,10 +722,8 @@ github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/ryanuber/go-glob v1.0.0 h1:iQh3xXAumdQ+4Ufa5b25cRpC5TYKlno6hsv6Cb3pkBk= github.com/ryanuber/go-glob v1.0.0/go.mod h1:807d1WSdnB0XRJzKNil9Om6lcp/3a0v4qIHxIXzX/Yc= -github.com/sagikazarmark/locafero v0.11.0 h1:1iurJgmM9G3PA/I+wWYIOw/5SyBtxapeHDcg+AAIFXc= -github.com/sagikazarmark/locafero v0.11.0/go.mod h1:nVIGvgyzw595SUSUE6tvCp3YYTeHs15MvlmU87WwIik= -github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ= -github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= +github.com/sagikazarmark/locafero v0.12.0 h1:/NQhBAkUb4+fH1jivKHWusDYFjMOOKU88eegjfxfHb4= +github.com/sagikazarmark/locafero v0.12.0/go.mod h1:sZh36u/YSZ918v0Io+U9ogLYQJ9tLLBmM4eneO6WwsI= github.com/sassoftware/relic v7.2.1+incompatible h1:Pwyh1F3I0r4clFJXkSI8bOyJINGqpgjJU3DYAZeI05A= github.com/sassoftware/relic v7.2.1+incompatible/go.mod h1:CWfAxv73/iLZ17rbyhIEq3K9hs5w6FpNMdUT//qR+zk= github.com/sassoftware/relic/v7 v7.6.2 h1:rS44Lbv9G9eXsukknS4mSjIAuuX+lMq/FnStgmZlUv4= @@ -731,6 +735,8 @@ github.com/secure-systems-lab/go-securesystemslib v0.11.0 h1:iuCR9kcMFD4QurdKrGv github.com/secure-systems-lab/go-securesystemslib v0.11.0/go.mod h1:+PMOTjUGwHj2vcZ+TFKlb1tXRbrdWE1LYDT5i9JC80Q= github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0= github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/sergi/go-diff v1.2.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM= github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw= github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4= @@ -750,8 +756,8 @@ github.com/sigstore/rekor-tiles/v2 v2.3.0 h1:HhMgH61UP0t899V8Fjt7pz1YdgOBptbaQdn github.com/sigstore/rekor-tiles/v2 v2.3.0/go.mod h1:DEFiKSyQ4nF75QRVNdOPaIH3cmvMkO2B6xDZjNYngPc= github.com/sigstore/sigstore v1.10.8 h1:1Mgkxvkw4AXMfIP1DOjc6kw0GkUgA8pGVpveN/EfOq4= github.com/sigstore/sigstore v1.10.8/go.mod h1:f9+B/4iaYimvUkySyb2mvc73n3RLqNn24grHZM/ET8M= -github.com/sigstore/sigstore-go v1.2.2 h1:xAJ8hxaoecC0HKBYVbrwUjkeAI+GJYu6vLqbxDlD2Q0= -github.com/sigstore/sigstore-go v1.2.2/go.mod h1:MIFwBxAHJD+/lKgZzt9n/4Zhq/3T2+EuGX8iGrIsZgU= +github.com/sigstore/sigstore-go v1.3.0 h1:hnIMHREyCNTYFtOE1o7ae3Axa9B5W5EjUSBJICP2NBE= +github.com/sigstore/sigstore-go v1.3.0/go.mod h1:AyRQXfpH89py1twjE3kEZxlRersng90GSYqQV9zGJE8= github.com/sigstore/sigstore/pkg/signature/kms/aws v1.10.8 h1:tofVQ+UWJgad/69I5zbqxdFCN5gpIn9tRQP7iBzIpBw= github.com/sigstore/sigstore/pkg/signature/kms/aws v1.10.8/go.mod h1:73AfJE8H6w5KGCFPBu4x/OG+i1Yxgmh0L/FtV7prd88= github.com/sigstore/sigstore/pkg/signature/kms/azure v1.10.8 h1:8Mt7J36GcUEmbiJaiFhz2tud5ZIgkfVVCe2H/WJCHmw= @@ -760,19 +766,17 @@ github.com/sigstore/sigstore/pkg/signature/kms/gcp v1.10.8 h1:MxpAIMZVzn0Tpbarc9 github.com/sigstore/sigstore/pkg/signature/kms/gcp v1.10.8/go.mod h1:bnAUEkFNam6STvkVZhptVwWzWR5pS24CEtQ+lhxu7S0= github.com/sigstore/sigstore/pkg/signature/kms/hashivault v1.10.8 h1:1DGe4/clcdOnkz5MINEczWlmEvjUtZd+AjPPT/cBhQ8= github.com/sigstore/sigstore/pkg/signature/kms/hashivault v1.10.8/go.mod h1:6IDFhpgxtzqbnzrFkyegbj7RfWwKeRrb3/+xAD1Wp+Y= -github.com/sigstore/timestamp-authority/v2 v2.1.2 h1:7DDhnknLL4w8VwomyvW2W8qblOS9LDR8oihna+jc7Ls= -github.com/sigstore/timestamp-authority/v2 v2.1.2/go.mod h1:o6rAVZceFyejClIj/uStRNIemP16bVMZtbMmhk6pr0U= +github.com/sigstore/timestamp-authority/v2 v2.1.3 h1:Fc+LjCTfik1lh3YLkaosENfkXa3R2Y1nswiUKutBdFA= +github.com/sigstore/timestamp-authority/v2 v2.1.3/go.mod h1:myoFOKJB/u5vNTFwvBBJVkG3NnOBeIJevbfjNeasLjo= github.com/sirupsen/logrus v1.4.1/go.mod h1:ni0Sbl8bgC9z8RoU9G6nDWqqs/fq4eDPysMBDgk/93Q= github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE= github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0= github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w= github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g= -github.com/skeema/knownhosts v1.3.1 h1:X2osQ+RAjK76shCbvhHHHVl3ZlgDm8apHEHFqRjnBY8= -github.com/skeema/knownhosts v1.3.1/go.mod h1:r7KTdC8l4uxWRyK2TpQZ/1o5HaSzh06ePQNxPwTcfiY= +github.com/skeema/knownhosts v1.3.2 h1:EDL9mgf4NzwMXCTfaxSD/o/a5fxDw/xL9nkU28JjdBg= +github.com/skeema/knownhosts v1.3.2/go.mod h1:bEg3iQAuw+jyiw+484wwFJoKSLwcfd7fqRy+N0QTiow= github.com/sourcegraph/annotate v0.0.0-20160123013949-f4cad6c6324d/go.mod h1:UdhH50NIW0fCiwBSr0co2m7BnFLdv4fQTgdqdJTHFeE= -github.com/sourcegraph/conc v0.3.1-0.20240121214520-5f936abd7ae8 h1:+jumHNA0Wrelhe64i8F6HNlS8pkoyMv5sreGx2Ry5Rw= -github.com/sourcegraph/conc v0.3.1-0.20240121214520-5f936abd7ae8/go.mod h1:3n1Cwaq1E1/1lhQhtRK2ts/ZwZEhjcQeJQ1RuC6Q/8U= github.com/sourcegraph/syntaxhighlight v0.0.0-20170531221838-bd320f5d308e/go.mod h1:HuIsMU8RRBOtsCgI77wP899iHVBQpCmg4ErYMZB+2IA= github.com/spf13/afero v1.15.0 h1:b/YBCLWAJdFWJTN9cLhiXXcD7mzKn9Dm86dNnfyQw1I= github.com/spf13/afero v1.15.0/go.mod h1:NC2ByUVxtQs4b3sIUphxK0NioZnmxgyCrfzeuq8lxMg= @@ -787,10 +791,10 @@ github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/viper v1.21.0 h1:x5S+0EU27Lbphp4UKm1C+1oQO+rKx36vfCoaVebLFSU= github.com/spf13/viper v1.21.0/go.mod h1:P0lhsswPGWD/1lZJ9ny3fYnVqxiegrlNrEmgLjbTCAY= -github.com/stacklok/toolhive-catalog v0.20260706.0 h1:ZgEtTogdPjl2Tc/zamYExRH609HDJxS+EbxgdjmMH/w= -github.com/stacklok/toolhive-catalog v0.20260706.0/go.mod h1:J+l8hyYSbIhLCe3IsBJuVpP1wQbsuUXRS5o+B9M0GpU= -github.com/stacklok/toolhive-core v0.0.28 h1:+ajc8S/YdPm59Ik10pQxDYk75HjpyzwUU/Ru1WiLOO8= -github.com/stacklok/toolhive-core v0.0.28/go.mod h1:t8dqTOjyT67i+xidAS1cRWh3K5qYW38R/NgP7xlvirE= +github.com/stacklok/toolhive-catalog v0.20260804.0 h1:SsHVjix5uStamO6q87dv731KeGHrEO3f3SeVBLyNFIQ= +github.com/stacklok/toolhive-catalog v0.20260804.0/go.mod h1:G33nbS4PplpBl7lk/hPiIYpsmBl5OATvpohpXJujSfg= +github.com/stacklok/toolhive-core v0.0.38 h1:hBfBXptxtLQjzne68PD7EHmFEO0Epumm832KgOFmo00= +github.com/stacklok/toolhive-core v0.0.38/go.mod h1:12yRNG27dBwaGkNo1MpuUfD/6kyBU+n9fhH01ZfpMm8= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE= @@ -814,8 +818,8 @@ github.com/sv-tools/openapi v0.4.0 h1:UhD9DVnGox1hfTePNclpUzUFgos57FvzT2jmcAuTOJ github.com/sv-tools/openapi v0.4.0/go.mod h1:kD/dG+KP0+Fom1r6nvcj/ORtLus8d8enXT6dyRZDirE= github.com/swaggo/swag/v2 v2.0.0-rc5 h1:fK7d6ET9rrEsdB8IyuwXREWMcyQN3N7gawGFbbrjgHk= github.com/swaggo/swag/v2 v2.0.0-rc5/go.mod h1:kCL8Fu4Zl8d5tB2Bgj96b8wRowwrwk175bZHXfuGVFI= -github.com/tailscale/hujson v0.0.0-20260302212456-ecc657c15afd h1:Rf9uhF1+VJ7ZHqxrG8pJ6YacmHvVCmByDmGbAWCc/gA= -github.com/tailscale/hujson v0.0.0-20260302212456-ecc657c15afd/go.mod h1:EbW0wDK/qEUYI0A5bqq0C2kF8JTQwWONmGDBbzsxxHo= +github.com/tailscale/hujson v0.0.0-20260727124030-b80ff77dac4f h1:9hiVElpCmKzsBKQHkBqZ8LGzt82iLfM8egxr4sew+Ys= +github.com/tailscale/hujson v0.0.0-20260727124030-b80ff77dac4f/go.mod h1:8/zr1Tv0+cKpVtGCEB/7YfRXr2TszsMxMXLaT8YuBgU= github.com/testcontainers/testcontainers-go v0.42.0 h1:He3IhTzTZOygSXLJPMX7n44XtK+qhjat1nI9cneBbUY= github.com/testcontainers/testcontainers-go v0.42.0/go.mod h1:vZjdY1YmUA1qEForxOIOazfsrdyORJAbhi0bp8plN30= github.com/tetratelabs/wabin v0.0.0-20230304001439-f6f874872834 h1:ZF+QBjOI+tILZjBaFj3HgFonKXUcwgJ4djLb6i42S3Q= @@ -837,12 +841,12 @@ github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/tink-crypto/tink-go-awskms/v3 v3.0.0 h1:XSohRhCkXAVI0iaCnWB/GS05TEmpnKurQmzaY1jzt3Y= github.com/tink-crypto/tink-go-awskms/v3 v3.0.0/go.mod h1:+7MXsShLzVbSQ6dI0Pe4JuZM52jD1jQ1itAygd/MDsA= -github.com/tink-crypto/tink-go-gcpkms/v2 v2.2.0 h1:3B9i6XBXNTRspfkTC0asN5W0K6GhOSgcujNiECNRNb0= -github.com/tink-crypto/tink-go-gcpkms/v2 v2.2.0/go.mod h1:jY5YN2BqD/KSCHM9SqZPIpJNG/u3zwfLXHgws4x2IRw= +github.com/tink-crypto/tink-go-gcpkms/v2 v2.3.0 h1:3s6YMgMOBZRU8qG6ybpKSF2Sau+y3sMvxR911M59SwA= +github.com/tink-crypto/tink-go-gcpkms/v2 v2.3.0/go.mod h1:X8UNvbQu2wanAGa8ixRUU/DWt1V2hUBfvPGy6s9nE2s= github.com/tink-crypto/tink-go-hcvault/v2 v2.5.0 h1:eXuNqgrcYelxU1MVikOJDP3wTS5lvihM4ntoAbAMfvs= github.com/tink-crypto/tink-go-hcvault/v2 v2.5.0/go.mod h1:3RhcxAqek6xUlRFmJifvU4CYLZN60KMQdIKqpZAZJG0= -github.com/tink-crypto/tink-go/v2 v2.6.0 h1:+KHNBHhWH33Vn+igZWcsgdEPUxKwBMEe0QC60t388v4= -github.com/tink-crypto/tink-go/v2 v2.6.0/go.mod h1:2WbBA6pfNsAfBwDCggboaHeB2X29wkU8XHtGwh2YIk8= +github.com/tink-crypto/tink-go/v2 v2.7.0 h1:k7QnUXJ1cRDpvoy/5l1FimZqMAArRff8vjUqzi5N04o= +github.com/tink-crypto/tink-go/v2 v2.7.0/go.mod h1:cWNpQ/yAT/QHzAV0kBGMOSJzzYTKofDZdJaUqOPPWCI= github.com/titanous/rocacheck v0.0.0-20171023193734-afe73141d399 h1:e/5i7d4oYZ+C1wj2THlRK+oAhjeS/TRQwMfkIuet3w0= github.com/titanous/rocacheck v0.0.0-20171023193734-afe73141d399/go.mod h1:LdwHTNJT99C5fTAzDz0ud328OgXz+gierycbcIx2fRs= github.com/tklauser/go-sysconf v0.3.16 h1:frioLaCQSsF5Cy1jgRBrzr6t502KIIwQ0MArYICU0nA= @@ -874,6 +878,16 @@ github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zI github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 h1:ilQV1hzziu+LLM3zUTJ0trRztfwgjqKnBWNtSRkbmwM= github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78/go.mod h1:aL8wCCfTfSfmXjznFBSZNN13rSJjlIOI1fUNAtF7rmI= +github.com/ysmood/fetchup v0.2.3 h1:ulX+SonA0Vma5zUFXtv52Kzip/xe7aj4vqT5AJwQ+ZQ= +github.com/ysmood/fetchup v0.2.3/go.mod h1:xhibcRKziSvol0H1/pj33dnKrYyI2ebIvz5cOOkYGns= +github.com/ysmood/goob v0.4.0 h1:HsxXhyLBeGzWXnqVKtmT9qM7EuVs/XOgkX7T6r1o1AQ= +github.com/ysmood/goob v0.4.0/go.mod h1:u6yx7ZhS4Exf2MwciFr6nIM8knHQIE22lFpWHnfql18= +github.com/ysmood/got v0.40.0 h1:ZQk1B55zIvS7zflRrkGfPDrPG3d7+JOza1ZkNxcc74Q= +github.com/ysmood/got v0.40.0/go.mod h1:W7DdpuX6skL3NszLmAsC5hT7JAhuLZhByVzHTq874Qg= +github.com/ysmood/gson v0.7.3 h1:QFkWbTH8MxyUTKPkVWAENJhxqdBa4lYTQWqZCiLG6kE= +github.com/ysmood/gson v0.7.3/go.mod h1:3Kzs5zDl21g5F/BlLTNcuAGAYLKt2lV5G8D1zF3RNmg= +github.com/ysmood/leakless v0.9.0 h1:qxCG5VirSBvmi3uynXFkcnLMzkphdh3xx5FtrORwDCU= +github.com/ysmood/leakless v0.9.0/go.mod h1:R8iAXPRaG97QJwqxs74RdwzcRHT1SWCGTNqY8q0JvMQ= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.4.1/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k= @@ -895,10 +909,10 @@ go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0. go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace v0.46.1/go.mod h1:GnOaBaFQ2we3b9AGWJpsBa7v1S5RlQzlC3O7dRMxZhM= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0 h1:CqXxU8VOmDefoh0+ztfGaymYbhdB/tT3zs79QaZTNGY= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0/go.mod h1:BuhAPThV8PBHBvg8ZzZ/Ok3idOdhWIodywz2xEcRbJo= -go.opentelemetry.io/contrib/propagators/b3 v1.21.0 h1:uGdgDPNzwQWRwCXJgw/7h29JaRqcq9B87Iv4hJDKAZw= -go.opentelemetry.io/contrib/propagators/b3 v1.21.0/go.mod h1:D9GQXvVGT2pzyTfp1QBOnD1rzKEWzKjjwu5q2mslCUI= -go.opentelemetry.io/contrib/propagators/jaeger v1.21.1 h1:f4beMGDKiVzg9IcX7/VuWVy+oGdjx3dNJ72YehmtY5k= -go.opentelemetry.io/contrib/propagators/jaeger v1.21.1/go.mod h1:U9jhkEl8d1LL+QXY7q3kneJWJugiN3kZJV2OWz3hkBY= +go.opentelemetry.io/contrib/propagators/b3 v1.40.0 h1:xariChe8OOVF3rNlfzGFgQc61npQmXhzZj/i82mxMfg= +go.opentelemetry.io/contrib/propagators/b3 v1.40.0/go.mod h1:72WvbdxbOfXaELEQfonFfOL6osvcVjI7uJEE8C2nkrs= +go.opentelemetry.io/contrib/propagators/jaeger v1.40.0 h1:aXl9uobjJs5vquMLt9ZkI/3zIuz8XQ3TqOKSWx0/xdU= +go.opentelemetry.io/contrib/propagators/jaeger v1.40.0/go.mod h1:ioMePqe6k6c/ovXSkmkMr1mbN5qRBGJxNTVop7/2XO0= go.opentelemetry.io/contrib/samplers/jaegerremote v0.15.1 h1:Qb+5A+JbIjXwO7l4HkRUhgIn4Bzz0GNS2q+qdmSx+0c= go.opentelemetry.io/contrib/samplers/jaegerremote v0.15.1/go.mod h1:G4vNCm7fRk0kjZ6pGNLo5SpLxAUvOfSrcaegnT8TPck= go.opentelemetry.io/otel v1.45.0 h1:pdrWmLHofpubmArBv1LgFSv1Z0Ie/ppdZzu+kUN5EeU= @@ -911,8 +925,8 @@ go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 h1:QRefszxJmfPdjXUUm3j go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0/go.mod h1:Tiz03lTBVBrm7eWZBOidzEaYaJa8tjwGUGv6d8mlTyk= go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 h1:QBajQ2SrwQijzHyZbQlPsuIzpl/ll8DY6wPWsajeGcI= go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0/go.mod h1:08ZQLjrPLQ6R4kAXvuOvODEer5Yh4CoFvll5qB2BCI8= -go.opentelemetry.io/otel/exporters/prometheus v0.66.0 h1:vkrK8PAznv2NKt2r+kdu252ccGzkEqLc2aSXbQIALYQ= -go.opentelemetry.io/otel/exporters/prometheus v0.66.0/go.mod h1:V/UB6D3vMF/UBOL5igAsAYnk1nG/bzYYTzvsB16cy7o= +go.opentelemetry.io/otel/exporters/prometheus v0.67.0 h1:7IefDa35e6V3NoiqIeLDMDxMFyZDk5qcoC0Ax4cC16E= +go.opentelemetry.io/otel/exporters/prometheus v0.67.0/go.mod h1:nsPI1awTg5Vmg1YrommL2mVarVGlqc4yXOoKAkPRD0c= go.opentelemetry.io/otel/exporters/zipkin v1.45.0 h1:KN3btaILMTxR4QDHVGAO87lq5ButzK7l+kIfLuxQ1oA= go.opentelemetry.io/otel/exporters/zipkin v1.45.0/go.mod h1:yNcodmUclM4InyWoOwX/YW4Jri0Gj5FWAlM+NqCrtqY= go.opentelemetry.io/otel/metric v1.45.0 h1:7Eg1uH7CJ5cXv9is6tnBe1FI6rj1nwUdbFypRm3br/M= @@ -962,6 +976,7 @@ golang.ngrok.com/ngrok/v2 v2.1.4 h1:0JQZRqzVGBYluIi5MuhxNYx653qxpN7AiNwNJzoa9DQ= golang.ngrok.com/ngrok/v2 v2.1.4/go.mod h1:1bwK0+ZB4RJCJdqaXs2mvdsjeSk+x4YrrLn8IqOrIGo= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190411191339-88737f569e3a/go.mod h1:WFFai1msRO1wXaEeE5yQxYXgSfI8pQAWXbQop6sCtWE= +golang.org/x/crypto v0.0.0-20190426145343-a29dc8fdc734/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20190820162420-60c769a6c586/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= @@ -977,10 +992,10 @@ golang.org/x/crypto v0.57.0 h1:3ZVCjf8Ggz7zneR/EHRVx68Ctf+2pmIMP2UFhh9cC6M= golang.org/x/crypto v0.57.0/go.mod h1:Fdz0i5U6CoizGwLda9DttjSk6qlZo25zYNtR+ycvuZA= golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f h1:W3F4c+6OLc6H2lb//N1q4WpJkhzJCK5J6kUi1NTVXfM= golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f/go.mod h1:J1xhfL/vlindoeF/aINzNzt2Bket5bjo9sdOYzOsU80= -golang.org/x/exp/event v0.0.0-20260603202125-055de637280b h1:IpzEcnKaJbfA/HT53pFMUAXzMPT5XdewwIPRg90OQvQ= -golang.org/x/exp/event v0.0.0-20260603202125-055de637280b/go.mod h1:RUiT8+7zwWmByxLGqLQEh4LSWvL7Mw8K/PMAHE2VtI0= -golang.org/x/exp/jsonrpc2 v0.0.0-20260611194520-c48552f49976 h1:r4lXELculxMIoA1AybwvwhaegK7P/Ee+0O7FkIrh5dk= -golang.org/x/exp/jsonrpc2 v0.0.0-20260611194520-c48552f49976/go.mod h1:QF08E0BddR88x8aRGsxrq0Nn/5InTRyK5FHHMGEJiRk= +golang.org/x/exp/event v0.0.0-20260611194520-c48552f49976 h1:aMc3mP8qjwSEdljznT5CYeyvpoFP7e3VCNqZZwUf3F0= +golang.org/x/exp/event v0.0.0-20260611194520-c48552f49976/go.mod h1:cV3mj8TmLcC/Zki//yCMYw/nSC1BpHamZNkMSU7lgPw= +golang.org/x/exp/jsonrpc2 v0.0.0-20260709172345-9ea1abe57597 h1:V+5CIAxiRMu5nXZ2eLlo3mQXPm4zrvbZI6ddjylxYLs= +golang.org/x/exp/jsonrpc2 v0.0.0-20260709172345-9ea1abe57597/go.mod h1:it3Sp7J2OmdyHxl5q3GMU4d2QZsFQGglbij9kqQVVaI= golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= @@ -1150,8 +1165,8 @@ k8s.io/client-go v0.35.3 h1:s1lZbpN4uI6IxeTM2cpdtrwHcSOBML1ODNTCCfsP1pg= k8s.io/client-go v0.35.3/go.mod h1:RzoXkc0mzpWIDvBrRnD+VlfXP+lRzqQjCmKtiwZ8Q9c= k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= k8s.io/klog/v2 v2.140.0/go.mod h1:o+/RWfJ6PwpnFn7OyAG3QnO47BFsymfEfrz6XyYSSp0= -k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912 h1:Y3gxNAuB0OBLImH611+UDZcmKS3g6CthxToOb37KgwE= -k8s.io/kube-openapi v0.0.0-20250910181357-589584f1c912/go.mod h1:kdmbQkyfwUagLfXIad1y2TdrjPFWp2Q89B3qkRwf/pQ= +k8s.io/kube-openapi v0.0.0-20251125145642-4e65d59e963e h1:iW9ChlU0cU16w8MpVYjXk12dqQ4BPFBEgif+ap7/hqQ= +k8s.io/kube-openapi v0.0.0-20251125145642-4e65d59e963e/go.mod h1:kdmbQkyfwUagLfXIad1y2TdrjPFWp2Q89B3qkRwf/pQ= k8s.io/utils v0.0.0-20260319190234-28399d86e0b5 h1:kBawHLSnx/mYHmRnNUf9d4CpjREbeZuxoSGOX/J+aYM= k8s.io/utils v0.0.0-20260319190234-28399d86e0b5/go.mod h1:xDxuJ0whA3d0I4mf/C4ppKHxXynQ+fxnkmQH0vTHnuk= modernc.org/cc/v4 v4.28.1 h1:XpLbkYVQ24E8tX5u8+yWGvaxerxkR/S4zqxI8ZoSBuc= diff --git a/pkg/api/discovery_permissions_windows_test.go b/pkg/api/discovery_permissions_windows_test.go new file mode 100644 index 0000000000..45f61fe740 --- /dev/null +++ b/pkg/api/discovery_permissions_windows_test.go @@ -0,0 +1,107 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build windows + +package api + +import ( + "context" + "encoding/json" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/adrg/xdg" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" + + "github.com/stacklok/toolhive/pkg/server/discovery" +) + +// TestWriteDiscoveryFile_RestrictsDirBeforeTrustingExistingFile pins the +// production ordering. writeDiscoveryFile must lock down the discovery chain +// before it acquires server.json.lock or calls Discover. A pre-planted file +// written while the chain was still loose must not be deleted during lockdown; +// startup fails closed when the planted record answers healthy. +// +//nolint:paralleltest // t.Setenv and xdg.Reload mutate process-wide state +func TestWriteDiscoveryFile_RestrictsDirBeforeTrustingExistingFile(t *testing.T) { + base := t.TempDir() + grantEveryone(t, base) + + t.Setenv("XDG_STATE_HOME", base) + xdg.Reload() + t.Cleanup(xdg.Reload) + // Guard against silently operating on the real %LOCALAPPDATA%: everything + // below rewrites ACLs. + require.Equal(t, base, xdg.StateHome, "XDG_STATE_HOME must redirect the discovery path") + + serverDir := filepath.Dir(discovery.FilePath()) + toolhiveDir := filepath.Dir(serverDir) + require.Equal(t, base, filepath.Dir(toolhiveDir)) + + // A server that answers /health with the planted nonce, so Discover + // classifies the planted file as StateRunning. + const plantedNonce = "planted-nonce" + healthSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set(discovery.NonceHeader, plantedNonce) + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(healthSrv.Close) + + // Plant the file the way an attacker with write access to a loose + // directory would: plain MkdirAll, so both directories inherit Everyone. + require.NoError(t, os.MkdirAll(serverDir, 0700)) + planted, err := json.Marshal(&discovery.ServerInfo{ + URL: healthSrv.URL, + PID: os.Getpid(), + Nonce: plantedNonce, + StartedAt: time.Now().UTC(), + }) + require.NoError(t, err) + require.NoError(t, os.WriteFile(discovery.FilePath(), planted, 0600)) + require.Contains(t, strings.ToUpper(dirSDDL(t, serverDir)), "WD", + "precondition: the planted directory must inherit Everyone (WD)") + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = listener.Close() }) + s := &Server{listener: listener, address: listener.Addr().String(), nonce: "our-nonce"} + + err = s.writeDiscoveryFile(context.Background()) + require.Error(t, err) + assert.Contains(t, err.Error(), "already running") + + // The chain must already be locked down, on the leaf and on the + // intermediate directory. + require.NoError(t, discovery.ValidateRestrictedDiscoveryDACL(toolhiveDir)) + require.NoError(t, discovery.ValidateRestrictedDiscoveryDACL(serverDir)) + + // The forged record must survive lockdown so a healthy predecessor is not erased. + onDisk, err := os.ReadFile(discovery.FilePath()) + require.NoError(t, err) + assert.Equal(t, string(planted), string(onDisk)) +} + +func grantEveryone(t *testing.T, path string) { + t.Helper() + // icacls is the product-path way to introduce a loose ACE; quoting keeps + // PowerShell from expanding (OI)/(CI). + out, err := exec.Command("icacls", path, "/grant", "*S-1-1-0:(OI)(CI)M").CombinedOutput() + require.NoError(t, err, "icacls grant Everyone failed: %s", out) +} + +func dirSDDL(t *testing.T, dir string) string { + t.Helper() + sd, err := windows.GetNamedSecurityInfo(dir, windows.SE_FILE_OBJECT, windows.DACL_SECURITY_INFORMATION) + require.NoError(t, err) + return sd.String() +} diff --git a/pkg/api/server.go b/pkg/api/server.go index dcb5028fda..d370a48f01 100644 --- a/pkg/api/server.go +++ b/pkg/api/server.go @@ -26,7 +26,6 @@ import ( "net" "net/http" "os" - "path/filepath" "strings" "time" @@ -37,6 +36,7 @@ import ( "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" + ociplugins "github.com/stacklok/toolhive-core/oci/plugins" ociskills "github.com/stacklok/toolhive-core/oci/skills" regtypes "github.com/stacklok/toolhive-core/registry/types" v1 "github.com/stacklok/toolhive/pkg/api/v1" @@ -49,6 +49,9 @@ import ( "github.com/stacklok/toolhive/pkg/fileutils" "github.com/stacklok/toolhive/pkg/groups" "github.com/stacklok/toolhive/pkg/networking" + "github.com/stacklok/toolhive/pkg/plugins" + "github.com/stacklok/toolhive/pkg/plugins/adapters" + "github.com/stacklok/toolhive/pkg/plugins/pluginsvc" "github.com/stacklok/toolhive/pkg/recovery" "github.com/stacklok/toolhive/pkg/registry" "github.com/stacklok/toolhive/pkg/server/discovery" @@ -78,21 +81,23 @@ const ( // ServerBuilder provides a fluent interface for building and configuring the API server type ServerBuilder struct { - address string - isUnixSocket bool - debugMode bool - enableDocs bool - nonce string - oidcConfig *auth.TokenValidatorConfig - otelEnabled bool - middlewares []func(http.Handler) http.Handler - customRoutes map[string]http.Handler - containerRuntime runtime.Runtime - clientManager client.Manager - workloadManager workloads.Manager - groupManager groups.Manager - skillManager skills.SkillService - skillStoreCloser io.Closer + address string + isUnixSocket bool + debugMode bool + enableDocs bool + nonce string + oidcConfig *auth.TokenValidatorConfig + otelEnabled bool + middlewares []func(http.Handler) http.Handler + customRoutes map[string]http.Handler + containerRuntime runtime.Runtime + clientManager client.Manager + workloadManager workloads.Manager + groupManager groups.Manager + skillManager skills.SkillService + skillStoreCloser io.Closer + pluginManager plugins.PluginService + pluginStoreCloser io.Closer } // NewServerBuilder creates a new ServerBuilder with default configuration @@ -283,46 +288,112 @@ func (b *ServerBuilder) createDefaultManagers(ctx context.Context) error { } if b.skillManager == nil { - store, storeErr := sqlite.NewDefaultSkillStore() - if storeErr != nil { - return fmt.Errorf("failed to create skill store: %w", storeErr) - } - b.skillStoreCloser = store - cm, cmErr := client.NewClientManager() - if cmErr != nil { - _ = store.Close() - return fmt.Errorf("failed to create client manager for skills: %w", cmErr) + if err := b.createDefaultSkillManager(); err != nil { + return err } + } - ociStore, ociErr := ociskills.NewStore(ociskills.DefaultStoreRoot()) - if ociErr != nil { - _ = store.Close() - return fmt.Errorf("failed to create OCI skill store: %w", ociErr) - } - ociRegistry, regErr := newOCIRegistryClient() - if regErr != nil { - _ = store.Close() - // ociStore is directory-backed with no open handles; no cleanup needed. - return fmt.Errorf("failed to create OCI registry client: %w", regErr) - } - packager := ociskills.NewPackager(ociStore) - - skillOpts := []skillsvc.Option{ - skillsvc.WithPathResolver(&clientPathAdapter{cm: cm}), - skillsvc.WithOCIStore(ociStore), - skillsvc.WithPackager(packager), - skillsvc.WithRegistryClient(ociRegistry), - skillsvc.WithGroupManager(b.groupManager), + if b.pluginManager == nil { + if err := b.createDefaultPluginManager(); err != nil { + return err } + } - skillOpts = append(skillOpts, - skillsvc.WithSkillLookup(lazySkillLookup{}), - skillsvc.WithGitResolver(gitresolver.NewResolver()), - ) + return nil +} - b.skillManager = skillsvc.New(store, skillOpts...) +// createDefaultSkillManager builds the default skill service (skillsvc.New) +// wired with a SQLite skill store, the OCI skill store/packager/registry, +// a path resolver backed by the client manager, the group manager, the +// lazy registry lookup, and the git resolver. +func (b *ServerBuilder) createDefaultSkillManager() error { + store, storeErr := sqlite.NewDefaultSkillStore() + if storeErr != nil { + return fmt.Errorf("failed to create skill store: %w", storeErr) + } + b.skillStoreCloser = store + cm, cmErr := client.NewClientManager() + if cmErr != nil { + _ = store.Close() + return fmt.Errorf("failed to create client manager for skills: %w", cmErr) } + ociStore, ociErr := ociskills.NewStore(ociskills.DefaultStoreRoot()) + if ociErr != nil { + _ = store.Close() + return fmt.Errorf("failed to create OCI skill store: %w", ociErr) + } + ociRegistry, regErr := newOCIRegistryClient() + if regErr != nil { + _ = store.Close() + // ociStore is directory-backed with no open handles; no cleanup needed. + return fmt.Errorf("failed to create OCI registry client: %w", regErr) + } + packager := ociskills.NewPackager(ociStore) + + skillOpts := []skillsvc.Option{ + skillsvc.WithPathResolver(&clientPathAdapter{cm: cm}), + skillsvc.WithOCIStore(ociStore), + skillsvc.WithPackager(packager), + skillsvc.WithRegistryClient(ociRegistry), + skillsvc.WithGroupManager(b.groupManager), + } + + skillOpts = append(skillOpts, + skillsvc.WithSkillLookup(lazySkillLookup{}), + skillsvc.WithGitResolver(gitresolver.NewResolver()), + ) + + b.skillManager = skillsvc.New(store, skillOpts...) + return nil +} + +// createDefaultPluginManager builds the default plugin service (pluginsvc.New) +// wired with a SQLite plugin store, the OCI plugin store/packager/registry, +// per-client materialization adapters, and the group manager. Mirrors the +// skill-manager block in createDefaultManagers. +func (b *ServerBuilder) createDefaultPluginManager() error { + store, storeErr := sqlite.NewDefaultPluginStore() + if storeErr != nil { + return fmt.Errorf("failed to create plugin store: %w", storeErr) + } + b.pluginStoreCloser = store + cm, cmErr := client.NewClientManager() + if cmErr != nil { + _ = store.Close() + return fmt.Errorf("failed to create client manager for plugins: %w", cmErr) + } + + ociStore, ociErr := ociplugins.NewStore(ociplugins.DefaultStoreRoot()) + if ociErr != nil { + _ = store.Close() + return fmt.Errorf("failed to create OCI plugin store: %w", ociErr) + } + ociRegistry, regErr := newPluginOCIRegistryClient() + if regErr != nil { + _ = store.Close() + // ociStore is directory-backed with no open handles; no cleanup needed. + return fmt.Errorf("failed to create plugin OCI registry client: %w", regErr) + } + packager := ociplugins.NewPackager(ociStore) + + materializers := map[string]plugins.MaterializationAdapter{ + string(client.ClaudeCode): adapters.NewClaudeCodeAdapter(cm), + string(client.Codex): adapters.NewCodexAdapter(cm), + } + + pluginOpts := []pluginsvc.Option{ + pluginsvc.WithStore(store), + pluginsvc.WithOCIStore(ociStore), + pluginsvc.WithPackager(packager), + pluginsvc.WithRegistryClient(ociRegistry), + pluginsvc.WithGroupManager(b.groupManager), + pluginsvc.WithMaterializers(materializers), + pluginsvc.WithClientManager(cm), + pluginsvc.WithPluginLookup(lazyPluginLookup{}), + } + + b.pluginManager = pluginsvc.New(pluginOpts...) return nil } @@ -338,6 +409,13 @@ func (b *ServerBuilder) setupDefaultRoutes(r *chi.Mux) { b.debugMode, )) + // Skills router does the same: install, sync, and upgrade pull OCI + // artifacts, so a flat 60s cap would sever them mid-transfer. + r.Mount("/api/v1beta/skills", v1.SkillsRouter(b.skillManager)) + + // Plugins router likewise: install, build, and push move OCI artifacts. + r.Mount("/api/v1beta/plugins", v1.PluginsRouter(b.pluginManager)) + // All other routes get standard timeout standardRouters := map[string]http.Handler{ "/health": v1.HealthcheckRouter(b.containerRuntime, b.nonce), @@ -347,7 +425,6 @@ func (b *ServerBuilder) setupDefaultRoutes(r *chi.Mux) { "/api/v1beta/clients": v1.ClientRouter(b.clientManager, b.workloadManager, b.groupManager), "/api/v1beta/secrets": v1.SecretsRouter(), "/api/v1beta/groups": v1.GroupsRouter(b.groupManager, b.workloadManager, b.clientManager), - "/api/v1beta/skills": v1.SkillsRouter(b.skillManager), "/registry": v1.RegistryV01Router(), } for prefix, router := range standardRouters { @@ -437,13 +514,14 @@ func getComponentAndVersionFromRequest(r *http.Request) (string, string, bool) { // Server represents a configured HTTP server type Server struct { - httpServer *http.Server - listener net.Listener - address string - isUnixSocket bool - addrType string - nonce string - storeCloser io.Closer + httpServer *http.Server + listener net.Listener + address string + isUnixSocket bool + addrType string + nonce string + skillStoreCloser io.Closer + pluginStoreCloser io.Closer } // NewServer creates a new Server instance from a pre-configured builder @@ -475,13 +553,14 @@ func NewServer(ctx context.Context, builder *ServerBuilder) (*Server, error) { } return &Server{ - httpServer: httpServer, - listener: listener, - address: builder.address, - isUnixSocket: builder.isUnixSocket, - addrType: addrType, - nonce: builder.nonce, - storeCloser: builder.skillStoreCloser, + httpServer: httpServer, + listener: listener, + address: builder.address, + isUnixSocket: builder.isUnixSocket, + addrType: addrType, + nonce: builder.nonce, + skillStoreCloser: builder.skillStoreCloser, + pluginStoreCloser: builder.pluginStoreCloser, }, nil } @@ -534,14 +613,23 @@ func (s *Server) writeDiscoveryFile(ctx context.Context) error { return nil } - // Ensure the discovery directory exists before acquiring the lock, - // since the lock file is created in the same directory. - discoveryPath := discovery.FilePath() - if err := os.MkdirAll(filepath.Dir(discoveryPath), 0700); err != nil { - return fmt.Errorf("failed to create discovery directory: %w", err) + // Create and lock down the discovery directory before acquiring the lock, + // since the lock file is created in the same directory and Discover below + // trusts whatever server.json it finds there. Restricting only on the write + // would be too late: a StateRunning result returns before the write, and + // the lock file itself would be taken in a directory other accounts can + // still write to. + secure, err := discovery.EnsureSecureDirEx() + if err != nil { + return err } + discoveryPath := discovery.FilePath() return fileutils.WithFileLock(discoveryPath, func() error { + if err := discovery.ReconcileDiscoveryAfterInsecureUpgrade(ctx, secure.RepairedInsecureChain); err != nil { + return err + } + // Guard against overwriting another server's discovery file. result, err := discovery.Discover(ctx) if err != nil { @@ -602,11 +690,16 @@ func (s *Server) cleanup() { slog.Warn("failed to remove discovery file", "error", err) } } - if s.storeCloser != nil { - if err := s.storeCloser.Close(); err != nil { + if s.skillStoreCloser != nil { + if err := s.skillStoreCloser.Close(); err != nil { slog.Warn("failed to close skill store", "error", err) } } + if s.pluginStoreCloser != nil { + if err := s.pluginStoreCloser.Close(); err != nil { + slog.Warn("failed to close plugin store", "error", err) + } + } if s.isUnixSocket { cleanupUnixSocket(s.address) } @@ -702,6 +795,53 @@ func ociRefHost(ref string) string { return ref } +// newPluginOCIRegistryClient creates an OCI registry client for plugin +// artifacts. In dev mode (TOOLHIVE_DEV=true), plain HTTP is used for loopback +// registries only via devModePluginRegistryClient. Mirrors +// newOCIRegistryClient (which is the skills analogue) but typed for the +// toolhive-core oci/plugins package. +func newPluginOCIRegistryClient() (ociplugins.RegistryClient, error) { + secure, err := ociplugins.NewRegistry() + if err != nil { + return nil, err + } + if os.Getenv("TOOLHIVE_DEV") != "true" { + return secure, nil + } + plain, err := ociplugins.NewRegistry(ociplugins.WithPlainHTTP(true)) + if err != nil { + return nil, err + } + return &devModePluginRegistryClient{secure: secure, plain: plain}, nil +} + +// devModePluginRegistryClient dispatches each Pull/Push to a plain-HTTP client +// when the target reference's host is loopback (a local test registry), and to +// a normal TLS client otherwise. Mirrors devModeOCIRegistryClient. +type devModePluginRegistryClient struct { + secure ociplugins.RegistryClient + plain ociplugins.RegistryClient +} + +func (c *devModePluginRegistryClient) Pull( + ctx context.Context, store *ociplugins.Store, ref string, +) (digest.Digest, error) { + return c.clientFor(ref).Pull(ctx, store, ref) +} + +func (c *devModePluginRegistryClient) Push( + ctx context.Context, store *ociplugins.Store, manifestDigest digest.Digest, ref string, +) error { + return c.clientFor(ref).Push(ctx, store, manifestDigest, ref) +} + +func (c *devModePluginRegistryClient) clientFor(ref string) ociplugins.RegistryClient { + if networking.IsLocalhost(ociRefHost(ref)) { + return c.plain + } + return c.secure +} + // lazySkillLookup implements skillsvc.SkillLookup by resolving the registry // provider on each call. This ensures that registry config changes (via // thv config set-registry or the API) are picked up without restarting @@ -717,6 +857,52 @@ func (lazySkillLookup) SearchSkills(query string) ([]regtypes.Skill, error) { return provider.SearchSkills(query) } +// lazyPluginLookup implements pluginsvc.PluginLookup by resolving the registry +// provider on each call, mirroring lazySkillLookup. Registry config changes are +// picked up without restarting the server. +type lazyPluginLookup struct{} + +func (lazyPluginLookup) SearchPlugins(_ context.Context, query string) ([]pluginsvc.PluginSearchHit, error) { + provider, err := registry.GetDefaultProviderWithConfig(config.NewDefaultProvider()) + if err != nil { + return nil, err + } + found, err := provider.SearchPlugins(query) + if err != nil { + return nil, err + } + return pluginHitsFromRegistry(found), nil +} + +// pluginHitsFromRegistry adapts registry plugin search results to the +// pluginsvc lookup shape. The OCI reference lives in SkillPackage.Identifier +// on the wire; pluginsvc consumes it as PluginPackage.Reference. Namespace, +// Version, and Digest are carried through so the install flow can disambiguate +// by namespace, honor an explicit version request, and pin to a verified +// digest. +func pluginHitsFromRegistry(regPlugins []regtypes.Plugin) []pluginsvc.PluginSearchHit { + hits := make([]pluginsvc.PluginSearchHit, 0, len(regPlugins)) + for i := range regPlugins { + p := ®Plugins[i] + pkgs := make([]pluginsvc.PluginPackage, 0, len(p.Packages)) + for _, sp := range p.Packages { + pkgs = append(pkgs, pluginsvc.PluginPackage{ + Reference: sp.Identifier, + Type: sp.RegistryType, + Digest: sp.Digest, + }) + } + hits = append(hits, pluginsvc.PluginSearchHit{ + Name: p.Name, + Namespace: p.Namespace, + Version: p.Version, + Description: p.Description, + Packages: pkgs, + }) + } + return hits +} + // clientPathAdapter adapts *client.ClientManager to the skills.PathResolver interface. type clientPathAdapter struct { cm *client.ClientManager diff --git a/pkg/api/server_test.go b/pkg/api/server_test.go index 57f8bc8fa6..12c72159af 100644 --- a/pkg/api/server_test.go +++ b/pkg/api/server_test.go @@ -17,6 +17,9 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + regtypes "github.com/stacklok/toolhive-core/registry/types" + pluginsmocks "github.com/stacklok/toolhive/pkg/plugins/mocks" + "github.com/stacklok/toolhive/pkg/plugins/pluginsvc" skillsmocks "github.com/stacklok/toolhive/pkg/skills/mocks" ) @@ -155,11 +158,13 @@ func TestServerBuilderExtensionPoints(t *testing.T) { func TestNewServer_ReadTimeoutConfigured(t *testing.T) { t.Parallel() - // Inject a skill manager so Build() skips creating the default SQLite skill - // store, which is shared on disk and races under parallel tests (SQLITE_BUSY). + // Inject mock skill and plugin managers so Build() skips creating the default + // SQLite stores, which share a DB file on disk and race under parallel tests + // (SQLITE_BUSY). ctrl := gomock.NewController(t) b := NewServerBuilder().WithAddress("127.0.0.1:0") b.skillManager = skillsmocks.NewMockSkillService(ctrl) + b.pluginManager = pluginsmocks.NewMockPluginService(ctrl) s, err := NewServer(context.Background(), b) require.NoError(t, err) @@ -169,3 +174,116 @@ func TestNewServer_ReadTimeoutConfigured(t *testing.T) { assert.Equal(t, readTimeout, s.httpServer.ReadTimeout) assert.Zero(t, s.httpServer.WriteTimeout) } + +func TestPluginHitsFromRegistry(t *testing.T) { + t.Parallel() + + skillPkg := func(identifier, registryType, digest string) regtypes.SkillPackage { + return regtypes.SkillPackage{Identifier: identifier, RegistryType: registryType, Digest: digest} + } + + tests := []struct { + name string + input []regtypes.Plugin + assert func(t *testing.T, hits []pluginsvc.PluginSearchHit) + }{ + { + name: "single plugin with one oci package", + input: []regtypes.Plugin{ + { + Name: "code-reviewer", + Namespace: "io.github.user", + Version: "1.0.0", + Description: "Reviews code for bugs", + Packages: []regtypes.SkillPackage{ + skillPkg("ghcr.io/org/code-reviewer:v1", "oci", "sha256:abc"), + }, + }, + }, + assert: func(t *testing.T, hits []pluginsvc.PluginSearchHit) { + t.Helper() + require.Len(t, hits, 1) + assert.Equal(t, "code-reviewer", hits[0].Name) + assert.Equal(t, "io.github.user", hits[0].Namespace) + assert.Equal(t, "1.0.0", hits[0].Version) + assert.Equal(t, "Reviews code for bugs", hits[0].Description) + require.Len(t, hits[0].Packages, 1) + assert.Equal(t, "ghcr.io/org/code-reviewer:v1", hits[0].Packages[0].Reference) + assert.Equal(t, "oci", hits[0].Packages[0].Type) + assert.Equal(t, "sha256:abc", hits[0].Packages[0].Digest) + }, + }, + { + name: "multiple packages (oci + git)", + input: []regtypes.Plugin{ + { + Name: "multi-pkg", + Namespace: "io.github.acme", + Version: "2.1.0", + Description: "Has multiple packages", + Packages: []regtypes.SkillPackage{ + skillPkg("ghcr.io/org/multi:latest", "oci", "sha256:def"), + skillPkg("https://github.com/org/repo.git", "git", ""), + }, + }, + }, + assert: func(t *testing.T, hits []pluginsvc.PluginSearchHit) { + t.Helper() + require.Len(t, hits, 1) + assert.Equal(t, "io.github.acme", hits[0].Namespace) + assert.Equal(t, "2.1.0", hits[0].Version) + require.Len(t, hits[0].Packages, 2) + assert.Equal(t, "ghcr.io/org/multi:latest", hits[0].Packages[0].Reference) + assert.Equal(t, "oci", hits[0].Packages[0].Type) + assert.Equal(t, "sha256:def", hits[0].Packages[0].Digest) + assert.Equal(t, "https://github.com/org/repo.git", hits[0].Packages[1].Reference) + assert.Equal(t, "git", hits[0].Packages[1].Type) + assert.Empty(t, hits[0].Packages[1].Digest) + }, + }, + { + name: "plugin with zero packages", + input: []regtypes.Plugin{ + {Name: "no-pkgs", Namespace: "io.github.bare", Version: "0.1.0", Description: "No packages"}, + }, + assert: func(t *testing.T, hits []pluginsvc.PluginSearchHit) { + t.Helper() + require.Len(t, hits, 1) + assert.Equal(t, "no-pkgs", hits[0].Name) + assert.Equal(t, "io.github.bare", hits[0].Namespace) + assert.Equal(t, "0.1.0", hits[0].Version) + assert.NotNil(t, hits[0].Packages) + assert.Empty(t, hits[0].Packages) + }, + }, + { + name: "empty input returns non-nil empty slice", + input: []regtypes.Plugin{}, + assert: func(t *testing.T, hits []pluginsvc.PluginSearchHit) { + t.Helper() + assert.NotNil(t, hits) + assert.Empty(t, hits) + }, + }, + { + name: "Description empty maps to empty string", + input: []regtypes.Plugin{ + {Name: "bare-plugin"}, + }, + assert: func(t *testing.T, hits []pluginsvc.PluginSearchHit) { + t.Helper() + require.Len(t, hits, 1) + assert.Equal(t, "bare-plugin", hits[0].Name) + assert.Empty(t, hits[0].Description) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + hits := pluginHitsFromRegistry(tt.input) + tt.assert(t, hits) + }) + } +} diff --git a/pkg/api/v1/plugins.go b/pkg/api/v1/plugins.go new file mode 100644 index 0000000000..2ff6aee30e --- /dev/null +++ b/pkg/api/v1/plugins.go @@ -0,0 +1,400 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/go-chi/chi/v5" + "github.com/go-chi/chi/v5/middleware" + + "github.com/stacklok/toolhive-core/httperr" + apierrors "github.com/stacklok/toolhive/pkg/api/errors" + "github.com/stacklok/toolhive/pkg/plugins" +) + +// PluginsRoutes defines the routes for plugin management. +type PluginsRoutes struct { + pluginService plugins.PluginService +} + +// PluginsRouter creates a new router for plugin management endpoints. +func PluginsRouter(pluginService plugins.PluginService) http.Handler { + routes := PluginsRoutes{ + pluginService: pluginService, + } + + // Mirrors WorkloadRouter and SkillsRouter: routes that move OCI artifacts + // get a timeout sized for the transfer, everything else keeps the short + // one. Without the split these routes inherit the flat 60s applied to + // standard routers, which severs a cold artifact pull mid-flight. + stdTimeout := middleware.Timeout(standardRouteTimeout) + longTimeout := middleware.Timeout(longRunningRouteTimeout) + + r := chi.NewRouter() + r.With(stdTimeout).Get("/", apierrors.ErrorHandler(routes.listPlugins)) + r.With(longTimeout).Post("/", apierrors.ErrorHandler(routes.installPlugin)) + r.With(stdTimeout).Delete("/{name}", apierrors.ErrorHandler(routes.uninstallPlugin)) + r.With(stdTimeout).Get("/{name}", apierrors.ErrorHandler(routes.getPluginInfo)) + r.With(stdTimeout).Post("/validate", apierrors.ErrorHandler(routes.validatePlugin)) + r.With(longTimeout).Post("/build", apierrors.ErrorHandler(routes.buildPlugin)) + r.With(longTimeout).Post("/push", apierrors.ErrorHandler(routes.pushPlugin)) + r.With(stdTimeout).Get("/builds", apierrors.ErrorHandler(routes.listBuilds)) + r.With(stdTimeout).Delete("/builds/{tag}", apierrors.ErrorHandler(routes.deleteBuild)) + r.With(stdTimeout).Get("/content", apierrors.ErrorHandler(routes.getPluginContent)) + + return r +} + +// listPlugins returns a list of installed plugins. +// +// @Summary List all installed plugins +// @Description Get a list of all installed plugins +// @Tags plugins +// @Produce json +// @Param scope query string false "Filter by scope (user or project)" Enums(user, project) +// @Param client query string false "Filter by client app" +// @Param project_root query string false "Filter by project root path" +// @Param group query string false "Filter by group name" +// @Success 200 {object} pluginListResponse +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins [get] +func (s *PluginsRoutes) listPlugins(w http.ResponseWriter, r *http.Request) error { + scope := plugins.Scope(r.URL.Query().Get("scope")) + projectRoot := r.URL.Query().Get("project_root") + client := r.URL.Query().Get("client") + group := r.URL.Query().Get("group") + + result, err := s.pluginService.List(r.Context(), plugins.ListOptions{ + Scope: scope, + ClientApp: client, + ProjectRoot: projectRoot, + Group: group, + }) + if err != nil { + return err + } + + if result == nil { + result = []plugins.InstalledPlugin{} + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(pluginListResponse{Plugins: result}) +} + +// installPlugin installs a plugin from a remote source. +// +// @Summary Install a plugin +// @Description Install a plugin from a remote source +// @Tags plugins +// @Accept json +// @Produce json +// @Param request body installPluginRequest true "Install request" +// @Success 201 {object} installPluginResponse +// @Header 201 {string} Location "URI of the installed plugin resource" +// @Failure 400 {string} string "Bad Request" +// @Failure 401 {string} string "Unauthorized (registry refused credentials)" +// @Failure 404 {string} string "Not Found (artifact not present in registry)" +// @Failure 409 {string} string "Conflict" +// @Failure 429 {string} string "Too Many Requests (registry rate limit)" +// @Failure 500 {string} string "Internal Server Error" +// @Failure 502 {string} string "Bad Gateway (upstream registry failure)" +// @Failure 504 {string} string "Gateway Timeout (upstream pull timed out)" +// @Router /api/v1beta/plugins [post] +func (s *PluginsRoutes) installPlugin(w http.ResponseWriter, r *http.Request) error { + var req installPluginRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + return httperr.WithCode( + fmt.Errorf("invalid request body: %w", err), + http.StatusBadRequest, + ) + } + + result, err := s.pluginService.Install(r.Context(), plugins.InstallOptions{ + Name: req.Name, + Version: req.Version, + Scope: req.Scope, + ProjectRoot: req.ProjectRoot, + Clients: req.Clients, + Force: req.Force, + Group: req.Group, + }) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Location", fmt.Sprintf("/api/v1beta/plugins/%s", result.Plugin.Metadata.Name)) + w.WriteHeader(http.StatusCreated) + return json.NewEncoder(w).Encode(installPluginResponse{Plugin: result.Plugin}) +} + +// uninstallPlugin removes an installed plugin. +// +// @Summary Uninstall a plugin +// @Description Remove an installed plugin +// @Tags plugins +// @Param name path string true "Plugin name" +// @Param scope query string false "Scope to uninstall from (user or project)" Enums(user, project) +// @Param project_root query string false "Project root path for project-scoped plugins" +// @Success 204 {string} string "No Content" +// @Failure 400 {string} string "Bad Request" +// @Failure 404 {string} string "Not Found" +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/{name} [delete] +func (s *PluginsRoutes) uninstallPlugin(w http.ResponseWriter, r *http.Request) error { + name := chi.URLParam(r, "name") + + if err := plugins.ValidatePluginName(name); err != nil { + return httperr.WithCode(err, http.StatusBadRequest) + } + + scope := plugins.Scope(r.URL.Query().Get("scope")) + projectRoot := r.URL.Query().Get("project_root") + + if err := s.pluginService.Uninstall(r.Context(), plugins.UninstallOptions{ + Name: name, + Scope: scope, + ProjectRoot: projectRoot, + }); err != nil { + return err + } + + w.WriteHeader(http.StatusNoContent) + return nil +} + +// getPluginInfo returns detailed information about a plugin. +// +// @Summary Get plugin details +// @Description Get detailed information about a specific plugin +// @Tags plugins +// @Produce json +// @Param name path string true "Plugin name" +// @Param scope query string false "Filter by scope (user or project)" Enums(user, project) +// @Param project_root query string false "Project root path for project-scoped plugins" +// @Success 200 {object} plugins.PluginInfo +// @Failure 400 {string} string "Bad Request" +// @Failure 404 {string} string "Not Found" +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/{name} [get] +func (s *PluginsRoutes) getPluginInfo(w http.ResponseWriter, r *http.Request) error { + name := chi.URLParam(r, "name") + + if err := plugins.ValidatePluginName(name); err != nil { + return httperr.WithCode(err, http.StatusBadRequest) + } + + scope := plugins.Scope(r.URL.Query().Get("scope")) + projectRoot := r.URL.Query().Get("project_root") + + info, err := s.pluginService.Info(r.Context(), plugins.InfoOptions{ + Name: name, + Scope: scope, + ProjectRoot: projectRoot, + }) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(info) +} + +// validatePlugin checks whether a plugin definition is valid. +// +// @Summary Validate a plugin +// @Description Validate a plugin definition +// @Tags plugins +// @Accept json +// @Produce json +// @Param request body validatePluginRequest true "Validate request" +// @Success 200 {object} plugins.ValidationResult +// @Failure 400 {string} string "Bad Request" +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/validate [post] +func (s *PluginsRoutes) validatePlugin(w http.ResponseWriter, r *http.Request) error { + var req validatePluginRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + return httperr.WithCode( + fmt.Errorf("invalid request body: %w", err), + http.StatusBadRequest, + ) + } + + if req.Path == "" { + return httperr.WithCode( + fmt.Errorf("path is required"), + http.StatusBadRequest, + ) + } + + result, err := s.pluginService.Validate(r.Context(), req.Path) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(result) +} + +// buildPlugin builds a plugin from a local directory into an OCI artifact. +// +// @Summary Build a plugin +// @Description Build a plugin from a local directory +// @Tags plugins +// @Accept json +// @Produce json +// @Param request body buildPluginRequest true "Build request" +// @Success 200 {object} plugins.BuildResult +// @Failure 400 {string} string "Bad Request" +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/build [post] +func (s *PluginsRoutes) buildPlugin(w http.ResponseWriter, r *http.Request) error { + var req buildPluginRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + return httperr.WithCode( + fmt.Errorf("invalid request body: %w", err), + http.StatusBadRequest, + ) + } + + if req.Path == "" { + return httperr.WithCode( + fmt.Errorf("path is required"), + http.StatusBadRequest, + ) + } + + result, err := s.pluginService.Build(r.Context(), plugins.BuildOptions{ + Path: req.Path, + Tag: req.Tag, + }) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(result) +} + +// pushPlugin pushes a built plugin artifact to a remote registry. +// +// @Summary Push a plugin +// @Description Push a built plugin artifact to a remote registry +// @Tags plugins +// @Accept json +// @Param request body pushPluginRequest true "Push request" +// @Success 204 {string} string "No Content" +// @Failure 400 {string} string "Bad Request" +// @Failure 404 {string} string "Not Found" +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/push [post] +func (s *PluginsRoutes) pushPlugin(w http.ResponseWriter, r *http.Request) error { + var req pushPluginRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + return httperr.WithCode( + fmt.Errorf("invalid request body: %w", err), + http.StatusBadRequest, + ) + } + + if req.Reference == "" { + return httperr.WithCode( + fmt.Errorf("reference is required"), + http.StatusBadRequest, + ) + } + + if err := s.pluginService.Push(r.Context(), plugins.PushOptions{ + Reference: req.Reference, + }); err != nil { + return err + } + + w.WriteHeader(http.StatusNoContent) + return nil +} + +// listBuilds returns a list of locally-built OCI plugin artifacts. +// +// @Summary List locally-built plugin artifacts +// @Description Get a list of all locally-built OCI plugin artifacts in the local store +// @Tags plugins +// @Produce json +// @Success 200 {object} pluginBuildListResponse +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/builds [get] +func (s *PluginsRoutes) listBuilds(w http.ResponseWriter, r *http.Request) error { + builds, err := s.pluginService.ListBuilds(r.Context()) + if err != nil { + return err + } + + if builds == nil { + builds = []plugins.LocalBuild{} + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(pluginBuildListResponse{Builds: builds}) +} + +// deleteBuild removes a locally-built OCI plugin artifact from the local store. +// +// @Summary Delete a locally-built plugin artifact +// @Description Remove a locally-built OCI plugin artifact and its blobs from the local store +// @Tags plugins +// @Param tag path string true "Artifact tag" +// @Success 204 {string} string "No Content" +// @Failure 404 {string} string "Not Found" +// @Failure 500 {string} string "Internal Server Error" +// @Router /api/v1beta/plugins/builds/{tag} [delete] +func (s *PluginsRoutes) deleteBuild(w http.ResponseWriter, r *http.Request) error { + tag := chi.URLParam(r, "tag") + if err := s.pluginService.DeleteBuild(r.Context(), tag); err != nil { + return err + } + w.WriteHeader(http.StatusNoContent) + return nil +} + +// getPluginContent retrieves the plugin.json body and file listing from an OCI artifact. +// +// @Summary Get plugin content +// @Description Retrieve the plugin.json body and file listing from an artifact +// @Description without installing it. Accepts OCI refs, git refs, or local tags. +// @Tags plugins +// @Produce json +// @Param ref query string true "OCI reference or local build tag" +// @Success 200 {object} plugins.PluginContent +// @Failure 400 {string} string "Bad Request" +// @Failure 401 {string} string "Unauthorized (registry refused credentials)" +// @Failure 404 {string} string "Not Found (artifact not present in registry)" +// @Failure 429 {string} string "Too Many Requests (registry rate limit)" +// @Failure 500 {string} string "Internal Server Error" +// @Failure 502 {string} string "Bad Gateway (upstream registry or git resolver failure)" +// @Failure 504 {string} string "Gateway Timeout (upstream pull timed out)" +// @Router /api/v1beta/plugins/content [get] +func (s *PluginsRoutes) getPluginContent(w http.ResponseWriter, r *http.Request) error { + ref := r.URL.Query().Get("ref") + if ref == "" { + return httperr.WithCode( + fmt.Errorf("ref query parameter is required"), + http.StatusBadRequest, + ) + } + + content, err := s.pluginService.GetContent(r.Context(), plugins.ContentOptions{ + Reference: ref, + }) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(content) +} diff --git a/pkg/api/v1/plugins_test.go b/pkg/api/v1/plugins_test.go new file mode 100644 index 0000000000..d621e994b0 --- /dev/null +++ b/pkg/api/v1/plugins_test.go @@ -0,0 +1,679 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/go-chi/chi/v5" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/plugins" + plugmocks "github.com/stacklok/toolhive/pkg/plugins/mocks" + "github.com/stacklok/toolhive/pkg/storage" +) + +func TestPluginsRouter(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + method string + path string + body string + setupMock func(*plugmocks.MockPluginService, string) + expectedStatus int + expectedBody string + }{ + // listPlugins + { + name: "list plugins success empty", + method: "GET", + path: "/", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().List(gomock.Any(), plugins.ListOptions{}). + Return([]plugins.InstalledPlugin{}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `{"plugins":[]}`, + }, + { + name: "list plugins success with results", + method: "GET", + path: "/", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().List(gomock.Any(), plugins.ListOptions{}). + Return([]plugins.InstalledPlugin{ + { + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + InstalledAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }, + }, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"my-plugin"`, + }, + { + name: "list plugins project scope missing project root", + method: "GET", + path: "/?scope=project", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().List(gomock.Any(), plugins.ListOptions{ + Scope: plugins.ScopeProject, + }).Return(nil, httperr.WithCode( + fmt.Errorf("project_root is required for project scope"), + http.StatusBadRequest, + )) + }, + expectedStatus: http.StatusBadRequest, + expectedBody: "project_root is required", + }, + { + name: "list plugins with project root filter", + method: "GET", + path: "/?scope=project&project_root={{project_root}}", + setupMock: func(svc *plugmocks.MockPluginService, projectRoot string) { + svc.EXPECT().List(gomock.Any(), plugins.ListOptions{ + Scope: plugins.ScopeProject, + ProjectRoot: projectRoot, + }).Return([]plugins.InstalledPlugin{}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `{"plugins":[]}`, + }, + { + name: "list plugins with client filter", + method: "GET", + path: "/?client=claude-code", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().List(gomock.Any(), plugins.ListOptions{ClientApp: "claude-code"}). + Return([]plugins.InstalledPlugin{}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `{"plugins":[]}`, + }, + { + name: "list plugins with group filter", + method: "GET", + path: "/?group=my-group", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().List(gomock.Any(), plugins.ListOptions{Group: "my-group"}). + Return([]plugins.InstalledPlugin{}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `{"plugins":[]}`, + }, + { + name: "list plugins error", + method: "GET", + path: "/", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().List(gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("database error")) + }, + expectedStatus: http.StatusInternalServerError, + expectedBody: "Internal Server Error", + }, + // installPlugin + { + name: "install plugin success", + method: "POST", + path: "/", + body: `{"name":"my-plugin"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Install(gomock.Any(), plugins.InstallOptions{Name: "my-plugin"}). + Return(&plugins.InstallResult{ + Plugin: plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusPending, + InstalledAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }, + }, nil) + }, + expectedStatus: http.StatusCreated, + expectedBody: `"my-plugin"`, + }, + { + name: "install plugin empty name", + method: "POST", + path: "/", + body: `{"name":""}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Install(gomock.Any(), plugins.InstallOptions{Name: ""}). + Return(nil, httperr.WithCode(fmt.Errorf("invalid plugin name: must not be empty"), http.StatusBadRequest)) + }, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid plugin name", + }, + { + name: "install plugin missing name field", + method: "POST", + path: "/", + body: `{}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Install(gomock.Any(), plugins.InstallOptions{Name: ""}). + Return(nil, httperr.WithCode(fmt.Errorf("invalid plugin name: must not be empty"), http.StatusBadRequest)) + }, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid plugin name", + }, + { + name: "install plugin malformed json", + method: "POST", + path: "/", + body: `{invalid`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid request body", + }, + { + name: "install plugin already exists", + method: "POST", + path: "/", + body: `{"name":"my-plugin"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Install(gomock.Any(), gomock.Any()). + Return(nil, storage.ErrAlreadyExists) + }, + expectedStatus: http.StatusConflict, + expectedBody: "resource already exists", + }, + { + name: "install plugin with clients", + method: "POST", + path: "/", + body: `{"name":"my-plugin","clients":["claude-code","codex"]}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Install(gomock.Any(), plugins.InstallOptions{ + Name: "my-plugin", + Clients: []string{"claude-code", "codex"}, + }).Return(&plugins.InstallResult{ + Plugin: plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Status: plugins.InstallStatusInstalled, + Clients: []string{"claude-code", "codex"}, + }, + }, nil) + }, + expectedStatus: http.StatusCreated, + expectedBody: `"my-plugin"`, + }, + { + name: "install plugin with version and scope", + method: "POST", + path: "/", + body: `{"name":"my-plugin","version":"1.2.0","scope":"project","project_root":"{{project_root}}"}`, + setupMock: func(svc *plugmocks.MockPluginService, projectRoot string) { + svc.EXPECT().Install(gomock.Any(), plugins.InstallOptions{ + Name: "my-plugin", + Version: "1.2.0", + Scope: plugins.ScopeProject, + ProjectRoot: projectRoot, + }).Return(&plugins.InstallResult{ + Plugin: plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin", Version: "1.2.0"}, + Scope: plugins.ScopeProject, + Status: plugins.InstallStatusPending, + }, + }, nil) + }, + expectedStatus: http.StatusCreated, + expectedBody: `"my-plugin"`, + }, + // uninstallPlugin + { + name: "uninstall plugin success", + method: "DELETE", + path: "/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Uninstall(gomock.Any(), plugins.UninstallOptions{Name: "my-plugin"}). + Return(nil) + }, + expectedStatus: http.StatusNoContent, + }, + { + name: "uninstall plugin invalid name", + method: "DELETE", + path: "/A", + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid", + }, + { + name: "uninstall plugin not found", + method: "DELETE", + path: "/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Uninstall(gomock.Any(), gomock.Any()). + Return(storage.ErrNotFound) + }, + expectedStatus: http.StatusNotFound, + expectedBody: "resource not found", + }, + { + name: "uninstall plugin with scope", + method: "DELETE", + path: "/my-plugin?scope=project&project_root={{project_root}}", + setupMock: func(svc *plugmocks.MockPluginService, projectRoot string) { + svc.EXPECT().Uninstall(gomock.Any(), plugins.UninstallOptions{ + Name: "my-plugin", + Scope: plugins.ScopeProject, + ProjectRoot: projectRoot, + }).Return(nil) + }, + expectedStatus: http.StatusNoContent, + }, + // getPluginInfo + { + name: "get plugin info found", + method: "GET", + path: "/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Info(gomock.Any(), plugins.InfoOptions{Name: "my-plugin"}). + Return(&plugins.PluginInfo{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + InstalledPlugin: &plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + }, + }, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"installed_plugin"`, + }, + { + name: "get plugin info not found", + method: "GET", + path: "/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Info(gomock.Any(), plugins.InfoOptions{Name: "my-plugin"}). + Return(nil, storage.ErrNotFound) + }, + expectedStatus: http.StatusNotFound, + expectedBody: "resource not found", + }, + { + name: "get plugin info invalid name", + method: "GET", + path: "/A", + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid", + }, + { + name: "get plugin info service error", + method: "GET", + path: "/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Info(gomock.Any(), plugins.InfoOptions{Name: "my-plugin"}). + Return(nil, fmt.Errorf("database error")) + }, + expectedStatus: http.StatusInternalServerError, + expectedBody: "Internal Server Error", + }, + { + name: "get plugin info with unmaterialized components and degraded clients", + method: "GET", + path: "/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Info(gomock.Any(), plugins.InfoOptions{Name: "my-plugin"}). + Return(&plugins.PluginInfo{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + InstalledPlugin: &plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Scope: plugins.ScopeProject, + Status: plugins.InstallStatusInstalled, + }, + UnmaterializedComponents: map[string][]plugins.ComponentType{ + "codex": {plugins.ComponentCommands, plugins.ComponentAgents}, + }, + ProjectScopeDegradedClients: []string{"codex"}, + }, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"unmaterialized_components"`, + }, + // validatePlugin + { + name: "validate plugin success", + method: "POST", + path: "/validate", + body: `{"path":"/tmp/plugin"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Validate(gomock.Any(), "/tmp/plugin"). + Return(&plugins.ValidationResult{Valid: true}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"valid":true`, + }, + { + name: "validate plugin bad request", + method: "POST", + path: "/validate", + body: `{invalid`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid request body", + }, + { + name: "validate plugin empty path", + method: "POST", + path: "/validate", + body: `{"path":""}`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "path is required", + }, + { + name: "validate plugin service error", + method: "POST", + path: "/validate", + body: `{"path":"/tmp/plugin"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Validate(gomock.Any(), "/tmp/plugin"). + Return(nil, fmt.Errorf("validation failed")) + }, + expectedStatus: http.StatusInternalServerError, + expectedBody: "Internal Server Error", + }, + // buildPlugin + { + name: "build plugin success", + method: "POST", + path: "/build", + body: `{"path":"/tmp/plugin","tag":"v1.0.0"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Build(gomock.Any(), plugins.BuildOptions{Path: "/tmp/plugin", Tag: "v1.0.0"}). + Return(&plugins.BuildResult{Reference: "v1.0.0"}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"reference":"v1.0.0"`, + }, + { + name: "build plugin bad request", + method: "POST", + path: "/build", + body: `{invalid`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid request body", + }, + { + name: "build plugin empty path", + method: "POST", + path: "/build", + body: `{"path":"","tag":"v1"}`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "path is required", + }, + // pushPlugin + { + name: "push plugin success", + method: "POST", + path: "/push", + body: `{"reference":"ghcr.io/test/plugin:v1"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Push(gomock.Any(), plugins.PushOptions{Reference: "ghcr.io/test/plugin:v1"}). + Return(nil) + }, + expectedStatus: http.StatusNoContent, + }, + { + name: "push plugin bad request", + method: "POST", + path: "/push", + body: `{invalid`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "invalid request body", + }, + { + name: "push plugin empty reference", + method: "POST", + path: "/push", + body: `{"reference":""}`, + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "reference is required", + }, + { + name: "push plugin service error", + method: "POST", + path: "/push", + body: `{"reference":"ghcr.io/test/plugin:v1"}`, + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().Push(gomock.Any(), plugins.PushOptions{Reference: "ghcr.io/test/plugin:v1"}). + Return(fmt.Errorf("push failed")) + }, + expectedStatus: http.StatusInternalServerError, + expectedBody: "Internal Server Error", + }, + // listBuilds + { + name: "list builds success empty", + method: "GET", + path: "/builds", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().ListBuilds(gomock.Any()). + Return([]plugins.LocalBuild{}, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `{"builds":[]}`, + }, + { + name: "list builds success with results", + method: "GET", + path: "/builds", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().ListBuilds(gomock.Any()). + Return([]plugins.LocalBuild{ + {Tag: "my-plugin", Digest: "sha256:abc123", Name: "my-plugin", Version: "1.0.0"}, + }, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"tag":"my-plugin"`, + }, + { + name: "list builds service error", + method: "GET", + path: "/builds", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().ListBuilds(gomock.Any()). + Return(nil, httperr.WithCode(fmt.Errorf("oci store not configured"), http.StatusInternalServerError)) + }, + expectedStatus: http.StatusInternalServerError, + expectedBody: "Internal Server Error", + }, + // deleteBuild + { + name: "delete build success", + method: "DELETE", + path: "/builds/my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().DeleteBuild(gomock.Any(), "my-plugin").Return(nil) + }, + expectedStatus: http.StatusNoContent, + }, + { + name: "delete build not found", + method: "DELETE", + path: "/builds/missing", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().DeleteBuild(gomock.Any(), "missing"). + Return(httperr.WithCode(fmt.Errorf("tag not found"), http.StatusNotFound)) + }, + expectedStatus: http.StatusNotFound, + }, + // getPluginContent + { + name: "get plugin content success", + method: "GET", + path: "/content?ref=my-plugin", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().GetContent(gomock.Any(), plugins.ContentOptions{Reference: "my-plugin"}). + Return(&plugins.PluginContent{ + Name: "my-plugin", + Manifest: `{"name":"my-plugin"}`, + Files: []plugins.PluginFileEntry{{Path: "plugin.json", Size: 42}}, + }, nil) + }, + expectedStatus: http.StatusOK, + expectedBody: `"my-plugin"`, + }, + { + name: "get plugin content missing ref", + method: "GET", + path: "/content", + setupMock: func(_ *plugmocks.MockPluginService, _ string) {}, + expectedStatus: http.StatusBadRequest, + expectedBody: "ref query parameter is required", + }, + { + name: "get plugin content service error", + method: "GET", + path: "/content?ref=missing", + setupMock: func(svc *plugmocks.MockPluginService, _ string) { + svc.EXPECT().GetContent(gomock.Any(), plugins.ContentOptions{Reference: "missing"}). + Return(nil, storage.ErrNotFound) + }, + expectedStatus: http.StatusNotFound, + expectedBody: "resource not found", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + path := tt.path + body := tt.body + projectRoot := "" + if strings.Contains(path, "{{project_root}}") || strings.Contains(body, "{{project_root}}") { + projectRoot = makeProjectRoot(t) + path = strings.ReplaceAll(path, "{{project_root}}", url.QueryEscape(projectRoot)) + body = strings.ReplaceAll(body, "{{project_root}}", projectRoot) + } + + ctrl := gomock.NewController(t) + mockSvc := plugmocks.NewMockPluginService(ctrl) + tt.setupMock(mockSvc, projectRoot) + + router := chi.NewRouter() + router.Mount("/", PluginsRouter(mockSvc)) + + req := httptest.NewRequest(tt.method, path, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + router.ServeHTTP(rec, req) + + assert.Equal(t, tt.expectedStatus, rec.Code) + if tt.expectedBody != "" { + assert.Contains(t, rec.Body.String(), tt.expectedBody) + } + }) + } +} + +func TestPluginsRouterGetInfoJSONKeys(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + mockSvc := plugmocks.NewMockPluginService(ctrl) + + mockSvc.EXPECT().Info(gomock.Any(), plugins.InfoOptions{Name: "my-plugin"}). + Return(&plugins.PluginInfo{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + InstalledPlugin: &plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Scope: plugins.ScopeProject, + Status: plugins.InstallStatusInstalled, + }, + UnmaterializedComponents: map[string][]plugins.ComponentType{ + "codex": {plugins.ComponentCommands, plugins.ComponentAgents}, + }, + ProjectScopeDegradedClients: []string{"codex"}, + }, nil) + + router := chi.NewRouter() + router.Mount("/", PluginsRouter(mockSvc)) + + req := httptest.NewRequest(http.MethodGet, "/my-plugin", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + body := rec.Body.String() + assert.Contains(t, body, `"unmaterialized_components"`) + assert.Contains(t, body, `"project_scope_degraded_clients"`) +} + +func TestPluginsListResponseFormat(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + mockSvc := plugmocks.NewMockPluginService(ctrl) + + mockSvc.EXPECT().List(gomock.Any(), gomock.Any()). + Return([]plugins.InstalledPlugin{ + { + Metadata: plugins.PluginMetadata{Name: "plugin-one", Version: "1.0.0"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + InstalledAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), + }, + }, nil) + + router := chi.NewRouter() + router.Mount("/", PluginsRouter(mockSvc)) + + req := httptest.NewRequest(http.MethodGet, "/", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusOK, rec.Code) + assert.Contains(t, rec.Body.String(), `"plugins"`) + assert.Contains(t, rec.Body.String(), `"plugin-one"`) +} + +func TestPluginsInstallLocationHeader(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + mockSvc := plugmocks.NewMockPluginService(ctrl) + + mockSvc.EXPECT().Install(gomock.Any(), plugins.InstallOptions{Name: "my-plugin"}). + Return(&plugins.InstallResult{ + Plugin: plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + }, + }, nil) + + router := chi.NewRouter() + router.Mount("/", PluginsRouter(mockSvc)) + + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"name":"my-plugin"}`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusCreated, rec.Code) + assert.Equal(t, "/api/v1beta/plugins/my-plugin", rec.Header().Get("Location")) +} diff --git a/pkg/api/v1/plugins_types.go b/pkg/api/v1/plugins_types.go new file mode 100644 index 0000000000..e32aac0b8d --- /dev/null +++ b/pkg/api/v1/plugins_types.go @@ -0,0 +1,78 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import "github.com/stacklok/toolhive/pkg/plugins" + +// pluginListResponse represents the response for listing plugins. +// +// @Description Response containing a list of installed plugins +type pluginListResponse struct { + // List of installed plugins + Plugins []plugins.InstalledPlugin `json:"plugins"` +} + +// installPluginRequest represents the request to install a plugin. +// +// @Description Request to install a plugin +type installPluginRequest struct { + // Name or OCI reference of the plugin to install + Name string `json:"name"` + // Version to install (empty means latest) + Version string `json:"version,omitempty"` + // Scope for the installation + Scope plugins.Scope `json:"scope,omitempty"` + // ProjectRoot is the project root path for project-scoped installs + ProjectRoot string `json:"project_root,omitempty"` + // Clients lists target client identifiers (e.g., "claude-code"), + // or ["all"] to target every plugin-supporting client. + // Omitting this field installs to all available clients. + Clients []string `json:"clients,omitempty"` + // Force allows overwriting unmanaged plugin directories + Force bool `json:"force,omitempty"` + // Group is the group name to add the plugin to after installation + Group string `json:"group,omitempty"` +} + +// installPluginResponse represents the response after installing a plugin. +// +// @Description Response after successfully installing a plugin +type installPluginResponse struct { + // The installed plugin + Plugin plugins.InstalledPlugin `json:"plugin"` +} + +// validatePluginRequest represents the request to validate a plugin. +// +// @Description Request to validate a plugin definition +type validatePluginRequest struct { + // Path to the plugin definition directory + Path string `json:"path"` +} + +// buildPluginRequest represents the request to build a plugin. +// +// @Description Request to build a plugin from a local directory +type buildPluginRequest struct { + // Path to the plugin definition directory + Path string `json:"path"` + // OCI tag for the built artifact + Tag string `json:"tag,omitempty"` +} + +// pushPluginRequest represents the request to push a plugin. +// +// @Description Request to push a built plugin artifact +type pushPluginRequest struct { + // OCI reference to push + Reference string `json:"reference"` +} + +// pluginBuildListResponse represents the response for listing locally-built OCI plugin artifacts. +// +// @Description Response containing a list of locally-built OCI plugin artifacts +type pluginBuildListResponse struct { + // List of locally-built OCI plugin artifacts + Builds []plugins.LocalBuild `json:"builds"` +} diff --git a/pkg/api/v1/registry_v01.go b/pkg/api/v1/registry_v01.go index 574bb98318..0e67f6adf7 100644 --- a/pkg/api/v1/registry_v01.go +++ b/pkg/api/v1/registry_v01.go @@ -32,6 +32,8 @@ func RegistryV01Router() http.Handler { r.Get("/servers/{serverName}/versions/latest", getServerV01) r.Get("/x/dev.toolhive/skills", listSkillsV01) r.Get("/x/dev.toolhive/skills/{namespace}/{skillName}", getSkillV01) + r.Get("/x/dev.toolhive/plugins", listPluginsV01) + r.Get("/x/dev.toolhive/plugins/{namespace}/{pluginName}", getPluginV01) }) return r } diff --git a/pkg/api/v1/registry_v01_plugins.go b/pkg/api/v1/registry_v01_plugins.go new file mode 100644 index 0000000000..5d14b5e0e1 --- /dev/null +++ b/pkg/api/v1/registry_v01_plugins.go @@ -0,0 +1,152 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "encoding/json" + "errors" + "log/slog" + "net/http" + "strings" + + "github.com/go-chi/chi/v5" + + types "github.com/stacklok/toolhive-core/registry/types" + regpkg "github.com/stacklok/toolhive/pkg/registry" + "github.com/stacklok/toolhive/pkg/registry/api" +) + +// listPluginsV01 handles GET /registry/{registryName}/v0.1/x/dev.toolhive/plugins +// +// @Summary List available registry plugins +// @Description Get a paginated list of plugins from the registry. Supports optional full-text search and pagination. +// @Tags registry-plugins +// @Produce json +// @Param registryName path string true "Registry name (currently ignored, uses the default provider)" +// @Param q query string false "Search filter — matches against plugin name, namespace, and description" +// @Param page query integer false "Page number, 1-based (default: 1)" +// @Param limit query integer false "Items per page, max 200 (default: 50)" +// @Success 200 {object} pluginsV01Response +// @Failure 500 {object} registryErrorResponse "Internal server error" +// @Failure 503 {object} registryErrorResponse "Registry authentication required or upstream registry unavailable" +// @Router /registry/{registryName}/v0.1/x/dev.toolhive/plugins [get] +func listPluginsV01(w http.ResponseWriter, r *http.Request) { + provider, ok := getRegistryProvider(w) + if !ok { + return + } + + plugins, err := provider.ListAvailablePlugins() + if err != nil { + slog.Error("failed to list plugins", "error", err) + writeJSONError(w, http.StatusInternalServerError, "internal_error", "Failed to list plugins") + return + } + if plugins == nil { + plugins = []types.Plugin{} + } + + // Apply search filter + if q := r.URL.Query().Get("q"); q != "" { + plugins = filterPluginsV01(plugins, q) + } + + // Paginate + page, limit := parsePaginationV01(r) + total := len(plugins) + start, end := paginateSlice(total, page, limit) + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(pluginsV01Response{ + Plugins: plugins[start:end], + Metadata: paginationV01Metadata{ + Total: total, + Page: page, + Limit: limit, + }, + }); err != nil { + slog.Error("failed to encode plugins response", "error", err) + } +} + +// getPluginV01 handles GET /registry/{registryName}/v0.1/x/dev.toolhive/plugins/{namespace}/{pluginName} +// +// @Summary Get a registry plugin +// @Description Retrieve a single plugin by its namespace and name from the registry. +// @Tags registry-plugins +// @Produce json +// @Param registryName path string true "Registry name (currently ignored, uses the default provider)" +// @Param namespace path string true "Plugin namespace in reverse-DNS format (e.g. io.github.stacklok)" +// @Param pluginName path string true "Plugin name" +// @Success 200 {object} types.Plugin +// @Failure 404 {object} registryErrorResponse "Plugin not found" +// @Failure 500 {object} registryErrorResponse "Internal server error" +// @Failure 503 {object} registryErrorResponse "Registry authentication required or upstream registry unavailable" +// @Router /registry/{registryName}/v0.1/x/dev.toolhive/plugins/{namespace}/{pluginName} [get] +func getPluginV01(w http.ResponseWriter, r *http.Request) { + provider, ok := getRegistryProvider(w) + if !ok { + return + } + getPluginV01WithProvider(w, r, provider) +} + +// getPluginV01WithProvider contains the Get-plugin logic against an explicit +// provider, separated from getRegistryProvider so tests can inject a mock +// provider without touching the process-wide singleton. +func getPluginV01WithProvider(w http.ResponseWriter, r *http.Request, provider regpkg.Provider) { + namespace := chi.URLParam(r, "namespace") + pluginName := chi.URLParam(r, "pluginName") + + plugin, err := provider.GetPlugin(namespace, pluginName) + if err != nil { + // Map upstream HTTP errors to appropriate responses + var httpErr *api.RegistryHTTPError + if errors.As(err, &httpErr) { + switch httpErr.StatusCode { + case http.StatusNotFound: + writeJSONError(w, http.StatusNotFound, "not_found", "Plugin not found") + return + case http.StatusUnauthorized, http.StatusForbidden: + writeRegistryAuthRequiredError(w) + return + } + } + slog.Error("failed to get plugin", "namespace", namespace, "name", pluginName, "error", err) + writeJSONError(w, http.StatusInternalServerError, "internal_error", "Failed to get plugin") + return + } + if plugin == nil { + writeJSONError(w, http.StatusNotFound, "not_found", "Plugin not found") + return + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(plugin); err != nil { + slog.Error("failed to encode plugin response", "error", err) + } +} + +func filterPluginsV01(plugins []types.Plugin, query string) []types.Plugin { + q := strings.ToLower(query) + result := make([]types.Plugin, 0) + for _, p := range plugins { + if strings.Contains(strings.ToLower(p.Name), q) || + strings.Contains(strings.ToLower(p.Namespace), q) || + strings.Contains(strings.ToLower(p.Description), q) { + result = append(result, p) + } + } + return result +} + +// pluginsV01Response is the response body for the v0.1 plugins list endpoint. +// +// @Description Paginated list of plugins from the registry +type pluginsV01Response struct { + // Plugins is the list of plugins on the current page + Plugins []types.Plugin `json:"plugins"` + // Metadata contains pagination information + Metadata paginationV01Metadata `json:"metadata"` +} diff --git a/pkg/api/v1/registry_v01_plugins_test.go b/pkg/api/v1/registry_v01_plugins_test.go new file mode 100644 index 0000000000..122ac7b34c --- /dev/null +++ b/pkg/api/v1/registry_v01_plugins_test.go @@ -0,0 +1,181 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/go-chi/chi/v5" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + types "github.com/stacklok/toolhive-core/registry/types" + "github.com/stacklok/toolhive/pkg/registry/api" + regmocks "github.com/stacklok/toolhive/pkg/registry/mocks" +) + +func TestFilterPluginsV01(t *testing.T) { + t.Parallel() + + plugins := []types.Plugin{ + {Namespace: "stacklok", Name: "code-review", Description: "Reviews code for issues"}, + {Namespace: "stacklok", Name: "commit", Description: "Creates git commits"}, + {Namespace: "other", Name: "weather", Description: "Weather data"}, + } + + tests := []struct { + query string + wantCount int + }{ + {"code", 1}, + {"CODE", 1}, // case-insensitive + {"Code-Review", 1}, // mixed case + {"stacklok", 2}, + {"weather", 1}, + {"commits", 1}, + {"nonexistent", 0}, + } + + for _, tt := range tests { + t.Run(tt.query, func(t *testing.T) { + t.Parallel() + result := filterPluginsV01(plugins, tt.query) + assert.Len(t, result, tt.wantCount) + }) + } +} + +func TestFilterPluginsV01_EmptyResult_NotNull(t *testing.T) { + t.Parallel() + + plugins := []types.Plugin{ + {Namespace: "stacklok", Name: "test", Description: "A test plugin"}, + } + + result := filterPluginsV01(plugins, "nonexistent") + assert.NotNil(t, result, "Filter result should be empty slice, not nil") + assert.Empty(t, result) + + // Verify JSON encoding produces [] not null + data, err := json.Marshal(result) + require.NoError(t, err) + assert.Equal(t, "[]", string(data)) +} + +func TestRegistryV01Router_ListPlugins(t *testing.T) { + t.Parallel() + + handler := RegistryV01Router() + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/default/v0.1/x/dev.toolhive/plugins") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Contains(t, resp.Header.Get("Content-Type"), "application/json") + + var body pluginsV01Response + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + // Should return plugins from the embedded catalog (may be empty in test env) + assert.NotNil(t, body.Plugins) + assert.GreaterOrEqual(t, body.Metadata.Total, 0) +} + +func TestRegistryV01Router_GetPlugin_NotFound(t *testing.T) { + t.Parallel() + + handler := RegistryV01Router() + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/default/v0.1/x/dev.toolhive/plugins/nonexistent/noplugin") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + assert.Contains(t, resp.Header.Get("Content-Type"), "application/json", + "Error responses should be JSON") + + var body registryErrorResponse + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + assert.Equal(t, "not_found", body.Code) +} + +func TestRegistryV01Router_ListPlugins_PaginationBeyondResults(t *testing.T) { + t.Parallel() + + handler := RegistryV01Router() + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/default/v0.1/x/dev.toolhive/plugins?page=999&limit=10") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body pluginsV01Response + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + assert.Empty(t, body.Plugins, "Page beyond results should return empty plugins") + assert.Equal(t, 999, body.Metadata.Page) + assert.GreaterOrEqual(t, body.Metadata.Total, 0) +} + +// TestRegistryV01Router_GetPlugin_RegistryAuthError confirms that an upstream +// *api.RegistryHTTPError with a 401/403 status is mapped to the structured +// registry_auth_required 503 response (writeRegistryAuthRequiredError), so +// desktop clients can branch on the code. The skills handler has the same +// mapping but lacks this test; plugins get explicit coverage here. +// +// The test wires getPluginV01WithProvider directly against a mock provider via +// a dedicated chi route, avoiding the process-wide default provider singleton +// so it stays deterministic and parallel-safe. +func TestRegistryV01Router_GetPlugin_RegistryAuthError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + statusCode int + }{ + {"unauthorized maps to registry_auth_required", http.StatusUnauthorized}, + {"forbidden maps to registry_auth_required", http.StatusForbidden}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + + provider := regmocks.NewMockProvider(ctrl) + provider.EXPECT().GetPlugin("io.test", "auth-plugin"). + Return(nil, &api.RegistryHTTPError{StatusCode: tt.statusCode, URL: "https://registry.test/plugin"}) + + r := chi.NewRouter() + r.Get("/x/dev.toolhive/plugins/{namespace}/{pluginName}", + func(w http.ResponseWriter, req *http.Request) { + getPluginV01WithProvider(w, req, provider) + }) + srv := httptest.NewServer(r) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/x/dev.toolhive/plugins/io.test/auth-plugin") + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + assert.Contains(t, resp.Header.Get("Content-Type"), "application/json") + + var body registryErrorResponse + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + assert.Equal(t, RegistryAuthRequiredCode, body.Code) + assert.Contains(t, body.Message, "Registry authentication required") + }) + } +} diff --git a/pkg/api/v1/skills.go b/pkg/api/v1/skills.go index b0f58c9001..2e39fea8ca 100644 --- a/pkg/api/v1/skills.go +++ b/pkg/api/v1/skills.go @@ -5,10 +5,12 @@ package v1 import ( "encoding/json" + "errors" "fmt" "net/http" "github.com/go-chi/chi/v5" + "github.com/go-chi/chi/v5/middleware" "github.com/stacklok/toolhive-core/httperr" apierrors "github.com/stacklok/toolhive/pkg/api/errors" @@ -18,25 +20,41 @@ import ( // SkillsRoutes defines the routes for skill management. type SkillsRoutes struct { skillService skills.SkillService + lockService skills.SkillLockService } -// SkillsRouter creates a new router for skill management endpoints. +// SkillsRouter creates a new router for skill management endpoints. If +// skillService's concrete implementation also satisfies skills.SkillLockService +// (as skillsvc.New's does), /sync and /upgrade are served; otherwise both +// return 501. func SkillsRouter(skillService skills.SkillService) http.Handler { routes := SkillsRoutes{ skillService: skillService, } + if lockSvc, ok := skillService.(skills.SkillLockService); ok { + routes.lockService = lockSvc + } + + // Mirrors WorkloadRouter: routes that move OCI artifacts get a timeout + // sized for the transfer, everything else keeps the short one. Without + // the split these routes inherit the flat 60s applied to standard + // routers, which severs a cold artifact pull mid-flight. + stdTimeout := middleware.Timeout(standardRouteTimeout) + longTimeout := middleware.Timeout(longRunningRouteTimeout) r := chi.NewRouter() - r.Get("/", apierrors.ErrorHandler(routes.listSkills)) - r.Post("/", apierrors.ErrorHandler(routes.installSkill)) - r.Delete("/{name}", apierrors.ErrorHandler(routes.uninstallSkill)) - r.Get("/{name}", apierrors.ErrorHandler(routes.getSkillInfo)) - r.Post("/validate", apierrors.ErrorHandler(routes.validateSkill)) - r.Post("/build", apierrors.ErrorHandler(routes.buildSkill)) - r.Post("/push", apierrors.ErrorHandler(routes.pushSkill)) - r.Get("/builds", apierrors.ErrorHandler(routes.listBuilds)) - r.Delete("/builds/{tag}", apierrors.ErrorHandler(routes.deleteBuild)) - r.Get("/content", apierrors.ErrorHandler(routes.getSkillContent)) + r.With(stdTimeout).Get("/", apierrors.ErrorHandler(routes.listSkills)) + r.With(longTimeout).Post("/", apierrors.ErrorHandler(routes.installSkill)) + r.With(stdTimeout).Delete("/{name}", apierrors.ErrorHandler(routes.uninstallSkill)) + r.With(stdTimeout).Get("/{name}", apierrors.ErrorHandler(routes.getSkillInfo)) + r.With(stdTimeout).Post("/validate", apierrors.ErrorHandler(routes.validateSkill)) + r.With(longTimeout).Post("/build", apierrors.ErrorHandler(routes.buildSkill)) + r.With(longTimeout).Post("/push", apierrors.ErrorHandler(routes.pushSkill)) + r.With(stdTimeout).Get("/builds", apierrors.ErrorHandler(routes.listBuilds)) + r.With(stdTimeout).Delete("/builds/{tag}", apierrors.ErrorHandler(routes.deleteBuild)) + r.With(stdTimeout).Get("/content", apierrors.ErrorHandler(routes.getSkillContent)) + r.With(longTimeout).Post("/sync", apierrors.ErrorHandler(routes.syncSkills)) + r.With(longTimeout).Post("/upgrade", apierrors.ErrorHandler(routes.upgradeSkills)) return r } @@ -103,13 +121,14 @@ func (s *SkillsRoutes) installSkill(w http.ResponseWriter, r *http.Request) erro } result, err := s.skillService.Install(r.Context(), skills.InstallOptions{ - Name: req.Name, - Version: req.Version, - Scope: req.Scope, - ProjectRoot: req.ProjectRoot, - Clients: req.Clients, - Force: req.Force, - Group: req.Group, + Name: req.Name, + Version: req.Version, + Scope: req.Scope, + ProjectRoot: req.ProjectRoot, + Clients: req.Clients, + Force: req.Force, + Group: req.Group, + AllowUnsigned: req.AllowUnsigned, }) if err != nil { return err @@ -365,3 +384,94 @@ func (s *SkillsRoutes) getSkillContent(w http.ResponseWriter, r *http.Request) e w.Header().Set("Content-Type", "application/json") return json.NewEncoder(w).Encode(content) } + +// syncSkills restores a project's installed skills to match its lock file. +// +// @Summary Sync project skills from the lock file +// @Description Restore a project's installed skills to match toolhive.lock.yaml +// @Tags skills +// @Accept json +// @Produce json +// @Param request body syncSkillsRequest true "Sync request" +// @Success 200 {object} skills.SyncResult +// @Failure 400 {string} string "Bad Request" +// @Failure 403 {string} string "Forbidden (feature not enabled)" +// @Failure 500 {string} string "Internal Server Error" +// @Failure 501 {string} string "Not Implemented" +// @Router /api/v1beta/skills/sync [post] +func (s *SkillsRoutes) syncSkills(w http.ResponseWriter, r *http.Request) error { + if s.lockService == nil { + return httperr.WithCode(errors.New("skill sync is not supported by this server"), http.StatusNotImplemented) + } + + var req syncSkillsRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + return httperr.WithCode( + fmt.Errorf("invalid request body: %w", err), + http.StatusBadRequest, + ) + } + + result, err := s.lockService.Sync(r.Context(), skills.SyncOptions{ + ProjectRoot: req.ProjectRoot, + Clients: req.Clients, + Prune: req.Prune, + Check: req.Check, + Adopt: req.Adopt, + AllowUnsigned: req.AllowUnsigned, + }) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(result) +} + +// upgradeSkills re-resolves a project's lock entries and installs newer +// content where available. With fail_on_changes set, it evaluates the plan +// and returns the outcomes without installing anything — clients derive the +// freshness verdict from the outcome statuses, mirroring sync's check mode. +// +// @Summary Upgrade project skills +// @Description Re-resolve a project's lock entries and install newer content where available +// @Tags skills +// @Accept json +// @Produce json +// @Param request body upgradeSkillsRequest true "Upgrade request" +// @Success 200 {object} skills.UpgradeResult +// @Failure 400 {string} string "Bad Request" +// @Failure 403 {string} string "Forbidden (feature not enabled)" +// @Failure 404 {string} string "Not Found (a requested name is not in the lock file)" +// @Failure 500 {string} string "Internal Server Error" +// @Failure 501 {string} string "Not Implemented" +// @Router /api/v1beta/skills/upgrade [post] +func (s *SkillsRoutes) upgradeSkills(w http.ResponseWriter, r *http.Request) error { + if s.lockService == nil { + return httperr.WithCode(errors.New("skill upgrade is not supported by this server"), http.StatusNotImplemented) + } + + var req upgradeSkillsRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + return httperr.WithCode( + fmt.Errorf("invalid request body: %w", err), + http.StatusBadRequest, + ) + } + + result, err := s.lockService.Upgrade(r.Context(), skills.UpgradeOptions{ + ProjectRoot: req.ProjectRoot, + Names: req.Names, + Preview: req.Preview, + FailOnChanges: req.FailOnChanges, + AllowRefChange: req.AllowRefChange, + AllowSignerChange: req.AllowSignerChange, + Clients: req.Clients, + }) + if err != nil { + return err + } + + w.Header().Set("Content-Type", "application/json") + return json.NewEncoder(w).Encode(result) +} diff --git a/pkg/api/v1/skills_sync_test.go b/pkg/api/v1/skills_sync_test.go new file mode 100644 index 0000000000..6923a9c1c7 --- /dev/null +++ b/pkg/api/v1/skills_sync_test.go @@ -0,0 +1,134 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + skillsmocks "github.com/stacklok/toolhive/pkg/skills/mocks" +) + +// skillServiceWithSync wraps a mocked SkillService and adds Sync and Upgrade +// methods, so SkillsRouter's opportunistic skills.SkillLockService type +// assertion succeeds — the same shape skillsvc.New's concrete service has. +type skillServiceWithSync struct { + skills.SkillService + syncFn func(ctx context.Context, opts skills.SyncOptions) (*skills.SyncResult, error) + upgradeFn func(ctx context.Context, opts skills.UpgradeOptions) (*skills.UpgradeResult, error) +} + +func (s *skillServiceWithSync) Sync(ctx context.Context, opts skills.SyncOptions) (*skills.SyncResult, error) { + return s.syncFn(ctx, opts) +} + +func (s *skillServiceWithSync) Upgrade(ctx context.Context, opts skills.UpgradeOptions) (*skills.UpgradeResult, error) { + if s.upgradeFn == nil { + return &skills.UpgradeResult{}, nil + } + return s.upgradeFn(ctx, opts) +} + +func TestSyncSkillsEndpoint(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + service skills.SkillService + body string + wantStatus int + wantContains string + }{ + { + name: "successful sync returns 200 with result", + service: &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + syncFn: func(_ context.Context, opts skills.SyncOptions) (*skills.SyncResult, error) { + assert.Equal(t, "/tmp/proj", opts.ProjectRoot) + assert.True(t, opts.Check) + return &skills.SyncResult{AlreadyCurrent: []string{"my-skill"}}, nil + }, + }, + body: `{"project_root":"/tmp/proj","check":true}`, + wantStatus: http.StatusOK, + wantContains: `"my-skill"`, + }, + { + name: "service without Sync support returns 501", + service: skillsmocks.NewMockSkillService(gomock.NewController(t)), + body: `{"project_root":"/tmp/proj"}`, + wantStatus: http.StatusNotImplemented, + }, + { + name: "invalid JSON body returns 400", + service: &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + syncFn: func(context.Context, skills.SyncOptions) (*skills.SyncResult, error) { + t.Fatal("Sync must not be called for an invalid body") + return nil, nil + }, + }, + body: `not json`, + wantStatus: http.StatusBadRequest, + }, + { + name: "sync error propagates with its status code", + service: &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + syncFn: func(context.Context, skills.SyncOptions) (*skills.SyncResult, error) { + return nil, httperr.WithCode(assert.AnError, http.StatusForbidden) + }, + }, + body: `{"project_root":"/tmp/proj"}`, + wantStatus: http.StatusForbidden, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + router := SkillsRouter(tt.service) + req := httptest.NewRequest(http.MethodPost, "/sync", bytes.NewBufferString(tt.body)) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, tt.wantStatus, rec.Code) + if tt.wantContains != "" { + assert.Contains(t, rec.Body.String(), tt.wantContains) + } + }) + } +} + +func TestSyncSkillsResponseIsValidJSON(t *testing.T) { + t.Parallel() + + svc := &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + syncFn: func(context.Context, skills.SyncOptions) (*skills.SyncResult, error) { + return &skills.SyncResult{Installed: []string{"a"}, Drifted: []string{"b"}}, nil + }, + } + router := SkillsRouter(svc) + req := httptest.NewRequest(http.MethodPost, "/sync", bytes.NewBufferString(`{"project_root":"/tmp/proj"}`)) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + var result skills.SyncResult + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &result)) + assert.Equal(t, []string{"a"}, result.Installed) + assert.Equal(t, []string{"b"}, result.Drifted) +} diff --git a/pkg/api/v1/skills_types.go b/pkg/api/v1/skills_types.go index 833a9b0aa5..755b06d2d6 100644 --- a/pkg/api/v1/skills_types.go +++ b/pkg/api/v1/skills_types.go @@ -31,6 +31,10 @@ type installSkillRequest struct { Clients []string `json:"clients,omitempty"` // Force allows overwriting unmanaged skill directories Force bool `json:"force,omitempty"` + // AllowUnsigned permits installing a project-scoped skill without a + // verified signature; the exception is recorded in the project's lock + // file. + AllowUnsigned bool `json:"allow_unsigned,omitempty"` // Group is the group name to add the skill to after installation Group string `json:"group,omitempty"` } @@ -69,6 +73,48 @@ type pushSkillRequest struct { Reference string `json:"reference"` } +// syncSkillsRequest represents the request to sync a project's skills. +// +// @Description Request to restore a project's installed skills to match its lock file +type syncSkillsRequest struct { + // ProjectRoot is the project root path whose lock file should be synced + ProjectRoot string `json:"project_root"` + // Clients lists target client identifiers. Empty means every + // skill-supporting client detected on this host. + Clients []string `json:"clients,omitempty"` + // Prune removes project-scoped skills installed but not present in the lock file + Prune bool `json:"prune,omitempty"` + // Check verifies on-disk content against the lock file without installing or writing anything + Check bool `json:"check,omitempty"` + // Adopt writes lock entries for existing unmanaged project-scope installs + Adopt bool `json:"adopt,omitempty"` + // AllowUnsigned permits adopting skills whose signature state cannot be + // established, recording them as unsigned + AllowUnsigned bool `json:"allow_unsigned,omitempty"` +} + +// upgradeSkillsRequest represents the request to upgrade a project's skills. +// +// @Description Request to re-resolve a project's lock entries and install newer content +type upgradeSkillsRequest struct { + // ProjectRoot is the project root path whose lock file should be upgraded + ProjectRoot string `json:"project_root"` + // Names restricts the upgrade to specific skill names. Empty means every entry. + Names []string `json:"names,omitempty"` + // Preview reports what would change without installing (still fetches to compare digests) + Preview bool `json:"preview,omitempty"` + // FailOnChanges exits with an error when any mutable source would upgrade + FailOnChanges bool `json:"fail_on_changes,omitempty"` + // AllowRefChange permits resolvedReference changes during upgrade + AllowRefChange bool `json:"allow_ref_change,omitempty"` + // AllowSignerChange permits upgrading to an artifact signed by a + // different identity than the recorded one + AllowSignerChange bool `json:"allow_signer_change,omitempty"` + // Clients lists target client identifiers. Empty means every + // skill-supporting client detected on this host. + Clients []string `json:"clients,omitempty"` +} + // buildListResponse represents the response for listing locally-built OCI skill artifacts. // // @Description Response containing a list of locally-built OCI skill artifacts diff --git a/pkg/api/v1/skills_upgrade_test.go b/pkg/api/v1/skills_upgrade_test.go new file mode 100644 index 0000000000..e9f5d6192c --- /dev/null +++ b/pkg/api/v1/skills_upgrade_test.go @@ -0,0 +1,119 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + skillsmocks "github.com/stacklok/toolhive/pkg/skills/mocks" +) + +func TestUpgradeSkillsEndpoint(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + service skills.SkillService + body string + wantStatus int + wantContains string + }{ + { + name: "successful upgrade returns 200 with outcomes", + service: &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + upgradeFn: func(_ context.Context, opts skills.UpgradeOptions) (*skills.UpgradeResult, error) { + assert.Equal(t, "/tmp/proj", opts.ProjectRoot) + assert.Equal(t, []string{"my-skill"}, opts.Names) + assert.True(t, opts.AllowRefChange) + return &skills.UpgradeResult{Outcomes: []skills.UpgradeOutcome{ + {Name: "my-skill", Status: skills.UpgradeStatusUpgraded}, + }}, nil + }, + }, + body: `{"project_root":"/tmp/proj","names":["my-skill"],"allow_ref_change":true}`, + wantStatus: http.StatusOK, + wantContains: `"upgraded"`, + }, + { + name: "service without Upgrade support returns 501", + service: skillsmocks.NewMockSkillService(gomock.NewController(t)), + body: `{"project_root":"/tmp/proj"}`, + wantStatus: http.StatusNotImplemented, + }, + { + name: "invalid JSON body returns 400", + service: &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + upgradeFn: func(context.Context, skills.UpgradeOptions) (*skills.UpgradeResult, error) { + t.Fatal("Upgrade must not be called for an invalid body") + return nil, nil + }, + }, + body: `not json`, + wantStatus: http.StatusBadRequest, + }, + { + name: "fail-on-changes conflict propagates as 409", + service: &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + upgradeFn: func(context.Context, skills.UpgradeOptions) (*skills.UpgradeResult, error) { + return nil, httperr.WithCode(assert.AnError, http.StatusConflict) + }, + }, + body: `{"project_root":"/tmp/proj","fail_on_changes":true}`, + wantStatus: http.StatusConflict, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + router := SkillsRouter(tt.service) + req := httptest.NewRequest(http.MethodPost, "/upgrade", bytes.NewBufferString(tt.body)) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, tt.wantStatus, rec.Code) + if tt.wantContains != "" { + assert.Contains(t, rec.Body.String(), tt.wantContains) + } + }) + } +} + +func TestUpgradeSkillsResponseIsValidJSON(t *testing.T) { + t.Parallel() + + svc := &skillServiceWithSync{ + SkillService: skillsmocks.NewMockSkillService(gomock.NewController(t)), + upgradeFn: func(context.Context, skills.UpgradeOptions) (*skills.UpgradeResult, error) { + return &skills.UpgradeResult{Outcomes: []skills.UpgradeOutcome{ + {Name: "a", Status: skills.UpgradeStatusUpToDate}, + }}, nil + }, + } + router := SkillsRouter(svc) + req := httptest.NewRequest(http.MethodPost, "/upgrade", bytes.NewBufferString(`{"project_root":"/tmp/proj"}`)) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + var result skills.UpgradeResult + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &result)) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, "a", result.Outcomes[0].Name) +} diff --git a/pkg/audit/auditor.go b/pkg/audit/auditor.go index 27e4efccf3..e19cee515a 100644 --- a/pkg/audit/auditor.go +++ b/pkg/audit/auditor.go @@ -16,13 +16,14 @@ import ( "strings" "time" + coreaudit "github.com/stacklok/toolhive-core/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/types" ) // LevelAudit is a custom audit log level - between Info and Warn -const LevelAudit = slog.Level(2) +const LevelAudit = coreaudit.LevelAudit // contextKey is an unexported type for context keys to avoid collisions type contextKey struct{} @@ -50,27 +51,7 @@ func BackendInfoFromContext(ctx context.Context) (*BackendInfo, bool) { // NewAuditLogger creates a new structured audit logger that writes to the specified writer. func NewAuditLogger(w io.Writer) *slog.Logger { - if w == nil { - w = os.Stdout - } - - handler := slog.NewJSONHandler(w, &slog.HandlerOptions{ - Level: LevelAudit, - ReplaceAttr: func(_ []string, a slog.Attr) slog.Attr { - // Replace the custom audit level with "AUDIT" string for better - // compatibility with log aggregation systems (Loki, Elasticsearch, etc.) - // that expect standard level names. This prevents audit events from - // appearing as "INFO+2" which breaks level-based filtering. - if a.Key == slog.LevelKey { - if level, ok := a.Value.Any().(slog.Level); ok && level == LevelAudit { - a.Value = slog.StringValue("AUDIT") - } - } - return a - }, - }) - - return slog.New(handler) + return coreaudit.NewAuditLogger(w) } // Auditor handles audit logging for HTTP requests. @@ -172,6 +153,12 @@ func (rw *responseWriter) Flush() { } } +// Unwrap exposes the underlying ResponseWriter so http.ResponseController +// can reach interfaces this wrapper does not re-implement (e.g. SetWriteDeadline). +func (rw *responseWriter) Unwrap() http.ResponseWriter { + return rw.ResponseWriter +} + // isMCPStreamOpenRequest returns true only for MCP "stream" opens: // - SSE transport's SSE endpoint (GET + Accept: text/event-stream) // - Streamable HTTP's GET stream (same header pattern) @@ -187,30 +174,62 @@ func (*Auditor) isMCPStreamOpenRequest(r *http.Request) bool { return strings.Contains(strings.ToLower(accept), "text/event-stream") } +// ensureAuditContext injects the mutable carriers the auditor reads after the +// inner chain returns: BackendInfo (backend routing), an auth.IdentityHolder +// (identity attached by an auth middleware running INSIDE audit), an +// mcp.ParsedRequestHolder (parsed MCP data from a parser running INSIDE +// audit), and an mcp.AuthzDenialMarker (pre-parse refusals flagged by the +// authz middleware running INSIDE audit). Each is only injected when absent +// so nested auditors share carriers. +func ensureAuditContext(r *http.Request) *http.Request { + ctx := r.Context() + changed := false + if _, ok := BackendInfoFromContext(ctx); !ok { + ctx = WithBackendInfo(ctx, &BackendInfo{}) + changed = true + } + if _, ok := auth.IdentityHolderFromContext(ctx); !ok { + ctx = auth.WithIdentityHolder(ctx, &auth.IdentityHolder{}) + changed = true + } + if _, ok := mcp.ParsedRequestHolderFromContext(ctx); !ok { + ctx = mcp.WithParsedRequestHolder(ctx, &mcp.ParsedRequestHolder{}) + changed = true + } + if _, ok := mcp.AuthzDenialMarkerFromContext(ctx); !ok { + ctx = mcp.WithAuthzDenialMarker(ctx, &mcp.AuthzDenialMarker{}) + changed = true + } + if !changed { + return r + } + return r.WithContext(ctx) +} + // Middleware creates an HTTP middleware that logs audit events. func (a *Auditor) Middleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Handle SSE endpoints specially - log the connection event immediately - // since SSE connections are long-lived and don't follow normal request/response pattern - if a.isMCPStreamOpenRequest(r) { - // Log SSE connection event immediately - a.logSSEConnectionEvent(r) + r = ensureAuditContext(r) - // Pass through to SSE handler without waiting - next.ServeHTTP(w, r) + // Handle MCP stream opens (SSE endpoint, streamable GET) specially: + // these connections are long-lived, so instead of waiting for the + // response to complete, the connection event is logged on the FIRST + // write. By then any inner auth middleware has run, so the event + // carries the authenticated identity (or records the 401/403 denial). + if a.isMCPStreamOpenRequest(r) { + sw := &streamOpenWriter{ResponseWriter: w, auditor: a, req: r} + // Deferred BEFORE ServeHTTP so a panic in an inner handler still + // produces the connection event during unwinding — some chains run + // the recovery middleware OUTSIDE audit (e.g. the vMCP Serve path), + // which would otherwise swallow the event entirely. If the stream + // already logged on first write this is a no-op. + defer sw.logOnce(http.StatusOK) + next.ServeHTTP(sw, r) return } startTime := time.Now() - // Add BackendInfo to context if not already present - // (backend enrichment middleware may have already added it) - if _, ok := BackendInfoFromContext(r.Context()); !ok { - backendInfo := &BackendInfo{} - ctx := WithBackendInfo(r.Context(), backendInfo) - r = r.WithContext(ctx) - } - // Capture request data if configured var requestData []byte if a.config.IncludeRequestData && r.Body != nil { @@ -262,6 +281,14 @@ func (a *Auditor) logAuditEvent(r *http.Request, rw *responseWriter, requestData // Determine outcome based on status code outcome := a.determineOutcome(rw.statusCode) + // A refusal by the authz middleware before message-level authorization + // could run (e.g. a non-JSON POST carrying a smuggled JSON-RPC body) + // writes a 400, which determineOutcome maps to a generic failure. The + // marker reclassifies it as a denial so blocked sweeps are alertable. + if marker, ok := mcp.AuthzDenialMarkerFromContext(r.Context()); ok && marker.Denied { + outcome = OutcomeDenied + } + // When HTTP status indicates success, check for JSON-RPC errors // hidden inside HTTP 200 responses. var mcpResponse *mcp.ParsedMCPResponse @@ -289,6 +316,9 @@ func (a *Auditor) logAuditEvent(r *http.Request, rw *responseWriter, requestData // Create the audit event event := NewAuditEvent(eventType, source, outcome, subjects, component) + // Attach the RFC 8693 delegation chain, if the caller used a delegated token + event.WithDelegationChain(a.extractDelegationChain(r)) + // Add target information target := a.extractTarget(r, eventType) if len(target) > 0 { @@ -319,10 +349,35 @@ func (a *Auditor) logAuditEvent(r *http.Request, rw *responseWriter, requestData event.LogTo(r.Context(), a.auditLogger, LevelAudit) } +// mcpMethodFor returns the parsed MCP method for the request, whether the +// parser ran outside audit (context value) or inside it (holder filled by +// the parser and read back after the inner chain returns). +func mcpMethodFor(r *http.Request) string { + if m := mcp.GetMCPMethod(r.Context()); m != "" { + return m + } + if holder, ok := mcp.ParsedRequestHolderFromContext(r.Context()); ok && holder.Parsed != nil { + return holder.Parsed.Method + } + return "" +} + +// mcpResourceIDFor returns the parsed MCP resource ID for the request, with +// the same context-then-holder fallback as mcpMethodFor. +func mcpResourceIDFor(r *http.Request) string { + if id := mcp.GetMCPResourceID(r.Context()); id != "" { + return id + } + if holder, ok := mcp.ParsedRequestHolderFromContext(r.Context()); ok && holder.Parsed != nil { + return holder.Parsed.ResourceID + } + return "" +} + // determineEventType determines the event type based on the HTTP request. func (a *Auditor) determineEventType(r *http.Request) string { - // First, try to get the parsed MCP method from context - if mcpMethod := mcp.GetMCPMethod(r.Context()); mcpMethod != "" { + // First, try to get the parsed MCP method + if mcpMethod := mcpMethodFor(r); mcpMethod != "" { return a.mapMCPMethodToEventType(mcpMethod) } @@ -488,12 +543,74 @@ func extractSubjectsFromIdentity(identity *auth.Identity) map[string]string { return subjects } +// extractDelegationChainFromIdentity returns the RFC 8693 delegation chain +// carried by the identity, bound to maxDepth, or nil when there is none. +// The identity layer (auth.claimsToIdentity) always parses at +// auth.DefaultMaxDelegationDepth, so a larger maxDepth requires re-parsing the +// raw "act" claim — rebinding an already-parsed chain can only shrink it. The +// re-parse is only trusted when it yields at least as many hops as are +// already parsed, so it can widen the chain but never narrow it. +func extractDelegationChainFromIdentity(identity *auth.Identity, maxDepth int) *coreaudit.DelegationChain { + if identity == nil { + return nil + } + rawAct := identity.Claims["act"] + chain := identity.DelegationChain + existingHops := 0 + if chain != nil { + existingHops = len(chain.Chain) + } + if rawAct != nil && (chain.IsZero() || + (chain.Truncated && maxDepth > existingHops)) { + parsed := coreaudit.ParseDelegationChain(rawAct, maxDepth) + if !parsed.IsZero() && len(parsed.Chain) >= existingHops { + return parsed + } + } + if !chain.IsZero() { + return rebindDelegationChain(chain, maxDepth) + } + return nil +} + +// rebindDelegationChain applies maxDepth to an already-parsed chain: hops +// beyond the cap are dropped (keeping the outermost hops, so Chain[0] — the +// current actor — is preserved), the chain is marked truncated, and the +// omitted count accounts for both hops truncated at parse time and hops +// dropped by this re-binding. +func rebindDelegationChain(chain *coreaudit.DelegationChain, maxDepth int) *coreaudit.DelegationChain { + if maxDepth <= 0 { + maxDepth = auth.DefaultMaxDelegationDepth + } + bound := *chain + if len(bound.Chain) > maxDepth { + bound.Omitted += len(bound.Chain) - maxDepth + bound.Chain = bound.Chain[:maxDepth] + bound.Truncated = true + } + return &bound +} + +// identityFromRequest resolves the authenticated identity for the request. +// The context value is present when an auth middleware runs OUTSIDE audit; +// the holder covers the audit-wraps-auth arrangement, where the identity +// attached for inner handlers is published back up via auth.WithIdentity. +func identityFromRequest(r *http.Request) (*auth.Identity, bool) { + if identity, ok := auth.IdentityFromContext(r.Context()); ok { + return identity, true + } + if holder, hok := auth.IdentityHolderFromContext(r.Context()); hok && holder.Identity != nil { + return holder.Identity, true + } + return nil, false +} + // extractSubjects extracts subject information from the HTTP request. func (*Auditor) extractSubjects(r *http.Request) map[string]string { subjects := make(map[string]string) - // Extract user information from Identity - if identity, ok := auth.IdentityFromContext(r.Context()); ok { + // Extract user information from Identity. + if identity, ok := identityFromRequest(r); ok { subjects = extractSubjectsFromIdentity(identity) } @@ -505,6 +622,16 @@ func (*Auditor) extractSubjects(r *http.Request) map[string]string { return subjects } +// extractDelegationChain extracts the RFC 8693 delegation chain from the +// HTTP request's identity, or nil when there is no identity or no chain. +func (a *Auditor) extractDelegationChain(r *http.Request) *coreaudit.DelegationChain { + identity, ok := identityFromRequest(r) + if !ok { + return nil + } + return extractDelegationChainFromIdentity(identity, a.config.MaxDelegationDepthOrDefault()) +} + // determineComponent determines the component name based on the request. func (a *Auditor) determineComponent(_ *http.Request) string { // Use the component from configuration if set @@ -524,12 +651,12 @@ func (*Auditor) extractTarget(r *http.Request, eventType string) map[string]stri target[TargetKeyMethod] = r.Method // Add MCP method if available from parsed data - if mcpMethod := mcp.GetMCPMethod(r.Context()); mcpMethod != "" { + if mcpMethod := mcpMethodFor(r); mcpMethod != "" { target[TargetKeyMethod] = mcpMethod } // Add resource ID if available from parsed data - if resourceID := mcp.GetMCPResourceID(r.Context()); resourceID != "" { + if resourceID := mcpResourceIDFor(r); resourceID != "" { target[TargetKeyName] = resourceID } @@ -615,8 +742,68 @@ func (a *Auditor) addEventData(event *AuditEvent, _ *http.Request, rw *responseW } } +// streamOpenWriter wraps the ResponseWriter for MCP stream-open requests +// (SSE endpoint, streamable GET). It logs the connection audit event exactly +// once, on the first WriteHeader/Write, so the event reflects the actual +// outcome (200 stream established, 401/403 denied by inner middleware) and +// carries the identity the inner auth middleware attached by that point. +type streamOpenWriter struct { + http.ResponseWriter + auditor *Auditor + req *http.Request + logged bool +} + +func (sw *streamOpenWriter) WriteHeader(statusCode int) { + // Informational (1xx) responses are not the final status — don't consume + // the one-shot connection event on them. + if statusCode >= http.StatusOK { + sw.logOnce(statusCode) + } + sw.ResponseWriter.WriteHeader(statusCode) +} + +func (sw *streamOpenWriter) Write(data []byte) (int, error) { + // An implicit WriteHeader(200) happens on first Write. + sw.logOnce(http.StatusOK) + return sw.ResponseWriter.Write(data) +} + +// Flush implements http.Flusher if the underlying ResponseWriter supports it. +// A flush commits the response headers with an implicit 200, so it counts as +// the stream being established — log the connection event here too, otherwise +// a handler that flushes before its first write would delay (or, on a stream +// that only ever flushes, lose) the event. +func (sw *streamOpenWriter) Flush() { + sw.logOnce(http.StatusOK) + if flusher, ok := sw.ResponseWriter.(http.Flusher); ok { + flusher.Flush() + } +} + +// Unwrap exposes the underlying ResponseWriter so http.ResponseController +// can reach interfaces this wrapper does not re-implement (e.g. SetWriteDeadline). +func (sw *streamOpenWriter) Unwrap() http.ResponseWriter { + return sw.ResponseWriter +} + +// logOnce logs the stream connection event with the given status on the first +// call; subsequent calls are no-ops. +func (sw *streamOpenWriter) logOnce(statusCode int) { + if sw.logged { + return + } + sw.logged = true + sw.auditor.logSSEConnectionEvent(sw.req, statusCode) +} + // logSSEConnectionEvent logs an audit event for SSE connection initiation. -func (a *Auditor) logSSEConnectionEvent(r *http.Request) { +func (a *Auditor) logSSEConnectionEvent(r *http.Request, statusCode int) { + // Honor the configured event-type filter, like logAuditEvent does. + if !a.config.ShouldAuditEvent(EventTypeSSEConnection) { + return + } + // Extract source information source := a.extractSource(r) @@ -627,7 +814,10 @@ func (a *Auditor) logSSEConnectionEvent(r *http.Request) { component := a.determineComponent(r) // Create the audit event for SSE connection - event := NewAuditEvent(EventTypeSSEConnection, source, OutcomeSuccess, subjects, component) + event := NewAuditEvent(EventTypeSSEConnection, source, a.determineOutcome(statusCode), subjects, component) + + // Attach the RFC 8693 delegation chain, if the caller used a delegated token + event.WithDelegationChain(a.extractDelegationChain(r)) // Add target information target := map[string]string{ diff --git a/pkg/audit/auditor_test.go b/pkg/audit/auditor_test.go index c874729e6a..ebc2ff828f 100644 --- a/pkg/audit/auditor_test.go +++ b/pkg/audit/auditor_test.go @@ -19,7 +19,9 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + coreaudit "github.com/stacklok/toolhive-core/audit" "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/mcp" ) func TestNewAuditor(t *testing.T) { @@ -401,6 +403,7 @@ func TestDetermineOutcome(t *testing.T) { {403, OutcomeDenied}, {400, OutcomeFailure}, {404, OutcomeFailure}, + {429, OutcomeFailure}, // Rate limiting is a load condition, not an identity/policy denial {499, OutcomeFailure}, {500, OutcomeError}, {503, OutcomeError}, @@ -579,6 +582,302 @@ func TestExtractSubjects(t *testing.T) { assert.Equal(t, "anonymous", subjects[SubjectKeyUser]) }) + + t.Run("with delegation chain", func(t *testing.T) { + t.Parallel() + claims := jwt.MapClaims{ + "sub": "user123", + "name": "John Doe", + "act": map[string]any{ + "sub": "agent-1", + "act": map[string]any{"sub": "agent-2"}, + }, + } + + req := httptest.NewRequest("GET", "/test", nil) + parsed := coreaudit.ParseDelegationChain(claims["act"], auth.DefaultMaxDelegationDepth) + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: claims["sub"].(string), + Name: claims["name"].(string), + Claims: claims, + DelegationChain: parsed, + }, + } + ctx := auth.WithIdentity(req.Context(), identity) + req = req.WithContext(ctx) + + subjects := auditor.extractSubjects(req) + assert.Equal(t, "user123", subjects[SubjectKeyUserID]) + assert.Equal(t, "John Doe", subjects[SubjectKeyUser]) + + chain := auditor.extractDelegationChain(req) + require.NotNil(t, chain) + require.Len(t, chain.Chain, 2) + assert.False(t, chain.Truncated) + assert.Equal(t, "agent-1", chain.Chain[0].Subject) + assert.Equal(t, "agent-2", chain.Chain[1].Subject) + }) + + t.Run("delegation chain respects configured max depth", func(t *testing.T) { + t.Parallel() + maxDepth := 1 + depthAuditor, err := NewAuditorWithTransport(&Config{MaxDelegationDepth: &maxDepth}, "sse") + require.NoError(t, err) + + claims := jwt.MapClaims{ + "sub": "user123", + "act": map[string]any{ + "sub": "agent-1", + "act": map[string]any{"sub": "agent-2"}, + }, + } + + req := httptest.NewRequest("GET", "/test", nil) + parsed := coreaudit.ParseDelegationChain(claims["act"], auth.DefaultMaxDelegationDepth) + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: claims, + DelegationChain: parsed, + }, + } + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + chain := depthAuditor.extractDelegationChain(req) + require.NotNil(t, chain) + require.Len(t, chain.Chain, 1, "chain must be capped at the configured max depth") + assert.Equal(t, "agent-1", chain.Chain[0].Subject) + assert.True(t, chain.Truncated, "dropped trailing actors must mark the chain truncated") + assert.Equal(t, 1, chain.Omitted, "one dropped actor must be reported") + }) + + t.Run("delegation chain widens through configured max depth", func(t *testing.T) { + t.Parallel() + maxDepth := 25 + depthAuditor, err := NewAuditorWithTransport(&Config{MaxDelegationDepth: &maxDepth}, "sse") + require.NoError(t, err) + + subs := []string{ + "agent-1", "agent-2", "agent-3", "agent-4", "agent-5", + "agent-6", "agent-7", "agent-8", "agent-9", "agent-10", + "agent-11", "agent-12", "agent-13", "agent-14", + } + claims := jwt.MapClaims{ + "sub": "user123", + "act": nestedActClaim(subs...), + } + + req := httptest.NewRequest("GET", "/test", nil) + // Shaped like what auth.claimsToIdentity produces for a 14-hop token: + // the identity-layer parse caps at auth.DefaultMaxDelegationDepth (10), + // truncating the chain even though the raw "act" claim has all 14 hops. + parsed := coreaudit.ParseDelegationChain(claims["act"], auth.DefaultMaxDelegationDepth) + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: claims, + DelegationChain: parsed, + }, + } + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + chain := depthAuditor.extractDelegationChain(req) + require.NotNil(t, chain) + require.Len(t, chain.Chain, len(subs), "configured depth must recover actors dropped by the identity-layer parse") + assert.False(t, chain.Truncated) + }) + + t.Run("no identity yields nil delegation chain", func(t *testing.T) { + t.Parallel() + req := httptest.NewRequest("GET", "/test", nil) + + assert.Nil(t, auditor.extractDelegationChain(req)) + }) +} + +func TestExtractDelegationChainFromIdentity(t *testing.T) { + t.Parallel() + + t.Run("nil identity", func(t *testing.T) { + t.Parallel() + assert.Nil(t, extractDelegationChainFromIdentity(nil, auth.DefaultMaxDelegationDepth)) + }) + + t.Run("identity without chain", func(t *testing.T) { + t.Parallel() + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "user123"}} + assert.Nil(t, extractDelegationChainFromIdentity(identity, auth.DefaultMaxDelegationDepth)) + }) + + t.Run("programmatic chain without raw act claim is re-bounded to max depth", func(t *testing.T) { + t.Parallel() + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + DelegationChain: &coreaudit.DelegationChain{ + Chain: []coreaudit.DelegatedActor{ + {Subject: "agent-1"}, + {Subject: "agent-2"}, + }, + }, + }, + } + chain := extractDelegationChainFromIdentity(identity, 1) + require.NotNil(t, chain) + require.Len(t, chain.Chain, 1, "the parsed chain must be capped at maxDepth") + assert.Equal(t, "agent-1", chain.Chain[0].Subject) + assert.True(t, chain.Truncated) + assert.Equal(t, 1, chain.Omitted) + }) + + t.Run("raw act claim without parsed chain is parsed with max depth", func(t *testing.T) { + t.Parallel() + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: map[string]any{ + "act": map[string]any{ + "sub": "agent-1", + "act": map[string]any{"sub": "agent-2"}, + }, + }, + }, + } + chain := extractDelegationChainFromIdentity(identity, 1) + require.NotNil(t, chain) + require.Len(t, chain.Chain, 1, "the raw act claim must be parsed with the depth cap") + assert.Equal(t, "agent-1", chain.Chain[0].Subject) + assert.True(t, chain.Truncated) + assert.Equal(t, 1, chain.Omitted) + }) + + t.Run("no chain and no act claim yields nil", func(t *testing.T) { + t.Parallel() + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: map[string]any{"sub": "user123"}, + }, + } + assert.Nil(t, extractDelegationChainFromIdentity(identity, auth.DefaultMaxDelegationDepth)) + }) + + t.Run("malformed act with zero actors is not discarded (re-parse guard)", func(t *testing.T) { + t.Parallel() + // No pre-parsed DelegationChain, so the ONLY path that can produce a + // result is the raw-act re-parse guard + // (coreaudit.ParseDelegationChain(rawAct, maxDepth)); the trailing + // "chain != nil" branch is unreachable here (chain is nil), so this + // isolates the re-parse guard's own Malformed handling. + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: map[string]any{"act": "not-an-object"}, + }, + } + chain := extractDelegationChainFromIdentity(identity, auth.DefaultMaxDelegationDepth) + require.NotNil(t, chain, "a malformed-but-actorless chain must still surface, not be swallowed by the zero-chain guards") + assert.True(t, chain.Malformed) + assert.Empty(t, chain.Chain) + }) + + t.Run("malformed act with zero actors is not discarded (trailing raw-claim branch)", func(t *testing.T) { + t.Parallel() + // No raw "act" claim is present (e.g. Claims not carrying the raw + // value, or already stripped), so the re-parse guard's rawAct != nil + // check is skipped entirely: this exercises the OTHER guard, the + // trailing "chain != nil && ..." branch that rebinds an + // already-parsed chain. + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + DelegationChain: &coreaudit.DelegationChain{ + Malformed: true, + }, + }, + } + chain := extractDelegationChainFromIdentity(identity, auth.DefaultMaxDelegationDepth) + require.NotNil(t, chain, "a malformed-but-actorless chain must still surface via the trailing rebind branch") + assert.True(t, chain.Malformed) + assert.Empty(t, chain.Chain) + }) + + // The following two cases share an identity shaped exactly like what + // auth.claimsToIdentity produces for a 14-hop token: the identity-layer + // parse always caps at auth.DefaultMaxDelegationDepth (10), so + // DelegationChain holds only the first 10 actors, truncated, with the + // full 14-hop chain still available in the raw "act" claim. + subs := []string{ + "agent-1", "agent-2", "agent-3", "agent-4", "agent-5", + "agent-6", "agent-7", "agent-8", "agent-9", "agent-10", + "agent-11", "agent-12", "agent-13", "agent-14", + } + truncatedIdentity := func() *auth.Identity { + actors := make([]coreaudit.DelegatedActor, auth.DefaultMaxDelegationDepth) + for i := range actors { + actors[i] = coreaudit.DelegatedActor{Subject: subs[i]} + } + return &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + DelegationChain: &coreaudit.DelegationChain{ + Chain: actors, + Truncated: true, + Omitted: len(subs) - auth.DefaultMaxDelegationDepth, + }, + Claims: map[string]any{"act": nestedActClaim(subs...)}, + }, + } + } + + t.Run("configured depth wider than identity-layer parse recovers dropped actors", func(t *testing.T) { + t.Parallel() + chain := extractDelegationChainFromIdentity(truncatedIdentity(), 25) + require.NotNil(t, chain) + require.Len(t, chain.Chain, len(subs)) + assert.False(t, chain.Truncated) + assert.Zero(t, chain.Omitted) + assert.Equal(t, subs[0], chain.Chain[0].Subject) + assert.Equal(t, subs[len(subs)-1], chain.Chain[len(subs)-1].Subject) + }) + + t.Run("configured depth narrower than identity-layer parse still rebinds", func(t *testing.T) { + t.Parallel() + chain := extractDelegationChainFromIdentity(truncatedIdentity(), 3) + require.NotNil(t, chain) + require.Len(t, chain.Chain, 3) + assert.True(t, chain.Truncated) + // 4 actors already dropped by the identity-layer parse (14 - 10), + // plus 7 more dropped by the narrower rebind (10 - 3). + assert.Equal(t, 11, chain.Omitted) + }) + + t.Run("truncated chain without raw act claim falls back to rebind", func(t *testing.T) { + t.Parallel() + identity := truncatedIdentity() + identity.Claims = nil // widening requested, but no raw claim to re-parse + + chain := extractDelegationChainFromIdentity(identity, 25) + require.NotNil(t, chain) + require.Len(t, chain.Chain, auth.DefaultMaxDelegationDepth) + assert.True(t, chain.Truncated, "without the raw claim, the identity-layer truncation cannot be recovered") + assert.Equal(t, 4, chain.Omitted) + }) +} + +// nestedActClaim builds a raw RFC 8693 "act" claim nesting subs in order, +// outermost first, matching the shape coreaudit.ParseDelegationChain expects. +func nestedActClaim(subs ...string) map[string]any { + var current map[string]any + for i := len(subs) - 1; i >= 0; i-- { + m := map[string]any{"sub": subs[i]} + if current != nil { + m["act"] = current + } + current = m + } + return current } func TestDetermineComponent(t *testing.T) { @@ -1063,3 +1362,478 @@ func TestAuditLoggerLevelFormat(t *testing.T) { assert.NotContains(t, logOutput, `"level":"AUDIT"`) }) } + +// newBufferAuditor returns an Auditor writing audit events to the returned +// buffer, for tests asserting on emitted events. +func newBufferAuditor(t *testing.T) (*Auditor, *bytes.Buffer) { + t.Helper() + return newBufferAuditorWithConfig(t, &Config{Component: "test"}) +} + +// newBufferAuditorWithConfig is newBufferAuditor for tests that need a +// non-default Config (e.g. MaxDelegationDepth). +func newBufferAuditorWithConfig(t *testing.T, cfg *Config) (*Auditor, *bytes.Buffer) { + t.Helper() + auditor, err := NewAuditorWithTransport(cfg, "streamable-http") + require.NoError(t, err) + var logBuf bytes.Buffer + auditor.auditLogger = NewAuditLogger(&logBuf) + return auditor, &logBuf +} + +// decodeAuditEvents parses the newline-delimited JSON events in buf. +func decodeAuditEvents(t *testing.T, buf *bytes.Buffer) []map[string]any { + t.Helper() + var events []map[string]any + for _, line := range strings.Split(strings.TrimSpace(buf.String()), "\n") { + if line == "" { + continue + } + var event map[string]any + require.NoError(t, json.Unmarshal([]byte(line), &event), "audit log line is not JSON: %s", line) + events = append(events, event) + } + return events +} + +// newToolsCallRequest builds a POST tools/call request suitable for the parser. +func newToolsCallRequest() *http.Request { + req := httptest.NewRequest("POST", "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"target_tool","arguments":{}}}`)) + req.Header.Set("Content-Type", "application/json") + return req +} + +// TestMiddlewareAuditsInnerChainOutcomes pins the audit-wraps-chain +// arrangement: auth and the MCP parser run INSIDE the audit middleware, and +// audit reads the identity and parsed MCP data back through the holder +// carriers it injects (auth.IdentityHolder, mcp.ParsedRequestHolder). This is +// what lets audit record rejections from any inner middleware — auth 401s, +// webhook denials, authz 403s — instead of only requests that reached its old +// position deep in the chain. +func TestMiddlewareAuditsInnerChainOutcomes(t *testing.T) { + t.Parallel() + + // innerAuth emulates an auth middleware running inside audit: it attaches + // the identity exactly as production middlewares do (via auth.WithIdentity, + // which also fills the IdentityHolder injected by audit). + innerAuth := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user-123", + Name: "Test User", + Claims: jwt.MapClaims{"sub": "user-123"}, + }} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) + } + + t.Run("identity attached by inner auth reaches the audit event", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + auditor.Middleware(innerAuth(handler)).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, OutcomeSuccess, events[0]["outcome"]) + subjects, ok := events[0]["subjects"].(map[string]any) + require.True(t, ok, "event must carry subjects") + assert.Equal(t, "user-123", subjects[SubjectKeyUserID], + "identity attached by an inner auth middleware must reach the audit event via the holder") + }) + + t.Run("inner 401 rejection is audited as denied with anonymous subject", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + reject := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "invalid token", http.StatusUnauthorized) + }) + auditor.Middleware(reject).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1, "an authentication failure must still produce an audit event") + assert.Equal(t, OutcomeDenied, events[0]["outcome"]) + subjects, ok := events[0]["subjects"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "anonymous", subjects[SubjectKeyUser], + "no identity exists when authentication fails") + }) + + t.Run("authz non-JSON refusal is audited as denied", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + // Mirror the authz middleware's refusal path: flag the injected + // marker and write the 400, as the explicit early return does. + refuse := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if marker, ok := mcp.AuthzDenialMarkerFromContext(r.Context()); ok { + marker.Denied = true + } + http.Error(w, "Invalid or malformed MCP request", http.StatusBadRequest) + }) + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"target_tool","arguments":{}}}`)) + req.Header.Set("Content-Type", "text/plain") + auditor.Middleware(refuse).ServeHTTP(httptest.NewRecorder(), req) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, OutcomeDenied, events[0]["outcome"], + "an authz refusal before message-level authorization must audit as denied, not a generic 400 failure") + }) + + t.Run("event type comes from inner parser via holder", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + // The parser runs INSIDE audit, as in the real chains. + auditor.Middleware(mcp.ParsingMiddleware(handler)).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, EventTypeMCPToolCall, events[0]["type"], + "parsed MCP data from an inner parser must drive the event type via the holder") + target, ok := events[0]["target"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "target_tool", target[TargetKeyName]) + }) +} + +// TestLogAuditEventDelegationChain pins the emitted delegation shape on the +// primary logAuditEvent path (every ordinary POST). The SSE path already has +// an equivalent pin (see "delegated token stream open carries the delegation +// chain" in TestStreamOpenAuditEvents) and workflow auditing has its own +// (workflow_auditor_test.go), but logAuditEvent had none, so a pkg/audit +// regression here was only ever caught by running pkg/authserver's suite. +// What this test uniquely pins beyond those two siblings is that the POST +// path honours Config.MaxDelegationDepth via +// Config.MaxDelegationDepthOrDefault at this specific call site. +func TestLogAuditEventDelegationChain(t *testing.T) { + t.Parallel() + + // twoHopIdentity builds an auth.Identity carrying a hand-built two-hop + // RFC 8693 chain, outermost (most recent) first: agent-2 delegated from + // agent-1, with agent-1 also carrying its issuer to pin the promoted + // per-hop "iss" field. + twoHopIdentity := func() *auth.Identity { + act := map[string]any{ + "sub": "agent-2", + "act": map[string]any{"sub": "agent-1", "iss": "https://issuer.example"}, + } + chain := coreaudit.ParseDelegationChain(act, auth.DefaultMaxDelegationDepth) + return &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user-123", + Claims: map[string]any{"act": act}, + DelegationChain: chain, + }} + } + + t.Run("full documented shape: two hops, one carrying iss", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + identity := twoHopIdentity() + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = auth.WithIdentity(r.Context(), identity) + w.WriteHeader(http.StatusOK) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + + chain, ok := events[0]["delegation"].(map[string]any) + require.True(t, ok, "the POST/logAuditEvent path must carry the delegation chain") + assert.Equal(t, false, chain["truncated"]) + assert.Equal(t, float64(0), chain["omitted"], + "omitted has no omitempty: it must be present and zero when nothing was dropped") + assert.Equal(t, false, chain["malformed"]) + _, hasReason := chain["malformedReason"] + assert.False(t, hasReason, "malformedReason must be omitted on a well-formed chain") + + hops, ok := chain["chain"].([]any) + require.True(t, ok, "chain should be an array") + require.Len(t, hops, 2) + + outer, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-2", outer["sub"], "chain[0] must be the outermost/most recent actor") + _, hasIss := outer["iss"] + assert.False(t, hasIss, "the outer hop here carries no issuer") + + inner, ok := hops[1].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-1", inner["sub"]) + assert.Equal(t, "https://issuer.example", inner["iss"], + "a hop's issuer must surface as the promoted per-hop iss field") + }) + + t.Run("MaxDelegationDepth truncates the chain and keeps the outermost actor", func(t *testing.T) { + t.Parallel() + maxDepth := 1 + auditor, logBuf := newBufferAuditorWithConfig(t, &Config{Component: "test", MaxDelegationDepth: &maxDepth}) + identity := twoHopIdentity() + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = auth.WithIdentity(r.Context(), identity) + w.WriteHeader(http.StatusOK) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + + chain, ok := events[0]["delegation"].(map[string]any) + require.True(t, ok) + assert.Equal(t, true, chain["truncated"]) + assert.Equal(t, float64(1), chain["omitted"], + "one of the two hops must be reported omitted") + + hops, ok := chain["chain"].([]any) + require.True(t, ok) + require.Len(t, hops, 1, "MaxDelegationDepth=1 must leave exactly one surviving hop") + surviving, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-2", surviving["sub"], "the surviving hop must be the outermost one") + }) + + t.Run("plain identity omits the delegation chain", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "user-123"}} + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = auth.WithIdentity(r.Context(), identity) + w.WriteHeader(http.StatusOK) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + + _, exists := events[0]["delegation"] + assert.False(t, exists, + "a non-delegated identity must not produce a delegation member on the POST path") + }) + + t.Run("malformed act claim is pinned as malformed:true with an empty chain", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + act := "not-an-object" + chain := coreaudit.ParseDelegationChain(act, auth.DefaultMaxDelegationDepth) + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user-123", + Claims: map[string]any{"act": act}, + DelegationChain: chain, + }} + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = auth.WithIdentity(r.Context(), identity) + w.WriteHeader(http.StatusOK) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newToolsCallRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + + logged, ok := events[0]["delegation"].(map[string]any) + require.True(t, ok, "a malformed act claim must still produce a delegation member") + assert.Equal(t, true, logged["malformed"]) + assert.Equal(t, string(coreaudit.MalformedReasonActNotObject), logged["malformedReason"]) + hops, ok := logged["chain"].([]any) + require.True(t, ok, "chain must marshal as a JSON array, not null") + assert.Len(t, hops, 0) + }) +} + +// TestStreamOpenAuditEvents pins the deferred stream-open logging: the +// connection event for SSE / streamable GET requests is logged on the FIRST +// response write, so it reflects the real outcome and the identity attached by +// inner middleware — instead of being logged on arrival with neither. +func TestStreamOpenAuditEvents(t *testing.T) { + t.Parallel() + + newStreamRequest := func() *http.Request { + req := httptest.NewRequest("GET", "/mcp", nil) + req.Header.Set("Accept", "text/event-stream") + return req + } + + t.Run("established stream logs one success event with identity", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Emulate inner auth attaching the identity before the stream starts. + // The returned context is intentionally discarded: only WithIdentity's + // side effect of filling the audit-injected IdentityHolder matters here. + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "user-123"}} + _ = auth.WithIdentity(r.Context(), id) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("event: message\ndata: {}\n\n")) + _, _ = w.Write([]byte("event: message\ndata: {}\n\n")) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1, "the connection event must be logged exactly once") + assert.Equal(t, EventTypeSSEConnection, events[0]["type"]) + assert.Equal(t, OutcomeSuccess, events[0]["outcome"]) + subjects, ok := events[0]["subjects"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "user-123", subjects[SubjectKeyUserID], + "deferring the log to first write makes the inner auth identity available") + }) + + t.Run("rejected stream open logs a denied event", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + reject := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "invalid token", http.StatusUnauthorized) + }) + auditor.Middleware(reject).ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, EventTypeSSEConnection, events[0]["type"]) + assert.Equal(t, OutcomeDenied, events[0]["outcome"]) + }) + + t.Run("handler that never writes still produces one event", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + silent := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}) + auditor.Middleware(silent).ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, OutcomeSuccess, events[0]["outcome"], + "net/http sends an implicit 200 when the handler writes nothing") + }) + + t.Run("flush before first write logs the connection event", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + // A handler that establishes the stream by flushing headers and then + // blocks (waiting for events to send) never hits WriteHeader/Write. + flusher := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + }) + auditor.Middleware(flusher).ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1, "the flush must not bypass the connection event") + assert.Equal(t, EventTypeSSEConnection, events[0]["type"]) + assert.Equal(t, OutcomeSuccess, events[0]["outcome"]) + }) + + t.Run("panic before first write still logs the connection event", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + panicker := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + panic("boom before first write") + }) + handler := auditor.Middleware(panicker) + require.Panics(t, func() { + handler.ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + }, "no recovery middleware on this chain: the panic must propagate") + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1, + "chains whose recovery middleware runs OUTSIDE audit (e.g. the vMCP Serve path) "+ + "must not lose the connection event to a panic") + assert.Equal(t, EventTypeSSEConnection, events[0]["type"]) + }) + + t.Run("delegated token stream open carries the delegation chain", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + act := map[string]any{ + "sub": "agent-1", + "act": map[string]any{"sub": "agent-2"}, + } + chain := coreaudit.ParseDelegationChain(act, auth.DefaultMaxDelegationDepth) + delegated := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user-123", + Claims: map[string]any{"act": act}, + DelegationChain: chain, + }, + } + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Emulate inner auth attaching the identity before the stream starts. + // The returned context is intentionally discarded: only WithIdentity's + // side effect of filling the audit-injected IdentityHolder matters here. + _ = auth.WithIdentity(r.Context(), delegated) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("event: message\ndata: {}\n\n")) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, EventTypeSSEConnection, events[0]["type"]) + + logged, ok := events[0]["delegation"].(map[string]any) + require.True(t, ok, "the SSE connection event must carry the delegation chain") + assert.Equal(t, false, logged["truncated"]) + hops, ok := logged["chain"].([]any) + require.True(t, ok) + require.Len(t, hops, 2) + first, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-1", first["sub"]) + second, ok := hops[1].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-2", second["sub"]) + }) + + t.Run("non-delegated identity stream open omits the delegation chain", func(t *testing.T) { + t.Parallel() + auditor, logBuf := newBufferAuditor(t) + + plain := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "user-123"}} + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Emulate inner auth attaching the identity before the stream starts. + // The returned context is intentionally discarded: only WithIdentity's + // side effect of filling the audit-injected IdentityHolder matters here. + _ = auth.WithIdentity(r.Context(), plain) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("event: message\ndata: {}\n\n")) + }) + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), newStreamRequest()) + + events := decodeAuditEvents(t, logBuf) + require.Len(t, events, 1) + assert.Equal(t, EventTypeSSEConnection, events[0]["type"]) + + subjects, ok := events[0]["subjects"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "user-123", subjects[SubjectKeyUserID], "the identity must have landed") + + _, exists := events[0]["delegation"] + assert.False(t, exists, + "an authenticated non-delegated identity must not produce a delegation member on the stream-open path") + }) +} diff --git a/pkg/audit/config.go b/pkg/audit/config.go index b7a7a0b728..f94f2b3956 100644 --- a/pkg/audit/config.go +++ b/pkg/audit/config.go @@ -11,6 +11,8 @@ import ( "log/slog" "os" "path/filepath" + + "github.com/stacklok/toolhive/pkg/auth" ) // Config represents the audit logging configuration. @@ -52,6 +54,13 @@ type Config struct { // +kubebuilder:default=1024 // +optional MaxDataSize int `json:"maxDataSize,omitempty" yaml:"maxDataSize,omitempty"` + // MaxDelegationDepth caps how many nested RFC 8693 "act" entries are + // recorded in an audit event's delegation chain. Deeper chains are + // truncated (marked with truncated=true). Defaults to 10 when unset. + // +kubebuilder:validation:Minimum=1 + // +kubebuilder:default=10 + // +optional + MaxDelegationDepth *int `json:"maxDelegationDepth,omitempty" yaml:"maxDelegationDepth,omitempty"` // LogFile specifies the file path for audit logs. If empty, logs to stdout. // +optional LogFile string `json:"logFile,omitempty" yaml:"logFile,omitempty"` @@ -94,6 +103,21 @@ func (c *Config) ShouldDetectApplicationErrors() bool { return *c.DetectApplicationErrors } +// MaxDelegationDepthOrDefault returns the configured cap on RFC 8693 +// delegation chain depth, or auth.DefaultMaxDelegationDepth (10) when unset +// or non-positive. +// +// The default is deliberately ToolHive's minting cap (10), NOT toolhive-core's +// audit.DefaultMaxDelegationDepth (16): core's value is a library-wide parse +// ceiling, while passing the mint cap here preserves the signal that a chain +// with truncated=true could not have come from ToolHive's own issuance path. +func (c *Config) MaxDelegationDepthOrDefault() int { + if c == nil || c.MaxDelegationDepth == nil || *c.MaxDelegationDepth <= 0 { + return auth.DefaultMaxDelegationDepth + } + return *c.MaxDelegationDepth +} + // LoadFromFile loads audit configuration from a file. func LoadFromFile(path string) (*Config, error) { // Clean the path to prevent directory traversal @@ -158,6 +182,10 @@ func (c *Config) Validate() error { return fmt.Errorf("maxDataSize cannot be negative") } + if c.MaxDelegationDepth != nil && *c.MaxDelegationDepth <= 0 { + return fmt.Errorf("maxDelegationDepth must be positive") + } + // Validate event types (basic validation - could be extended) validEventTypes := map[string]bool{ EventTypeMCPInitialize: true, diff --git a/pkg/audit/config_test.go b/pkg/audit/config_test.go index 2ec0a477d4..cdd9064666 100644 --- a/pkg/audit/config_test.go +++ b/pkg/audit/config_test.go @@ -16,6 +16,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/auth" ) func TestDefaultConfig(t *testing.T) { @@ -447,6 +449,62 @@ func TestGetLogWriter_WithActualFile(t *testing.T) { }) } +func TestValidateMaxDelegationDepth(t *testing.T) { + t.Parallel() + + depth := func(d int) *int { return &d } + + tests := []struct { + name string + value *int + wantErr bool + }{ + {name: "nil is valid", value: nil, wantErr: false}, + {name: "positive is valid", value: depth(5), wantErr: false}, + {name: "zero is rejected", value: depth(0), wantErr: true}, + {name: "negative is rejected", value: depth(-3), wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + config := &Config{MaxDelegationDepth: tt.value} + err := config.Validate() + if tt.wantErr { + assert.Error(t, err) + assert.Contains(t, err.Error(), "maxDelegationDepth must be positive") + } else { + assert.NoError(t, err) + } + }) + } +} + +func TestMaxDelegationDepthOrDefault(t *testing.T) { + t.Parallel() + + depth := func(d int) *int { return &d } + + tests := []struct { + name string + cfg *Config + want int + }{ + {name: "nil receiver", cfg: nil, want: auth.DefaultMaxDelegationDepth}, + {name: "nil MaxDelegationDepth field", cfg: &Config{}, want: auth.DefaultMaxDelegationDepth}, + {name: "zero MaxDelegationDepth", cfg: &Config{MaxDelegationDepth: depth(0)}, want: auth.DefaultMaxDelegationDepth}, + {name: "negative MaxDelegationDepth", cfg: &Config{MaxDelegationDepth: depth(-3)}, want: auth.DefaultMaxDelegationDepth}, + {name: "positive MaxDelegationDepth is honored", cfg: &Config{MaxDelegationDepth: depth(5)}, want: 5}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, tt.cfg.MaxDelegationDepthOrDefault()) + }) + } +} + // waitForAuditLog polls the audit log file until content is available or timeout is reached. // This is more reliable than a fixed sleep for async log writes. func waitForAuditLog(t *testing.T, logFilePath string, timeout time.Duration) []byte { diff --git a/pkg/audit/event.go b/pkg/audit/event.go index b93aff1eea..e2f41f7315 100644 --- a/pkg/audit/event.go +++ b/pkg/audit/event.go @@ -11,6 +11,8 @@ import ( "time" "github.com/google/uuid" + + coreaudit "github.com/stacklok/toolhive-core/audit" ) // The following code is adapted from github.com/metal-toolbox/auditevent @@ -64,6 +66,12 @@ type AuditEvent struct { // e.g. who triggered the event? Additional information // may be added, such as group membership and/or role Subjects map[string]string `json:"subjects"` + // DelegationChain: when the caller authenticated with an RFC 8693 + // delegated token, this records the full chain of acting parties + // (outermost/most recent first) extracted from the token's "act" claim. + // The type and the "delegation" wire key are toolhive-core's canonical + // delegation-chain schema. + DelegationChain *coreaudit.DelegationChain `json:"delegation,omitempty"` // Component: allows to determine in which component the event occurred // (Answering the "Where" question of section c in the NIST SP 800-53 // Revision 5.1 Control AU-3). @@ -144,6 +152,20 @@ func NewAuditEventWithID( } } +// WithDelegationChain sets the RFC 8693 delegation chain of the event. +// A nil or zero chain (no hops, no truncation, no malformation — see +// [coreaudit.DelegationChain.IsZero]) clears any previously set chain. A +// hopless but malformed chain is kept: "the caller asserted delegation we +// could not read" is audit-relevant and distinct from "no delegation". +func (e *AuditEvent) WithDelegationChain(chain *coreaudit.DelegationChain) *AuditEvent { + if chain.IsZero() { + e.DelegationChain = nil + return e + } + e.DelegationChain = chain + return e +} + // WithTarget sets the target of the event. func (e *AuditEvent) WithTarget(target map[string]string) *AuditEvent { e.Target = target @@ -181,6 +203,14 @@ func (e *AuditEvent) LogTo(ctx context.Context, logger *slog.Logger, level slog. slog.Any("subjects", e.Subjects), } + // Add delegation chain if present and non-zero. The chain resolves through + // coreaudit's [log/slog.LogValuer] implementation, which mirrors its JSON + // shape (so the slog and JSON wire forms are structurally identical) and + // never exposes the per-hop in-memory extra claims (PII guard). + if !e.DelegationChain.IsZero() { + attrs = append(attrs, slog.Any("delegation", e.DelegationChain)) + } + // Add target if present if e.Target != nil { attrs = append(attrs, slog.Any("target", e.Target)) diff --git a/pkg/audit/mcp_events.go b/pkg/audit/mcp_events.go index a06ad5ff58..4ca1923f15 100644 --- a/pkg/audit/mcp_events.go +++ b/pkg/audit/mcp_events.go @@ -2,138 +2,144 @@ // SPDX-License-Identifier: Apache-2.0 // Package audit provides MCP-specific audit event types and constants. +// +// The constants are aliases of toolhive-core's audit package, kept so +// existing toolhive call sites don't need import churn. New code should +// import github.com/stacklok/toolhive-core/audit directly. package audit +import coreaudit "github.com/stacklok/toolhive-core/audit" + // MCP-specific event types based on the Model Context Protocol specification const ( // EventTypeMCPInitialize represents an MCP initialization event - EventTypeMCPInitialize = "mcp_initialize" + EventTypeMCPInitialize = coreaudit.EventTypeMCPInitialize // EventTypeSSEConnection represents an SSE connection event - EventTypeSSEConnection = "sse_connection" + EventTypeSSEConnection = coreaudit.EventTypeSSEConnection // EventTypeMCPToolCall represents an MCP tool call event - EventTypeMCPToolCall = "mcp_tool_call" + EventTypeMCPToolCall = coreaudit.EventTypeMCPToolCall // EventTypeMCPToolsList represents an MCP tools list event - EventTypeMCPToolsList = "mcp_tools_list" + EventTypeMCPToolsList = coreaudit.EventTypeMCPToolsList // EventTypeMCPResourceRead represents an MCP resource read event - EventTypeMCPResourceRead = "mcp_resource_read" + EventTypeMCPResourceRead = coreaudit.EventTypeMCPResourceRead // EventTypeMCPResourcesList represents an MCP resources list event - EventTypeMCPResourcesList = "mcp_resources_list" + EventTypeMCPResourcesList = coreaudit.EventTypeMCPResourcesList // EventTypeMCPPromptGet represents an MCP prompt get event - EventTypeMCPPromptGet = "mcp_prompt_get" + EventTypeMCPPromptGet = coreaudit.EventTypeMCPPromptGet // EventTypeMCPPromptsList represents an MCP prompts list event - EventTypeMCPPromptsList = "mcp_prompts_list" + EventTypeMCPPromptsList = coreaudit.EventTypeMCPPromptsList // EventTypeMCPNotification represents an MCP notification event - EventTypeMCPNotification = "mcp_notification" + EventTypeMCPNotification = coreaudit.EventTypeMCPNotification // EventTypeMCPPing represents an MCP ping event - EventTypeMCPPing = "mcp_ping" + EventTypeMCPPing = coreaudit.EventTypeMCPPing // EventTypeMCPLogging represents an MCP logging event - EventTypeMCPLogging = "mcp_logging" + EventTypeMCPLogging = coreaudit.EventTypeMCPLogging // EventTypeMCPCompletion represents an MCP completion event - EventTypeMCPCompletion = "mcp_completion" + EventTypeMCPCompletion = coreaudit.EventTypeMCPCompletion // EventTypeMCPRootsListChanged represents an MCP roots list changed notification - EventTypeMCPRootsListChanged = "mcp_roots_list_changed" + EventTypeMCPRootsListChanged = coreaudit.EventTypeMCPRootsListChanged // Workflow-specific event types for vMCP composite workflow execution // EventTypeWorkflowStarted represents workflow execution start - EventTypeWorkflowStarted = "vmcp_workflow_started" + EventTypeWorkflowStarted = coreaudit.EventTypeWorkflowStarted // EventTypeWorkflowCompleted represents successful workflow completion - EventTypeWorkflowCompleted = "vmcp_workflow_completed" + EventTypeWorkflowCompleted = coreaudit.EventTypeWorkflowCompleted // EventTypeWorkflowFailed represents workflow failure - EventTypeWorkflowFailed = "vmcp_workflow_failed" + EventTypeWorkflowFailed = coreaudit.EventTypeWorkflowFailed // EventTypeWorkflowTimedOut represents workflow timeout - EventTypeWorkflowTimedOut = "vmcp_workflow_timed_out" + EventTypeWorkflowTimedOut = coreaudit.EventTypeWorkflowTimedOut // EventTypeWorkflowStepStarted represents workflow step execution start - EventTypeWorkflowStepStarted = "vmcp_workflow_step_started" + EventTypeWorkflowStepStarted = coreaudit.EventTypeWorkflowStepStarted // EventTypeWorkflowStepCompleted represents successful step completion - EventTypeWorkflowStepCompleted = "vmcp_workflow_step_completed" + EventTypeWorkflowStepCompleted = coreaudit.EventTypeWorkflowStepCompleted // EventTypeWorkflowStepFailed represents step failure - EventTypeWorkflowStepFailed = "vmcp_workflow_step_failed" + EventTypeWorkflowStepFailed = coreaudit.EventTypeWorkflowStepFailed // EventTypeWorkflowStepSkipped represents conditional step skip - EventTypeWorkflowStepSkipped = "vmcp_workflow_step_skipped" + EventTypeWorkflowStepSkipped = coreaudit.EventTypeWorkflowStepSkipped // Fallback event types for unrecognized or generic requests // EventTypeMCPRequest represents a generic MCP request when specific type cannot be determined - EventTypeMCPRequest = "mcp_request" + EventTypeMCPRequest = coreaudit.EventTypeMCPRequest // EventTypeHTTPRequest represents a generic HTTP request (non-MCP) - EventTypeHTTPRequest = "http_request" + EventTypeHTTPRequest = coreaudit.EventTypeHTTPRequest ) // MCP target types for audit events const ( // TargetTypeTool represents a tool target - TargetTypeTool = "tool" + TargetTypeTool = coreaudit.TargetTypeTool // TargetTypeResource represents a resource target - TargetTypeResource = "resource" + TargetTypeResource = coreaudit.TargetTypeResource // TargetTypePrompt represents a prompt target - TargetTypePrompt = "prompt" + TargetTypePrompt = coreaudit.TargetTypePrompt // TargetTypeServer represents a server target - TargetTypeServer = "server" + TargetTypeServer = coreaudit.TargetTypeServer // TargetTypeWorkflow represents a workflow target - TargetTypeWorkflow = "workflow" + TargetTypeWorkflow = coreaudit.TargetTypeWorkflow // TargetTypeWorkflowStep represents a workflow step target - TargetTypeWorkflowStep = "workflow_step" + TargetTypeWorkflowStep = coreaudit.TargetTypeWorkflowStep ) // MCP-specific target field keys const ( // TargetKeyType is the key for the target type in the target map - TargetKeyType = "type" + TargetKeyType = coreaudit.TargetKeyType // TargetKeyName is the key for the target name in the target map - TargetKeyName = "name" + TargetKeyName = coreaudit.TargetKeyName // TargetKeyURI is the key for the target URI in the target map - TargetKeyURI = "uri" + TargetKeyURI = coreaudit.TargetKeyURI // TargetKeyMethod is the key for the MCP method in the target map - TargetKeyMethod = "method" + TargetKeyMethod = coreaudit.TargetKeyMethod // TargetKeyEndpoint is the key for the endpoint in the target map - TargetKeyEndpoint = "endpoint" + TargetKeyEndpoint = coreaudit.TargetKeyEndpoint // TargetKeyWorkflowID is the key for the unique workflow execution ID - TargetKeyWorkflowID = "workflow_id" + TargetKeyWorkflowID = coreaudit.TargetKeyWorkflowID // TargetKeyWorkflowName is the key for the workflow definition name - TargetKeyWorkflowName = "workflow_name" + TargetKeyWorkflowName = coreaudit.TargetKeyWorkflowName // TargetKeyStepID is the key for the step identifier - TargetKeyStepID = "step_id" + TargetKeyStepID = coreaudit.TargetKeyStepID // TargetKeyStepType is the key for the step type (tool, elicitation) - TargetKeyStepType = "step_type" + TargetKeyStepType = coreaudit.TargetKeyStepType // TargetKeyToolName is the key for the tool being called (for tool steps) - TargetKeyToolName = "tool_name" + TargetKeyToolName = coreaudit.TargetKeyToolName ) // MCP-specific subject field keys const ( // SubjectKeyUser is the key for the user in the subjects map - SubjectKeyUser = "user" + SubjectKeyUser = coreaudit.SubjectKeyUser // SubjectKeyUserID is the key for the user ID in the subjects map - SubjectKeyUserID = "user_id" + SubjectKeyUserID = coreaudit.SubjectKeyUserID // SubjectKeyClientName is the key for the client name in the subjects map - SubjectKeyClientName = "client_name" + SubjectKeyClientName = coreaudit.SubjectKeyClientName // SubjectKeyClientVersion is the key for the client version in the subjects map - SubjectKeyClientVersion = "client_version" + SubjectKeyClientVersion = coreaudit.SubjectKeyClientVersion ) // MCP-specific source field keys for EventSource.Extra const ( // SourceExtraKeyUserAgent is the key for the user agent in the source extra map - SourceExtraKeyUserAgent = "user_agent" + SourceExtraKeyUserAgent = coreaudit.SourceExtraKeyUserAgent // SourceExtraKeyRequestID is the key for the request ID in the source extra map - SourceExtraKeyRequestID = "request_id" + SourceExtraKeyRequestID = coreaudit.SourceExtraKeyRequestID // SourceExtraKeySessionID is the key for the session ID in the source extra map - SourceExtraKeySessionID = "session_id" + SourceExtraKeySessionID = coreaudit.SourceExtraKeySessionID ) // MCP-specific metadata field keys for EventMetadata.Extra const ( // MetadataExtraKeyMCPVersion is the key for the MCP version in the metadata extra map - MetadataExtraKeyMCPVersion = "mcp_version" + MetadataExtraKeyMCPVersion = coreaudit.MetadataExtraKeyMCPVersion // MetadataExtraKeyTransport is the key for the transport type in the metadata extra map - MetadataExtraKeyTransport = "transport" + MetadataExtraKeyTransport = coreaudit.MetadataExtraKeyTransport // MetadataExtraKeyDuration is the key for the request duration in the metadata extra map - MetadataExtraKeyDuration = "duration_ms" + MetadataExtraKeyDuration = coreaudit.MetadataExtraKeyDuration // MetadataExtraKeyResponseSize is the key for the response size in the metadata extra map - MetadataExtraKeyResponseSize = "response_size_bytes" + MetadataExtraKeyResponseSize = coreaudit.MetadataExtraKeyResponseSize // MetadataExtraKeyRetryCount is the key for the number of retries performed - MetadataExtraKeyRetryCount = "retry_count" + MetadataExtraKeyRetryCount = coreaudit.MetadataExtraKeyRetryCount // MetadataExtraKeyStepCount is the key for the total number of steps in a workflow - MetadataExtraKeyStepCount = "step_count" + MetadataExtraKeyStepCount = coreaudit.MetadataExtraKeyStepCount // MetadataExtraKeyTimeout is the key for the workflow timeout in milliseconds - MetadataExtraKeyTimeout = "timeout_ms" + MetadataExtraKeyTimeout = coreaudit.MetadataExtraKeyTimeout ) diff --git a/pkg/audit/workflow_auditor.go b/pkg/audit/workflow_auditor.go index 8966a17cf0..e87a3ce94c 100644 --- a/pkg/audit/workflow_auditor.go +++ b/pkg/audit/workflow_auditor.go @@ -60,16 +60,7 @@ func (w *WorkflowAuditor) LogWorkflowStarted( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowStarted, - source, - OutcomeSuccess, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowStarted, OutcomeSuccess) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -111,16 +102,7 @@ func (w *WorkflowAuditor) LogWorkflowCompleted( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowCompleted, - source, - OutcomeSuccess, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowCompleted, OutcomeSuccess) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -163,16 +145,7 @@ func (w *WorkflowAuditor) LogWorkflowFailed( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowFailed, - source, - OutcomeFailure, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowFailed, OutcomeFailure) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -202,16 +175,7 @@ func (w *WorkflowAuditor) LogWorkflowTimedOut( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowTimedOut, - source, - OutcomeFailure, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowTimedOut, OutcomeFailure) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -241,16 +205,7 @@ func (w *WorkflowAuditor) LogStepStarted( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowStepStarted, - source, - OutcomeSuccess, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowStepStarted, OutcomeSuccess) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -278,16 +233,7 @@ func (w *WorkflowAuditor) LogStepCompleted( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowStepCompleted, - source, - OutcomeSuccess, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowStepCompleted, OutcomeSuccess) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -317,16 +263,7 @@ func (w *WorkflowAuditor) LogStepFailed( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowStepFailed, - source, - OutcomeFailure, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowStepFailed, OutcomeFailure) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -354,16 +291,7 @@ func (w *WorkflowAuditor) LogStepSkipped( return } - source := w.extractSource(ctx) - subjects := w.extractSubjects(ctx) - - event := NewAuditEvent( - EventTypeWorkflowStepSkipped, - source, - OutcomeSuccess, - subjects, - w.component, - ) + event := w.newEvent(ctx, EventTypeWorkflowStepSkipped, OutcomeSuccess) target := map[string]string{ TargetKeyWorkflowID: workflowID, @@ -382,6 +310,22 @@ func (w *WorkflowAuditor) LogStepSkipped( event.LogTo(ctx, w.auditLogger, LevelAudit) } +// newEvent creates an audit event of the given type and outcome with the +// source, subjects, and RFC 8693 delegation chain extracted from the context. +// Every Log* method MUST build its event through this helper so that the +// delegation chain cannot be forgotten at any individual call site. +func (w *WorkflowAuditor) newEvent(ctx context.Context, eventType, outcome string) *AuditEvent { + event := NewAuditEvent( + eventType, + w.extractSource(ctx), + outcome, + w.extractSubjects(ctx), + w.component, + ) + w.attachDelegation(ctx, event) + return event +} + // extractSource extracts source information from context. // For workflows, source is always local since they're internal orchestration. func (*WorkflowAuditor) extractSource(_ context.Context) EventSource { @@ -408,3 +352,15 @@ func (*WorkflowAuditor) extractSubjects(ctx context.Context) map[string]string { return subjects } + +// attachDelegation attaches the RFC 8693 delegation chain from the context's +// identity to the event. It is a no-op when there is no identity or the +// identity carries no delegation chain. +func (w *WorkflowAuditor) attachDelegation(ctx context.Context, event *AuditEvent) { + identity, ok := auth.IdentityFromContext(ctx) + if !ok { + return + } + event.WithDelegationChain( + extractDelegationChainFromIdentity(identity, w.config.MaxDelegationDepthOrDefault())) +} diff --git a/pkg/audit/workflow_auditor_test.go b/pkg/audit/workflow_auditor_test.go index 7a0ec9581d..4b8f49dd2c 100644 --- a/pkg/audit/workflow_auditor_test.go +++ b/pkg/audit/workflow_auditor_test.go @@ -15,6 +15,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + coreaudit "github.com/stacklok/toolhive-core/audit" "github.com/stacklok/toolhive/pkg/auth" ) @@ -594,6 +595,78 @@ func TestWorkflowAuditor_ExtractSubjects(t *testing.T) { } } +func TestWorkflowAuditor_DelegationChain(t *testing.T) { + t.Parallel() + + delegatedChain := coreaudit.ParseDelegationChain( + map[string]any{ + "sub": "agent-1", + "act": map[string]any{"sub": "agent-2"}, + }, auth.DefaultMaxDelegationDepth) + delegatedIdentity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user-delegated", + Name: "Delegated User", + Claims: map[string]any{ + "act": map[string]any{ + "sub": "agent-1", + "act": map[string]any{"sub": "agent-2"}, + }, + }, + DelegationChain: delegatedChain, + }, + } + + // Every Log* method builds its event through newEvent, so testing the + // helper covers the delegation attachment for all of them by construction. + t.Run("newEvent attaches delegation chain", func(t *testing.T) { + t.Parallel() + auditor, _ := createTestAuditor(t, DefaultConfig()) + + ctx := auth.WithIdentity(context.Background(), delegatedIdentity) + event := auditor.newEvent(ctx, EventTypeWorkflowStarted, OutcomeSuccess) + + require.NotNil(t, event.DelegationChain, "delegation chain should be attached") + assert.Equal(t, delegatedChain, event.DelegationChain) + }) + + t.Run("no identity omits delegation chain", func(t *testing.T) { + t.Parallel() + auditor, writer := createTestAuditor(t, DefaultConfig()) + + auditor.LogWorkflowStarted(context.Background(), "wf-1", "wf", nil, time.Second) + + require.NotEmpty(t, writer.logs, "expected log entry") + entry := parseLogEntry(t, writer.getLastLog()) + + _, exists := entry["delegation"] + assert.False(t, exists, "delegation should be omitted without an identity") + }) + + t.Run("configured max depth truncates chain", func(t *testing.T) { + t.Parallel() + maxDepth := 1 + auditor, writer := createTestAuditor(t, &Config{MaxDelegationDepth: &maxDepth}) + + ctx := auth.WithIdentity(context.Background(), delegatedIdentity) + auditor.LogWorkflowStarted(ctx, "wf-1", "wf", nil, time.Second) + + require.NotEmpty(t, writer.logs, "expected log entry") + entry := parseLogEntry(t, writer.getLastLog()) + + chain, ok := entry["delegation"].(map[string]any) + require.True(t, ok) + assert.Equal(t, true, chain["truncated"]) + hops, ok := chain["chain"].([]any) + require.True(t, ok) + require.Len(t, hops, 1) + first, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-1", first["sub"]) + assert.Equal(t, float64(1), chain["omitted"]) + }) +} + func TestWorkflowAuditor_ExtractSource(t *testing.T) { t.Parallel() diff --git a/pkg/audit/zz_generated.deepcopy.go b/pkg/audit/zz_generated.deepcopy.go index 7af0eff790..1305e4a068 100644 --- a/pkg/audit/zz_generated.deepcopy.go +++ b/pkg/audit/zz_generated.deepcopy.go @@ -40,6 +40,11 @@ func (in *Config) DeepCopyInto(out *Config) { *out = new(bool) **out = **in } + if in.MaxDelegationDepth != nil { + in, out := &in.MaxDelegationDepth, &out.MaxDelegationDepth + *out = new(int) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Config. diff --git a/pkg/auth/awssts/middleware.go b/pkg/auth/awssts/middleware.go index 325499c6f5..6e2ead67c7 100644 --- a/pkg/auth/awssts/middleware.go +++ b/pkg/auth/awssts/middleware.go @@ -12,7 +12,9 @@ import ( "io" "log/slog" "net/http" + "net/textproto" "net/url" + "strings" "github.com/aws/aws-sdk-go-v2/aws" @@ -227,6 +229,37 @@ func createAWSStsMiddlewareFunc( } } +// hopByHopHeaders lists hop-by-hop headers that httputil.ReverseProxy +// removes before forwarding. This matches the standard library's +// hopHeaders list in net/http/httputil/reverseproxy.go. +var hopByHopHeaders = []string{ + "Connection", + "Proxy-Connection", + "Keep-Alive", + "Proxy-Authenticate", + "Proxy-Authorization", + "Te", + "Trailer", + "Transfer-Encoding", + "Upgrade", +} + +// stripHopByHopHeaders removes hop-by-hop headers from a header set, +// including any header named in the Connection value (RFC 7230 §6.1, +// RFC 9110 §7.6.1). +func stripHopByHopHeaders(header http.Header) { + for _, f := range header.Values("Connection") { + for _, name := range strings.Split(f, ",") { + if name = textproto.TrimString(name); name != "" { + header.Del(name) + } + } + } + for _, h := range hopByHopHeaders { + header.Del(h) + } +} + // signRequestForTarget signs the request with SigV4 for the given target host // without permanently modifying r.Host or r.URL. When targetURL is non-nil, a // clone is used for signing so that only the SigV4 headers are copied back to @@ -255,6 +288,12 @@ func signRequestForTarget(r *http.Request, signer *RequestSigner, creds *aws.Cre _ = r.Body.Close() } + // Strip hop-by-hop headers from the actual request before cloning. + // httputil.ReverseProxy removes these before forwarding; if they + // remain on r they can trigger removal of newly-signed headers + // (e.g. Connection: X-Amz-Date would delete the SigV4 X-Amz-Date). + stripHopByHopHeaders(r.Header) + // Build a signing-only clone with the target host. signingReq := r.Clone(r.Context()) signingReq.URL.Scheme = targetURL.Scheme diff --git a/pkg/auth/awssts/middleware_test.go b/pkg/auth/awssts/middleware_test.go index b7602af0d5..d17a0d46a4 100644 --- a/pkg/auth/awssts/middleware_test.go +++ b/pkg/auth/awssts/middleware_test.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "net/http/httptest" + "net/http/httputil" "net/url" "strings" "testing" @@ -383,6 +384,173 @@ func TestMiddlewareFunc_ProxyHeadersExcludedFromSignature(t *testing.T) { } } +func TestMiddlewareFunc_HopByHopHeadersExcludedFromSignature(t *testing.T) { + t.Parallel() + + expiration := time.Now().Add(time.Hour) + successResponse := &sts.AssumeRoleWithWebIdentityOutput{ + Credentials: &ststypes.Credentials{ + AccessKeyId: aws.String("AKIATEST"), SecretAccessKey: aws.String("secret"), + SessionToken: aws.String("session"), Expiration: &expiration, + }, + } + + targetURL, err := url.Parse("https://aws-mcp.us-east-1.api.aws") + require.NoError(t, err) + + exchanger := &Exchanger{client: &mockSTSClient{response: successResponse}} + roleMapper, err := NewRoleMapper(&Config{ + Region: "us-east-1", + FallbackRoleArn: "arn:aws:iam::123456789012:role/TestRole", + }) + require.NoError(t, err) + signer, err := newRequestSigner("us-east-1") + require.NoError(t, err) + + middlewareFunc := createAWSStsMiddlewareFunc(exchanger, roleMapper, signer, "sub", 3600, targetURL) + + tests := []struct { + name string + setHeaders func(r *http.Request) + excluded []string + extraCheck func(t *testing.T, r *http.Request, signedHeaders string) + }{ + { + name: "standard hop-by-hop headers", + setHeaders: func(r *http.Request) { + r.Header.Set("Connection", "keep-alive") + r.Header.Set("Keep-Alive", "timeout=5, max=1000") + r.Header.Set("Proxy-Connection", "keep-alive") + r.Header.Set("Proxy-Authenticate", "Basic") + r.Header.Set("Proxy-Authorization", "Basic dGVzdDp0ZXN0") + r.Header.Set("Te", "trailers") + r.Header.Set("Trailer", "X-Custom") + r.Header.Set("Transfer-Encoding", "chunked") + r.Header.Set("Upgrade", "websocket") + }, + excluded: []string{ + "connection", "keep-alive", "proxy-connection", + "proxy-authenticate", "proxy-authorization", + "te", "trailer", "transfer-encoding", "upgrade", + }, + }, + { + name: "Connection names a custom header", + setHeaders: func(r *http.Request) { + r.Header.Set("Connection", "X-Test-Hop") + r.Header.Set("X-Test-Hop", "value") + }, + excluded: []string{"connection", "x-test-hop"}, + }, + { + name: "Connection names X-Amz-Date", + setHeaders: func(r *http.Request) { + r.Header.Set("Connection", "X-Amz-Date") + }, + excluded: []string{"connection"}, + extraCheck: func(t *testing.T, r *http.Request, signedHeaders string) { + t.Helper() + require.Contains(t, signedHeaders, "x-amz-date", + "X-Amz-Date must be in SignedHeaders after ReverseProxy") + require.NotEmpty(t, r.Header.Get("X-Amz-Date"), + "X-Amz-Date must survive ReverseProxy") + }, + }, + { + name: "Connection names Authorization", + setHeaders: func(r *http.Request) { + r.Header.Set("Connection", "Authorization") + }, + excluded: []string{"connection"}, + extraCheck: func(t *testing.T, r *http.Request, _ string) { + t.Helper() + require.NotEmpty(t, r.Header.Get("Authorization"), + "Authorization must survive ReverseProxy") + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var receivedReq *http.Request + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedReq = r + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + backendURL, err := url.Parse(backend.URL) + require.NoError(t, err) + + proxy := &httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.SetURL(backendURL) + pr.SetXForwarded() + }, + } + + req := httptest.NewRequest(http.MethodPost, "http://localhost:8080/mcp/v1", strings.NewReader(`{}`)) + req.Header.Set("Authorization", "Bearer test-jwt-token") + tt.setHeaders(req) + + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: map[string]interface{}{"sub": "user123"}, + }} + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + rec := httptest.NewRecorder() + middlewareFunc(proxy).ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, "middleware rejected request") + require.NotNil(t, receivedReq, "backend should have received the request") + + authz := receivedReq.Header.Get("Authorization") + require.Contains(t, authz, "SignedHeaders=") + + signedHeaders := extractSignedHeaders(authz) + + for _, h := range tt.excluded { + for _, signed := range strings.Split(signedHeaders, ";") { + assert.NotEqual(t, h, signed, + "hop-by-hop header %q must not appear in SignedHeaders", h) + } + } + + // Every signed header must still be present on the request + // received by the backend, except host and content-length which + // are carried by the transport layer. + for _, name := range strings.Split(signedHeaders, ";") { + switch name { + case "host", "content-length": + continue + default: + assert.NotEmpty(t, receivedReq.Header.Values(name), + "signed header %q missing after ReverseProxy", name) + } + } + + if tt.extraCheck != nil { + tt.extraCheck(t, receivedReq, signedHeaders) + } + }) + } +} + +func extractSignedHeaders(authz string) string { + i := strings.Index(authz, "SignedHeaders=") + if i < 0 { + return "" + } + s := authz[i+len("SignedHeaders="):] + if j := strings.Index(s, ","); j >= 0 { + s = s[:j] + } + return s +} + // TestMiddlewareFunc_RoleMapperFailure tests that the middleware returns 403 // when the role mapper cannot determine an IAM role for the request. func TestMiddlewareFunc_RoleMapperFailure(t *testing.T) { diff --git a/pkg/auth/context.go b/pkg/auth/context.go index ea806db36b..d6510854dc 100644 --- a/pkg/auth/context.go +++ b/pkg/auth/context.go @@ -9,8 +9,22 @@ import ( "errors" "github.com/golang-jwt/jwt/v5" + + "github.com/stacklok/toolhive-core/audit" ) +// DefaultMaxDelegationDepth is ToolHive's default cap on how many RFC 8693 +// delegation hops are retained when parsing an "act" claim. It is aligned +// with the issuance-side cap in pkg/authserver/server/tokenexchange/handler.go +// (maxDelegationDepth) and the claim-nesting cap in +// pkg/authz/authorizers/cedar/entity.go (maxClaimNestingDepth), both 10. +// +// This is deliberately LOWER than toolhive-core's audit.DefaultMaxDelegationDepth +// (16): core's value is a library-wide parse ceiling, while 10 is ToolHive's own +// minting cap. Parsing at the mint cap preserves the signal that a chain with +// Truncated=true could not have been minted by ToolHive's own issuance path. +const DefaultMaxDelegationDepth = 10 + // IdentityContextKey is the key used to store Identity in the request context. // This provides type-safe context storage and retrieval for authenticated identities. // @@ -24,6 +38,16 @@ type IdentityContextKey struct{} // This function is typically called by authentication middleware after successful // authentication to make the identity available to downstream handlers. // +// Side effect: if the context carries an IdentityHolder (injected by the audit +// middleware wrapping this call), the identity is also published into it. Two +// invariants follow for callers on request-derived contexts: +// - Last write wins: the holder reports the most recently attached identity, +// so audit events carry that principal. Do not attach a different principal +// (e.g. a re-minted service identity) on a request-derived context unless +// that is the principal audit should report. +// - Call only on the request goroutine: the audit middleware reads the holder +// when the response is written, so off-goroutine writes would race. +// // Example: // // identity := &Identity{PrincipalInfo: PrincipalInfo{Subject: "user123", Name: "Alice"}} @@ -32,9 +56,46 @@ func WithIdentity(ctx context.Context, identity *Identity) context.Context { if identity == nil { return ctx } + // Also publish the identity to an IdentityHolder if one is present, so + // middleware wrapping the auth middleware (e.g. audit) can observe the + // identity even though the derived context only flows downstream. + if holder, ok := IdentityHolderFromContext(ctx); ok { + holder.Identity = identity + } return context.WithValue(ctx, IdentityContextKey{}, identity) } +// IdentityHolderContextKey is the key used to store an IdentityHolder in the +// request context. +type IdentityHolderContextKey struct{} + +// IdentityHolder is a mutable carrier that lets middleware running OUTSIDE the +// auth middleware observe the authenticated identity. Context values only flow +// downstream, so a wrapper such as the audit middleware cannot read the +// identity that auth attaches for inner handlers. The wrapper injects an empty +// holder via WithIdentityHolder before calling the inner chain; WithIdentity +// fills it when the identity is attached; the wrapper reads it after the inner +// chain returns. +// +// The holder is written and read by the single request goroutine (writes +// happen-before the wrapper's post-ServeHTTP read), mirroring the audit +// package's BackendInfo pattern, so no synchronization is needed. +type IdentityHolder struct { + Identity *Identity +} + +// WithIdentityHolder returns a new context carrying the given IdentityHolder. +func WithIdentityHolder(ctx context.Context, holder *IdentityHolder) context.Context { + return context.WithValue(ctx, IdentityHolderContextKey{}, holder) +} + +// IdentityHolderFromContext retrieves the IdentityHolder from the context. +// Returns (nil, false) if no holder is present. +func IdentityHolderFromContext(ctx context.Context) (*IdentityHolder, bool) { + holder, ok := ctx.Value(IdentityHolderContextKey{}).(*IdentityHolder) + return holder, ok && holder != nil +} + // IdentityFromContext retrieves an Identity from the context. // Returns the identity and true if a non-nil identity is present, nil and false otherwise. // A typed-nil *Identity stored directly in the context (bypassing WithIdentity) is @@ -156,6 +217,19 @@ func claimsToIdentity(claims jwt.MapClaims, token string) (*Identity, error) { identity.Email = email } + // Parse the RFC 8693 "act" claim (if present) into a delegation chain. + // This is primarily for audit: per RFC 8693 §4.1, an authorizer applying + // access control MUST consider only the current actor (chain.Chain[0]); + // any nested actors in chain.Chain[1:] are informational only. + // The parse cannot fail by design: a malformed act claim must never fail + // authentication, only be recorded (Malformed/MalformedReason on the chain). + if act, ok := claims["act"]; ok && act != nil { + chain := audit.ParseDelegationChain(act, DefaultMaxDelegationDepth) + if !chain.IsZero() { + identity.DelegationChain = chain + } + } + return identity, nil } diff --git a/pkg/auth/context_test.go b/pkg/auth/context_test.go index 6300056032..89c5a612cb 100644 --- a/pkg/auth/context_test.go +++ b/pkg/auth/context_test.go @@ -10,6 +10,8 @@ import ( "github.com/golang-jwt/jwt/v5" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/audit" ) // TestIdentityContext_StoreAndRetrieve verifies basic context storage and retrieval functionality. @@ -126,6 +128,59 @@ func TestClaimsToIdentity_PopulatesPlatformUserID(t *testing.T) { assert.Equal(t, "user123", id.PlatformUserID) } +// TestClaimsToIdentity_ParsesActClaim verifies that an RFC 8693 "act" claim is +// parsed into the identity's delegation chain (outermost actor first), and +// that tokens without an act claim carry an empty chain. +func TestClaimsToIdentity_ParsesActClaim(t *testing.T) { + t.Parallel() + + t.Run("nested act claim", func(t *testing.T) { + t.Parallel() + claims := jwt.MapClaims{ + "sub": "user123", + "act": map[string]any{ + "sub": "agent-1", + "iss": "https://issuer.example", + "act": map[string]any{"sub": "agent-2"}, + }, + } + + id, err := claimsToIdentity(claims, "tok") + require.NoError(t, err) + + require.NotNil(t, id.DelegationChain) + require.Len(t, id.DelegationChain.Chain, 2) + assert.False(t, id.DelegationChain.Truncated) + assert.Equal(t, "agent-1", id.DelegationChain.Chain[0].Subject) + assert.Equal(t, "https://issuer.example", id.DelegationChain.Chain[0].Issuer) + assert.Equal(t, "agent-2", id.DelegationChain.Chain[1].Subject) + }) + + t.Run("no act claim", func(t *testing.T) { + t.Parallel() + id, err := claimsToIdentity(jwt.MapClaims{"sub": "user123"}, "tok") + require.NoError(t, err) + assert.Nil(t, id.DelegationChain) + }) + + t.Run("nil act claim", func(t *testing.T) { + t.Parallel() + id, err := claimsToIdentity(jwt.MapClaims{"sub": "user123", "act": nil}, "tok") + require.NoError(t, err) + assert.Nil(t, id.DelegationChain) + }) + + t.Run("non-map act claim is surfaced as malformed, not silently dropped", func(t *testing.T) { + t.Parallel() + id, err := claimsToIdentity(jwt.MapClaims{"sub": "user123", "act": "agent-1"}, "tok") + require.NoError(t, err) + require.NotNil(t, id.DelegationChain, "a malformed act must still produce a chain so the issue is visible") + assert.True(t, id.DelegationChain.Malformed) + assert.Equal(t, audit.MalformedReasonActNotObject, id.DelegationChain.MalformedReason) + assert.Empty(t, id.DelegationChain.Chain) + }) +} + // TestPlatformUserContext_StoreAndRetrieve verifies the dedicated platform-user key // round-trips and that an empty userID leaves the context (and the identity key) // untouched. diff --git a/pkg/auth/dcr/request.go b/pkg/auth/dcr/request.go index 055056c160..435202a638 100644 --- a/pkg/auth/dcr/request.go +++ b/pkg/auth/dcr/request.go @@ -43,13 +43,16 @@ type Request struct { // RFC 8252 §7.3, so the redirect-URI defaulting is never // exercised on this path. // - // The resolver's cache key is (Issuer, RedirectURI, ScopesHash); the - // two consumers do not collide on the key because the embedded - // authserver's RedirectURI lives on the authserver's origin and the - // CLI's lives on a loopback (RFC 8252 §7.3) — even when Issuer - // happens to match between the two profiles, the RedirectURI - // component keeps the cache key distinct. See the cross-consumer - // caveat on dcrFlight in resolver.go for the wider invariant. + // The resolver's cache key is (Issuer, UpstreamID, RedirectURI, + // ScopesHash), where UpstreamID is derived by the resolver from + // DiscoveryURL / RegistrationEndpoint (issue #5823). The two consumers + // do not collide on the key because the embedded authserver's + // RedirectURI lives on the authserver's origin and the CLI's lives on a + // loopback (RFC 8252 §7.3) — even when Issuer happens to match between + // the two profiles, the RedirectURI component keeps the cache key + // distinct; UpstreamID additionally keeps two upstreams within a single + // consumer apart. See the cross-consumer caveat on dcrFlight in + // resolver.go for the wider invariant. // // Issuer is *not* used for RFC 8414 §3.3 metadata verification — // that uses an issuer recovered from DiscoveryURL inside the @@ -71,7 +74,9 @@ type Request struct { // document. The resolver fetches this URL exactly once and reads // registration_endpoint, authorization_endpoint, token_endpoint, // token_endpoint_auth_methods_supported, scopes_supported, and - // code_challenge_methods_supported from it. + // code_challenge_methods_supported from it. On this branch the resolver + // also derives the cache key's UpstreamID from the upstream issuer + // recovered from this URL. // // Mutually exclusive with RegistrationEndpoint. DiscoveryURL string @@ -80,7 +85,8 @@ type Request struct { // On this branch the caller is expected to also supply AuthorizationEndpoint // and TokenEndpoint explicitly; the resolver's auth-method selection // defaults to client_secret_basic because no server-capability fields - // are available unless CodeChallengeMethodsSupported is also set. + // are available unless CodeChallengeMethodsSupported is also set. This + // URL is also the cache key's UpstreamID on this branch. // // Mutually exclusive with DiscoveryURL. RegistrationEndpoint string @@ -132,4 +138,16 @@ type Request struct { // code_challenge_methods_supported the S256 gate cannot be evaluated, // so the resolver refuses to register as a public client. PublicClient bool + + // AllowPrivateIPs permits both of the resolver's outbound calls — the + // discovery fetch to DiscoveryURL and the registration POST to the + // resolved registration endpoint — to connect to private IP ranges. + // See networking.NewHostScopedClientBuilder for the CWE-918 guard policy + // this widens (loopback exemption, dial-time re-check, etc.). Set this to + // true only when the upstream authorization server is reachable solely + // over a private address (e.g. an in-cluster IdP with no public + // endpoint); defaults to false. Mirrors the AllowPrivateIPs posture + // already carried by the OAuth2 and OIDC upstream configs so a DCR + // upstream is not the one outbound-facing path without an SSRF guard. + AllowPrivateIPs bool } diff --git a/pkg/auth/dcr/resolver.go b/pkg/auth/dcr/resolver.go index 6ffba819f0..ba1888554f 100644 --- a/pkg/auth/dcr/resolver.go +++ b/pkg/auth/dcr/resolver.go @@ -34,12 +34,12 @@ // # Concurrency // // The package maintains a process-global singleflight keyed on the tuple -// (issuer, redirectURI, scopesHash) so concurrent ResolveCredentials calls -// across all consumers in a single process coalesce when their cache keys -// match. Consumers that share any of those three values will share a flight -// — the deduplication is a feature for the embedded authserver but means -// callers cannot assume per-call-site flight isolation. See the dcrFlight -// doc comment below for the rationale. +// (issuer, upstreamID, redirectURI, scopesHash) so concurrent +// ResolveCredentials calls across all consumers in a single process +// coalesce when their cache keys match. Consumers that share all four +// values will share a flight — the deduplication is a feature for the +// embedded authserver but means callers cannot assume per-call-site flight +// isolation. See the dcrFlight doc comment below for the rationale. // // See issue #5145 for the design discussion that motivated lifting this out // of pkg/authserver/runner. @@ -88,8 +88,9 @@ import ( // // Cross-consumer caveat: because dcrFlight is package-global, two // consumers that happen to construct identical Keys (same issuer, same -// redirect URI, same scopes hash) will share a single in-flight -// registration even if they semantically want different client profiles. +// upstream ID, same redirect URI, same scopes hash) will share a single +// in-flight registration even if they semantically want different client +// profiles. // The two current call sites do not collide by construction — the embedded // authserver's redirect URI lives on the AS origin, and the CLI flow's // redirect URI lives on a loopback (http://localhost:{port}/callback per @@ -103,13 +104,30 @@ import ( var dcrFlight singleflight.Group // flightKeyOf canonicalises a Key into the singleflight string used by -// dcrFlight. The "\n" separator is safe because newline is not a valid -// byte in any of the three components: URI reference characters in -// Issuer and RedirectURI (RFC 3986 §2), and hex digits in ScopesHash -// (the form storage.ScopesHash always emits). Exposed as a function so -// tests and future inspection helpers can compute the exact key the -// resolver would route through dcrFlight without re-implementing the -// concatenation. +// dcrFlight. Each field is length-prefixed (the same scheme redisDCRKey +// uses in pkg/authserver/storage) so the key is unambiguous regardless of +// field contents; ScopesHash is a SHA-256 hex digest with no colons, so it +// is appended without a prefix, exactly as redisDCRKey does. Exposed as a +// function so tests and future inspection helpers can compute the exact key +// the resolver would route through dcrFlight without re-implementing the +// encoding. +// +// Length-prefixing rather than a "separator byte that never appears": +// UpstreamID is NOT guaranteed to be a clean URL when this key is built. On +// the CLI/remote path it is req.RegistrationEndpoint copied verbatim from +// the upstream's discovery-document JSON, and validateUpstreamEndpointURL +// does not run until later, inside the singleflight body +// (resolveDCREndpoints) — after this key already exists. A newline (or any +// byte) in a discovery-supplied registration_endpoint would therefore reach +// a naive "\n"-joined key, where a crafted value could byte-collide with a +// concurrently-registering upstream and, as the singleflight follower, +// observe that upstream's Resolution. Length-prefixing removes the +// assumption entirely and aligns the flight key with the persisted DCRKey. +// +// UpstreamID keeps two upstreams that share Issuer, RedirectURI, and +// scopes from coalescing into a single flight (issue #5823) — the same +// distinction the persisted DCRKey draws, so the singleflight and cache +// layers stay aligned. // // PublicClient is intentionally NOT part of the flight key — the // dcrFlight doc above explains why: today's two consumers register on @@ -119,7 +137,11 @@ var dcrFlight singleflight.Group // DCRKey from cross-profile coalescence; the two layers are aligned // rather than asymmetric. func flightKeyOf(key Key) string { - return key.Issuer + "\n" + key.RedirectURI + "\n" + key.ScopesHash + return fmt.Sprintf("%d:%s:%d:%s:%d:%s:%s", + len(key.Issuer), key.Issuer, + len(key.UpstreamID), key.UpstreamID, + len(key.RedirectURI), key.RedirectURI, + key.ScopesHash) } // defaultUpstreamRedirectPath is the redirect path derived from the issuer @@ -235,6 +257,7 @@ type Resolution struct { const ( dcrStepValidate = "validate" dcrStepResolveRedirect = "resolve_redirect_uri" + dcrStepResolveUpstream = "resolve_upstream_id" dcrStepCacheRead = "cache_read" dcrStepMetadata = "metadata_discovery" dcrStepSelectAuthMethod = "select_auth_method" @@ -326,9 +349,19 @@ func ResolveCredentials( fmt.Errorf("resolve redirect uri: %w", err)) } + // Identify the upstream authorization server so two upstreams that share + // req.Issuer, redirectURI, and scopes — the common case within a single + // embedded authserver — do not collide on one cache entry (issue #5823). + upstreamID, err := resolveUpstreamKeyIdentity(req) + if err != nil { + return nil, newDCRStepError(dcrStepResolveUpstream, req.Issuer, redirectURI, + fmt.Errorf("resolve upstream identity: %w", err)) + } + scopes := slices.Clone(req.Scopes) key := Key{ Issuer: req.Issuer, + UpstreamID: upstreamID, RedirectURI: redirectURI, ScopesHash: storage.ScopesHash(scopes), } @@ -445,6 +478,7 @@ func registerAndCache( //nolint:gosec // G706: client_id is public metadata per RFC 7591. slog.Debug("dcr: registered new client", "local_issuer", req.Issuer, + "upstream_id", key.UpstreamID, "redirect_uri", redirectURI, "client_id", resolution.ClientID, ) @@ -658,6 +692,7 @@ func lookupCachedResolution( //nolint:gosec // G706: client_id is public metadata per RFC 7591. slog.Debug("dcr: cache hit ignored; cached secret expired per upstream client_secret_expires_at", "local_issuer", localIssuer, + "upstream_id", key.UpstreamID, "redirect_uri", redirectURI, "client_id", cached.ClientID, "client_secret_expires_at", cached.ClientSecretExpiresAt.UTC().Format(time.RFC3339), @@ -671,6 +706,7 @@ func lookupCachedResolution( //nolint:gosec // G706: client_id is public metadata per RFC 7591. slog.Debug("dcr: cache hit", "local_issuer", localIssuer, + "upstream_id", key.UpstreamID, "redirect_uri", redirectURI, "client_id", cached.ClientID, "dcr_age_days", ageDays, @@ -683,6 +719,7 @@ func lookupCachedResolution( "consider rotating the registration via RFC 7592 deregistration "+ "and re-registering at next startup", "local_issuer", localIssuer, + "upstream_id", key.UpstreamID, "redirect_uri", redirectURI, "client_id", cached.ClientID, "dcr_age_days", ageDays, @@ -750,7 +787,10 @@ func performRegistration( registrationEndpoint, redirectURI, authMethod string, scopes []string, ) (*oauthproto.DynamicClientRegistrationResponse, error) { - httpClient := newDCRHTTPClient(req.InitialAccessToken) + httpClient, err := newDCRHTTPClient(req.InitialAccessToken, registrationEndpoint, req.AllowPrivateIPs) + if err != nil { + return nil, fmt.Errorf("dcr: build registration http client: %w", err) + } clientName := req.ClientName if clientName == "" { @@ -891,10 +931,77 @@ func resolveDCREndpoints( return nil, err } - metadata, err := oauthproto.FetchAuthorizationServerMetadataFromURL(ctx, req.DiscoveryURL, upstreamIssuer, nil) + // Dial the discovery fetch through the private-IP-guarded client so a + // DiscoveryURL that resolves to a private/loopback/link-local address — + // or rebinds to one after the caller validated it — is refused at connect + // time (CWE-918). Guarded by default; req.AllowPrivateIPs opts in to + // private ranges for an in-cluster upstream. + discoveryHost, err := hostFromURL(req.DiscoveryURL) + if err != nil { + return nil, err + } + discoveryClient, err := newGuardedDCRClient(discoveryHost, req.AllowPrivateIPs) + if err != nil { + return nil, fmt.Errorf("dcr: build discovery http client: %w", err) + } + // The discovery URL is operator-configured, but the document it returns — + // and any 30x the upstream serves in response — is upstream-controlled. + // Restrict the fetch to same-host redirects so a malicious redirect cannot + // walk the request onto another host (CWE-918); this matters even when the + // dial guard is relaxed for a loopback discovery host. The registration + // client refuses redirects outright because it carries the bearer token; + // the discovery GET carries no secret, so same-host redirects are allowed — + // matching the CLI discovery clients in pkg/auth/discovery. + discoveryClient.CheckRedirect = networking.SameHostRedirectPolicy() + + metadata, err := oauthproto.FetchAuthorizationServerMetadataFromURL(ctx, req.DiscoveryURL, upstreamIssuer, discoveryClient) return endpointsFromMetadata(metadata, err, upstreamIssuer) } +// resolveUpstreamKeyIdentity returns the stable identifier for the upstream +// authorization server a registration is bound to, used as the +// Key.UpstreamID cache component. It disambiguates upstreams that share the +// caller's Issuer, RedirectURI, and scope set — the common case inside a +// single embedded authserver — so they receive distinct dynamically- +// registered clients instead of colliding on one cache entry (issue #5823). +// +// The identity source mirrors resolveDCREndpoints' own branch order so the +// key names the exact server the client registers against: +// +// - RegistrationEndpoint set: the registration endpoint URL itself. +// - DiscoveryURL set: the upstream issuer recovered from the discovery URL +// via deriveExpectedIssuerFromDiscoveryURL — the same value used for +// RFC 8414 §3.3 metadata verification. +// +// validateResolveInputs guarantees exactly one of the two is set, so the +// final return is only reached on the DiscoveryURL branch. The derivation is +// pure URL parsing (no network I/O), so computing it before the cache lookup +// keeps the cache-hit path free of I/O; a malformed DiscoveryURL that fails +// here could never have produced a stored entry to hit anyway, because a +// successful registration requires the same derivation to succeed first. +// +// Residual limitation: on the DiscoveryURL branch the identity is the derived +// upstream *issuer*, not the discovery URL. For a custom (non-well-known) +// discovery URL, deriveExpectedIssuerFromDiscoveryURL falls back to the bare +// origin and discards the path, so two upstreams on the same host with +// distinct custom metadata paths, the same scopes, and both declaring that +// origin as their metadata "issuer" derive the same UpstreamID and still +// collide. This is intentional: UpstreamID names the authorization server +// (aligned with the RFC 8414 §3.3 issuer used for metadata verification), so +// two configs that resolve to the same issuer are the same AS and correctly +// share one registration. The collision is therefore confined to the +// nonstandard case where a single issuer serves distinct registration +// endpoints under custom (non-well-known) discovery paths: both upstreams +// resolve to that one issuer, so their UpstreamIDs match. Upstreams that +// resolve to different issuers derive different UpstreamIDs and never collide +// — the key itself keeps them apart, independent of §3.3 verification. +func resolveUpstreamKeyIdentity(req *Request) (string, error) { + if req.RegistrationEndpoint != "" { + return req.RegistrationEndpoint, nil + } + return deriveExpectedIssuerFromDiscoveryURL(req.DiscoveryURL) +} + // deriveExpectedIssuerFromDiscoveryURL recovers the issuer identifier the // upstream is expected to claim in its RFC 8414 / OIDC Discovery document, // given an operator-configured DiscoveryURL. @@ -1029,6 +1136,12 @@ func endpointsFromMetadata( return nil, fmt.Errorf("synthesise registration endpoint: %w", err) } registrationEndpoint = synth + } else if err := validateUpstreamEndpointURL(registrationEndpoint, "registration_endpoint"); err != nil { + // Unlike the synthesised branch above, this value came straight from + // the discovery document — validate it the same as + // authorization_endpoint/token_endpoint rather than letting it reach + // hostFromURL unvalidated. + return nil, fmt.Errorf("dcr: discovered %w", err) } return &dcrEndpoints{ @@ -1263,27 +1376,33 @@ var errDCRRedirectRefused = errors.New( "to avoid forwarding the RFC 7591 initial access token to a foreign origin") // newDCRHTTPClient returns the http.Client to pass to -// oauthproto.RegisterClientDynamically. The client always blocks HTTP -// redirects so that an upstream cannot use a 30x to coerce us into +// oauthproto.RegisterClientDynamically for a registration POST to +// registrationEndpoint. The client dials through the private-IP-guarded +// transport built by newGuardedDCRClient (CWE-918 SSRF protection) and always +// blocks HTTP redirects so that an upstream cannot use a 30x to coerce us into // re-issuing the registration request (and any attached // Authorization: Bearer header) against a different origin. RFC 7591 §3 // does not require redirect support, so refusing them is safe. // -// When initialAccessToken is non-empty the client also wraps the canonical -// DCR client's transport with a bearerTokenTransport that injects the -// Authorization header. The combination of the bearer transport plus the -// redirect block is what prevents the token-leak class of bug. -// -// The timeout policy is sourced from oauthproto.NewDefaultDCRClient so -// future tightening of those bounds propagates automatically. -func newDCRHTTPClient(initialAccessToken string) *http.Client { - client := oauthproto.NewDefaultDCRClient() +// When initialAccessToken is non-empty the client also wraps the guarded +// transport with a bearerTokenTransport that injects the Authorization header. +// The combination of the bearer transport plus the redirect block is what +// prevents the token-leak class of bug. +func newDCRHTTPClient(initialAccessToken, registrationEndpoint string, allowPrivateIPs bool) (*http.Client, error) { + host, err := hostFromURL(registrationEndpoint) + if err != nil { + return nil, err + } + client, err := newGuardedDCRClient(host, allowPrivateIPs) + if err != nil { + return nil, err + } client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { return errDCRRedirectRefused } if initialAccessToken == "" { - return client + return client, nil } next := client.Transport @@ -1294,5 +1413,44 @@ func newDCRHTTPClient(initialAccessToken string) *http.Client { token: initialAccessToken, next: next, } - return client + return client, nil +} + +// newGuardedDCRClient builds the private-IP-guarded *http.Client used for both +// of the resolver's outbound calls — the discovery fetch and the registration +// POST. See networking.NewHostScopedClientBuilder for the CWE-918 guard policy +// this applies (allowPrivateIPs semantics, loopback exemption, HTTPS +// enforcement). Keep-alive is disabled so the dial-time check re-runs on +// every request rather than being bypassed by a pooled connection. +// +// The returned client has no CheckRedirect policy: each caller layers its own +// (the registration client refuses all redirects to protect the bearer token; +// the discovery client restricts them to the same host). The dial guard alone +// does not stop a redirect to a different public host, so the redirect policy +// is a required complement, not an optional one. +func newGuardedDCRClient(host string, allowPrivateIPs bool) (*http.Client, error) { + return networking.NewHostScopedClientBuilder(host, allowPrivateIPs, false). + WithDisableKeepAlives(true). + Build() +} + +// hostFromURL extracts the host[:port] component used to scope the guarded +// HTTP client. Every URL reaching this helper has already passed +// scheme-and-host validation at the resolver's entry points +// (validateUpstreamEndpointURL for the registration endpoint — both the +// caller-supplied case and the metadata-discovered case handled in +// endpointsFromMetadata — and FetchAuthorizationServerMetadataFromURL for the +// discovery URL; the synthesised-endpoint case derives its host from an +// upstream issuer that already passed the RFC 8414 §3.3 issuer-match check), +// so a parse failure or empty host here signals an internal inconsistency +// rather than untrusted input. +func hostFromURL(rawURL string) (string, error) { + u, err := url.Parse(rawURL) + if err != nil { + return "", fmt.Errorf("dcr: parse url for http client host: %w", err) + } + if u.Host == "" { + return "", fmt.Errorf("dcr: url missing host: %q", rawURL) + } + return u.Host, nil } diff --git a/pkg/auth/dcr/resolver_test.go b/pkg/auth/dcr/resolver_test.go index 2c163e110d..caf7f07481 100644 --- a/pkg/auth/dcr/resolver_test.go +++ b/pkg/auth/dcr/resolver_test.go @@ -22,6 +22,7 @@ import ( "github.com/stretchr/testify/require" "github.com/stacklok/toolhive/pkg/authserver/storage" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/oauthproto" ) @@ -58,6 +59,11 @@ type dcrTestHandlerConfig struct { // does not expire". clientIDIssuedAt int64 clientSecretExpiresAt int64 + + // clientID overrides the client_id returned by the registration + // endpoint. Defaults to "test-client-id". Distinct values let a test + // prove two upstreams received distinct dynamically-registered clients. + clientID string } // newDCRTestServer mounts RFC 8414 metadata and a DCR endpoint on a single @@ -107,11 +113,15 @@ func newDCRTestServer(t *testing.T, cfg dcrTestHandlerConfig) *httptest.Server { http.Error(w, err.Error(), http.StatusBadRequest) return } + clientID := cfg.clientID + if clientID == "" { + clientID = "test-client-id" + } resp := oauthproto.DynamicClientRegistrationResponse{ - ClientID: "test-client-id", + ClientID: clientID, ClientSecret: "test-client-secret", RegistrationAccessToken: "test-reg-token", - RegistrationClientURI: server.URL + "/register/test-client-id", + RegistrationClientURI: server.URL + "/register/" + clientID, TokenEndpointAuthMethod: req.TokenEndpointAuthMethod, ClientIDIssuedAt: cfg.clientIDIssuedAt, ClientSecretExpiresAt: cfg.clientSecretExpiresAt, @@ -145,8 +155,12 @@ func TestResolveDCRCredentials_CacheHitShortCircuits(t *testing.T) { // Pre-populate the cache with a resolution matching the key we will // look up. redirectURI := issuer + "/oauth/callback" + // UpstreamID is the upstream issuer the resolver derives from the + // DiscoveryURL below; for a well-known discovery URL that is the issuer + // origin itself. The preloaded key must match it for the cache to hit. key := Key{ Issuer: issuer, + UpstreamID: issuer, RedirectURI: redirectURI, ScopesHash: storage.ScopesHash([]string{"openid", "profile"}), } @@ -213,14 +227,229 @@ func TestResolveDCRCredentials_RegistersOnCacheMiss(t *testing.T) { assert.Equal(t, []string{issuer + "/oauth/callback"}, dcrReq.RedirectURIs) assert.ElementsMatch(t, []string{"openid", "profile"}, []string(dcrReq.Scopes)) - // Cache was populated. + // Cache was populated. UpstreamID is the issuer origin the resolver + // derives from the oauth-authorization-server discovery URL. cached, ok, err := cache.Get(context.Background(), - Key{Issuer: issuer, RedirectURI: issuer + "/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid", "profile"})}) + Key{Issuer: issuer, UpstreamID: issuer, RedirectURI: issuer + "/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid", "profile"})}) require.NoError(t, err) require.True(t, ok) assert.Equal(t, "test-client-id", cached.ClientID) } +// TestResolveDCRCredentials_DistinctUpstreamsSameScopesDoNotCollide is the +// regression test for issue #5823. Within one embedded authserver every +// OAuth2 upstream shares the authserver's own Issuer and the single defaulted +// {issuer}/oauth/callback RedirectURI, so the cache key differed only by the +// scopes hash. Two upstreams configured with equal scopes therefore collided: +// the second read back the first's dynamically-registered client and never +// registered with its own authorization server. The UpstreamID key component +// disambiguates them. +func TestResolveDCRCredentials_DistinctUpstreamsSameScopesDoNotCollide(t *testing.T) { + t.Parallel() + + var regCountA, regCountB int32 + serverA := newDCRTestServer(t, dcrTestHandlerConfig{ + clientID: "client-a", + tokenEndpointAuthMethodsSupported: []string{"client_secret_basic"}, + scopesSupported: []string{"openid", "profile"}, + observeRegistration: func(*http.Request, []byte) { atomic.AddInt32(®CountA, 1) }, + }) + serverB := newDCRTestServer(t, dcrTestHandlerConfig{ + clientID: "client-b", + tokenEndpointAuthMethodsSupported: []string{"client_secret_basic"}, + scopesSupported: []string{"openid", "profile"}, + observeRegistration: func(*http.Request, []byte) { atomic.AddInt32(®CountB, 1) }, + }) + + // One shared cache, one shared local issuer, one shared scope set — the + // embedded-authserver shape. Only the upstream (DiscoveryURL) differs. + cache := newMemoryDCRStore(t) + const localIssuer = "https://authserver.example.com" + ctx := context.Background() + reqFor := func(discoveryHost string) *Request { + return &Request{ + Issuer: localIssuer, + Scopes: []string{"openid", "profile"}, + DiscoveryURL: discoveryHost + "/.well-known/oauth-authorization-server", + } + } + + resA, err := ResolveCredentials(ctx, reqFor(serverA.URL), cache) + require.NoError(t, err) + resB, err := ResolveCredentials(ctx, reqFor(serverB.URL), cache) + require.NoError(t, err) + + // The crux: the second upstream registered with its OWN authorization + // server and received its own client, rather than reading back the + // first's cached credentials. + assert.Equal(t, "client-a", resA.ClientID) + assert.Equal(t, "client-b", resB.ClientID, + "second upstream must register with its own AS, not reuse the first's cached client (#5823)") + assert.Equal(t, serverA.URL+"/token", resA.TokenEndpoint) + assert.Equal(t, serverB.URL+"/token", resB.TokenEndpoint, + "second upstream's endpoints must come from its own AS metadata") + assert.Equal(t, int32(1), atomic.LoadInt32(®CountA), "upstream A must register exactly once") + assert.Equal(t, int32(1), atomic.LoadInt32(®CountB), + "upstream B must register exactly once; a cache collision would leave this at zero") + + // Both resolutions coexist in the cache under keys that differ only by + // UpstreamID (shared Issuer, RedirectURI, and scopes hash). + redirectURI := localIssuer + "/oauth/callback" + scopesHash := storage.ScopesHash([]string{"openid", "profile"}) + cachedA, okA, err := cache.Get(ctx, + Key{Issuer: localIssuer, UpstreamID: serverA.URL, RedirectURI: redirectURI, ScopesHash: scopesHash}) + require.NoError(t, err) + require.True(t, okA) + assert.Equal(t, "client-a", cachedA.ClientID) + cachedB, okB, err := cache.Get(ctx, + Key{Issuer: localIssuer, UpstreamID: serverB.URL, RedirectURI: redirectURI, ScopesHash: scopesHash}) + require.NoError(t, err) + require.True(t, okB) + assert.Equal(t, "client-b", cachedB.ClientID) +} + +// TestResolveDCRCredentials_ConcurrentDistinctUpstreamsBothRegister guards the +// singleflight (flight-key) split, complementing the sequential +// TestResolveDCRCredentials_DistinctUpstreamsSameScopesDoNotCollide which only +// proves the persistent cache-key split. Two upstreams that share the caller's +// Issuer, RedirectURI, and scopes but differ by UpstreamID must not coalesce +// into one dcrFlight — each must register with its own AS. A barrier in the +// registration handlers forces both calls to be in-flight simultaneously, so a +// flight-key collision would surface as one call becoming the other's follower +// (its registration count staying at zero). +func TestResolveDCRCredentials_ConcurrentDistinctUpstreamsBothRegister(t *testing.T) { + t.Parallel() + + // Barrier: release both registration handlers only once both have been + // entered, forcing genuine in-flight overlap. The timeout keeps the test + // fail-fast — if the two calls coalesced there would be only one handler + // to enter, so the sole handler proceeds after the wait and the reg-count + // assertions below catch the collision instead of the test hanging. + var regCountA, regCountB, arrivals int32 + bothArrived := make(chan struct{}) + barrier := func() { + if atomic.AddInt32(&arrivals, 1) == 2 { + close(bothArrived) + } + select { + case <-bothArrived: + case <-time.After(3 * time.Second): + } + } + + serverA := newDCRTestServer(t, dcrTestHandlerConfig{ + clientID: "client-a", + tokenEndpointAuthMethodsSupported: []string{"client_secret_basic"}, + scopesSupported: []string{"openid", "profile"}, + observeRegistration: func(*http.Request, []byte) { atomic.AddInt32(®CountA, 1); barrier() }, + }) + serverB := newDCRTestServer(t, dcrTestHandlerConfig{ + clientID: "client-b", + tokenEndpointAuthMethodsSupported: []string{"client_secret_basic"}, + scopesSupported: []string{"openid", "profile"}, + observeRegistration: func(*http.Request, []byte) { atomic.AddInt32(®CountB, 1); barrier() }, + }) + + cache := newMemoryDCRStore(t) + const localIssuer = "https://authserver.example.com" + reqFor := func(discoveryHost string) *Request { + return &Request{ + Issuer: localIssuer, + Scopes: []string{"openid", "profile"}, + DiscoveryURL: discoveryHost + "/.well-known/oauth-authorization-server", + } + } + + var wg sync.WaitGroup + wg.Add(2) + results := make([]*Resolution, 2) + errs := make([]error, 2) + go func() { + defer wg.Done() + results[0], errs[0] = ResolveCredentials(context.Background(), reqFor(serverA.URL), cache) + }() + go func() { + defer wg.Done() + results[1], errs[1] = ResolveCredentials(context.Background(), reqFor(serverB.URL), cache) + }() + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent ResolveCredentials calls; possible flight-key deadlock") + } + + require.NoError(t, errs[0]) + require.NoError(t, errs[1]) + assert.Equal(t, "client-a", results[0].ClientID) + assert.Equal(t, "client-b", results[1].ClientID, + "each concurrent upstream must receive its own client, not the other flight's leader result") + assert.Equal(t, int32(1), atomic.LoadInt32(®CountA), "upstream A must register exactly once") + assert.Equal(t, int32(1), atomic.LoadInt32(®CountB), + "upstream B must register exactly once; a flight-key collision would coalesce the two calls and leave this at zero") +} + +// TestResolveDCRCredentials_DistinctRegistrationEndpointsDoNotCollide covers +// the RegistrationEndpoint branch of resolveUpstreamKeyIdentity: two upstreams +// configured with explicit (distinct) registration endpoints but the same +// caller Issuer, RedirectURI, and scopes must still register separately. The +// UpstreamID is the registration endpoint URL on this branch. +func TestResolveDCRCredentials_DistinctRegistrationEndpointsDoNotCollide(t *testing.T) { + t.Parallel() + + var regCountA, regCountB int32 + serverA := newDCRTestServer(t, dcrTestHandlerConfig{ + clientID: "client-a", + observeRegistration: func(*http.Request, []byte) { atomic.AddInt32(®CountA, 1) }, + }) + serverB := newDCRTestServer(t, dcrTestHandlerConfig{ + clientID: "client-b", + observeRegistration: func(*http.Request, []byte) { atomic.AddInt32(®CountB, 1) }, + }) + + cache := newMemoryDCRStore(t) + const localIssuer = "https://authserver.example.com" + ctx := context.Background() + // The RegistrationEndpoint branch performs no discovery, so the caller + // must supply authorization/token endpoints explicitly. + reqFor := func(server *httptest.Server) *Request { + return &Request{ + Issuer: localIssuer, + Scopes: []string{"openid", "profile"}, + RegistrationEndpoint: server.URL + "/register", + AuthorizationEndpoint: server.URL + "/authorize", + TokenEndpoint: server.URL + "/token", + } + } + + resA, err := ResolveCredentials(ctx, reqFor(serverA), cache) + require.NoError(t, err) + resB, err := ResolveCredentials(ctx, reqFor(serverB), cache) + require.NoError(t, err) + + assert.Equal(t, "client-a", resA.ClientID) + assert.Equal(t, "client-b", resB.ClientID, + "second upstream must register with its own registration endpoint, not reuse the first's cached client (#5823)") + assert.Equal(t, int32(1), atomic.LoadInt32(®CountA), "upstream A must register exactly once") + assert.Equal(t, int32(1), atomic.LoadInt32(®CountB), + "upstream B must register exactly once; a cache collision would leave this at zero") + + // The two entries coexist under keys whose only differing component is + // UpstreamID (the registration endpoint URL). + redirectURI := localIssuer + "/oauth/callback" + scopesHash := storage.ScopesHash([]string{"openid", "profile"}) + _, okA, err := cache.Get(ctx, + Key{Issuer: localIssuer, UpstreamID: serverA.URL + "/register", RedirectURI: redirectURI, ScopesHash: scopesHash}) + require.NoError(t, err) + assert.True(t, okA, "upstream A entry must be keyed by its registration endpoint") + _, okB, err := cache.Get(ctx, + Key{Issuer: localIssuer, UpstreamID: serverB.URL + "/register", RedirectURI: redirectURI, ScopesHash: scopesHash}) + require.NoError(t, err) + assert.True(t, okB, "upstream B entry must be keyed by its registration endpoint") +} + func TestResolveDCRCredentials_ExplicitEndpointsOverride(t *testing.T) { t.Parallel() @@ -343,6 +572,235 @@ func TestResolveDCRCredentials_DoesNotForwardBearerOnRedirect(t *testing.T) { atomic.LoadInt32(&foreignHits), foreignAuthHeaders) } +// TestResolveDCRCredentials_BlocksPrivateIPTargets pins the CWE-918 SSRF guard +// added for issue #5825: both of the resolver's outbound calls — the discovery +// fetch and the registration POST — are dialed through a private-IP-guarded +// client, so a registration endpoint that resolves to a private or link-local +// address is refused at connect time. The guard fires before any bytes leave +// the host, so these cases need no live server behind the private target. +// AllowPrivateIPs defaults to false; loopback stays permitted for development. +func TestResolveDCRCredentials_BlocksPrivateIPTargets(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + newReq func(t *testing.T) *Request + }{ + { + // Direct RegistrationEndpoint branch: an operator-configured + // endpoint that resolves (or rebinds) to an RFC-1918 address. + name: "direct registration endpoint on a private IP", + newReq: func(_ *testing.T) *Request { + // Issuer is unique per test: the process-global dcrFlight + // singleflight keys on (Issuer, RedirectURI, ScopesHash), so a + // shared synthetic issuer would coalesce this call with another + // parallel test and return that test's result. + return &Request{ + Issuer: "https://block-direct.example.test", + Scopes: []string{"openid"}, + RegistrationEndpoint: "https://10.255.255.1/register", + } + }, + }, + { + // Discovery-indirection branch: the discovery document is served + // from an allowed (loopback) host, but it points + // registration_endpoint at a link-local metadata address the + // metadata itself controls. HTTPS so it clears scheme validation + // and reaches the guarded dial. + name: "registration endpoint from discovery on a link-local IP", + newReq: func(t *testing.T) *Request { + t.Helper() + mux := http.NewServeMux() + var server *httptest.Server + mux.HandleFunc("/.well-known/oauth-authorization-server", + func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(oauthproto.AuthorizationServerMetadata{ + Issuer: server.URL, + AuthorizationEndpoint: server.URL + "/authorize", + TokenEndpoint: server.URL + "/token", + JWKSURI: server.URL + "/jwks", + RegistrationEndpoint: "https://169.254.169.254/register", + }) + }) + server = httptest.NewServer(mux) + t.Cleanup(server.Close) + return &Request{ + Issuer: server.URL, + Scopes: []string{"openid"}, + DiscoveryURL: server.URL + "/.well-known/oauth-authorization-server", + } + }, + }, + { + // Discovery-fetch branch itself: DiscoveryURL points directly at a + // private IP, so the *discovery* client's own dial guard + // (newGuardedDCRClient(discoveryHost, ...) in resolveDCREndpoints) + // must refuse it. The two cases above only prove the guard on the + // registration client; this proves it on the discovery client. No + // live server is needed — the guard fires before any request is + // sent, and deriveExpectedIssuerFromDiscoveryURL resolves the + // issuer from the URL string alone. + name: "discovery URL itself on a private IP", + newReq: func(_ *testing.T) *Request { + return &Request{ + Issuer: "https://block-discovery.example.test", + Scopes: []string{"openid"}, + DiscoveryURL: "https://10.255.255.1/.well-known/oauth-authorization-server", + } + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + _, err := ResolveCredentials(context.Background(), tc.newReq(t), newMemoryDCRStore(t)) + require.Error(t, err) + assert.ErrorContains(t, err, networking.ErrPrivateIpAddress, + "a private/link-local registration target must be refused at connect time") + }) + } +} + +// TestResolveDCRCredentials_AllowPrivateIPsHonored proves req.AllowPrivateIPs +// is threaded through to the guarded client: with it set, the resolver dials +// the private target instead of refusing it at the guard. The target is a +// non-routable RFC 5737 documentation address (TEST-NET-1), so the dial fails +// with a network error rather than the guard error — and can never reach a +// real host. +func TestResolveDCRCredentials_AllowPrivateIPsHonored(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + // Issuer is unique per test so the process-global dcrFlight singleflight + // (keyed on Issuer, RedirectURI, ScopesHash) cannot coalesce this call + // with another parallel test and hand back that test's result. + req := &Request{ + Issuer: "https://allow-private.example.test", + Scopes: []string{"openid"}, + RegistrationEndpoint: "https://192.0.2.1/register", + AllowPrivateIPs: true, + } + + _, err := ResolveCredentials(ctx, req, newMemoryDCRStore(t)) + require.Error(t, err, "the dead documentation address cannot complete registration") + assert.NotContains(t, err.Error(), networking.ErrPrivateIpAddress, + "AllowPrivateIPs=true must lift the guard so the dial is attempted, not refused") +} + +// TestResolveDCRCredentials_AllowPrivateIPsHonoredViaDiscovery proves +// req.AllowPrivateIPs also lifts the guard on the discovery-indirection +// path — not just the direct RegistrationEndpoint branch covered by +// TestResolveDCRCredentials_AllowPrivateIPsHonored above: with it set, a +// registration_endpoint the discovery document points at a private IP is +// dialed instead of refused at the guard. +func TestResolveDCRCredentials_AllowPrivateIPsHonoredViaDiscovery(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + mux := http.NewServeMux() + var server *httptest.Server + mux.HandleFunc("/.well-known/oauth-authorization-server", + func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(oauthproto.AuthorizationServerMetadata{ + Issuer: server.URL, + AuthorizationEndpoint: server.URL + "/authorize", + TokenEndpoint: server.URL + "/token", + JWKSURI: server.URL + "/jwks", + // Non-routable RFC 5737 documentation address (TEST-NET-1): the + // dial fails with a network error rather than reaching a real + // host, same as TestResolveDCRCredentials_AllowPrivateIPsHonored. + RegistrationEndpoint: "https://192.0.2.1/register", + }) + }) + server = httptest.NewServer(mux) + t.Cleanup(server.Close) + + req := &Request{ + Issuer: server.URL, + Scopes: []string{"openid"}, + DiscoveryURL: server.URL + "/.well-known/oauth-authorization-server", + AllowPrivateIPs: true, + } + + _, err := ResolveCredentials(ctx, req, newMemoryDCRStore(t)) + require.Error(t, err, "the dead documentation address cannot complete registration") + assert.NotContains(t, err.Error(), networking.ErrPrivateIpAddress, + "AllowPrivateIPs=true must lift the guard on the discovery-indirection path too") +} + +// TestResolveDCRCredentials_DiscoveryRefusesCrossHostRedirect pins that the +// discovery fetch installs SameHostRedirectPolicy: a discovery endpoint that +// 30x-redirects to a different host must not be followed, so a malicious +// upstream cannot walk the metadata fetch onto an unintended origin (CWE-918). +// A different port counts as a different host, so two loopback httptest servers +// suffice. +func TestResolveDCRCredentials_DiscoveryRefusesCrossHostRedirect(t *testing.T) { + t.Parallel() + + var foreignHits int32 + foreign := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + atomic.AddInt32(&foreignHits, 1) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(foreign.Close) + + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, foreign.URL+"/.well-known/oauth-authorization-server", http.StatusFound) + }) + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + + req := &Request{ + Issuer: server.URL, + Scopes: []string{"openid"}, + DiscoveryURL: server.URL + "/.well-known/oauth-authorization-server", + } + + _, err := ResolveCredentials(context.Background(), req, newMemoryDCRStore(t)) + require.Error(t, err, "discovery must fail when the endpoint redirects cross-host") + assert.EqualValues(t, 0, atomic.LoadInt32(&foreignHits), + "the discovery client must not follow a cross-host redirect") +} + +func TestHostFromURL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + rawURL string + want string + wantErr bool + }{ + {name: "https with port", rawURL: "https://idp.example.com:8443/register", want: "idp.example.com:8443"}, + {name: "https without port", rawURL: "https://idp.example.com/register", want: "idp.example.com"}, + {name: "loopback with port", rawURL: "http://127.0.0.1:5000/register", want: "127.0.0.1:5000"}, + {name: "missing host", rawURL: "/register", wantErr: true}, + {name: "malformed url", rawURL: "://bad", wantErr: true}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, err := hostFromURL(tc.rawURL) + if tc.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tc.want, got) + }) + } +} + func TestResolveDCRCredentials_AuthMethodPreference(t *testing.T) { t.Parallel() @@ -1682,6 +2140,28 @@ func TestDcrStepError(t *testing.T) { require.True(t, errors.As(err, &stepErr)) assert.Equal(t, dcrStepValidate, stepErr.Step) }) + + t.Run("malformed DiscoveryURL fails at dcrStepResolveUpstream", func(t *testing.T) { + t.Parallel() + + // A DiscoveryURL with no scheme/host passes validateResolveInputs + // (which only checks non-emptiness) and the Issuer-based redirect + // derivation, then fails in resolveUpstreamKeyIdentity — + // deriveExpectedIssuerFromDiscoveryURL rejects the missing origin. + // This is a new failure mode introduced with UpstreamID: a malformed + // DiscoveryURL now surfaces here rather than later in + // resolveDCREndpoints, so the step must be dcrStepResolveUpstream. + req := &Request{ + Issuer: "https://authserver.example.com", + Scopes: []string{"openid"}, + DiscoveryURL: "not-a-url", + } + _, err := ResolveCredentials(context.Background(), req, newMemoryDCRStore(t)) + require.Error(t, err) + var stepErr *dcrStepError + require.True(t, errors.As(err, &stepErr)) + assert.Equal(t, dcrStepResolveUpstream, stepErr.Step) + }) } func TestSanitizeErrorForLog(t *testing.T) { diff --git a/pkg/auth/dcr/store.go b/pkg/auth/dcr/store.go index 8694e6bbd6..c6b56ce470 100644 --- a/pkg/auth/dcr/store.go +++ b/pkg/auth/dcr/store.go @@ -30,8 +30,8 @@ const dcrStaleAgeThreshold = 90 * 24 * time.Hour type Key = storage.DCRKey // CredentialStore caches RFC 7591 Dynamic Client Registration resolutions -// keyed by the (Issuer, RedirectURI, ScopesHash) tuple. Implementations must -// be safe for concurrent use. +// keyed by the (Issuer, UpstreamID, RedirectURI, ScopesHash) tuple. +// Implementations must be safe for concurrent use. // // This is the runner-facing interface used by the DCR resolver. It is a // narrow re-projection of storage.DCRCredentialStore that exchanges diff --git a/pkg/auth/dcr/store_test.go b/pkg/auth/dcr/store_test.go index 51e7980ea6..956663dee7 100644 --- a/pkg/auth/dcr/store_test.go +++ b/pkg/auth/dcr/store_test.go @@ -25,6 +25,7 @@ func TestStorageBackedStore_PutGet_RoundTrip(t *testing.T) { key := Key{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://toolhive.example.com/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid", "profile"}), } @@ -73,24 +74,38 @@ func TestStorageBackedStore_DistinctKeysDoNotCollide(t *testing.T) { keyA := Key{ Issuer: "https://idp-a.example.com", + UpstreamID: "https://upstream-a.example.com", RedirectURI: "https://toolhive.example.com/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid"}), } keyB := Key{ Issuer: "https://idp-b.example.com", + UpstreamID: "https://upstream-a.example.com", RedirectURI: "https://toolhive.example.com/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid"}), } keyC := Key{ Issuer: "https://idp-a.example.com", + UpstreamID: "https://upstream-a.example.com", RedirectURI: "https://other.example.com/callback", ScopesHash: storage.ScopesHash([]string{"openid"}), } keyD := Key{ Issuer: "https://idp-a.example.com", + UpstreamID: "https://upstream-a.example.com", RedirectURI: "https://toolhive.example.com/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid", "email"}), } + // keyE shares Issuer, RedirectURI, and scopes with keyA and differs only + // by UpstreamID — the exact shape of two OAuth2 upstreams inside one + // embedded authserver. Before UpstreamID was part of the key these two + // collided (issue #5823); keyE must resolve to its own entry. + keyE := Key{ + Issuer: "https://idp-a.example.com", + UpstreamID: "https://upstream-b.example.com", + RedirectURI: "https://toolhive.example.com/oauth/callback", + ScopesHash: storage.ScopesHash([]string{"openid"}), + } // The persisted *storage.DCRCredentials shape requires non-empty // AuthorizationEndpoint / TokenEndpoint per validateDCRCredentialsForStore; @@ -108,6 +123,7 @@ func TestStorageBackedStore_DistinctKeysDoNotCollide(t *testing.T) { require.NoError(t, store.Put(ctx, keyB, resolution("b"))) require.NoError(t, store.Put(ctx, keyC, resolution("c"))) require.NoError(t, store.Put(ctx, keyD, resolution("d"))) + require.NoError(t, store.Put(ctx, keyE, resolution("e"))) for _, tc := range []struct { key Key @@ -117,6 +133,7 @@ func TestStorageBackedStore_DistinctKeysDoNotCollide(t *testing.T) { {keyB, "b"}, {keyC, "c"}, {keyD, "d"}, + {keyE, "e"}, } { got, ok, err := store.Get(ctx, tc.key) require.NoError(t, err) @@ -133,6 +150,7 @@ func TestStorageBackedStore_Put_OverwritesExisting(t *testing.T) { key := Key{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x.example.com/cb", ScopesHash: storage.ScopesHash([]string{"openid"}), } @@ -189,6 +207,7 @@ func TestStorageBackedStore_GetReturnsDefensiveCopy(t *testing.T) { key := Key{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x.example.com/cb", ScopesHash: storage.ScopesHash([]string{"openid"}), } @@ -241,6 +260,7 @@ func TestStorageBackedStore_ConcurrentAccess(t *testing.T) { overlappingKey := func(i int) Key { return Key{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: fmt.Sprintf("overlap-%d", i%4), } @@ -248,6 +268,7 @@ func TestStorageBackedStore_ConcurrentAccess(t *testing.T) { disjointKey := func(worker, i int) Key { return Key{ Issuer: fmt.Sprintf("https://idp-%d.example.com", worker), + UpstreamID: "https://upstream.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: fmt.Sprintf("disjoint-%d", i), } @@ -329,6 +350,7 @@ func TestResolutionCredentialsRoundTrip(t *testing.T) { key := Key{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid", "profile"}), } @@ -476,6 +498,7 @@ func TestInMemoryStore_PutGetCloseShareBackend(t *testing.T) { ctx := context.Background() key := Key{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://toolhive.example.com/oauth/callback", ScopesHash: storage.ScopesHash([]string{"openid"}), } diff --git a/pkg/auth/discovery/dcr_resolver_test.go b/pkg/auth/discovery/dcr_resolver_test.go index 2886d78ac2..cbd2c1e94f 100644 --- a/pkg/auth/discovery/dcr_resolver_test.go +++ b/pkg/auth/discovery/dcr_resolver_test.go @@ -17,6 +17,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/oauthproto" ) @@ -80,6 +81,11 @@ type dcrTestServerConfig struct { // handler so concurrent goroutines can pile up at the singleflight // before any can finish. registrationDelay time.Duration + + // Optional RFC 7591 response metadata used by renewal happy-path tests. + clientSecretExpiresAt int64 + registrationAccessToken string + registrationClientURI string } // newDCRDiscoveryServer mounts /.well-known/openid-configuration and a @@ -129,6 +135,9 @@ func newDCRDiscoveryServer(t *testing.T, cfg dcrTestServerConfig) *httptest.Serv } resp := oauthproto.DynamicClientRegistrationResponse{ ClientID: "cli-registered-client", + ClientSecretExpiresAt: cfg.clientSecretExpiresAt, + RegistrationAccessToken: cfg.registrationAccessToken, + RegistrationClientURI: cfg.registrationClientURI, TokenEndpointAuthMethod: "none", } w.Header().Set("Content-Type", "application/json") @@ -172,6 +181,8 @@ func TestHandleDynamicRegistration_InheritsS256Gating(t *testing.T) { config := &OAuthFlowConfig{ Scopes: []string{"openid", "profile"}, CallbackPort: 8765, + // Loopback test server: guard would otherwise refuse to dial it. + AllowPrivateIPs: true, } err := handleDynamicRegistration(context.Background(), server.URL, config) @@ -206,8 +217,9 @@ func TestHandleDynamicRegistration_InheritsRedirectRefusal(t *testing.T) { }) config := &OAuthFlowConfig{ - Scopes: []string{"openid", "profile"}, - CallbackPort: 8765, + Scopes: []string{"openid", "profile"}, + CallbackPort: 8765, + AllowPrivateIPs: true, // loopback test server; guard would otherwise refuse to dial it } err := handleDynamicRegistration(context.Background(), server.URL, config) @@ -218,6 +230,54 @@ func TestHandleDynamicRegistration_InheritsRedirectRefusal(t *testing.T) { "foreign origin must receive zero requests; the redirect refusal prevents the leak") } +// TestHandleDynamicRegistration_MetadataRefetchBlocksPrivateIP pins the +// CWE-918 guard on resolveDCRCredentials's own metadata re-fetch (used to +// recover code_challenge_methods_supported on the pre-discovered path): an +// issuer that resolves to a private IP must be refused at connect time, the +// same as the resolver's own outbound calls in pkg/auth/dcr. +func TestHandleDynamicRegistration_MetadataRefetchBlocksPrivateIP(t *testing.T) { + t.Parallel() + + config := &OAuthFlowConfig{ + Scopes: []string{"openid", "profile"}, + CallbackPort: 8765, + AuthorizeURL: "https://10.255.255.1/authorize", + TokenURL: "https://10.255.255.1/token", + RegistrationEndpoint: "https://10.255.255.1/register", + } + + err := handleDynamicRegistration(context.Background(), "https://10.255.255.1", config) + require.Error(t, err, "an issuer resolving to a private IP must be refused at connect time") + assert.ErrorContains(t, err, networking.ErrPrivateIpAddress, + "the metadata re-fetch must be guarded the same as the resolver's own outbound calls") +} + +// TestHandleDynamicRegistration_MetadataRefetchAllowPrivateIPsHonored proves +// OAuthFlowConfig.AllowPrivateIPs lifts the guard on the same metadata +// re-fetch pinned as blocked-by-default above. The target is a non-routable +// RFC 5737 documentation address (TEST-NET-1), so the dial fails with a +// network error rather than the guard error — and can never reach a real host. +func TestHandleDynamicRegistration_MetadataRefetchAllowPrivateIPsHonored(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + config := &OAuthFlowConfig{ + Scopes: []string{"openid", "profile"}, + CallbackPort: 8765, + AuthorizeURL: "https://192.0.2.1/authorize", + TokenURL: "https://192.0.2.1/token", + RegistrationEndpoint: "https://192.0.2.1/register", + AllowPrivateIPs: true, + } + + err := handleDynamicRegistration(ctx, "https://192.0.2.1", config) + require.Error(t, err, "the dead documentation address cannot complete registration") + assert.NotContains(t, err.Error(), networking.ErrPrivateIpAddress, + "AllowPrivateIPs=true must lift the guard on the metadata re-fetch") +} + // TestHandleDynamicRegistration_InheritsSingleflightDedup verifies that // N concurrent handleDynamicRegistration calls for the same (issuer, // scopes, redirectURI) tuple coalesce into exactly one upstream /register @@ -256,8 +316,9 @@ func TestHandleDynamicRegistration_InheritsSingleflightDedup(t *testing.T) { // Each goroutine gets its own OAuthFlowConfig so we can assert // the resolution was applied to every caller's config. cfg := &OAuthFlowConfig{ - Scopes: []string{"openid", "profile"}, - CallbackPort: 8765, + Scopes: []string{"openid", "profile"}, + CallbackPort: 8765, + AllowPrivateIPs: true, // loopback test server; guard would otherwise refuse to dial it } results[i] = cfg go func(idx int) { @@ -336,8 +397,9 @@ func TestHandleDynamicRegistration_NonRootIssuerRFC8414PathInsertion(t *testing. t.Cleanup(server.Close) config := &OAuthFlowConfig{ - Scopes: []string{"openid", "profile"}, - CallbackPort: 8765, + Scopes: []string{"openid", "profile"}, + CallbackPort: 8765, + AllowPrivateIPs: true, // loopback test server; guard would otherwise refuse to dial it } err := handleDynamicRegistration(context.Background(), server.URL+"/oauth", config) @@ -475,20 +537,29 @@ func TestHandleDynamicRegistration_SynthesisesEndpointWhenMetadataOmitsIt(t *tes "synthesised endpoint must be contacted exactly once") } -// TestHandleDynamicRegistration_PopulatesEndpoints verifies the contract -// the rest of the CLI flow depends on: handleDynamicRegistration writes -// the resolved AuthorizeURL / TokenURL onto OAuthFlowConfig so -// createOAuthConfig can construct the OAuth2 flow without re-discovery. -func TestHandleDynamicRegistration_PopulatesEndpoints(t *testing.T) { +// TestHandleDynamicRegistration_PopulatesEndpointsAndRenewalMetadata verifies +// the contract the rest of the CLI flow depends on: handleDynamicRegistration +// writes the resolved endpoints and DCR renewal metadata onto OAuthFlowConfig. +func TestHandleDynamicRegistration_PopulatesEndpointsAndRenewalMetadata(t *testing.T) { t.Parallel() + const ( + callbackPort = 8765 + secretExpiryUnix = int64(1_893_456_000) + registrationToken = "registration-access-token" + registrationClientURI = "https://issuer.example/register/cli-registered-client" + ) server := newDCRDiscoveryServer(t, dcrTestServerConfig{ codeChallengeMethodsSupported: []string{"S256"}, + clientSecretExpiresAt: secretExpiryUnix, + registrationAccessToken: registrationToken, + registrationClientURI: registrationClientURI, }) config := &OAuthFlowConfig{ - Scopes: []string{"openid", "profile"}, - CallbackPort: 8765, + Scopes: []string{"openid", "profile"}, + CallbackPort: callbackPort, + AllowPrivateIPs: true, // loopback test server; guard would otherwise refuse to dial it } err := handleDynamicRegistration(context.Background(), server.URL, config) @@ -499,4 +570,9 @@ func TestHandleDynamicRegistration_PopulatesEndpoints(t *testing.T) { "handleDynamicRegistration must populate OAuthFlowConfig.AuthorizeURL") assert.Equal(t, server.URL+"/token", config.TokenURL, "handleDynamicRegistration must populate OAuthFlowConfig.TokenURL") + assert.Equal(t, time.Unix(secretExpiryUnix, 0).UTC(), config.SecretExpiry) + assert.Equal(t, registrationToken, config.RegistrationAccessToken) + assert.Equal(t, registrationClientURI, config.RegistrationClientURI) + assert.Equal(t, "none", config.TokenEndpointAuthMethod) + assert.Equal(t, callbackPort, config.RegisteredCallbackPort) } diff --git a/pkg/auth/discovery/discovery.go b/pkg/auth/discovery/discovery.go index d487cc0ce7..1fe56af665 100644 --- a/pkg/auth/discovery/discovery.go +++ b/pkg/auth/discovery/discovery.go @@ -519,6 +519,24 @@ type OAuthFlowConfig struct { Resource string // RFC 8707 resource indicator (optional) OAuthParams map[string]string ScopeParamName string // Override scope query parameter name (e.g., "user_scope" for Slack) + + // AllowPrivateIPs permits the Dynamic Client Registration calls (discovery + // fetch and registration POST) to reach private/loopback/link-local + // addresses. Callers set it from networking.TargetIsPrivate on the remote + // target: when ToolHive is pointed at a private remote the upstream IdP may + // legitimately be private too, so its DCR calls must be allowed; otherwise + // they are guarded to contain SSRF (CWE-918). This mirrors the + // blockPrivateIPs decision the flow already applies to its other discovery + // fetches. Defaults to false (guarded). + AllowPrivateIPs bool + + // DCR renewal metadata — populated by handleDynamicRegistration and threaded + // into OAuthFlowResult so callers can persist the data for RFC 7592 operations. + SecretExpiry time.Time // zero means the secret never expires + RegistrationAccessToken string //nolint:gosec // G117: field legitimately holds sensitive data + RegistrationClientURI string + TokenEndpointAuthMethod string + RegisteredCallbackPort int } // OAuthFlowResult contains the result of an OAuth flow @@ -534,6 +552,16 @@ type OAuthFlowResult struct { // DCR client credentials for persistence (obtained during Dynamic Client Registration) ClientID string ClientSecret string //nolint:gosec // G117: field legitimately holds sensitive data + + // DCR renewal metadata (RFC 7591 §3.2.1 / RFC 7592). + // SecretExpiry is zero when the provider did not issue an expiring secret. + // RegistrationAccessToken and RegistrationClientURI are empty when the + // provider does not support RFC 7592 management operations. + SecretExpiry time.Time + RegistrationAccessToken string //nolint:gosec // G117: field legitimately holds sensitive data + RegistrationClientURI string + TokenEndpointAuthMethod string + RegisteredCallbackPort int } func shouldDynamicallyRegisterClient(config *OAuthFlowConfig) bool { @@ -606,20 +634,12 @@ func PerformOAuthFlow(ctx context.Context, issuer string, config *OAuthFlowConfi // exists. // // One consequence of option (b) is that the resolver's RFC 7591 §3.2.1 -// expiry-driven refetch does NOT participate in the CLI's cross- -// invocation persistence loop: each PerformOAuthFlow call builds a fresh -// in-memory store, so a "cached but expired" entry from the previous -// invocation never reaches the resolver. Cross-invocation expiry is also -// NOT enforced by the remote handler's gate today — -// HasCachedClientCredentials only checks CachedClientID != "" and does -// not consult CachedSecretExpiry, so an expired-but-still-cached client -// gets reused on the next invocation and surfaces as a token-endpoint -// failure rather than a clean DCR re-registration. Tightening the gate -// to also check CachedSecretExpiry is open follow-up work; the -// behaviour today is "cross-invocation expiry is unhandled". Within a -// single invocation, the resolver's expiry check is still in the loop -// and would fire if the same call site somehow registered, persisted, -// and re-queried the in-memory store — but the CLI never does this today. +// expiry-driven refetch does NOT participate in the CLI's cross-invocation +// persistence loop: each PerformOAuthFlow call builds a fresh in-memory store, +// so a cached entry from a previous invocation never reaches the resolver. +// Cross-invocation client-secret expiry is handled instead by the remote +// handler, which consults CachedSecretExpiry and renews through RFC 7592 before +// cached credentials are used. // // Wrapping the remote handler's secretProvider into a dcr.CredentialStore // adapter (option (a)) would close that loop and is the natural follow-up; @@ -667,6 +687,18 @@ func handleDynamicRegistration(ctx context.Context, issuer string, config *OAuth config.TokenURL = resolution.TokenEndpoint } + // Store DCR renewal metadata for RFC 7592 operations. + // A zero ClientSecretExpiresAt means the secret never expires (RFC 7591 §3.2.1). + config.SecretExpiry = resolution.ClientSecretExpiresAt + config.RegistrationAccessToken = resolution.RegistrationAccessToken + config.RegistrationClientURI = resolution.RegistrationClientURI + config.TokenEndpointAuthMethod = resolution.TokenEndpointAuthMethod + config.RegisteredCallbackPort = config.CallbackPort + + if resolution.RegistrationAccessToken != "" { + slog.Debug("DCR response includes registration access token for RFC 7592 operations") + } + return nil } @@ -707,12 +739,14 @@ func resolveDCRCredentials( // dcr.Request.Issuer carries the caller's logical scope for cache // keying. The CLI flow has no separate logical issuer of its own, so // it deliberately reuses the upstream's issuer URL here. This is - // safe because the resolver's cache key is (Issuer, RedirectURI, - // ScopesHash) and the CLI always supplies an explicit loopback - // RedirectURI per RFC 8252 §7.3 — even if a future embedded-authserver - // upstream happened to share Issuer with a CLI invocation, the - // distinct RedirectURI keeps the cache keys apart. See the Issuer - // field doc on dcr.Request for the wider semantics. + // safe because the resolver's cache key is (Issuer, UpstreamID, + // RedirectURI, ScopesHash) and the CLI always supplies an explicit + // loopback RedirectURI per RFC 8252 §7.3 — even if a future + // embedded-authserver upstream happened to share Issuer with a CLI + // invocation, the distinct RedirectURI keeps the cache keys apart. + // UpstreamID (derived by the resolver from the RegistrationEndpoint set + // below, or a DiscoveryURL) further distinguishes distinct upstreams. + // See the Issuer field doc on dcr.Request for the wider semantics. req := &dcr.Request{ Issuer: issuer, RedirectURI: redirectURI, @@ -720,6 +754,11 @@ func resolveDCRCredentials( AuthorizationEndpoint: discoveredDoc.AuthorizationEndpoint, TokenEndpoint: discoveredDoc.TokenEndpoint, PublicClient: true, + // Carry the flow's private-IP decision so DCR shares the same SSRF + // posture as the flow's other discovery fetches; without this the DCR + // calls would always be guarded and a legitimately private upstream + // would be refused. See OAuthFlowConfig.AllowPrivateIPs. + AllowPrivateIPs: config.AllowPrivateIPs, } // Fetch AS metadata using the multi-URL fallback so non-root issuers @@ -728,7 +767,26 @@ func resolveDCRCredentials( // defence-in-depth for a future refactor that produces an empty-issuer // doc and preserves the pre-existing RegistrationEndpoint-direct path. if discoveredDoc.Issuer != "" { - fullMeta, metaErr := oauthproto.FetchAuthorizationServerMetadata(ctx, discoveredDoc.Issuer, nil) + // Guard this fetch the same way pkg/auth/dcr's own outbound calls are + // guarded (CWE-918): discoveredDoc.Issuer can be untrusted server input + // on some discovery branches (see discoverIssuerAndScopes in + // pkg/auth/remote/handler.go), so a nil client here would reopen the + // discovery-indirection and DNS-rebinding vectors this PR closes + // elsewhere. See networking.NewHostScopedClientBuilder for the guard + // policy and pkg/auth/dcr's newGuardedDCRClient for the same pattern. + metaHost, parseErr := url.Parse(discoveredDoc.Issuer) + if parseErr != nil { + return nil, fmt.Errorf("dynamic client registration failed: parse issuer for http client: %w", parseErr) + } + metaClient, clientErr := networking.NewHostScopedClientBuilder(metaHost.Host, config.AllowPrivateIPs, false). + WithDisableKeepAlives(true). + Build() + if clientErr != nil { + return nil, fmt.Errorf("dynamic client registration failed: build metadata http client: %w", clientErr) + } + metaClient.CheckRedirect = networking.SameHostRedirectPolicy() + + fullMeta, metaErr := oauthproto.FetchAuthorizationServerMetadata(ctx, discoveredDoc.Issuer, metaClient) if metaErr != nil && !errors.Is(metaErr, oauthproto.ErrRegistrationEndpointMissing) { return nil, fmt.Errorf("dynamic client registration failed: discover authorization server metadata: %w", metaErr) } @@ -810,8 +868,10 @@ func getDiscoveryDocument( }, nil } - // Fall back to discovering endpoints - return oauth.DiscoverOIDCEndpoints(ctx, issuer) + // Fall back to discovering endpoints. This fetch precedes the DCR + // registration this function's callers are about to perform, so it shares + // the same private-IP policy (CWE-918) rather than defaulting to unguarded. + return oauth.DiscoverOIDCEndpoints(ctx, issuer, !config.AllowPrivateIPs) } // createOAuthConfig creates the OAuth configuration based on available endpoints @@ -891,15 +951,30 @@ func newOAuthFlow(ctx context.Context, oauthConfig *oauth.Config, config *OAuthF } source := flow.TokenSource() + return buildOAuthFlowResult(source, oauthConfig, tokenResult, config), nil +} + +func buildOAuthFlowResult( + tokenSource oauth2.TokenSource, + oauthConfig *oauth.Config, + tokenResult *oauth.TokenResult, + config *OAuthFlowConfig, +) *OAuthFlowResult { return &OAuthFlowResult{ - TokenSource: source, + TokenSource: tokenSource, Config: oauthConfig, AccessToken: tokenResult.AccessToken, RefreshToken: tokenResult.RefreshToken, Expiry: tokenResult.Expiry, ClientID: oauthConfig.ClientID, ClientSecret: oauthConfig.ClientSecret, - }, nil + // DCR renewal metadata — populated only when dynamic registration was performed. + SecretExpiry: config.SecretExpiry, + RegistrationAccessToken: config.RegistrationAccessToken, + RegistrationClientURI: config.RegistrationClientURI, + TokenEndpointAuthMethod: config.TokenEndpointAuthMethod, + RegisteredCallbackPort: config.RegisteredCallbackPort, + } } // FetchResourceMetadata fetches RFC 9728 protected-resource metadata from a diff --git a/pkg/auth/discovery/discovery_test.go b/pkg/auth/discovery/discovery_test.go index 70f9c03305..0e00faa4fa 100644 --- a/pkg/auth/discovery/discovery_test.go +++ b/pkg/auth/discovery/discovery_test.go @@ -17,6 +17,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive/pkg/auth/oauth" "github.com/stacklok/toolhive/pkg/networking" ) @@ -1324,8 +1325,9 @@ func TestHandleDynamicRegistration_MissingRegistrationEndpoint(t *testing.T) { t.Cleanup(server.Close) config := &OAuthFlowConfig{ - Scopes: []string{"openid", "profile"}, - CallbackPort: 8765, + Scopes: []string{"openid", "profile"}, + CallbackPort: 8765, + AllowPrivateIPs: true, // loopback test server; guard would otherwise refuse to dial it } err := handleDynamicRegistration(context.Background(), server.URL, config) @@ -1339,3 +1341,24 @@ func TestHandleDynamicRegistration_MissingRegistrationEndpoint(t *testing.T) { assert.Contains(t, err.Error(), "--remote-auth-client-id") assert.Contains(t, err.Error(), "--remote-auth-client-secret") } + +func TestBuildOAuthFlowResult_CopiesDCRRenewalMetadata(t *testing.T) { + t.Parallel() + + secretExpiry := time.Date(2030, time.January, 2, 3, 4, 5, 0, time.UTC) + config := &OAuthFlowConfig{ + SecretExpiry: secretExpiry, + RegistrationAccessToken: "registration-access-token", + RegistrationClientURI: "https://issuer.example/register/client-id", + TokenEndpointAuthMethod: "client_secret_basic", + RegisteredCallbackPort: 49152, + } + + result := buildOAuthFlowResult(nil, &oauth.Config{}, &oauth.TokenResult{}, config) + + assert.Equal(t, config.SecretExpiry, result.SecretExpiry) + assert.Equal(t, config.RegistrationAccessToken, result.RegistrationAccessToken) + assert.Equal(t, config.RegistrationClientURI, result.RegistrationClientURI) + assert.Equal(t, config.TokenEndpointAuthMethod, result.TokenEndpointAuthMethod) + assert.Equal(t, config.RegisteredCallbackPort, result.RegisteredCallbackPort) +} diff --git a/pkg/auth/identity.go b/pkg/auth/identity.go index 66bb9bccbd..0f00f446e8 100644 --- a/pkg/auth/identity.go +++ b/pkg/auth/identity.go @@ -8,6 +8,7 @@ import ( "encoding/json" "fmt" + "github.com/stacklok/toolhive-core/audit" upstreamtoken "github.com/stacklok/toolhive/pkg/auth/upstreamtoken" ) @@ -62,6 +63,13 @@ type PrincipalInfo struct { // Claims contains additional claims from the auth token. // This preserves all JWT claims for authorization policies. Claims map[string]any `json:"claims,omitempty"` + + // DelegationChain is the parsed RFC 8693 "act" claim: the ordered list of + // parties that acted on behalf of this identity's subject (outermost/most + // recent first). Nil when the token carries no act claim. The "delegation" + // tag matches the audit-event wire key so webhook payloads and audit logs + // share one spelling of the same structure. + DelegationChain *audit.DelegationChain `json:"delegation,omitempty"` } // Identity represents an authenticated user or service account. @@ -121,7 +129,7 @@ type Identity struct { // is active and the JWT contains a token session ID (tsid claim). // Each value is the rotated ID token when a refresh produced one // (OIDC Core 1.0 §12.2), otherwise the original JWT captured at the initial - // OIDC login; it is not independently validated for freshness. + // OIDC login. // // State semantics: // - nil — no tsid claim was present on the incoming JWT @@ -129,8 +137,29 @@ type Identity struct { // - empty map — tsid claim was valid but no providers had an // ID token stored for the session. // - populated map — keys are upstream provider names; values are the - // stored ID token JWTs (may be expired; callers MUST - // validate the `exp` claim before use). + // stored ID token JWTs, which may be expired. + // + // Whether `exp` matters depends on what the caller does with the token. Nothing + // in this package enforces it; the two current consumers use the token + // differently and neither validates it locally: + // + // - PRESENTING it as a credential — e.g. as the subject_token of an RFC 8693 + // exchange (pkg/vmcp/auth/strategies) — is subject to expiry, because the + // party receiving it will reject an expired assertion. That consumer + // deliberately does not pre-check `exp` either: it lets the IdP reject the + // exchange and surfaces the resulting `invalid_grant`. So such a caller must + // either check `exp` first or be prepared to handle that failure. + // - READING identity claims out of it — as the Cedar authorizer does for the + // profile claims an upstream access token omits — is not subject to expiry at + // all. The token is evidence that the upstream asserted something at login, + // and expiry says nothing about whether that assertion happened. Rejecting it + // once expired would flip an authorization decision mid-session, since + // sessions outlive ID tokens by days; and the claims are no staler than the + // alternative, since the profile claims mirrored into the ToolHive-issued + // token are likewise captured once at login and never refreshed. + // + // Either way the token is NOT re-validated for signature or freshness here, so + // treat its claims as login-time facts rather than current state. // // Redacted in MarshalJSON() to prevent token leakage. // MUST NOT be mutated after the Identity is placed in the request context. @@ -156,17 +185,18 @@ func (i *Identity) MarshalJSON() ([]byte, error) { // Create a safe representation with lowercase field names and redacted token type SafeIdentity struct { - Subject string `json:"subject"` - PlatformUserID string `json:"platformUserId,omitempty"` - Name string `json:"name"` - Email string `json:"email"` - Groups []string `json:"groups"` - Claims map[string]any `json:"claims"` - Token string `json:"token"` - TokenType string `json:"tokenType"` - Metadata map[string]string `json:"metadata"` - UpstreamTokens map[string]string `json:"upstreamTokens,omitempty"` - UpstreamIDTokens map[string]string `json:"upstreamIDTokens,omitempty"` + Subject string `json:"subject"` + PlatformUserID string `json:"platformUserId,omitempty"` + Name string `json:"name"` + Email string `json:"email"` + Groups []string `json:"groups"` + Claims map[string]any `json:"claims"` + DelegationChain *audit.DelegationChain `json:"delegation,omitempty"` + Token string `json:"token"` + TokenType string `json:"tokenType"` + Metadata map[string]string `json:"metadata"` + UpstreamTokens map[string]string `json:"upstreamTokens,omitempty"` + UpstreamIDTokens map[string]string `json:"upstreamIDTokens,omitempty"` } const redacted = "REDACTED" @@ -221,6 +251,7 @@ func (i *Identity) MarshalJSON() ([]byte, error) { Email: i.Email, Groups: i.Groups, Claims: claims, + DelegationChain: i.DelegationChain, Token: token, TokenType: i.TokenType, Metadata: i.Metadata, diff --git a/pkg/auth/identity_test.go b/pkg/auth/identity_test.go index 07480c08fc..dd64b82c7d 100644 --- a/pkg/auth/identity_test.go +++ b/pkg/auth/identity_test.go @@ -10,6 +10,8 @@ import ( "github.com/golang-jwt/jwt/v5" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/audit" ) func TestClaimsToIdentity(t *testing.T) { @@ -511,6 +513,61 @@ func TestIdentity_MarshalJSON(t *testing.T) { assert.Equal(t, "org456", claims["org_id"], "org_id claim must be preserved") }, }, + { + name: "includes_delegation_chain", + identity: &Identity{ + PrincipalInfo: PrincipalInfo{ + Subject: "user123", + DelegationChain: audit.ParseDelegationChain(map[string]any{ + "sub": "agent-1", + "iss": "https://issuer.example", + "act": map[string]any{"sub": "agent-2"}, + }, DefaultMaxDelegationDepth), + }, + }, + wantErr: false, + checkFunc: func(t *testing.T, data []byte) { + t.Helper() + + var result map[string]any + require.NoError(t, json.Unmarshal(data, &result)) + + chain, ok := result["delegation"].(map[string]any) + require.True(t, ok, "delegation should be a map") + assert.Equal(t, false, chain["truncated"]) + assert.Equal(t, float64(0), chain["omitted"]) + assert.Equal(t, false, chain["malformed"]) + + hops, ok := chain["chain"].([]any) + require.True(t, ok, "chain should be an array") + require.Len(t, hops, 2) + + first, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-1", first["sub"]) + assert.Equal(t, "https://issuer.example", first["iss"]) + + second, ok := hops[1].(map[string]any) + require.True(t, ok) + assert.Equal(t, "agent-2", second["sub"]) + }, + }, + { + name: "omits_empty_delegation_chain", + identity: &Identity{ + PrincipalInfo: PrincipalInfo{Subject: "user123"}, + }, + wantErr: false, + checkFunc: func(t *testing.T, data []byte) { + t.Helper() + + var result map[string]any + require.NoError(t, json.Unmarshal(data, &result)) + + _, exists := result["delegation"] + assert.False(t, exists, "delegation key should be omitted when no chain is present") + }, + }, { name: "marshals_claims_without_tsid_unchanged", identity: &Identity{ diff --git a/pkg/auth/oauth/oidc.go b/pkg/auth/oauth/oidc.go index d4de7156e9..08fed930d4 100644 --- a/pkg/auth/oauth/oidc.go +++ b/pkg/auth/oauth/oidc.go @@ -20,9 +20,18 @@ import ( "github.com/stacklok/toolhive/pkg/oauthproto" ) -// DiscoverOIDCEndpoints discovers OAuth endpoints from an OIDC issuer -func DiscoverOIDCEndpoints(ctx context.Context, issuer string) (*oauthproto.OIDCDiscoveryDocument, error) { - return discoverOIDCEndpointsWithClient(ctx, issuer, nil, false) +// DiscoverOIDCEndpoints discovers OAuth endpoints from an OIDC issuer. +// +// issuer originates from untrusted remote-server discovery in some callers +// (e.g. the CLI DCR flow's well-known fallback) and from an operator-supplied +// flag in others; blockPrivateIPs lets each caller apply its own SSRF policy +// (CWE-918) rather than this function silently choosing one for everyone. +func DiscoverOIDCEndpoints( + ctx context.Context, + issuer string, + blockPrivateIPs bool, +) (*oauthproto.OIDCDiscoveryDocument, error) { + return discoverOIDCEndpointsWithClient(ctx, issuer, nil, false, blockPrivateIPs) } // DiscoverActualIssuer discovers the actual issuer from a URL that might be different from the issuer itself @@ -41,17 +50,18 @@ func DiscoverActualIssuer( return discoverOIDCEndpointsWithClientAndValidation(ctx, metadataURL, nil, false, false, blockPrivateIPs) } -// discoverOIDCEndpointsWithClient discovers OAuth endpoints from an OIDC issuer with a custom HTTP client (private for testing) +// discoverOIDCEndpointsWithClient discovers OAuth endpoints from an OIDC +// issuer with a custom HTTP client (private for testing). blockPrivateIPs is +// only consulted when client is nil — a caller-supplied client carries its +// own dial policy. func discoverOIDCEndpointsWithClient( ctx context.Context, issuer string, client networking.HTTPClient, insecureAllowHTTP bool, + blockPrivateIPs bool, ) (*oauthproto.OIDCDiscoveryDocument, error) { - // A caller-supplied client carries its own dial policy; the default-client - // private-IP guard only applies when client is nil, so blockPrivateIPs is - // irrelevant here. - return discoverOIDCEndpointsWithClientAndValidation(ctx, issuer, client, true, insecureAllowHTTP, false) + return discoverOIDCEndpointsWithClientAndValidation(ctx, issuer, client, true, insecureAllowHTTP, blockPrivateIPs) } // discoverOIDCEndpointsWithClientAndValidation discovers OAuth endpoints with optional issuer validation @@ -242,8 +252,10 @@ func createOAuthConfigFromOIDCWithClient( resource string, client networking.HTTPClient, ) (*Config, error) { - // Discover OIDC endpoints (insecureAllowHTTP is false for OAuth config creation) - doc, err := discoverOIDCEndpointsWithClient(ctx, issuer, client, false) + // Discover OIDC endpoints (insecureAllowHTTP is false for OAuth config creation). + // blockPrivateIPs=false here preserves this call path's existing behavior; + // it is unrelated to the CLI DCR fallback fixed in DiscoverOIDCEndpoints. + doc, err := discoverOIDCEndpointsWithClient(ctx, issuer, client, false, false) if err != nil { return nil, fmt.Errorf("failed to discover OIDC endpoints: %w", err) } diff --git a/pkg/auth/oauth/oidc_test.go b/pkg/auth/oauth/oidc_test.go index 4a523abbb7..b3a51be416 100644 --- a/pkg/auth/oauth/oidc_test.go +++ b/pkg/auth/oauth/oidc_test.go @@ -1123,7 +1123,7 @@ func TestDiscoverOIDCEndpoints_Production(t *testing.T) { }, } } - doc, err := discoverOIDCEndpointsWithClient(ctx, issuer, client, false) + doc, err := discoverOIDCEndpointsWithClient(ctx, issuer, client, false, false) if tt.expectError { require.Error(t, err) diff --git a/pkg/auth/remote/config.go b/pkg/auth/remote/config.go index 8675036dda..733e301024 100644 --- a/pkg/auth/remote/config.go +++ b/pkg/auth/remote/config.go @@ -11,7 +11,7 @@ import ( "strings" "time" - "github.com/stacklok/toolhive-core/registry/types" + registry "github.com/stacklok/toolhive-core/registry/types" httpval "github.com/stacklok/toolhive-core/validation/http" ) @@ -68,7 +68,8 @@ type Config struct { // ClientSecretExpiresAt indicates when the client secret expires (if provided by the DCR server). // A zero value means the secret does not expire. CachedSecretExpiry time.Time `json:"cached_secret_expiry,omitempty" yaml:"cached_secret_expiry,omitempty"` - // RegistrationAccessToken is used to update/delete the client registration. + // CachedRegTokenRef is a secret manager reference to the registration_access_token + // returned in the DCR response. Used for RFC 7592 client update operations. // Stored as a secret reference since it's sensitive. CachedRegTokenRef string `json:"cached_reg_token_ref,omitempty" yaml:"cached_reg_token_ref,omitempty"` @@ -78,6 +79,17 @@ type Config struct { // rotation clears CachedClientID without touching the stable CIMD URL. // Read by resolveClientCredentials to send the correct client_id on token refresh. CachedCIMDClientID string `json:"cached_cimd_client_id,omitempty" yaml:"cached_cimd_client_id,omitempty"` + // CachedRegClientURI is the registration_client_uri from the DCR response. + // This is the endpoint used for RFC 7592 client read/update/delete operations. + // Stored as plain text since it is not sensitive. + CachedRegClientURI string `json:"cached_reg_client_uri,omitempty" yaml:"cached_reg_client_uri,omitempty"` + // CachedTokenEndpointAuthMethod is the auth method used for the token endpoint + // (e.g., "client_secret_basic", "none"). Persisted for RFC 7592 updates. + CachedTokenEndpointAuthMethod string `json:"cached_token_auth_method,omitempty" yaml:"cached_token_auth_method,omitempty"` + // CachedDCRCallbackPort is the callback port that was actually registered + // during DCR. It may differ from CallbackPort when the requested port was + // unavailable and a fallback port was selected. + CachedDCRCallbackPort int `json:"cached_dcr_callback_port,omitempty" yaml:"cached_dcr_callback_port,omitempty"` } // BearerTokenEnvVarName is the environment variable name used for bearer token authentication. @@ -184,6 +196,9 @@ func (c *Config) ClearCachedClientCredentials() { c.CachedClientSecretRef = "" c.CachedSecretExpiry = time.Time{} c.CachedRegTokenRef = "" + c.CachedRegClientURI = "" + c.CachedTokenEndpointAuthMethod = "" + c.CachedDCRCallbackPort = 0 } // LogContext returns the upstream issuer and resolved client_id for use as diff --git a/pkg/auth/remote/handler.go b/pkg/auth/remote/handler.go index f92e7a04ec..b3b4334b9a 100644 --- a/pkg/auth/remote/handler.go +++ b/pkg/auth/remote/handler.go @@ -9,6 +9,7 @@ import ( "fmt" "log/slog" "strings" + "time" "golang.org/x/oauth2" @@ -25,6 +26,7 @@ type Handler struct { tokenPersister TokenPersister clientCredentialsPersister ClientCredentialsPersister secretProvider secrets.Provider + httpClient networking.HTTPClient } // NewHandler creates a new remote authentication handler @@ -51,6 +53,11 @@ func (h *Handler) SetClientCredentialsPersister(persister ClientCredentialsPersi h.clientCredentialsPersister = persister } +// SetHTTPClient sets the HTTP client used for RFC 7592 registration management requests. +func (h *Handler) SetHTTPClient(client networking.HTTPClient) { + h.httpClient = client +} + // Authenticate is the main entry point for remote MCP server authentication func (h *Handler) Authenticate(ctx context.Context, remoteURL string) (oauth2.TokenSource, error) { // Priority 1: Bearer token authentication (if configured) @@ -85,7 +92,7 @@ func (h *Handler) Authenticate(ctx context.Context, remoteURL string) (oauth2.To } // Discover OAuth endpoints once (used by both cached token restore and fresh OAuth) - issuer, scopes, authServerInfo, err := h.discoverIssuerAndScopes(ctx, authInfo, remoteURL) + issuer, scopes, authServerInfo, allowPrivateIPs, err := h.discoverIssuerAndScopes(ctx, authInfo, remoteURL) if err != nil { return nil, err } @@ -103,8 +110,13 @@ func (h *Handler) Authenticate(ctx context.Context, remoteURL string) (oauth2.To } } - // Priority 3: Fresh OAuth authentication flow - return h.performOAuthFlow(ctx, issuer, scopes, authServerInfo) + // Priority 3: Fresh OAuth authentication flow. Reuse the same target-private + // decision the discovery fetches above already computed, so DCR (discovery + + // registration) is allowed to reach a private upstream only when the + // operator-configured remote is itself private, and guarded otherwise + // (CWE-918) — without a second, later TargetIsPrivate DNS lookup that could + // disagree with the first. + return h.performOAuthFlow(ctx, issuer, scopes, authServerInfo, allowPrivateIPs) } // validateBearerRequirement checks if Bearer auth is required without OAuth fallback @@ -132,6 +144,7 @@ func (h *Handler) performOAuthFlow( issuer string, scopes []string, authServerInfo *discovery.AuthServerInfo, + allowPrivateIPs bool, ) (oauth2.TokenSource, error) { slog.Debug("Starting OAuth authentication flow", "issuer", issuer) @@ -139,7 +152,7 @@ func (h *Handler) performOAuthFlow( // Priority 1: Pre-configured credentials — set by buildOAuthFlowConfig from h.config.ClientID/ClientSecret. // Priority 2: CIMD — AS advertises support and no credentials are set; use metadata URL as client_id. // Priority 3: DCR — PerformOAuthFlow handles this when ClientID is still empty after the above. - flowConfig := h.buildOAuthFlowConfig(scopes, authServerInfo) + flowConfig := h.buildOAuthFlowConfig(scopes, authServerInfo, allowPrivateIPs) if shouldUseCIMD(authServerInfo, flowConfig) { flowConfig.ClientID = oauthproto.ToolHiveClientMetadataDocumentURL slog.Debug("Using CIMD client_id", "url", oauthproto.ToolHiveClientMetadataDocumentURL) @@ -162,19 +175,24 @@ func (h *Handler) performOAuthFlow( } // buildOAuthFlowConfig creates the OAuth flow configuration -func (h *Handler) buildOAuthFlowConfig(scopes []string, authServerInfo *discovery.AuthServerInfo) *discovery.OAuthFlowConfig { +func (h *Handler) buildOAuthFlowConfig( + scopes []string, + authServerInfo *discovery.AuthServerInfo, + allowPrivateIPs bool, +) *discovery.OAuthFlowConfig { flowConfig := &discovery.OAuthFlowConfig{ - ClientID: h.config.ClientID, - ClientSecret: h.config.ClientSecret, - AuthorizeURL: h.config.AuthorizeURL, - TokenURL: h.config.TokenURL, - Scopes: scopes, - CallbackPort: h.config.CallbackPort, - Timeout: h.config.Timeout, - SkipBrowser: h.config.SkipBrowser, - Resource: h.config.Resource, - OAuthParams: h.config.OAuthParams, - ScopeParamName: h.config.ScopeParamName, + ClientID: h.config.ClientID, + ClientSecret: h.config.ClientSecret, + AuthorizeURL: h.config.AuthorizeURL, + TokenURL: h.config.TokenURL, + Scopes: scopes, + CallbackPort: h.config.CallbackPort, + Timeout: h.config.Timeout, + SkipBrowser: h.config.SkipBrowser, + Resource: h.config.Resource, + OAuthParams: h.config.OAuthParams, + ScopeParamName: h.config.ScopeParamName, + AllowPrivateIPs: allowPrivateIPs, } // If we have discovered endpoints from the authorization server metadata, @@ -208,7 +226,15 @@ func (h *Handler) wrapWithPersistence(result *discovery.OAuthFlowResult) oauth2. // CIMD client IDs (HTTPS URLs) are stable constants and are stored separately below. if h.clientCredentialsPersister != nil && result.ClientID != "" && !oauthproto.IsClientIDMetadataDocumentURL(result.ClientID) { - if err := h.clientCredentialsPersister(result.ClientID, result.ClientSecret); err != nil { + if err := h.clientCredentialsPersister( + result.ClientID, + result.ClientSecret, + result.SecretExpiry, + result.RegistrationAccessToken, + result.RegistrationClientURI, + result.TokenEndpointAuthMethod, + result.RegisteredCallbackPort, + ); err != nil { slog.Warn("Failed to persist DCR client credentials", "error", err) } else { slog.Debug("Successfully persisted DCR client credentials for future restarts") @@ -279,6 +305,27 @@ func (h *Handler) tryRestoreFromCachedTokens( return nil, fmt.Errorf("secret provider not configured, cannot restore cached tokens") } + // Check if the cached client secret is expired before attempting token refresh. + // If it has fully expired and renewal also fails we must force a fresh OAuth flow. + if h.isSecretExpiredOrExpiringSoon() { + slog.Debug("Cached client secret is expiring or expired; attempting renewal before token restore", + "expiry", h.config.CachedSecretExpiry) + if renewErr := h.renewClientSecret(ctx, issuer); renewErr != nil { + slog.Warn("Client secret renewal failed", "error", renewErr) + // Hard-fail only when the secret is already past its expiry. + // If we are still in the buffer window the existing secret may work. + if !h.config.CachedSecretExpiry.IsZero() && time.Now().After(h.config.CachedSecretExpiry) { + return nil, fmt.Errorf( + "client secret expired at %v and renewal failed: %w", + h.config.CachedSecretExpiry, renewErr) + } + // Still within buffer — log and continue with the existing (still-valid) secret + slog.Warn("Proceeding with expiring client secret after failed renewal attempt") + } else { + slog.Debug("Successfully renewed client secret before token restore") + } + } + refreshToken, err := h.secretProvider.GetSecret(ctx, h.config.CachedRefreshTokenRef) if err != nil { return nil, fmt.Errorf("failed to retrieve cached refresh token: %w", err) @@ -352,13 +399,18 @@ func (h *Handler) discoverIssuerAndScopes( ctx context.Context, authInfo *discovery.AuthInfo, remoteURL string, -) (string, []string, *discovery.AuthServerInfo, error) { +) (string, []string, *discovery.AuthServerInfo, bool, error) { // Decide once whether discovery fetches derived from untrusted server input // (realm, resource_metadata, authorization_servers) may reach private // addresses. If the operator-configured target is itself internal they may; // otherwise block them to contain SSRF (CWE-918). The configured-issuer path // below is exempt because that URL comes from the operator, not the server. - blockPrivateIPs := !networking.TargetIsPrivate(ctx, remoteURL) + // Callers also thread this decision into the DCR flow (see performOAuthFlow), + // so it's computed exactly once per Authenticate call rather than being + // re-derived from a second, later TargetIsPrivate DNS lookup that could + // disagree with this one. + allowPrivateIPs := networking.TargetIsPrivate(ctx, remoteURL) + blockPrivateIPs := !allowPrivateIPs // Priority 1: Use configured issuer if available. Fetch discovery to populate // AuthServerInfo (including ClientIDMetadataDocumentSupported) even when the @@ -366,7 +418,7 @@ func (h *Handler) discoverIssuerAndScopes( if h.config.Issuer != "" { slog.Debug("Using configured issuer", "issuer", h.config.Issuer) authServerInfo, _ := discovery.ValidateAndDiscoverAuthServer(ctx, h.config.Issuer, false) - return h.config.Issuer, h.config.Scopes, authServerInfo, nil + return h.config.Issuer, h.config.Scopes, authServerInfo, allowPrivateIPs, nil } // Priority 2: Try to derive from realm (RFC 8414). Fetch discovery for the @@ -376,7 +428,7 @@ func (h *Handler) discoverIssuerAndScopes( if derivedIssuer != "" { slog.Debug("Derived issuer from realm", "issuer", derivedIssuer) authServerInfo, _ := discovery.ValidateAndDiscoverAuthServer(ctx, derivedIssuer, blockPrivateIPs) - return derivedIssuer, h.config.Scopes, authServerInfo, nil + return derivedIssuer, h.config.Scopes, authServerInfo, allowPrivateIPs, nil } } @@ -384,7 +436,7 @@ func (h *Handler) discoverIssuerAndScopes( if authInfo.ResourceMetadata != "" { issuer, scopes, authServerInfo, err := h.tryDiscoverFromResourceMetadata(ctx, authInfo.ResourceMetadata, blockPrivateIPs) if err == nil { - return issuer, scopes, authServerInfo, nil + return issuer, scopes, authServerInfo, allowPrivateIPs, nil } slog.Debug("Resource metadata discovery failed, falling through to well-known discovery", "error", err) } @@ -393,7 +445,7 @@ func (h *Handler) discoverIssuerAndScopes( // This handles cases where the issuer differs from the server URL (e.g., Atlassian) issuer, scopes, authServerInfo, err := h.tryDiscoverFromWellKnown(ctx, remoteURL, blockPrivateIPs) if err == nil { - return issuer, scopes, authServerInfo, nil + return issuer, scopes, authServerInfo, allowPrivateIPs, nil } slog.Debug("Could not discover from well-known endpoint", "error", err) @@ -401,11 +453,11 @@ func (h *Handler) discoverIssuerAndScopes( derivedIssuer := discovery.DeriveIssuerFromURL(remoteURL) if derivedIssuer != "" { slog.Debug("Using derived issuer from URL", "issuer", derivedIssuer) - return derivedIssuer, h.config.Scopes, nil, nil + return derivedIssuer, h.config.Scopes, nil, allowPrivateIPs, nil } // No issuer could be determined - return "", nil, nil, fmt.Errorf("could not determine OAuth issuer. Please provide issuer in configuration, " + + return "", nil, nil, false, fmt.Errorf("could not determine OAuth issuer. Please provide issuer in configuration, " + "or ensure the server provides a valid realm parameter or resource_metadata URL in the WWW-Authenticate header") } diff --git a/pkg/auth/remote/handler_test.go b/pkg/auth/remote/handler_test.go index 2aa4000139..3a98f0cd14 100644 --- a/pkg/auth/remote/handler_test.go +++ b/pkg/auth/remote/handler_test.go @@ -274,7 +274,7 @@ func TestDiscoverIssuerAndScopes(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) defer cancel() - issuer, scopes, authServerInfo, err := handler.discoverIssuerAndScopes( + issuer, scopes, authServerInfo, _, err := handler.discoverIssuerAndScopes( ctx, authInfo, remoteURL, @@ -388,7 +388,7 @@ func TestDiscoverIssuerAndScopes_Security(t *testing.T) { } ctx := t.Context() - issuer, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") + issuer, _, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") require.NoError(t, err) // The path traversal should be normalized @@ -410,7 +410,7 @@ func TestDiscoverIssuerAndScopes_Security(t *testing.T) { defer mockServer.Close() ctx := t.Context() - issuer, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, mockServer.URL) + issuer, _, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, mockServer.URL) require.NoError(t, err) // Should not use the insecure realm, should fall through @@ -444,7 +444,7 @@ func TestDiscoverIssuerAndScopes_Security(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 1*time.Second) defer cancel() - issuer, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") + issuer, _, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") // Should not hang or crash; Priority 3 fails gracefully and falls through to URL-derived issuer require.NoError(t, err) @@ -532,7 +532,7 @@ func TestDiscoveryPriorityChain(t *testing.T) { } ctx := context.Background() - issuer, scopes, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") + issuer, scopes, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") require.NoError(t, err) assert.Equal(t, "https://configured.example.com", issuer) @@ -551,7 +551,7 @@ func TestDiscoveryPriorityChain(t *testing.T) { } ctx := context.Background() - issuer, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") + issuer, _, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") require.NoError(t, err) assert.Equal(t, "https://realm.example.com/oauth", issuer) @@ -569,7 +569,7 @@ func TestDiscoveryPriorityChain(t *testing.T) { } ctx := context.Background() - issuer, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") + issuer, _, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com") require.NoError(t, err) // Should fall through to URL-derived issuer @@ -587,7 +587,7 @@ func TestDiscoveryPriorityChain(t *testing.T) { } ctx := context.Background() - issuer, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com/path") + issuer, _, _, _, err := handler.discoverIssuerAndScopes(ctx, authInfo, "https://server.example.com/path") require.NoError(t, err) assert.Equal(t, "https://server.example.com", issuer) @@ -1011,3 +1011,21 @@ func TestResolveClientCredentials(t *testing.T) { }) } } + +// TestBuildOAuthFlowConfig_ThreadsAllowPrivateIPs pins that performOAuthFlow's +// allowPrivateIPs parameter reaches the built OAuthFlowConfig. Without this, +// a future refactor could silently drop the SSRF-guard decision computed in +// Authenticate/discoverIssuerAndScopes without any test failing. +func TestBuildOAuthFlowConfig_ThreadsAllowPrivateIPs(t *testing.T) { + t.Parallel() + + h := &Handler{config: &Config{}} + + flowConfig := h.buildOAuthFlowConfig([]string{"openid"}, nil, true) + assert.True(t, flowConfig.AllowPrivateIPs, + "allowPrivateIPs=true must reach OAuthFlowConfig.AllowPrivateIPs") + + flowConfig = h.buildOAuthFlowConfig([]string{"openid"}, nil, false) + assert.False(t, flowConfig.AllowPrivateIPs, + "allowPrivateIPs=false must reach OAuthFlowConfig.AllowPrivateIPs") +} diff --git a/pkg/auth/remote/persisting_token_source.go b/pkg/auth/remote/persisting_token_source.go index 3ba90b6253..11e92081e9 100644 --- a/pkg/auth/remote/persisting_token_source.go +++ b/pkg/auth/remote/persisting_token_source.go @@ -21,8 +21,26 @@ import ( type TokenPersister func(refreshToken string, expiry time.Time) error // ClientCredentialsPersister is called when DCR client credentials need to be persisted. -// This is used to store client_id and client_secret obtained during Dynamic Client Registration. -type ClientCredentialsPersister func(clientID, clientSecret string) error +// This is used to store client_id, client_secret, and renewal metadata obtained during +// Dynamic Client Registration (RFC 7591) and needed for secret renewal (RFC 7592). +// +// Parameters: +// - clientID: the registered client ID (public, stored as plain text) +// - clientSecret: the registered client secret (sensitive, stored via secret manager) +// - secretExpiry: when the client secret expires; zero value means it never expires +// - registrationAccessToken: bearer token for RFC 7592 management operations (sensitive) +// - registrationClientURI: endpoint for RFC 7592 client update/read operations (plain text) +// - tokenEndpointAuthMethod: the auth method used for the token endpoint (e.g., "client_secret_basic", "none") +// - registeredCallbackPort: the callback port used in the original DCR redirect URI +type ClientCredentialsPersister func( + clientID string, + clientSecret string, + secretExpiry time.Time, + registrationAccessToken string, + registrationClientURI string, + tokenEndpointAuthMethod string, + registeredCallbackPort int, +) error // PersistingTokenSource wraps an oauth2.TokenSource and persists tokens // whenever they are refreshed. This enables session restoration across diff --git a/pkg/auth/remote/secret_renewal.go b/pkg/auth/remote/secret_renewal.go new file mode 100644 index 0000000000..f950e8ad4f --- /dev/null +++ b/pkg/auth/remote/secret_renewal.go @@ -0,0 +1,330 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package remote + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "path" + "strings" + "time" + + "github.com/stacklok/toolhive/pkg/networking" + "github.com/stacklok/toolhive/pkg/oauthproto" +) + +// secretExpiryBuffer is the lead time before expiry at which we proactively +// renew the client secret (RFC 7592). Renewal is attempted when the secret +// expires within this window, not only after expiry. +const secretExpiryBuffer = 24 * time.Hour + +// clientUpdateRequest is the body sent in a RFC 7592 §2.2 PUT request. +// Per the spec, all client metadata fields that were provided during +// registration must be included in the update request body. +type clientUpdateRequest struct { + ClientID string `json:"client_id"` + ClientName string `json:"client_name,omitempty"` + RedirectURIs []string `json:"redirect_uris,omitempty"` + GrantTypes []string `json:"grant_types,omitempty"` + ResponseTypes []string `json:"response_types,omitempty"` + TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty"` +} + +// clientUpdateResponse is the body returned by a RFC 7592 §2.1 response. +// The provider may rotate the registration_access_token; if present we must +// replace the stored one. +type clientUpdateResponse struct { + // Required fields mirrored from registration response + ClientID string `json:"client_id"` + ClientSecret string `json:"client_secret,omitempty"` //nolint:gosec // G117: field holds sensitive data + + // Expiry fields + ClientSecretExpiresAt int64 `json:"client_secret_expires_at,omitempty"` + + // Management fields — registration_access_token may be rotated + RegistrationAccessToken string `json:"registration_access_token,omitempty"` //nolint:gosec + RegistrationClientURI string `json:"registration_client_uri,omitempty"` +} + +// isSecretExpiredOrExpiringSoon returns true when the cached client secret is +// already expired or will expire within secretExpiryBuffer. +// A zero CachedSecretExpiry means the secret never expires, so this returns false. +func (h *Handler) isSecretExpiredOrExpiringSoon() bool { + expiry := h.config.CachedSecretExpiry + if expiry.IsZero() { + return false // Non-expiring secret + } + return time.Now().After(expiry.Add(-secretExpiryBuffer)) +} + +// renewClientSecret attempts to renew the client secret using RFC 7592 §2.2. +// It retrieves the stored registration_access_token and sends a PUT request +// to the registration_client_uri with the current client metadata. +// +// On success the handler's config is updated with the new secret, expiry, and +// (if rotated) the new registration_access_token. +// +// Callers should log a warning and continue if renewal fails — the existing +// secret may still be valid for some time, or the provider may not support renewal. +func (h *Handler) renewClientSecret(ctx context.Context, issuer string) error { + if err := h.validateRenewalPrerequisites(); err != nil { + return err + } + + // Retrieve the registration access token from the secret manager + regAccessToken, err := h.secretProvider.GetSecret(ctx, h.config.CachedRegTokenRef) + if err != nil { + return fmt.Errorf("failed to retrieve registration access token: %w", err) + } + + slog.Debug("Attempting RFC 7592 client secret renewal", + "registration_client_uri", h.config.CachedRegClientURI) + + // Validate the registration_client_uri before using it + if err := validateRegistrationClientURI(h.config.CachedRegClientURI); err != nil { + return fmt.Errorf("invalid registration_client_uri: %w", err) + } + + // Build the update request body with the current client metadata. + // Per RFC 7592 §2.2, the request MUST include all client metadata fields + // that were provided during the initial registration. + // ToolHive's CLI DCR flow registers this fixed metadata set; if future DCR + // callers support custom metadata, persist and replay those original values. + updateReq := clientUpdateRequest{ + ClientID: h.config.CachedClientID, + ClientName: oauthproto.ToolHiveMCPClientName, + RedirectURIs: []string{fmt.Sprintf("http://localhost:%d/callback", h.registeredDCRCallbackPort())}, + GrantTypes: []string{"authorization_code", "refresh_token"}, + ResponseTypes: []string{"code"}, + TokenEndpointAuthMethod: h.config.CachedTokenEndpointAuthMethod, + } + + reqBody, err := json.Marshal(updateReq) + if err != nil { + return fmt.Errorf("failed to marshal client update request: %w", err) + } + + // Create PUT request per RFC 7592 §2.2 + req, err := http.NewRequestWithContext( + ctx, + http.MethodPut, + h.config.CachedRegClientURI, + bytes.NewReader(reqBody), + ) + if err != nil { + return fmt.Errorf("failed to create client update request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+regAccessToken) + + httpClient := h.httpClient + if httpClient == nil { + httpClient = newRenewalHTTPClient(ctx, issuer) + } + + resp, err := httpClient.Do(req) + if err != nil { + return fmt.Errorf("client update request failed: %w", err) + } + defer func() { + if closeErr := resp.Body.Close(); closeErr != nil { + slog.Debug("Failed to close renewal response body", "error", closeErr) + } + }() + + if resp.StatusCode != http.StatusOK { + errorBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + _, _ = io.Copy(io.Discard, resp.Body) + return fmt.Errorf("client update request returned HTTP %d: %s", resp.StatusCode, string(errorBody)) + } + + // Parse the renewal response + const maxResponseSize = 1024 * 1024 // 1 MB + var updateResp clientUpdateResponse + if err := json.NewDecoder(io.LimitReader(resp.Body, maxResponseSize)).Decode(&updateResp); err != nil { + return fmt.Errorf("failed to decode client update response: %w", err) + } + + if updateResp.ClientID == "" { + return fmt.Errorf("client update response missing client_id") + } + if updateResp.ClientSecret == "" { + return fmt.Errorf("client update response missing client_secret") + } + + return h.persistRenewedSecret(updateResp) +} + +func (h *Handler) validateRenewalPrerequisites() error { + if h.config.CachedRegClientURI == "" { + return fmt.Errorf("registration_client_uri missing; cannot renew secret (RFC 7592 unsupported)") + } + if h.config.CachedRegTokenRef == "" { + return fmt.Errorf("registration_access_token missing; cannot renew secret (RFC 7592 unsupported)") + } + if h.secretProvider == nil { + return fmt.Errorf("secret provider not configured; cannot retrieve registration access token") + } + return nil +} + +func (h *Handler) registeredDCRCallbackPort() int { + if h.config.CachedDCRCallbackPort != 0 { + return h.config.CachedDCRCallbackPort + } + return h.config.CallbackPort +} + +func (h *Handler) persistRenewedSecret(updateResp clientUpdateResponse) error { + if h.clientCredentialsPersister == nil { + return fmt.Errorf("client credentials persister not configured; cannot save renewed secret") + } + + var newExpiry time.Time + if updateResp.ClientSecretExpiresAt > 0 { + newExpiry = time.Unix(updateResp.ClientSecretExpiresAt, 0) + } + + // Use the rotated registration_access_token if provided; fall back to existing. + newRegToken := updateResp.RegistrationAccessToken + newRegURI := updateResp.RegistrationClientURI + if newRegURI != "" { + if err := validateRegistrationClientURIRotation(h.config.CachedRegClientURI, newRegURI); err != nil { + return fmt.Errorf("invalid rotated registration_client_uri: %w", err) + } + } + if newRegURI == "" { + newRegURI = h.config.CachedRegClientURI + } + + if err := h.clientCredentialsPersister( + updateResp.ClientID, + updateResp.ClientSecret, + newExpiry, + newRegToken, + newRegURI, + h.config.CachedTokenEndpointAuthMethod, + h.registeredDCRCallbackPort(), + ); err != nil { + return fmt.Errorf("failed to persist renewed client secret: %w", err) + } + + h.config.CachedClientID = updateResp.ClientID + h.config.CachedSecretExpiry = newExpiry + h.config.CachedRegClientURI = newRegURI + if h.config.CachedDCRCallbackPort == 0 { + h.config.CachedDCRCallbackPort = h.config.CallbackPort + } + + slog.Debug("Successfully renewed client secret via RFC 7592", + "client_id", updateResp.ClientID, + "new_expiry_zero", newExpiry.IsZero(), + "reg_token_rotated", newRegToken != "") + + return nil +} + +// validateRegistrationClientURI validates that the registration_client_uri is +// a valid HTTPS URL (or localhost for development). +func validateRegistrationClientURI(registrationClientURI string) error { + if registrationClientURI == "" { + return fmt.Errorf("registration_client_uri is empty") + } + + parsedURL, err := url.Parse(registrationClientURI) + if err != nil { + return fmt.Errorf("invalid registration_client_uri URL: %w", err) + } + if parsedURL.Scheme == "" || parsedURL.Host == "" { + return fmt.Errorf("registration_client_uri must include a scheme and host: %s", registrationClientURI) + } + + if parsedURL.Scheme != "https" && !networking.IsLocalhost(parsedURL.Host) { + return fmt.Errorf("registration_client_uri must use HTTPS: %s", registrationClientURI) + } + cleanPath := path.Clean(parsedURL.Path) + if parsedURL.Path == "" || cleanPath == "." || cleanPath == "/" { + return fmt.Errorf("registration_client_uri must include a non-root path: %s", registrationClientURI) + } + + return nil +} + +// validateRegistrationClientURIRotation prevents a renewal response from +// moving the bearer-authenticated management endpoint to another origin. +func validateRegistrationClientURIRotation(currentURI, rotatedURI string) error { + if err := validateRegistrationClientURI(rotatedURI); err != nil { + return err + } + + current, err := url.Parse(currentURI) + if err != nil { + return fmt.Errorf("invalid current registration_client_uri: %w", err) + } + rotated, err := url.Parse(rotatedURI) + if err != nil { + return fmt.Errorf("invalid rotated registration_client_uri: %w", err) + } + if !sameOrigin(current, rotated) { + return fmt.Errorf("rotated registration_client_uri must remain on origin %s://%s", current.Scheme, current.Host) + } + + return nil +} + +func sameOrigin(first, second *url.URL) bool { + return strings.EqualFold(first.Scheme, second.Scheme) && + strings.EqualFold(first.Hostname(), second.Hostname()) && + effectivePort(first) == effectivePort(second) +} + +func effectivePort(parsedURL *url.URL) string { + if port := parsedURL.Port(); port != "" { + return port + } + if strings.EqualFold(parsedURL.Scheme, "https") { + return "443" + } + if strings.EqualFold(parsedURL.Scheme, "http") { + return "80" + } + return "" +} + +// newRenewalHTTPClient derives renewal transport behavior from the shared +// OAuth client while adding protections required for a bearer-authenticated +// request to an endpoint supplied by an authorization server. +func newRenewalHTTPClient(ctx context.Context, issuer string) *http.Client { + baseClient := oauthproto.DefaultHTTPClient() + client := *baseClient + client.CheckRedirect = networking.SameHostRedirectPolicy() + + // Internal IdPs are an intentional deployment mode. Base the exception on + // the configured issuer rather than the server-returned management URI. + if networking.TargetIsPrivate(ctx, issuer) { + return &client + } + + if transport, ok := baseClient.Transport.(*http.Transport); ok { + protectedTransport := transport.Clone() + // A proxy resolves and connects to the destination on ToolHive's behalf, + // bypassing the destination-IP check below. Public management endpoints + // therefore use direct dials; trusted private issuers retain base-client + // proxy behavior through the branch above. + protectedTransport.Proxy = nil + protectedTransport.DialContext = networking.NewPrivateIPBlockingDialContext() + protectedTransport.DisableKeepAlives = true + client.Transport = protectedTransport + } + + return &client +} diff --git a/pkg/auth/remote/secret_renewal_test.go b/pkg/auth/remote/secret_renewal_test.go new file mode 100644 index 0000000000..f46cc4b504 --- /dev/null +++ b/pkg/auth/remote/secret_renewal_test.go @@ -0,0 +1,681 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package remote + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/networking" + "github.com/stacklok/toolhive/pkg/oauthproto" + secretmocks "github.com/stacklok/toolhive/pkg/secrets/mocks" +) + +func newTestSecretProvider(t *testing.T, values map[string]string) *secretmocks.MockProvider { + t.Helper() + + ctrl := gomock.NewController(t) + provider := secretmocks.NewMockProvider(ctrl) + provider.EXPECT().GetSecret(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, name string) (string, error) { + value, ok := values[name] + if !ok { + return "", fmt.Errorf("secret %q not found", name) + } + return value, nil + }, + ).AnyTimes() + return provider +} + +// TestIsSecretExpiredOrExpiringSoon tests the expiry helper on various time scenarios. +func TestIsSecretExpiredOrExpiringSoon(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + expiry time.Time + wantExpired bool + }{ + { + name: "zero expiry means never expires", + expiry: time.Time{}, + wantExpired: false, + }, + { + name: "expiry far in the future — not expiring", + expiry: time.Now().Add(48 * time.Hour), + wantExpired: false, + }, + { + name: "expiry within 24h buffer — expiring soon", + expiry: time.Now().Add(12 * time.Hour), + wantExpired: true, + }, + { + name: "expiry in the past — already expired", + expiry: time.Now().Add(-1 * time.Hour), + wantExpired: true, + }, + { + name: "expiry exactly at buffer boundary — expiring soon", + expiry: time.Now().Add(secretExpiryBuffer - time.Minute), + wantExpired: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + h := &Handler{ + config: &Config{ + CachedSecretExpiry: tt.expiry, + }, + } + assert.Equal(t, tt.wantExpired, h.isSecretExpiredOrExpiringSoon()) + }) + } +} + +// TestValidateRegistrationClientURI tests URI validation. +func TestValidateRegistrationClientURI(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + uri string + wantErr bool + }{ + { + name: "empty URI", + uri: "", + wantErr: true, + }, + { + name: "valid HTTPS URI", + uri: "https://example.com/oauth/register/client-id", + wantErr: false, + }, + { + name: "HTTP URI for non-localhost is rejected", + uri: "http://example.com/oauth/register/client-id", + wantErr: true, + }, + { + name: "localhost HTTP is allowed (development)", + uri: "http://localhost:8080/oauth/register/client-id", + wantErr: false, + }, + { + name: "127.0.0.1 HTTP is allowed (development)", + uri: "http://127.0.0.1:8080/oauth/register/client-id", + wantErr: false, + }, + { + name: "invalid URL", + uri: "://bad-url", + wantErr: true, + }, + { + name: "root path URI is rejected", + uri: "https://example.com/", + wantErr: true, + }, + { + name: "non-canonical root path URI is rejected", + uri: "https://example.com//", + wantErr: true, + }, + { + name: "missing host is rejected", + uri: "https:///oauth/register/client-id", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := validateRegistrationClientURI(tt.uri) + if tt.wantErr { + assert.Error(t, err) + } else { + assert.NoError(t, err) + } + }) + } +} + +// TestRenewClientSecret_MissingConfig tests early-exit conditions. +func TestRenewClientSecret_MissingConfig(t *testing.T) { + t.Parallel() + + t.Run("missing registration_client_uri", func(t *testing.T) { + t.Parallel() + + h := &Handler{ + config: &Config{ + CachedRegClientURI: "", + CachedRegTokenRef: "some-ref", + }, + } + err := h.renewClientSecret(context.Background(), "") + require.Error(t, err) + assert.Contains(t, err.Error(), "registration_client_uri missing") + }) + + t.Run("missing registration_token_ref", func(t *testing.T) { + t.Parallel() + + h := &Handler{ + config: &Config{ + CachedRegClientURI: "https://example.com/register/client-id", + CachedRegTokenRef: "", + }, + } + err := h.renewClientSecret(context.Background(), "") + require.Error(t, err) + assert.Contains(t, err.Error(), "registration_access_token missing") + }) + + t.Run("missing secret provider", func(t *testing.T) { + t.Parallel() + + h := &Handler{ + config: &Config{ + CachedRegClientURI: "https://example.com/register/client-id", + CachedRegTokenRef: "some-ref", + }, + secretProvider: nil, // no provider + } + err := h.renewClientSecret(context.Background(), "") + require.Error(t, err) + assert.Contains(t, err.Error(), "secret provider not configured") + }) +} + +// TestRenewClientSecret_Success tests the happy path with a mock RFC 7592 server. +func TestRenewClientSecret_Success(t *testing.T) { + t.Parallel() + + newSecret := "new-client-secret-xyz" + newExpiry := time.Now().Add(24 * time.Hour * 30).Unix() + newRegToken := "new-registration-access-token" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // RFC 7592 §2.2: must be PUT with Bearer auth + assert.Equal(t, http.MethodPut, r.Method) + assert.Contains(t, r.Header.Get("Authorization"), "Bearer reg-access-token") + assert.Equal(t, "application/json", r.Header.Get("Content-Type")) + + var updateReq clientUpdateRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&updateReq)) + assert.Equal(t, []string{"http://localhost:9876/callback"}, updateReq.RedirectURIs) + + // Return the updated registration response + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "client_id": "test-client-id", + "client_secret": newSecret, + "client_secret_expires_at": newExpiry, + "registration_access_token": newRegToken, + "registration_client_uri": "http://" + r.Host + r.URL.Path, + }) + })) + t.Cleanup(server.Close) + + // Set up persister capture + var persistedClientID, persistedSecret, persistedRegToken, persistedRegURI string + var persistedExpiry time.Time + + h := &Handler{ + config: &Config{ + CachedClientID: "test-client-id", + CachedRegClientURI: server.URL + "/register/test-client-id", + CachedRegTokenRef: "reg-token-secret-ref", + CallbackPort: 8666, + CachedDCRCallbackPort: 9876, + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "reg-token-secret-ref": "reg-access-token", + }), + clientCredentialsPersister: func( + clientID, secret string, + expiry time.Time, + regToken, regURI, _ string, + callbackPort int, + ) error { + persistedClientID = clientID + persistedSecret = secret + persistedExpiry = expiry + persistedRegToken = regToken + persistedRegURI = regURI + assert.Equal(t, 9876, callbackPort) + return nil + }, + } + + err := h.renewClientSecret(context.Background(), server.URL) + require.NoError(t, err) + + assert.Equal(t, "test-client-id", persistedClientID) + assert.Equal(t, newSecret, persistedSecret) + assert.Equal(t, newRegToken, persistedRegToken) + assert.False(t, persistedExpiry.IsZero(), "expiry should be set") + assert.NotEmpty(t, persistedRegURI) +} + +// TestRenewClientSecret_ServerError tests error propagation when the server returns non-200. +func TestRenewClientSecret_ServerError(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + })) + t.Cleanup(server.Close) + + h := &Handler{ + config: &Config{ + CachedClientID: "test-client-id", + CachedRegClientURI: server.URL + "/register/test-client-id", + CachedRegTokenRef: "reg-token-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "reg-token-ref": "bad-token", + }), + clientCredentialsPersister: func(_, _ string, _ time.Time, _, _, _ string, _ int) error { + return nil + }, + } + + err := h.renewClientSecret(context.Background(), server.URL) + require.Error(t, err) + assert.Contains(t, err.Error(), "401") +} + +// TestRenewClientSecret_NoPersister tests failure when persister is not set. +func TestRenewClientSecret_NoPersister(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "client_id": "test-client-id", + "client_secret": "new-secret", + }) + })) + t.Cleanup(server.Close) + + h := &Handler{ + config: &Config{ + CachedClientID: "test-client-id", + CachedRegClientURI: server.URL + "/register/test-client-id", + CachedRegTokenRef: "reg-token-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "reg-token-ref": "some-token", + }), + clientCredentialsPersister: nil, // no persister + } + + err := h.renewClientSecret(context.Background(), server.URL) + require.Error(t, err) + assert.Contains(t, err.Error(), "client credentials persister not configured") +} + +// TestRenewClientSecret_ZeroExpiryInResponse tests that a zero client_secret_expires_at +// is correctly interpreted as a non-expiring secret. +func TestRenewClientSecret_ZeroExpiryInResponse(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "client_id": "test-client-id", + "client_secret": "new-secret", + "client_secret_expires_at": 0, // never expires + }) + })) + t.Cleanup(server.Close) + + var capturedExpiry time.Time + + h := &Handler{ + config: &Config{ + CachedClientID: "test-client-id", + CachedRegClientURI: server.URL + "/register/test-client-id", + CachedRegTokenRef: "reg-token-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "reg-token-ref": "some-token", + }), + clientCredentialsPersister: func(_, _ string, expiry time.Time, _, _, _ string, _ int) error { + capturedExpiry = expiry + return nil + }, + } + + err := h.renewClientSecret(context.Background(), server.URL) + require.NoError(t, err) + assert.True(t, capturedExpiry.IsZero(), "zero client_secret_expires_at must produce zero time.Time") +} + +func TestRenewClientSecret_MalformedJSON(t *testing.T) { + t.Parallel() + + svc := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{invalid-json`)) + })) + t.Cleanup(svc.Close) + + h := &Handler{ + config: &Config{ + CachedRegClientURI: svc.URL + "/register/test-client-id", + CachedRegTokenRef: "rat-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{"rat-ref": "rat-token"}), + } + + err := h.renewClientSecret(context.Background(), svc.URL) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to decode client update response") +} + +func TestRenewClientSecret_MissingFields(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + response string + wantErr string + }{ + { + name: "missing_client_id", + response: `{"client_secret": "new-secret"}`, + wantErr: "client update response missing client_id", + }, + { + name: "missing_client_secret", + response: `{"client_id": "test-client-id"}`, + wantErr: "client update response missing client_secret", + }, + } + + for _, tt := range tests { + tt := tt // capture loop variable + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + svc := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(tt.response)) + })) + t.Cleanup(svc.Close) + + h := &Handler{ + config: &Config{ + CachedRegClientURI: svc.URL + "/register/test-client-id", + CachedRegTokenRef: "rat-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{"rat-ref": "rat-token"}), + } + + err := h.renewClientSecret(context.Background(), svc.URL) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestHandler_Restore_RenewSuccess(t *testing.T) { + t.Parallel() + + // Initial setup: secret expiring in 1 hour + expiry := time.Now().Add(1 * time.Hour) + svc := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"client_id": "test-client", "client_secret": "new-secret", "client_secret_expires_at": 4102444800}`)) + })) + t.Cleanup(svc.Close) + + var persistedID, persistedSecret string + renewalRequests := 0 + h := &Handler{ + config: &Config{ + CachedClientID: "test-client", + CachedSecretExpiry: expiry, + CachedRegClientURI: svc.URL + "/register/test-client", + CachedRegTokenRef: "rat-ref", + CachedRefreshTokenRef: "refresh-token-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "rat-ref": "rat-token", + "refresh-token-ref": "some-refresh-token", + }), + clientCredentialsPersister: func(id, secret string, _ time.Time, _, _, _ string, _ int) error { + persistedID = id + persistedSecret = secret + renewalRequests++ + return nil + }, + } + + // Calling tryRestoreFromCachedTokens should trigger renewal because of the 1h expiry. + // We expect an error because it will try to refresh the token and fail (no token endpoint). + _, err := h.tryRestoreFromCachedTokens(context.Background(), svc.URL, nil, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "cached tokens are invalid or expired") + + // But renewal DID happen + assert.Equal(t, "test-client", persistedID) + assert.Equal(t, "new-secret", persistedSecret) + assert.Equal(t, 1, renewalRequests) + assert.False(t, h.isSecretExpiredOrExpiringSoon()) +} + +func TestHandler_Restore_RenewFail_Soft(t *testing.T) { + t.Parallel() + + // Initial setup: secret expiring in 1 hour + expiry := time.Now().Add(1 * time.Hour) + var renewalPUTs atomic.Int32 + svc := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPut { + renewalPUTs.Add(1) + } + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(svc.Close) + + h := &Handler{ + config: &Config{ + CachedClientID: "test-client", + CachedSecretExpiry: expiry, + CachedRegClientURI: svc.URL + "/register/test-client", + CachedRegTokenRef: "rat-ref", + CachedRefreshTokenRef: "refresh-token-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "rat-ref": "rat-token", + "refresh-token-ref": "some-refresh-token", + }), + clientCredentialsPersister: func(_, _ string, _ time.Time, _, _, _ string, _ int) error { return nil }, + } + + // Renewal fails, but since it's only "expiring soon", each restore should + // continue to token refresh after making exactly one renewal request. + for attempt := int32(1); attempt <= 2; attempt++ { + _, err := h.tryRestoreFromCachedTokens(context.Background(), svc.URL, nil, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "cached tokens are invalid or expired") + assert.Equal(t, attempt, renewalPUTs.Load()) + } +} + +func TestHandler_Restore_RenewFail_Hard(t *testing.T) { + t.Parallel() + + // Initial setup: secret already expired + expiry := time.Now().Add(-1 * time.Hour) + svc := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(svc.Close) + + h := &Handler{ + config: &Config{ + CachedClientID: "test-client", + CachedSecretExpiry: expiry, + CachedRegClientURI: svc.URL + "/register/test-client", + CachedRegTokenRef: "rat-ref", + CachedRefreshTokenRef: "refresh-token-ref", + }, + secretProvider: newTestSecretProvider(t, map[string]string{ + "rat-ref": "rat-token", + "refresh-token-ref": "some-refresh-token", + }), + clientCredentialsPersister: func(string, string, time.Time, string, string, string, int) error { return nil }, + } + + // Renewal fails and it's fully expired -> fatal error + _, err := h.tryRestoreFromCachedTokens(context.Background(), svc.URL, nil, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "client secret expired at") + assert.Contains(t, err.Error(), "and renewal failed") +} + +func TestPersistRenewedSecret_RegistrationClientURIRotation(t *testing.T) { + t.Parallel() + + const currentURI = "https://issuer.example/register/test-client" + tests := []struct { + name string + rotatedURI string + wantErr string + }{ + { + name: "same origin is accepted", + rotatedURI: "https://issuer.example/manage/test-client", + }, + { + name: "equivalent default port is accepted", + rotatedURI: "https://issuer.example:443/manage/test-client", + }, + { + name: "cross origin is rejected", + rotatedURI: "https://attacker.example/manage/test-client", + wantErr: "must remain on origin", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + persistCalls := 0 + persistedURI := "" + h := &Handler{ + config: &Config{CachedRegClientURI: currentURI}, + clientCredentialsPersister: func(_, _ string, _ time.Time, _, regURI, _ string, _ int) error { + persistCalls++ + persistedURI = regURI + return nil + }, + } + + err := h.persistRenewedSecret(clientUpdateResponse{ + ClientID: "test-client", + ClientSecret: "rotated-secret", + RegistrationClientURI: tt.rotatedURI, + }) + if tt.wantErr != "" { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + assert.Zero(t, persistCalls) + assert.Equal(t, currentURI, h.config.CachedRegClientURI) + return + } + + require.NoError(t, err) + assert.Equal(t, 1, persistCalls) + assert.Equal(t, tt.rotatedURI, persistedURI) + assert.Equal(t, tt.rotatedURI, h.config.CachedRegClientURI) + }) + } +} + +func TestNewRenewalHTTPClient(t *testing.T) { + t.Parallel() + + baseClient := oauthproto.DefaultHTTPClient() + baseTransport, ok := baseClient.Transport.(*http.Transport) + require.True(t, ok) + require.NotNil(t, baseTransport.Proxy) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(server.Close) + + localhostClient := newRenewalHTTPClient( + context.Background(), + server.URL, + ) + assert.Equal(t, baseClient.Timeout, localhostClient.Timeout) + assert.Same(t, baseTransport, localhostClient.Transport) + require.NotNil(t, localhostClient.CheckRedirect) + + resp, err := localhostClient.Get(server.URL) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + assert.Equal(t, http.StatusNoContent, resp.StatusCode) + + originalRequest, err := http.NewRequest(http.MethodGet, "https://issuer.example/register", nil) + require.NoError(t, err) + sameHostRequest, err := http.NewRequest(http.MethodGet, "https://issuer.example/rotated", nil) + require.NoError(t, err) + crossHostRequest, err := http.NewRequest(http.MethodGet, "https://attacker.example/rotated", nil) + require.NoError(t, err) + assert.NoError(t, localhostClient.CheckRedirect(sameHostRequest, []*http.Request{originalRequest})) + assert.ErrorIs(t, + localhostClient.CheckRedirect(crossHostRequest, []*http.Request{originalRequest}), + networking.ErrRedirectRefused, + ) + + publicClient := newRenewalHTTPClient( + context.Background(), + "https://issuer.example", + ) + assert.Equal(t, baseClient.Timeout, publicClient.Timeout) + protectedTransport, ok := publicClient.Transport.(*http.Transport) + require.True(t, ok) + assert.NotSame(t, baseTransport, protectedTransport) + assert.Nil(t, protectedTransport.Proxy) + assert.NotNil(t, protectedTransport.DialContext) + assert.True(t, protectedTransport.DisableKeepAlives) + assert.Equal(t, baseTransport.TLSHandshakeTimeout, protectedTransport.TLSHandshakeTimeout) + assert.False(t, baseTransport.DisableKeepAlives, "the shared OAuth transport must not be mutated") + + _, err = publicClient.Get(server.URL) + require.Error(t, err) + assert.Contains(t, err.Error(), networking.ErrPrivateIpAddress) +} diff --git a/pkg/auth/token.go b/pkg/auth/token.go index 0360f35905..9201717b92 100644 --- a/pkg/auth/token.go +++ b/pkg/auth/token.go @@ -734,8 +734,30 @@ func (v *TokenValidator) ensureJWKSRegistered(ctx context.Context) error { registrationCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() - // Attempt registration - if err := v.jwksClient.Register(registrationCtx, v.jwksURL); err != nil { + // Attempt registration. The CA-aware client must be passed per-resource: + // jwx >= 3.1.0 injects its own default client at the resource level when + // none is given here, which takes precedence over the client-level one + // configured in NewTokenValidator and silently drops custom CA support. + err := v.jwksClient.Register(registrationCtx, v.jwksURL, jwk.WithHTTPClient(v.client)) + switch { + case err == nil: + // Registered and the first fetch succeeded. + case errors.Is(err, httprc.ErrNotReady()): + // The resource is registered and httprc will keep fetching it in the + // background; only the first fetch has not completed within the + // registration budget. Treat it as registered: Lookup returns a + // not-ready error until a background fetch succeeds. Retrying Register + // instead would fail with ErrResourceAlreadyExists forever. + //nolint:gosec // G706: JWKS URL is from server configuration or OIDC discovery + slog.Debug( + "JWKS URL registered but first fetch not ready; fetching continues in background", + "jwks_url", v.jwksURL, "error", err, + ) + case errors.Is(err, httprc.ErrResourceAlreadyExists()): + // The URL is already in httprc's resource map, e.g. a previous attempt + // returned ErrNotReady before this state was tracked, or OIDC + // re-discovery produced the same URL after resetting the flag. + default: v.jwksRegistrationErr = fmt.Errorf("failed to register JWKS URL: %w", err) // Do NOT set jwksRegistered = true -- allow retry on next call return v.jwksRegistrationErr diff --git a/pkg/auth/token_test.go b/pkg/auth/token_test.go index b900884a82..6481638963 100644 --- a/pkg/auth/token_test.go +++ b/pkg/auth/token_test.go @@ -3180,3 +3180,67 @@ func TestValidateToken_DiscoveryFailsWithKeyProvider(t *testing.T) { require.Contains(t, err.Error(), "local key provider could not resolve key") }) } + +func TestEnsureJWKSRegistered_NonFatalRegistrationErrors(t *testing.T) { + t.Parallel() + + t.Run("ErrNotReady marks the JWKS as registered", func(t *testing.T) { + t.Parallel() + // A JWKS endpoint that never succeeds keeps the resource from + // becoming ready within the registration budget. + jwksServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(jwksServer.Close) + caCertPath := writeTestServerCert(t, jwksServer) + + validator, err := NewTokenValidator(context.Background(), TokenValidatorConfig{ + Issuer: "test-issuer", + Audience: "test-audience", + JWKSURL: jwksServer.URL, + ClientID: "test-client", + CACertPath: caCertPath, + AllowPrivateIP: true, + }) + require.NoError(t, err) + + // Bound the ready-wait well below the 5s registration budget. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, validator.ensureJWKSRegistered(ctx), + "ErrNotReady must be treated as registered-but-pending") + require.True(t, validator.jwksRegistered) + + // Lookup surfaces not-ready until a background fetch succeeds. + _, err = validator.jwksClient.Lookup(context.Background(), jwksServer.URL) + require.Error(t, err) + }) + + t.Run("ErrResourceAlreadyExists marks the JWKS as registered", func(t *testing.T) { + t.Parallel() + jwksServer, caCertPath := createTestJWKSServer(t, jwk.NewSet()) + t.Cleanup(jwksServer.Close) + + validator, err := NewTokenValidator(context.Background(), TokenValidatorConfig{ + Issuer: "test-issuer", + Audience: "test-audience", + JWKSURL: jwksServer.URL, + ClientID: "test-client", + CACertPath: caCertPath, + AllowPrivateIP: true, + }) + require.NoError(t, err) + require.NoError(t, validator.ensureJWKSRegistered(context.Background())) + + // Simulate the reset done after OIDC re-discovery: the flag is + // cleared but the URL is still in httprc's resource map, so the + // next registration attempt returns ErrResourceAlreadyExists. + validator.jwksRegistrationMu.Lock() + validator.jwksRegistered = false + validator.jwksRegistrationMu.Unlock() + + require.NoError(t, validator.ensureJWKSRegistered(context.Background()), + "re-registering an already-registered URL must not fail") + require.True(t, validator.jwksRegistered) + }) +} diff --git a/pkg/auth/upstreamtoken/service.go b/pkg/auth/upstreamtoken/service.go index 503e15ab6d..3c948b590e 100644 --- a/pkg/auth/upstreamtoken/service.go +++ b/pkg/auth/upstreamtoken/service.go @@ -82,7 +82,9 @@ func (s *InProcessService) GetValidTokens(ctx context.Context, sessionID, provid // the rotated ID token when a refresh produced one (OIDC Core 1.0 §12.2), // otherwise the original JWT captured at the initial OIDC login; it is not // independently validated for freshness and may be empty if the upstream login -// never yielded one. Callers MUST check its exp claim before use. +// never yielded one. Callers that PRESENT it as a credential must expect expiry +// to be rejected; callers that only READ ITS CLAIMS are unaffected by expiry. See +// Identity.UpstreamIDTokens for both consumers. // // Returns an empty map and nil failed slice (not error) for unknown sessions. func (s *InProcessService) GetAllUpstreamCredentials( @@ -131,9 +133,15 @@ func (s *InProcessService) GetAllUpstreamCredentials( result[providerName] = *refreshed } - // TODO(auth): the "check exp" contract on UpstreamCredential.IDToken is - // documented but not enforced here. Enforcement belongs at the RFC 8693 - // token-exchange consumer when it receives the credential. + // The returned ID tokens are deliberately not checked for expiry here. Neither + // of today's two consumers wants that pre-check: one presents the token as an + // RFC 8693 subject_token and relies on the IdP to reject an expired assertion + // (pkg/vmcp/auth/strategies surfaces the resulting invalid_grant), and the other + // only reads its claims, for which expiry is irrelevant. Enforcing it here would + // break the second while duplicating what the IdP already does for the first. + // A future consumer that presents the token to a sink which does NOT validate + // `exp` would change that calculus — such a caller must check it itself. + // See Identity.UpstreamIDTokens for both. return result, failed, nil } diff --git a/pkg/auth/upstreamtoken/types.go b/pkg/auth/upstreamtoken/types.go index df26f3fc2d..90beaaac04 100644 --- a/pkg/auth/upstreamtoken/types.go +++ b/pkg/auth/upstreamtoken/types.go @@ -23,13 +23,21 @@ const TokenSessionIDClaimKey = "tsid" // the original JWT captured at the initial OIDC login (OIDC Core 1.0 §3.1.3.7). // It is not independently validated for freshness. // -// Callers MUST check its `exp` claim before using it (e.g. as the subject_token -// of an RFC 8693 token exchange), as it may be expired. Note also that the ID +// Callers that PRESENT it as a credential — e.g. as the subject_token of an +// RFC 8693 token exchange — are subject to expiry and must either check its `exp` +// claim first or handle the resulting rejection. Today's such consumer +// (pkg/vmcp/auth/strategies) deliberately takes the second option: it lets the IdP +// reject the exchange and surfaces the `invalid_grant`. Note also that the ID // token's `aud` is this auth server's client registration with the issuing // upstream, not the token-exchange endpoint — whether a target authorization // server accepts it as a subject token is governed by that server's policy and // is typically limited to the same issuer/audience. // +// Callers that only READ ITS CLAIMS are a different case and deliberately do not +// check `exp`: the token is then evidence of what the upstream asserted at login, +// and expiry says nothing about whether that assertion happened. See +// Identity.UpstreamIDTokens for the two consumers and why they differ. +// // IDToken may be empty when the upstream login did not return an id_token // (e.g. the provider was not asked for the openid scope). An empty IDToken // is a legitimate state, not an error. Note that Identity.UpstreamTokens and @@ -57,9 +65,10 @@ type TokenReader interface { // // Each returned IDToken is the rotated ID token when a refresh produced // one (OIDC Core 1.0 §12.2), otherwise the original JWT captured at OIDC - // login (OIDC Core 1.0 §3.1.3.7). Callers MUST check each ID token's `exp` - // claim before using it for e.g. RFC 8693 subject-token exchange, as it may - // be expired. + // login (OIDC Core 1.0 §3.1.3.7), and it may be expired. Callers that PRESENT + // it as a credential (e.g. RFC 8693 subject-token exchange) must check `exp` + // first or handle the resulting rejection; callers that only READ ITS CLAIMS + // are unaffected by expiry. See UpstreamCredential for both consumers. // // Returns an empty map and nil failed slice (not error) for unknown sessions. GetAllUpstreamCredentials(ctx context.Context, sessionID string) ( @@ -78,8 +87,10 @@ type Service interface { // - ErrNoRefreshToken if the access token is expired and no refresh token is available // - ErrRefreshFailed if the refresh attempt fails (e.g., revoked refresh token) // - // The returned UpstreamCredential.IDToken may be empty or expired; callers - // MUST check its `exp` claim before using it (e.g. as the subject_token of - // an RFC 8693 token exchange). See UpstreamCredential for details. + // The returned UpstreamCredential.IDToken may be empty or expired. Callers that + // PRESENT it as a credential (e.g. as the subject_token of an RFC 8693 token + // exchange) must check its `exp` claim first or handle the resulting rejection; + // callers that only READ ITS CLAIMS are unaffected by expiry. See + // UpstreamCredential for details. GetValidTokens(ctx context.Context, sessionID, providerName string) (*UpstreamCredential, error) } diff --git a/pkg/authserver/config.go b/pkg/authserver/config.go index 92ff27a8e2..2ca7e54249 100644 --- a/pkg/authserver/config.go +++ b/pkg/authserver/config.go @@ -10,13 +10,16 @@ import ( "log/slog" "net/url" "regexp" + "slices" "strings" "time" + oauthserver "github.com/stacklok/toolhive/pkg/authserver/server" servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto" "github.com/stacklok/toolhive/pkg/authserver/server/handlers" "github.com/stacklok/toolhive/pkg/authserver/server/keys" "github.com/stacklok/toolhive/pkg/authserver/server/registration" + "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange" "github.com/stacklok/toolhive/pkg/authserver/storage" "github.com/stacklok/toolhive/pkg/authserver/upstream" "github.com/stacklok/toolhive/pkg/networking" @@ -63,6 +66,11 @@ type RunConfig struct { // If nil, defaults are applied (access: 1h, refresh: 7d, authCode: 10m). TokenLifespans *TokenLifespanRunConfig `json:"token_lifespans,omitempty" yaml:"token_lifespans,omitempty"` + // DelegationTokenLifespan is the maximum lifetime for delegated tokens issued + // via RFC 8693 token exchange. Specified as a Go duration string (e.g., "15m"). + // If empty, defaults to 15 minutes. + DelegationTokenLifespan string `json:"delegation_token_lifespan,omitempty" yaml:"delegation_token_lifespan,omitempty"` + // Upstreams configures connections to upstream Identity Providers. // At least one upstream is required - the server delegates authentication to these providers. // Multiple upstreams are supported for sequential authorization chains. @@ -111,6 +119,25 @@ type RunConfig struct { // Production deployments reachable outside the cluster MUST use https://. //nolint:lll // field tags require full JSON+YAML names InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"` + + // TrustedIssuers lists external OIDC issuers whose tokens are accepted as + // subject tokens during RFC 8693 token exchange. Empty (the default) means + // only self-issued subject tokens are accepted. + // + // Prerequisite: the token-exchange grant requires a confidential client, + // and no supported deployment path provisions one today (DCR and CIMD + // clients are both public-only, and there is no client-seeding field on + // this RunConfig), so this grant is not yet usable end to end for + // self-issued or external subject tokens alike. Tracked in + // https://github.com/stacklok/toolhive/issues/6082. + // + // See tokenexchange.TrustedIssuer for the per-issuer field reference, and + // docs/arch/17-token-exchange-delegation.md for the trust model, consent + // signals, and operator-facing constraints (audience/scope bounding, + // subject namespace qualification, required client binding) that aren't + // visible from the config shape alone. + //nolint:lll // field tags require full JSON+YAML names + TrustedIssuers []tokenexchange.TrustedIssuer `json:"trusted_issuers,omitempty" yaml:"trusted_issuers,omitempty"` } // Validate checks that the on-disk RunConfig is internally consistent. Called @@ -123,6 +150,16 @@ func (c *RunConfig) Validate() error { return fmt.Errorf("cimd: %w", err) } } + // Also checked by Config.Validate() (same reasoning as + // validateBaselineClientScopes below): buildUpstreamConfigs performs live + // RFC 7591 registration against upstream IdPs before authserver.New ever + // reaches Config.Validate(), so a malformed trusted-issuer config must + // fail here, before that side-effecting work runs — not on a crash loop + // after it, which would also orphan an upstream registration on every + // restart with the default in-memory DCR store. + if err := validateTrustedIssuers(c.TrustedIssuers, c.Issuer); err != nil { + return err + } return c.validateBaselineClientScopes() } @@ -394,10 +431,14 @@ type OAuth2UpstreamRunConfig struct { DCRConfig *DCRUpstreamConfig `json:"dcr_config,omitempty" yaml:"dcr_config,omitempty"` // AllowPrivateIPs permits the upstream provider's HTTP client to connect to - // private IP ranges (RFC-1918, link-local). Use only when the upstream is - // hosted inside the same cluster and has no public endpoint. HTTP-scheme - // restrictions are unchanged — HTTPS is still required for non-localhost hosts. - // Defaults to false. + // private IP ranges (RFC-1918, link-local). When DCRConfig is set, this + // also gates the DCR discovery and registration calls made on this + // upstream's behalf (see pkg/authserver/runner/dcr_adapter.go), so a + // single flag covers the whole upstream rather than needing a separate + // DCR-specific setting. Use only when the upstream is hosted inside the + // same cluster and has no public endpoint. HTTP-scheme restrictions are + // unchanged — HTTPS is still required for non-localhost hosts. Defaults + // to false. AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"` // InsecureAllowHTTP permits plain-HTTP authorization and token endpoint URLs @@ -612,6 +653,11 @@ type Config struct { // If zero, defaults to 10 minutes. AuthCodeLifespan time.Duration + // DelegationTokenLifespan is the maximum lifetime for delegated tokens issued + // via RFC 8693 token exchange. The actual lifetime is the minimum of this value + // and the subject token's remaining lifetime. If zero, defaults to 15 minutes. + DelegationTokenLifespan time.Duration + // Upstreams contains configurations for connecting to upstream IDPs. // At least one upstream is required - the server delegates authentication to the upstream IDP. // Multiple upstreams form a sequential authorization chain. @@ -670,6 +716,13 @@ type Config struct { // Only set this for in-cluster Kubernetes deployments on a trusted network. // Production deployments reachable outside the cluster MUST use https://. InsecureAllowHTTP bool + + // TrustedIssuers lists external OIDC issuers whose tokens are accepted as + // subject tokens during RFC 8693 token exchange. See the identically + // named field on RunConfig for the full doc comment, including the + // fail-closed AllowedActors semantics and the audience/scope constraints + // operators must account for. + TrustedIssuers []tokenexchange.TrustedIssuer } // Validate checks that the Config is valid. @@ -726,13 +779,23 @@ func (c *Config) Validate() error { } } - if c.CIMDEnabled && c.CIMDCacheMaxSize < 1 { - return fmt.Errorf("cimd.cache_max_size must be >= 1 when CIMD is enabled") + if err := c.validateCIMDBounds(); err != nil { + return err } - if c.CIMDEnabled && c.CIMDCacheFallbackTTL < 0 { - return fmt.Errorf("cimd.cache_fallback_ttl must be non-negative when CIMD is enabled") + + if err := c.validateDelegationTokenLifespan(); err != nil { + return err } + // RunConfig.Validate() also runs this (see the comment there for why: + // buildUpstreamConfigs's live DCR registration happens before this method + // is reached), but a caller that constructs Config directly bypasses that, + // same as the BaselineClientScopes check above. + if err := validateTrustedIssuers(c.TrustedIssuers, c.Issuer); err != nil { + return err + } + c.warnTrustedIssuerAudiences() + slog.Debug("authserver config validation passed", "issuer", c.Issuer, "upstream_count", len(c.Upstreams), @@ -740,6 +803,132 @@ func (c *Config) Validate() error { return nil } +// validateCIMDBounds rejects invalid CIMD cache bounds when CIMD is enabled. +// When CIMD is disabled the cache fields are ignored. +func (c *Config) validateCIMDBounds() error { + if !c.CIMDEnabled { + return nil + } + if c.CIMDCacheMaxSize < 1 { + return fmt.Errorf("cimd.cache_max_size must be >= 1 when CIMD is enabled") + } + if c.CIMDCacheFallbackTTL < 0 { + return fmt.Errorf("cimd.cache_fallback_ttl must be non-negative when CIMD is enabled") + } + return nil +} + +// validateDelegationTokenLifespan rejects negative or excessively long delegation +// token lifespans. Capped at oauthserver.MaxAccessTokenLifespan (the same ceiling +// the token-exchange Factory enforces) so validation and construction agree on a +// single source of truth — delegated tokens should be short-lived. Zero is +// accepted; applyDefaults substitutes the default. +func (c *Config) validateDelegationTokenLifespan() error { + if c.DelegationTokenLifespan < 0 { + return fmt.Errorf("delegation token lifespan must not be negative") + } + if c.DelegationTokenLifespan > oauthserver.MaxAccessTokenLifespan { + return fmt.Errorf("delegation token lifespan must not exceed %v", oauthserver.MaxAccessTokenLifespan) + } + return nil +} + +// validateTrustedIssuers checks every configured TrustedIssuer as early as +// RunConfig.Validate can catch it, so a bad actor_claim or duplicate +// issuer_url fails before buildUpstreamConfigs performs live RFC 7591 +// registration against upstream IdPs — not after it, on a crash loop that +// orphans an upstream registration on every restart. Shared by +// RunConfig.Validate() and Config.Validate() (mirrors +// validateBaselineClientScopes); NewMultiIssuerTokenValidator repeats these +// checks again at server startup as defence in depth. +// +// issuer_url is checked by validateTrustedIssuerURL, jwks_url (when set) by +// validateJWKSEndpointURL — see their doc comments for the URL rules each +// enforces. The remaining structural checks (required fields, self-issuer +// collision, duplicate issuers, ActorClaim reachability) run via +// tokenexchange.ValidateTrustedIssuers. +func validateTrustedIssuers(issuers []tokenexchange.TrustedIssuer, selfIssuer string) error { + for _, ti := range issuers { + if err := validateTrustedIssuerURL(ti.IssuerURL, ti.InsecureAllowHTTP); err != nil { + return fmt.Errorf("trusted_issuers: issuer_url %q: %w", ti.IssuerURL, err) + } + // AllowPrivateIPs without a hand-configured jwks_url would let OIDC + // discovery — a document fetched from, and thus influenceable by, + // the external issuer itself — choose the private target the dial + // is allowed to reach. Requiring jwks_url pins that target to + // operator-supplied config instead. + // + // This is the fail-fast layer, not the only one: validateTrustedIssuer + // (multi_issuer_validator.go) enforces the same invariant inside + // NewMultiIssuerTokenValidator, so a caller constructing a validator + // without routing through Config.Validate is still covered. Note that + // ensureRegistered's ValidateJWKSURL does NOT cover it — that check is + // gated on net.ParseIP, so it only rejects private IP *literals*, and a + // discovery document advertising a private *hostname* passes it + // cleanly. Checking here and in the constructor is deliberate + // duplication, not redundancy. + if ti.AllowPrivateIPs && ti.JWKSURL == "" { + return fmt.Errorf( + "trusted_issuers: issuer_url %q: allow_private_ips requires jwks_url to be set explicitly; "+ + "otherwise OIDC discovery — fetched from the external issuer — would choose the private target", + ti.IssuerURL) + } + if ti.JWKSURL != "" { + if err := validateJWKSEndpointURL(ti.JWKSURL, ti.InsecureAllowHTTP, ti.AllowPrivateIPs); err != nil { + return fmt.Errorf("trusted_issuers: jwks_url %q: %w", ti.JWKSURL, err) + } + } + } + if err := tokenexchange.ValidateTrustedIssuers(issuers, selfIssuer); err != nil { + return fmt.Errorf("trusted_issuers: %w", err) + } + return nil +} + +// validateJWKSEndpointURL checks that rawURL parses, has a host, uses the +// "https" scheme (or "http" when insecureAllowHTTP is set), and — when the +// host is an IP literal — is not a private or loopback address unless +// allowPrivateIPs permits it. Unlike validateIssuerURL, it does not enforce +// OIDC issuer-identifier rules (no query/fragment/trailing-slash) since a +// JWKS endpoint legitimately carries those. +// +// Delegates to tokenexchange.ValidateJWKSURL, the same predicate the runtime +// choke point (ensureRegistered, called on every JWKS fetch) enforces — the two +// were previously separate implementations that had drifted apart (a +// runtime check laxer than this one would silently defeat this config-time +// guard), so this is now the single source of truth for both. +// +// Deliberately not networking.ValidateEndpointURL / +// ValidateEndpointURLWithInsecure: both also honor the +// INSECURE_DISABLE_URL_VALIDATION environment variable, which would let an +// unrelated env var silently disable this SSRF-relevant scheme check; the +// insecure variant also skips the parse/host check entirely rather than +// only relaxing the scheme. This helper takes its "insecure" bits solely +// from the issuer's own explicit InsecureAllowHTTP/AllowPrivateIPs fields. +func validateJWKSEndpointURL(rawURL string, insecureAllowHTTP, allowPrivateIPs bool) error { + return tokenexchange.ValidateJWKSURL(rawURL, insecureAllowHTTP, allowPrivateIPs) +} + +// warnTrustedIssuerAudiences logs a warning for each TrustedIssuer whose +// ExpectedAudience is absent from AllowedAudiences. This is not a hard +// error: a subject token may carry additional audiences beyond +// ExpectedAudience, so the mismatch is sufficient-but-not-necessary for +// every exchange from that issuer to fail — an operator may know a specific +// subject token will present a matching aud even though ExpectedAudience +// itself is not in AllowedAudiences. But when it's not intentional, this is +// the invalid_target footgun documented on the TrustedIssuers field: warn so +// it surfaces at startup instead of at the first exchange attempt. +func (c *Config) warnTrustedIssuerAudiences() { + for _, ti := range c.TrustedIssuers { + if !slices.Contains(c.AllowedAudiences, ti.ExpectedAudience) { + slog.Warn("trusted issuer's expected_audience is not in allowed_audiences; "+ + "token exchange will fail with invalid_target unless subject tokens from it "+ + "carry an additional audience matching one", + "issuer", ti.IssuerURL, "expected_audience", ti.ExpectedAudience) + } + } +} + // Validate checks that the OAuth2UpstreamRunConfig is internally consistent. // It enforces the mutual exclusivity of ClientID and DCRConfig: exactly one must // be set. A ClientID is required for pre-provisioned clients; a DCRConfig is @@ -957,6 +1146,10 @@ func (c *Config) applyDefaults() error { c.AuthCodeLifespan = 10 * time.Minute slog.Debug("applied default auth code lifespan", "duration", c.AuthCodeLifespan) } + if c.DelegationTokenLifespan == 0 { + c.DelegationTokenLifespan = 15 * time.Minute + slog.Debug("applied default delegation token lifespan", "duration", c.DelegationTokenLifespan) + } if c.HMACSecrets == nil { secret := make([]byte, servercrypto.MinSecretLength) if _, err := rand.Read(secret); err != nil { @@ -991,7 +1184,45 @@ func (c *Config) applyDefaults() error { // MUST use the "https" scheme, except for localhost during development. // When insecureAllowHTTP is true, http:// is also permitted for non-localhost // hosts (for in-cluster Kubernetes deployments on trusted networks). +// +// This server's own issuer is additionally held to a no-trailing-slash rule +// that OIDC itself does not require (see validateIssuerURLCore's +// allowTrailingSlash parameter) — defensible here only because we control +// this value, unlike a trusted external issuer (validateTrustedIssuerURL). func validateIssuerURL(issuer string, insecureAllowHTTP bool) error { + return validateIssuerURLCore(issuer, insecureAllowHTTP, true, false) +} + +// validateTrustedIssuerURL is like validateIssuerURL but never exempts +// localhost from the HTTPS requirement: a trusted external issuer is not +// this server's own issuer, so it must not inherit the same-host +// development convenience validateIssuerURL grants the server's own issuer +// and AuthorizationEndpointBaseURL. Without this, "issuer_url: +// http://localhost:9000" with insecure_allow_http: false would pass config +// validation here yet fail at runtime, since the per-issuer HTTP client is +// still built with InsecureAllowHTTP=false (see NewMultiIssuerTokenValidator) +// — jwks_url has no such exemption, so the two would otherwise disagree. +// +// Unlike validateIssuerURL, a trailing slash is accepted: OIDC Discovery §3 +// forbids query and fragment components on an issuer identifier, but not a +// trailing slash — §4.1 only requires one be trimmed before the well-known +// discovery path is appended, which presupposes a trailing-slash issuer is +// legal in the first place, and §4.3 requires the discovery document's +// "issuer" to match the token's "iss" verbatim. Microsoft Entra ID v1 — the +// default for a newly registered API — issues +// "iss": "https://sts.windows.net/{tenant}/" with a trailing slash, so +// rejecting it here would make v1 tokens impossible to configure at all. +func validateTrustedIssuerURL(issuer string, insecureAllowHTTP bool) error { + return validateIssuerURLCore(issuer, insecureAllowHTTP, false, true) +} + +// validateIssuerURLCore is the shared implementation behind validateIssuerURL +// and validateTrustedIssuerURL. localhostExempt controls whether a loopback +// host is treated as HTTPS-exempt regardless of insecureAllowHTTP. +// allowTrailingSlash controls whether a trailing slash on the issuer is +// accepted — see validateTrustedIssuerURL's doc comment for why the trusted- +// issuer path must allow it while this server's own issuer does not. +func validateIssuerURLCore(issuer string, insecureAllowHTTP, localhostExempt, allowTrailingSlash bool) error { if issuer == "" { return fmt.Errorf("issuer is required") } @@ -1016,20 +1247,34 @@ func validateIssuerURL(issuer string, insecureAllowHTTP bool) error { if parsed.Fragment != "" { return fmt.Errorf("must not contain fragment component") } + // Userinfo is rejected on two independent grounds. It cannot work: OIDC + // Discovery 1.0 Section 4.3 compares the discovery document's "issuer" + // against this value by exact string match, and no provider echoes back + // embedded credentials, so such an issuer always fails discovery. And it + // must not be stored: a password here would sit in the RunConfig and be + // echoed by the validation errors and startup warnings that quote the + // issuer URL. Rejecting it outright beats redacting it at every use. + // Note that parsed.User is non-nil even for "https://user@host" with no + // password, which is equally unusable as an issuer identifier. + if parsed.User != nil { + return fmt.Errorf("must not contain userinfo (credentials in the URL)") + } - // HTTPS is required unless it's a loopback address (for development) or - // insecureAllowHTTP is explicitly set for trusted in-cluster deployments. + // HTTPS is required unless it's a loopback address (for development, and + // only when localhostExempt) or insecureAllowHTTP is explicitly set. if parsed.Scheme != "https" { if parsed.Scheme != "http" { return fmt.Errorf("scheme must be https (or http for localhost)") } - if !networking.IsLocalhost(parsed.Host) && !insecureAllowHTTP { + if !insecureAllowHTTP && (!localhostExempt || !networking.IsLocalhost(parsed.Host)) { return fmt.Errorf("http scheme is only allowed for localhost, use https for %s", parsed.Hostname()) } } - // Issuer must not have trailing slash per OIDC spec - if strings.HasSuffix(issuer, "/") { + // Not an OIDC requirement — see validateTrustedIssuerURL's doc comment. + // ToolHive's own issuer is held to this stricter, self-imposed rule + // since we control the value; a trusted external issuer is not. + if !allowTrailingSlash && strings.HasSuffix(issuer, "/") { return fmt.Errorf("must not have trailing slash") } diff --git a/pkg/authserver/config_test.go b/pkg/authserver/config_test.go index 850a7d0347..b067163f71 100644 --- a/pkg/authserver/config_test.go +++ b/pkg/authserver/config_test.go @@ -5,6 +5,7 @@ package authserver import ( "bytes" + "log/slog" "strings" "testing" "time" @@ -14,6 +15,7 @@ import ( servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto" "github.com/stacklok/toolhive/pkg/authserver/server/keys" "github.com/stacklok/toolhive/pkg/authserver/server/registration" + "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange" "github.com/stacklok/toolhive/pkg/authserver/upstream" ) @@ -47,6 +49,15 @@ func TestValidateIssuerURL(t *testing.T) { {name: "missing host", issuer: "https://", wantErr: true, errMsg: "host is required"}, {name: "query component", issuer: "https://example.com?foo=bar", wantErr: true, errMsg: "must not contain query"}, {name: "fragment component", issuer: "https://example.com#section", wantErr: true, errMsg: "must not contain fragment"}, + { + name: "userinfo with password", issuer: "https://user:hunter2@example.com", + wantErr: true, errMsg: "must not contain userinfo", + }, + { + // url.Parse populates User for a bare username too. + name: "userinfo without password", issuer: "https://user@example.com", + wantErr: true, errMsg: "must not contain userinfo", + }, {name: "http non-localhost", issuer: "http://example.com", wantErr: true, errMsg: "http scheme is only allowed for localhost"}, {name: "ftp scheme", issuer: "ftp://example.com", wantErr: true, errMsg: "scheme must be https"}, {name: "trailing slash", issuer: "https://example.com/", wantErr: true, errMsg: "must not have trailing slash"}, @@ -264,6 +275,9 @@ func TestConfigApplyDefaults(t *testing.T) { func assertError(t *testing.T, err error, wantErr bool, errMsg string) { t.Helper() if wantErr { + if errMsg == "" { + t.Fatal("wantErr is true but errMsg is empty: strings.Contains(x, \"\") is always true, so this case would pass unconditionally") + } if err == nil { t.Errorf("expected error containing %q, got nil", errMsg) } else if !strings.Contains(err.Error(), errMsg) { @@ -669,3 +683,495 @@ func TestConfigApplyDefaults_CIMD(t *testing.T) { }) } } + +// TestConfigValidate_DelegationTokenLifespan covers the RFC 8693 delegation +// token lifespan bounds added to Config.Validate: zero is accepted (it is +// defaulted later by applyDefaults), values in (0, 24h] are accepted, and +// negative or over-24h values are rejected. +func TestConfigValidate_DelegationTokenLifespan(t *testing.T) { + t.Parallel() + + // base returns a minimally-valid Config so each case isolates the + // DelegationTokenLifespan check from unrelated validation failures. + base := func() Config { + return Config{ + Issuer: "https://example.com", + KeyProvider: keys.NewGeneratingProvider(keys.DefaultAlgorithm), + HMACSecrets: &servercrypto.HMACSecrets{Current: make([]byte, 32)}, + Upstreams: []UpstreamConfig{{ + Name: "default", + Type: UpstreamProviderTypeOAuth2, + OAuth2Config: &upstream.OAuth2Config{ + CommonOAuthConfig: upstream.CommonOAuthConfig{ClientID: "c", RedirectURI: "https://example.com/cb"}, + AuthorizationEndpoint: "https://idp.example.com/authorize", + TokenEndpoint: "https://idp.example.com/token", + }, + }}, + AllowedAudiences: []string{"https://mcp.example.com"}, + } + } + + tests := []struct { + name string + lifespan time.Duration + wantErr bool + errMsg string + }{ + {name: "zero accepted (defaulted later)", lifespan: 0}, + {name: "valid 15m", lifespan: 15 * time.Minute}, + {name: "valid 1h", lifespan: time.Hour}, + {name: "valid 24h boundary", lifespan: 24 * time.Hour}, + {name: "negative rejected", lifespan: -time.Second, wantErr: true, errMsg: "delegation token lifespan must not be negative"}, + {name: "over 24h rejected", lifespan: 24*time.Hour + time.Second, wantErr: true, errMsg: "delegation token lifespan must not exceed 24h"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + cfg := base() + cfg.DelegationTokenLifespan = tt.lifespan + assertError(t, cfg.Validate(), tt.wantErr, tt.errMsg) + }) + } +} + +// TestConfigApplyDefaults_DelegationTokenLifespan verifies that applyDefaults +// fills a zero DelegationTokenLifespan with the 15-minute default and preserves +// a caller-supplied value. +func TestConfigApplyDefaults_DelegationTokenLifespan(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input time.Duration + want time.Duration + }{ + {name: "zero gets 15m default", input: 0, want: 15 * time.Minute}, + {name: "custom value preserved", input: 5 * time.Minute, want: 5 * time.Minute}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + cfg := Config{Issuer: "https://example.com", DelegationTokenLifespan: tt.input} + require.NoError(t, cfg.applyDefaults()) + require.Equal(t, tt.want, cfg.DelegationTokenLifespan) + }) + } +} + +// TestConfigValidate_TrustedIssuers covers validateTrustedIssuers as reached +// from Config.Validate: the URL-shape checks (validateTrustedIssuerURL on +// issuer_url, validateJWKSEndpointURL on jwks_url) and the structural checks +// delegated to tokenexchange.ValidateTrustedIssuers. +func TestConfigValidate_TrustedIssuers(t *testing.T) { + t.Parallel() + + // base returns a minimally-valid Config (Issuer "https://example.com", + // AllowedAudiences ["https://mcp.example.com"]) so each case isolates the + // TrustedIssuers check from unrelated validation failures. + base := func() Config { + return Config{ + Issuer: "https://example.com", + KeyProvider: keys.NewGeneratingProvider(keys.DefaultAlgorithm), + HMACSecrets: &servercrypto.HMACSecrets{Current: make([]byte, 32)}, + Upstreams: []UpstreamConfig{{ + Name: "default", + Type: UpstreamProviderTypeOAuth2, + OAuth2Config: &upstream.OAuth2Config{ + CommonOAuthConfig: upstream.CommonOAuthConfig{ClientID: "c", RedirectURI: "https://example.com/cb"}, + AuthorizationEndpoint: "https://idp.example.com/authorize", + TokenEndpoint: "https://idp.example.com/token", + }, + }}, + AllowedAudiences: []string{"https://mcp.example.com"}, + } + } + + tests := []struct { + name string + issuers []tokenexchange.TrustedIssuer + wantErr bool + errMsg string + }{ + { + name: "no trusted issuers is byte-identical to before", + issuers: nil, + }, + { + name: "issuer_url bad scheme rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "htps://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "issuer_url", + }, + { + name: "issuer_url empty rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "issuer is required", + }, + { + name: "issuer_url http without per-issuer insecure_allow_http rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "http://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "http scheme is only allowed for localhost", + }, + { + name: "issuer_url http with per-issuer insecure_allow_http accepted", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "http://idp.example.com", ExpectedAudience: "https://mcp.example.com", InsecureAllowHTTP: true, AllowedDelegateClients: []string{"*"}}, + }, + }, + { + // Unlike Config.Issuer, a trusted issuer gets no localhost + // exemption: it isn't this server's own issuer, so the same + // same-host development convenience doesn't apply — see + // validateTrustedIssuerURL's doc comment. Without + // insecure_allow_http, http://localhost must be rejected here + // the same as any other http issuer_url. + name: "issuer_url http localhost rejected without per-issuer insecure_allow_http", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "http://localhost:8080", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "http scheme is only allowed for localhost", + }, + { + name: "issuer_url http localhost accepted with per-issuer insecure_allow_http", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "http://localhost:8080", ExpectedAudience: "https://mcp.example.com", InsecureAllowHTTP: true, AllowedDelegateClients: []string{"*"}}, + }, + }, + { + // Azure AD B2C's real jwks_uri carries a query string + // (?p=B2C_1_...); jwks_url must be validated as an ordinary + // endpoint URL, not an OIDC issuer identifier, or a legitimate + // production IdP would be rejected. + name: "jwks_url with query string accepted", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "https://idp.example.com/keys?p=B2C_1_signin", + AllowedDelegateClients: []string{"*"}, + }, + }, + }, + { + name: "jwks_url with trailing slash accepted", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "https://idp.example.com/keys/", + AllowedDelegateClients: []string{"*"}, + }, + }, + }, + { + name: "jwks_url bad scheme rejected", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "ftp://idp.example.com/keys", + AllowedDelegateClients: []string{"*"}, + }, + }, + wantErr: true, + errMsg: "jwks_url", + }, + { + name: "jwks_url http rejected without per-issuer insecure_allow_http", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "http://idp.example.com/keys", + AllowedDelegateClients: []string{"*"}, + }, + }, + wantErr: true, + errMsg: "jwks_url", + }, + { + name: "jwks_url http accepted with per-issuer insecure_allow_http", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "http://idp.example.com/keys", + InsecureAllowHTTP: true, + AllowedDelegateClients: []string{"*"}, + }, + }, + }, + { + name: "jwks_url private IP literal rejected without allow_private_ips", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "https://10.0.0.5/keys", + AllowedDelegateClients: []string{"*"}, + }, + }, + wantErr: true, + errMsg: "private or loopback", + }, + { + name: "jwks_url private IP literal accepted with allow_private_ips", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "https://10.0.0.5/keys", + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }, + }, + }, + { + // Mirrors the CRD's Kubebuilder CEL rule requiring jwksUrl + // whenever allowPrivateIPs is set (mcpexternalauthconfig_types.go): + // without a hand-configured jwks_url, OIDC discovery — a document + // fetched from, and thus influenceable by, the external issuer + // itself — would choose the private JWKS dial target. A + // hand-written RunConfig must not be able to bypass what the CRD + // path already guarantees. + name: "allow_private_ips without jwks_url rejected", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }, + }, + wantErr: true, + errMsg: "allow_private_ips requires jwks_url", + }, + { + name: "missing expected_audience rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "expected_audience is required", + }, + { + name: "issuer_url equal to Config.Issuer rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "must not equal the authorization server's own issuer", + }, + { + name: "duplicate issuer_url rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "configured more than once", + }, + { + name: "actor_claim sub rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", ActorClaim: "sub", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "actor_claim", + }, + { + // Unlike Config.Issuer, a trusted issuer_url permits a trailing + // slash: Microsoft Entra ID v1 (the default for a newly + // registered API) issues "iss": "https://sts.windows.net/{tenant}/" + // with one, and OIDC Discovery has no rule against it — see + // validateTrustedIssuerURL's doc comment. + name: "issuer_url with trailing slash accepted", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://sts.windows.net/11111111-2222-3333-4444-555555555555/", + ExpectedAudience: "https://mcp.example.com", + AllowedDelegateClients: []string{"*"}, + }, + }, + }, + { + // A may_act-only issuer is legitimate: an empty AllowedActors + // means every non-may_act token from it is rejected at + // validation time, not that the config itself is invalid. + name: "empty allowed_actors accepted", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedActors: nil, AllowedDelegateClients: []string{"*"}}, + }, + }, + { + // #5989 hardening reaches Config.Validate through the same + // shared validateTrustedIssuers -> tokenexchange.ValidateTrustedIssuers + // path as the constructor-level check. + name: "absent allowed_delegate_clients rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com"}, + }, + wantErr: true, + errMsg: "allowed_delegate_clients is required", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + cfg := base() + cfg.TrustedIssuers = tt.issuers + assertError(t, cfg.Validate(), tt.wantErr, tt.errMsg) + }) + } +} + +// TestRunConfigValidate_TrustedIssuers asserts that RunConfig.Validate +// itself catches every TrustedIssuers failure mode reachable through +// validateTrustedIssuers — the four structural checks, the issuer_url shape +// check, and (mirroring TestConfigValidate_TrustedIssuers) the jwks_url +// private-IP guard — not only Config.Validate, because it's the same +// shared function both call. This matters because buildUpstreamConfigs +// performs live RFC 7591 registration against upstream IdPs before +// Config.Validate is ever reached (see the comment on RunConfig.Validate). +// A bad trusted issuer caught only at the Config layer would orphan an +// upstream client registration on every restart of the resulting crash loop. +func TestRunConfigValidate_TrustedIssuers(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + issuers []tokenexchange.TrustedIssuer + wantErr bool + errMsg string + }{ + { + name: "no trusted issuers passes", + issuers: nil, + }, + { + name: "malformed issuer_url rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "htps://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "issuer_url", + }, + { + name: "missing expected_audience rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "expected_audience is required", + }, + { + name: "issuer_url equal to RunConfig.Issuer rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "must not equal the authorization server's own issuer", + }, + { + name: "duplicate issuer_url rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "configured more than once", + }, + { + name: "actor_claim sub rejected", + issuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", ActorClaim: "sub", AllowedDelegateClients: []string{"*"}}, + }, + wantErr: true, + errMsg: "actor_claim", + }, + { + name: "jwks_url private IP literal rejected without allow_private_ips", + issuers: []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + JWKSURL: "https://10.0.0.5/keys", + AllowedDelegateClients: []string{"*"}, + }, + }, + wantErr: true, + errMsg: "private or loopback", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + cfg := RunConfig{Issuer: "https://example.com", TrustedIssuers: tt.issuers} + assertError(t, cfg.Validate(), tt.wantErr, tt.errMsg) + }) + } +} + +// TestConfig_WarnTrustedIssuerAudiences pins the invalid_target footgun +// warning: it must fire at startup for an ExpectedAudience absent from +// AllowedAudiences, and stay silent when the audience is present. +// +// See TestNewEmbeddedAuthServer_TrustedIssuers in +// runner/embeddedauthserver_test.go for the same pattern with the rationale +// spelled out. +// +//nolint:paralleltest // captures the package-global slog.Default() +func TestConfig_WarnTrustedIssuerAudiences(t *testing.T) { + tests := []struct { + name string + audiences []string + wantWarn bool + }{ + { + name: "expected_audience absent from allowed_audiences warns", + audiences: []string{"https://other.example.com"}, + wantWarn: true, + }, + { + name: "expected_audience present in allowed_audiences is silent", + audiences: []string{"https://mcp.example.com"}, + wantWarn: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + cfg := Config{ + AllowedAudiences: tt.audiences, + TrustedIssuers: []tokenexchange.TrustedIssuer{ + {IssuerURL: "https://idp.example.com", ExpectedAudience: "https://mcp.example.com", AllowedDelegateClients: []string{"*"}}, + }, + } + cfg.warnTrustedIssuerAudiences() + + if tt.wantWarn { + require.Contains(t, buf.String(), "trusted issuer's expected_audience is not in allowed_audiences") + } else { + require.Empty(t, buf.String()) + } + }) + } +} diff --git a/pkg/authserver/integration_delegation_audit_test.go b/pkg/authserver/integration_delegation_audit_test.go new file mode 100644 index 0000000000..5623787a29 --- /dev/null +++ b/pkg/authserver/integration_delegation_audit_test.go @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package authserver + +import ( + "context" + "crypto/rsa" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/go-jose/go-jose/v4/jwt" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/authserver/server/registration" + "github.com/stacklok/toolhive/pkg/oauthproto" +) + +// readSingleAuditEvent reads the audit log file at path and decodes exactly +// one newline-delimited JSON event, failing the test if there isn't exactly one. +func readSingleAuditEvent(t *testing.T, path string) map[string]any { + t.Helper() + + data, err := os.ReadFile(path) + require.NoError(t, err) + require.NotEmpty(t, strings.TrimSpace(string(data)), "audit log is empty; expected one event") + + lines := strings.Split(strings.TrimSpace(string(data)), "\n") + require.Len(t, lines, 1, "expected exactly one audit event, got: %q", string(data)) + + var event map[string]any + require.NoError(t, json.Unmarshal([]byte(lines[0]), &event)) + return event +} + +// signTestJWT signs a JWT with key using the standard test issuer/audience and +// a 30-minute expiry, embedding subject and any extraClaims (e.g. "client_id" +// for a delegation act claim). extraClaims may be nil. +func signTestJWT(t *testing.T, key *rsa.PrivateKey, subject string, extraClaims map[string]any) string { + t.Helper() + + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.RS256, Key: key}, + (&jose.SignerOptions{}).WithType("JWT").WithHeader("kid", "test-key"), + ) + require.NoError(t, err) + + now := time.Now() + builder := jwt.Signed(signer). + Claims(jwt.Claims{ + Issuer: testIssuer, + Subject: subject, + Audience: jwt.Audience{testAudience}, + Expiry: jwt.NewNumericDate(now.Add(30 * time.Minute)), + IssuedAt: jwt.NewNumericDate(now), + }) + if extraClaims != nil { + builder = builder.Claims(extraClaims) + } + token, err := builder.Serialize() + require.NoError(t, err) + return token +} + +// exchangeForDelegatedToken performs an RFC 8693 token exchange against +// serverURL using the given confidential client credentials and subject +// token, and returns the resulting delegated access token. +func exchangeForDelegatedToken(t *testing.T, serverURL, clientID, clientSecret, subjectToken string) string { + t.Helper() + + resp := makeTokenRequest(t, serverURL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {clientID}, + "client_secret": {clientSecret}, + }) + defer resp.Body.Close() + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange should succeed, got %d (body: %v)", resp.StatusCode, body) + delegated, ok := body["access_token"].(string) + require.True(t, ok, "access_token should be a string") + require.NotEmpty(t, delegated) + return delegated +} + +// auditEventForToken runs token through the audit(auth(stub)) middleware +// chain — audit outermost, auth innermost — and returns the single resulting +// audit event. The stub handler does not re-publish the identity itself: this +// pins that TokenValidator.Middleware alone publishes the validated identity +// to the audit holder via the IdentityHolder read-back. +func auditEventForToken(t *testing.T, key *rsa.PrivateKey, token, validationMsg string) map[string]any { + t.Helper() + return auditEventForTokenWithConfig(t, key, token, validationMsg, audit.Config{}) +} + +// auditEventForTokenWithConfig is auditEventForToken for tests that need a +// non-default audit.Config (e.g. MaxDelegationDepth). cfg is taken by value +// and its LogFile field is set on the local copy, so the caller's config is +// never mutated. +func auditEventForTokenWithConfig( + t *testing.T, key *rsa.PrivateKey, token, validationMsg string, cfg audit.Config, +) map[string]any { + t.Helper() + + // Build a validator that trusts the server's issuer and keys, exactly like + // the middleware in front of a protected resource server. The in-process + // key provider avoids self-referential HTTP JWKS fetches. + validator, err := auth.NewTokenValidator(context.Background(), auth.TokenValidatorConfig{ + Issuer: testIssuer, + Audience: testAudience, + }, auth.WithKeyProvider(&testKeyProvider{key: key})) + require.NoError(t, err) + + cfg.LogFile = filepath.Join(t.TempDir(), "audit.log") + auditor, err := audit.NewAuditorWithTransport(&cfg, "streamable-http") + require.NoError(t, err) + t.Cleanup(func() { + assert.NoError(t, auditor.Close()) + }) + + stub := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + req := httptest.NewRequest(http.MethodPost, "/messages", strings.NewReader("{}")) + req.Header.Set("Authorization", "Bearer "+token) + rr := httptest.NewRecorder() + auditor.Middleware(validator.Middleware(stub)).ServeHTTP(rr, req) + + require.Equal(t, http.StatusOK, rr.Code, validationMsg) + + return readSingleAuditEvent(t, cfg.LogFile) +} + +// TestIntegration_TokenExchange_AuditDelegationChain proves the end-to-end +// audit story for RFC 8693 delegation: a delegated token minted by the real +// token-exchange handler is validated by the auth middleware, and the audit +// middleware emits an event whose delegation chain records the acting agent. +func TestIntegration_TokenExchange_AuditDelegationChain(t *testing.T) { + t.Parallel() + + const ( + agentClientID = "test-audit-agent-client" + agentClientSecret = "test-audit-agent-secret" + delegatedUserSub = "audit-delegated-user-sub" + ) + + agentClient, err := registration.New(registration.Config{ + ID: agentClientID, + Secret: agentClientSecret, + Public: false, + GrantTypes: []string{oauthproto.GrantTypeTokenExchange}, + Scopes: registration.DefaultScopes, + Audience: []string{testAudience}, + }) + require.NoError(t, err) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, withExtraClient(agentClient)) + + // Mint the subject token (user) signed by the server's key. + subjectToken := signTestJWT(t, ts.PrivateKey, delegatedUserSub, map[string]any{ + "client_id": agentClientID, + }) + + // Exchange it for a delegated token carrying act.sub = agent. + delegated := exchangeForDelegatedToken(t, ts.Server.URL, agentClientID, agentClientSecret, subjectToken) + + // Chain: audit (outer) -> auth -> bare stub handler. + event := auditEventForToken(t, ts.PrivateKey, delegated, + "delegated token must validate through the auth middleware") + + subjects, ok := event["subjects"].(map[string]any) + require.True(t, ok, "subjects should be a map") + assert.Equal(t, delegatedUserSub, subjects["user_id"], + "audit subject must be the delegated user, not the agent") + + chain, ok := event["delegation"].(map[string]any) + require.True(t, ok, "delegation must be present in the audit event") + assert.Equal(t, false, chain["truncated"]) + assert.Equal(t, float64(0), chain["omitted"]) + assert.Equal(t, false, chain["malformed"]) + hops, ok := chain["chain"].([]any) + require.True(t, ok, "chain should be an array") + require.Len(t, hops, 1, "a single exchange yields a single hop") + hop, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, agentClientID, hop["sub"], + "the delegation chain must record the acting agent client") +} + +// TestIntegration_AuditMiddleware_NonDelegatedTokenOmitsDelegationChain is the +// negative companion to TestIntegration_TokenExchange_AuditDelegationChain: a +// plain (non-delegated) token through the same audit(auth(stub)) chain must +// produce an audit event with NO delegation member. +func TestIntegration_AuditMiddleware_NonDelegatedTokenOmitsDelegationChain(t *testing.T) { + t.Parallel() + + const plainUserSub = "audit-plain-user-sub" + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m) + + // Mint a plain subject token signed by the server's key (no act claim). + plainToken := signTestJWT(t, ts.PrivateKey, plainUserSub, nil) + + event := auditEventForToken(t, ts.PrivateKey, plainToken, + "plain token must validate through the auth middleware") + + subjects, ok := event["subjects"].(map[string]any) + require.True(t, ok, "subjects should be a map") + assert.Equal(t, plainUserSub, subjects["user_id"], + "audit subject must be the authenticated user") + + _, exists := event["delegation"] + assert.False(t, exists, + "a non-delegated token must not produce a delegation member") +} + +// mintTwoHopDelegatedToken produces a token whose second hop is genuine +// production output of a real RFC 8693 token-exchange call — firstAgentClientID +// "delegates" to secondAgentClientID — while the first hop is a hand-signed +// stand-in for an earlier exchange this test never performs. What this pins +// is the handler's nesting direction: the second (real) actor must land at +// chain[0], outermost, with the hand-signed first actor at chain[1]. +// +// This AS can honour an RFC 8693 §4.4 may_act claim — checkDelegationConsent +// (handler.go) treats it as the authoritative consent signal and skips the +// client_id binding entirely when present — but no code path anywhere writes +// may_act, and the exchange does not propagate an inbound one: +// pkg/authserver/server/tokenexchange/validator.go:276-281 routes may_act +// into the structured MayAct field so it never lands in Extra, and +// handler.go:160 only ever copies Extra["act"]. So the subject token here +// hand-signs the one claim a may_act-emitting issuer would have written +// (plus a prior "act" entry standing in for that earlier hop); everything +// downstream of it — the consent check, the act-claim nesting, the issued +// token — is genuine production output of a single real exchange. +func mintTwoHopDelegatedToken(t *testing.T) (ts *testServerWithUpstream, delegated, firstAgentClientID, secondAgentClientID string) { + t.Helper() + + const ( + firstAgent = "test-first-agent-client" + secondAgent = "test-second-agent-client" + secondAgentSecret = "test-second-agent-secret" //nolint:gosec // test fixture, not a real credential + delegatedUserSub = "audit-multi-actor-user-sub" + ) + + secondAgentClient, err := registration.New(registration.Config{ + ID: secondAgent, + Secret: secondAgentSecret, + Public: false, + GrantTypes: []string{oauthproto.GrantTypeTokenExchange}, + Scopes: registration.DefaultScopes, + Audience: []string{testAudience}, + }) + require.NoError(t, err) + + m := startMockOIDC(t) + ts = setupTestServerWithMockOIDC(t, m, withExtraClient(secondAgentClient)) + + subjectToken := signTestJWT(t, ts.PrivateKey, delegatedUserSub, map[string]any{ + "client_id": firstAgent, // ignored: may_act wins + "act": map[string]any{"sub": firstAgent}, // prior hop + "may_act": map[string]any{"sub": secondAgent}, // RFC 8693 §4.4 consent + }) + + delegated = exchangeForDelegatedToken(t, ts.Server.URL, secondAgent, secondAgentSecret, subjectToken) + return ts, delegated, firstAgent, secondAgent +} + +// TestIntegration_TokenExchange_MultiActorDelegationChain proves the +// handler's nesting direction end to end: hop 2 (secondAgentClientID) is +// real handler output from an actual token-exchange call, hop 1 +// (firstAgentClientID) is the hand-signed stand-in described on +// mintTwoHopDelegatedToken, and the audit event for the resulting token +// orders the two hops outermost (most recent) first. +func TestIntegration_TokenExchange_MultiActorDelegationChain(t *testing.T) { + t.Parallel() + + ts, delegated, firstAgentClientID, secondAgentClientID := mintTwoHopDelegatedToken(t) + + event := auditEventForToken(t, ts.PrivateKey, delegated, + "re-delegated token must validate through the auth middleware") + + chain, ok := event["delegation"].(map[string]any) + require.True(t, ok, "delegation must be present in the audit event") + assert.Equal(t, false, chain["truncated"]) + hops, ok := chain["chain"].([]any) + require.True(t, ok, "chain should be an array") + require.Len(t, hops, 2, "a re-delegated token yields two hops") + + hop0, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, secondAgentClientID, hop0["sub"], + "chain[0] (outermost) must be the client that performed the second exchange") + hop1, ok := hops[1].(map[string]any) + require.True(t, ok) + assert.Equal(t, firstAgentClientID, hop1["sub"], + "chain[1] (innermost) must be the earlier, first-hop actor") +} + +// TestIntegration_TokenExchange_TruncatedDelegationChainAuditEvent proves +// that an auditor configured with MaxDelegationDepth: 1 truncates a genuine +// 2-hop delegation chain down to its outermost (most recent) actor. +func TestIntegration_TokenExchange_TruncatedDelegationChainAuditEvent(t *testing.T) { + t.Parallel() + + ts, delegated, _, secondAgentClientID := mintTwoHopDelegatedToken(t) + + maxDepth := 1 + event := auditEventForTokenWithConfig(t, ts.PrivateKey, delegated, + "re-delegated token must validate through the auth middleware", + audit.Config{MaxDelegationDepth: &maxDepth}) + + chain, ok := event["delegation"].(map[string]any) + require.True(t, ok, "delegation must be present in the audit event") + assert.Equal(t, true, chain["truncated"]) + assert.Equal(t, float64(1), chain["omitted"], + "one of the two hops must be reported omitted") + + hops, ok := chain["chain"].([]any) + require.True(t, ok, "chain should be an array") + require.Len(t, hops, 1, "MaxDelegationDepth=1 must leave exactly one surviving hop") + hop0, ok := hops[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, secondAgentClientID, hop0["sub"], + "the surviving hop must be the outermost one") +} diff --git a/pkg/authserver/integration_test.go b/pkg/authserver/integration_test.go index 56c5132708..2644d3b939 100644 --- a/pkg/authserver/integration_test.go +++ b/pkg/authserver/integration_test.go @@ -36,8 +36,10 @@ import ( "github.com/stacklok/toolhive/pkg/authserver/server/keys" "github.com/stacklok/toolhive/pkg/authserver/server/registration" "github.com/stacklok/toolhive/pkg/authserver/server/session" + "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange" "github.com/stacklok/toolhive/pkg/authserver/storage" "github.com/stacklok/toolhive/pkg/authserver/upstream" + "github.com/stacklok/toolhive/pkg/oauthproto" ) const ( @@ -95,6 +97,15 @@ type testServerOptions struct { // which delegates to setupTestServer. setupTestServerWithRTProxy does not // accept testServerOption at all, so this field does not apply to it. upstreamFilter handlers.UpstreamFilter + // extraClients, when non-empty, are registered in storage in addition to the + // default public PKCE test client. Used to install a confidential client for + // flows the default public client cannot exercise (e.g. RFC 8693 token + // exchange, which requires a confidential acting client). + extraClients []fosite.Client + // trustedIssuers, when non-empty, is passed through to Config.TrustedIssuers, + // enabling RFC 8693 token exchange with subject tokens from external OIDC + // issuers in addition to self-issued ones. + trustedIssuers []tokenexchange.TrustedIssuer } // testServerOption is a functional option for test server setup. @@ -130,6 +141,23 @@ func withUpstreamFilter(f handlers.UpstreamFilter) testServerOption { } } +// withExtraClient registers an additional client in storage alongside the +// default public PKCE test client. +func withExtraClient(c fosite.Client) testServerOption { + return func(opts *testServerOptions) { + opts.extraClients = append(opts.extraClients, c) + } +} + +// withTrustedIssuers configures Config.TrustedIssuers, enabling token +// exchange with subject tokens from the given external OIDC issuers in +// addition to self-issued ones. +func withTrustedIssuers(issuers []tokenexchange.TrustedIssuer) testServerOption { + return func(opts *testServerOptions) { + opts.trustedIssuers = issuers + } +} + // withRedisBackedStorage swaps the default in-memory storage for a // miniredis-backed *RedisStorage. This exercises the same Lua scripts and // Redis-shape key layout used in production, while remaining hermetic and @@ -226,6 +254,11 @@ func setupTestServer(t *testing.T, opts ...testServerOption) *testServer { }) require.NoError(t, err) + // Register any extra clients (e.g. a confidential client for token exchange). + for _, c := range options.extraClients { + require.NoError(t, stor.RegisterClient(ctx, c)) + } + // 5. Build upstream config for newServer // When no upstream is provided, use a dummy config that satisfies validation // Note: Uses HTTPS to pass config validation @@ -255,6 +288,7 @@ func setupTestServer(t *testing.T, opts ...testServerOption) *testServer { Upstreams: []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: upstreamCfg}}, UpstreamFilter: options.upstreamFilter, AllowedAudiences: []string{"https://mcp.example.com"}, + TrustedIssuers: options.trustedIssuers, } // 7. Create server using newServer with test options @@ -585,6 +619,971 @@ func TestIntegration_TokenEndpoint_RefreshToken(t *testing.T) { require.GreaterOrEqual(t, replayResp.StatusCode, 400, "old refresh token must be rejected after rotation") } +// ============================================================================ +// RFC 8693 Token Exchange Wiring Tests +// ============================================================================ + +// TestIntegration_TokenExchange_PublicClientRejected proves that public clients +// are barred from the RFC 8693 token-exchange grant (only confidential clients +// may act on a user's behalf), and that the grant is wired into the fosite +// provider and reachable at the token endpoint. +// +// It does not assert a full delegated-token issuance: the handler requires a +// confidential client (RFC 8693 §2.1) and the shared test harness only +// registers a public client. Instead it relies on a decisive dispatch +// discriminator observable at the token endpoint: +// +// - With the token-exchange factory registered, a token-exchange request is +// routed to the handler, whose first guard rejects the public client with +// error=invalid_grant and a "token-exchange"-specific hint. +// - Without the factory, no handler claims grant_type=token-exchange and +// fosite returns error=invalid_request (see fosite NewAccessRequest: an +// unmatched grant yields ErrInvalidRequest). +// +// So invalid_grant with a token-exchange hint proves the handler is registered +// and executed — exactly the regression a dropped factory would introduce. +func TestIntegration_TokenExchange_PublicClientRejected(t *testing.T) { + t.Parallel() + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m) + + verifier := servercrypto.GeneratePKCEVerifier() + challenge := servercrypto.ComputePKCEChallenge(verifier) + + // Mint a genuine server-issued access token to use as the subject_token, so + // the request is a well-formed RFC 8693 exchange up to the client-type gate. + authCode, _ := completeAuthorizationFlow(t, ts.Server.URL, authorizationParams{ + ClientID: testClientID, + RedirectURI: testRedirectURI, + State: "token-exchange-dispatch", + Challenge: challenge, + Scope: "openid profile", + ResponseType: "code", + }) + tokenData := exchangeCodeForTokens(t, ts.Server.URL, authCode, verifier, testAudience) + subjectToken, ok := tokenData["access_token"].(string) + require.True(t, ok, "access_token should be a string") + require.NotEmpty(t, subjectToken) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {testClientID}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + errCode, _ := body["error"].(string) + errDesc, _ := body["error_description"].(string) + + require.Equal(t, http.StatusBadRequest, resp.StatusCode, + "token-exchange from a public client should be a 400, got %d (body: %v)", resp.StatusCode, body) + // The decisive assertion: the token-exchange handler ran and rejected the + // public client (invalid_grant), rather than the request falling through + // unhandled (invalid_request), which is what a missing factory would produce. + require.Equal(t, "invalid_grant", errCode, + "expected invalid_grant from the token-exchange handler; invalid_request would mean "+ + "the token-exchange factory is not wired into the provider") + assert.Contains(t, errDesc, "token-exchange", + "the rejection must originate from the token-exchange handler specifically") +} + +// TestIntegration_TokenExchange_ConfidentialClientHappyPath drives a full RFC 8693 +// delegation exchange over HTTP through the real fosite provider: a confidential +// acting client authenticates with client_secret_post, presents a server-signed +// subject token, and receives a delegated access token. +// +// Unlike the unit tests in the tokenexchange package (which call the handler +// directly with mock strategy/storage), this exercises the complete glued path — +// fosite NewAccessRequest -> client authentication -> handler dispatch -> +// PopulateTokenEndpointResponse JSON issuance — proving the token exchange is not +// just registered but functional end-to-end. +func TestIntegration_TokenExchange_ConfidentialClientHappyPath(t *testing.T) { + t.Parallel() + + const ( + agentClientID = "test-agent-client" + agentClientSecret = "test-agent-secret" + delegatedUserSub = "delegated-user-sub" + ) + + // A confidential client registered for the token-exchange grant is the acting + // agent. The handler rejects public clients, so this must be confidential. + agentClient, err := registration.New(registration.Config{ + ID: agentClientID, + Secret: agentClientSecret, + Public: false, + GrantTypes: []string{oauthproto.GrantTypeTokenExchange}, + Scopes: registration.DefaultScopes, + Audience: []string{testAudience}, + }) + require.NoError(t, err) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, withExtraClient(agentClient)) + + // Mint a subject token signed by the server's own key (the validator verifies + // against the server's JWKS). client_id must equal the acting client so the + // RFC 8693 §4.4 delegation-consent check (client_id binding) passes. + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.RS256, Key: ts.PrivateKey}, + (&jose.SignerOptions{}).WithType("JWT").WithHeader("kid", "test-key"), + ) + require.NoError(t, err) + + now := time.Now() + subjectToken, err := jwt.Signed(signer). + Claims(jwt.Claims{ + Issuer: testIssuer, + Subject: delegatedUserSub, + Audience: jwt.Audience{testAudience}, + Expiry: jwt.NewNumericDate(now.Add(30 * time.Minute)), + IssuedAt: jwt.NewNumericDate(now), + }). + Claims(map[string]any{ + "client_id": agentClientID, + "name": "Delegated User", + "email": "deleg@example.com", + }). + Serialize() + require.NoError(t, err) + + // No resource/audience is sent: the server defaults to its sole allowed + // audience (testAudience), which the agent client is registered for. + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange should succeed, got %d (body: %v)", resp.StatusCode, body) + + // RFC 8693 §2.2.1 requires issued_token_type in the response. + assert.Equal(t, oauthproto.TokenTypeAccessToken, body["issued_token_type"], + "response must advertise the issued token type") + + delegated, ok := body["access_token"].(string) + require.True(t, ok, "access_token should be a string") + require.NotEmpty(t, delegated) + + // The delegated token is a JWT signed by the server; verify and inspect claims. + parsed, err := jwt.ParseSigned(delegated, []jose.SignatureAlgorithm{jose.RS256}) + require.NoError(t, err) + var claims map[string]any + require.NoError(t, parsed.Claims(ts.PrivateKey.Public(), &claims)) + + // The delegated token carries the USER as subject and the AGENT as the + // RFC 8693 §4.1 "act" (acting party). + assert.Equal(t, delegatedUserSub, claims["sub"], "delegated token subject must be the user") + assert.Equal(t, testIssuer, claims["iss"]) + act, ok := claims["act"].(map[string]any) + require.True(t, ok, "delegated token must carry an 'act' claim") + assert.Equal(t, agentClientID, act["sub"], "act.sub must identify the acting agent client") + + // Audience: no resource/audience was requested, so grantDefaultAudience must + // bind the token to the server's sole allowed audience. + aud, ok := claims["aud"].([]interface{}) + require.True(t, ok, "aud claim should be an array") + require.Len(t, aud, 1, "aud should have exactly one audience") + assert.Equal(t, testAudience, aud[0], "delegated token audience should default to the sole allowed audience") + + // Lifetime cap: the delegated token's exp must be min(subject_remaining=30m, + // delegationLifespan=15m default) ≈ 15m — NOT the subject token's 30m. This + // verifies the handler's min() cap and rules out a "subject token echoed + // back" regression, which would otherwise satisfy sub/iss/signature. + exp, ok := claims["exp"].(float64) + require.True(t, ok, "exp claim should be a number") + assert.WithinDuration(t, now.Add(15*time.Minute), time.Unix(int64(exp), 0), 2*time.Minute, + "delegated token exp must be capped at the 15m delegation lifespan, not the subject token's 30m") +} + +// TestIntegration_TokenExchange_SelfIssuedSubjectTokenScopeFromScp proves +// that a token-exchange request carrying a requested scope succeeds against +// a subject token minted by the server's own authorization_code grant, not +// just a hand-built JWT. +// +// ToolHive's own issued access tokens carry granted scopes as a "scp" array +// claim, not the RFC 9068 "scope" string claim — fosite's default JWT claims +// strategy writes "scp" unless ScopeField is explicitly String or Both, which +// this server does not set (see TestIntegration_FullPKCEFlow's own "scp" +// assertion). Every other token-exchange test in this file hand-mints its +// subject token with an explicit "scope" claim, which masks this: assignClaim +// previously only read "scope", so a genuine self-issued access token used as +// subject_token always resolved to Scopes == "", and grantScopes rejected any +// requested scope with invalid_scope. +func TestIntegration_TokenExchange_SelfIssuedSubjectTokenScopeFromScp(t *testing.T) { + t.Parallel() + + const ( + agentClientID = "test-agent-client-scp" + agentClientSecret = "test-agent-secret-scp" + ) + + // The acting agent is also the client that logs in and obtains the + // subject token: it must be confidential (token exchange requires it) and + // registered for both authorization_code (to mint a genuine access token) + // and token-exchange (to perform the exchange as itself). + agentClient, err := registration.New(registration.Config{ + ID: agentClientID, + Secret: agentClientSecret, + Public: false, + RedirectURIs: []string{testRedirectURI}, + GrantTypes: []string{"authorization_code", oauthproto.GrantTypeTokenExchange}, + Scopes: registration.DefaultScopes, + Audience: []string{testAudience}, + }) + require.NoError(t, err) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, withExtraClient(agentClient)) + + verifier := servercrypto.GeneratePKCEVerifier() + challenge := servercrypto.ComputePKCEChallenge(verifier) + + authCode, _ := completeAuthorizationFlow(t, ts.Server.URL, authorizationParams{ + ClientID: agentClientID, + RedirectURI: testRedirectURI, + State: "scp-scope-state", + Challenge: challenge, + Scope: "openid profile", + ResponseType: "code", + }) + + // Mint the subject token via the real authorization_code grant (a + // confidential client, so client_secret is required), rather than + // hand-building a JWT — this is what makes the resulting token carry + // "scp", not "scope". + tokenResp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {"authorization_code"}, + "code": {authCode}, + "redirect_uri": {testRedirectURI}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + "code_verifier": {verifier}, + }) + defer tokenResp.Body.Close() + tokenBody := parseTokenResponse(t, tokenResp) + require.Equal(t, http.StatusOK, tokenResp.StatusCode, + "authorization_code exchange should succeed, got %d (body: %v)", tokenResp.StatusCode, tokenBody) + subjectToken, ok := tokenBody["access_token"].(string) + require.True(t, ok, "access_token should be a string") + require.NotEmpty(t, subjectToken) + + // Exchange the genuine access token for a delegated token, requesting a + // scope that was granted to it ("profile"). Before the assignClaim fix, + // this fails with invalid_scope because Scopes never picks up "scp". + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + "scope": {"profile"}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange requesting a scope carried by the subject token's 'scp' claim should succeed, "+ + "got %d (body: %v)", resp.StatusCode, body) + assert.Equal(t, "profile", body["scope"], "delegated token should be granted the requested scope") +} + +// ============================================================================ +// RFC 8693 Token Exchange: Trusted External Issuer Integration Tests +// ============================================================================ + +// externalIdPKeyID is the "kid" advertised in the test-local external IdP's JWKS. +const externalIdPKeyID = "external-idp-key" + +// startExternalIdPServer starts a test-local external OIDC issuer serving +// both a discovery document and a JWKS endpoint, signed by key — deliberately +// NOT the authorization server's own key. The discovery document echoes +// r.Host as its own issuer so it stays self-consistent regardless of which +// random port httptest.NewServer binds to. +// +// The returned counter increments on every discovery-document hit, so a +// subtest configuring an explicit jwks_url (which should skip discovery +// entirely) can assert it stayed at zero, and a subtest relying on discovery +// can assert it didn't. +func startExternalIdPServer(t *testing.T, key *rsa.PrivateKey) (*httptest.Server, *atomic.Int64) { + t.Helper() + + jwks := jose.JSONWebKeySet{Keys: []jose.JSONWebKey{{ + Key: key.Public(), + KeyID: externalIdPKeyID, + Algorithm: string(jose.RS256), + Use: "sig", + }}} + + var discoveryHits atomic.Int64 + + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) { + discoveryHits.Add(1) + base := "http://" + r.Host + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "issuer": base, + "jwks_uri": base + "/jwks", + }) + }) + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(jwks) + }) + + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv, &discoveryHits +} + +// signExternalToken signs a JWT with key — the external IdP's key, distinct +// from the authorization server's own signing key — for use as a subject +// token presented during token exchange. +func signExternalToken(t *testing.T, key *rsa.PrivateKey, claims jwt.Claims, extra map[string]any) string { + t.Helper() + + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.RS256, Key: key}, + (&jose.SignerOptions{}).WithType("JWT").WithHeader("kid", externalIdPKeyID), + ) + require.NoError(t, err) + + builder := jwt.Signed(signer).Claims(claims) + if extra != nil { + builder = builder.Claims(extra) + } + raw, err := builder.Serialize() + require.NoError(t, err) + return raw +} + +// TestIntegration_TokenExchange_TrustedExternalIssuer drives RFC 8693 token +// exchange over HTTP with subject tokens from a trusted external OIDC +// issuer — proving the fail-closed consent policy (an allowlisted actor +// claim, or an authoritative may_act, or rejection) at the HTTP level, not +// just in the tokenexchange package's own unit tests. +func TestIntegration_TokenExchange_TrustedExternalIssuer(t *testing.T) { + t.Parallel() + + const ( + agentClientID = "external-agent-client" + agentClientSecret = "external-agent-secret" + allowedActor = "external-agent-azp" + externalUserSub = "external-user-sub" + ) + + // The acting agent is a confidential ToolHive client registered for the + // token-exchange grant. Its value is safe to reuse across parallel + // subtests: each subtest registers it into its own storage instance. + newAgentClient := func(t *testing.T) fosite.Client { + t.Helper() + c, err := registration.New(registration.Config{ + ID: agentClientID, + Secret: agentClientSecret, + Public: false, + GrantTypes: []string{oauthproto.GrantTypeTokenExchange}, + Scopes: registration.DefaultScopes, + Audience: []string{testAudience}, + }) + require.NoError(t, err) + return c + } + + // externalClaims returns standard claims for a subject token from the + // external IdP. The audience must equal testAudience — the server's sole + // AllowedAudience — because ensureAudienceSubsetOfSubject bounds the + // granted (default) audience by the subject token's own "aud". + externalClaims := func(issuer string) jwt.Claims { + now := time.Now() + return jwt.Claims{ + Subject: externalUserSub, + Issuer: issuer, + Audience: jwt.Audience{testAudience}, + Expiry: jwt.NewNumericDate(now.Add(30 * time.Minute)), + IssuedAt: jwt.NewNumericDate(now), + } + } + + t.Run("allowlisted actor happy path", func(t *testing.T) { + t.Parallel() + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + // JWKSURL is required whenever AllowPrivateIPs is set (see + // validateTrustedIssuers in pkg/authserver/config.go): the + // private dial target must come from operator config, not an + // OIDC discovery document. idpServer is loopback, so + // AllowPrivateIPs is unavoidable here; "explicit jwks_url + // resolution path" below is the dedicated test for the + // discovery-vs-explicit distinction this used to also cover. + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + subjectToken := signExternalToken(t, externalKey, externalClaims(idpServer.URL), map[string]any{ + "azp": allowedActor, + }) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange should succeed, got %d (body: %v)", resp.StatusCode, body) + assert.Equal(t, oauthproto.TokenTypeAccessToken, body["issued_token_type"]) + + tokenType, ok := body["token_type"].(string) + require.True(t, ok, "token_type should be a string") + assert.Equal(t, "bearer", strings.ToLower(tokenType)) + + delegated, ok := body["access_token"].(string) + require.True(t, ok, "access_token should be a string") + require.NotEmpty(t, delegated) + + parsed, err := jwt.ParseSigned(delegated, []jose.SignatureAlgorithm{jose.RS256}) + require.NoError(t, err) + var claims map[string]any + require.NoError(t, parsed.Claims(ts.PrivateKey.Public(), &claims)) + + assert.Equal(t, idpServer.URL+"#"+externalUserSub, claims["sub"], + "delegated token subject must be the external issuer's URL, qualifying the external user's "+ + "subject so it can never collide with a native ToolHive user's UUID") + assert.Equal(t, testIssuer, claims["iss"], + "delegated token must carry ToolHive's own issuer, never the external one") + + aud, ok := claims["aud"].([]interface{}) + require.True(t, ok, "aud claim should be an array") + require.Len(t, aud, 1) + assert.Equal(t, testAudience, aud[0]) + + act, ok := claims["act"].(map[string]any) + require.True(t, ok, "delegated token must carry an 'act' claim") + assert.Equal(t, agentClientID, act["sub"], "outermost act.sub must be the ToolHive acting client") + + nested, ok := act["act"].(map[string]any) + require.True(t, ok, "external delegation must nest the issuer/actor provenance record") + assert.Equal(t, allowedActor, nested["sub"], "nested act.sub is the external actor claim value") + assert.Equal(t, idpServer.URL, nested["iss"], "nested act.iss is the external issuer") + }) + + t.Run("non-allowlisted actor rejected", func(t *testing.T) { + t.Parallel() + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + const rejectedActor = "some-other-client" + subjectToken := signExternalToken(t, externalKey, externalClaims(idpServer.URL), map[string]any{ + "azp": rejectedActor, + }) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusBadRequest, resp.StatusCode, + "token exchange with a non-allowlisted actor should be a 400, got %d (body: %v)", resp.StatusCode, body) + // The validator rejects the token before checkDelegationConsent runs, so + // RFC 8693 §2.2.2's invalid_request applies — not invalid_grant. + assert.Equal(t, "invalid_request", body["error"]) + + errDesc, _ := body["error_description"].(string) + assert.Contains(t, errDesc, "invalid or could not be verified", + "the handler's fixed hint must not be replaced by a more specific — and leakier — message") + assert.NotContains(t, errDesc, rejectedActor, + "the error must not leak the rejected actor claim value") + }) + + t.Run("may_act path skips the allowlist", func(t *testing.T) { + t.Parallel() + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + JWKSURL: idpServer.URL + "/jwks", + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + // AllowedActors deliberately empty: may_act must be honored + // without any actor being allowlisted. + AllowedDelegateClients: []string{"*"}, + }}), + ) + + // No "azp" claim at all — only may_act, naming the ToolHive agent + // client directly as the party authorized to act. + subjectToken := signExternalToken(t, externalKey, externalClaims(idpServer.URL), map[string]any{ + "may_act": map[string]any{"sub": agentClientID, "iss": testIssuer}, + }) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange via may_act should succeed, got %d (body: %v)", resp.StatusCode, body) + + delegated, ok := body["access_token"].(string) + require.True(t, ok) + require.NotEmpty(t, delegated) + + parsed, err := jwt.ParseSigned(delegated, []jose.SignatureAlgorithm{jose.RS256}) + require.NoError(t, err) + var claims map[string]any + require.NoError(t, parsed.Claims(ts.PrivateKey.Public(), &claims)) + + act, ok := claims["act"].(map[string]any) + require.True(t, ok) + assert.Equal(t, agentClientID, act["sub"]) + + // may_act carries no ExternalActor (see ValidatedClaims.ExternalActor's + // doc comment), but the external issuer must still be recorded — this + // is the path that bypasses the allowlist entirely, so it needs the + // audit trail at least as much as the allowlist path does. + nested, ok := act["act"].(map[string]any) + require.True(t, ok, "external issuer must still be nested even without an allowlisted actor") + assert.Equal(t, idpServer.URL, nested["iss"]) + _, hasSub := nested["sub"] + assert.False(t, hasSub, "no client-namespace actor claim exists to report on the may_act path") + }) + + t.Run("allowed delegate clients binds the allowlisted actor to a specific ToolHive client", func(t *testing.T) { + t.Parallel() + + const ( + otherAgentClientID = "other-external-agent-client" + otherAgentClientSecret = "other-external-agent-secret" + ) + + // A second confidential client, also registered for the token-exchange + // grant and otherwise identical to the primary agent client — the only + // difference the test exercises is that it is NOT in this issuer's + // AllowedDelegateClients. Without that field (or if the binding check + // were removed), this client would succeed exactly like the primary + // one, since both hold the grant and both present the same + // allowlisted external actor claim. + otherAgentClient, err := registration.New(registration.Config{ + ID: otherAgentClientID, + Secret: otherAgentClientSecret, + Public: false, + GrantTypes: []string{oauthproto.GrantTypeTokenExchange}, + Scopes: registration.DefaultScopes, + Audience: []string{testAudience}, + }) + require.NoError(t, err) + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withExtraClient(otherAgentClient), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + AllowedDelegateClients: []string{agentClientID}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + }}), + ) + + subjectToken := signExternalToken(t, externalKey, externalClaims(idpServer.URL), map[string]any{ + "azp": allowedActor, + }) + + exchange := func(clientID, clientSecret string) *http.Response { + return makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {clientID}, + "client_secret": {clientSecret}, + }) + } + + permittedResp := exchange(agentClientID, agentClientSecret) + defer permittedResp.Body.Close() + permittedBody := parseTokenResponse(t, permittedResp) + require.Equal(t, http.StatusOK, permittedResp.StatusCode, + "the allowlisted delegate client should succeed, got %d (body: %v)", + permittedResp.StatusCode, permittedBody) + + rejectedResp := exchange(otherAgentClientID, otherAgentClientSecret) + defer rejectedResp.Body.Close() + rejectedBody := parseTokenResponse(t, rejectedResp) + require.Equal(t, http.StatusBadRequest, rejectedResp.StatusCode, + "a client absent from AllowedDelegateClients should be a 400, got %d (body: %v)", + rejectedResp.StatusCode, rejectedBody) + assert.Equal(t, "invalid_grant", rejectedBody["error"]) + errDesc, _ := rejectedBody["error_description"].(string) + assert.Contains(t, errDesc, "not authorized to exchange subject tokens") + }) + + t.Run("untrusted issuer rejected before any JWKS fetch", func(t *testing.T) { + t.Parallel() + + const untrustedIssuer = "https://untrusted-issuer.example.com" + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, discoveryHits := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + // The sole registered trusted issuer is idpServer.URL — a live, + // fetchable JWKS signed with the SAME key the subject token below + // uses. This makes the rejection discriminating rather than + // coincidental: the subject token's signature WOULD verify + // successfully against this real JWKS if the issuer-map lookup were + // ever bypassed (a fallback to the sole configured issuer, a wildcard + // match, or deleted "iss"-based routing) — so a 400 here can only be + // explained by the "iss" string itself failing to match idpServer.URL, + // not by an unverifiable signature or a validator that was never + // constructed in the first place. + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + // Signed with idpServer's real key, but claims an issuer that was + // never registered in TrustedIssuers. + subjectToken := signExternalToken(t, externalKey, externalClaims(untrustedIssuer), map[string]any{ + "azp": allowedActor, + }) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusBadRequest, resp.StatusCode, + "token exchange from an untrusted issuer should be a 400, got %d (body: %v)", resp.StatusCode, body) + assert.Equal(t, "invalid_request", body["error"]) + + errDesc, _ := body["error_description"].(string) + assert.Contains(t, errDesc, "invalid or could not be verified", + "the handler's fixed hint must not be replaced by a more specific — and leakier — message") + assert.NotContains(t, errDesc, untrustedIssuer, + "the error must not leak the untrusted issuer URL") + + // The issuer-map miss must short-circuit before any JWKS fetch. + // Note: JWKSURL is now preconfigured above (see comment on the + // "allowlisted actor happy path" subtest), which already skips + // discovery on its own — so this assertion holding is necessary but + // not on its own sufficient proof of the short-circuit; the 400 + // response and error content below are the discriminating checks. + assert.Zero(t, discoveryHits.Load(), "issuer-map miss must precede any JWKS discovery fetch") + }) + + t.Run("self-issued subject token still works alongside trusted issuers", func(t *testing.T) { + t.Parallel() + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.RS256, Key: ts.PrivateKey}, + (&jose.SignerOptions{}).WithType("JWT").WithHeader("kid", "test-key"), + ) + require.NoError(t, err) + + now := time.Now() + subjectToken, err := jwt.Signed(signer). + Claims(jwt.Claims{ + Issuer: testIssuer, + Subject: "self-issued-delegated-user", + Audience: jwt.Audience{testAudience}, + Expiry: jwt.NewNumericDate(now.Add(30 * time.Minute)), + IssuedAt: jwt.NewNumericDate(now), + }). + Claims(map[string]any{"client_id": agentClientID}). + Serialize() + require.NoError(t, err) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "self-issued token exchange should still succeed with trusted issuers configured, "+ + "got %d (body: %v)", resp.StatusCode, body) + assert.Equal(t, oauthproto.TokenTypeAccessToken, body["issued_token_type"]) + }) + + t.Run("explicit jwks_url resolution path", func(t *testing.T) { + t.Parallel() + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, discoveryHits := startExternalIdPServer(t, externalKey) + + // Prove the counter actually increments before relying on its + // zero-ness below — otherwise a discovery handler that silently + // stopped counting would make the "must skip discovery" assertion + // vacuously true. + discResp, err := http.Get(idpServer.URL + "/.well-known/openid-configuration") //nolint:noctx + require.NoError(t, err) + discResp.Body.Close() + require.Equal(t, int64(1), discoveryHits.Load(), "discovery counter must be live") + discoveryHits.Store(0) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + // Set directly rather than relying on discovery: this exercises + // ensureRegistered's pre-configured-URL branch instead of + // discoverJWKSURL. + JWKSURL: idpServer.URL + "/jwks", + ExpectedAudience: testAudience, + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + subjectToken := signExternalToken(t, externalKey, externalClaims(idpServer.URL), map[string]any{ + "azp": allowedActor, + }) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange with an explicit jwks_url should succeed, got %d (body: %v)", resp.StatusCode, body) + assert.Equal(t, oauthproto.TokenTypeAccessToken, body["issued_token_type"]) + + assert.Zero(t, discoveryHits.Load(), + "a preconfigured jwks_url must skip OIDC discovery entirely") + }) + + t.Run("external token's exp bounds the delegated token's lifetime", func(t *testing.T) { + t.Parallel() + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: testAudience, + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + // 5m remaining lifetime, well under the 15m default delegation + // lifespan — so the delegated token's exp must track the subject + // token, not the (longer) configured cap. + now := time.Now() + subjClaims := jwt.Claims{ + Subject: externalUserSub, + Issuer: idpServer.URL, + Audience: jwt.Audience{testAudience}, + Expiry: jwt.NewNumericDate(now.Add(5 * time.Minute)), + IssuedAt: jwt.NewNumericDate(now), + } + subjectToken := signExternalToken(t, externalKey, subjClaims, map[string]any{"azp": allowedActor}) + + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusOK, resp.StatusCode, + "token exchange should succeed, got %d (body: %v)", resp.StatusCode, body) + + delegated, ok := body["access_token"].(string) + require.True(t, ok) + require.NotEmpty(t, delegated) + + parsed, err := jwt.ParseSigned(delegated, []jose.SignatureAlgorithm{jose.RS256}) + require.NoError(t, err) + var claims map[string]any + require.NoError(t, parsed.Claims(ts.PrivateKey.Public(), &claims)) + + exp, ok := claims["exp"].(float64) + require.True(t, ok, "exp claim should be a number") + assert.WithinDuration(t, now.Add(5*time.Minute), time.Unix(int64(exp), 0), 2*time.Minute, + "delegated token exp must track the external subject token's 5m remaining lifetime, "+ + "not the 15m default delegation lifespan") + }) + + t.Run("invalid_target when the external aud isn't a ToolHive-allowed audience", func(t *testing.T) { + t.Parallel() + + // A realistic Entra-style app-ID audience: legitimate as the trusted + // issuer's own ExpectedAudience, but not one of ToolHive's + // AllowedAudiences — the footgun documented on RunConfig.TrustedIssuers. + const foreignAudience = "api://some-app-id" + + externalKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + idpServer, _ := startExternalIdPServer(t, externalKey) + + m := startMockOIDC(t) + ts := setupTestServerWithMockOIDC(t, m, + withExtraClient(newAgentClient(t)), + withTrustedIssuers([]tokenexchange.TrustedIssuer{{ + IssuerURL: idpServer.URL, + ExpectedAudience: foreignAudience, + JWKSURL: idpServer.URL + "/jwks", + AllowedActors: []string{allowedActor}, + InsecureAllowHTTP: true, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{"*"}, + }}), + ) + + now := time.Now() + subjClaims := jwt.Claims{ + Subject: externalUserSub, + Issuer: idpServer.URL, + Audience: jwt.Audience{foreignAudience}, + Expiry: jwt.NewNumericDate(now.Add(30 * time.Minute)), + IssuedAt: jwt.NewNumericDate(now), + } + subjectToken := signExternalToken(t, externalKey, subjClaims, map[string]any{"azp": allowedActor}) + + // No resource/audience requested: the handler defaults to ToolHive's + // sole AllowedAudience (testAudience), which this subject token's aud + // does not cover. + resp := makeTokenRequest(t, ts.Server.URL, url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + "client_id": {agentClientID}, + "client_secret": {agentClientSecret}, + }) + defer resp.Body.Close() + + body := parseTokenResponse(t, resp) + require.Equal(t, http.StatusBadRequest, resp.StatusCode, + "an external aud outside ToolHive's AllowedAudiences must be rejected, got %d (body: %v)", + resp.StatusCode, body) + assert.Equal(t, "invalid_target", body["error"]) + errDesc, _ := body["error_description"].(string) + assert.Contains(t, errDesc, "not covered by the subject token") + }) +} + // ============================================================================ // Full PKCE Flow Integration Tests with Mock Upstream IDP (using mockoidc) // ============================================================================ diff --git a/pkg/authserver/runner/dcr_adapter.go b/pkg/authserver/runner/dcr_adapter.go index 37dfd54507..4682129d92 100644 --- a/pkg/authserver/runner/dcr_adapter.go +++ b/pkg/authserver/runner/dcr_adapter.go @@ -79,6 +79,10 @@ func newDCRRequest(rc *authserver.OAuth2UpstreamRunConfig, localIssuer string) ( AuthorizationEndpoint: rc.AuthorizationEndpoint, TokenEndpoint: rc.TokenEndpoint, InitialAccessToken: initialAccessToken, + // Reuse the upstream's private-IP policy so the DCR discovery and + // registration calls share the same SSRF posture as its token and + // userinfo calls (see upstream.OAuth2Config.AllowPrivateIPs). + AllowPrivateIPs: rc.AllowPrivateIPs, }, nil } @@ -143,28 +147,34 @@ func consumeResolution(rc authserver.OAuth2UpstreamRunConfig, res *dcr.Resolutio } // applyResolutionToOAuth2Config returns a copy of cfg with the DCR- -// resolved ClientSecret overlaid onto it. This is the companion to -// consumeResolution: where that function writes fields representable in -// the file-or-env run-config model, this one writes the inline-only -// ClientSecret directly on the runtime config. +// resolved ClientSecret and TokenEndpointAuthMethod overlaid onto it. This +// is the companion to consumeResolution: where that function writes fields +// representable in the file-or-env run-config model, this one writes the +// inline-only fields directly on the runtime config. // // cfg is taken by value and the modified copy is returned, mirroring // consumeResolution. The no-mutation contract is compile-time enforced // by the signature rather than a prose discipline. // +// TokenEndpointAuthMethod carries the method the upstream negotiated during +// registration (RFC 7591); see upstream.authStyleFromMethod for why it matters. +// // The split exists because buildPureOAuth2Config intentionally retains a // narrow file-or-env contract (no DCR awareness) and because OAuth2's // ClientSecret on the run-config is modelled as a reference rather than // an inline string. Any future output path from OAuth2UpstreamRunConfig // to upstream.OAuth2Config must call BOTH consumeResolution (run-config // side) AND applyResolutionToOAuth2Config (built-config side) to get a -// fully-resolved DCR client. Forgetting the second call leaves -// ClientSecret empty and produces silent auth failures at request time — -// the type system does not enforce the pair, so the invariant lives here. +// fully-resolved DCR client. Forgetting the second call leaves both +// ClientSecret and TokenEndpointAuthMethod empty, producing silent auth +// failures at request time — a client that negotiated client_secret_basic +// reverts to POST-body credentials and is rejected. The type system does +// not enforce the pair, so the invariant lives here. func applyResolutionToOAuth2Config(cfg upstream.OAuth2Config, res *dcr.Resolution) upstream.OAuth2Config { if res == nil { return cfg } cfg.ClientSecret = res.ClientSecret + cfg.TokenEndpointAuthMethod = res.TokenEndpointAuthMethod return cfg } diff --git a/pkg/authserver/runner/dcr_adapter_test.go b/pkg/authserver/runner/dcr_adapter_test.go index 3a4e7815c1..aa3bc5dbca 100644 --- a/pkg/authserver/runner/dcr_adapter_test.go +++ b/pkg/authserver/runner/dcr_adapter_test.go @@ -174,11 +174,16 @@ func TestApplyResolutionToOAuth2Config(t *testing.T) { cfg := upstream.OAuth2Config{} res := &dcr.Resolution{ - ClientSecret: "dcr-issued-secret", + ClientSecret: "dcr-issued-secret", + TokenEndpointAuthMethod: "client_secret_basic", } out := applyResolutionToOAuth2Config(cfg, res) assert.Equal(t, "dcr-issued-secret", out.ClientSecret) + // The negotiated auth method must be carried through so newBaseOAuth2Provider + // presents credentials the way the upstream registered the client (issue #5865). + assert.Equal(t, "client_secret_basic", out.TokenEndpointAuthMethod) assert.Equal(t, "", cfg.ClientSecret, "caller's cfg must not be mutated") + assert.Equal(t, "", cfg.TokenEndpointAuthMethod, "caller's cfg must not be mutated") // nil resolution is a no-op rather than a crash. unchanged := applyResolutionToOAuth2Config(cfg, nil) @@ -247,6 +252,7 @@ func TestNewDCRRequest(t *testing.T) { wantInitialAccessToken string wantDiscoveryURL string wantRegistration string + wantAllowPrivateIPs bool }{ { name: "discovery_url branch resolves file-based initial access token", @@ -274,6 +280,24 @@ func TestNewDCRRequest(t *testing.T) { wantIssuer: "https://thv.example.com", wantRegistration: "https://idp.example.com/register", }, + { + // Pins the AllowPrivateIPs one-line copy from + // OAuth2UpstreamRunConfig onto dcr.Request: without this test, a + // future refactor could silently drop or invert the SSRF-guard + // decision without any test failing. + name: "AllowPrivateIPs propagates from run-config", + rc: &authserver.OAuth2UpstreamRunConfig{ + Scopes: []string{"openid"}, + DCRConfig: &authserver.DCRUpstreamConfig{ + RegistrationEndpoint: "https://idp.example.com/register", + }, + AllowPrivateIPs: true, + }, + localIssuer: "https://thv.example.com", + wantIssuer: "https://thv.example.com", + wantRegistration: "https://idp.example.com/register", + wantAllowPrivateIPs: true, + }, { name: "nil run-config rejected", rc: nil, @@ -302,6 +326,7 @@ func TestNewDCRRequest(t *testing.T) { assert.Equal(t, tc.wantInitialAccessToken, req.InitialAccessToken) assert.Equal(t, tc.wantDiscoveryURL, req.DiscoveryURL) assert.Equal(t, tc.wantRegistration, req.RegistrationEndpoint) + assert.Equal(t, tc.wantAllowPrivateIPs, req.AllowPrivateIPs) }) } } diff --git a/pkg/authserver/runner/embeddedauthserver.go b/pkg/authserver/runner/embeddedauthserver.go index 70102ee937..f765b5a936 100644 --- a/pkg/authserver/runner/embeddedauthserver.go +++ b/pkg/authserver/runner/embeddedauthserver.go @@ -204,7 +204,16 @@ func NewEmbeddedAuthServerWithStorage( return nil, fmt.Errorf("failed to build upstream configs: %w", err) } - // 6. Build the resolved Config. + // 6. Parse delegation token lifespan if configured. + var delegationLifespan time.Duration + if cfg.DelegationTokenLifespan != "" { + delegationLifespan, err = time.ParseDuration(cfg.DelegationTokenLifespan) + if err != nil { + return nil, fmt.Errorf("invalid delegation token lifespan: %w", err) + } + } + + // 7. Build the resolved Config. // // Defensive copies of the scope/audience slices: cfg is operator-supplied // input that may be retained or mutated by the caller (e.g. tests, a @@ -223,6 +232,7 @@ func NewEmbeddedAuthServerWithStorage( AccessTokenLifespan: accessLifespan, RefreshTokenLifespan: refreshLifespan, AuthCodeLifespan: authCodeLifespan, + DelegationTokenLifespan: delegationLifespan, Upstreams: upstreams, ScopesSupported: slices.Clone(cfg.ScopesSupported), BaselineClientScopes: slices.Clone(cfg.BaselineClientScopes), @@ -231,9 +241,14 @@ func NewEmbeddedAuthServerWithStorage( CIMDCacheMaxSize: cimdCacheMaxSize, CIMDCacheFallbackTTL: cimdCacheFallbackTTL, InsecureAllowHTTP: cfg.InsecureAllowHTTP, + // slices.Clone is shallow: each TrustedIssuer's own AllowedActors + // slice is still shared with cfg. NewMultiIssuerTokenValidator's + // constructor clones AllowedActors per issuer before use, so the + // authorization-critical data is protected without a deep copy here. + TrustedIssuers: slices.Clone(cfg.TrustedIssuers), } - // 7. Create the auth server. authserver.New also asserts the DCR + // 8. Create the auth server. authserver.New also asserts the DCR // capability internally so its DCRStore() accessor returns the same // asserted handle this constructor used for buildUpstreamConfigs. server, err := authserver.New(ctx, resolvedCfg, stor) @@ -368,8 +383,12 @@ func loadHMACSecrets(files []string) (*servercrypto.HMACSecrets, error) { return nil, fmt.Errorf("failed to read HMAC secret from %s: %w", files[0], err) } - // Trim whitespace (Kubernetes Secret mounts may include trailing newlines) - current = bytes.TrimSpace(current) + // HMAC secrets are raw binary (see MCPExternalAuthConfig.HMACSecretRefs doc) — do not + // trim; any byte, including whitespace-valued ones, may be genuine key material. + if len(current) < servercrypto.MinSecretLength { + return nil, fmt.Errorf("HMAC secret in %s is %d bytes, but must be at least %d bytes", + files[0], len(current), servercrypto.MinSecretLength) + } secrets := &servercrypto.HMACSecrets{ Current: current, @@ -385,7 +404,11 @@ func loadHMACSecrets(files []string) (*servercrypto.HMACSecrets, error) { if err != nil { return nil, fmt.Errorf("failed to read rotated HMAC secret from %s: %w", file, err) } - secrets.Rotated = append(secrets.Rotated, bytes.TrimSpace(secret)) + if len(secret) < servercrypto.MinSecretLength { + return nil, fmt.Errorf("rotated HMAC secret in %s is %d bytes, but must be at least %d bytes", + file, len(secret), servercrypto.MinSecretLength) + } + secrets.Rotated = append(secrets.Rotated, secret) } return secrets, nil diff --git a/pkg/authserver/runner/embeddedauthserver_test.go b/pkg/authserver/runner/embeddedauthserver_test.go index a95b62679c..89d09184f9 100644 --- a/pkg/authserver/runner/embeddedauthserver_test.go +++ b/pkg/authserver/runner/embeddedauthserver_test.go @@ -19,6 +19,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "sync" "sync/atomic" "testing" "time" @@ -29,6 +30,7 @@ import ( "github.com/stacklok/toolhive/pkg/authserver" servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto" "github.com/stacklok/toolhive/pkg/authserver/server/keys" + "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange" "github.com/stacklok/toolhive/pkg/authserver/storage" "github.com/stacklok/toolhive/pkg/oauthproto" ) @@ -167,12 +169,12 @@ func TestLoadHMACSecrets(t *testing.T) { assert.Equal(t, []byte(rotatedSecret), secrets.Rotated[0]) }) - t.Run("trims whitespace from secrets", func(t *testing.T) { + t.Run("preserves leading and trailing whitespace bytes in secrets", func(t *testing.T) { t.Parallel() tmpDir := t.TempDir() secretFile := filepath.Join(tmpDir, "hmac-secret") - secretValue := " secret-with-whitespace \n" + secretValue := " secret-with-whitespace-that-is-at-least-32-bytes \n" err := os.WriteFile(secretFile, []byte(secretValue), 0600) require.NoError(t, err) @@ -181,7 +183,88 @@ func TestLoadHMACSecrets(t *testing.T) { require.NoError(t, err) require.NotNil(t, secrets) - assert.Equal(t, []byte("secret-with-whitespace"), secrets.Current) + assert.Equal(t, []byte(secretValue), secrets.Current) + }) + + t.Run("does not corrupt binary secrets with whitespace-valued edge bytes", func(t *testing.T) { + t.Parallel() + + tmpDir := t.TempDir() + secretFile := filepath.Join(tmpDir, "hmac-secret") + + original := make([]byte, 32) + for i := range original { + original[i] = 0xAB + } + original[0] = 0x20 + original[31] = 0x0A + + require.NoError(t, os.WriteFile(secretFile, original, 0600)) + + secrets, err := loadHMACSecrets([]string{secretFile}) + require.NoError(t, err) + require.NotNil(t, secrets) + + require.Len(t, secrets.Current, 32) + assert.Equal(t, original, secrets.Current) + }) + + t.Run("does not corrupt rotated binary secrets with whitespace-valued edge bytes", func(t *testing.T) { + t.Parallel() + + tmpDir := t.TempDir() + currentFile := filepath.Join(tmpDir, "hmac-current") + rotatedFile := filepath.Join(tmpDir, "hmac-rotated") + + require.NoError(t, os.WriteFile(currentFile, []byte("current-secret-32-bytes-minimum!"), 0600)) + + original := make([]byte, 32) + for i := range original { + original[i] = 0xAB + } + original[0] = 0x20 + original[31] = 0x0A + + require.NoError(t, os.WriteFile(rotatedFile, original, 0600)) + + secrets, err := loadHMACSecrets([]string{currentFile, rotatedFile}) + require.NoError(t, err) + require.NotNil(t, secrets) + + require.Len(t, secrets.Rotated, 1) + assert.Equal(t, original, secrets.Rotated[0]) + }) + + t.Run("short current secret returns error naming file and byte counts", func(t *testing.T) { + t.Parallel() + + tmpDir := t.TempDir() + secretFile := filepath.Join(tmpDir, "hmac-secret") + + require.NoError(t, os.WriteFile(secretFile, make([]byte, 31), 0600)) + + _, err := loadHMACSecrets([]string{secretFile}) + require.Error(t, err) + assert.Contains(t, err.Error(), secretFile) + assert.Contains(t, err.Error(), "31") + assert.Contains(t, err.Error(), "32") + }) + + t.Run("short rotated secret returns error naming file and byte counts", func(t *testing.T) { + t.Parallel() + + tmpDir := t.TempDir() + currentFile := filepath.Join(tmpDir, "hmac-current") + rotatedFile := filepath.Join(tmpDir, "hmac-rotated") + + require.NoError(t, os.WriteFile(currentFile, []byte("current-secret-32-bytes-minimum!"), 0600)) + require.NoError(t, os.WriteFile(rotatedFile, make([]byte, 31), 0600)) + + _, err := loadHMACSecrets([]string{currentFile, rotatedFile}) + require.Error(t, err) + assert.Contains(t, err.Error(), rotatedFile) + assert.Contains(t, err.Error(), "31") + assert.Contains(t, err.Error(), "32") }) t.Run("skips empty paths in rotated files", func(t *testing.T) { @@ -1668,6 +1751,7 @@ func TestBuildUpstreamConfigs_DCR(t *testing.T) { redirectURI := server.URL + "/oauth/callback" key := storage.DCRKey{ Issuer: server.URL, + UpstreamID: server.URL, RedirectURI: redirectURI, ScopesHash: storage.ScopesHash([]string{"openid", "profile"}), } @@ -1809,6 +1893,7 @@ func TestNewEmbeddedAuthServer_DCRBoot(t *testing.T) { redirectURI := server.URL + "/oauth/callback" key := storage.DCRKey{ Issuer: server.URL, + UpstreamID: server.URL, RedirectURI: redirectURI, ScopesHash: storage.ScopesHash([]string{"openid", "profile"}), } @@ -1989,6 +2074,7 @@ func TestEmbeddedAuthServer_DCRStorePersistsAcrossClose(t *testing.T) { redirectURI := server.URL + "/oauth/callback" key := storage.DCRKey{ Issuer: server.URL, + UpstreamID: server.URL, RedirectURI: redirectURI, ScopesHash: storage.ScopesHash([]string{"openid", "profile"}), } @@ -2027,6 +2113,32 @@ func (s *urlErrorOnCloseStorage) Close() error { return s.closeErr } +// syncBuffer is a concurrency-safe io.Writer over a bytes.Buffer, used by the +// slog-capturing tests below. slog.SetDefault is process-global, so while a +// capturing handler is installed, any goroutine in the test binary — not just +// the capturing test's own subtests — can write a record into it; a plain +// bytes.Buffer is a data race waiting for some concurrently running test to +// log, which is exactly what it became once a parallel parent was added. +// Mirrors the identically-named helper in pkg/authz/authorizers/cedar/ +// core_test.go; kept as a separate copy since it is an unexported test type +// and not worth exporting across packages. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + // TestNewEmbeddedAuthServer_DeferredCleanupSanitizesLog pins the post-#5196 // invariant that the deferred-cleanup slog.Warn at the top of // NewEmbeddedAuthServerWithStorage routes both closeErr and retErr through @@ -2042,11 +2154,11 @@ func (s *urlErrorOnCloseStorage) Close() error { // - the captured log record DOES contain the host components, so // operators retain enough context to correlate the failure. // -// NOT t.Parallel(): the test swaps slog.Default() to capture output and -// restores it via t.Cleanup. Running in parallel would race with any other -// test in this package that emits a log record. Confirmed against the -// paralleltest rule on a sample run — every other test failed with a -// data-race report on slog's internal default-logger handle. +// NOT t.Parallel(): the test swaps slog.Default() and restores it via +// t.Cleanup. The capturing buffer is a syncBuffer, so a concurrent writer is +// not itself a data race — the reason to stay sequential is the swap/restore +// pair, which two overlapping capturing tests would interleave, leaving +// whichever restored last as the default and clobbering the other's capture. // //nolint:paralleltest // see comment above; mutates the package-global slog.Default() func TestNewEmbeddedAuthServer_DeferredCleanupSanitizesLog(t *testing.T) { @@ -2099,7 +2211,7 @@ func TestNewEmbeddedAuthServer_DeferredCleanupSanitizesLog(t *testing.T) { // Capture slog output by swapping the default logger for the duration // of this test. Restore on cleanup so parallel tests are unaffected. - var buf bytes.Buffer + var buf syncBuffer prev := slog.Default() slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) t.Cleanup(func() { slog.SetDefault(prev) }) @@ -2125,6 +2237,104 @@ func TestNewEmbeddedAuthServer_DeferredCleanupSanitizesLog(t *testing.T) { "closeErr host must remain in the Warn record after sanitisation") } +// TestNewEmbeddedAuthServer_TrustedIssuers pins the RunConfig.TrustedIssuers +// -> Config.TrustedIssuers conversion (embeddedauthserver.go's resolvedCfg, +// ~line 244). "Construction succeeds" alone doesn't discriminate: deleting +// the resolvedCfg.TrustedIssuers assignment entirely still builds a server +// (it degrades to the empty-TrustedIssuers case, which is itself valid), so +// the no-trusted-issuers control subtest can't disagree with a broken one. +// +// Instead this uses an observable side effect that only fires if the field +// actually reached the Factory closure and NewMultiIssuerTokenValidator's +// constructor (not merely RunConfig.Validate/Config.Validate, both of which +// run before resolvedCfg is built): NewMultiIssuerTokenValidator logs a +// slog.Warn ("Trusted issuer has no allowed actors configured") for each +// TrustedIssuer with an empty AllowedActors. Confirmed by mutation: deleting +// `TrustedIssuers: slices.Clone(cfg.TrustedIssuers)` from resolvedCfg's +// construction makes this test fail (the warning never fires because +// MultiIssuerTokenValidator is never even built — Factory falls back to the +// self-issued validator). +// +// What this does NOT cover: resolvedCfg has no exported accessor, so the +// mutate-the-caller's-slice-after-construction claim (slices.Clone protects +// the outer []TrustedIssuer from a later append/removal on cfg, but is +// shallow — a mutation reaching through a still-shared AllowedActors slice +// would not be caught by this outer clone) is not observable through this +// package's public API. Proving it would require a live token-exchange HTTP +// round trip against an external issuer (Step 7's HTTP-level integration +// tests), or a production-code accessor this task's scope excludes. +// NewMultiIssuerTokenValidator's own AllowedActors clone (see +// multi_issuer_validator.go) closes that gap at the layer where it actually +// matters for concurrent request handling. +func TestNewEmbeddedAuthServer_TrustedIssuers(t *testing.T) { + t.Parallel() + + base := func() *authserver.RunConfig { + return &authserver.RunConfig{ + SchemaVersion: authserver.CurrentSchemaVersion, + Issuer: "https://auth.example.com", + Upstreams: []authserver.UpstreamRunConfig{ + { + Name: "test-upstream", + Type: authserver.UpstreamProviderTypeOAuth2, + OAuth2Config: &authserver.OAuth2UpstreamRunConfig{ + AuthorizationEndpoint: "https://example.com/authorize", + TokenEndpoint: "https://example.com/token", + ClientID: "test-client-id", + RedirectURI: "https://auth.example.com/oauth/callback", + }, + }, + }, + AllowedAudiences: []string{"https://mcp.example.com"}, + } + } + + t.Run("no trusted issuers builds server (control case)", func(t *testing.T) { + t.Parallel() + + cfg := base() + + srv, err := NewEmbeddedAuthServer(context.Background(), cfg) + require.NoError(t, err) + require.NotNil(t, srv) + require.NoError(t, srv.Close()) + }) + + // NOT t.Parallel(): captures the package-global slog.Default() via + // slog.SetDefault, which is atomic, and restores it via t.Cleanup. + // The buffer itself is the only unsynchronized object in that swap, and + // it's mutex-guarded (syncBuffer), so records emitted by any other test + // running concurrently in this parallel phase land harmlessly. That's + // fine here because the assertion below is Contains-only, which + // tolerates foreign records mixed into the buffer. + //nolint:paralleltest // see comment above + t.Run("valid trusted issuers reach the Factory closure and validator constructor", func(t *testing.T) { + cfg := base() + cfg.TrustedIssuers = []tokenexchange.TrustedIssuer{ + { + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + // AllowedActors intentionally empty: this is what makes + // NewMultiIssuerTokenValidator log its warning, the + // observable proof this test relies on. + AllowedDelegateClients: []string{"*"}, + }, + } + + var buf syncBuffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + srv, err := NewEmbeddedAuthServer(context.Background(), cfg) + require.NoError(t, err) + require.NotNil(t, srv) + t.Cleanup(func() { _ = srv.Close() }) + + assert.Contains(t, buf.String(), "Trusted issuer has no allowed actors configured") + }) +} + func TestResolveCIMDConfig(t *testing.T) { t.Parallel() diff --git a/pkg/authserver/server/crypto/keys.go b/pkg/authserver/server/crypto/keys.go index 111a2678ef..0493b48d62 100644 --- a/pkg/authserver/server/crypto/keys.go +++ b/pkg/authserver/server/crypto/keys.go @@ -26,7 +26,6 @@ import ( "encoding/pem" "fmt" "os" - "strings" "github.com/go-jose/go-jose/v4" ) @@ -227,62 +226,3 @@ func DeriveSigningKeyParams(key crypto.Signer, keyID, algorithm string) (*Signin return params, nil } - -// loadHMACSecretFile loads an HMAC secret from a file (internal helper). -// Returns nil if path is empty (triggers random generation in toInternalConfig). -// The secret must be at least 32 bytes after trimming whitespace. -func loadHMACSecretFile(secretPath string) ([]byte, error) { - if secretPath == "" { - return nil, nil - } - - data, err := os.ReadFile(secretPath) // #nosec G304 - secretPath is provided by user via CLI flag or config - if err != nil { - return nil, fmt.Errorf("failed to read HMAC secret file: %w", err) - } - - // Trim whitespace (common in Kubernetes Secret mounts which often add trailing newlines) - secret := []byte(strings.TrimSpace(string(data))) - - if len(secret) < MinSecretLength { - return nil, fmt.Errorf("HMAC secret must be at least %d bytes", MinSecretLength) - } - - return secret, nil -} - -// LoadHMACSecrets loads HMAC secrets from file paths for rotation support. -// paths[0] is the current (signing) secret; paths[1:] are rotated (verification) secrets. -// Returns nil if paths is empty (caller should generate random secret). -func LoadHMACSecrets(paths []string) (*HMACSecrets, error) { - if len(paths) == 0 { - return nil, nil - } - - // Load current secret (required, cannot be empty path) - if paths[0] == "" { - return nil, fmt.Errorf("current HMAC secret path cannot be empty") - } - current, err := loadHMACSecretFile(paths[0]) - if err != nil { - return nil, fmt.Errorf("failed to load current HMAC secret: %w", err) - } - - // Load rotated secrets (optional, skip empty paths) - var rotated [][]byte - for i, path := range paths[1:] { - if path == "" { - continue - } - secret, err := loadHMACSecretFile(path) - if err != nil { - return nil, fmt.Errorf("failed to load rotated HMAC secret [%d]: %w", i+1, err) - } - rotated = append(rotated, secret) - } - - return &HMACSecrets{ - Current: current, - Rotated: rotated, - }, nil -} diff --git a/pkg/authserver/server/crypto/keys_test.go b/pkg/authserver/server/crypto/keys_test.go index 358fd82774..59949d6722 100644 --- a/pkg/authserver/server/crypto/keys_test.go +++ b/pkg/authserver/server/crypto/keys_test.go @@ -25,7 +25,6 @@ import ( "encoding/pem" "os" "path/filepath" - "strings" "testing" "github.com/stretchr/testify/assert" @@ -286,137 +285,6 @@ func TestDeriveKeyID(t *testing.T) { assert.NotEqual(t, id1, id3, "different keys should produce different IDs") } -func TestLoadHMACSecrets(t *testing.T) { - t.Parallel() - - validSecret := strings.Repeat("a", 32) - validSecret2 := strings.Repeat("b", 32) - tooShortSecret := strings.Repeat("a", 31) - - tests := []struct { - name string - setup func(t *testing.T, dir string) []string - wantCurrent []byte - wantRotated [][]byte - wantErr string - }{ - { - name: "empty paths", - setup: func(_ *testing.T, _ string) []string { return []string{} }, - wantCurrent: nil, - wantRotated: nil, - }, - { - name: "single secret", - setup: func(_ *testing.T, dir string) []string { - return []string{writeFileNamed(t, dir, "current", validSecret)} - }, - wantCurrent: []byte(validSecret), - wantRotated: nil, - }, - { - name: "with rotated secrets", - setup: func(_ *testing.T, dir string) []string { - return []string{ - writeFileNamed(t, dir, "current", validSecret), - writeFileNamed(t, dir, "rotated1", validSecret2), - } - }, - wantCurrent: []byte(validSecret), - wantRotated: [][]byte{[]byte(validSecret2)}, - }, - { - name: "empty current path", - setup: func(_ *testing.T, _ string) []string { - return []string{""} - }, - wantErr: "current HMAC secret path cannot be empty", - }, - { - name: "invalid current secret file", - setup: func(_ *testing.T, _ string) []string { - return []string{"/nonexistent/secret"} - }, - wantErr: "failed to load current", - }, - { - name: "invalid rotated secret", - setup: func(_ *testing.T, dir string) []string { - return []string{ - writeFileNamed(t, dir, "current", validSecret), - "/nonexistent/rotated", - } - }, - wantErr: "failed to load rotated HMAC secret [1]", - }, - { - name: "skip empty rotated paths", - setup: func(_ *testing.T, dir string) []string { - return []string{ - writeFileNamed(t, dir, "current", validSecret), - "", - writeFileNamed(t, dir, "rotated2", validSecret2), - } - }, - wantCurrent: []byte(validSecret), - wantRotated: [][]byte{[]byte(validSecret2)}, - }, - { - name: "whitespace trimmed", - setup: func(_ *testing.T, dir string) []string { - return []string{ - writeFileNamed(t, dir, "current", " "+validSecret+" \n\n"), - writeFileNamed(t, dir, "rotated", "\t"+validSecret2+"\n"), - } - }, - wantCurrent: []byte(validSecret), - wantRotated: [][]byte{[]byte(validSecret2)}, - }, - { - name: "current too short", - setup: func(_ *testing.T, dir string) []string { - return []string{writeFileNamed(t, dir, "current", tooShortSecret)} - }, - wantErr: "HMAC secret must be at least", - }, - { - name: "rotated too short", - setup: func(_ *testing.T, dir string) []string { - return []string{ - writeFileNamed(t, dir, "current", validSecret), - writeFileNamed(t, dir, "rotated", tooShortSecret), - } - }, - wantErr: "failed to load rotated HMAC secret [1]", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - dir := t.TempDir() - paths := tt.setup(t, dir) - - secrets, err := LoadHMACSecrets(paths) - - if tt.wantErr != "" { - require.Error(t, err) - assert.Contains(t, err.Error(), tt.wantErr) - assert.Nil(t, secrets) - } else { - require.NoError(t, err) - if tt.wantCurrent == nil { - assert.Nil(t, secrets) - } else { - require.NotNil(t, secrets) - assert.Equal(t, tt.wantCurrent, secrets.Current) - assert.Equal(t, tt.wantRotated, secrets.Rotated) - } - } - }) - } -} - // Helpers func writePEM(t *testing.T, dir, pemType string, der []byte) string { @@ -426,10 +294,3 @@ func writePEM(t *testing.T, dir, pemType string, der []byte) string { require.NoError(t, os.WriteFile(path, data, 0600)) return path } - -func writeFileNamed(t *testing.T, dir, name, content string) string { - t.Helper() - path := filepath.Join(dir, name) - require.NoError(t, os.WriteFile(path, []byte(content), 0600)) - return path -} diff --git a/pkg/authserver/server/handlers/token.go b/pkg/authserver/server/handlers/token.go index ec5066e5fa..369cd46690 100644 --- a/pkg/authserver/server/handlers/token.go +++ b/pkg/authserver/server/handlers/token.go @@ -6,9 +6,6 @@ package handlers import ( "log/slog" "net/http" - "strings" - - "github.com/ory/fosite" "github.com/stacklok/toolhive/pkg/authserver/server" "github.com/stacklok/toolhive/pkg/authserver/server/session" @@ -30,7 +27,9 @@ func (h *Handler) TokenHandler(w http.ResponseWriter, req *http.Request) { // Parse and validate the access request accessRequest, err := h.provider.NewAccessRequest(ctx, req, sess) if err != nil { - logTokenEndpointError("request_validation", req, err) + slog.Error("failed to create access request", + "error", err, + ) h.provider.WriteAccessError(ctx, w, accessRequest, err) return } @@ -50,8 +49,7 @@ func (h *Handler) TokenHandler(w http.ResponseWriter, req *http.Request) { server.ErrInvalidTarget.WithHint("Multiple resource parameters are not supported")) return } - switch { - case len(resources) == 1 && resources[0] != "": + if len(resources) == 1 && resources[0] != "" { resource := resources[0] // Validate URI format per RFC 8707 if err := server.ValidateAudienceURI(resource); err != nil { @@ -77,25 +75,25 @@ func (h *Handler) TokenHandler(w http.ResponseWriter, req *http.Request) { "resource", resource, ) accessRequest.GrantAudience(resource) - case len(h.config.AllowedAudiences) == 1: - defaultAudience := h.config.AllowedAudiences[0] - slog.Debug("granting default audience for token request without resource", //nolint:gosec // G706: configured audience URI - "audience", defaultAudience, - ) - accessRequest.GrantAudience(defaultAudience) - case len(h.config.AllowedAudiences) > 1: - slog.Debug("resource parameter required when multiple audiences are configured", //nolint:gosec // G706: count is an integer - "count", len(h.config.AllowedAudiences), + } else if accessRequest.GetGrantTypes().ExactOne("authorization_code") && len(h.config.AllowedAudiences) == 1 { + // No resource parameter provided (or provided as empty) during an authorization_code + // exchange; default to the sole allowed audience. The len == 1 guard makes the + // intended audience unambiguous and the index access safe. We restrict this defaulting + // to authorization_code grants: for refresh_token grants, fosite already carries the + // originally-granted audience forward through the session, so re-granting here would + // conflict with fosite's audience matching strategy. + slog.Debug("no resource parameter, defaulting to sole allowed audience", + "audience", h.config.AllowedAudiences[0], ) - h.provider.WriteAccessError(ctx, w, accessRequest, - server.ErrInvalidTarget.WithHint("Token request must include resource when multiple audiences are configured")) - return + accessRequest.GrantAudience(h.config.AllowedAudiences[0]) } // Generate the access response (tokens) response, err := h.provider.NewAccessResponse(ctx, accessRequest) if err != nil { - logTokenEndpointError("response_generation", req, err) + slog.Error("failed to create access response", + "error", err, + ) h.provider.WriteAccessError(ctx, w, accessRequest, err) return } @@ -117,63 +115,3 @@ func (h *Handler) TokenHandler(w http.ResponseWriter, req *http.Request) { // Write the token response h.provider.WriteAccessResponse(ctx, w, accessRequest, response) } - -func logTokenEndpointError(phase string, req *http.Request, err error) { - slog.Error("token endpoint request failed", tokenEndpointErrorAttrs(phase, req, err)...) -} - -// tokenEndpointErrorAttrs builds the structured log attributes for a failed -// token endpoint request. -func tokenEndpointErrorAttrs(phase string, req *http.Request, err error) []any { - oauthErr := fosite.ErrorToRFC6749Error(err) - grantType := req.PostFormValue("grant_type") - - attrs := []any{ - "phase", phase, - "error", err, - "oauth_error", oauthErr.ErrorField, - } - if grantType != "" { - attrs = append(attrs, "grant_type", grantType) - } - if clientID := req.PostFormValue("client_id"); clientID != "" { - attrs = append(attrs, "client_id", clientID) - } - if oauthErr.HintField != "" { - attrs = append(attrs, "oauth_hint", oauthErr.HintField) - } - if oauthErr.DebugField != "" { - attrs = append(attrs, "oauth_debug", oauthErr.DebugField) - } - if grantType == "refresh_token" { - attrs = append(attrs, "refresh_failure_classification", classifyRefreshTokenFailure(oauthErr)) - } - - return attrs -} - -func classifyRefreshTokenFailure(err *fosite.RFC6749Error) string { - if err == nil { - return "unknown" - } - - hint := strings.ToLower(err.HintField) - switch { - case strings.Contains(hint, "already used"): - return "rotated_or_replayed" - case strings.Contains(hint, "malformed or not valid"): - return "storage_lookup_miss" - case strings.Contains(hint, "expired"): - return "expired" - case strings.Contains(hint, "does not match the id during the initial token issuance"): - return "client_mismatch" - case strings.Contains(hint, "multiple concurrent requests"): - return "concurrent_refresh_race" - case strings.Contains(hint, "failed to refresh token"): - return "refresh_rotation_retryable_failure" - case err.ErrorField != "": - return err.ErrorField - default: - return "unknown" - } -} diff --git a/pkg/authserver/server/handlers/token_test.go b/pkg/authserver/server/handlers/token_test.go index bef8ba602c..f317a44a55 100644 --- a/pkg/authserver/server/handlers/token_test.go +++ b/pkg/authserver/server/handlers/token_test.go @@ -4,9 +4,7 @@ package handlers import ( - "bytes" "encoding/json" - "log/slog" "net/http" "net/http/httptest" "net/url" @@ -15,15 +13,12 @@ import ( "time" "github.com/go-jose/go-jose/v4" - "github.com/go-jose/go-jose/v4/jwt" - "github.com/ory/fosite" + josejwt "github.com/go-jose/go-jose/v4/jwt" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto" - "github.com/stacklok/toolhive/pkg/authserver/server/session" "github.com/stacklok/toolhive/pkg/authserver/storage" - "github.com/stacklok/toolhive/pkg/authserver/upstream" ) func TestTokenHandler_MissingGrantType(t *testing.T) { @@ -150,95 +145,6 @@ func TestTokenHandler_InvalidClient(t *testing.T) { assert.Contains(t, rec.Body.String(), "invalid_client") } -func TestClassifyRefreshTokenFailure(t *testing.T) { - t.Parallel() - - testCases := []struct { - name string - err *fosite.RFC6749Error - expected string - }{ - { - name: "reuse detected", - err: fosite.ErrInvalidGrant.WithHint("The refresh token was already used."), - expected: "rotated_or_replayed", - }, - { - name: "lookup miss", - err: fosite.ErrInvalidGrant.WithHint("The refresh token is malformed or not valid."), - expected: "storage_lookup_miss", - }, - { - name: "expired", - err: fosite.ErrInvalidGrant.WithHint("The refresh token expired."), - expected: "expired", - }, - { - name: "client mismatch", - err: fosite.ErrInvalidGrant.WithHint("The OAuth 2.0 Client ID from this request does not match the ID during the initial token issuance."), - expected: "client_mismatch", - }, - { - name: "concurrent race", - err: fosite.ErrInvalidRequest.WithHint("Failed to refresh token because of multiple concurrent requests using the same token. Please retry the request."), - expected: "concurrent_refresh_race", - }, - { - name: "generic fallback", - err: fosite.ErrUnauthorizedClient.WithHint("The OAuth 2.0 Client is not allowed to use authorization grant 'refresh_token'."), - expected: "unauthorized_client", - }, - } - - for _, tc := range testCases { - tc := tc - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - assert.Equal(t, tc.expected, classifyRefreshTokenFailure(tc.err)) - }) - } -} - -func TestTokenHandler_InvalidRefreshTokenLogsClassification(t *testing.T) { - t.Parallel() - handler, _, _ := handlerTestSetup(t) - - newRequest := func() *http.Request { - form := url.Values{ - "grant_type": {"refresh_token"}, - "client_id": {testAuthClientID}, - "refresh_token": {"invalid-refresh-token"}, - } - req := httptest.NewRequest(http.MethodPost, "/oauth/token", strings.NewReader(form.Encode())) - req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - return req - } - rec := httptest.NewRecorder() - - handler.TokenHandler(rec, newRequest()) - - assert.Equal(t, http.StatusBadRequest, rec.Code) - assert.Contains(t, rec.Body.String(), "invalid_grant") - - // Rebuild the validation error that TokenHandler logs and render its log - // attributes with a local logger. The process-wide slog default is not - // replaced, so the test can run in parallel with the others. - req := newRequest() - _, err := handler.provider.NewAccessRequest(req.Context(), req, session.New("", "", "", session.UserClaims{})) - require.Error(t, err) - - var logBuf bytes.Buffer - logger := slog.New(slog.NewTextHandler(&logBuf, &slog.HandlerOptions{Level: slog.LevelDebug})) - logger.Error("token endpoint request failed", tokenEndpointErrorAttrs("request_validation", req, err)...) - - logOutput := logBuf.String() - assert.Contains(t, logOutput, "token endpoint request failed") - assert.Contains(t, logOutput, "phase=request_validation") - assert.Contains(t, logOutput, "grant_type=refresh_token") - assert.Contains(t, logOutput, "oauth_error=invalid_grant") - assert.Contains(t, logOutput, "refresh_failure_classification=storage_lookup_miss") -} - func TestTokenHandler_Success(t *testing.T) { t.Parallel() handler, storState, _ := handlerTestSetup(t) @@ -359,7 +265,7 @@ func TestTokenHandler_AudienceClaim(t *testing.T) { require.True(t, ok, "access_token should be a string") require.NotEmpty(t, accessToken) - parsedToken, err := jwt.ParseSigned(accessToken, []jose.SignatureAlgorithm{jose.RS256}) + parsedToken, err := josejwt.ParseSigned(accessToken, []jose.SignatureAlgorithm{jose.RS256}) require.NoError(t, err) var claims map[string]any @@ -373,87 +279,6 @@ func TestTokenHandler_AudienceClaim(t *testing.T) { } } -func TestTokenHandler_DefaultsSingleAudienceWhenResourceOmitted(t *testing.T) { - t.Parallel() - handler, storState, _ := handlerTestSetup(t) - - authorizeCode := simulateAuthorizeFlow(t, handler, storState) - - form := url.Values{ - "grant_type": {"authorization_code"}, - "client_id": {testAuthClientID}, - "redirect_uri": {testAuthRedirectURI}, - "code": {authorizeCode}, - "code_verifier": {testPKCEVerifier}, - } - req := httptest.NewRequest(http.MethodPost, "/oauth/token", strings.NewReader(form.Encode())) - req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - rec := httptest.NewRecorder() - - handler.TokenHandler(rec, req) - - require.Equal(t, http.StatusOK, rec.Code, "expected 200 OK, got %d: %s", rec.Code, rec.Body.String()) - - tokenData := parseTokenResponse(t, rec.Body.String()) - accessToken, ok := tokenData["access_token"].(string) - require.True(t, ok, "access_token should be a string") - require.NotEmpty(t, accessToken, "access_token should not be empty") - - parsedToken, err := jwt.ParseSigned(accessToken, []jose.SignatureAlgorithm{jose.RS256}) - require.NoError(t, err, "should be able to parse JWT") - - var claims map[string]interface{} - err = parsedToken.Claims(handler.config.SigningKey.Public(), &claims) - require.NoError(t, err, "JWT signature should be valid") - - aud, ok := claims["aud"].([]interface{}) - require.True(t, ok, "aud claim should be an array") - require.Len(t, aud, 1, "aud should have exactly one audience") - assert.Equal(t, "https://api.example.com", aud[0], "default audience should match the sole allowed audience") -} - -func TestTokenHandler_RequiresResourceWhenMultipleAudiencesConfigured(t *testing.T) { - t.Parallel() - - provider, oauth2Config, stor, storState := baseTestSetup(t) - oauth2Config.AllowedAudiences = []string{"https://api.example.com", "https://mcp.example.com"} - - mockUpstream := &mockIDPProvider{ - providerType: upstream.ProviderTypeOAuth2, - authorizationURL: "https://idp.example.com/authorize", - exchangeResult: &upstream.Identity{ - Tokens: &upstream.Tokens{ - AccessToken: "upstream-access-token", - RefreshToken: "upstream-refresh-token", - IDToken: "upstream-id-token", - ExpiresAt: time.Now().Add(time.Hour), - }, - Subject: "user-123", - }, - } - handler, err := NewHandler(provider, oauth2Config, stor, []NamedUpstream{{Name: "test-upstream", Provider: mockUpstream}}) - require.NoError(t, err) - - authorizeCode := simulateAuthorizeFlow(t, handler, storState) - - form := url.Values{ - "grant_type": {"authorization_code"}, - "client_id": {testAuthClientID}, - "redirect_uri": {testAuthRedirectURI}, - "code": {authorizeCode}, - "code_verifier": {testPKCEVerifier}, - } - req := httptest.NewRequest(http.MethodPost, "/oauth/token", strings.NewReader(form.Encode())) - req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - rec := httptest.NewRecorder() - - handler.TokenHandler(rec, req) - - assert.Equal(t, http.StatusBadRequest, rec.Code) - assert.Contains(t, rec.Body.String(), "invalid_target") - assert.Contains(t, rec.Body.String(), "Token request must include resource") -} - func TestTokenHandler_RouteRegistered(t *testing.T) { t.Parallel() handler, _, _ := handlerTestSetup(t) @@ -519,12 +344,3 @@ func simulateAuthorizeFlow(t *testing.T, handler *Handler, storState *testStorag return code } - -func parseTokenResponse(t *testing.T, body string) map[string]interface{} { - t.Helper() - - var tokenData map[string]interface{} - err := json.Unmarshal([]byte(body), &tokenData) - require.NoError(t, err, "token response should be valid JSON") - return tokenData -} diff --git a/pkg/authserver/server/provider.go b/pkg/authserver/server/provider.go index 92dac49960..9da6209d57 100644 --- a/pkg/authserver/server/provider.go +++ b/pkg/authserver/server/provider.go @@ -79,7 +79,7 @@ type AuthorizationServerConfig struct { // The strategy parameter is typed as any because fosite uses different strategy // interfaces for different flows (e.g., oauth2.CoreStrategy, openid.OpenIDConnectTokenStrategy) // that do not share a common base interface. -type Factory func(config *AuthorizationServerConfig, storage fosite.Storage, strategy any) any +type Factory func(config *AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error) // AuthorizationServerParams contains the configuration needed to create an AuthorizationServerConfig. // This is a minimal subset of the authserver.Config fields needed for OAuth2. @@ -273,35 +273,49 @@ func NewAuthorizationServer( storage fosite.Storage, strategy any, factories ...Factory, -) fosite.OAuth2Provider { +) (fosite.OAuth2Provider, error) { fositeConfig := config.Config provider := fosite.NewOAuth2Provider(storage, fositeConfig) for _, factory := range factories { - result := factory(config, storage, strategy) + result, err := factory(config, storage, strategy) + if err != nil { + return nil, fmt.Errorf("authorization server factory failed: %w", err) + } + + var matched bool if ah, ok := result.(fosite.AuthorizeEndpointHandler); ok { fositeConfig.AuthorizeEndpointHandlers.Append(ah) + matched = true } if th, ok := result.(fosite.TokenEndpointHandler); ok { fositeConfig.TokenEndpointHandlers.Append(th) + matched = true } if ti, ok := result.(fosite.TokenIntrospector); ok { fositeConfig.TokenIntrospectionHandlers.Append(ti) + matched = true } if rh, ok := result.(fosite.RevocationHandler); ok { fositeConfig.RevocationHandlers.Append(rh) + matched = true } if ph, ok := result.(fosite.PushedAuthorizeEndpointHandler); ok { fositeConfig.PushedAuthorizeEndpointHandlers.Append(ph) + matched = true + } + + if result != nil && !matched { + return nil, fmt.Errorf("authorization server factory returned unrecognized handler type %T", result) } } - return provider + return provider, nil } // GetSigningKey returns the config's signing key. diff --git a/pkg/authserver/server/provider_test.go b/pkg/authserver/server/provider_test.go index d633d412ee..4699defa0c 100644 --- a/pkg/authserver/server/provider_test.go +++ b/pkg/authserver/server/provider_test.go @@ -564,7 +564,7 @@ func (*mockTokenIntrospector) IntrospectToken(_ context.Context, _ string, _ fos // mockRevocationHandler implements fosite.RevocationHandler for testing. type mockRevocationHandler struct{} -func (*mockRevocationHandler) RevokeToken(_ context.Context, _ string, _ string, _ fosite.Client) error { +func (*mockRevocationHandler) RevokeToken(_ context.Context, _ string, _ fosite.TokenType, _ fosite.Client) error { return nil } @@ -597,72 +597,79 @@ func TestNewAuthorizationServer(t *testing.T) { t.Run("creates provider with no factories", func(t *testing.T) { t.Parallel() - provider := NewAuthorizationServer(createConfig(t), storage, nil) + provider, err := NewAuthorizationServer(createConfig(t), storage, nil) + require.NoError(t, err) require.NotNil(t, provider) }) t.Run("creates provider with authorize handler factory", func(t *testing.T) { t.Parallel() - factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return &mockAuthorizeHandler{} + factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return &mockAuthorizeHandler{}, nil } - provider := NewAuthorizationServer(createConfig(t), storage, nil, factory) + provider, err := NewAuthorizationServer(createConfig(t), storage, nil, factory) + require.NoError(t, err) require.NotNil(t, provider) }) t.Run("creates provider with token handler factory", func(t *testing.T) { t.Parallel() - factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return &mockTokenHandler{} + factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return &mockTokenHandler{}, nil } - provider := NewAuthorizationServer(createConfig(t), storage, nil, factory) + provider, err := NewAuthorizationServer(createConfig(t), storage, nil, factory) + require.NoError(t, err) require.NotNil(t, provider) }) t.Run("creates provider with multiple factories", func(t *testing.T) { t.Parallel() - authorizeFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return &mockAuthorizeHandler{} + authorizeFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return &mockAuthorizeHandler{}, nil } - tokenFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return &mockTokenHandler{} + tokenFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return &mockTokenHandler{}, nil } - introspectorFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return &mockTokenIntrospector{} + introspectorFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return &mockTokenIntrospector{}, nil } - revocationFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return &mockRevocationHandler{} + revocationFactory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return &mockRevocationHandler{}, nil } - provider := NewAuthorizationServer(createConfig(t), storage, nil, + provider, err := NewAuthorizationServer(createConfig(t), storage, nil, authorizeFactory, tokenFactory, introspectorFactory, revocationFactory) + require.NoError(t, err) require.NotNil(t, provider) }) t.Run("handles factory returning nil", func(t *testing.T) { t.Parallel() - factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return nil + factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return nil, nil } - provider := NewAuthorizationServer(createConfig(t), storage, nil, factory) + provider, err := NewAuthorizationServer(createConfig(t), storage, nil, factory) + require.NoError(t, err) require.NotNil(t, provider) }) - t.Run("handles factory returning non-handler type", func(t *testing.T) { + t.Run("errors on factory returning non-handler type", func(t *testing.T) { t.Parallel() - factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) any { - return "not a handler" + factory := func(_ *AuthorizationServerConfig, _ fosite.Storage, _ any) (any, error) { + return "not a handler", nil } - provider := NewAuthorizationServer(createConfig(t), storage, nil, factory) - require.NotNil(t, provider) + provider, err := NewAuthorizationServer(createConfig(t), storage, nil, factory) + require.Error(t, err) + require.Contains(t, err.Error(), "string") + require.Nil(t, provider) }) } diff --git a/pkg/authserver/server/tokenexchange/factory.go b/pkg/authserver/server/tokenexchange/factory.go new file mode 100644 index 0000000000..99c3f58cac --- /dev/null +++ b/pkg/authserver/server/tokenexchange/factory.go @@ -0,0 +1,78 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "fmt" + "time" + + "github.com/ory/fosite" + "github.com/ory/fosite/handler/oauth2" + + "github.com/stacklok/toolhive/pkg/authserver/server" +) + +// Factory returns a server.Factory that creates a token exchange Handler. +// The delegationLifespan parameter sets the maximum lifetime for delegated tokens; +// the actual lifetime is the minimum of this value and the subject token's remaining lifetime. +// Returns an error if delegationLifespan is not in (0, server.MaxAccessTokenLifespan]: a zero +// or negative value would produce delegated tokens with an expiry already in the past, and a +// value above the access token ceiling would only be caught at request time by the per-request cap. +// +// When trustedIssuers is non-empty, subject tokens are validated by a +// MultiIssuerTokenValidator wrapping the self-issued validator; otherwise the +// self-issued validator is used directly, preserving prior behavior exactly. +// Each TrustedIssuer carries its own InsecureAllowHTTP/AllowPrivateIPs (see +// NewMultiIssuerTokenValidator) — this Factory takes no validator-wide +// equivalent, so a self-issuer setting can never reach the external path +// through here. +func Factory(delegationLifespan time.Duration, trustedIssuers []TrustedIssuer) (server.Factory, error) { + if delegationLifespan <= 0 || delegationLifespan > server.MaxAccessTokenLifespan { + return nil, fmt.Errorf("tokenexchange: delegationLifespan must be between %v and %v, got %v", + time.Duration(0), server.MaxAccessTokenLifespan, delegationLifespan) + } + return func(config *server.AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error) { + selfValidator, err := NewSelfIssuedTokenValidator(config.PublicJWKS(), config.GetAccessTokenIssuer(), config.AllowedAudiences) + if err != nil { + return nil, fmt.Errorf("tokenexchange: failed to create subject token validator: %w", err) + } + + // IIFE keeps validator a single immutable assignment rather than a + // mutable var reassigned across branches (go-style): reassigning it + // in place risked ending up with a non-nil SubjectTokenValidator + // wrapping a nil *MultiIssuerTokenValidator on the error path. + validator, err := func() (SubjectTokenValidator, error) { + if len(trustedIssuers) == 0 { + return selfValidator, nil + } + return NewMultiIssuerTokenValidator(selfValidator, config.GetAccessTokenIssuer(), trustedIssuers) + }() + if err != nil { + return nil, fmt.Errorf("tokenexchange: trusted_issuers: %w", err) + } + + // Use the embedded *fosite.Config for HandleHelper and handlerConfig + // because AuthorizationServerConfig shadows GetAccessTokenLifespan() without + // a context parameter, which doesn't satisfy fosite's provider interfaces. + atStrategy, ok := strategy.(oauth2.AccessTokenStrategy) + if !ok { + return nil, fmt.Errorf("tokenexchange: strategy does not implement oauth2.AccessTokenStrategy (got %T)", strategy) + } + atStorage, ok := storage.(oauth2.AccessTokenStorage) + if !ok { + return nil, fmt.Errorf("tokenexchange: storage does not implement oauth2.AccessTokenStorage (got %T)", storage) + } + return &Handler{ + HandleHelper: &oauth2.HandleHelper{ + AccessTokenStrategy: atStrategy, + AccessTokenStorage: atStorage, + Config: config.Config, + }, + validator: validator, + delegationLifespan: delegationLifespan, + config: config.Config, + allowedAudiences: config.AllowedAudiences, + }, nil + }, nil +} diff --git a/pkg/authserver/server/tokenexchange/factory_test.go b/pkg/authserver/server/tokenexchange/factory_test.go new file mode 100644 index 0000000000..bd3c75f28d --- /dev/null +++ b/pkg/authserver/server/tokenexchange/factory_test.go @@ -0,0 +1,192 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "context" + "crypto/rand" + "crypto/rsa" + "testing" + "time" + + "github.com/ory/fosite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/authserver/server" + servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto" +) + +func TestFactory(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + delegationLifespan time.Duration + wantErr bool + }{ + { + name: "zero delegationLifespan returns error", + delegationLifespan: 0, + wantErr: true, + }, + { + name: "negative delegationLifespan returns error", + delegationLifespan: -time.Minute, + wantErr: true, + }, + { + name: "positive delegationLifespan succeeds", + delegationLifespan: 15 * time.Minute, + }, + { + name: "delegationLifespan at max access token lifespan succeeds", + delegationLifespan: server.MaxAccessTokenLifespan, + }, + { + name: "delegationLifespan above max access token lifespan returns error", + delegationLifespan: server.MaxAccessTokenLifespan + time.Hour, + wantErr: true, + }, + { + name: "delegationLifespan of 48h returns error", + delegationLifespan: 48 * time.Hour, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + f, err := Factory(tt.delegationLifespan, nil) + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), "delegationLifespan must be between") + assert.Nil(t, f) + return + } + require.NoError(t, err) + assert.NotNil(t, f) + }) + } +} + +// buildTestAuthServerConfig returns a minimally-valid +// *server.AuthorizationServerConfig for exercising the closure Factory +// returns. AllowedAudiences is non-empty so it doubles as a valid +// ExpectedAudience target for a TrustedIssuer in the tests below. +func buildTestAuthServerConfig(t *testing.T) *server.AuthorizationServerConfig { + t.Helper() + + rsaKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + + cfg, err := server.NewAuthorizationServerConfig(&server.AuthorizationServerParams{ + Issuer: "https://auth.example.com", + AccessTokenLifespan: time.Hour, + RefreshTokenLifespan: time.Hour * 24, + AuthCodeLifespan: time.Minute * 10, + HMACSecrets: servercrypto.NewHMACSecrets([]byte("test-secret-with-32-bytes-long!!")), + SigningKeyID: "key-1", + SigningKeyAlgorithm: "RS256", + SigningKey: rsaKey, + AllowedAudiences: []string{"https://mcp.example.com"}, + }) + require.NoError(t, err) + return cfg +} + +// fakeClientManager is a no-op fosite.ClientManager (the entirety of +// fosite.Storage) so fakeFactoryStorage satisfies the Factory closure's +// storage fosite.Storage parameter without needing a real client store — +// the closure only type-asserts storage to oauth2.AccessTokenStorage, it +// never calls a ClientManager method. +type fakeClientManager struct{} + +func (fakeClientManager) GetClient(context.Context, string) (fosite.Client, error) { + return nil, fosite.ErrNotFound +} +func (fakeClientManager) ClientAssertionJWTValid(context.Context, string) error { return nil } +func (fakeClientManager) SetClientAssertionJWT(context.Context, string, time.Time) error { + return nil +} + +// fakeFactoryStorage combines the no-op ClientManager with the package's +// existing mockAccessTokenStorage so it satisfies both fosite.Storage (the +// closure's declared parameter type) and oauth2.AccessTokenStorage (what the +// closure actually type-asserts against). +type fakeFactoryStorage struct { + fakeClientManager + *mockAccessTokenStorage +} + +// TestFactory_ValidatorSelection asserts which SubjectTokenValidator the +// closure returned by Factory builds into the Handler: the self-issued +// validator when trustedIssuers is empty, the multi-issuer validator when +// it isn't, and a hard error — not a silent downgrade to the self-issued +// validator — when a configured TrustedIssuer is itself invalid. +func TestFactory_ValidatorSelection(t *testing.T) { + t.Parallel() + + validIssuer := TrustedIssuer{ + IssuerURL: "https://idp.example.com", + ExpectedAudience: "https://mcp.example.com", + AllowedDelegateClients: []string{anyDelegateClient}, + } + invalidIssuer := TrustedIssuer{ + IssuerURL: "https://idp.example.com", + // ExpectedAudience deliberately empty: invalid per validateTrustedIssuer. + AllowedDelegateClients: []string{anyDelegateClient}, + } + + tests := []struct { + name string + trustedIssuers []TrustedIssuer + wantErr string + wantValidator any // nil when wantErr is set + }{ + { + name: "no trusted issuers builds self-issued validator", + trustedIssuers: nil, + wantValidator: &SelfIssuedTokenValidator{}, + }, + { + name: "valid trusted issuer builds multi-issuer validator", + trustedIssuers: []TrustedIssuer{validIssuer}, + wantValidator: &MultiIssuerTokenValidator{}, + }, + { + name: "invalid trusted issuer fails closed, not silently downgraded", + trustedIssuers: []TrustedIssuer{invalidIssuer}, + wantErr: "trusted_issuers", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + f, err := Factory(15*time.Minute, tt.trustedIssuers) + require.NoError(t, err) + + cfg := buildTestAuthServerConfig(t) + storage := &fakeFactoryStorage{mockAccessTokenStorage: &mockAccessTokenStorage{}} + strategy := &mockAccessTokenStrategy{} + + result, err := f(cfg, storage, strategy) + if tt.wantErr != "" { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + assert.Nil(t, result) + return + } + require.NoError(t, err) + + handler, ok := result.(*Handler) + require.True(t, ok, "Factory closure must return *Handler, got %T", result) + assert.IsType(t, tt.wantValidator, handler.validator) + }) + } +} diff --git a/pkg/authserver/server/tokenexchange/handler.go b/pkg/authserver/server/tokenexchange/handler.go new file mode 100644 index 0000000000..304fa9eea7 --- /dev/null +++ b/pkg/authserver/server/tokenexchange/handler.go @@ -0,0 +1,685 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "context" + "fmt" + "log/slog" + "net/url" + "slices" + "strings" + "time" + + "github.com/ory/fosite" + "github.com/ory/fosite/handler/oauth2" + "github.com/ory/x/errorsx" + + coreaudit "github.com/stacklok/toolhive-core/audit" + "github.com/stacklok/toolhive/pkg/authserver/server" + "github.com/stacklok/toolhive/pkg/authserver/server/session" + "github.com/stacklok/toolhive/pkg/oauthproto" +) + +// Compile-time check that Handler implements fosite.TokenEndpointHandler. +var _ fosite.TokenEndpointHandler = (*Handler)(nil) + +// maxDelegationDepth bounds how many nested "act" entries a subject token's +// delegation chain may carry. Without a cap, a client could repeatedly +// re-exchange a delegated token to grow the chain without limit, bloating the +// issued token and the load it places on downstream resource servers that +// parse it. +const maxDelegationDepth = 10 + +// anyDelegateClient is the TrustedIssuer.AllowedDelegateClients entry that +// explicitly opts an issuer into "any ToolHive confidential client holding +// the token-exchange grant may act as its delegate" — the behavior that used +// to be the implicit default whenever the field was left empty. See +// checkDelegationConsent's doc comment for why that default was removed: +// permissiveness must now be declared with this wildcard, not left to an +// absent field. +const anyDelegateClient = "*" + +// Handler implements RFC 8693 token exchange for user-to-agent delegation. +// +// When an authenticated OAuth client (the acting agent) presents a user's JWT +// as subject_token, the handler validates the token and issues a delegated JWT +// with sub=user and an act claim containing the client's identity, per +// RFC 8693 Section 4.1. +// +// Subject tokens are intentionally reusable within their lifetime: per +// RFC 8693's security considerations, a token exchange does not invalidate +// the subject token, so the same subject token may be exchanged more than +// once. Replay is bounded by the delegated token's lifetime cap +// (min(subject-remaining, delegation)), not by single-use tracking; per-jti +// single-use enforcement is deferred to the broader M2M/sender-constrained- +// token effort. +type Handler struct { + *oauth2.HandleHelper + validator SubjectTokenValidator + delegationLifespan time.Duration + config tokenExchangeConfig + allowedAudiences []string +} + +// tokenExchangeConfig defines the configuration interface needed by the handler. +type tokenExchangeConfig interface { + fosite.ScopeStrategyProvider + fosite.AudienceStrategyProvider + fosite.AccessTokenLifespanProvider +} + +// CanHandleTokenEndpointRequest returns true if the request's grant_type is +// the RFC 8693 token exchange grant type. +func (*Handler) CanHandleTokenEndpointRequest(_ context.Context, requester fosite.AccessRequester) bool { + return requester.GetGrantTypes().ExactOne(oauthproto.GrantTypeTokenExchange) +} + +// CanSkipClientAuth returns false because client authentication is required +// for all token exchange requests. +func (*Handler) CanSkipClientAuth(_ context.Context, _ fosite.AccessRequester) bool { + return false +} + +// HandleTokenEndpointRequest validates the token exchange request parameters, +// verifies the subject token, and constructs a delegated session with the act claim. +// +// The delegated token's lifetime is the minimum of the subject token's remaining +// lifetime and the configured delegation lifespan. +func (h *Handler) HandleTokenEndpointRequest(ctx context.Context, requester fosite.AccessRequester) error { + if !h.CanHandleTokenEndpointRequest(ctx, requester) { + return errorsx.WithStack(fosite.ErrUnknownRequest) + } + + client := requester.GetClient() + + // The client MUST be confidential — only authenticated confidential clients + // may act on behalf of a user. + if client.IsPublic() { + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The OAuth 2.0 Client is marked as public and is thus not allowed to use authorization grant 'token-exchange'.")) + } + + // Reject clients not registered for this grant type up front, before any + // subject-token validation work. PopulateTokenEndpointResponse repeats + // this check as defense-in-depth. + if !client.GetGrantTypes().Has(oauthproto.GrantTypeTokenExchange) { + return errorsx.WithStack(fosite.ErrUnauthorizedClient.WithHint( + "The OAuth 2.0 Client is not allowed to use authorization grant 'token-exchange'.")) + } + + // The authenticated client is the acting party ("actor"). The client is + // already authenticated by fosite's client authentication strategy before + // this handler runs. + actorID := client.GetID() + + subjectToken, err := validateExchangeParams(requester.GetRequestForm()) + if err != nil { + return err + } + + // Validate the subject token against the server's own JWKS. + validatedClaims, err := h.validator.Validate(ctx, subjectToken) + if err != nil { + slog.Debug("Subject token validation failed", + "error", err, + "actor", actorID, + ) + // RFC 8693 §2.2.2 mandates invalid_request here: "if either the + // subject_token or actor_token are invalid for any reason, or are + // unacceptable based on policy ... The value of the error parameter + // MUST be the invalid_request error code." checkDelegationConsent + // below deliberately returns invalid_grant instead for its own + // failures — a documented deviation: RFC 6749 §5.2's invalid_grant + // covers "was issued to another client", §2.2.2 permits other error + // codes for other failures, and both Keycloak and Hydra follow this + // same split. + return errorsx.WithStack(fosite.ErrInvalidRequest.WithHint( + "The subject token is invalid or could not be verified.")) + } + + if err := checkDelegationConsent(validatedClaims, actorID); err != nil { + return err + } + + if err := h.grantScopes(ctx, requester, client, validatedClaims); err != nil { + return err + } + + if err := h.grantAndBoundAudiences(ctx, requester, client, validatedClaims); err != nil { + return err + } + + // Build the delegated session with the user's identity and the agent's act claim. + delegatedSession := session.New( + delegatedSubject(validatedClaims), + "", // No IDP session link for delegated tokens. + actorID, + session.UserClaims{ + Name: validatedClaims.Name, + Email: validatedClaims.Email, + }, + ) + + act, err := buildActClaim(validatedClaims, actorID) + if err != nil { + return err + } + delegatedSession.JWTClaims.Extra["act"] = act + + // Compute the delegated token lifetime: the shorter of the subject token's + // remaining lifetime and the configured delegation lifespan. + lifetime, err := h.computeLifetime(validatedClaims.Expiry) + if err != nil { + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The subject token has expired.")) + } + delegatedSession.SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(lifetime)) + + requester.SetSession(delegatedSession) + + slog.Debug("Token exchange request validated", + "subject", validatedClaims.Subject, + "actor", actorID, + "issuer", validatedClaims.Issuer, + "subject_token_client", validatedClaims.ExternalActor, + "lifetime", lifetime.String(), + ) + + return nil +} + +// PopulateTokenEndpointResponse issues the delegated access token and sets +// the RFC 8693 issued_token_type in the response. +func (h *Handler) PopulateTokenEndpointResponse( + ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder, +) error { + if !h.CanHandleTokenEndpointRequest(ctx, requester) { + return errorsx.WithStack(fosite.ErrUnknownRequest) + } + + if !requester.GetClient().GetGrantTypes().Has(oauthproto.GrantTypeTokenExchange) { + return errorsx.WithStack(fosite.ErrUnauthorizedClient.WithHint( + "The OAuth 2.0 Client is not allowed to use authorization grant 'token-exchange'.")) + } + + // Use the session's ExpiresAt (set during HandleTokenEndpointRequest) as the + // authoritative lifetime. This respects the min(subject_remaining, delegation) + // bound computed earlier. Fall back to the configured access token lifespan + // only if no expiry was set on the session. + atLifespan := h.config.GetAccessTokenLifespan(ctx) + if sessionExpiry := requester.GetSession().GetExpiresAt(fosite.AccessToken); !sessionExpiry.IsZero() { + remaining := time.Until(sessionExpiry) + if remaining <= 0 { + // The session's bound has already elapsed since HandleTokenEndpointRequest + // computed it. Fail closed instead of signing a token whose exp is already + // in the past while reporting a negative expires_in — fosite's JWT strategy + // sets exp directly from this session's expiry (GetExpiresAt), independent + // of atLifespan, so letting this fall through would not simply hand out a + // longer-lived token, it would issue an already-expired one. + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The delegated token's bound lifetime has already elapsed.")) + } + if remaining < atLifespan { + atLifespan = remaining + } + } + + if _, err := h.IssueAccessToken(ctx, atLifespan, requester, responder); err != nil { + return err + } + + // Per RFC 8693 Section 2.2.1, the response MUST include issued_token_type. + responder.SetExtra("issued_token_type", oauthproto.TokenTypeAccessToken) + + return nil +} + +// computeLifetime returns the minimum of the subject token's remaining lifetime +// and the configured delegation lifespan. Returns an error if the subject token +// has already expired. +func (h *Handler) computeLifetime(subjectExpiry time.Time) (time.Duration, error) { + remaining := time.Until(subjectExpiry) + if remaining <= 0 { + return 0, fmt.Errorf("subject token expired %v ago", -remaining) + } + + if remaining < h.delegationLifespan { + return remaining, nil + } + return h.delegationLifespan, nil +} + +// validateExchangeParams validates the required RFC 8693 form parameters and +// returns the subject token on success. +func validateExchangeParams(form url.Values) (string, error) { + subjectToken := form.Get("subject_token") + if subjectToken == "" { + return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint( + "The 'subject_token' parameter is required for token exchange.")) + } + + subjectTokenType := form.Get("subject_token_type") + if subjectTokenType == "" { + return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint( + "The 'subject_token_type' parameter is required for token exchange.")) + } + + if subjectTokenType != oauthproto.TokenTypeAccessToken && subjectTokenType != oauthproto.TokenTypeJWT { + return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf( + "The 'subject_token_type' value %q is not supported. Use %q or %q.", + subjectTokenType, oauthproto.TokenTypeAccessToken, oauthproto.TokenTypeJWT)) + } + + // Reject actor_token parameters for now — the acting party identity is + // derived from the authenticated OAuth client. A later commit adds + // actor_token support for asserting a distinct actor. + if form.Get("actor_token") != "" || form.Get("actor_token_type") != "" { + return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHint( + "The 'actor_token' and 'actor_token_type' parameters are not yet supported.")) + } + + // Validate requested_token_type per RFC 8693 Section 2.1: if the client + // requests a token type the server does not support, the request must fail. + requestedTokenType := form.Get("requested_token_type") + if requestedTokenType != "" && requestedTokenType != oauthproto.TokenTypeAccessToken { + return "", errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf( + "The 'requested_token_type' value %q is not supported. This server only issues %q.", + requestedTokenType, oauthproto.TokenTypeAccessToken)) + } + + return subjectToken, nil +} + +// delegatedSubject returns the "sub" to embed in the delegated token. +// +// Self-issued tokens pass through unchanged: their subject is already a +// UUID from UserResolver.ResolveUser's (providerID, providerSubject) +// mapping (pkg/authserver/server/handlers/user.go), ToolHive's own native +// subject namespace. +// +// A subject token from a trusted external issuer is qualified as +// "#" instead, because the issued token's "sub" is the only +// signal Cedar's extractClientIDFromClaims (pkg/authz/authorizers/cedar/core.go) +// keys on — it reads "sub" alone, not "iss" — and the external issuer +// chooses that value, not ToolHive. Without qualification, an external +// issuer could pick a "sub" equal to a native user's UUID and mint a +// delegated token indistinguishable from that user's, entirely without +// malicious intent on either issuer's part: two trusted issuers merely +// choosing overlapping subject values collide with no attacker involved. +// +// "#" is an unambiguous separator: an issuer URL cannot itself contain a +// fragment (the OIDC/OAuth "iss" value is a URL without a fragment +// component, and this validator's own issuer comparison is an exact string +// match against that value), so strings.Cut(sub, "#") on the qualified +// subject always recovers (issuerURL, externalSub) uniquely. A native UUID +// can never collide with a qualified value, since a UUID contains no "#". +// +// Re-exchanging a previously-issued delegated token does not double-qualify: +// such a token's "iss" is always this server's own issuer (fosite stamps +// "iss" from the server's config at issuance — see session.New's doc +// comment), so Validate routes it to the self-issued path, which never +// calls this branch. +func delegatedSubject(validatedClaims *ValidatedClaims) string { + if validatedClaims.ExternalIssuer == "" { + return validatedClaims.Subject + } + return validatedClaims.ExternalIssuer + "#" + validatedClaims.Subject +} + +// buildActClaim assembles the RFC 8693 Section 4.1 "act" claim for the +// delegated token. The outermost act.sub is always actorID (the ToolHive +// client) — every downstream consumer reads that as "who is acting", and it +// must not change regardless of how the subject token was obtained. +// +// Extracted from HandleTokenEndpointRequest rather than inlined: the external +// provenance nesting and the prior-chain depth gate below add five branches +// that have nothing to do with the surrounding request plumbing, and two of +// them reject the request outright. +func buildActClaim(validatedClaims *ValidatedClaims, actorID string) (map[string]any, error) { + act := map[string]any{"sub": actorID} + nestUnder := act + newLevels := 1 + + // When the subject token came from a trusted external issuer + // (ValidatedClaims.ExternalIssuer is set only there — see + // multi_issuer_validator.go), nest that issuer one level in, together + // with the allowlisted actor when one was resolved. RFC 8693 §4.1 + // anticipates exactly this: "the combination of the two claims 'iss' and + // 'sub' might be necessary to uniquely identify an actor." Without this, + // the issued token would carry no record that the delegation originated + // externally at all — including for a may_act-bearing external token, + // which leaves ExternalActor unset because may_act.sub already names the + // delegate directly via the actorID binding above. That token still has + // an external issuer worth recording, so nesting here is keyed on + // ExternalIssuer, not ExternalActor: the allowlist path's accepted + // any-ToolHive-client scope limitation (see checkDelegationConsent) + // depends on this provenance being auditable after the fact, and a + // may_act-bearing external token — the path that bypasses the allowlist + // entirely — needs it at least as much. + if validatedClaims.ExternalIssuer != "" { + external := map[string]any{"iss": validatedClaims.ExternalIssuer} + // ExternalActor is only present on the allowlist path (see its doc + // comment) — a may_act-bearing external token has no client-namespace + // actor claim to report, so the nested entry there carries "iss" only. + if validatedClaims.ExternalActor != "" { + external["sub"] = validatedClaims.ExternalActor + } + act["act"] = external + nestUnder = external + newLevels = 2 + } + + // If the subject token itself carries an "act" claim (i.e. it is a + // previously-delegated token being re-exchanged), nest it so the full + // delegation chain remains auditable rather than being discarded. + // newLevels accounts for the external wrapper above, if any, so the + // resulting chain never exceeds maxDelegationDepth regardless of how + // many levels this exchange itself adds. + if priorAct, ok := validatedClaims.Extra["act"]; ok && priorAct != nil { + // Parse with the shared audit-side parser rather than a bespoke walker: it + // reports both the chain depth and any RFC 8693 Section 4.1 conformance + // violation, so a non-object act cannot slip past the depth gate and be + // re-minted. A JSON-null act is filtered by the guard above: it asserts no + // delegation and must not read as malformed. + chain := coreaudit.ParseDelegationChain(priorAct, maxDelegationDepth) + if chain.Malformed { + // RFC 8693 Section 2.2.2 nominally calls for invalid_request on a bad + // subject_token, but also allows that "other error codes may also be + // used, as appropriate": invalid_grant keeps this consistent with the + // depth, consent, and expiry gates in this same handler, all of which + // reject a structurally unacceptable subject token that way. + // + // MalformedReason is a closed, value-free enum; it is surfaced to the + // client and MUST stay that way — never interpolate claim contents here. + return nil, errorsx.WithStack(fosite.ErrInvalidGrant.WithHintf( + "The subject token's delegation chain is malformed (%s).", chain.MalformedReason)) + } + // len(Chain) is the prior chain's depth: the parser appends exactly one + // hop per level and stops at maxDelegationDepth, so it is + // min(depth, maxDelegationDepth). newLevels is what this exchange adds + // on top — one for the acting client, two when an external issuer is + // also nested — so the resulting chain never exceeds the cap. + if len(chain.Chain)+newLevels > maxDelegationDepth { + return nil, errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The subject token's delegation chain is too deep.")) + } + // Nest priorAct verbatim rather than re-serializing from chain: core keeps + // each hop's extra claims in an unexported map, so rebuilding would silently + // drop the history trail that Section 4.1 asks us to preserve. The cost is + // that unknown hop claims — and the non-identity claims Section 4.1 calls + // "not meaningful" inside act — pass through re-signed, which sits in + // tension with Section 6 data minimization. Accepted here; bounding the + // subtree is tracked separately. + // + // nestUnder is the innermost entry this exchange added: the acting + // client normally, or the external-issuer entry when one was nested + // above, so the prior chain hangs below the provenance rather than + // displacing it. + nestUnder["act"] = priorAct + } + return act, nil +} + +// delegateClientAllowed reports whether actorID may use an issuer's +// AllowedDelegateClients-bound consent grant (a may_act.sub match or an +// AllowedActors match): either the list contains the wildcard +// anyDelegateClient, or it contains actorID directly. An empty list denies +// every client — validateTrustedIssuer rejects that configuration outright +// for any TrustedIssuer, so the empty case here is an extra fail-closed +// backstop, not the path an operator is expected to reach. Callers must only +// invoke this for a claim that actually went through that per-issuer +// validation (ExternalIssuer != ""); a self-issued token's +// AllowedDelegateClients is always nil, since that field has no equivalent +// concept on the self-issued path. +func delegateClientAllowed(allowedDelegateClients []string, actorID string) bool { + return slices.Contains(allowedDelegateClients, anyDelegateClient) || + slices.Contains(allowedDelegateClients, actorID) +} + +// checkDelegationConsent enforces the RFC 8693 §4.4 delegation consent check. +// +// There are three consent sources, checked in order: +// +// 1. may_act: if present, it is the authoritative consent signal — only the +// party named in may_act.sub may delegate. The client_id binding is +// skipped in this case because may_act enables cross-client delegation +// (the token was issued to client A but authorizes client B to act). +// ValidatedClaims.AllowedDelegateClients is additionally enforced here, +// but ONLY when ExternalIssuer is set: that field is a TrustedIssuer +// concept with no self-issued equivalent, so a self-issued may_act (a +// ToolHive user token delegating to a ToolHive client directly) is +// bound by may_act.sub alone, the same as before this check existed. +// On the external path, may_act bypasses AllowedActors, not per-client +// containment — delegateClientAllowed still applies there. +// +// 2. ExternalActor: if may_act is absent but the multi-issuer validator has +// already established consent for this token (multi_issuer_validator.go), +// it did so by matching the subject token's actor claim against that +// issuer's operator-configured AllowedActors. That actor claim — even +// when configured as "client_id" — names a client in the EXTERNAL +// issuer's namespace, not ToolHive's, so it must never be compared +// against actorID the way case 3 below compares ValidatedClaims.ClientID. +// This case MUST be checked before case 3: it is not a fallback for an +// empty ClientID, it is a distinct, already-verified consent signal that +// takes priority whenever it is set, even if ClientID also happens to be +// populated. +// +// ValidatedClaims.AllowedDelegateClients binds this to specific ToolHive +// clients: actorID must appear in it (or it must contain the wildcard +// anyDelegateClient), checked below by delegateClientAllowed. Without +// this, the allowlist would authorize "this external client's tokens may +// be exchanged", not "...by this particular ToolHive client" — every +// ToolHive confidential client holding the token-exchange grant would be +// delegation-equivalent with respect to an allowlisted external actor, so +// compromise of the weakest such client would be as good as compromise +// of all of them (see #5989). validateTrustedIssuer rejects an empty +// AllowedDelegateClients at construction for exactly this reason: an +// operator must declare permissiveness with the wildcard, not get it by +// omission. Either way this remains bounded by: the calling client must +// already possess a valid subject token (it cannot forge one), and +// grantScopes/grantAndBoundAudiences still narrow the result to what +// both the client and the subject token are authorized for. +// +// 3. client_id: if neither of the above applies, fall back to client_id +// binding — the subject token's client_id must match the authenticated +// client. This prevents a stolen subject token from being exchanged by a +// different client. +// +// If none of the three consent sources apply, the subject token carries no +// verifiable binding to any client at all — this fails closed (CWE-863) +// rather than allowing an unbound token through. +func checkDelegationConsent(validatedClaims *ValidatedClaims, actorID string) error { + switch { + case validatedClaims.MayAct != nil: + if validatedClaims.MayAct.Sub != actorID { + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The subject token does not authorize this client to act on behalf of the subject.")) + } + if validatedClaims.ExternalIssuer != "" && !delegateClientAllowed(validatedClaims.AllowedDelegateClients, actorID) { + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "This client is not authorized to exchange subject tokens from the external actor's issuer.")) + } + case validatedClaims.ExternalActor != "": + // Consent was already established by the external-issuer validator: the + // subject token's actor claim matched this issuer's operator-configured + // AllowedActors. That claim lives in the external issuer's client + // namespace, not ToolHive's, so — even when ClientID is also populated + // (ActorClaim: "client_id") — it must never be compared against + // actorID. This case must be checked before the client_id cases below, + // not merged with them. + // + // AllowedDelegateClients binds this allowlisted actor to a specific set + // of ToolHive clients — see delegateClientAllowed. + if !delegateClientAllowed(validatedClaims.AllowedDelegateClients, actorID) { + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "This client is not authorized to exchange subject tokens from the external actor's issuer.")) + } + case validatedClaims.ClientID != "" && validatedClaims.ClientID != actorID: + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The subject token was issued to a different client.")) + case validatedClaims.ClientID == "": + return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint( + "The subject token carries no verifiable client binding: it has neither a 'may_act' claim nor a 'client_id' claim.")) + } + return nil +} + +// grantScopes validates that each requested scope is allowed for both the +// client and the subject token, granting the intersection. The delegated +// token's scope set is the intersection of the client's registered scopes +// and the subject token's granted scopes, preventing a client from +// escalating privileges beyond what the user authorized. A subject token +// without a scope claim grants no scopes. +func (h *Handler) grantScopes( + ctx context.Context, requester fosite.AccessRequester, client fosite.Client, validatedClaims *ValidatedClaims, +) error { + subjectScopes := strings.Fields(validatedClaims.Scopes) + subjectScopeSet := make(map[string]bool, len(subjectScopes)) + for _, s := range subjectScopes { + subjectScopeSet[s] = true + } + + for _, scope := range requester.GetRequestedScopes() { + if !h.config.GetScopeStrategy(ctx)(client.GetScopes(), scope) { + return errorsx.WithStack(fosite.ErrInvalidScope.WithHintf( + "The OAuth 2.0 Client is not allowed to request scope '%s'.", scope)) + } + if !subjectScopeSet[scope] { + return errorsx.WithStack(fosite.ErrInvalidScope.WithHintf( + "The scope '%s' was not granted by the subject token.", scope)) + } + requester.GrantScope(scope) + } + return nil +} + +// grantAndBoundAudiences resolves the delegated token's audience from the +// request (explicit "audience", RFC 8707 "resource", or the configured +// default) and then bounds the result by the subject token. +// +// The delegated token must not target a resource the subject token was not +// itself valid for: every granted audience must be covered by the subject +// token's audience — delegation narrows, never broadens, the resource +// boundary. This mirrors grantScopes, which bounds scopes by the subject +// token; without it a client registered for audiences A and B could exchange a +// user token minted for A into a token for B, an escalation the user never +// consented to. +func (h *Handler) grantAndBoundAudiences( + ctx context.Context, requester fosite.AccessRequester, client fosite.Client, validatedClaims *ValidatedClaims, +) error { + if err := h.grantAudiences(ctx, requester, client); err != nil { + return err + } + if err := h.grantResourceAudience(ctx, requester, client); err != nil { + return err + } + if err := h.grantDefaultAudience(ctx, requester, client); err != nil { + return err + } + return ensureAudienceSubsetOfSubject(requester.GetGrantedAudience(), validatedClaims.Audience) +} + +// grantAudiences validates that the requested audience is allowed for this +// client and grants it. +func (h *Handler) grantAudiences(ctx context.Context, requester fosite.AccessRequester, client fosite.Client) error { + if err := h.config.GetAudienceStrategy(ctx)(client.GetAudience(), requester.GetRequestedAudience()); err != nil { + return errorsx.WithStack(err) + } + for _, aud := range requester.GetRequestedAudience() { + requester.GrantAudience(aud) + } + return nil +} + +// grantResourceAudience validates the RFC 8707 resource parameter against +// both the server's allowedAudiences and the client's own registered +// audiences, and grants it as an additional audience claim, binding the +// issued token to a specific resource server (e.g., an MCP server). +// +// Per RFC 8707 §2, a request MAY carry multiple resource parameters; this +// handler deliberately narrows that to at most one, since a delegated MCP +// token is expected to target a single resource. +func (h *Handler) grantResourceAudience(ctx context.Context, requester fosite.AccessRequester, client fosite.Client) error { + resources := requester.GetRequestForm()["resource"] + if len(resources) > 1 { + return errorsx.WithStack(server.ErrInvalidTarget.WithHint( + "Multiple resource parameters are not supported.")) + } + if len(resources) == 1 && resources[0] != "" { + resource := resources[0] + if err := server.ValidateAudienceURI(resource); err != nil { + return errorsx.WithStack(err) + } + if err := server.ValidateAudienceAllowed(resource, h.allowedAudiences); err != nil { + return errorsx.WithStack(err) + } + // The resource parameter is RFC 8707's mechanism for requesting an + // audience, so it must be subject to the same per-client audience + // registration as the "audience" parameter (grantAudiences) — otherwise + // a client could bypass its registered audiences simply by using + // "resource" instead of "audience". + if err := h.config.GetAudienceStrategy(ctx)(client.GetAudience(), []string{resource}); err != nil { + return errorsx.WithStack(server.ErrInvalidTarget.WithHintf( + "The client is not permitted to request a token for resource %q.", resource)) + } + requester.GrantAudience(resource) + } + return nil +} + +// grantDefaultAudience ensures the delegated token carries at least one +// audience, since RFC 9068 §4 makes "aud" a required claim. grantAudiences and +// grantResourceAudience only grant an audience when the client explicitly +// requests one; if neither was requested, this defaults to the server's sole +// configured audience when unambiguous (mirroring handlers/token.go's +// authorization_code fallback), or rejects the request when no unambiguous +// default exists rather than silently issuing an audience-less token. The +// default audience is checked against the client's registered audiences +// like any explicitly requested audience, so a client is never granted an +// audience it wasn't registered for. +func (h *Handler) grantDefaultAudience(ctx context.Context, requester fosite.AccessRequester, client fosite.Client) error { + if len(requester.GetGrantedAudience()) > 0 { + return nil + } + if len(h.allowedAudiences) != 1 { + return errorsx.WithStack(fosite.ErrInvalidRequest.WithHint( + "An explicit 'resource' or 'audience' parameter is required because no unambiguous default audience is configured.")) + } + defaultAudience := h.allowedAudiences[0] + if err := h.config.GetAudienceStrategy(ctx)(client.GetAudience(), []string{defaultAudience}); err != nil { + return errorsx.WithStack(fosite.ErrInvalidRequest.WithHintf( + "The client is not permitted to request a token for the default audience %q; "+ + "an explicit 'resource' or 'audience' parameter is required.", defaultAudience)) + } + slog.Debug("no resource parameter, defaulting to sole allowed audience", + "audience", defaultAudience, + ) + requester.GrantAudience(defaultAudience) + return nil +} + +// ensureAudienceSubsetOfSubject verifies that every audience granted to the +// delegated token is covered by the subject token's own audience. A subject +// token always carries at least one audience, so an empty subject audience +// here can only reject — but what guarantees that non-empty audience differs +// by path. On the self-issued path, SelfIssuedTokenValidator rejects tokens +// whose aud does not intersect this server's AllowedAudiences. On the +// external path, validateExternalToken instead requires aud to contain the +// issuer's own configured ExpectedAudience, which has no required +// relationship to AllowedAudiences — see the TrustedIssuers field doc +// comment in pkg/authserver/config.go for the operator-facing consequence. +func ensureAudienceSubsetOfSubject(granted, subjectAud []string) error { + subj := make(map[string]bool, len(subjectAud)) + for _, a := range subjectAud { + subj[a] = true + } + for _, a := range granted { + if !subj[a] { + return errorsx.WithStack(server.ErrInvalidTarget.WithHintf( + "The delegated token audience %q is not covered by the subject token's audience.", a)) + } + } + return nil +} diff --git a/pkg/authserver/server/tokenexchange/handler_test.go b/pkg/authserver/server/tokenexchange/handler_test.go new file mode 100644 index 0000000000..a5f1f9a389 --- /dev/null +++ b/pkg/authserver/server/tokenexchange/handler_test.go @@ -0,0 +1,1640 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "context" + "errors" + "fmt" + "net/url" + "testing" + "time" + + "github.com/go-jose/go-jose/v4/jwt" + "github.com/ory/fosite" + "github.com/ory/fosite/handler/oauth2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/authserver/server" + "github.com/stacklok/toolhive/pkg/authserver/server/session" + "github.com/stacklok/toolhive/pkg/oauthproto" +) + +const testAgentClientID = "devops-agent" + +// newTestHandler creates a TokenExchangeHandler wired to the given test JWKS and +// delegation lifespan. It does not need a full fosite strategy stack because +// HandleTokenEndpointRequest (which is the focus of these tests) does not issue tokens. +func newTestHandler(t *testing.T, tj *testJWKS, delegationLifespan time.Duration) *Handler { + t.Helper() + + validator, err := NewSelfIssuedTokenValidator(tj.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + return &Handler{ + // HandleHelper is nil — we only test HandleTokenEndpointRequest, not + // PopulateTokenEndpointResponse, so IssueAccessToken is never called. + HandleHelper: nil, + validator: validator, + delegationLifespan: delegationLifespan, + allowedAudiences: []string{testIssuer}, + config: &mockConfig{ + scopeStrategy: fosite.ExactScopeStrategy, + audienceStrategy: fosite.DefaultAudienceMatchingStrategy, + }, + } +} + +// newAccessRequest builds a fosite.AccessRequest for token exchange tests. +func newAccessRequest(t *testing.T, client fosite.Client, formValues url.Values) *fosite.AccessRequest { + t.Helper() + + req := fosite.NewAccessRequest(&session.Session{}) + req.GrantTypes = fosite.Arguments{oauthproto.GrantTypeTokenExchange} + req.Client = client + req.Form = formValues + return req +} + +// defaultClient returns a confidential client configured for token exchange. +func defaultClient() *fosite.DefaultClient { + return &fosite.DefaultClient{ + ID: testAgentClientID, + GrantTypes: fosite.Arguments{oauthproto.GrantTypeTokenExchange}, + Scopes: fosite.Arguments{"openid", "profile"}, + Audience: []string{testIssuer}, + Public: false, + } +} + +// defaultFormValues returns the minimum form values for a valid token exchange request. +func defaultFormValues(t *testing.T, tj *testJWKS) url.Values { + t.Helper() + + token := tj.signToken(t, validClaims(), validExtraClaims()) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } +} + +// actFormValues returns valid token exchange form values whose subject token +// carries the given "act" claim. +func actFormValues(t *testing.T, tj *testJWKS, act any) url.Values { + t.Helper() + + extra := validExtraClaims() + extra["act"] = act + form := defaultFormValues(t, tj) + form.Set("subject_token", tj.signToken(t, validClaims(), extra)) + return form +} + +// nestedActChain builds a chain of depth nested "act" maps, each holding a +// distinct "sub", for exercising the delegation-depth cap. depth must be at +// least 1. +func nestedActChain(depth int) map[string]any { + chain := map[string]any{"sub": fmt.Sprintf("agent-%d", depth-1)} + for i := depth - 2; i >= 0; i-- { + chain = map[string]any{"sub": fmt.Sprintf("agent-%d", i), "act": chain} + } + return chain +} + +func TestTokenExchangeHandler_CanHandleTokenEndpointRequest(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + h := newTestHandler(t, tj, 15*time.Minute) + + t.Run("matches token exchange grant type", func(t *testing.T) { + t.Parallel() + req := fosite.NewAccessRequest(&session.Session{}) + req.GrantTypes = fosite.Arguments{oauthproto.GrantTypeTokenExchange} + assert.True(t, h.CanHandleTokenEndpointRequest(context.Background(), req)) + }) + + t.Run("rejects other grant types", func(t *testing.T) { + t.Parallel() + req := fosite.NewAccessRequest(&session.Session{}) + req.GrantTypes = fosite.Arguments{"client_credentials"} + assert.False(t, h.CanHandleTokenEndpointRequest(context.Background(), req)) + }) +} + +func TestTokenExchangeHandler_CanSkipClientAuth(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + h := newTestHandler(t, tj, 15*time.Minute) + + req := fosite.NewAccessRequest(&session.Session{}) + assert.False(t, h.CanSkipClientAuth(context.Background(), req)) +} + +func TestTokenExchangeHandler_HandleTokenEndpointRequest(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + + tests := []struct { + name string + ctx func(t *testing.T) context.Context + client func() *fosite.DefaultClient + form func(t *testing.T) url.Values + grantTypes fosite.Arguments // Override grant type; nil means oauthproto.GrantTypeTokenExchange + lifespan time.Duration + wantErr bool + wantFositeIs *fosite.RFC6749Error // Check errors.Is against this fosite sentinel + hintContains string // Check the fosite error's Reason/Hint field + check func(t *testing.T, req *fosite.AccessRequest) + }{ + { + name: "valid exchange sets act claim and session", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { t.Helper(); return defaultFormValues(t, tj) }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + + // Verify subject is the user from the subject token. + assert.Equal(t, "user-123", sess.JWTClaims.Subject) + + // Verify the act claim contains the authenticated client's ID. + actClaim, exists := sess.JWTClaims.Extra["act"] + require.True(t, exists, "act claim must be present") + actMap, ok := actClaim.(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + + // Verify user claims were propagated. + assert.Equal(t, "Test User", sess.JWTClaims.Extra["name"]) + assert.Equal(t, "test@example.com", sess.JWTClaims.Extra["email"]) + + // Verify expiry was set on the session. + expiry := sess.GetExpiresAt(fosite.AccessToken) + assert.False(t, expiry.IsZero(), "access token expiry should be set") + }, + }, + { + name: "missing subject_token", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(_ *testing.T) url.Values { + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidRequest, + hintContains: "subject_token", + }, + { + name: "missing subject_token_type", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + token := tj.signToken(t, validClaims(), validExtraClaims()) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidRequest, + hintContains: "subject_token_type", + }, + { + name: "invalid subject_token_type", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + token := tj.signToken(t, validClaims(), validExtraClaims()) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {"urn:ietf:params:oauth:token-type:refresh_token"}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidRequest, + hintContains: "subject_token_type", + }, + { + name: "public client rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { c := defaultClient(); c.Public = true; return c }, + form: func(t *testing.T) url.Values { t.Helper(); return defaultFormValues(t, tj) }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "public", + }, + { + name: "subject token client_id mismatch", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.ID = "different-client" + return c + }, + form: func(t *testing.T) url.Values { t.Helper(); return defaultFormValues(t, tj) }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "different client", + }, + { + name: "invalid subject_token — bad JWT", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(_ *testing.T) url.Values { + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {"not-a-valid-jwt"}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidRequest, + hintContains: "subject token is invalid", + }, + { + name: "unsupported requested_token_type rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("requested_token_type", "urn:ietf:params:oauth:token-type:id_token") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidRequest, + hintContains: "requested_token_type", + }, + { + name: "delegation lifetime capped by subject token expiry", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + // Subject token expires in 5 minutes. + claims := validClaims() + claims.Expiry = jwt.NewNumericDate(time.Now().Add(5 * time.Minute)) + token := tj.signToken(t, claims, validExtraClaims()) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, // Handler's max is 15 min, but subject expires in 5. + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok) + + expiry := sess.GetExpiresAt(fosite.AccessToken) + // The delegated token should expire within ~5 minutes, not 15. + remaining := time.Until(expiry) + assert.Less(t, remaining, 6*time.Minute, + "delegated token lifetime should be capped by subject token expiry") + assert.Greater(t, remaining, 4*time.Minute, + "delegated token lifetime should be close to subject token remaining time") + }, + }, + { + name: "wrong grant type returns ErrUnknownRequest", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { t.Helper(); return defaultFormValues(t, tj) }, + grantTypes: fosite.Arguments{"client_credentials"}, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrUnknownRequest, + }, + { + name: "requested scope not allowed by client", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Scopes = fosite.Arguments{"openid"} // No "admin" scope. + return c + }, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("scope", "admin") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidScope, + hintContains: "scope", + }, + { + name: "scope escalation blocked by subject token", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Scopes = fosite.Arguments{"openid", "profile", "admin"} + return c + }, + form: func(t *testing.T) url.Values { + t.Helper() + // Subject token has scope="openid" only. + claims := validClaims() + extra := validExtraClaims() + extra["scope"] = "openid" + token := tj.signToken(t, claims, extra) + f := url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + f.Set("scope", "admin") // Client has admin, but subject token doesn't. + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidScope, + hintContains: "subject token", + }, + { + name: "may_act claim authorizes this client", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": testAgentClientID} + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + assert.Equal(t, "user-123", sess.JWTClaims.Subject) + actClaim, exists := sess.JWTClaims.Extra["act"] + require.True(t, exists, "act claim must be present") + actMap, ok := actClaim.(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + }, + }, + { + name: "may_act claim does not authorize this client", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": "different-agent"} + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "does not authorize", + }, + { + name: "may_act authorizes cross-client delegation", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := validExtraClaims() + extra["client_id"] = "original-client" // different from the authenticated client + extra["may_act"] = map[string]any{"sub": testAgentClientID} + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + assert.Equal(t, "user-123", sess.JWTClaims.Subject) + actClaim, exists := sess.JWTClaims.Extra["act"] + require.True(t, exists, "act claim must be present") + actMap, ok := actClaim.(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + }, + }, + { + name: "subject token with no may_act and no client_id rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := map[string]any{ + "name": "Test User", + "email": "test@example.com", + } + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "no verifiable client binding", + }, + { + name: "subject token's own act claim is nested under the new act claim", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := validExtraClaims() + extra["act"] = map[string]any{"sub": "original-agent"} + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + actClaim, exists := sess.JWTClaims.Extra["act"] + require.True(t, exists, "act claim must be present") + actMap, ok := actClaim.(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + priorAct, ok := actMap["act"].(map[string]any) + require.True(t, ok, "prior act claim should be nested under the new act claim") + assert.Equal(t, "original-agent", priorAct["sub"]) + }, + }, + { + name: "requested audience not allowed by client", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Audience = []string{"https://allowed.example.com"} + return c + }, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("audience", "https://not-allowed.example.com") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + }, + { + // The subject token is valid only for testIssuer, but the client is + // registered for a second audience. Requesting that second audience + // must be rejected: delegation cannot broaden the resource boundary. + name: "audience escalation blocked by subject token", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Audience = []string{testIssuer, "https://other.example.com"} + return c + }, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("audience", "https://other.example.com") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: server.ErrInvalidTarget, + hintContains: "not covered by the subject token", + }, + { + // The subject token carries both audiences, so requesting one of them + // is within the delegated boundary and must succeed. + name: "audience within subject token granted", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Audience = []string{testIssuer, "https://other.example.com"} + return c + }, + form: func(t *testing.T) url.Values { + t.Helper() + claims := validClaims() + claims.Audience = jwt.Audience{testIssuer, "https://other.example.com"} + token := tj.signToken(t, claims, validExtraClaims()) + f := url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + f.Set("audience", "https://other.example.com") + return f + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + assert.Contains(t, req.GetGrantedAudience(), "https://other.example.com", + "audience covered by the subject token should be granted") + }, + }, + { + name: "valid resource parameter granted as audience", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("resource", testIssuer) + return f + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + audiences := req.GetGrantedAudience() + assert.Contains(t, audiences, testIssuer, + "resource parameter should be granted as audience") + }, + }, + { + name: "resource not in allowed audiences rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("resource", "https://evil.example.com") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + hintContains: "not a registered audience", + }, + { + name: "multiple resource parameters rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Add("resource", testIssuer) + f.Add("resource", "https://other.example.com") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + hintContains: "Multiple resource parameters", + }, + { + name: "resource with invalid URI rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("resource", "not-a-uri") + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + hintContains: "absolute URI", + }, + { + // Server-wide allowedAudiences (set by newTestHandler to testIssuer) + // permits testIssuer, but the client is not registered for it — the + // per-client audience strategy must still reject the resource param. + name: "resource not permitted for this client rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Audience = []string{"https://other.example.com"} + return c + }, + form: func(t *testing.T) url.Values { + t.Helper() + f := defaultFormValues(t, tj) + f.Set("resource", testIssuer) + return f + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: server.ErrInvalidTarget, + hintContains: "not permitted", + }, + { + name: "confidential client without token-exchange grant type rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.GrantTypes = fosite.Arguments{"client_credentials"} + return c + }, + form: func(t *testing.T) url.Values { t.Helper(); return defaultFormValues(t, tj) }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrUnauthorizedClient, + hintContains: "not allowed to use authorization grant", + }, + { + name: "delegation chain at max depth rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := validExtraClaims() + extra["act"] = nestedActChain(maxDelegationDepth) + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "too deep", + }, + { + name: "delegation chain under max depth is nested", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + extra := validExtraClaims() + extra["act"] = nestedActChain(maxDelegationDepth - 1) + token := tj.signToken(t, validClaims(), extra) + return url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + actClaim, exists := sess.JWTClaims.Extra["act"] + require.True(t, exists, "act claim must be present") + actMap, ok := actClaim.(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + assert.Equal(t, nestedActChain(maxDelegationDepth-1), actMap["act"], + "the prior act chain should be nested unchanged under the new act claim") + }, + }, + { + // RFC 8693 Section 4.1 requires act to be a JSON object. Rejecting keeps + // us from re-minting the violation. A non-object array or scalar reaches + // the same parser branch, so the string case covers both. + name: "non-object act rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + return actFormValues(t, tj, "some-agent") + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "act_not_object", + }, + { + name: "non-object nested act rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + return actFormValues(t, tj, map[string]any{"sub": "agent-0", "act": "agent-1"}) + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "nested_act_not_object", + }, + { + // RFC 7519 Section 4.1.2 fixes sub as a string, so a hop carrying a + // non-string sub is non-conformant even though the chain is an object. + name: "non-string sub in act hop rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + return actFormValues(t, tj, map[string]any{"sub": 42}) + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "sub_not_string", + }, + { + // RFC 7519 Section 4.1.1 fixes iss as a string too. sub is kept valid + // here to isolate the reason: the parser reports the first violation it + // finds and checks iss before sub, so a hop with both would report only + // iss_not_string and leave this branch unpinned. + name: "non-string iss in act hop rejected", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + return actFormValues(t, tj, map[string]any{"iss": 42, "sub": "agent-0"}) + }, + lifespan: 15 * time.Minute, + wantErr: true, + wantFositeIs: fosite.ErrInvalidGrant, + hintContains: "iss_not_string", + }, + { + // A JSON-null nested act ends the chain; it asserts no further + // delegation and must not read as malformed. + name: "null nested act is accepted and nested unchanged", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + return actFormValues(t, tj, map[string]any{"sub": "agent-0", "act": nil}) + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + actMap, ok := sess.JWTClaims.Extra["act"].(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + assert.Equal(t, map[string]any{"sub": "agent-0", "act": nil}, actMap["act"], + "the prior act claim should be nested verbatim, null member included") + }, + }, + { + // Tolerated by design: RFC 8693 Section 4.1 only requires an object, and + // sub is conventional rather than mandatory, so an empty hop is + // conformant. Pinned so a future tightening of the shared parser breaks + // a test rather than a caller. + name: "empty act object is accepted", + ctx: func(_ *testing.T) context.Context { return context.Background() }, + client: defaultClient, + form: func(t *testing.T) url.Values { + t.Helper() + return actFormValues(t, tj, map[string]any{}) + }, + lifespan: 15 * time.Minute, + check: func(t *testing.T, req *fosite.AccessRequest) { + t.Helper() + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + actMap, ok := sess.JWTClaims.Extra["act"].(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, actMap["sub"]) + assert.Equal(t, map[string]any{}, actMap["act"]) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + h := newTestHandler(t, tj, tt.lifespan) + ctx := tt.ctx(t) + client := tt.client() + form := tt.form(t) + + req := newAccessRequest(t, client, form) + + // Override grant type if specified in the test case. + if tt.grantTypes != nil { + req.GrantTypes = tt.grantTypes + } + + // Set requested scopes from the form if present. + if scopeStr := form.Get("scope"); scopeStr != "" { + req.SetRequestedScopes(fosite.Arguments{scopeStr}) + } + + // Set requested audience from the form if present. + if audStr := form.Get("audience"); audStr != "" { + req.SetRequestedAudience(fosite.Arguments{audStr}) + } + + err := h.HandleTokenEndpointRequest(ctx, req) + + if tt.wantErr { + require.Error(t, err) + if tt.wantFositeIs != nil { + assert.True(t, errors.Is(err, tt.wantFositeIs), + "expected error to match %v, got: %v", tt.wantFositeIs.ErrorField, err) + } + if tt.hintContains != "" { + var rfcErr *fosite.RFC6749Error + require.True(t, errors.As(err, &rfcErr), "expected fosite RFC6749Error") + assert.Contains(t, rfcErr.Reason(), tt.hintContains, + "fosite error hint should contain %q", tt.hintContains) + } + return + } + + require.NoError(t, err) + if tt.check != nil { + tt.check(t, req) + } + }) + } +} + +func TestTokenExchangeHandler_DefaultAudience(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + + tests := []struct { + name string + allowedAudiences []string + client func() *fosite.DefaultClient // defaults to defaultClient if nil + wantErr bool + hintContains string // defaults to "unambiguous default audience" if empty + wantAudience string + }{ + { + name: "single allowed audience defaults when none requested", + allowedAudiences: []string{testIssuer}, + wantAudience: testIssuer, + }, + { + name: "no allowed audiences configured rejects the request", + allowedAudiences: nil, + wantErr: true, + }, + { + name: "multiple allowed audiences rejects the ambiguous request", + allowedAudiences: []string{testIssuer, "https://other.example.com"}, + wantErr: true, + }, + { + // The sole configured audience is unambiguous server-wide, but the + // client is not registered for it — the per-client check must still + // reject rather than silently granting an unregistered audience. + name: "default audience not permitted for this client rejected", + allowedAudiences: []string{testIssuer}, + client: func() *fosite.DefaultClient { + c := defaultClient() + c.Audience = []string{"https://other.example.com"} + return c + }, + wantErr: true, + hintContains: "not permitted to request a token for the default audience", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // newTestHandler wires the subject-token validator's allowed + // audiences to testIssuer regardless of tt.allowedAudiences — + // that governs the subject token's own "aud" claim, a separate + // concern from the delegated token's audience under test here. + h := newTestHandler(t, tj, 15*time.Minute) + h.allowedAudiences = tt.allowedAudiences + + client := defaultClient + if tt.client != nil { + client = tt.client + } + req := newAccessRequest(t, client(), defaultFormValues(t, tj)) + err := h.HandleTokenEndpointRequest(context.Background(), req) + + if tt.wantErr { + require.Error(t, err) + var rfcErr *fosite.RFC6749Error + require.True(t, errors.As(err, &rfcErr), "expected fosite RFC6749Error") + hintContains := tt.hintContains + if hintContains == "" { + hintContains = "unambiguous default audience" + } + assert.Contains(t, rfcErr.Reason(), hintContains) + return + } + require.NoError(t, err) + assert.Contains(t, req.GetGrantedAudience(), tt.wantAudience) + }) + } +} + +func TestCheckDelegationConsent(t *testing.T) { + t.Parallel() + + const actorID = testAgentClientID + + tests := []struct { + name string + claims *ValidatedClaims + wantErr bool + errContains string + }{ + { + name: "may_act matching actorID accepted", + claims: &ValidatedClaims{MayAct: &MayActClaim{Sub: actorID}}, + }, + { + name: "may_act not matching actorID rejected", + claims: &ValidatedClaims{MayAct: &MayActClaim{Sub: "different-agent"}}, + wantErr: true, + errContains: "does not authorize", + }, + { + // #5989 fix: AllowedDelegateClients must be enforced on the + // may_act path too, not only ExternalActor's — an issuer that + // can emit may_act but has no AllowedActors-equivalent consent + // path (Entra, Okta) would otherwise have no per-client + // containment at all. ExternalIssuer must be set for this check + // to apply at all — see checkDelegationConsent's doc comment for + // why a self-issued may_act (ExternalIssuer == "") skips it. + name: "may_act matching actorID, actorID in AllowedDelegateClients accepted", + claims: &ValidatedClaims{ + MayAct: &MayActClaim{Sub: actorID}, + ExternalIssuer: "https://idp.example.com", + AllowedDelegateClients: []string{"some-other-client", actorID}, + }, + }, + { + name: "may_act matching actorID, actorID not in AllowedDelegateClients rejected", + claims: &ValidatedClaims{ + MayAct: &MayActClaim{Sub: actorID}, + ExternalIssuer: "https://idp.example.com", + AllowedDelegateClients: []string{"some-other-client"}, + }, + wantErr: true, + errContains: "not authorized to exchange subject tokens", + }, + { + // A self-issued may_act (ExternalIssuer unset) has no + // AllowedDelegateClients equivalent, so an empty/absent value + // here must not reject it — unlike the external-issuer case + // below, where the same emptiness is a hard rejection. + name: "may_act matching actorID, ExternalIssuer unset, AllowedDelegateClients not enforced", + claims: &ValidatedClaims{ + MayAct: &MayActClaim{Sub: actorID}, + }, + }, + { + name: "ExternalActor set, actorID in wildcard AllowedDelegateClients accepted", + claims: &ValidatedClaims{ + ExternalActor: "ext-agent", + AllowedDelegateClients: []string{anyDelegateClient}, + }, + }, + { + // Guards the switch ordering in checkDelegationConsent: the + // ExternalActor case must be checked before the client_id cases, + // so a populated (and mismatched) ClientID must not cause + // rejection when ExternalActor is already set. + name: "ExternalActor set with a differing ClientID still accepted", + claims: &ValidatedClaims{ + ExternalActor: "ext-agent", + ClientID: "some-other-client", + AllowedDelegateClients: []string{anyDelegateClient}, + }, + }, + { + // may_act must keep winning even when ExternalActor is also set + // (the external validator only sets ExternalActor when MayAct is + // nil, but checkDelegationConsent's own ordering must not depend + // on that invariant holding). + name: "ExternalActor set but MayAct mismatched is rejected", + claims: &ValidatedClaims{ + ExternalActor: "ext-agent", + MayAct: &MayActClaim{Sub: "different-agent"}, + }, + wantErr: true, + errContains: "does not authorize", + }, + { + // #5989 hardening: AllowedDelegateClients is now required at + // construction (validateTrustedIssuer), so nil/empty must be + // rejected at runtime too, not treated as "any client" — this + // pins the fail-closed default this PR introduces. + name: "ExternalActor set, nil AllowedDelegateClients rejected", + claims: &ValidatedClaims{ + ExternalActor: "ext-agent", + // AllowedDelegateClients intentionally nil. + }, + wantErr: true, + errContains: "not authorized to exchange subject tokens", + }, + { + name: "ExternalActor set, actorID in AllowedDelegateClients accepted", + claims: &ValidatedClaims{ + ExternalActor: "ext-agent", + AllowedDelegateClients: []string{"some-other-client", actorID}, + }, + }, + { + name: "ExternalActor set, actorID not in AllowedDelegateClients rejected", + claims: &ValidatedClaims{ + ExternalActor: "ext-agent", + AllowedDelegateClients: []string{"some-other-client"}, + }, + wantErr: true, + errContains: "not authorized to exchange subject tokens", + }, + { + name: "client_id matching actorID accepted", + claims: &ValidatedClaims{ClientID: actorID}, + }, + { + name: "client_id not matching actorID rejected", + claims: &ValidatedClaims{ClientID: "different-client"}, + wantErr: true, + errContains: "different client", + }, + { + name: "no may_act, no ExternalActor, no client_id rejected", + claims: &ValidatedClaims{}, + wantErr: true, + errContains: "no verifiable client binding", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := checkDelegationConsent(tt.claims, actorID) + if tt.wantErr { + require.Error(t, err) + var rfcErr *fosite.RFC6749Error + require.True(t, errors.As(err, &rfcErr), "expected fosite RFC6749Error") + assert.Contains(t, rfcErr.Reason(), tt.errContains) + return + } + require.NoError(t, err) + }) + } +} + +// newTestHandlerWithValidator creates a Handler wired to a caller-provided +// validator, for tests that need external-issuer delegation coverage that +// newTestHandler's self-issued-only validator cannot produce. +func newTestHandlerWithValidator(validator SubjectTokenValidator) *Handler { + return &Handler{ + validator: validator, + delegationLifespan: 15 * time.Minute, + allowedAudiences: []string{testIssuer}, + config: &mockConfig{ + scopeStrategy: fosite.ExactScopeStrategy, + audienceStrategy: fosite.DefaultAudienceMatchingStrategy, + }, + } +} + +// TestTokenExchangeHandler_ActChainProvenance covers the act-chain shape that +// HandleTokenEndpointRequest builds, for the case a token exchange terminates +// on (ValidatedClaims.ExternalActor set) that TestTokenExchangeHandler_ +// HandleTokenEndpointRequest's self-issued-only table cannot exercise. +func TestTokenExchangeHandler_ActChainProvenance(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + multiValidator := newMultiValidator(t, tj, trustedIssuers) + + // externalToken signs a subject token from the trusted external issuer + // with an allowlisted actor (so ExternalActor is set) and an optional + // prior act chain. The audience includes testIssuer alongside the + // issuer's required ExpectedAudience so the delegated token's + // default-audience grant (testIssuer) stays covered by the subject + // token's own audience (ensureAudienceSubsetOfSubject). + externalToken := func(t *testing.T, priorAct any) string { + t.Helper() + claims := externalClaims() + claims.Audience = jwt.Audience{testExternalAudience, testIssuer} + extra := map[string]any{"azp": "ext-agent"} + if priorAct != nil { + extra["act"] = priorAct + } + return externalJWKS.signToken(t, claims, extra) + } + + requestWith := func(t *testing.T, h *Handler, token string) (*fosite.AccessRequest, error) { + t.Helper() + form := url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + req := newAccessRequest(t, defaultClient(), form) + return req, h.HandleTokenEndpointRequest(context.Background(), req) + } + + t.Run("external delegation nests actor and issuer under the client's act entry", func(t *testing.T) { + t.Parallel() + h := newTestHandlerWithValidator(multiValidator) + + req, err := requestWith(t, h, externalToken(t, nil)) + require.NoError(t, err) + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + assert.Equal(t, testExternalIssuer+"#"+"ext-user-456", sess.JWTClaims.Subject, + "external subject must be qualified with the issuer URL, not copied verbatim") + act, ok := sess.JWTClaims.Extra["act"].(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, act["sub"], "outermost act.sub is always the ToolHive client") + + nested, ok := act["act"].(map[string]any) + require.True(t, ok, "external actor/issuer must be nested under act.act") + assert.Equal(t, "ext-agent", nested["sub"]) + assert.Equal(t, testExternalIssuer, nested["iss"]) + assert.Nil(t, nested["act"], "no prior chain to nest further") + }) + + t.Run("self-issued delegation stays a single level with no external nesting", func(t *testing.T) { + t.Parallel() + h := newTestHandler(t, tj, 15*time.Minute) + + req, err := requestWith(t, h, tj.signToken(t, validClaims(), validExtraClaims())) + require.NoError(t, err) + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + assert.Equal(t, "user-123", sess.JWTClaims.Subject, + "self-issued subject must pass through unchanged, never qualified") + act, ok := sess.JWTClaims.Extra["act"].(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, act["sub"]) + assert.Nil(t, act["act"], "self-issued delegation must not nest an external actor") + }) + + // #5989 fix: an external issuer choosing a "sub" equal to a native + // ToolHive user's UUID must never produce a delegated token whose "sub" + // collides with that UUID — Cedar's extractClientIDFromClaims + // (pkg/authz/authorizers/cedar/core.go) keys on "sub" alone, and native + // ToolHive subjects are UUIDs from UserResolver.ResolveUser, never an + // upstream provider's raw subject value. + t.Run("external subject equal to a plausible native UUID is qualified, not collided", func(t *testing.T) { + t.Parallel() + h := newTestHandlerWithValidator(multiValidator) + + const nativeLikeUUID = "8a34c9e0-6b1a-4e5f-9c2d-1f0a7b3d5e6c" + claims := externalClaims() + claims.Subject = nativeLikeUUID + claims.Audience = jwt.Audience{testExternalAudience, testIssuer} + token := externalJWKS.signToken(t, claims, map[string]any{"azp": "ext-agent"}) + + req, err := requestWith(t, h, token) + require.NoError(t, err) + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + assert.NotEqual(t, nativeLikeUUID, sess.JWTClaims.Subject, + "the qualified subject must never equal the bare external sub, "+ + "or it could collide with a native user sharing that UUID") + assert.Equal(t, testExternalIssuer+"#"+nativeLikeUUID, sess.JWTClaims.Subject) + }) + + // The documented dangerous config (see checkDelegationConsent and + // resolveAllowedActor's doc comments): a trusted external issuer emitting + // a well-formed may_act naming this client bypasses the allowlist + // entirely. Only covered at the checkDelegationConsent unit level until + // now — this exercises it through the full handler. + t.Run("external token carrying may_act still nests the issuer, with no actor to report", func(t *testing.T) { + t.Parallel() + h := newTestHandlerWithValidator(multiValidator) + + claims := externalClaims() + claims.Audience = jwt.Audience{testExternalAudience, testIssuer} + token := externalJWKS.signToken(t, claims, map[string]any{ + "may_act": map[string]any{"sub": testAgentClientID, "iss": testIssuer}, + }) + + req, err := requestWith(t, h, token) + require.NoError(t, err) + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + act, ok := sess.JWTClaims.Extra["act"].(map[string]any) + require.True(t, ok, "act claim must be a map") + assert.Equal(t, testAgentClientID, act["sub"]) + + // The path that bypasses the allowlist entirely still records where + // the delegation came from (ValidatedClaims.ExternalIssuer is set + // unconditionally by validateExternalToken) — it just has no + // client-namespace actor claim to report, since may_act.sub already + // named the delegate directly via the outermost act.sub above. + nested, ok := act["act"].(map[string]any) + require.True(t, ok, "external issuer must be nested even on the may_act path") + assert.Equal(t, testExternalIssuer, nested["iss"]) + _, hasSub := nested["sub"] + assert.False(t, hasSub, "no ExternalActor exists on the may_act path") + }) + + // maxDelegationDepth (10) bounds the prior chain depth + newLevels. The + // external wrapper adds a level of its own (newLevels=2) versus the + // self-issued path (newLevels=1), so the external path's prior chain must + // be one level shallower to fit: + // self-issued: depth 9 accepted (9+1=10); depth 10 rejected (10+1=11>10) + // — see "delegation chain under max depth is nested" / + // "delegation chain at max depth rejected" in + // TestTokenExchangeHandler_HandleTokenEndpointRequest. + // external: depth 8 accepted (8+2=10); depth 9 rejected (9+2=11>10) + // Split into two subtests (rather than one asserting both sides) so a + // failure on the accept half can't hide a regression on the reject half. + t.Run("external delegation at depth 8 fits exactly and preserves the full nested chain", func(t *testing.T) { + t.Parallel() + h := newTestHandlerWithValidator(multiValidator) + + req, err := requestWith(t, h, externalToken(t, nestedActChain(8))) + require.NoError(t, err, "depth 8 plus the external wrapper must fit exactly at maxDelegationDepth") + + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok, "session should be *session.Session") + act, ok := sess.JWTClaims.Extra["act"].(map[string]any) + require.True(t, ok, "act claim must be a map") + external, ok := act["act"].(map[string]any) + require.True(t, ok, "external actor/issuer must be nested under act.act") + assert.Equal(t, "ext-agent", external["sub"]) + assert.Equal(t, testExternalIssuer, external["iss"]) + // The prior chain must be preserved under the external wrapper, not + // dropped or overwritten by it — this is the provenance record + // checkDelegationConsent's doc comment relies on being auditable. + assert.Equal(t, nestedActChain(8), external["act"], + "the prior act chain must be nested under the external wrapper, unchanged") + }) + + t.Run("external delegation at depth 9 is rejected as too deep", func(t *testing.T) { + t.Parallel() + h := newTestHandlerWithValidator(multiValidator) + + _, err := requestWith(t, h, externalToken(t, nestedActChain(9))) + require.Error(t, err) + assert.True(t, errors.Is(err, fosite.ErrInvalidGrant)) + var rfcErr *fosite.RFC6749Error + require.True(t, errors.As(err, &rfcErr), "expected fosite RFC6749Error") + assert.Contains(t, rfcErr.Reason(), "too deep") + }) +} + +// mockConfig implements the tokenExchangeConfig interface for testing. +type mockConfig struct { + scopeStrategy fosite.ScopeStrategy + audienceStrategy fosite.AudienceMatchingStrategy + accessLifespan time.Duration +} + +func (m *mockConfig) GetScopeStrategy(_ context.Context) fosite.ScopeStrategy { + return m.scopeStrategy +} + +func (m *mockConfig) GetAudienceStrategy(_ context.Context) fosite.AudienceMatchingStrategy { + return m.audienceStrategy +} + +func (m *mockConfig) GetAccessTokenLifespan(_ context.Context) time.Duration { + if m.accessLifespan > 0 { + return m.accessLifespan + } + return 15 * time.Minute +} + +// mockAccessTokenStrategy is a minimal oauth2.AccessTokenStrategy that issues +// deterministic opaque tokens without any JWT signing machinery. +type mockAccessTokenStrategy struct { + generateErr error +} + +func (*mockAccessTokenStrategy) AccessTokenSignature(_ context.Context, token string) string { + return "sig-" + token +} + +func (m *mockAccessTokenStrategy) GenerateAccessToken(_ context.Context, _ fosite.Requester) (string, string, error) { + if m.generateErr != nil { + return "", "", m.generateErr + } + return "test-access-token", "test-signature", nil +} + +func (*mockAccessTokenStrategy) ValidateAccessToken(_ context.Context, _ fosite.Requester, _ string) error { + return nil +} + +// mockAccessTokenStorage is a minimal oauth2.AccessTokenStorage that records +// the persisted session signature and never fails. +type mockAccessTokenStorage struct { + createdSig string + createErr error + createdReqs []fosite.Requester +} + +func (m *mockAccessTokenStorage) CreateAccessTokenSession(_ context.Context, signature string, req fosite.Requester) error { + if m.createErr != nil { + return m.createErr + } + m.createdSig = signature + m.createdReqs = append(m.createdReqs, req) + return nil +} + +func (*mockAccessTokenStorage) GetAccessTokenSession(_ context.Context, _ string, _ fosite.Session) (fosite.Requester, error) { + return nil, fosite.ErrNotFound +} + +func (*mockAccessTokenStorage) DeleteAccessTokenSession(_ context.Context, _ string) error { + return nil +} + +// mockHandleHelperConfig implements oauth2.HandleHelperConfigProvider, providing +// both access and refresh token lifespans required by HandleHelper. +type mockHandleHelperConfig struct { + accessLifespan time.Duration + refreshLifespan time.Duration +} + +func (m *mockHandleHelperConfig) GetAccessTokenLifespan(_ context.Context) time.Duration { + if m.accessLifespan > 0 { + return m.accessLifespan + } + return 15 * time.Minute +} + +func (m *mockHandleHelperConfig) GetRefreshTokenLifespan(_ context.Context) time.Duration { + if m.refreshLifespan > 0 { + return m.refreshLifespan + } + return 30 * 24 * time.Hour +} + +// newTestHandlerWithHelper builds a Handler wired with a real oauth2.HandleHelper +// backed by mock strategy/storage, enabling PopulateTokenEndpointResponse tests. +// The delegation lifespan is fixed at 15 minutes; only accessLifespan varies +// across callers. +func newTestHandlerWithHelper(t *testing.T, tj *testJWKS, accessLifespan time.Duration, + strategy *mockAccessTokenStrategy, storage *mockAccessTokenStorage) *Handler { + t.Helper() + + const delegationLifespan = 15 * time.Minute + + validator, err := NewSelfIssuedTokenValidator(tj.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + cfg := &mockConfig{ + scopeStrategy: fosite.ExactScopeStrategy, + audienceStrategy: fosite.DefaultAudienceMatchingStrategy, + accessLifespan: accessLifespan, + } + + return &Handler{ + HandleHelper: &oauth2.HandleHelper{ + AccessTokenStrategy: strategy, + AccessTokenStorage: storage, + Config: &mockHandleHelperConfig{ + accessLifespan: accessLifespan, + }, + }, + validator: validator, + delegationLifespan: delegationLifespan, + allowedAudiences: []string{testIssuer}, + config: cfg, + } +} + +func TestTokenExchangeHandler_PopulateTokenEndpointResponse(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + + t.Run("wrong grant type returns ErrUnknownRequest", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{} + storage := &mockAccessTokenStorage{} + h := newTestHandlerWithHelper(t, tj, 15*time.Minute, strategy, storage) + + req := newAccessRequest(t, defaultClient(), defaultFormValues(t, tj)) + // Override grant type so CanHandleTokenEndpointRequest returns false. + req.GrantTypes = fosite.Arguments{"client_credentials"} + + responder := fosite.NewAccessResponse() + err := h.PopulateTokenEndpointResponse(context.Background(), req, responder) + + require.Error(t, err) + assert.True(t, errors.Is(err, fosite.ErrUnknownRequest), + "expected ErrUnknownRequest, got: %v", err) + // No token should have been issued on this error path. + assert.Empty(t, responder.GetAccessToken()) + assert.Empty(t, storage.createdSig) + }) + + t.Run("client without token-exchange grant returns ErrUnauthorizedClient", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{} + storage := &mockAccessTokenStorage{} + h := newTestHandlerWithHelper(t, tj, 15*time.Minute, strategy, storage) + + // Client is allowed client_credentials only, not token-exchange. + client := defaultClient() + client.GrantTypes = fosite.Arguments{"client_credentials"} + req := newAccessRequest(t, client, defaultFormValues(t, tj)) + req.SetSession(&session.Session{}) + + responder := fosite.NewAccessResponse() + err := h.PopulateTokenEndpointResponse(context.Background(), req, responder) + + require.Error(t, err) + assert.True(t, errors.Is(err, fosite.ErrUnauthorizedClient), + "expected ErrUnauthorizedClient, got: %v", err) + // No token should have been issued. + assert.Empty(t, responder.GetAccessToken()) + assert.Empty(t, storage.createdSig) + }) + + t.Run("successful populate sets issued_token_type and access token", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{} + storage := &mockAccessTokenStorage{} + // Access token lifespan larger than the delegation lifespan so the + // session expiry (capped at delegation lifespan) is the binding lifetime. + h := newTestHandlerWithHelper(t, tj, time.Hour, strategy, storage) + + req := newAccessRequest(t, defaultClient(), defaultFormValues(t, tj)) + require.NoError(t, h.HandleTokenEndpointRequest(context.Background(), req)) + + responder := fosite.NewAccessResponse() + require.NoError(t, h.PopulateTokenEndpointResponse(context.Background(), req, responder)) + + // RFC 8693 Section 2.2.1: response MUST include issued_token_type. + assert.Equal(t, oauthproto.TokenTypeAccessToken, responder.GetExtra("issued_token_type")) + // The mock strategy issues a fixed opaque token; verify it propagated. + assert.Equal(t, "test-access-token", responder.GetAccessToken()) + assert.Equal(t, "bearer", responder.GetTokenType()) + // The session should have been persisted via CreateAccessTokenSession. + assert.Equal(t, "test-signature", storage.createdSig) + require.Len(t, storage.createdReqs, 1) + }) + + t.Run("lifetime capped by subject token expiry", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{} + storage := &mockAccessTokenStorage{} + // Subject token expires in 5 minutes; delegation max is 15; access + // token lifespan default is 1h. The effective issued lifetime must be + // capped at the subject token's remaining life (~5 minutes). + h := newTestHandlerWithHelper(t, tj, time.Hour, strategy, storage) + + claims := validClaims() + claims.Expiry = jwt.NewNumericDate(time.Now().Add(5 * time.Minute)) + token := tj.signToken(t, claims, validExtraClaims()) + form := url.Values{ + "grant_type": {oauthproto.GrantTypeTokenExchange}, + "subject_token": {token}, + "subject_token_type": {oauthproto.TokenTypeAccessToken}, + } + + req := newAccessRequest(t, defaultClient(), form) + require.NoError(t, h.HandleTokenEndpointRequest(context.Background(), req)) + + responder := fosite.NewAccessResponse() + require.NoError(t, h.PopulateTokenEndpointResponse(context.Background(), req, responder)) + + // getExpiresIn in HandleHelper computes expires_in from the session's + // ExpiresAt minus now. The capping logic in PopulateTokenEndpointResponse + // sets that ExpiresAt from the delegation lifetime computation (min of + // subject remaining and delegationLifespan). Subject expires in ~5 min, + // delegationLifespan is 15 min, so ExpiresIn should be ~5 min (< 6 min). + expiresIn := responder.ToMap()["expires_in"] + require.NotNil(t, expiresIn) + // fosite serializes expires_in as an int (seconds). Tolerate the type. + var secs int64 + switch v := expiresIn.(type) { + case int: + secs = int64(v) + case int64: + secs = v + case float64: + secs = int64(v) + default: + t.Fatalf("unexpected expires_in type %T: %v", expiresIn, expiresIn) + } + assert.Greater(t, secs, int64((4 * time.Minute).Seconds()), + "expires_in should be capped near subject token remaining (~5m), got %d", secs) + assert.Less(t, secs, int64((6 * time.Minute).Seconds()), + "expires_in should be capped near subject token remaining (~5m), got %d", secs) + }) + + t.Run("AccessTokenStrategy error propagates", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{ + generateErr: errors.New("strategy failure"), + } + storage := &mockAccessTokenStorage{} + h := newTestHandlerWithHelper(t, tj, time.Hour, strategy, storage) + + req := newAccessRequest(t, defaultClient(), defaultFormValues(t, tj)) + require.NoError(t, h.HandleTokenEndpointRequest(context.Background(), req)) + + responder := fosite.NewAccessResponse() + err := h.PopulateTokenEndpointResponse(context.Background(), req, responder) + + require.Error(t, err) + assert.Contains(t, err.Error(), "strategy failure") + // No session should have been persisted on failure. + assert.Empty(t, storage.createdSig) + }) + + t.Run("stale session expiry fails closed instead of falling back to full lifespan", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{} + storage := &mockAccessTokenStorage{} + h := newTestHandlerWithHelper(t, tj, time.Hour, strategy, storage) + + req := newAccessRequest(t, defaultClient(), defaultFormValues(t, tj)) + require.NoError(t, h.HandleTokenEndpointRequest(context.Background(), req)) + + // Simulate the narrow window where the session's bound expiry (set by + // HandleTokenEndpointRequest) elapses before PopulateTokenEndpointResponse runs. + sess, ok := req.GetSession().(*session.Session) + require.True(t, ok) + sess.SetExpiresAt(fosite.AccessToken, time.Now().Add(-time.Second)) + + responder := fosite.NewAccessResponse() + err := h.PopulateTokenEndpointResponse(context.Background(), req, responder) + + require.Error(t, err) + assert.True(t, errors.Is(err, fosite.ErrInvalidGrant), + "expected ErrInvalidGrant, got: %v", err) + assert.Empty(t, responder.GetAccessToken()) + assert.Empty(t, storage.createdSig) + }) + + t.Run("AccessTokenStorage error propagates", func(t *testing.T) { + t.Parallel() + + strategy := &mockAccessTokenStrategy{} + storage := &mockAccessTokenStorage{ + createErr: errors.New("storage failure"), + } + h := newTestHandlerWithHelper(t, tj, time.Hour, strategy, storage) + + req := newAccessRequest(t, defaultClient(), defaultFormValues(t, tj)) + require.NoError(t, h.HandleTokenEndpointRequest(context.Background(), req)) + + responder := fosite.NewAccessResponse() + err := h.PopulateTokenEndpointResponse(context.Background(), req, responder) + + require.Error(t, err) + assert.Contains(t, err.Error(), "storage failure") + // issued_token_type must NOT be set when token issuance failed. + assert.Nil(t, responder.GetExtra("issued_token_type")) + }) +} diff --git a/pkg/authserver/server/tokenexchange/multi_issuer_validator.go b/pkg/authserver/server/tokenexchange/multi_issuer_validator.go new file mode 100644 index 0000000000..62cb4ede03 --- /dev/null +++ b/pkg/authserver/server/tokenexchange/multi_issuer_validator.go @@ -0,0 +1,1004 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/url" + "slices" + "strings" + "sync" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/go-jose/go-jose/v4/jwt" + "github.com/lestrrat-go/httprc/v3" + "github.com/lestrrat-go/jwx/v3/jwk" + + "github.com/stacklok/toolhive/pkg/networking" + "github.com/stacklok/toolhive/pkg/oauthproto" +) + +const ( + // httpTimeout is the timeout for HTTP requests to external OIDC endpoints. + httpTimeout = 10 * time.Second + + // maxResponseBodySize is the maximum size of HTTP response bodies read + // from external OIDC discovery documents AND JWKS fetches (1 MiB). This + // prevents resource exhaustion from unexpectedly large responses. The + // discovery read enforces it directly via io.LimitReader; the JWKS read + // goes through jwx instead, so it is enforced by wrapping every issuer's + // *http.Client transport with limitedBodyTransport — see where that + // client is built in NewMultiIssuerTokenValidator. + maxResponseBodySize = 1 << 20 + + // maxJWKSKeys caps the number of keys accepted from an external JWKS to + // prevent CPU amplification from a hostile endpoint serving many keys. + maxJWKSKeys = 100 + + // minKidRefreshInterval bounds how often refreshOnUnknownKid forces an + // issuer's jwk.Cache to fetch its JWKS ahead of jwx's own background + // refresh schedule. verifySignature's kidMatched check runs before the + // subject token's signature is trusted, so without this floor a client + // presenting a syntactically valid JWT that merely names a made-up kid + // could force a fresh fetch to the external IdP on every attempt. + minKidRefreshInterval = 30 * time.Second + + // jwksFetchFailureBackoff bounds how often ensureRegistered retries a + // JWKS fetch for an issuer that has never yet succeeded. Key resolution + // runs before the subject token's signature is checked, so without this + // an authenticated client holding the token-exchange grant could force a + // real outbound fetch to a broken external IdP on every single request. + jwksFetchFailureBackoff = 30 * time.Second + + // jwksRefreshInterval is the fixed interval at which each issuer's + // jwk.Cache re-fetches its JWKS in the background. It is passed to + // Register via jwk.WithConstantInterval, which makes the resource + // ignore the response's Cache-Control/Expires headers entirely rather + // than merely bounding them with WithMaxInterval — deliberately: + // absent this override, httprc derives the interval from those + // headers, clamped to [DefaultMinInterval, DefaultMaxInterval] = + // [15m, 30 days], so a hostile or misconfigured issuer could otherwise + // extend our own key-retention window up to a month simply by setting + // a long max-age. + jwksRefreshInterval = 5 * time.Minute + + // defaultActorClaim is the claim read to identify the client that + // requested an external subject token, when a TrustedIssuer does not + // configure ActorClaim. This matches Microsoft Entra v2 and many other + // OIDC providers' convention for the authorized-party claim. + // + // "client_id" is the normative claim: RFC 8693 §4.3 names it as the + // party a token was issued to, and RFC 9068 §2.2 makes it a REQUIRED + // access-token claim. "azp" is the default here only because Entra, + // Okta, and other major providers omit "client_id" from their access + // tokens in practice — "azp" itself appears nowhere in RFC 9068; it is + // an OIDC Core claim defined for ID tokens, not access tokens. + defaultActorClaim = "azp" + + // externalClockSkewLeeway widens "nbf"/"iat"/"exp" acceptance for external + // subject tokens to tolerate clock skew between the external IdP and + // ToolHive. validateExternalToken independently rejects an expired + // subject token by comparing its exp against time.Now() with zero + // tolerance, so this leeway only protects against spurious + // not-yet-valid/issued-in-the-future rejections — it never causes an + // expired token to be accepted. + externalClockSkewLeeway = 60 * time.Second +) + +// actorClaimsNotInExtra are claims assignClaim drops or reroutes, so they +// never reach ValidatedClaims.Extra. Only "client_id" has an explicit +// fallback in resolveAllowedActor; configuring ActorClaim to any of these +// would make every external token look like it's missing the claim. +var actorClaimsNotInExtra = []string{ + "sub", "iss", "aud", "exp", "iat", "nbf", "jti", "name", "email", "scope", "scp", "may_act", +} + +// Compile-time check that MultiIssuerTokenValidator implements SubjectTokenValidator. +var _ SubjectTokenValidator = (*MultiIssuerTokenValidator)(nil) + +// TrustedIssuer configures an external OIDC issuer whose tokens are +// accepted as subject tokens during token exchange. +// +// This type is reused verbatim as the wire schema for +// authserver.RunConfig.TrustedIssuers (deliberately, to avoid a parallel +// type that drifts — see the go-style rule against that). Its JSON/YAML +// tags are therefore part of the serialized RunConfig, which is reflected +// into docs/server/swagger.*; adding, renaming, or retagging a field here is +// a schema change, not a purely internal one. +type TrustedIssuer struct { + // IssuerURL is the expected "iss" claim value (exact match). + IssuerURL string `json:"issuer_url" yaml:"issuer_url"` + // ExpectedAudience is the expected "aud" claim value that must appear + // in the token's audience list (a resource/API identifier, not a + // client ID — required, but not enforced; see looksLikeResourceIdentifier). + // See docs/arch/17-token-exchange-delegation.md ("ID/access-token + // discrimination") for why and its limits. + ExpectedAudience string `json:"expected_audience" yaml:"expected_audience"` + // JWKSURL is the URL to fetch the issuer's JSON Web Key Set from. + // If empty, it is resolved via OIDC discovery at {IssuerURL}/.well-known/openid-configuration. + JWKSURL string `json:"jwks_url,omitempty" yaml:"jwks_url,omitempty"` + // InsecureAllowHTTP permits plain-HTTP OIDC discovery and JWKS fetches + // for THIS issuer only. Development and testing only — never set in + // production. Does not relax the private-IP guard; see AllowPrivateIPs. + // Deliberately per-issuer: this server's own InsecureAllowHTTP must not + // silently permit plaintext discovery for every trusted external issuer + // too — a network attacker who can intercept that traffic could + // substitute a JWKS and forge subject tokens for that issuer's + // namespace. + InsecureAllowHTTP bool `json:"insecure_allow_http,omitempty" yaml:"insecure_allow_http,omitempty"` + // AllowPrivateIPs permits OIDC discovery and JWKS fetches for THIS + // issuer to resolve to a private or loopback address. Use only when the + // issuer is hosted inside the same cluster and has no public endpoint. + AllowPrivateIPs bool `json:"allow_private_ips,omitempty" yaml:"allow_private_ips,omitempty"` + // ActorClaim names the claim identifying the client that requested the + // subject token from THIS EXTERNAL ISSUER (used by AllowedActors below). + // Values are in the external issuer's namespace, NOT ToolHive client + // IDs. Defaults to "azp"; use "appid" for Microsoft Entra v1, "cid" for + // Okta. The special value "client_id" reads ValidatedClaims.ClientID + // instead of Extra (assignClaim routes it to that field) — it is still + // the external token's client_id claim, not a ToolHive one. + ActorClaim string `json:"actor_claim,omitempty" yaml:"actor_claim,omitempty"` + // AllowedActors is the allowlist of ActorClaim values authorized to + // exchange a subject token from this issuer when it carries no + // "may_act" claim; empty means only may_act-bearing tokens are + // accepted. By itself names no ToolHive client — see + // AllowedDelegateClients and docs/arch/17-token-exchange-delegation.md + // ("Accepted limitations" #1). + AllowedActors []string `json:"allowed_actors,omitempty" yaml:"allowed_actors,omitempty"` + // AllowedDelegateClients restricts which ToolHive client IDs may + // exchange a subject token from this issuer, for BOTH consent paths. + // Required (validateTrustedIssuer rejects empty/absent); "*" permits + // any confidential client holding the grant. See + // docs/arch/17-token-exchange-delegation.md ("Accepted limitations" #1). + //nolint:lll // field tags require full JSON+YAML names + AllowedDelegateClients []string `json:"allowed_delegate_clients,omitempty" yaml:"allowed_delegate_clients,omitempty"` +} + +// MultiIssuerTokenValidator validates subject tokens from the authorization +// server itself or from configured external OIDC issuers, delegating +// self-issued tokens to SelfIssuedTokenValidator and resolving external +// issuers' JWKS (via OIDC discovery if needed) to verify signature and +// claims. +// +// A valid signature and audience alone would authorize ToolHive as a +// resource, not any particular client, as a delegate — a confused-deputy +// risk (CWE-863). validateExternalToken therefore requires a "may_act" +// claim or a matching AllowedActors entry (surfaced as +// ValidatedClaims.ExternalActor) before returning successfully. See +// docs/arch/17-token-exchange-delegation.md for the full consent-signal +// precedence and trust model. +type MultiIssuerTokenValidator struct { + selfIssuer string + selfValidator *SelfIssuedTokenValidator + issuers map[string]*externalIssuerConfig +} + +// externalIssuerConfig holds the configuration and cached state for an external +// OIDC issuer. The embedded TrustedIssuer is treated as immutable after +// construction: resolveAllowedActor reads TrustedIssuer.AllowedActors (and +// discoverJWKSURL/fetchJWKS read httpClient) on every validation without +// holding mu, since mu protects JWKS state only. Mutating TrustedIssuer +// fields in place after NewMultiIssuerTokenValidator returns is a data race. +type externalIssuerConfig struct { + TrustedIssuer + + // httpClient is dedicated to this issuer, built once at construction + // time from its own InsecureAllowHTTP/AllowPrivateIPs. A single + // validator-wide client couldn't enforce per-issuer SSRF/transport + // policy: http.Client.CheckRedirect and Transport.DialContext have no + // way to know which issuer's fetch they are guarding. + httpClient *http.Client + + // jwksCache is this issuer's own jwk.Cache, registered with httpClient + // above via jwk.WithHTTPClient (see registerOrRefresh). A cache per + // issuer, rather than one shared across every configured issuer, is what + // makes two issuers resolving to the same jwks_url (e.g. two Microsoft + // Entra v1 tenants, which share one tenant-independent JWKS endpoint) a + // non-event: httprc keys a cached resource by URL alone and only honors + // jwk.WithHTTPClient on a URL's first Register call, so a shared cache + // would have the second such issuer silently inherit the first one's + // *http.Client — defeating InsecureAllowHTTP/AllowPrivateIPs's per-issuer + // guarantee for it. Splitting the cache per issuer makes that collision + // unrepresentable instead of guarding against it. + jwksCache *jwk.Cache + + mu sync.Mutex + // jwksURL is resolved from OIDC discovery, or copied from + // TrustedIssuer.JWKSURL when hand-configured. Once set, it is never + // cleared: unlike the JWKS document itself (which jwksCache keeps fresh + // on its own schedule — see ensureRegistered), a change to + // the *endpoint URL* an issuer serves its keys from requires a process + // restart to pick up. Neither Microsoft Entra nor Okta documents rotating + // that URL, only the keys served at it, and every other TrustedIssuer + // field already requires a restart to take effect, so this is + // consistent with the rest of this config's lifecycle. + jwksURL string + // fetched is true once at least one JWKS fetch has succeeded for this + // issuer since process start. Until then, ensureRegistered forces a + // synchronous fetch on every call, since there is no cached value yet to + // fall back on and jwx's own background schedule offers no way to wait + // for its result. + fetched bool + // lastKidRefresh is the last time refreshOnUnknownKid forced a fetch; + // see minKidRefreshInterval. + lastKidRefresh time.Time + // fetchFailedAt and fetchErr gate retries once registered but never fetched: + // key resolution runs before signature verification, so without this an + // authenticated client can otherwise drive one real outbound fetch per + // token-exchange request just by naming an issuer whose endpoint is + // down. fetchErr is served directly while the gate is closed, so the + // caller still gets a specific error instead of a generic "try later". + // Not consulted once fetched is true: a healthy issuer, or one serving + // stale-but-valid keys through a later background-refresh failure, must + // never be gated. + fetchFailedAt time.Time + fetchErr error +} + +// limitedBodyTransport wraps an http.RoundTripper to cap every response body +// at max bytes, via http.MaxBytesReader rather than io.LimitReader: the +// latter truncates silently, which would let a caller parse a cut-off JWKS +// document as if it were complete, where the former surfaces a +// *http.MaxBytesError instead. Its error text ("http: request body too +// large") is written for the request-body case MaxBytesReader was designed +// for, so it reads oddly for a capped response — an accepted rough edge +// rather than justifying a custom ReadCloser. +// +// The nil first argument (http.ResponseWriter) is safe: MaxBytesReader only +// reaches it through a type assertion (`l.w.(requestTooLarger)`) used to tell +// a real server connection to close early, which — on a nil interface value +// — safely evaluates to false rather than panicking (net/http/request.go). +type limitedBodyTransport struct { + base http.RoundTripper + max int64 +} + +// RoundTrip delegates to base and then caps the returned body. The cap +// applies to every response, including a non-2xx one whose body the caller +// doesn't intend to parse — draining it is only ever io.Discard-ed, not +// unbounded, so this errs on the safe side rather than special-casing status +// codes. +func (t *limitedBodyTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil { + return nil, err + } + resp.Body = http.MaxBytesReader(nil, resp.Body, t.max) + return resp, nil +} + +// NewMultiIssuerTokenValidator creates a validator that accepts tokens from +// the authorization server itself and from the provided trusted external +// issuers. Returns an error if selfValidator is nil, selfIssuer is empty, +// any TrustedIssuer is invalid (see validateTrustedIssuer), or an issuer's +// dedicated HTTP client cannot be built (see newExternalIssuerConfig for +// why each issuer gets its own). +func NewMultiIssuerTokenValidator( + selfValidator *SelfIssuedTokenValidator, + selfIssuer string, + trustedIssuers []TrustedIssuer, +) (*MultiIssuerTokenValidator, error) { + if selfValidator == nil { + return nil, errors.New("selfValidator must not be nil") + } + if selfIssuer == "" { + return nil, errors.New("selfIssuer must not be empty") + } + + issuers := make(map[string]*externalIssuerConfig, len(trustedIssuers)) + for _, ti := range trustedIssuers { + if err := validateTrustedIssuer(ti, selfIssuer, issuers); err != nil { + return nil, err + } + if len(ti.AllowedActors) == 0 { + slog.Warn("Trusted issuer has no allowed actors configured; "+ + "only may_act-bearing subject tokens from it will be accepted", + "issuer", ti.IssuerURL, + ) + } + if !looksLikeResourceIdentifier(ti.ExpectedAudience) { + slog.Warn("Trusted issuer's expected_audience does not look like a resource identifier "+ + "(no URI scheme, e.g. \"https://\" or \"api://\"); it may actually be a client ID. "+ + "rejectIDTokenClaims only catches an ID token that carries at_hash/c_hash, which OIDC "+ + "does not require, so this issuer's audience check may not distinguish an ID token "+ + "from an access token", + "issuer", ti.IssuerURL, "expected_audience", ti.ExpectedAudience, + ) + } + + issuerConfig, err := newExternalIssuerConfig(ti) + if err != nil { + return nil, err + } + issuers[ti.IssuerURL] = issuerConfig + } + + return &MultiIssuerTokenValidator{ + selfIssuer: selfIssuer, + selfValidator: selfValidator, + issuers: issuers, + }, nil +} + +// newExternalIssuerConfig builds the *externalIssuerConfig for a single +// already-validated TrustedIssuer: a dedicated HTTP client (scoped to that +// issuer's own InsecureAllowHTTP/AllowPrivateIPs), its body-size-capped +// transport, and its own jwk.Cache. Called once per issuer from +// NewMultiIssuerTokenValidator's constructor loop, after validateTrustedIssuer +// and the startup warnings have already run for ti. +func newExternalIssuerConfig(ti TrustedIssuer) (*externalIssuerConfig, error) { + // Clone AllowedActors and AllowedDelegateClients so a caller mutating + // their original slices in place (e.g. a future config reload) cannot + // race with the unsynchronized reads in resolveAllowedActor and + // validateExternalToken, which run on every validation without + // holding externalIssuerConfig.mu (that mutex guards JWKS state only). + ti.AllowedActors = slices.Clone(ti.AllowedActors) + ti.AllowedDelegateClients = slices.Clone(ti.AllowedDelegateClients) + + // Deliberately networking.NewHttpClientBuilder(), not + // NewHostScopedClientBuilder: that helper ORs + // INSECURE_DISABLE_URL_VALIDATION and an auto-localhost exemption into + // BOTH the HTTP-scheme and private-IP gates, so an unrelated env var — + // or a trusted issuer that merely happens to be on localhost — would + // silently widen AllowPrivateIPs regardless of what the operator set, + // defeating the point of splitting the two flags per issuer. + // + // Keep-alive connections are disabled: this client dials jwks_uri, a host taken + // from an untrusted discovery document, only on the fixed + // jwksRefreshInterval schedule plus occasional on-demand refreshes — + // no hot path here to trade the per-dial SSRF check away for. + httpClient, err := networking.NewHttpClientBuilder(). + WithInsecureAllowHTTP(ti.InsecureAllowHTTP). + WithPrivateIPs(ti.AllowPrivateIPs). + WithTimeout(httpTimeout). + WithDisableKeepAlives(true). + Build() + if err != nil { + return nil, fmt.Errorf("issuer_url %q: failed to build HTTP client: %w", ti.IssuerURL, err) + } + // Guard against a discovery/JWKS redirect hop landing on a + // different, unvetted host — the same policy the transparent proxy + // data path applies to a response derived from an untrusted remote + // server (see SameHostRedirectPolicy's doc comment). + httpClient.CheckRedirect = networking.SameHostRedirectPolicy() + + // Cap every response body this client reads at maxResponseBodySize — + // discovery already enforces this itself via io.LimitReader + // (discoverJWKSURL), but the JWKS fetch is handed to jwx's jwk.Cache + // below, which has no equivalent cap of its own (httprc.MaxBufferSize + // is ~1000 MiB, and its transformer does an unbounded io.ReadAll under + // that ceiling before parsing). Wrapped OUTSIDE httpClient.Transport + // (which Build() always sets — see networking's builder) so the + // private-IP dial guard and ValidatingTransport's scheme check still + // run first, on the inner, unwrapped transport. + httpClient.Transport = &limitedBodyTransport{ + base: httpClient.Transport, + max: maxResponseBodySize, + } + + // One jwk.Cache per issuer (see externalIssuerConfig.jwksCache's doc + // comment for why), each running its own background worker pool + // (jwk.NewCache -> httprc.Client.Start) for the life of the process. + // WithWorkers(1) caps that pool to one worker per issuer instead of + // httprc's default five — budget roughly three goroutines per issuer + // including its controller loop and wait-group waiter. + // context.Background() is deliberate: there's no per-call context to + // root this in, and the loop is meant to outlive any single call, + // stopped only via jwk.Cache.Shutdown — which nothing here calls, + // matching pkg/auth/token.go's TokenValidator. + jwksCache, err := jwk.NewCache(context.Background(), httprc.NewClient(httprc.WithWorkers(1))) + if err != nil { + return nil, fmt.Errorf("issuer_url %q: failed to create JWKS cache: %w", ti.IssuerURL, err) + } + + return &externalIssuerConfig{ + TrustedIssuer: ti, + jwksURL: ti.JWKSURL, + httpClient: httpClient, + jwksCache: jwksCache, + }, nil +} + +// ValidateTrustedIssuers runs every structural check +// NewMultiIssuerTokenValidator performs on trustedIssuers — required fields, +// self-issuer collision, duplicate issuers, and ActorClaim reachability — +// without constructing a validator or any per-issuer HTTP client. Config +// validation calls this to fail before the live upstream DCR registration +// and storage creation that run between RunConfig.Validate and server +// construction; NewMultiIssuerTokenValidator repeats the same checks at +// server startup as defence in depth. Both route through validateTrustedIssuer, +// so the two can't drift out of sync. +func ValidateTrustedIssuers(trustedIssuers []TrustedIssuer, selfIssuer string) error { + issuers := make(map[string]*externalIssuerConfig, len(trustedIssuers)) + for _, ti := range trustedIssuers { + if err := validateTrustedIssuer(ti, selfIssuer, issuers); err != nil { + return err + } + issuers[ti.IssuerURL] = &externalIssuerConfig{TrustedIssuer: ti} + } + return nil +} + +// Validate parses the raw JWT to extract the issuer claim, then routes validation +// to either the self-issued validator or the appropriate external issuer validator. +// Returns an error if the issuer is not trusted. +func (v *MultiIssuerTokenValidator) Validate(ctx context.Context, rawToken string) (*ValidatedClaims, error) { + // Parse the JWT without verification to peek at the issuer claim. + // This is safe because we verify the signature in a subsequent step. + issuer, err := peekIssuer(rawToken) + if err != nil { + return nil, fmt.Errorf("failed to determine token issuer: %w", err) + } + + if issuer == v.selfIssuer { + return v.selfValidator.Validate(ctx, rawToken) + } + + issuerConfig, ok := v.issuers[issuer] + if !ok { + return nil, fmt.Errorf("untrusted issuer: %q", issuer) + } + + return v.validateExternalToken(ctx, rawToken, issuerConfig) +} + +// validateExternalToken verifies a JWT from a trusted external issuer by +// fetching the issuer's JWKS (with caching) and validating the signature and claims. +func (v *MultiIssuerTokenValidator) validateExternalToken( + ctx context.Context, + rawToken string, + issuerConfig *externalIssuerConfig, +) (*ValidatedClaims, error) { + parsedToken, err := jwt.ParseSigned(rawToken, allowedSignatureAlgorithms) + if err != nil { + return nil, fmt.Errorf("subject token is not a valid JWT: %w", err) + } + + jwks, err := v.lookupJWKS(ctx, issuerConfig) + if err != nil { + return nil, fmt.Errorf("failed to fetch JWKS for issuer %s: %w", issuerConfig.IssuerURL, err) + } + + standardClaims, extraClaims, kidMatched, err := verifySignature(parsedToken, jwks) + if err != nil && !kidMatched { + // The token's kid isn't among the keys we have cached — possibly a + // legitimate rotation the issuer's cache hasn't caught up with yet + // (jwx's own background refresh floor is 15 minutes by default). + // Force an immediate re-fetch and retry once before giving up; a + // spoofed-kid attempt (kidMatched true) never reaches this branch, + // since a refresh can't change what signature the token was made + // with. + v.refreshOnUnknownKid(ctx, issuerConfig) + if refreshedJWKS, lookupErr := v.lookupJWKS(ctx, issuerConfig); lookupErr == nil { + standardClaims, extraClaims, _, err = verifySignature(parsedToken, refreshedJWKS) + } + } + if err != nil { + return nil, err + } + + // Validate issuer and audience, tolerating externalClockSkewLeeway on + // nbf/iat/exp. Expiry is enforced strictly (no leeway) below. + expected := jwt.Expected{ + Issuer: issuerConfig.IssuerURL, + AnyAudience: jwt.Audience{issuerConfig.ExpectedAudience}, + } + if err := standardClaims.ValidateWithLeeway(expected, externalClockSkewLeeway); err != nil { + return nil, fmt.Errorf("subject token claims validation failed: %w", err) + } + + // Subject is required for delegation. + if standardClaims.Subject == "" { + return nil, fmt.Errorf("subject token is missing required 'sub' claim") + } + + // Expiry is required so the delegated token can be bounded by the subject + // token's remaining lifetime. + if standardClaims.Expiry == nil { + return nil, errors.New("subject token is missing required 'exp' claim") + } + + // An expired subject token is never acceptable, despite the leeway + // above: it cannot bound the delegated token's lifetime. + if standardClaims.Expiry.Time().Before(time.Now()) { + return nil, errors.New("subject token has expired") + } + + // Reject a subject token that is actually an ID token — see + // rejectIDTokenClaims's doc comment. + if err := rejectIDTokenClaims(extraClaims); err != nil { + return nil, err + } + + // selfIssuer (not issuerConfig.IssuerURL) is the correct comparison for + // may_act's optional "iss": that member identifies the namespace of + // may_act.sub, always compared against a ToolHive client ID regardless + // of which issuer validated the surrounding token. requireIss is true + // here — unlike the self-issued path, the external path cannot leave + // "iss" optional; see validateMayActShape's doc comment. + if err := validateMayActShape(extraClaims, v.selfIssuer, true); err != nil { + return nil, err + } + + claims := buildValidatedClaims(standardClaims, extraClaims) + + // Recorded for every external token, even a may_act-bearing one that + // leaves ExternalActor unset — without this, a may_act-authorized + // external token would be indistinguishable from a self-issued exchange + // downstream, which is backwards: this is the path that most needs an + // audit trail. Set from issuerConfig, already matched against the + // validated "iss" claim above — never from token content. + claims.ExternalIssuer = issuerConfig.IssuerURL + + // Set unconditionally, including for a may_act-bearing token: an issuer + // that can emit may_act but has no way to populate AllowedActors (e.g. + // Entra, Okta) would otherwise have no per-client containment on that + // path (see #5989). checkDelegationConsent enforces it against the + // authenticated client; this validator never sees that client ID. + claims.AllowedDelegateClients = issuerConfig.AllowedDelegateClients + + // A may_act claim is authoritative and enforced by checkDelegationConsent + // against the authenticated client — validateMayActShape above has + // already confirmed its own "iss" (if any) names this server, so + // may_act.sub is in ToolHive's own client namespace. AllowedActors below + // is skipped entirely when MayAct is set. Otherwise, the resolved actor + // claim (even "client_id") is in the EXTERNAL issuer's namespace and + // must match AllowedActors — it can never be compared against the + // authenticated ToolHive client directly. + if claims.MayAct == nil { + actor, err := resolveAllowedActor(issuerConfig, claims) + if err != nil { + return nil, err + } + claims.ExternalActor = actor + } + + return claims, nil +} + +// ensureRegistered resolves issuerConfig.jwksURL — via OIDC discovery on +// first use — and registers it with issuerConfig's own jwk.Cache (see +// registerOrRefresh for the Register-vs-Refresh decision). Discovery +// happens at most once per issuer for the life of the process; see +// externalIssuerConfig's jwksURL doc comment. +// +// The JWKS fetch is retried until one succeeds (fetched stays false +// otherwise), gated by jwksFetchFailureBackoff: key resolution runs before +// the subject token's signature is checked, so without this gate an +// authenticated client could force a real outbound attempt on every +// request to an issuer whose endpoint is down. The gate only applies +// before the first successful fetch — once fetched is true, a healthy +// issuer or one serving stale-but-valid keys through a later refresh +// failure is unaffected. While closed, the last error is replayed +// directly rather than retried. +// +// Every background refresh, not just the first request-triggered one, +// goes through the maxResponseBodySize-capped transport (see +// limitedBodyTransport) despite having no request or client authentication +// of its own — without that cap a compromised issuer could force an +// oversized allocation on every autonomous refresh, indefinitely. +func (v *MultiIssuerTokenValidator) ensureRegistered(ctx context.Context, issuerConfig *externalIssuerConfig) error { + issuerConfig.mu.Lock() + defer issuerConfig.mu.Unlock() + + if issuerConfig.fetched { + return nil + } + if time.Since(issuerConfig.fetchFailedAt) < jwksFetchFailureBackoff { + return issuerConfig.fetchErr + } + + if err := v.registerOrRefresh(ctx, issuerConfig); err != nil { + issuerConfig.fetchErr = err + issuerConfig.fetchFailedAt = time.Now() + return err + } + issuerConfig.fetched = true + issuerConfig.fetchErr = nil + return nil +} + +// registerOrRefresh performs the actual discovery/registration/fetch +// attempt for issuerConfig; ensureRegistered holds issuerConfig.mu across +// this call, single-flighting it per issuer so concurrent validations +// don't pile up N redundant fetches. +// +// Whether jwksURL is already registered is asked of +// issuerConfig.jwksCache.IsRegistered directly, never remembered in a +// field: Register's own registration step is a channel send to the +// cache's backend goroutine and can fail after enqueue but before receipt +// (context deadline), in which case nothing was actually registered. A +// locally remembered "we called Register" flag can't distinguish that +// from "registered, only the fetch failed", and would wrongly keep +// retrying via Refresh — which errors on a URL the cache never heard of — +// forever after. Asking the cache directly is authoritative either way. +// +// IsRegistered makes no network call but isn't free: it's a round-trip +// over that same channel, so it blocks if the backend is busy and returns +// false (not an error) on context expiry. Its own timeout below matters +// for that reason — a false from an expired context just routes to +// Register, whose "already registered" error is transient and absorbed by +// ensureRegistered's backoff gate. +func (v *MultiIssuerTokenValidator) registerOrRefresh(ctx context.Context, issuerConfig *externalIssuerConfig) error { + // Detach from the caller's request context throughout this function: net/http + // cancels ctx when the client disconnects, and this runs before the subject + // token's signature is even checked, so an aborted connection must not cut off + // work other in-flight validations of this issuer are waiting on (mu, held by + // the caller), nor let repeating the abort drive unbounded outbound requests + // to the external IdP. + detached := context.WithoutCancel(ctx) + + if issuerConfig.jwksURL == "" { + discoverCtx, cancel := context.WithTimeout(detached, httpTimeout) + jwksURL, err := v.discoverJWKSURL(discoverCtx, issuerConfig) + cancel() + if err != nil { + return fmt.Errorf("OIDC discovery failed for %s: %w", issuerConfig.IssuerURL, err) + } + issuerConfig.jwksURL = jwksURL + } + + // Validate the JWKS URL here, in the single choke point every + // registration passes through — whether it was hand-configured on + // TrustedIssuer or just discovered above. A configured JWKSURL never + // reaches discoverJWKSURL, so checking only there would leave + // hand-configured URLs unvalidated. + if err := ValidateJWKSURL(issuerConfig.jwksURL, issuerConfig.InsecureAllowHTTP, issuerConfig.AllowPrivateIPs); err != nil { + return fmt.Errorf("jwks_url for issuer %s is invalid: %w", issuerConfig.IssuerURL, err) + } + + registeredCtx, cancel := context.WithTimeout(detached, httpTimeout) + registered := issuerConfig.jwksCache.IsRegistered(registeredCtx, issuerConfig.jwksURL) + cancel() + + if registered { + // Already registered with this issuer's own cache, on a prior call + // whose own fetch never completed successfully — the only way this + // can be true, now that each issuer has its own cache. Register + // would error on an already-tracked URL, so retry via Refresh + // instead. + fetchCtx, cancel := context.WithTimeout(detached, httpTimeout) + defer cancel() + if _, err := issuerConfig.jwksCache.Refresh(fetchCtx, issuerConfig.jwksURL); err != nil { + return fmt.Errorf("failed to fetch JWKS for issuer %s: %w", issuerConfig.IssuerURL, err) + } + return nil + } + + // A newly created httprc.Resource is always scheduled to fetch + // immediately, so Register's own default WithWaitReady(true) blocks on + // that single automatic fetch. An explicit Refresh call right here would + // race it and issue a genuine second outbound request — that's why the + // registered branch above, not this one, is where Refresh is used. + fetchCtx, cancel := context.WithTimeout(detached, httpTimeout) + defer cancel() + if err := issuerConfig.jwksCache.Register(fetchCtx, issuerConfig.jwksURL, + jwk.WithHTTPClient(issuerConfig.httpClient), jwk.WithConstantInterval(jwksRefreshInterval)); err != nil { + return fmt.Errorf("failed to register JWKS for issuer %s: %w", issuerConfig.IssuerURL, err) + } + return nil +} + +// lookupJWKS returns issuerConfig's current JWKS, registering and fetching it +// first if this is the first use (see ensureRegistered). jwk.Cache serves the +// last successfully fetched Set even while a later background refresh is +// failing (httprc only stores a value after a successful fetch), so a +// transient outage at an issuer that has already been reached once no longer +// surfaces as a validation failure. +// Each call converts the cached Set into a fresh *jose.JSONWebKeySet — the +// two libraries' key representations aren't shared, so nothing here is +// visible to, or mutable by, any other concurrent caller. +func (v *MultiIssuerTokenValidator) lookupJWKS( + ctx context.Context, + issuerConfig *externalIssuerConfig, +) (*jose.JSONWebKeySet, error) { + if err := v.ensureRegistered(ctx, issuerConfig); err != nil { + return nil, err + } + + set, err := issuerConfig.jwksCache.Lookup(ctx, issuerConfig.jwksURL) + if err != nil { + return nil, fmt.Errorf("failed to lookup JWKS: %w", err) + } + + jwks, err := bridgeJWKSet(set) + if err != nil { + return nil, err + } + + if len(jwks.Keys) == 0 { + return nil, errors.New("JWKS contains no keys") + } + if len(jwks.Keys) > maxJWKSKeys { + return nil, fmt.Errorf("JWKS contains too many keys: %d (max %d)", len(jwks.Keys), maxJWKSKeys) + } + + return jwks, nil +} + +// bridgeJWKSet converts a jwx jwk.Set into the go-jose jose.JSONWebKeySet +// shape verifySignature works with, by round-tripping through the standard +// JWKS JSON both types serialize to and from. This is the simplest correct +// conversion between the two libraries' key representations — it doesn't +// require hand-mapping every key type's fields (RSA, EC, OKP, ...) between +// jwk.Key and jose.JSONWebKey. +func bridgeJWKSet(set jwk.Set) (*jose.JSONWebKeySet, error) { + raw, err := json.Marshal(set) + if err != nil { + return nil, fmt.Errorf("failed to marshal JWKS: %w", err) + } + var jwks jose.JSONWebKeySet + if err := json.Unmarshal(raw, &jwks); err != nil { + return nil, fmt.Errorf("failed to parse JWKS: %w", err) + } + return &jwks, nil +} + +// refreshOnUnknownKid forces issuerConfig's own jwk.Cache to re-fetch its +// JWKS immediately, ahead of jwx's own background refresh schedule, when a +// subject token names a kid the last cached JWKS doesn't have — the +// situation a legitimate key rotation produces. Gated by minKidRefreshInterval +// and single-flighted via issuerConfig.mu, the same mutex ensureRegistered +// uses for its own initial fetch. +// +// Errors are logged, not returned: the caller (validateExternalToken) has +// already failed signature verification once and will simply fail again if +// the refresh didn't produce a usable key, which is the correct outcome for +// a genuinely invalid token. +func (*MultiIssuerTokenValidator) refreshOnUnknownKid(ctx context.Context, issuerConfig *externalIssuerConfig) { + issuerConfig.mu.Lock() + defer issuerConfig.mu.Unlock() + + if time.Since(issuerConfig.lastKidRefresh) < minKidRefreshInterval { + return + } + issuerConfig.lastKidRefresh = time.Now() + + fetchCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*httpTimeout) + defer cancel() + if _, err := issuerConfig.jwksCache.Refresh(fetchCtx, issuerConfig.jwksURL); err != nil { + slog.Debug("JWKS refresh on unknown kid failed", "issuer", issuerConfig.IssuerURL, "error", err) + } +} + +// peekIssuer parses a JWT without signature verification to extract the "iss" claim. +func peekIssuer(rawToken string) (string, error) { + token, err := jwt.ParseSigned(rawToken, allowedSignatureAlgorithms) + if err != nil { + return "", fmt.Errorf("subject token is not a valid JWT: %w", err) + } + + var claims jwt.Claims + if err := token.UnsafeClaimsWithoutVerification(&claims); err != nil { + return "", fmt.Errorf("failed to extract claims from subject token: %w", err) + } + + if claims.Issuer == "" { + return "", fmt.Errorf("subject token is missing 'iss' claim") + } + + return claims.Issuer, nil +} + +// discoverJWKSURL performs OIDC discovery to resolve the JWKS URL for an issuer. +// It fetches the OpenID Connect discovery document at {issuerURL}/.well-known/openid-configuration +// and extracts the jwks_uri field, using issuerConfig's own dedicated HTTP client. +// +// IssuerURL itself may legitimately carry a trailing slash (see +// validateTrustedIssuerURL in pkg/authserver/config.go — Microsoft Entra ID +// v1 is one real-world example), so it is trimmed before the well-known path +// is appended here, per OIDC Discovery §4.1. The comparison against the +// discovery document's own "issuer" below intentionally uses the untrimmed +// IssuerURL: per §4.3 that comparison must be exact, since it stands in for +// the eventual match against the token's "iss". +func (*MultiIssuerTokenValidator) discoverJWKSURL(ctx context.Context, issuerConfig *externalIssuerConfig) (string, error) { + discoveryURL := strings.TrimSuffix(issuerConfig.IssuerURL, "/") + "/.well-known/openid-configuration" + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, discoveryURL, nil) + if err != nil { + return "", fmt.Errorf("failed to create discovery request: %w", err) + } + + resp, err := issuerConfig.httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("discovery request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + _, _ = io.Copy(io.Discard, resp.Body) + return "", fmt.Errorf("discovery endpoint returned status %d", resp.StatusCode) + } + + body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBodySize)) + if err != nil { + return "", fmt.Errorf("failed to read discovery response: %w", err) + } + + var doc oauthproto.OIDCDiscoveryDocument + if err := json.Unmarshal(body, &doc); err != nil { + return "", fmt.Errorf("failed to parse discovery document: %w", err) + } + + if doc.Issuer != issuerConfig.IssuerURL { + return "", fmt.Errorf("discovery document issuer %q does not match expected issuer %q", doc.Issuer, issuerConfig.IssuerURL) + } + + if doc.JWKSURI == "" { + return "", fmt.Errorf("discovery document missing 'jwks_uri'") + } + + // The returned URL is validated by the caller (ensureRegistered), which is + // the single choke point covering both discovered and pre-configured + // JWKS URLs. + return doc.JWKSURI, nil +} + +// ValidateJWKSURL checks that jwksURL parses, has a host, uses HTTPS unless +// insecureAllowHTTP permits plain HTTP — and only exactly the "http" scheme, +// not any other non-https scheme such as "file" or "ftp" — and, when the +// host is an IP literal, is not a private or loopback address unless +// allowPrivateIPs permits that. Both flags come from the specific +// TrustedIssuer being fetched (see ensureRegistered), never from a validator-wide +// or self-issuer setting. This prevents SSRF attacks where a compromised +// discovery document — or a hand-configured jwks_url — points to internal +// services. +// +// This is the single implementation shared by the runtime choke point above +// (ensureRegistered, on every fetch) and pkg/authserver/config.go's config-time +// check (validateJWKSEndpointURL): the two must not drift out of sync, or a +// laxer runtime check would silently defeat the config-time guard. +func ValidateJWKSURL(jwksURL string, insecureAllowHTTP, allowPrivateIPs bool) error { + u, err := url.Parse(jwksURL) + if err != nil { + return fmt.Errorf("invalid URL: %w", err) + } + + if u.Host == "" { + return errors.New("host is required") + } + + // Unlike issuer_url, a jwks_url carrying userinfo would actually work — + // net/http turns it into a Basic auth header on every JWKS fetch — which + // is precisely why it is rejected rather than tolerated: it would put a + // live credential in the RunConfig, in this function's error strings, and + // in any log that quotes the URL. A JWKS endpoint is public by + // definition (it serves verification keys), so there is no legitimate + // reason to authenticate to one. + if u.User != nil { + return errors.New("must not contain userinfo (credentials in the URL)") + } + + if u.Scheme != "https" && (u.Scheme != "http" || !insecureAllowHTTP) { + return fmt.Errorf("must use HTTPS, got %q", u.Scheme) + } + + host := u.Hostname() + ip := net.ParseIP(host) + if ip != nil && !allowPrivateIPs && networking.IsPrivateIP(ip) { + return errors.New("must not point to a private or loopback address") + } + + return nil +} + +// validateTrustedIssuer checks a single TrustedIssuer for structural validity +// before it is admitted into issuers: required fields, no collision with +// selfIssuer or an already-registered issuer, an ActorClaim that +// resolveAllowedActor can actually read from Extra (or ClientID), and a +// well-formed AllowedDelegateClients (non-empty, no empty entries, and the +// wildcard anyDelegateClient never combined with specific client IDs). +// +// Error messages name TrustedIssuer's wire keys (issuer_url, +// expected_audience, actor_claim), not its Go field names: TrustedIssuer is +// the serialized RunConfig.TrustedIssuers schema (see its doc comment), so an +// operator's YAML uses these keys and should see them echoed back, not Go +// identifiers they never wrote. +func validateTrustedIssuer(ti TrustedIssuer, selfIssuer string, issuers map[string]*externalIssuerConfig) error { + if ti.IssuerURL == "" { + return errors.New("issuer_url is required") + } + if ti.ExpectedAudience == "" { + return fmt.Errorf("issuer_url %q: expected_audience is required", ti.IssuerURL) + } + if ti.IssuerURL == selfIssuer { + return fmt.Errorf("issuer_url %q: must not equal the authorization server's own issuer; "+ + "self-issued tokens are already handled separately", ti.IssuerURL) + } + if _, dup := issuers[ti.IssuerURL]; dup { + return fmt.Errorf("issuer_url %q: configured more than once", ti.IssuerURL) + } + if ti.ActorClaim != "" && slices.Contains(actorClaimsNotInExtra, ti.ActorClaim) { + return fmt.Errorf( + "issuer_url %q: actor_claim %q is not supported "+ + `(use "client_id" or a non-registered claim such as "azp", "appid", "cid")`, + ti.IssuerURL, ti.ActorClaim) + } + if len(ti.AllowedDelegateClients) == 0 { + return fmt.Errorf( + "issuer_url %q: allowed_delegate_clients is required; set it to [%q] to permit any ToolHive "+ + "confidential client holding the token-exchange grant, or list specific client IDs to bind "+ + "delegation to them — an issuer with no entries can never authorize a delegated exchange", + ti.IssuerURL, anyDelegateClient) + } + if slices.Contains(ti.AllowedDelegateClients, "") { + return fmt.Errorf( + "issuer_url %q: allowed_delegate_clients must not contain an empty client ID", ti.IssuerURL) + } + if slices.Contains(ti.AllowedDelegateClients, anyDelegateClient) && len(ti.AllowedDelegateClients) > 1 { + return fmt.Errorf( + "issuer_url %q: allowed_delegate_clients must not combine the wildcard %q with specific client IDs", + ti.IssuerURL, anyDelegateClient) + } + // AllowPrivateIPs without a hand-configured jwks_url would let OIDC + // discovery — a document fetched from, and thus influenceable by, the + // external issuer itself — choose the private target the dial is + // allowed to reach. Requiring jwks_url pins that target to + // operator-supplied config. Mirrors the config-time check in + // pkg/authserver/config.go's validateTrustedIssuers; duplicated here so + // a caller that builds the validator directly (factory, tests) without + // running Config.Validate cannot bypass it. + if ti.AllowPrivateIPs && ti.JWKSURL == "" { + return fmt.Errorf( + "issuer_url %q: allow_private_ips requires jwks_url to be set explicitly; "+ + "otherwise OIDC discovery — fetched from the external issuer — would choose the private target", + ti.IssuerURL) + } + return nil +} + +// looksLikeResourceIdentifier reports whether aud is shaped like a resource/API +// identifier (an absolute URI, e.g. "https://api.example.com" or +// "api://") rather than a bare opaque string or GUID, which is more +// likely a client ID. This is a heuristic warning signal only — Entra v1 +// legitimately uses a bare app-ID GUID as an access token's "aud" (see +// TrustedIssuer.ExpectedAudience's doc comment), so a bare value is not +// rejected, only flagged for the operator to double-check. +func looksLikeResourceIdentifier(aud string) bool { + return strings.Contains(aud, "://") +} + +// resolveAllowedActor resolves the issuer's configured actor claim from +// claims and checks it against issuerConfig.AllowedActors, returning the +// matched value on success. Called only when the subject token carries no +// may_act claim — see validateExternalToken. +func resolveAllowedActor(issuerConfig *externalIssuerConfig, claims *ValidatedClaims) (string, error) { + claimName := issuerConfig.ActorClaim + if claimName == "" { + claimName = defaultActorClaim + } + + // "client_id" is routed to ValidatedClaims.ClientID by assignClaim, so it + // never appears in Extra; without this fallback a plausible operator + // config (ActorClaim: "client_id") would silently reject all traffic. + var raw any + if claimName == "client_id" { + raw = claims.ClientID + } else { + raw = claims.Extra[claimName] + } + + actor, ok := raw.(string) + if !ok || actor == "" { + return "", fmt.Errorf( + "subject token from issuer %q is missing or has an invalid %q claim required for delegation consent", + issuerConfig.IssuerURL, claimName) + } + + if !slices.Contains(issuerConfig.AllowedActors, actor) { + return "", fmt.Errorf( + "subject token from issuer %q names actor %q in claim %q, which is not in the allowed actors list", + issuerConfig.IssuerURL, actor, claimName) + } + + return actor, nil +} diff --git a/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go b/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go new file mode 100644 index 0000000000..62b7c7726e --- /dev/null +++ b/pkg/authserver/server/tokenexchange/multi_issuer_validator_test.go @@ -0,0 +1,2344 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/go-jose/go-jose/v4/jwt" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/networking" +) + +const ( + testExternalIssuer = "https://keycloak.example.com/realms/test" + testExternalAudience = "toolhive-authserver" +) + +// startJWKSServer creates a test HTTP server that serves a JWKS endpoint. +// The returned server must be closed by the caller. +func startJWKSServer(t *testing.T, tj *testJWKS) *httptest.Server { + t.Helper() + + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + // Serve only the public keys. + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(tj.publicJWKS()) + }) + + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +// newMultiValidator creates a MultiIssuerTokenValidator configured for testing. +// The external JWKS URL is pre-resolved (no discovery needed) unless jwksURL is empty. +func newMultiValidator( + t *testing.T, + selfJWKS *testJWKS, + trustedIssuers []TrustedIssuer, +) *MultiIssuerTokenValidator { + t.Helper() + + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + // Copy rather than mutate the caller's slice: give every test issuer + // permissive flags so its httptest discovery/JWKS server, which runs on + // loopback over plain HTTP, is reachable. + issuers := make([]TrustedIssuer, len(trustedIssuers)) + for i, ti := range trustedIssuers { + ti.InsecureAllowHTTP = true + ti.AllowPrivateIPs = true + issuers[i] = ti + } + + v, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, issuers) + require.NoError(t, err) + return v +} + +// externalClaims returns standard JWT claims for a token issued by the external issuer. +func externalClaims() jwt.Claims { + now := time.Now() + return jwt.Claims{ + Subject: "ext-user-456", + Issuer: testExternalIssuer, + Audience: jwt.Audience{testExternalAudience}, + Expiry: jwt.NewNumericDate(now.Add(time.Hour)), + IssuedAt: jwt.NewNumericDate(now), + NotBefore: jwt.NewNumericDate(now.Add(-time.Minute)), + ID: "jti-ext-789", + } +} + +func TestMultiIssuerTokenValidator_Validate(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + tests := []struct { + name string + trustedIssuers []TrustedIssuer + token func(t *testing.T) string + wantErr bool + errContains string + check func(t *testing.T, vc *ValidatedClaims) + }{ + { + name: "self-issued token routes to self validator", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return selfJWKS.signToken(t, validClaims(), validExtraClaims()) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "user-123", vc.Subject) + assert.Equal(t, testIssuer, vc.Issuer) + assert.Equal(t, []string{testIssuer}, vc.Audience) + assert.Equal(t, "Test User", vc.Name) + assert.Equal(t, "test@example.com", vc.Email) + assert.Empty(t, vc.ExternalIssuer, "self-issued path must never set ExternalIssuer") + }, + }, + { + name: "external token accepted", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]interface{}{ + "name": "External User", + "email": "ext@keycloak.example.com", + "azp": "ext-agent", + }) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "ext-user-456", vc.Subject) + assert.Equal(t, testExternalIssuer, vc.Issuer) + assert.Equal(t, []string{testExternalAudience}, vc.Audience) + assert.Equal(t, "jti-ext-789", vc.JWTID) + assert.Equal(t, "External User", vc.Name) + assert.Equal(t, "ext@keycloak.example.com", vc.Email) + assert.False(t, vc.Expiry.IsZero()) + assert.False(t, vc.IssuedAt.IsZero()) + assert.Equal(t, "ext-agent", vc.ExternalActor, "actor claim matched via default azp") + assert.Equal(t, testExternalIssuer, vc.ExternalIssuer) + }, + }, + { + name: "external token surfaces configured AllowedDelegateClients", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{"toolhive-agent-a", "toolhive-agent-b"}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]interface{}{"azp": "ext-agent"}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "ext-agent", vc.ExternalActor) + assert.Equal(t, []string{"toolhive-agent-a", "toolhive-agent-b"}, vc.AllowedDelegateClients, + "the validator surfaces the issuer's configured allowlist for the handler to enforce") + }, + }, + { + // #5989 hardening: AllowedDelegateClients is required (validateTrustedIssuer + // rejects an empty one), so the permissive case is now the wildcard, + // declared explicitly, rather than an absent field. + name: "external token surfaces the wildcard AllowedDelegateClients when the issuer opts into it", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]interface{}{"azp": "ext-agent"}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, []string{anyDelegateClient}, vc.AllowedDelegateClients) + }, + }, + { + // #5989 fix: an issuer emitting may_act with no way to populate + // AllowedActors (e.g. Entra, Okta) previously had no per-client + // containment at all. AllowedDelegateClients must now surface on + // the may_act path too, for checkDelegationConsent to enforce. + name: "external token with may_act still surfaces configured AllowedDelegateClients", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{"toolhive-agent-a"}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), + map[string]any{"may_act": map[string]any{"sub": "some-toolhive-client", "iss": testIssuer}}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Equal(t, []string{"toolhive-agent-a"}, vc.AllowedDelegateClients, + "may_act bypasses the AllowedActors allowlist, but not per-client containment") + }, + }, + { + name: "external token custom actor claim appid accepted", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + ActorClaim: "appid", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"appid": "ext-agent"}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "ext-agent", vc.ExternalActor) + }, + }, + { + name: "external token custom actor claim cid accepted", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + ActorClaim: "cid", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"cid": "ext-agent"}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "ext-agent", vc.ExternalActor) + }, + }, + { + name: "external token actor claim client_id resolves from ClientID field", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + ActorClaim: "client_id", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + // "client_id" is routed to ValidatedClaims.ClientID by assignClaim, so + // it never lands in Extra — resolveAllowedActor must fall back to + // reading the structured field instead. + return externalJWKS.signToken(t, externalClaims(), map[string]any{"client_id": "ext-agent"}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "ext-agent", vc.ClientID) + assert.Equal(t, "ext-agent", vc.ExternalActor) + }, + }, + { + name: "external token actor not in allowed actors rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": "some-other-client"}) + }, + wantErr: true, + errContains: "not in the allowed actors list", + }, + { + name: "external token missing actor claim rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), nil) + }, + wantErr: true, + errContains: "required for delegation consent", + }, + { + name: "external token actor claim is a number rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": 123}) + }, + wantErr: true, + errContains: "required for delegation consent", + }, + { + name: "external token actor claim is a bool rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": true}) + }, + wantErr: true, + errContains: "required for delegation consent", + }, + { + name: "external token actor claim is an array rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": []string{"ext-agent"}}) + }, + wantErr: true, + errContains: "required for delegation consent", + }, + { + name: "external token actor claim is a nested object rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), + map[string]any{"azp": map[string]any{"id": "ext-agent"}}) + }, + wantErr: true, + errContains: "required for delegation consent", + }, + { + name: "external token actor claim empty string rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": ""}) + }, + wantErr: true, + errContains: "required for delegation consent", + }, + { + name: "external token no may_act and empty allowed actors rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + // AllowedActors intentionally empty. + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": "ext-agent"}) + }, + wantErr: true, + errContains: "not in the allowed actors list", + }, + { + name: "external token with may_act and empty allowed actors accepted", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + // AllowedActors intentionally empty: may_act is authoritative and + // skips the allowlist entirely. + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), + map[string]any{"may_act": map[string]any{"sub": "some-toolhive-client", "iss": testIssuer}}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Equal(t, "some-toolhive-client", vc.MayAct.Sub) + assert.Empty(t, vc.ExternalActor, "allowlist is skipped whenever may_act is present") + assert.Equal(t, testExternalIssuer, vc.ExternalIssuer, + "provenance must still be recorded on the may_act path, which bypasses the allowlist") + }, + }, + { + // #5989 fix: the external path cannot leave may_act.iss optional + // the way the self-issued path does (see + // TestSelfIssuedTokenValidator_Validate's "token with may_act + // claim extracts MayAct" for the self-issued equivalent, which + // omits iss and is accepted). Without this, an external issuer + // could authorize ANY ToolHive client via a bare may_act.sub, + // bypassing AllowedActors/AllowedDelegateClients entirely. + name: "external token may_act missing iss rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), + map[string]any{"may_act": map[string]any{"sub": "some-toolhive-client"}}) + }, + wantErr: true, + errContains: "missing required 'iss'", + }, + { + name: "external token may_act wins even when azp is not allowlisted", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"someone-else"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "not-in-the-allowlist", + "may_act": map[string]any{"sub": "some-toolhive-client", "iss": testIssuer}, + }) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Empty(t, vc.ExternalActor) + assert.Equal(t, testExternalIssuer, vc.ExternalIssuer) + }, + }, + { + name: "external token malformed may_act — JSON string, not object — rejected outright", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + // azp is allowlisted, so a regression that fell through to the + // allowlist instead of rejecting would wrongly accept this. + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "may_act": "agent-a", + }) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "external token malformed may_act — non-string sub — rejected outright", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "may_act": map[string]any{"sub": 123}, + }) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "external token malformed may_act — empty sub — rejected outright", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "may_act": map[string]any{"sub": ""}, + }) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "external token malformed may_act — array sub — rejected outright", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "may_act": map[string]any{"sub": []string{"agent-a"}}, + }) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "external token malformed may_act — object with no sub key — rejected outright", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "may_act": map[string]any{"subject": "agent-a"}, + }) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + // Proves validateMayActShape is passed v.selfIssuer, not + // issuerConfig.IssuerURL: may_act.sub is always compared against + // a ToolHive client ID by checkDelegationConsent, regardless of + // which external issuer validated the surrounding token, so its + // optional iss must be checked against THIS server's own + // issuer. + name: "external token may_act.iss matching this server's own issuer accepted", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), + map[string]any{"may_act": map[string]any{"sub": "some-toolhive-client", "iss": testIssuer}}) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Equal(t, testIssuer, vc.MayAct.Iss) + }, + }, + { + // The external issuer's own URL is never a valid may_act.iss + // value: a regression that compared against issuerConfig.IssuerURL + // instead of v.selfIssuer would wrongly accept this. + name: "external token may_act.iss matching the external issuer's own URL rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), + map[string]any{"may_act": map[string]any{"sub": "some-toolhive-client", "iss": testExternalIssuer}}) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "external token carrying c_hash (an ID token) rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "c_hash": "some-hash-value", + }) + }, + wantErr: true, + errContains: "'c_hash'", + }, + { + name: "external token wrong audience", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + claims := externalClaims() + claims.Audience = jwt.Audience{"wrong-audience"} + return externalJWKS.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + name: "unknown issuer rejected", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + claims := externalClaims() + claims.Issuer = "https://evil.example.com" + return externalJWKS.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "untrusted issuer", + }, + { + name: "external token bad signature", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + // Sign with a different key than the one the JWKS server serves. + wrongJWKS := newTestJWKS(t) + return wrongJWKS.signToken(t, externalClaims(), nil) + }, + wantErr: true, + errContains: "signature verification failed", + }, + { + name: "self-issued token signed by external key fails", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + // Token claims say iss=self, but signed by the external key. + // Routes to self validator, which rejects the signature. + return externalJWKS.signToken(t, validClaims(), nil) + }, + wantErr: true, + errContains: "signature verification failed", + }, + { + name: "external token missing subject", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + claims := externalClaims() + claims.Subject = "" + return externalJWKS.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "missing required 'sub' claim", + }, + { + name: "external token missing exp claim", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + claims := externalClaims() + claims.Expiry = nil + return externalJWKS.signToken(t, claims, map[string]any{"azp": "ext-agent"}) + }, + wantErr: true, + errContains: "missing required 'exp' claim", + }, + { + // Distinct from "malformed token" below: that hits the JWT parse + // failure inside peekIssuer, this hits its separate empty-issuer + // check on an otherwise well-formed token. Both wrap to the same + // outer "failed to determine token issuer" message, so assert on + // the distinct inner text instead. + name: "token missing iss claim", + token: func(t *testing.T) string { + t.Helper() + claims := externalClaims() + claims.Issuer = "" + return externalJWKS.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "missing 'iss' claim", + }, + { + // Proves the allowlist is skipped (not just "would have failed + // anyway") whenever may_act is present: both azp and + // AllowedActors would satisfy resolveAllowedActor if it ran, so + // a regression that calls it unconditionally and only gates the + // error on MayAct==nil would still populate ExternalActor here. + name: "external token may_act present alongside an allowlisted azp still yields empty ExternalActor", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + return externalJWKS.signToken(t, externalClaims(), map[string]any{ + "azp": "ext-agent", + "may_act": map[string]any{"sub": "some-toolhive-client", "iss": testIssuer}, + }) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Empty(t, vc.ExternalActor, "allowlist must be skipped entirely when may_act is present") + }, + }, + { + name: "external token expired", + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + token: func(t *testing.T) string { + t.Helper() + claims := externalClaims() + claims.Expiry = jwt.NewNumericDate(time.Now().Add(-time.Hour)) + claims.IssuedAt = jwt.NewNumericDate(time.Now().Add(-2 * time.Hour)) + return externalJWKS.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + name: "malformed token", + trustedIssuers: nil, + token: func(_ *testing.T) string { + return "not-a-jwt" + }, + wantErr: true, + errContains: "failed to determine token issuer", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + validator := newMultiValidator(t, selfJWKS, tt.trustedIssuers) + rawToken := tt.token(t) + + result, err := validator.Validate(context.Background(), rawToken) + + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errContains) + assert.Nil(t, result) + return + } + + require.NoError(t, err) + require.NotNil(t, result) + if tt.check != nil { + tt.check(t, result) + } + }) + } +} + +// TestMultiIssuerTokenValidator_TrailingSlashIssuerExactMatch covers a +// trailing-slash issuer_url (e.g. Microsoft Entra ID v1's +// "https://sts.windows.net/{tenant}/") being accepted as an exact "iss" +// match. JWKSURL is set explicitly rather than left for discovery to resolve: +// validateTrustedIssuer rejects {AllowPrivateIPs: true, JWKSURL: ""} (the +// combination newMultiValidator forces on every test issuer to reach its +// loopback server), so an empty JWKSURL would fail construction here. +// discoverJWKSURL's own trailing-slash trim is covered directly by +// TestMultiIssuerTokenValidator_DiscoverJWKSURL instead; this test is only +// about the trailing-slash IssuerURL matching the token's "iss" claim exactly +// (per OIDC Core §3.1.3.3), which happens regardless of how JWKSURL was +// resolved. +func TestMultiIssuerTokenValidator_TrailingSlashIssuerExactMatch(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + trailingSlashIssuer := "https://tenant.example.com/realm/" + trustedIssuers := []TrustedIssuer{{ + IssuerURL: trailingSlashIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + claims := jwt.Claims{ + Subject: "trailing-slash-user", + Issuer: trailingSlashIssuer, + Audience: jwt.Audience{testExternalAudience}, + Expiry: jwt.NewNumericDate(time.Now().Add(time.Hour)), + IssuedAt: jwt.NewNumericDate(time.Now()), + NotBefore: jwt.NewNumericDate(time.Now().Add(-time.Minute)), + ID: "jti-trailing-slash-001", + } + rawToken := externalJWKS.signToken(t, claims, map[string]any{"azp": "ext-agent"}) + + result, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "trailing-slash-user", result.Subject) + assert.Equal(t, trailingSlashIssuer, result.Issuer) +} + +func TestMultiIssuerTokenValidator_JWKSCaching(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + + // Track how many times the JWKS endpoint is hit. + var fetchCount atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + fetchCount.Add(1) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(externalJWKS.publicJWKS()) + }) + jwksServer := httptest.NewServer(mux) + t.Cleanup(jwksServer.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + // Validate two tokens — the JWKS should be fetched only once (cached). + for i := range 2 { + claims := externalClaims() + claims.ID = fmt.Sprintf("jti-cache-%d", i) + rawToken := externalJWKS.signToken(t, claims, map[string]any{"azp": "ext-agent"}) + + result, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "ext-user-456", result.Subject) + } + + assert.Equal(t, int32(1), fetchCount.Load(), "JWKS should be fetched only once due to caching") +} + +// TestMultiIssuerTokenValidator_JWKSRefreshIntervalIsPinned confirms +// registerOrRefresh actually threads jwk.WithConstantInterval(jwksRefreshInterval) +// through to the underlying httprc.Resource, even though the JWKS endpoint +// advertises a much longer Cache-Control max-age. Without the constant +// interval, httprc would derive the refresh schedule from that header +// instead — see jwksRefreshInterval's doc comment for why an external +// issuer must not get to choose how long we keep its keys cached. +// +// This can't be observed by waiting for a second fetch without a wall-clock +// sleep (forbidden by this repo's testing rules), so it asserts directly on +// the registered resource's ConstantInterval() instead of on fetch timing. +func TestMultiIssuerTokenValidator_JWKSRefreshIntervalIsPinned(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + // A long max-age that would push httprc's derived interval far past + // jwksRefreshInterval if the constant interval were not applied. + w.Header().Set("Cache-Control", "max-age=2592000") // 30 days + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(externalJWKS.publicJWKS()) + }) + jwksServer := httptest.NewServer(mux) + t.Cleanup(jwksServer.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + rawToken := externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": "ext-agent"}) + _, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + + issuerConfig := validator.issuers[testExternalIssuer] + resource, err := issuerConfig.jwksCache.LookupResource(context.Background(), issuerConfig.jwksURL) + require.NoError(t, err) + assert.Equal(t, jwksRefreshInterval, resource.ConstantInterval(), + "registered resource must ignore the endpoint's own Cache-Control max-age") +} + +func TestNewMultiIssuerTokenValidator_Validation(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + tests := []struct { + name string + selfValidator *SelfIssuedTokenValidator + selfIssuer string + trustedIssuers []TrustedIssuer + errContains string + }{ + { + name: "nil selfValidator rejected", + selfIssuer: testIssuer, + errContains: "selfValidator must not be nil", + }, + { + name: "empty selfIssuer rejected", + selfValidator: selfValidator, + selfIssuer: "", + errContains: "selfIssuer must not be empty", + }, + { + name: "empty IssuerURL rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: "", + ExpectedAudience: testExternalAudience, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "issuer_url is required", + }, + { + name: "empty ExpectedAudience rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: "", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "expected_audience is required", + }, + { + name: "IssuerURL equal to selfIssuer rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testIssuer, + ExpectedAudience: testExternalAudience, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "must not equal the authorization server's own issuer", + }, + { + name: "duplicate IssuerURL rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{ + {IssuerURL: testExternalIssuer, ExpectedAudience: testExternalAudience, AllowedDelegateClients: []string{anyDelegateClient}}, + {IssuerURL: testExternalIssuer, ExpectedAudience: testExternalAudience, AllowedDelegateClients: []string{anyDelegateClient}}, + }, + errContains: "configured more than once", + }, + { + name: `ActorClaim "sub" rejected — assignClaim never leaves it in Extra`, + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + ActorClaim: "sub", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "actor_claim", + }, + { + name: `ActorClaim "scope" rejected — rerouted to a structured field`, + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + ActorClaim: "scope", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "actor_claim", + }, + { + name: `ActorClaim "may_act" rejected — rerouted to a structured field`, + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + ActorClaim: "may_act", + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "actor_claim", + }, + { + name: "AllowedDelegateClients with an empty entry rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + AllowedDelegateClients: []string{"toolhive-agent-a", ""}, + }}, + errContains: "allowed_delegate_clients must not contain an empty client ID", + }, + { + // #5989 hardening: permissiveness must be declared with the + // wildcard, not obtained by omission — pins the fail-closed + // default this PR introduces. + name: "AllowedDelegateClients absent rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + // AllowedDelegateClients intentionally unset. + }}, + errContains: "allowed_delegate_clients is required", + }, + { + name: "AllowedDelegateClients empty (non-nil) slice rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + AllowedDelegateClients: []string{}, + }}, + errContains: "allowed_delegate_clients is required", + }, + { + // The wildcard authorizes every ToolHive client explicitly; mixing + // it with specific IDs would either be redundant or, worse, give + // the false impression that only the listed IDs are bound. + // Rejecting outright avoids silently ignoring the specific IDs. + name: "AllowedDelegateClients wildcard mixed with a specific client ID rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + AllowedDelegateClients: []string{anyDelegateClient, "toolhive-agent-a"}, + }}, + errContains: "must not combine the wildcard", + }, + { + // AllowPrivateIPs without a hand-configured jwks_url would let + // OIDC discovery choose the private target the dial is allowed to + // reach; validateTrustedIssuer must reject it so a caller that + // skips Config.Validate (factory, tests) cannot bypass the + // config-time check in pkg/authserver/config.go. + name: "AllowPrivateIPs without JWKSURL rejected", + selfValidator: selfValidator, + selfIssuer: testIssuer, + trustedIssuers: []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + AllowPrivateIPs: true, + AllowedDelegateClients: []string{anyDelegateClient}, + }}, + errContains: "allow_private_ips requires jwks_url", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + v, err := NewMultiIssuerTokenValidator(tt.selfValidator, tt.selfIssuer, tt.trustedIssuers) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errContains) + assert.Nil(t, v) + }) + } +} + +func TestNewMultiIssuerTokenValidator_EmptyAllowedActorsAccepted(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + // Empty AllowedActors must be accepted by the constructor: a may_act-only + // issuer (no allowlisted actors at all) is a legitimate configuration. + v, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + AllowedDelegateClients: []string{anyDelegateClient}, + }}) + require.NoError(t, err) + assert.NotNil(t, v) +} + +// syncBuffer is a concurrency-safe io.Writer over a bytes.Buffer, used by +// TestNewMultiIssuerTokenValidator_AudienceShapeWarning to capture +// slog.Default() output. slog.SetDefault is process-global, so a plain +// bytes.Buffer would race against any other goroutine that logs while a +// capturing handler is installed. Mirrors the identically-named helper in +// pkg/authserver/runner/embeddedauthserver_test.go; kept as a separate copy +// since it is an unexported test type not worth exporting across packages. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +// TestNewMultiIssuerTokenValidator_AudienceShapeWarning pins +// looksLikeResourceIdentifier's startup warning: a TrustedIssuer whose +// ExpectedAudience has no URI scheme (e.g. a bare client-ID-shaped value) +// logs a slog.Warn naming the issuer, since that audience shape cannot +// distinguish an ID token from an access token for that issuer (the token +// itself may carry neither "at_hash" nor "c_hash", which OIDC permits — see +// rejectIDTokenClaims's doc comment). A URI-shaped audience (has a "://" +// scheme) must NOT trigger the warning. +// +// NOT t.Parallel(): swaps the package-global slog.Default() and restores it +// via t.Cleanup; see syncBuffer's doc comment for why a concurrent writer +// alone would be safe but an overlapping swap/restore pair would not be. +// +//nolint:paralleltest // mutates the package-global slog.Default() +func TestNewMultiIssuerTokenValidator_AudienceShapeWarning(t *testing.T) { + tests := []struct { + name string + expectedAudience string + wantWarning bool + }{ + { + name: "bare client-ID-shaped audience warns", + expectedAudience: "toolhive-authserver", + wantWarning: true, + }, + { + name: "bare GUID audience warns", + expectedAudience: "3f2504e0-4f89-11d3-9a0c-0305e82c3301", + wantWarning: true, + }, + { + name: "https URI audience does not warn", + expectedAudience: "https://api.example.com", + wantWarning: false, + }, + { + name: "api URI audience does not warn", + expectedAudience: "api://3f2504e0-4f89-11d3-9a0c-0305e82c3301", + wantWarning: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + selfJWKS := newTestJWKS(t) + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + + var buf syncBuffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + v, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: tt.expectedAudience, + AllowedActors: []string{"ext-agent"}, // avoid the unrelated AllowedActors warning + AllowedDelegateClients: []string{anyDelegateClient}, + }}) + require.NoError(t, err) + require.NotNil(t, v) + + logged := buf.String() + const wantSubstring = "does not look like a resource identifier" + if tt.wantWarning { + assert.Contains(t, logged, wantSubstring) + assert.Contains(t, logged, testExternalIssuer) + } else { + assert.NotContains(t, logged, wantSubstring) + } + }) + } +} + +func TestNewMultiIssuerTokenValidator_ClonesAllowedActors(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + allowedActors := []string{"ext-agent"} + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: allowedActors, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + // Mutate the caller's slice after construction. If the constructor didn't + // clone it, this would change what the validator accepts. + allowedActors[0] = "someone-else" + + rawToken := externalJWKS.signToken(t, externalClaims(), map[string]any{"azp": "ext-agent"}) + result, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err, "post-construction mutation of the caller's slice must not affect validation") + require.NotNil(t, result) + assert.Equal(t, "ext-agent", result.ExternalActor) +} + +// TestMultiIssuerTokenValidator_ClockSkewLeeway exercises externalClockSkewLeeway +// (60s): nbf/iat tolerate it, but exp is enforced strictly by an independent +// check in validateExternalToken, since go-jose's leeway would otherwise widen +// exp too and let a genuinely expired subject token through. +func TestMultiIssuerTokenValidator_ClockSkewLeeway(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + tests := []struct { + name string + claims func() jwt.Claims + wantErr bool + errContains string + }{ + { + name: "nbf ~30s in the future is within leeway", + claims: func() jwt.Claims { + c := externalClaims() + c.NotBefore = jwt.NewNumericDate(time.Now().Add(30 * time.Second)) + return c + }, + }, + { + name: "nbf ~5m in the future exceeds leeway", + claims: func() jwt.Claims { + c := externalClaims() + c.NotBefore = jwt.NewNumericDate(time.Now().Add(5 * time.Minute)) + return c + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + // go-jose's leeway also widens exp, so a naive implementation would + // accept this. validateExternalToken's independent strict check + // must reject it instead. + name: "expired ~30s ago is rejected by the strict expiry check, not the leeway", + claims: func() jwt.Claims { + c := externalClaims() + c.Expiry = jwt.NewNumericDate(time.Now().Add(-30 * time.Second)) + return c + }, + wantErr: true, + errContains: "subject token has expired", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rawToken := externalJWKS.signToken(t, tt.claims(), map[string]any{"azp": "ext-agent"}) + result, err := validator.Validate(context.Background(), rawToken) + + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errContains) + assert.Nil(t, result) + return + } + require.NoError(t, err) + require.NotNil(t, result) + }) + } +} + +// TestMultiIssuerTokenValidator_DiscoverJWKSURL exercises discoverJWKSURL +// directly via the externalIssuerConfig seam, bypassing constructor +// validation so an empty JWKSURL can be tested at all. +// NewMultiIssuerTokenValidator can no longer exercise OIDC discovery through +// its own constructor: validateTrustedIssuer rejects {AllowPrivateIPs: true, +// JWKSURL: ""}, and every test issuer built through it runs on a loopback +// httptest server, which requires AllowPrivateIPs: true to be reachable at +// all. discoverJWKSURL itself only reads issuerConfig.IssuerURL and +// .httpClient, so it needs neither a real validator nor a registered issuer. +func TestMultiIssuerTokenValidator_DiscoverJWKSURL(t *testing.T) { + t.Parallel() + + newDiscoveryServer := func(t *testing.T, handler http.HandlerFunc) *httptest.Server { + t.Helper() + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/openid-configuration", handler) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv + } + + tests := []struct { + name string + // build returns the issuerConfig to call discoverJWKSURL with, and + // (only for the no-error cases) the jwks_uri it must return. + build func(t *testing.T) (issuerConfig *externalIssuerConfig, wantJWKSURI string) + errContains string + }{ + { + name: "happy path returns the advertised jwks_uri", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, r *http.Request) { + base := "http://" + r.Host + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "issuer": base, + "jwks_uri": base + "/jwks", + }) + }) + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, srv.URL + "/jwks" + }, + }, + { + // Proves two things at once: the well-known path is built by + // trimming the trailing slash (a client that refuses to follow + // redirects would see a 301 from http.ServeMux's path-cleaning if + // the trim were dropped, since the request path would then + // contain "//"), and the issuer-match comparison below that use + // the *untrimmed* IssuerURL (the discovery document's own + // "issuer" here still carries the trailing slash). + name: "trailing-slash issuer_url is trimmed before the well-known path is appended", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, r *http.Request) { + base := "http://" + r.Host + "/" + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "issuer": base, + "jwks_uri": base + "jwks", + }) + }) + client := srv.Client() + client.CheckRedirect = func(*http.Request, []*http.Request) error { + return fmt.Errorf("unexpected redirect: issuer_url trailing slash was not trimmed") + } + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL + "/", AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: client, + }, srv.URL + "/jwks" + }, + }, + { + name: "invalid issuer_url fails to build the discovery request", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + return &externalIssuerConfig{ + // A raw control character is rejected by url.Parse inside + // http.NewRequestWithContext, before any network I/O. + TrustedIssuer: TrustedIssuer{IssuerURL: "http://example.com/\x00", AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: &http.Client{}, + }, "" + }, + errContains: "failed to create discovery request", + }, + { + name: "unreachable issuer fails the discovery request", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := httptest.NewServer(http.NotFoundHandler()) + srv.Close() // closed before use: the port is now guaranteed unreachable. + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, "" + }, + errContains: "discovery request failed", + }, + { + name: "non-200 response", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + }) + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, "" + }, + errContains: "discovery endpoint returned status 500", + }, + { + name: "malformed JSON body fails to parse", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte("{not-json")) + }) + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, "" + }, + // The distinguishing substring from json.Unmarshal itself, not + // just the wrapping "failed to parse discovery document": the + // oversized-body case below wraps the same outer message but + // fails for a different underlying reason. + errContains: "invalid character", + }, + { + // A body over maxResponseBodySize (1 MiB) is not its own error + // branch: io.LimitReader silently truncates it, and the + // truncated bytes then fail to parse as JSON — so this exercises + // the same "failed to parse discovery document" wrapper as the + // malformed-JSON case above, but for a distinct underlying + // reason (an unterminated value, not a syntax error), which is + // what's asserted on below. + name: "oversized discovery document is truncated and fails JSON parsing", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, r *http.Request) { + issuer := "http://" + r.Host + w.Header().Set("Content-Type", "application/json") + // Deliberately WELL-FORMED and complete: issuer matches and + // jwks_uri is present, so if the whole body were read this + // document would parse and discovery would succeed. Only the + // maxResponseBodySize cap cutting inside the padding makes it + // fail. That is what pins the cap — an unterminated body + // would fail the same way with the cap deleted. + _, _ = w.Write([]byte(`{"issuer":"` + issuer + `","jwks_uri":"` + issuer + `/jwks","padding":"`)) + _, _ = w.Write([]byte(strings.Repeat("a", 2*maxResponseBodySize))) + _, _ = w.Write([]byte(`"}`)) + }) + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, "" + }, + errContains: "unexpected end of JSON input", + }, + { + name: "issuer mismatch is rejected", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "issuer": "https://different-issuer.example.com", + "jwks_uri": "http://" + r.Host + "/jwks", + }) + }) + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, "" + }, + errContains: "does not match expected issuer", + }, + { + name: "missing jwks_uri is rejected", + build: func(t *testing.T) (*externalIssuerConfig, string) { + t.Helper() + srv := newDiscoveryServer(t, func(w http.ResponseWriter, r *http.Request) { + // Issuer must match so discovery reaches the jwks_uri check. + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "issuer": "http://" + r.Host, + }) + }) + return &externalIssuerConfig{ + TrustedIssuer: TrustedIssuer{IssuerURL: srv.URL, AllowedDelegateClients: []string{anyDelegateClient}}, + httpClient: srv.Client(), + }, "" + }, + errContains: "missing 'jwks_uri'", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + issuerConfig, wantJWKSURI := tt.build(t) + gotJWKSURI, err := (&MultiIssuerTokenValidator{}).discoverJWKSURL(context.Background(), issuerConfig) + + if tt.errContains != "" { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errContains) + assert.Empty(t, gotJWKSURI) + return + } + require.NoError(t, err) + assert.Equal(t, wantJWKSURI, gotJWKSURI) + }) + } +} + +// TestMultiIssuerTokenValidator_DiscoveryRefusesPrivateAddress proves the +// dial-time private-IP guard on the strict (AllowPrivateIPs: false) path: an +// issuer with an empty JWKSURL forces OIDC discovery, and every test server +// in this file runs on loopback — so with AllowPrivateIPs left false, +// discovery must be refused before any response is even read, regardless of +// what the discovery endpoint would have returned. The content-level failure +// modes of discoverJWKSURL itself (non-200, malformed doc, issuer mismatch, +// missing jwks_uri) are exercised directly by +// TestMultiIssuerTokenValidator_DiscoverJWKSURL instead, since none of them +// can be reached through this constructor path anymore. +func TestMultiIssuerTokenValidator_DiscoveryRefusesPrivateAddress(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + + // The handler would, if reached, serve a valid discovery document — this + // confirms the failure happens at the dial, not because the endpoint + // itself is broken. + mux := http.NewServeMux() + mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) { + base := "http://" + r.Host + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "issuer": base, + "jwks_uri": base + "/jwks", + }) + }) + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + + // JWKSURL is left empty to force discovery. AllowPrivateIPs is + // deliberately false: newMultiValidator would force it true to reach the + // loopback server, defeating the point of this test, so the constructor + // is called directly here instead. + trustedIssuers := []TrustedIssuer{{ + IssuerURL: server.URL, + ExpectedAudience: testExternalAudience, + InsecureAllowHTTP: true, + AllowPrivateIPs: false, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + validator, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, trustedIssuers) + require.NoError(t, err) + + claims := externalClaims() + claims.Issuer = server.URL + rawToken := externalJWKS.signToken(t, claims, nil) + + result, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), "OIDC discovery failed") + assert.Contains(t, err.Error(), networking.ErrPrivateIpAddress, + "the failure must be the dial-time private-IP guard specifically, not just any discovery error") + assert.Nil(t, result) +} + +func TestMultiIssuerTokenValidator_KidMismatch(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, externalJWKS) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: jwksServer.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + // Sign with a key whose public half is NOT in the served JWKS and whose + // kid does not match any served key. The kid lookup misses, so the + // validator falls back to trying every served key and fails verification. + unknownKey := newECDSAJWK(t, "unknown-kid") + claims := externalClaims() + rawToken := signWithJWK(t, unknownKey, jose.ES256, claims) + + result, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), "signature verification failed") + assert.Nil(t, result) +} + +// TestValidateJWKSURL exercises ValidateJWKSURL directly: this is the SSRF +// guard applied in ensureRegistered to every JWKS URL for a given issuer, whether +// hand-configured on TrustedIssuer or resolved via discovery, and shared +// verbatim with pkg/authserver/config.go's config-time check +// (validateJWKSEndpointURL) so the two can't drift out of sync. The +// equivalent check on redirect hops (networking.SameHostRedirectPolicy) and +// the dial-time IP guard (networking.NewHostScopedClientBuilder) are +// exercised via the networking package's own tests, not here. These cases +// all pass insecureAllowHTTP=false, allowPrivateIPs=false — the strict +// defaults — since every other test in this file goes through +// newMultiValidator, which sets both permissive flags on its test issuers to +// reach their httptest servers over plain HTTP on loopback; the +// insecureAllowHTTP=true cases below are the exception, covering the laxity +// that must not extend beyond "http" to every other non-https scheme. +func TestValidateJWKSURL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + url string + insecureAllowHTTP bool + wantErr string + }{ + {name: "https accepted", url: "https://issuer.example.com/jwks"}, + {name: "http rejected", url: "http://issuer.example.com/jwks", wantErr: "must use HTTPS"}, + { + name: "userinfo with password rejected", + url: "https://user:hunter2@issuer.example.com/jwks", + wantErr: "must not contain userinfo", + }, + { + // url.Parse sets User for a bare username too, and net/http would + // still send it as a Basic auth header. + name: "userinfo without password rejected", + url: "https://user@issuer.example.com/jwks", + wantErr: "must not contain userinfo", + }, + { + name: "userinfo rejected even with insecureAllowHTTP", + url: "http://user:hunter2@issuer.example.com/jwks", + insecureAllowHTTP: true, + wantErr: "must not contain userinfo", + }, + { + name: "http accepted with insecureAllowHTTP", + url: "http://issuer.example.com/jwks", + insecureAllowHTTP: true, + }, + { + name: "ftp rejected even with insecureAllowHTTP", + url: "ftp://issuer.example.com/jwks", + insecureAllowHTTP: true, + wantErr: "must use HTTPS", + }, + { + name: "no scheme rejected even with insecureAllowHTTP", + url: "//issuer.example.com/jwks", + insecureAllowHTTP: true, + wantErr: "must use HTTPS", + }, + {name: "loopback IP literal rejected", url: "https://127.0.0.1/jwks", wantErr: "private or loopback"}, + {name: "private IP literal rejected", url: "https://10.1.2.3/jwks", wantErr: "private or loopback"}, + {name: "malformed URL rejected", url: "://not-a-url", wantErr: "invalid URL"}, + {name: "missing host rejected", url: "https:///jwks", wantErr: "host is required"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := ValidateJWKSURL(tt.url, tt.insecureAllowHTTP, false) + if tt.wantErr == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +// TestMultiIssuerTokenValidator_FetchJWKS exercises ensureRegistered's and +// lookupJWKS's error paths through the full Validate path, bypassing OIDC +// discovery via a preconfigured JWKSURL. Registration and the JWKS's own +// zero-keys/too-many-keys checks now go through the issuer's own jwk.Cache and +// lookupJWKS respectively rather than a private HTTP fetch. +// +// For "non-200 response" and "malformed JSON body", verified empirically: +// httprc.Resource's `ready` channel only ever closes on a *successful* +// fetch, so Register's WithWaitReady(true) wait (ensureRegistered's +// first-ever-registration path) can't distinguish "still fetching" from +// "fetch failed" — it just blocks until fetchCtx's own deadline and returns +// a generic timeout, discarding the real HTTP/parse error, which is why +// these two cases assert on "context deadline exceeded" rather than on +// jwx's own status-code or parse-error text (that text only ever surfaces +// on a *later* validation of the same never-yet-succeeded issuer, via +// ensureRegistered's Refresh-based retry branch — not exercised by a single +// Validate call here). "zero keys" and "too many keys" are unaffected: a +// JWKS that parses but fails this file's own key-count checks still +// registers successfully, so lookupJWKS's checks run immediately. +func TestMultiIssuerTokenValidator_FetchJWKS(t *testing.T) { + t.Parallel() + + // Built once: maxJWKSKeys+1 distinct public keys for the "too many keys" case. + tooManyKeys := make([]jose.JSONWebKey, maxJWKSKeys+1) + for i := range tooManyKeys { + key := newECDSAJWK(t, fmt.Sprintf("k%d", i)) + tooManyKeys[i] = key.Public() + } + + // A complete, otherwise-valid JWKS document padded past + // maxResponseBodySize with an unrelated field — see its use below for why + // that matters. + oversizedJWKSDoc := func() []byte { + paddingKey := newECDSAJWK(t, "padding-kid") + raw, err := json.Marshal(jose.JSONWebKeySet{Keys: []jose.JSONWebKey{paddingKey.Public()}}) + require.NoError(t, err) + var doc map[string]any + require.NoError(t, json.Unmarshal(raw, &doc)) + doc["padding"] = strings.Repeat("a", 4*1024*1024) + padded, err := json.Marshal(doc) + require.NoError(t, err) + return padded + }() + + tests := []struct { + name string + handler http.HandlerFunc + wantErr string + }{ + { + name: "non-200 response", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + }, + wantErr: "context deadline exceeded", + }, + { + name: "malformed JSON body", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte("{not-json")) + }, + wantErr: "context deadline exceeded", + }, + { + name: "zero keys", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(jose.JSONWebKeySet{Keys: []jose.JSONWebKey{}}) + }, + wantErr: "contains no keys", + }, + { + name: "too many keys", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(jose.JSONWebKeySet{Keys: tooManyKeys}) + }, + wantErr: "too many keys", + }, + { + // Pins limitedBodyTransport's cap on the JWKS fetch path (jwx's + // own httprc.MaxBufferSize ceiling is ~1000 MiB, far too high to + // bound anything here on its own). Deliberately WELL-FORMED and + // complete, unlike the other failure cases above: the padding + // field is oversized but the document would parse successfully + // if read in full, so only limitedBodyTransport cutting the read + // short makes this fail — the same technique + // TestMultiIssuerTokenValidator_DiscoverJWKSURL's oversized-body + // case uses for the discovery path. The 4 MiB padding size is a + // fixed literal independent of maxResponseBodySize (1 MiB): sizing + // it as a multiple of that constant would make a broken cap and a + // shrunken constant fail identically, hiding a regression in the + // cap itself. + name: "oversized JWKS response is rejected rather than parsed", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(oversizedJWKSDoc) + }, + wantErr: "context deadline exceeded", + }, + } + + selfJWKS := newTestJWKS(t) + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + mux := http.NewServeMux() + mux.HandleFunc("/jwks", tt.handler) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: srv.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + rawToken := signExternalToken(t, newECDSAJWK(t, "any-kid"), externalClaims()) + _, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +// signExternalToken signs claims with the given key for the external-issuer +// path, adding a may_act claim so the token satisfies delegation consent +// without needing an AllowedActors allowlist match — signWithJWK doesn't +// support extra claims like azp. +func signExternalToken(t *testing.T, key jose.JSONWebKey, claims jwt.Claims) string { + t.Helper() + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.ES256, Key: key}, + (&jose.SignerOptions{}).WithType("JWT"), + ) + require.NoError(t, err) + raw, err := jwt.Signed(signer). + Claims(claims). + Claims(map[string]any{"may_act": map[string]any{"sub": "some-toolhive-client", "iss": testIssuer}}). + Serialize() + require.NoError(t, err) + return raw +} + +// TestMultiIssuerTokenValidator_KeyRotationRefreshesImmediately is the test +// that demonstrates the point of the cache rewrite: a subject token signed +// with a newly rotated key, carrying a kid the validator has never seen, must +// validate successfully on the very first attempt — not after jwx's own +// 15-minute-minimum background refresh, and not only on a second, separate +// request. +func TestMultiIssuerTokenValidator_KeyRotationRefreshesImmediately(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + keyV1 := newECDSAJWK(t, "v1") + keyV2 := newECDSAJWK(t, "v2") + + var mu sync.Mutex + currentKey := keyV1 + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + mu.Lock() + key := currentKey + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(publicJWKSOf(key)) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: srv.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + // Prime the cache: validate successfully against v1. + _, err := validator.Validate(context.Background(), signExternalToken(t, keyV1, externalClaims())) + require.NoError(t, err) + + // Rotate: the server now serves only v2's public key, under a new kid. + mu.Lock() + currentKey = keyV2 + mu.Unlock() + + claims := externalClaims() + claims.ID = "jti-post-rotation" + tokenV2 := signExternalToken(t, keyV2, claims) + + done := make(chan struct { + result *ValidatedClaims + err error + }, 1) + go func() { + result, err := validator.Validate(context.Background(), tokenV2) + done <- struct { + result *ValidatedClaims + err error + }{result, err} + }() + + select { + case out := <-done: + require.NoError(t, out.err) + require.NotNil(t, out.result) + assert.Equal(t, "ext-user-456", out.result.Subject) + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for validation after key rotation") + } +} + +// TestMultiIssuerTokenValidator_UnknownKidRefreshIsRateLimited proves +// refreshOnUnknownKid's minKidRefreshInterval gate: repeated subject tokens +// naming a kid absent from the cached JWKS must not force a fetch per +// request — only the first ever unknown-kid attempt (within the interval) +// may do so. +func TestMultiIssuerTokenValidator_UnknownKidRefreshIsRateLimited(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + keyV1 := newECDSAJWK(t, "v1") + + var fetchCount atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + fetchCount.Add(1) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(publicJWKSOf(keyV1)) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: srv.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + // Prime the cache with a first, successful validation. + _, err := validator.Validate(context.Background(), signExternalToken(t, keyV1, externalClaims())) + require.NoError(t, err) + primedFetches := fetchCount.Load() + + // Repeatedly present a token naming a kid absent from the served JWKS. + unknownKey := newECDSAJWK(t, "unknown-kid") + for i := range 5 { + claims := externalClaims() + claims.ID = fmt.Sprintf("jti-unknown-%d", i) + rawToken := signExternalToken(t, unknownKey, claims) + _, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + } + + // Only the first unknown-kid attempt should have forced a refresh; the + // gate must hold for the remaining four within minKidRefreshInterval. + assert.Equal(t, primedFetches+1, fetchCount.Load(), + "repeated unknown-kid tokens within the window must not each force a fetch") +} + +// TestMultiIssuerTokenValidator_NeverFetchedRetryIsRateLimited proves +// ensureRegistered's jwksFetchFailureBackoff gate: an issuer whose endpoint +// has never once succeeded must not be re-fetched on every request — key +// resolution runs before signature verification, so without this gate any +// client holding a subject token naming this issuer could drive one real +// outbound fetch per request. +func TestMultiIssuerTokenValidator_NeverFetchedRetryIsRateLimited(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + + var fetchCount atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + fetchCount.Add(1) + w.WriteHeader(http.StatusInternalServerError) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: srv.URL + "/jwks", + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + rawToken := signExternalToken(t, newECDSAJWK(t, "any-kid"), externalClaims()) + + // First attempt: no cached value yet, so this one genuinely fetches (and + // blocks on Register's own wait — see ensureRegistered's doc comment). + _, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + require.Equal(t, int32(1), fetchCount.Load()) + + // Further attempts within jwksFetchFailureBackoff must replay the stored + // error instead of fetching again. + for range 3 { + _, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + } + assert.Equal(t, int32(1), fetchCount.Load(), + "repeated requests against a never-successfully-fetched issuer must not each force a fetch") +} + +// TestMultiIssuerTokenValidator_SharedJWKSURL_SamePolicy proves that two +// issuers resolving to the identical jwksURL under the identical HTTP +// transport policy both validate successfully — each through its own +// jwk.Cache and *http.Client (see externalIssuerConfig.jwksCache). This is +// the common real-world case: Microsoft Entra v1 tenants share one +// tenant-independent JWKS endpoint, so two Entra tenants configured as +// separate trusted issuers collide on the same jwks_url by construction. +func TestMultiIssuerTokenValidator_SharedJWKSURL_SamePolicy(t *testing.T) { + t.Parallel() + + const ( + issuerAURL = "https://issuer-a.example.com" + issuerBURL = "https://issuer-b.example.com" + audienceA = "aud-a" + audienceB = "aud-b" + ) + + selfJWKS := newTestJWKS(t) + sharedJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, sharedJWKS) + + trustedIssuers := []TrustedIssuer{ + { + IssuerURL: issuerAURL, + ExpectedAudience: audienceA, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"agent-a"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }, + { + IssuerURL: issuerBURL, + ExpectedAudience: audienceB, + JWKSURL: jwksServer.URL + "/jwks", + AllowedActors: []string{"agent-b"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }, + } + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + tokenFor := func(issuer, audience, actor, jti string) string { + now := time.Now() + claims := jwt.Claims{ + Subject: "shared-user", + Issuer: issuer, + Audience: jwt.Audience{audience}, + Expiry: jwt.NewNumericDate(now.Add(time.Hour)), + IssuedAt: jwt.NewNumericDate(now), + NotBefore: jwt.NewNumericDate(now.Add(-time.Minute)), + ID: jti, + } + return sharedJWKS.signToken(t, claims, map[string]any{"azp": actor}) + } + + resultA, err := validator.Validate(context.Background(), tokenFor(issuerAURL, audienceA, "agent-a", "jti-a")) + require.NoError(t, err, "first issuer to register the shared jwks_url must validate") + assert.Equal(t, issuerAURL, resultA.ExternalIssuer) + + resultB, err := validator.Validate(context.Background(), tokenFor(issuerBURL, audienceB, "agent-b", "jti-b")) + require.NoError(t, err, "second issuer sharing the same jwks_url under the same policy must also validate") + assert.Equal(t, issuerBURL, resultB.ExternalIssuer) +} + +// TestMultiIssuerTokenValidator_SharedJWKSURL_DifferingPolicy proves the +// property a per-issuer jwk.Cache adds over a shared one: two issuers +// resolving to the same jwks_url but configuring DIFFERENT +// insecure_allow_http/allow_private_ips settings both validate +// independently, each fetching through its own dedicated *http.Client. A +// shared cache could not do this — httprc keys a cached resource by URL +// alone and only honors jwk.WithHTTPClient on a URL's first Register call, +// so the second issuer would have silently inherited the first one's client +// and transport policy. Splitting the cache per issuer removes that +// collision instead of merely guarding against it. +func TestMultiIssuerTokenValidator_SharedJWKSURL_DifferingPolicy(t *testing.T) { + t.Parallel() + + const ( + issuerAURL = "https://issuer-a.example.com" + issuerBURL = "https://issuer-b.example.com" + audienceA = "aud-a" + audienceB = "aud-b" + ) + + selfJWKS := newTestJWKS(t) + sharedJWKS := newTestJWKS(t) + jwksServer := startJWKSServer(t, sharedJWKS) + sharedJWKSURL := jwksServer.URL + "/jwks" + + // Deliberately NOT newMultiValidator: that helper forces + // InsecureAllowHTTP and AllowPrivateIPs to true on every issuer so its + // loopback httptest servers are reachable, which would erase the very + // difference this test exists to exercise. Both issuers share one + // plain-HTTP loopback jwks_url and allow private IPs, and differ ONLY in + // InsecureAllowHTTP — so each is judged against its own transport policy: + // A is refused for its own reason (no HTTP permitted), B succeeds. + // + // Under a shared cache B could not succeed here: the policy-claim guard + // rejected any second issuer whose policy differed from the URL's first + // claimant, and without that guard B would have silently inherited A's + // client. Per-issuer caches make both outcomes independent. + selfValidator, err := NewSelfIssuedTokenValidator(selfJWKS.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + validator, err := NewMultiIssuerTokenValidator(selfValidator, testIssuer, []TrustedIssuer{ + { + IssuerURL: issuerAURL, + ExpectedAudience: audienceA, + JWKSURL: sharedJWKSURL, + AllowedActors: []string{"agent-a"}, + AllowPrivateIPs: true, + InsecureAllowHTTP: false, + AllowedDelegateClients: []string{anyDelegateClient}, + }, + { + IssuerURL: issuerBURL, + ExpectedAudience: audienceB, + JWKSURL: sharedJWKSURL, + AllowedActors: []string{"agent-b"}, + AllowPrivateIPs: true, + InsecureAllowHTTP: true, + AllowedDelegateClients: []string{anyDelegateClient}, + }, + }) + require.NoError(t, err) + + tokenFor := func(issuer, audience, actor, jti string) string { + now := time.Now() + claims := jwt.Claims{ + Subject: "shared-user", + Issuer: issuer, + Audience: jwt.Audience{audience}, + Expiry: jwt.NewNumericDate(now.Add(time.Hour)), + IssuedAt: jwt.NewNumericDate(now), + NotBefore: jwt.NewNumericDate(now.Add(-time.Minute)), + ID: jti, + } + return sharedJWKS.signToken(t, claims, map[string]any{"azp": actor}) + } + + // A is judged against its OWN policy: it forbids plain HTTP, so its fetch + // of the shared http:// jwks_url is refused. Not a policy-conflict error — + // A is simply misconfigured for this URL. + _, err = validator.Validate(context.Background(), tokenFor(issuerAURL, audienceA, "agent-a", "jti-a")) + require.Error(t, err, "the issuer forbidding plain HTTP must be refused for its own jwks_url") + assert.Contains(t, err.Error(), "must use HTTPS", + "the refusal must come from issuer A's own transport policy") + + // B shares that exact URL but permits HTTP, and succeeds — the outcome a + // shared cache could not produce, since A reached the URL first. + resultB, err := validator.Validate(context.Background(), tokenFor(issuerBURL, audienceB, "agent-b", "jti-b")) + require.NoError(t, err, "the issuer permitting HTTP must validate independently, "+ + "neither blocked by nor inheriting issuer A's stricter policy") + assert.Equal(t, issuerBURL, resultB.ExternalIssuer) +} + +// TestMultiIssuerTokenValidator_RetryAfterFetchFailureRefreshes proves the +// regression-safety half of the ensureRegistered rewrite (removing +// externalIssuerConfig.added in favor of asking issuerConfig.jwksCache.IsRegistered +// directly): once a JWKS fetch has failed but the resource was genuinely +// registered with the issuer's own cache, a later retry — once +// jwksFetchFailureBackoff has elapsed — must refresh the existing +// registration rather than attempt to register it again, which would fail +// with "already registered". +func TestMultiIssuerTokenValidator_RetryAfterFetchFailureRefreshes(t *testing.T) { + t.Parallel() + + selfJWKS := newTestJWKS(t) + externalJWKS := newTestJWKS(t) + + var succeed atomic.Bool + var fetchCount atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) { + fetchCount.Add(1) + if !succeed.Load() { + w.WriteHeader(http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(externalJWKS.publicJWKS()) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + trustedIssuers := []TrustedIssuer{{ + IssuerURL: testExternalIssuer, + ExpectedAudience: testExternalAudience, + JWKSURL: srv.URL + "/jwks", + AllowedActors: []string{"ext-agent"}, + AllowedDelegateClients: []string{anyDelegateClient}, + }} + validator := newMultiValidator(t, selfJWKS, trustedIssuers) + + tokenWithID := func(jti string) string { + claims := externalClaims() + claims.ID = jti + return externalJWKS.signToken(t, claims, map[string]any{"azp": "ext-agent"}) + } + + // First attempt: the endpoint is broken. This genuinely registers the + // resource with the issuer's own cache (per httprc.Controller.Add's own + // ordering), but the fetch itself fails. + _, err := validator.Validate(context.Background(), tokenWithID("jti-1")) + require.Error(t, err) + require.Equal(t, int32(1), fetchCount.Load()) + + // Force the backoff gate open, as if jwksFetchFailureBackoff had elapsed, + // without waiting the real 30s out. + issuerConfig := validator.issuers[testExternalIssuer] + issuerConfig.mu.Lock() + issuerConfig.fetchFailedAt = time.Now().Add(-jwksFetchFailureBackoff - time.Second) + issuerConfig.mu.Unlock() + + // The endpoint now recovers. The retry must go through Refresh — a + // second Register call against the same URL would fail with "already + // registered". + succeed.Store(true) + result, err := validator.Validate(context.Background(), tokenWithID("jti-2")) + require.NoError(t, err, "retry after a registered-but-failed fetch must refresh, not re-register") + require.NotNil(t, result) +} + +// TestLimitedBodyTransport asserts directly on the body cap that protects the +// JWKS fetch path. A direct test is necessary rather than sufficient coverage +// via Validate: jwx surfaces every fetch failure as its own WaitReady timeout, +// so the cap's error never reaches a caller and cannot be distinguished there +// from a 500, a parse failure, or a kid mismatch. Asserting on the cap itself +// is the only way to pin it — the oversized case in +// TestMultiIssuerTokenValidator_FetchJWKS proves the fetch fails, not why. +func TestLimitedBodyTransport(t *testing.T) { + t.Parallel() + + const bodyCap = 1024 + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(strings.Repeat("a", 8*1024))) + })) + t.Cleanup(srv.Close) + + client := srv.Client() + client.Transport = &limitedBodyTransport{base: client.Transport, max: bodyCap} + + resp, err := client.Get(srv.URL) + require.NoError(t, err, "the cap applies to reading the body, not to the round trip") + t.Cleanup(func() { _ = resp.Body.Close() }) + + body, err := io.ReadAll(resp.Body) + require.Error(t, err, "reading past the cap must fail rather than truncate silently: "+ + "a truncated JWKS would be parsed as though it were the whole document") + assert.LessOrEqual(t, int64(len(body)), int64(bodyCap), + "no more than the cap may be delivered before the error") +} diff --git a/pkg/authserver/server/tokenexchange/types.go b/pkg/authserver/server/tokenexchange/types.go new file mode 100644 index 0000000000..d5f7d9846e --- /dev/null +++ b/pkg/authserver/server/tokenexchange/types.go @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import "context" + +// SubjectTokenValidator validates subject tokens presented during RFC 8693 token exchange. +type SubjectTokenValidator interface { + Validate(ctx context.Context, rawToken string) (*ValidatedClaims, error) +} diff --git a/pkg/authserver/server/tokenexchange/validator.go b/pkg/authserver/server/tokenexchange/validator.go new file mode 100644 index 0000000000..665dba033a --- /dev/null +++ b/pkg/authserver/server/tokenexchange/validator.go @@ -0,0 +1,545 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package tokenexchange implements RFC 8693 token exchange for the authorization server. +// It provides validation of subject tokens that were issued by the same authorization +// server, enabling agents to act on behalf of users through delegation. +package tokenexchange + +import ( + "context" + "errors" + "fmt" + "slices" + "strings" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/go-jose/go-jose/v4/jwt" +) + +// allowedSignatureAlgorithms lists the signing algorithms accepted during JWT verification. +// This restricts which algorithms the validator will accept, preventing algorithm +// confusion attacks where an attacker might try to use a weaker algorithm. +var allowedSignatureAlgorithms = []jose.SignatureAlgorithm{ + jose.ES256, + jose.ES384, + jose.RS256, + jose.RS384, + jose.RS512, + jose.EdDSA, +} + +// ValidatedClaims holds the verified claims extracted from a subject token. +// All fields are populated from a successfully validated JWT that was issued +// by this authorization server. +type ValidatedClaims struct { + // Subject is the user identity from the "sub" claim (required for delegation). + Subject string + // Issuer is the token issuer from the "iss" claim. + Issuer string + // Audience is the list of intended recipients from the "aud" claim. + Audience []string + // Expiry is the token expiration time from the "exp" claim. + Expiry time.Time + // IssuedAt is the token issuance time from the "iat" claim. + IssuedAt time.Time + // JWTID is the unique token identifier from the "jti" claim. + JWTID string + // Name is the user's display name from the custom "name" claim. + Name string + // Email is the user's email address from the custom "email" claim. + Email string + // ClientID is the OAuth client ID from the custom "client_id" claim. + ClientID string + // Scopes is the space-delimited scope string assembled from the "scope" + // or "scp" claim. RFC 9068 §2.2.1 spells it "scope" as a JSON string, but + // fosite's default JWT claims strategy (token/jwt/claims_jwt.go) writes + // scopes as a JSON array under "scp" instead, unless ScopeField is + // explicitly set to String or Both — this server does not set it, so a + // genuine ToolHive-issued access token used as a subject token carries + // "scp", not "scope". When both are present, "scope" wins. Empty if the + // subject token carries neither claim. + Scopes string + // MayAct holds the authorized actor from the "may_act" claim (RFC 8693 §4.4). + // Nil when the subject token does not carry a may_act claim. + MayAct *MayActClaim + // ExternalActor is the client identity that the external-issuer validation + // path has already authorized for delegation, via a per-issuer actor-claim + // allowlist match. It is set ONLY by that path (multi_issuer_validator.go), + // after the resolved actor claim is confirmed present in the issuer's + // AllowedActors — never populated from token claims by buildValidatedClaims + // or assignClaim. It is empty for self-issued tokens and for external + // tokens that carry a may_act claim (may_act is authoritative there + // instead). A non-empty value means the validator has already authorized + // this token's actor for delegation; it is not itself a raw claim. + ExternalActor string + // ExternalIssuer is set by the external-issuer validation path + // (validateExternalToken in multi_issuer_validator.go) to that issuer's + // IssuerURL, for EVERY external token it validates — unlike + // ExternalActor, this is set regardless of whether the token carries a + // may_act claim. It exists so the handler can record provenance (which + // issuer a delegation actually originated from) even for a + // may_act-bearing external token, which leaves ExternalActor unset. + // Empty for self-issued tokens. Like ExternalActor, it is never + // populated from token claims by buildValidatedClaims or assignClaim — + // it comes from the already-validated issuer config the token matched, + // not from anything token-supplied, so it cannot be spoofed via claims. + ExternalIssuer string + // AllowedDelegateClients is set for EVERY external token — like + // ExternalIssuer and unlike ExternalActor, it does not depend on whether + // the token carries a may_act claim (see validateExternalToken). That + // matters: the may_act path bypasses the AllowedActors allowlist + // entirely, so it is the path that most needs this restriction to still + // apply. It is set to that issuer's configured + // TrustedIssuer.AllowedDelegateClients, and is never populated + // from token claims by buildValidatedClaims or assignClaim, and the + // validator never compares it against anything: the validator does not + // know the authenticated ToolHive client, so checkDelegationConsent + // (handler.go) is the one that checks actorID against this list. Nil + // means the issuer did not configure AllowedDelegateClients, which is + // permissive — any ToolHive client may use the allowlisted external + // actor, same as before this field existed. + AllowedDelegateClients []string + // Extra contains all non-standard claims not captured by other fields. + Extra map[string]any +} + +// MayActClaim represents the RFC 8693 §4.4 may_act claim from a subject token. +// It identifies the party authorized to act on behalf of the subject. +type MayActClaim struct { + Sub string `json:"sub"` + // Iss, when present, qualifies Sub's namespace. RFC 8693 §4.4: "the + // combination of the two claims 'iss' and 'sub' are sometimes necessary + // to uniquely identify an authorized actor." Without it, Sub alone could + // name a party in some other issuer's namespace while still being + // compared against a ToolHive client ID by checkDelegationConsent — a + // namespace-confusion bug, not just a missing feature. validateMayActShape + // requires Iss, when present, to equal this authorization server's own + // issuer, so by the time checkDelegationConsent reads Sub it is + // guaranteed to be in ToolHive's own client namespace. + Iss string `json:"iss,omitempty"` +} + +// Compile-time check that SelfIssuedTokenValidator implements SubjectTokenValidator. +var _ SubjectTokenValidator = (*SelfIssuedTokenValidator)(nil) + +// SelfIssuedTokenValidator validates subject tokens presented during RFC 8693 token exchange. +// It verifies that the token was issued by this authorization server by checking +// the signature against the server's own JWKS, and validates standard JWT claims. +type SelfIssuedTokenValidator struct { + publicJWKS *jose.JSONWebKeySet + issuer string + allowedAudiences []string +} + +// NewSelfIssuedTokenValidator creates a new validator for subject tokens. +// The jwks parameter must be non-nil and contain only the authorization server's +// public signing keys (e.g. AuthorizationServerConfig.PublicJWKS) — the validator +// only ever verifies signatures, so it must not be handed private key material. +// The issuer parameter is the expected "iss" claim value. allowedAudiences is the +// set of audiences this server accepts in a subject token's "aud" claim; per the +// same secure default as AuthorizationServerConfig.AllowedAudiences, an empty +// allowedAudiences rejects every subject token rather than skipping the check. +func NewSelfIssuedTokenValidator( + jwks *jose.JSONWebKeySet, issuer string, allowedAudiences []string, +) (*SelfIssuedTokenValidator, error) { + if jwks == nil { + return nil, fmt.Errorf("JWKS must not be nil") + } + if issuer == "" { + return nil, fmt.Errorf("issuer must not be empty") + } + return &SelfIssuedTokenValidator{ + publicJWKS: jwks, + issuer: issuer, + allowedAudiences: allowedAudiences, + }, nil +} + +// Validate parses and verifies a raw JWT subject token. +// It checks the signature against the server's JWKS, validates issuer and audience, +// ensures the token is not expired, and requires a subject claim for delegation. +// +// The subject token's "aud" claim is checked against allowedAudiences, but this +// validator deliberately does not require that the authorization server itself +// (i.e. the token endpoint) be among those audiences. RFC 8693 leaves subject-token +// validation criteria out of scope, and ToolHive's vMCP flow legitimately exchanges +// tokens addressed to a downstream/upstream resource rather than the AS. The residual +// cross-resource risk is mitigated elsewhere: the token is still pinned to the +// server-wide allowedAudiences, and the handler enforces the requested resource +// against the client's registered audiences. +// +// Returns the validated claims on success, or a descriptive error on failure. +func (v *SelfIssuedTokenValidator) Validate(_ context.Context, rawToken string) (*ValidatedClaims, error) { + parsedToken, err := jwt.ParseSigned(rawToken, allowedSignatureAlgorithms) + if err != nil { + return nil, fmt.Errorf("subject token is not a valid JWT: %w", err) + } + + // kidMatched is discarded here: it exists to let the external-issuer path + // (validateExternalToken in multi_issuer_validator.go) decide whether an + // unknown kid is worth an on-demand JWKS refresh. This validator's JWKS + // never comes over HTTP — it's this authorization server's own + // PublicJWKS() — so there is nothing to refresh. + standardClaims, extraClaims, _, err := verifySignature(parsedToken, v.publicJWKS) + if err != nil { + return nil, err + } + + // Validate issuer and expiry. Audience is intentionally excluded from + // jwt.Expected here — go-jose's AnyAudience check is skipped entirely when + // empty, which is the opposite of this server's secure default (empty + // AllowedAudiences means no audience is permitted). Audience is checked + // explicitly below so an empty allowlist fails closed. + expected := jwt.Expected{ + Issuer: v.issuer, + } + if err := standardClaims.ValidateWithLeeway(expected, 0); err != nil { + return nil, fmt.Errorf("subject token claims validation failed: %w", err) + } + + if err := validateAudience(standardClaims.Audience, v.allowedAudiences); err != nil { + return nil, fmt.Errorf("subject token claims validation failed: %w", err) + } + + // Expiry is required for delegation — a subject token without an expiry + // cannot be safely bounded by the delegation lifetime. + if standardClaims.Expiry == nil { + return nil, fmt.Errorf("subject token is missing required 'exp' claim") + } + + // Subject is required for delegation — the token must identify the user. + if standardClaims.Subject == "" { + return nil, fmt.Errorf("subject token is missing required 'sub' claim") + } + + // Reject a subject token that is actually an ID token — see + // rejectIDTokenClaims's doc comment. + if err := rejectIDTokenClaims(extraClaims); err != nil { + return nil, err + } + + // If may_act is present, it must be well-formed. A present-but-malformed + // may_act is treated as an invalid token, not an absent one — fail closed + // to prevent silent downgrade to client_id. requireIss is false here: see + // validateMayActShape's doc comment for why the self-issued path doesn't + // need it. + if err := validateMayActShape(extraClaims, v.issuer, false); err != nil { + return nil, err + } + + return buildValidatedClaims(standardClaims, extraClaims), nil +} + +// verifySignature attempts to verify the JWT signature using the provided public keys. +// If the JWT header names a key ID (kid) that matches a key in the JWKS, only that +// key is tried: a token naming a kid but signed by a different key is a spoofed-kid +// attempt and must fail, not fall back to verifying against the rest of the keyset. +// When there is no kid, or the kid doesn't match any key in the JWKS, every key whose +// declared "use"/"alg" (RFC 7517 §4.2/§4.4) is compatible with this signature is +// tried in turn — a key marked for a different use (e.g. "enc") or a different +// algorithm is skipped rather than tried, per RFC 8725 §3.12's recommendation to +// scope keys to a single purpose as a substitution defense. Returns on the first +// successful verification. On failure, the last verification error is wrapped. +// +// The returned kidMatched is true only when the token's kid header was found among +// publicJWKS's keys — regardless of whether the signature itself then verified. The +// external-issuer path (multi_issuer_validator.go) uses this to tell "this key +// genuinely isn't in our cached JWKS yet, possibly a rotation" (kidMatched false) +// apart from "this kid exists here but the signature is wrong" (kidMatched true, +// which a refresh can never fix and must not retry). +func verifySignature( + token *jwt.JSONWebToken, + publicJWKS *jose.JSONWebKeySet, +) (jwt.Claims, map[string]any, bool, error) { + candidates := publicJWKS.Keys + var kidMatched bool + if len(token.Headers) > 0 && token.Headers[0].KeyID != "" { + if matched := publicJWKS.Key(token.Headers[0].KeyID); len(matched) > 0 { + candidates = matched + kidMatched = true + } + } + + var headerAlg string + if len(token.Headers) > 0 { + headerAlg = token.Headers[0].Algorithm + } + + var lastErr error + var filtered int + for _, key := range candidates { + if key.Use != "" && key.Use != "sig" { + filtered++ + continue + } + if key.Algorithm != "" && key.Algorithm != headerAlg { + filtered++ + continue + } + var claims jwt.Claims + extra := map[string]any{} + err := token.Claims(key, &claims, &extra) + if err == nil { + return claims, extra, kidMatched, nil + } + lastErr = err + } + if lastErr != nil { + return jwt.Claims{}, nil, kidMatched, fmt.Errorf("subject token signature verification failed: %w", lastErr) + } + if filtered > 0 { + return jwt.Claims{}, nil, kidMatched, fmt.Errorf( + "subject token signature verification failed: no compatible keys in JWKS "+ + "(%d key(s) present but all skipped by use/alg filter)", + len(candidates)) + } + return jwt.Claims{}, nil, kidMatched, fmt.Errorf("subject token signature verification failed: no keys in JWKS") +} + +// validateAudience checks that tokenAudience intersects allowedAudiences. +// Per this server's secure default (see AuthorizationServerConfig.AllowedAudiences), +// an empty allowedAudiences means no audience is permitted, so the check fails +// closed rather than skipping validation. +func validateAudience(tokenAudience jwt.Audience, allowedAudiences []string) error { + if len(allowedAudiences) == 0 { + return fmt.Errorf("no audiences are configured on this server") + } + for _, aud := range tokenAudience { + if slices.Contains(allowedAudiences, aud) { + return nil + } + } + return fmt.Errorf("token audience %v does not match any allowed audience", []string(tokenAudience)) +} + +// buildValidatedClaims constructs a ValidatedClaims from standard and extra claims. +func buildValidatedClaims( + standard jwt.Claims, + extra map[string]any, +) *ValidatedClaims { + vc := &ValidatedClaims{ + Subject: standard.Subject, + Issuer: standard.Issuer, + Audience: []string(standard.Audience), + JWTID: standard.ID, + Extra: make(map[string]any), + } + + if standard.Expiry != nil { + vc.Expiry = standard.Expiry.Time() + } + if standard.IssuedAt != nil { + vc.IssuedAt = standard.IssuedAt.Time() + } + + // Extract well-known custom claims and collect the rest into Extra. + // Standard JWT registered claims are filtered out since they are already + // captured in the structured fields above. + for k, val := range extra { + assignClaim(vc, k, val) + } + + // Fall back to "scp" only if "scope" (handled above) didn't already set + // Scopes. This must run after the loop, not inside assignClaim's per-key + // switch, because map iteration order is random — assignClaim can't tell + // whether a not-yet-seen "scope" claim will still show up. + // + // Without this fallback, a genuine self-issued subject token (which + // carries "scp", not "scope") reads as scope-less and any scoped + // exchange fails with invalid_scope — scoped delegation against a + // self-issued token was effectively unusable before this fallback. + if vc.Scopes == "" { + if scp, ok := extra["scp"]; ok { + vc.Scopes = scpToScopeString(scp) + } + } + + return vc +} + +// scpToScopeString normalizes an "scp" claim value into a space-delimited +// scope string. fosite's default JWT claims strategy writes "scp" as a JSON +// array (see the Scopes field doc comment), which json.Unmarshal decodes as +// []any — but a plain space-delimited string is accepted too, in case an +// issuer writes it that way instead. Non-string array elements are skipped +// rather than rejected: a malformed entry degrades to a smaller scope set +// instead of failing subject-token validation outright. +func scpToScopeString(val any) string { + switch v := val.(type) { + case string: + return v + case []any: + scopes := make([]string, 0, len(v)) + for _, item := range v { + if s, ok := item.(string); ok { + scopes = append(scopes, s) + } + } + return strings.Join(scopes, " ") + default: + return "" + } +} + +// assignClaim routes one non-standard JWT claim onto its structured +// ValidatedClaims field, or into Extra if it isn't a well-known claim. +// Registered JWT claims (sub, iss, aud, exp, iat, nbf, jti) are dropped — +// they're already captured in buildValidatedClaims's structured fields. +// "scp" is likewise dropped here, though its value is read separately (see +// buildValidatedClaims's post-loop fallback) since its precedence relative +// to "scope" can't be decided from a single, order-independent key/value +// pair. Keep actorClaimsNotInExtra (multi_issuer_validator.go) in sync with +// the cases here. +func assignClaim(vc *ValidatedClaims, key string, val any) { + switch key { + case "name": + if s, ok := val.(string); ok { + vc.Name = s + } + case "email": + if s, ok := val.(string); ok { + vc.Email = s + } + case "client_id": + if s, ok := val.(string); ok { + vc.ClientID = s + } + case "scope": + if s, ok := val.(string); ok { + vc.Scopes = s + } + case "scp": + // Handled by buildValidatedClaims after its call to assignClaim for + // every key completes, so "scope" (if present) wins regardless of + // map iteration order. Case exists here only to keep "scp" out of + // Extra, like every other well-known claim below. + case "may_act": + if m, ok := val.(map[string]any); ok { + if s, ok := m["sub"].(string); ok { + mayAct := &MayActClaim{Sub: s} + if iss, ok := m["iss"].(string); ok { + mayAct.Iss = iss + } + vc.MayAct = mayAct + } + } + case "sub", "iss", "aud", "exp", "iat", "nbf", "jti": + // Skip registered JWT claims — already in structured fields. + default: + vc.Extra[key] = val + } +} + +// validateMayActShape rejects a present-but-malformed may_act claim. Both +// validation paths (self-issued here, external in multi_issuer_validator.go) +// call this before buildValidatedClaims: assignClaim silently drops a +// malformed may_act rather than surfacing it, and each path's fallback +// consent signal when may_act is absent (client_id here, the actor allowlist +// externally) is weaker than what a well-formed may_act would have granted. +// Letting a malformed claim fall through would silently widen consent +// instead of rejecting the token. +// +// selfIssuer is this authorization server's own issuer identifier. When +// may_act carries an "iss" member, RFC 8693 §4.4 uses it together with "sub" +// to identify the actor's namespace; this server only ever grants delegation +// to its own clients (checkDelegationConsent compares may_act.sub against a +// ToolHive client ID), so an "iss" naming any other issuer would mean sub is +// being read out of the wrong namespace. Requiring iss == selfIssuer when +// present — same fail-closed treatment as a malformed sub — prevents that +// namespace confusion. +// +// requireIss makes "iss" mandatory rather than merely constrained when +// present. The self-issued path passes false: the subject token's own "iss" +// already equals selfIssuer (Validate routed it here on exactly that basis), +// so may_act.iss would be redundant to require. The external path +// (multi_issuer_validator.go) passes true: there, an external issuer that +// emits no "iss" on may_act could otherwise authorize ANY ToolHive client by +// naming its ID in a bare may_act.sub, entirely bypassing this issuer's +// AllowedActors/AllowedDelegateClients allowlists — the one consent path +// that does. No major IdP is known to emit may_act at all (Microsoft Entra, +// Okta, and Google do not; Keycloak emits "act", a different claim), so +// reaching this path already requires an operator configuring a custom +// claim mapper for a nonstandard claim — an operator capable of adding "sub" +// to that mapper is assumed capable of adding "iss" too. That assumption is +// not verified against any real IdP's behavior in this repo; it is the +// rationale for requireIss, not a tested fact. +func validateMayActShape(extra map[string]any, selfIssuer string, requireIss bool) error { + raw, ok := extra["may_act"] + if !ok || raw == nil { + return nil + } + m, ok := raw.(map[string]any) + if !ok { + return errors.New("subject token has malformed 'may_act' claim: expected a JSON object") + } + if sub, ok := m["sub"].(string); !ok || sub == "" { + return errors.New("subject token has malformed 'may_act' claim: missing or invalid 'sub'") + } + rawIss, present := m["iss"] + if !present { + if requireIss { + return errors.New("subject token has malformed 'may_act' claim: missing required 'iss'") + } + return nil + } + iss, ok := rawIss.(string) + if !ok || iss == "" { + return errors.New("subject token has malformed 'may_act' claim: invalid 'iss'") + } + if iss != selfIssuer { + return fmt.Errorf( + "subject token has malformed 'may_act' claim: 'iss' %q does not match this authorization server's issuer", + iss) + } + return nil +} + +// rejectIDTokenClaims rejects a subject token that carries "at_hash" or +// "c_hash" — OIDC Core §3.3.2.11 / §3.3.2.10 define both for ID tokens only, +// never present on an access token. This is a sound POSITIVE detector (their +// presence proves an ID token) but is NOT EXHAUSTIVE: both claims are +// OPTIONAL even on a valid authorization-code-flow ID token, so an ID token +// that omits them passes this check undetected. A "typ: at+jwt" header check +// is not viable as a replacement: Entra v1/v2, Okta, and Google all emit bare +// "JWT" for access tokens too. Deliberately does NOT check "nonce": Microsoft +// Graph access tokens carry one, so rejecting on its presence would reject +// legitimate Entra access tokens. +// +// The other, complementary layer is the audience-shape check on the external +// path: TrustedIssuer.ExpectedAudience (multi_issuer_validator.go) is +// documented as MUST be a resource identifier, not a client ID, since an ID +// token's "aud" is the requesting client's own ID. That check is +// operator-dependent, not self-enforcing — see ExpectedAudience's doc +// comment and looksLikeResourceIdentifier's startup warning. Between the two, +// ID/access-token discrimination for an external issuer is partial: an +// operator who mis-sets ExpectedAudience to a client ID, presented with a +// real IdP whose ID tokens omit at_hash/c_hash (both legal per OIDC Core), +// has no remaining defense here. The blast radius is bounded elsewhere: an +// admitted ID token carries no "scope"/"scp" claim, so grantScopes +// (handler.go) either rejects the exchange with invalid_scope or grants a +// zero-scope delegated token. +// +// A third discriminator was considered and rejected: an ID token's "aud" +// MUST equal its own "azp"/authorized-party claim (OIDC Core §2, §3.1.3.7 +// step 9), so rejecting a token whose "aud" equals the issuer's configured +// actor claim looked like a candidate. It is not shipped, because it also +// fires on a legitimate access token: an app whose API and client share one +// registration — e.g. an Entra v1 app that calls its own exposed API, or any +// setup where ExpectedAudience is misconfigured to the app's bare +// client/resource GUID, which is exactly the ambiguous case this file's +// audience-shape check already warns about — issues access tokens whose +// "aud" is that shared ID and whose azp/appid is the very same ID. Rejecting +// on aud == azp would then reject that provider's genuine access tokens, not +// just its ID tokens. A false-positive risk of that shape is worse than the +// gap it would close. +func rejectIDTokenClaims(extra map[string]any) error { + if _, ok := extra["at_hash"]; ok { + return errors.New("subject token carries an 'at_hash' claim: ID tokens are not accepted as subject tokens") + } + if _, ok := extra["c_hash"]; ok { + return errors.New("subject token carries a 'c_hash' claim: ID tokens are not accepted as subject tokens") + } + return nil +} diff --git a/pkg/authserver/server/tokenexchange/validator_test.go b/pkg/authserver/server/tokenexchange/validator_test.go new file mode 100644 index 0000000000..56877792d4 --- /dev/null +++ b/pkg/authserver/server/tokenexchange/validator_test.go @@ -0,0 +1,644 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package tokenexchange + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/rsa" + "testing" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/go-jose/go-jose/v4/jwt" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testIssuer = "https://auth.example.com" + +// testJWKS holds a signing key and JWKS for test token generation. +type testJWKS struct { + privateKey *ecdsa.PrivateKey + jwk jose.JSONWebKey + jwks *jose.JSONWebKeySet +} + +// newTestJWKS generates a fresh ECDSA P-256 key pair and wraps it in a JWKS. +func newTestJWKS(t *testing.T) *testJWKS { + t.Helper() + + privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + + jwk := jose.JSONWebKey{ + Key: privateKey, + KeyID: "test-key-1", + Algorithm: string(jose.ES256), + Use: "sig", + } + + return &testJWKS{ + privateKey: privateKey, + jwk: jwk, + jwks: &jose.JSONWebKeySet{Keys: []jose.JSONWebKey{jwk}}, + } +} + +// publicJWKS returns the public-only JWKS, mirroring how Factory constructs +// the validator from AuthorizationServerConfig.PublicJWKS(). +func (tj *testJWKS) publicJWKS() *jose.JSONWebKeySet { + return &jose.JSONWebKeySet{Keys: []jose.JSONWebKey{tj.jwk.Public()}} +} + +// signToken creates a signed JWT with the given claims using the test JWKS key. +func (tj *testJWKS) signToken(t *testing.T, claims jwt.Claims, extraClaims map[string]any) string { + t.Helper() + + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.ES256, Key: tj.jwk}, + (&jose.SignerOptions{}).WithType("JWT"), + ) + require.NoError(t, err) + + builder := jwt.Signed(signer).Claims(claims) + if extraClaims != nil { + builder = builder.Claims(extraClaims) + } + + raw, err := builder.Serialize() + require.NoError(t, err) + + return raw +} + +// validClaims returns a set of standard claims that pass all validation checks. +func validClaims() jwt.Claims { + now := time.Now() + return jwt.Claims{ + Subject: "user-123", + Issuer: testIssuer, + Audience: jwt.Audience{testIssuer}, + Expiry: jwt.NewNumericDate(now.Add(time.Hour)), + IssuedAt: jwt.NewNumericDate(now), + NotBefore: jwt.NewNumericDate(now.Add(-time.Minute)), + ID: "jti-abc-123", + } +} + +// validExtraClaims returns extra claims that match the session package claim keys. +func validExtraClaims() map[string]any { + return map[string]any{ + "name": "Test User", + "email": "test@example.com", + "client_id": testAgentClientID, + "tsid": "session-xyz", + } +} + +func TestSelfIssuedTokenValidator_NewValidation(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + + t.Run("nil JWKS returns error", func(t *testing.T) { + t.Parallel() + _, err := NewSelfIssuedTokenValidator(nil, testIssuer, []string{testIssuer}) + require.Error(t, err) + assert.Contains(t, err.Error(), "JWKS must not be nil") + }) + + t.Run("empty issuer returns error", func(t *testing.T) { + t.Parallel() + _, err := NewSelfIssuedTokenValidator(tj.publicJWKS(), "", []string{testIssuer}) + require.Error(t, err) + assert.Contains(t, err.Error(), "issuer must not be empty") + }) + + t.Run("valid params succeed", func(t *testing.T) { + t.Parallel() + v, err := NewSelfIssuedTokenValidator(tj.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + assert.NotNil(t, v) + }) +} + +func TestSelfIssuedTokenValidator_Validate(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + + tests := []struct { + name string + token func(t *testing.T) string + wantErr bool + errContains string + check func(t *testing.T, vc *ValidatedClaims) + }{ + { + name: "valid token with all claims", + token: func(t *testing.T) string { + t.Helper() + return tj.signToken(t, validClaims(), validExtraClaims()) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + assert.Equal(t, "user-123", vc.Subject) + assert.Equal(t, testIssuer, vc.Issuer) + assert.Equal(t, []string{testIssuer}, vc.Audience) + assert.Equal(t, "jti-abc-123", vc.JWTID) + assert.Equal(t, "Test User", vc.Name) + assert.Equal(t, "test@example.com", vc.Email) + assert.Equal(t, testAgentClientID, vc.ClientID) + assert.False(t, vc.Expiry.IsZero()) + assert.False(t, vc.IssuedAt.IsZero()) + // tsid should be in Extra since it's not a well-known claim field + assert.Equal(t, "session-xyz", vc.Extra["tsid"]) + // Standard JWT claims should NOT be in Extra + _, hasSub := vc.Extra["sub"] + assert.False(t, hasSub, "standard claim 'sub' should not be in Extra") + _, hasIss := vc.Extra["iss"] + assert.False(t, hasIss, "standard claim 'iss' should not be in Extra") + }, + }, + { + name: "expired token", + token: func(t *testing.T) string { + t.Helper() + claims := validClaims() + claims.Expiry = jwt.NewNumericDate(time.Now().Add(-time.Hour)) + claims.IssuedAt = jwt.NewNumericDate(time.Now().Add(-2 * time.Hour)) + return tj.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + name: "wrong issuer", + token: func(t *testing.T) string { + t.Helper() + claims := validClaims() + claims.Issuer = "https://evil.example.com" + return tj.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + name: "wrong audience", + token: func(t *testing.T) string { + t.Helper() + claims := validClaims() + claims.Audience = jwt.Audience{"https://other.example.com"} + return tj.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + name: "not yet valid — nbf in the future", + token: func(t *testing.T) string { + t.Helper() + claims := validClaims() + claims.NotBefore = jwt.NewNumericDate(time.Now().Add(time.Hour)) + return tj.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "claims validation failed", + }, + { + name: "bad signature — signed with different key", + token: func(t *testing.T) string { + t.Helper() + otherJWKS := newTestJWKS(t) + return otherJWKS.signToken(t, validClaims(), nil) + }, + wantErr: true, + errContains: "signature verification failed", + }, + { + name: "malformed token", + token: func(_ *testing.T) string { + return "not-a-jwt" + }, + wantErr: true, + errContains: "not a valid JWT", + }, + { + name: "missing subject claim", + token: func(t *testing.T) string { + t.Helper() + claims := validClaims() + claims.Subject = "" + return tj.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "missing required 'sub' claim", + }, + { + name: "missing exp claim", + token: func(t *testing.T) string { + t.Helper() + claims := validClaims() + claims.Expiry = nil + return tj.signToken(t, claims, nil) + }, + wantErr: true, + errContains: "missing required 'exp' claim", + }, + { + name: "token with may_act claim extracts MayAct", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": "authorized-agent"} + return tj.signToken(t, validClaims(), extra) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Equal(t, "authorized-agent", vc.MayAct.Sub) + }, + }, + { + name: "malformed may_act — not an object", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = "not-an-object" + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "malformed may_act — missing sub", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"foo": "bar"} + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "malformed may_act — non-string sub", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": 123} + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "may_act with iss matching this server's own issuer accepted", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": "authorized-agent", "iss": testIssuer} + return tj.signToken(t, validClaims(), extra) + }, + check: func(t *testing.T, vc *ValidatedClaims) { + t.Helper() + require.NotNil(t, vc.MayAct) + assert.Equal(t, "authorized-agent", vc.MayAct.Sub) + assert.Equal(t, testIssuer, vc.MayAct.Iss) + }, + }, + { + // RFC 8693 §4.4: sub is only meaningful together with iss. This + // server only ever grants delegation to its own clients, so an + // iss naming a different issuer is namespace confusion, not a + // legitimate qualifier — reject rather than silently ignore it. + name: "malformed may_act — iss not matching this server's own issuer", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": "authorized-agent", "iss": "https://evil.example.com"} + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "malformed may_act — empty iss", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["may_act"] = map[string]any{"sub": "authorized-agent", "iss": ""} + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "malformed 'may_act' claim", + }, + { + name: "ID token masquerading as subject token rejected — at_hash present", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["at_hash"] = "some-hash-value" + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "'at_hash'", + }, + { + name: "ID token masquerading as subject token rejected — c_hash present", + token: func(t *testing.T) string { + t.Helper() + extra := validExtraClaims() + extra["c_hash"] = "some-hash-value" + return tj.signToken(t, validClaims(), extra) + }, + wantErr: true, + errContains: "'c_hash'", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + validator, err := NewSelfIssuedTokenValidator(tj.publicJWKS(), testIssuer, []string{testIssuer}) + require.NoError(t, err) + rawToken := tt.token(t) + + result, err := validator.Validate(context.Background(), rawToken) + + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errContains) + assert.Nil(t, result) + return + } + + require.NoError(t, err) + require.NotNil(t, result) + if tt.check != nil { + tt.check(t, result) + } + }) + } +} + +func TestSelfIssuedTokenValidator_AudienceValidation(t *testing.T) { + t.Parallel() + + tj := newTestJWKS(t) + + tests := []struct { + name string + allowedAudiences []string + tokenAudience jwt.Audience + wantErr bool + }{ + { + name: "empty allowedAudiences rejects even a matching-looking token", + allowedAudiences: nil, + tokenAudience: jwt.Audience{testIssuer}, + wantErr: true, + }, + { + name: "token audience intersects allowedAudiences", + allowedAudiences: []string{"https://other.example.com", testIssuer}, + tokenAudience: jwt.Audience{testIssuer}, + wantErr: false, + }, + { + name: "token audience does not intersect allowedAudiences", + allowedAudiences: []string{"https://other.example.com"}, + tokenAudience: jwt.Audience{testIssuer}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + validator, err := NewSelfIssuedTokenValidator(tj.publicJWKS(), testIssuer, tt.allowedAudiences) + require.NoError(t, err) + + claims := validClaims() + claims.Audience = tt.tokenAudience + rawToken := tj.signToken(t, claims, nil) + + result, err := validator.Validate(context.Background(), rawToken) + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), "claims validation failed") + assert.Nil(t, result) + return + } + require.NoError(t, err) + require.NotNil(t, result) + }) + } +} + +// newECDSAJWK generates a fresh ECDSA P-256 key with the given key ID. The +// returned JWK carries the private key (in its Key field), so it can be used +// directly for both signing (signWithJWK) and JWKS construction. +func newECDSAJWK(t *testing.T, kid string) jose.JSONWebKey { + t.Helper() + privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + return jose.JSONWebKey{ + Key: privateKey, + KeyID: kid, + Algorithm: string(jose.ES256), + Use: "sig", + } +} + +// publicJWKSOf builds a JWKS containing only the public portion of each key, +// mirroring how Factory constructs the validator from PublicJWKS(). +func publicJWKSOf(keys ...jose.JSONWebKey) *jose.JSONWebKeySet { + public := make([]jose.JSONWebKey, len(keys)) + for i, k := range keys { + public[i] = k.Public() + } + return &jose.JSONWebKeySet{Keys: public} +} + +// signWithJWK signs claims with the given signing key (no extra claims). +func signWithJWK(t *testing.T, signingKey jose.JSONWebKey, alg jose.SignatureAlgorithm, claims jwt.Claims) string { + t.Helper() + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: alg, Key: signingKey}, + (&jose.SignerOptions{}).WithType("JWT"), + ) + require.NoError(t, err) + raw, err := jwt.Signed(signer).Claims(claims).Serialize() + require.NoError(t, err) + return raw +} + +func TestSelfIssuedTokenValidator_MultiKeyJWKS(t *testing.T) { + t.Parallel() + + t.Run("token verified with kid-matched key", func(t *testing.T) { + t.Parallel() + + jwk1 := newECDSAJWK(t, "test-key-1") + jwk2 := newECDSAJWK(t, "test-key-2") + jwks := publicJWKSOf(jwk1, jwk2) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + rawToken := signWithJWK(t, jwk2, jose.ES256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "user-123", result.Subject) + }) + + t.Run("token verified by full iteration when kid absent", func(t *testing.T) { + t.Parallel() + + jwk1 := newECDSAJWK(t, "test-key-1") + jwk2 := newECDSAJWK(t, "test-key-2") + jwks := publicJWKSOf(jwk1, jwk2) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + // Sign with jwk1 (whose public half is in the JWKS) but omit the kid + // from the signer so the token header carries no kid — the validator + // must then fall back to iterating all keys rather than kid-matching. + signingKeyNoKID := jwk1 + signingKeyNoKID.KeyID = "" + rawToken := signWithJWK(t, signingKeyNoKID, jose.ES256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "user-123", result.Subject) + }) + + t.Run("token signed with different key in JWKS fails", func(t *testing.T) { + t.Parallel() + + jwk1 := newECDSAJWK(t, "test-key-1") + jwks := publicJWKSOf(jwk1) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + // Key B is not in the JWKS. + otherJWK := newECDSAJWK(t, "other-key") + rawToken := signWithJWK(t, otherJWK, jose.ES256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), "signature verification failed") + assert.Nil(t, result) + }) + + t.Run("token claiming a kid it is not signed with fails", func(t *testing.T) { + t.Parallel() + + jwk1 := newECDSAJWK(t, "test-key-1") + jwk2 := newECDSAJWK(t, "test-key-2") + jwks := publicJWKSOf(jwk1, jwk2) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + // Sign with jwk2's key material but claim jwk1's kid in the header. + // A validator that fell back to the full keyset after a kid match + // would wrongly accept this via jwk2; verifySignature must try only + // the matched key (jwk1) and fail. + spoofedKey := jwk2 + spoofedKey.KeyID = jwk1.KeyID + rawToken := signWithJWK(t, spoofedKey, jose.ES256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), "signature verification failed") + assert.Nil(t, result) + }) + + t.Run("RSA-signed token accepted", func(t *testing.T) { + t.Parallel() + + rsaKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + rsaJWK := jose.JSONWebKey{ + Key: rsaKey, + KeyID: "rsa-key-1", + Algorithm: string(jose.RS256), + Use: "sig", + } + jwks := publicJWKSOf(rsaJWK) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + rawToken := signWithJWK(t, rsaJWK, jose.RS256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "user-123", result.Subject) + }) + + t.Run("key marked for a non-signing use is never tried", func(t *testing.T) { + t.Parallel() + + // Same key material a signature could actually verify against, but + // declared "enc" in the JWKS — RFC 7517 §4.2 reserves it for a + // different purpose. Without the use filter, this key would still + // verify the signature below and wrongly accept the token. + jwk := newECDSAJWK(t, "enc-key") + jwk.Use = "enc" + jwks := publicJWKSOf(jwk) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + // No kid on the token, so verifySignature must fall back to trying + // every key in the JWKS — and skip this one. + signingKeyNoKID := jwk + signingKeyNoKID.KeyID = "" + rawToken := signWithJWK(t, signingKeyNoKID, jose.ES256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), "signature verification failed") + assert.Nil(t, result) + }) + + t.Run("key with a mismatched declared algorithm is never tried", func(t *testing.T) { + t.Parallel() + + // The JWKS entry claims RS256, but the token header (and the actual + // signing) uses ES256 — a misconfigured or spoofed JWKS entry. + // Without the alg filter, go-jose would still attempt (and, for an + // EC key, fail) this candidate; the point here is that it must be + // skipped rather than tried at all. + jwk := newECDSAJWK(t, "mismatched-alg-key") + jwk.Algorithm = string(jose.RS256) + jwks := publicJWKSOf(jwk) + + validator, err := NewSelfIssuedTokenValidator(jwks, testIssuer, []string{testIssuer}) + require.NoError(t, err) + + signingKeyNoKID := jwk + signingKeyNoKID.KeyID = "" + rawToken := signWithJWK(t, signingKeyNoKID, jose.ES256, validClaims()) + + result, err := validator.Validate(context.Background(), rawToken) + require.Error(t, err) + assert.Contains(t, err.Error(), "signature verification failed") + assert.Nil(t, result) + }) +} diff --git a/pkg/authserver/server_impl.go b/pkg/authserver/server_impl.go index 007446ce12..458ac1b273 100644 --- a/pkg/authserver/server_impl.go +++ b/pkg/authserver/server_impl.go @@ -16,6 +16,7 @@ import ( oauthserver "github.com/stacklok/toolhive/pkg/authserver/server" "github.com/stacklok/toolhive/pkg/authserver/server/handlers" + "github.com/stacklok/toolhive/pkg/authserver/server/tokenexchange" "github.com/stacklok/toolhive/pkg/authserver/storage" "github.com/stacklok/toolhive/pkg/authserver/upstream" ) @@ -177,28 +178,17 @@ func newServer(ctx context.Context, cfg Config, stor storage.Storage, opts ...se // Wrap storage with the CIMD decorator before constructing the fosite provider // so that GetClient calls for HTTPS client_id values are intercepted at the // fosite level (not just the handler level). - if cfg.CIMDEnabled { - if len(cfg.BaselineClientScopes) > 0 { - slog.Warn("CIMD is enabled with baseline_client_scopes configured; "+ - "any third-party client resolved via CIMD will also receive these scopes — "+ - "ensure they are scopes you would grant by default to any unknown client", - "baseline_client_scopes", cfg.BaselineClientScopes) - } - stor, err = storage.NewCIMDStorageDecorator(stor, storage.CIMDDecoratorConfig{ - Enabled: true, - CacheMaxSize: cfg.CIMDCacheMaxSize, - FallbackTTL: cfg.CIMDCacheFallbackTTL, - ScopesSupported: cfg.ScopesSupported, - BaselineClientScopes: cfg.BaselineClientScopes, - }) - if err != nil { - return nil, fmt.Errorf("failed to initialize CIMD storage decorator: %w", err) - } + stor, err = decorateStorageForCIMD(cfg, stor) + if err != nil { + return nil, err } // Create fosite provider with the (possibly decorated) storage. slog.Debug("creating fosite OAuth2 provider") - fositeProvider := createProvider(authServerConfig, stor) + fositeProvider, err := buildProvider(cfg, authServerConfig, stor) + if err != nil { + return nil, fmt.Errorf("failed to create fosite OAuth2 provider: %w", err) + } // Give the handler a refresher so the authorization chain can transparently // refresh an expired upstream leg during login instead of skipping it and @@ -228,6 +218,44 @@ func newServer(ctx context.Context, cfg Config, stor storage.Storage, opts ...se }, nil } +// decorateStorageForCIMD wraps stor with the CIMD decorator when CIMD is enabled, +// so GetClient calls for HTTPS client_id values are intercepted at the fosite +// level (not just the handler level). Returns stor unchanged when CIMD is disabled. +func decorateStorageForCIMD(cfg Config, stor storage.Storage) (storage.Storage, error) { + if !cfg.CIMDEnabled { + return stor, nil + } + if len(cfg.BaselineClientScopes) > 0 { + slog.Warn("CIMD is enabled with baseline_client_scopes configured; "+ + "any third-party client resolved via CIMD will also receive these scopes — "+ + "ensure they are scopes you would grant by default to any unknown client", + "baseline_client_scopes", cfg.BaselineClientScopes) + } + decorated, err := storage.NewCIMDStorageDecorator(stor, storage.CIMDDecoratorConfig{ + Enabled: true, + CacheMaxSize: cfg.CIMDCacheMaxSize, + FallbackTTL: cfg.CIMDCacheFallbackTTL, + ScopesSupported: cfg.ScopesSupported, + BaselineClientScopes: cfg.BaselineClientScopes, + }) + if err != nil { + return nil, fmt.Errorf("failed to initialize CIMD storage decorator: %w", err) + } + return decorated, nil +} + +// buildProvider assembles the fosite OAuth2 provider, registering the RFC 8693 +// token-exchange handler as an extension grant alongside the standard grants. +func buildProvider( + cfg Config, authServerConfig *oauthserver.AuthorizationServerConfig, stor storage.Storage, +) (fosite.OAuth2Provider, error) { + tokenExchangeFactory, err := tokenexchange.Factory(cfg.DelegationTokenLifespan, cfg.TrustedIssuers) + if err != nil { + return nil, fmt.Errorf("failed to create token exchange factory: %w", err) + } + return createProvider(authServerConfig, stor, tokenExchangeFactory) +} + // buildHandlerOptions assembles the handlers.Option list for NewHandler: the // refresher is always wired, and the filter is added only when the caller's // Config sets one so a nil Config.UpstreamFilter preserves the pre-filter @@ -294,9 +322,10 @@ func (s *server) Close() error { // createProvider creates a fosite OAuth2Provider configured for the authorization code flow. // -// Fosite is an OAuth 2.0 framework that implements the protocol details. The compose package -// provides a builder pattern to wire together configuration, storage, token strategies, -// and grant type handlers into a single OAuth2Provider that can handle all OAuth endpoints. +// Fosite is an OAuth 2.0 framework that implements the protocol details. We use +// server.NewAuthorizationServer which accepts server.Factory functions to register +// grant type handlers. The standard compose factories are wrapped via wrapComposeFactory +// and any extra factories (e.g., token exchange) are appended. // // The provider is configured with: // - JWT strategy for access tokens (asymmetric signing, distributed validation via JWKS) @@ -304,7 +333,12 @@ func (s *server) Close() error { // - Authorization code grant (RFC 6749 Section 4.1) // - Refresh token grant (RFC 6749 Section 6) // - PKCE (RFC 7636) for public client security -func createProvider(authServerConfig *oauthserver.AuthorizationServerConfig, stor storage.Storage) fosite.OAuth2Provider { +// - Any extra factories passed in (e.g., RFC 8693 token exchange) +func createProvider( + authServerConfig *oauthserver.AuthorizationServerConfig, + stor storage.Storage, + extraFactories ...oauthserver.Factory, +) (fosite.OAuth2Provider, error) { slog.Debug("configuring fosite OAuth2 provider", "key_id", authServerConfig.SigningKey.KeyID, "algorithm", authServerConfig.SigningKey.Algorithm, @@ -334,18 +368,21 @@ func createProvider(authServerConfig *oauthserver.AuthorizationServerConfig, sto authServerConfig.Config, ) - // compose.Compose wires together all the pieces into an OAuth2Provider: - // - Config: token lifespans, issuer URL, HMAC secret - // - Storage: where to persist authorization codes, tokens, and client data - // - Strategy: how to generate and validate tokens - // - Factories: which OAuth grant types to enable (each adds handlers for specific flows) - return compose.Compose( - authServerConfig.Config, + commonStrategy := &compose.CommonStrategy{CoreStrategy: jwtStrategy} + + // Wrap fosite's compose factories to match server.Factory signature. + factories := []oauthserver.Factory{ + wrapComposeFactory(compose.OAuth2AuthorizeExplicitFactory), // Authorization code grant + wrapComposeFactory(compose.OAuth2RefreshTokenGrantFactory), // Refresh token grant + wrapComposeFactory(compose.OAuth2PKCEFactory), // PKCE for public clients + } + factories = append(factories, extraFactories...) + + return oauthserver.NewAuthorizationServer( + authServerConfig, stor, - &compose.CommonStrategy{CoreStrategy: jwtStrategy}, - compose.OAuth2AuthorizeExplicitFactory, // Authorization code grant - compose.OAuth2RefreshTokenGrantFactory, // Refresh token grant - compose.OAuth2PKCEFactory, // PKCE for public clients + commonStrategy, + factories..., ) } @@ -376,3 +413,13 @@ func runLegacyMigration(ctx context.Context, stor storage.Storage, upstreams []U } return nil } + +// wrapComposeFactory adapts a compose.Factory to a server.Factory. +// Compose factories take (fosite.Configurator, interface{}, interface{}) while +// server factories take (*AuthorizationServerConfig, fosite.Storage, any). +// The embedded *fosite.Config satisfies fosite.Configurator. +func wrapComposeFactory(cf compose.Factory) oauthserver.Factory { + return func(config *oauthserver.AuthorizationServerConfig, storage fosite.Storage, strategy any) (any, error) { + return cf(config.Config, storage, strategy), nil + } +} diff --git a/pkg/authserver/storage/memory.go b/pkg/authserver/storage/memory.go index 639ad3dc23..0e79a584c5 100644 --- a/pkg/authserver/storage/memory.go +++ b/pkg/authserver/storage/memory.go @@ -1281,6 +1281,7 @@ func (s *MemoryStorage) GetDCRCredentials(_ context.Context, key DCRKey) (*DCRCr if !ok { slog.Debug("dcr credentials not found", "issuer", key.Issuer, + "upstream_id", key.UpstreamID, "redirect_uri", key.RedirectURI, ) return nil, fmt.Errorf("%w: %w", ErrNotFound, fosite.ErrNotFound.WithHint("DCR credentials not found")) diff --git a/pkg/authserver/storage/memory_test.go b/pkg/authserver/storage/memory_test.go index 968238d342..3497dc3aa4 100644 --- a/pkg/authserver/storage/memory_test.go +++ b/pkg/authserver/storage/memory_test.go @@ -1766,6 +1766,7 @@ func TestMemoryStorage_DCRCredentials_RoundTrip(t *testing.T) { withStorage(t, func(ctx context.Context, s *MemoryStorage) { key := DCRKey{ Issuer: "https://thv.example.com", + UpstreamID: "https://idp.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: ScopesHash([]string{"openid", "profile"}), } @@ -1796,9 +1797,10 @@ func TestMemoryStorage_DCRCredentials_RoundTrip(t *testing.T) { func TestMemoryStorage_DCRCredentials_DistinctKeysDoNotCollide(t *testing.T) { withStorage(t, func(ctx context.Context, s *MemoryStorage) { - mkKey := func(issuer, redirect string, scopes []string) DCRKey { + mkKey := func(issuer, upstreamID, redirect string, scopes []string) DCRKey { return DCRKey{ Issuer: issuer, + UpstreamID: upstreamID, RedirectURI: redirect, ScopesHash: ScopesHash(scopes), } @@ -1812,10 +1814,14 @@ func TestMemoryStorage_DCRCredentials_DistinctKeysDoNotCollide(t *testing.T) { } } entries := []*DCRCredentials{ - mkCreds(mkKey("https://idp-a.example.com", "https://x/cb", []string{"openid"}), "a"), - mkCreds(mkKey("https://idp-b.example.com", "https://x/cb", []string{"openid"}), "b"), - mkCreds(mkKey("https://idp-a.example.com", "https://y/cb", []string{"openid"}), "c"), - mkCreds(mkKey("https://idp-a.example.com", "https://x/cb", []string{"openid", "email"}), "d"), + mkCreds(mkKey("https://idp-a.example.com", "https://up-a", "https://x/cb", []string{"openid"}), "a"), + mkCreds(mkKey("https://idp-b.example.com", "https://up-a", "https://x/cb", []string{"openid"}), "b"), + mkCreds(mkKey("https://idp-a.example.com", "https://up-a", "https://y/cb", []string{"openid"}), "c"), + mkCreds(mkKey("https://idp-a.example.com", "https://up-a", "https://x/cb", []string{"openid", "email"}), "d"), + // Differs from entry "a" only by UpstreamID — two OAuth2 upstreams + // inside one embedded authserver sharing issuer, redirect, and + // scopes. Before UpstreamID joined the key these collided (#5823). + mkCreds(mkKey("https://idp-a.example.com", "https://up-b", "https://x/cb", []string{"openid"}), "e"), } for _, e := range entries { require.NoError(t, s.StoreDCRCredentials(ctx, e)) @@ -1834,6 +1840,7 @@ func TestMemoryStorage_DCRCredentials_OverwriteSemantics(t *testing.T) { withStorage(t, func(ctx context.Context, s *MemoryStorage) { key := DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), } @@ -1864,7 +1871,7 @@ func TestMemoryStorage_DCRCredentials_NotFound(t *testing.T) { // TestMemoryStorage_DCRCredentials_StoreInvalidInputRejected pins the // fail-loud-on-invalid-input contract: nil creds, an unpopulated Key -// (empty Issuer, RedirectURI, or ScopesHash), and missing RFC 7591 +// (empty Issuer, UpstreamID, RedirectURI, or ScopesHash), and missing RFC 7591 // mandatory response fields (ClientID, AuthorizationEndpoint, // TokenEndpoint) must be rejected with fosite.ErrInvalidRequest rather // than producing a working-looking write that fails downstream at the @@ -1879,6 +1886,7 @@ func TestMemoryStorage_DCRCredentials_StoreInvalidInputRejected(t *testing.T) { return &DCRCredentials{ Key: DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), }, @@ -1903,6 +1911,13 @@ func TestMemoryStorage_DCRCredentials_StoreInvalidInputRejected(t *testing.T) { return c }, }, + { + name: "empty upstream_id", + mutator: func(c *DCRCredentials) *DCRCredentials { + c.Key.UpstreamID = "" + return c + }, + }, { name: "empty redirect_uri", mutator: func(c *DCRCredentials) *DCRCredentials { @@ -1961,6 +1976,7 @@ func TestMemoryStorage_DCRCredentials_GetReturnsDefensiveCopy(t *testing.T) { withStorage(t, func(ctx context.Context, s *MemoryStorage) { key := DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), } @@ -1988,6 +2004,7 @@ func TestMemoryStorage_DCRCredentials_StoreCopyIsolatesCaller(t *testing.T) { withStorage(t, func(ctx context.Context, s *MemoryStorage) { key := DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), } @@ -2015,6 +2032,7 @@ func TestMemoryStorage_DCRCredentials_ExcludedFromCleanupExpired(t *testing.T) { withStorage(t, func(ctx context.Context, s *MemoryStorage) { key := DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), } @@ -2057,6 +2075,7 @@ func TestMemoryStorage_DCRCredentials_ConcurrentAccess(t *testing.T) { overlappingKey := func(i int) DCRKey { return DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: fmt.Sprintf("overlap-%d", i%4), } @@ -2064,6 +2083,7 @@ func TestMemoryStorage_DCRCredentials_ConcurrentAccess(t *testing.T) { disjointKey := func(worker, i int) DCRKey { return DCRKey{ Issuer: fmt.Sprintf("https://idp-%d.example.com", worker), + UpstreamID: "https://upstream.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: fmt.Sprintf("disjoint-%d", i), } diff --git a/pkg/authserver/storage/redis.go b/pkg/authserver/storage/redis.go index ea4d3ec421..73ca8669ab 100644 --- a/pkg/authserver/storage/redis.go +++ b/pkg/authserver/storage/redis.go @@ -1232,6 +1232,7 @@ type storedDCRCredentials struct { // Embed the canonical key so a row recovered without its lookup key // (e.g. via SCAN during diagnostics) still self-identifies. KeyIssuer string `json:"key_issuer"` + KeyUpstreamID string `json:"key_upstream_id"` KeyRedirectURI string `json:"key_redirect_uri"` KeyScopesHash string `json:"key_scopes_hash"` @@ -1268,6 +1269,7 @@ func (s *storedDCRCredentials) toDCRCredentials() *DCRCredentials { return &DCRCredentials{ Key: DCRKey{ Issuer: s.KeyIssuer, + UpstreamID: s.KeyUpstreamID, RedirectURI: s.KeyRedirectURI, ScopesHash: s.KeyScopesHash, }, @@ -1314,6 +1316,7 @@ func (s *RedisStorage) StoreDCRCredentials(ctx context.Context, creds *DCRCreden stored := storedDCRCredentials{ KeyIssuer: creds.Key.Issuer, + KeyUpstreamID: creds.Key.UpstreamID, KeyRedirectURI: creds.Key.RedirectURI, KeyScopesHash: creds.Key.ScopesHash, ProviderName: creds.ProviderName, @@ -1361,7 +1364,7 @@ func (s *RedisStorage) StoreDCRCredentials(ctx context.Context, creds *DCRCreden // Returns ErrNotFound (wrapped) when no entry exists. The returned value is a // fresh struct decoded from JSON, which acts as a defensive copy. // -// An unpopulated key (empty Issuer, RedirectURI, or ScopesHash) cannot match +// An unpopulated key (empty Issuer, UpstreamID, RedirectURI, or ScopesHash) cannot match // any stored row because StoreDCRCredentials rejects such keys, so a Get // against one is a normal miss — ErrNotFound — matching // MemoryStorage.GetDCRCredentials and the DCRCredentialStore interface contract. diff --git a/pkg/authserver/storage/redis_integration_test.go b/pkg/authserver/storage/redis_integration_test.go index 5cd08f5e19..ab43f00a58 100644 --- a/pkg/authserver/storage/redis_integration_test.go +++ b/pkg/authserver/storage/redis_integration_test.go @@ -1462,8 +1462,8 @@ func TestIntegration_DCRCredentials_RoundTrip(t *testing.T) { func TestIntegration_DCRCredentials_DistinctKeysCoexist(t *testing.T) { withIntegrationStorage(t, func(ctx context.Context, s *RedisStorage) { - mkKey := func(issuer, redirect string, scopes []string) DCRKey { - return DCRKey{Issuer: issuer, RedirectURI: redirect, ScopesHash: ScopesHash(scopes)} + mkKey := func(issuer, upstreamID, redirect string, scopes []string) DCRKey { + return DCRKey{Issuer: issuer, UpstreamID: upstreamID, RedirectURI: redirect, ScopesHash: ScopesHash(scopes)} } mk := func(key DCRKey, clientID string) *DCRCredentials { return &DCRCredentials{ @@ -1474,10 +1474,13 @@ func TestIntegration_DCRCredentials_DistinctKeysCoexist(t *testing.T) { } } entries := []*DCRCredentials{ - mk(mkKey("https://idp-a.example.com", "https://x/cb", []string{"openid"}), "a"), - mk(mkKey("https://idp-b.example.com", "https://x/cb", []string{"openid"}), "b"), - mk(mkKey("https://idp-a.example.com", "https://y/cb", []string{"openid"}), "c"), - mk(mkKey("https://idp-a.example.com", "https://x/cb", []string{"openid", "email"}), "d"), + mk(mkKey("https://idp-a.example.com", "https://up-a", "https://x/cb", []string{"openid"}), "a"), + mk(mkKey("https://idp-b.example.com", "https://up-a", "https://x/cb", []string{"openid"}), "b"), + mk(mkKey("https://idp-a.example.com", "https://up-a", "https://y/cb", []string{"openid"}), "c"), + mk(mkKey("https://idp-a.example.com", "https://up-a", "https://x/cb", []string{"openid", "email"}), "d"), + // Differs from entry "a" only by UpstreamID — the two-upstreams-in- + // one-authserver shape that collided before #5823. + mk(mkKey("https://idp-a.example.com", "https://up-b", "https://x/cb", []string{"openid"}), "e"), } for _, e := range entries { require.NoError(t, s.StoreDCRCredentials(ctx, e)) diff --git a/pkg/authserver/storage/redis_keys.go b/pkg/authserver/storage/redis_keys.go index 45f8978d7a..42600ed336 100644 --- a/pkg/authserver/storage/redis_keys.go +++ b/pkg/authserver/storage/redis_keys.go @@ -95,20 +95,31 @@ func redisProviderKey(prefix, providerID, providerSubject string) string { } // redisDCRKey generates a Redis key for a DCR credential entry, identifying the -// (Issuer, RedirectURI, ScopesHash) tuple that DCRKey canonicalises. +// (Issuer, UpstreamID, RedirectURI, ScopesHash) tuple that DCRKey canonicalises. // -// Format: "{prefix}dcr:::::" +// Format: "{prefix}dcr:::::::" // -// The first two segments are length-prefixed to handle colons in RedirectURI -// (and, for symmetry, Issuer) without ambiguity, mirroring redisProviderKey. -// ScopesHash is expected to be a SHA-256 hex digest produced by -// storage.ScopesHash — only [0-9a-f] and never colon-bearing — so it is +// The first three segments are length-prefixed to handle colons in UpstreamID +// and RedirectURI (and, for symmetry, Issuer) without ambiguity, mirroring +// redisProviderKey. ScopesHash is expected to be a SHA-256 hex digest produced +// by storage.ScopesHash — only [0-9a-f] and never colon-bearing — so it is // appended without a length prefix. The format is robust for that domain; // validateDCRCredentialsForStore (called by every Store path) already // rejects an empty ScopesHash, and callers are required to compute the hash // via storage.ScopesHash. Length-prefix collision-safety is preserved on // the leading segments either way. // +// The UpstreamID segment (issue #5823) means two upstreams that share the +// consumer's Issuer and RedirectURI but differ only by their authorization +// server no longer map to the same Redis key. Adding it changes the key +// format, so entries written by an older binary (without the segment) will +// miss on lookup and be harmlessly re-registered under the new key. Orphaned +// old rows self-evict only when the upstream asserted an RFC 7591 §3.2.1 +// client_secret_expires_at (StoreDCRCredentials sets a matching Redis TTL for +// those); entries for non-expiring secrets — the common §3.2.1 "never" case — +// carry no TTL and persist indefinitely, so they require manual cleanup. +// There is no automated one-shot migration for this key-format change today. +// // The public-vs-confidential client distinction is intentionally NOT // encoded here — see DCRKey's doc for the rationale. Today's two consumers // register on disjoint RedirectURI address spaces (AS-origin vs RFC 8252 @@ -117,9 +128,10 @@ func redisProviderKey(prefix, providerID, providerSubject string) string { // add the distinguishing component back to the key format alongside an // explicit migration story for existing Redis-cached entries. func redisDCRKey(prefix string, key DCRKey) string { - return fmt.Sprintf("%s%s:%d:%s:%d:%s:%s", + return fmt.Sprintf("%s%s:%d:%s:%d:%s:%d:%s:%s", prefix, KeyTypeDCR, len(key.Issuer), key.Issuer, + len(key.UpstreamID), key.UpstreamID, len(key.RedirectURI), key.RedirectURI, key.ScopesHash) } diff --git a/pkg/authserver/storage/redis_test.go b/pkg/authserver/storage/redis_test.go index 23f7459e77..21922995ce 100644 --- a/pkg/authserver/storage/redis_test.go +++ b/pkg/authserver/storage/redis_test.go @@ -2066,12 +2066,14 @@ func TestRedisKeyGeneration(t *testing.T) { t.Parallel() result := redisDCRKey("test:auth:", DCRKey{ Issuer: "https://thv.example.com", + UpstreamID: "https://idp.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: "abc123", }) - // 23 = len("https://thv.example.com"), 38 = len("https://thv.example.com/oauth/callback") + // 23 = len("https://thv.example.com"), 23 = len("https://idp.example.com"), + // 38 = len("https://thv.example.com/oauth/callback") assert.Equal(t, - "test:auth:dcr:23:https://thv.example.com:38:https://thv.example.com/oauth/callback:abc123", + "test:auth:dcr:23:https://thv.example.com:23:https://idp.example.com:38:https://thv.example.com/oauth/callback:abc123", result) }) } @@ -2084,8 +2086,8 @@ func TestRedisKeyGeneration(t *testing.T) { func TestRedisDCRKey_Distinct(t *testing.T) { t.Parallel() - mk := func(issuer, redirect, scopes string) DCRKey { - return DCRKey{Issuer: issuer, RedirectURI: redirect, ScopesHash: scopes} + mk := func(issuer, upstreamID, redirect, scopes string) DCRKey { + return DCRKey{Issuer: issuer, UpstreamID: upstreamID, RedirectURI: redirect, ScopesHash: scopes} } tests := []struct { @@ -2094,33 +2096,56 @@ func TestRedisDCRKey_Distinct(t *testing.T) { }{ { name: "different issuer", - a: mk("https://idp-a.example.com", "https://x/cb", "h1"), - b: mk("https://idp-b.example.com", "https://x/cb", "h1"), + a: mk("https://idp-a.example.com", "https://up", "https://x/cb", "h1"), + b: mk("https://idp-b.example.com", "https://up", "https://x/cb", "h1"), + }, + { + // Two upstreams that share the consumer's issuer, redirect, and + // scopes and differ only by upstream — the #5823 shape. + name: "different upstream_id", + a: mk("https://idp.example.com", "https://up-a", "https://x/cb", "h1"), + b: mk("https://idp.example.com", "https://up-b", "https://x/cb", "h1"), }, { name: "different redirect_uri", - a: mk("https://idp.example.com", "https://x/cb", "h1"), - b: mk("https://idp.example.com", "https://y/cb", "h1"), + a: mk("https://idp.example.com", "https://up", "https://x/cb", "h1"), + b: mk("https://idp.example.com", "https://up", "https://y/cb", "h1"), }, { name: "different scopes hash", - a: mk("https://idp.example.com", "https://x/cb", "h1"), - b: mk("https://idp.example.com", "https://x/cb", "h2"), + a: mk("https://idp.example.com", "https://up", "https://x/cb", "h1"), + b: mk("https://idp.example.com", "https://up", "https://x/cb", "h2"), }, { // Without length prefixing, ("ab", "cd") and ("a", "bcd") would // both yield ":ab:cd:" as the issuer/redirect segment after a // fmt.Sprintf collapse. The length prefix prevents that. name: "redirect_uri-issuer boundary collision (length-prefix property)", - a: mk("ab", "cd", "h1"), - b: mk("a", "bcd", "h1"), + a: mk("ab", "up", "cd", "h1"), + b: mk("a", "up", "bcd", "h1"), + }, + { + // UpstreamID legitimately contains colons and can abut the issuer + // boundary; the length prefix keeps ("ab","cd") and ("a","bcd") + // distinct across the issuer/upstream seam too. + name: "issuer-upstream_id boundary collision (length-prefix property)", + a: mk("ab", "cd", "https://x/cb", "h1"), + b: mk("a", "bcd", "https://x/cb", "h1"), + }, + { + // UpstreamID also abuts RedirectURI; the length prefix must keep + // ("ab","cd") and ("a","bcd") distinct across the second new seam + // introduced by inserting UpstreamID between Issuer and RedirectURI. + name: "upstream_id-redirect_uri boundary collision (length-prefix property)", + a: mk("https://idp.example.com", "ab", "cd", "h1"), + b: mk("https://idp.example.com", "a", "bcd", "h1"), }, { // RedirectURI legitimately contains colons (e.g. ":443"). A plain // "%s:%s:%s" key would be ambiguous; the length prefix is not. name: "colons inside redirect_uri", - a: mk("https://idp.example.com", "https://x.example.com:443/cb", "h1"), - b: mk("https://idp.example.com", "https://x.example.com/cb:443", "h1"), + a: mk("https://idp.example.com", "https://up", "https://x.example.com:443/cb", "h1"), + b: mk("https://idp.example.com", "https://up", "https://x.example.com/cb:443", "h1"), }, } @@ -2142,6 +2167,7 @@ func TestRedisDCRKey_Deterministic(t *testing.T) { key := DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: ScopesHash([]string{"openid", "profile", "email"}), } @@ -2389,6 +2415,7 @@ func TestRedisStorage_GetLatestUpstreamTokensForUser(t *testing.T) { func dcrFixtureKey() DCRKey { return DCRKey{ Issuer: "https://thv.example.com", + UpstreamID: "https://idp.example.com", RedirectURI: "https://thv.example.com/oauth/callback", ScopesHash: ScopesHash([]string{"openid", "profile"}), } @@ -2488,8 +2515,8 @@ func TestRedisStorage_DCRCredentials_NotFound(t *testing.T) { func TestRedisStorage_DCRCredentials_DistinctKeysCoexist(t *testing.T) { withRedisStorage(t, func(ctx context.Context, s *RedisStorage, _ *miniredis.Miniredis) { - mkKey := func(issuer, redirect string, scopes []string) DCRKey { - return DCRKey{Issuer: issuer, RedirectURI: redirect, ScopesHash: ScopesHash(scopes)} + mkKey := func(issuer, upstreamID, redirect string, scopes []string) DCRKey { + return DCRKey{Issuer: issuer, UpstreamID: upstreamID, RedirectURI: redirect, ScopesHash: ScopesHash(scopes)} } mk := func(key DCRKey, clientID string) *DCRCredentials { return &DCRCredentials{ @@ -2500,10 +2527,13 @@ func TestRedisStorage_DCRCredentials_DistinctKeysCoexist(t *testing.T) { } } entries := []*DCRCredentials{ - mk(mkKey("https://idp-a.example.com", "https://x/cb", []string{"openid"}), "a"), - mk(mkKey("https://idp-b.example.com", "https://x/cb", []string{"openid"}), "b"), - mk(mkKey("https://idp-a.example.com", "https://y/cb", []string{"openid"}), "c"), - mk(mkKey("https://idp-a.example.com", "https://x/cb", []string{"openid", "email"}), "d"), + mk(mkKey("https://idp-a.example.com", "https://up-a", "https://x/cb", []string{"openid"}), "a"), + mk(mkKey("https://idp-b.example.com", "https://up-a", "https://x/cb", []string{"openid"}), "b"), + mk(mkKey("https://idp-a.example.com", "https://up-a", "https://y/cb", []string{"openid"}), "c"), + mk(mkKey("https://idp-a.example.com", "https://up-a", "https://x/cb", []string{"openid", "email"}), "d"), + // Differs from entry "a" only by UpstreamID — the two-upstreams-in- + // one-authserver shape that collided before #5823. + mk(mkKey("https://idp-a.example.com", "https://up-b", "https://x/cb", []string{"openid"}), "e"), } for _, e := range entries { require.NoError(t, s.StoreDCRCredentials(ctx, e)) @@ -2532,6 +2562,7 @@ func TestRedisStorage_DCRCredentials_StoreInvalidInputRejected(t *testing.T) { return &DCRCredentials{ Key: DCRKey{ Issuer: "https://idp.example.com", + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), }, @@ -2556,6 +2587,13 @@ func TestRedisStorage_DCRCredentials_StoreInvalidInputRejected(t *testing.T) { return c }, }, + { + name: "empty upstream_id", + mutator: func(c *DCRCredentials) *DCRCredentials { + c.Key.UpstreamID = "" + return c + }, + }, { name: "empty redirect_uri", mutator: func(c *DCRCredentials) *DCRCredentials { @@ -2775,6 +2813,7 @@ func runDCRConcurrentAccess( case dcrConcurrentDisjointKeys: return DCRKey{ Issuer: fmt.Sprintf("https://idp-%d.example.com", gid), + UpstreamID: "https://upstream.example.com", RedirectURI: "https://x/cb", ScopesHash: ScopesHash([]string{"openid"}), } diff --git a/pkg/authserver/storage/types.go b/pkg/authserver/storage/types.go index 91490cd0ab..7b236d6cd6 100644 --- a/pkg/authserver/storage/types.go +++ b/pkg/authserver/storage/types.go @@ -130,9 +130,9 @@ type DCRKey struct { // issuer, because the CLI has no separate local issuer of its own. // The cache is keyed by this value because two different consumers // registering against the same upstream are distinct OAuth clients and - // must not share credentials. The (Issuer, RedirectURI, ScopesHash) - // tuple keeps the two consumer profiles' entries apart via the - // RedirectURI component (the embedded authserver registers an + // must not share credentials. The (Issuer, UpstreamID, RedirectURI, + // ScopesHash) tuple keeps the two consumer profiles' entries apart via + // the RedirectURI component (the embedded authserver registers an // AS-origin callback while the CLI registers a loopback callback per // RFC 8252 §7.3 — the two address spaces are disjoint), so a collision // between profiles is impossible by construction even when the @@ -146,6 +146,31 @@ type DCRKey struct { // explicit migration story. Issuer string + // UpstreamID identifies the upstream authorization server this + // registration is bound to, disambiguating upstreams that share the + // consumer's Issuer, RedirectURI, and scope set. Within a single + // embedded authserver every OAuth2 upstream shares the authserver's own + // Issuer and the one defaulted {issuer}/oauth/callback RedirectURI, so + // before this component existed two upstreams configured with equal + // scopes collided on one cache entry — the second read back the first's + // dynamically-registered client_id / client_secret and never registered + // with its own authorization server (issue #5823). + // + // The value is the upstream's registration identity: the upstream issuer + // recovered from the consumer's discovery URL, or the registration + // endpoint URL when the consumer configured one directly. It is derived + // by the dcr resolver from the Request; do not hand-build it at call + // sites. + // + // On the discovery-URL path this is the derived issuer, so it identifies + // the authorization server rather than the exact discovery URL: two + // configs that resolve to the same issuer intentionally share one + // registration. It does not fully disambiguate the nonstandard case of a + // single issuer exposing distinct registration endpoints under custom + // (non-well-known) discovery paths — see resolveUpstreamKeyIdentity in + // pkg/auth/dcr for the full rationale. + UpstreamID string + // RedirectURI is the redirect URI registered with the upstream // authorization server. Embedded-authserver callers register an // AS-origin callback; CLI callers register an RFC 8252 loopback @@ -198,13 +223,16 @@ func ScopesHash(scopes []string) string { // from drifting across implementations: a record that fails loud against one // backend cannot silently persist against another. // -// Rejected inputs: nil creds, an unpopulated Key (empty Issuer, RedirectURI, -// or ScopesHash), and missing RFC 7591 mandatory response fields (ClientID, -// AuthorizationEndpoint, TokenEndpoint). An empty ScopesHash is rejected -// because the canonical digest of any scope set — including the empty-scope -// set via ScopesHash(nil) — is non-empty, so an empty string can only be a -// caller bug. ClientSecret is left permissive because RFC 7591 §2 public -// clients (auth method "none") legitimately register without a secret. +// Rejected inputs: nil creds, an unpopulated Key (empty Issuer, UpstreamID, +// RedirectURI, or ScopesHash), and missing RFC 7591 mandatory response fields +// (ClientID, AuthorizationEndpoint, TokenEndpoint). An empty ScopesHash is +// rejected because the canonical digest of any scope set — including the +// empty-scope set via ScopesHash(nil) — is non-empty, so an empty string can +// only be a caller bug. An empty UpstreamID is rejected because the resolver +// always derives it from the Request (issue #5823); an empty value here would +// re-open the cross-upstream collision it was added to close. ClientSecret is +// left permissive because RFC 7591 §2 public clients (auth method "none") +// legitimately register without a secret. func validateDCRCredentialsForStore(creds *DCRCredentials) error { if creds == nil { return fosite.ErrInvalidRequest.WithHint("dcr credentials cannot be nil") @@ -212,6 +240,9 @@ func validateDCRCredentialsForStore(creds *DCRCredentials) error { if creds.Key.Issuer == "" { return fosite.ErrInvalidRequest.WithHint("dcr credentials key issuer cannot be empty") } + if creds.Key.UpstreamID == "" { + return fosite.ErrInvalidRequest.WithHint("dcr credentials key upstream_id cannot be empty") + } if creds.Key.RedirectURI == "" { return fosite.ErrInvalidRequest.WithHint("dcr credentials key redirect_uri cannot be empty") } @@ -263,7 +294,7 @@ func validateDCRCredentialsForStore(creds *DCRCredentials) error { // TestResolutionCredentialsRoundTrip in // pkg/auth/dcr/store_test.go. type DCRCredentials struct { - // Key is the canonical cache key: (Issuer, RedirectURI, ScopesHash). + // Key is the canonical cache key: (Issuer, UpstreamID, RedirectURI, ScopesHash). Key DCRKey // ProviderName is the upstream's UpstreamRunConfig.Name. Debug / audit @@ -340,7 +371,7 @@ type DCRCredentials struct { // DCRKey is embedded as DCRCredentials.Key so the persisted blob is // self-describing: a Redis SCAN, an admin-tool dump, or a cross-replica // reconciliation path can identify a record's logical cache slot -// (Issuer, RedirectURI, ScopesHash) from the value alone, without +// (Issuer, UpstreamID, RedirectURI, ScopesHash) from the value alone, without // reconstructing it from a separately-passed key. This is a deliberate // asymmetry with the rest of the package — callers must populate creds.Key // before Store, and implementations validate it (see MemoryStorage docs diff --git a/pkg/authserver/upstream/oauth2.go b/pkg/authserver/upstream/oauth2.go index 11153f9dd9..bb44b31643 100644 --- a/pkg/authserver/upstream/oauth2.go +++ b/pkg/authserver/upstream/oauth2.go @@ -28,7 +28,6 @@ import ( "maps" "net/http" "net/url" - "os" "strings" "golang.org/x/oauth2" @@ -148,6 +147,18 @@ type OAuth2Config struct { // TokenEndpoint is the URL for the OAuth token endpoint. TokenEndpoint string `json:"token_endpoint" yaml:"token_endpoint"` + // TokenEndpointAuthMethod is the RFC 7591 client authentication method used + // at the token endpoint; see authStyleFromMethod for the mapping to + // oauth2.AuthStyle and the rationale. When empty, the historical default + // (POST body) is used. + // + // Only the DCR path populates this, via applyResolutionToOAuth2Config. + // OAuth2UpstreamRunConfig has no corresponding field, so a statically- + // configured upstream cannot set it and always gets the default — an + // intentional limitation scoped to issue #5865 (DCR-negotiated clients). + //nolint:lll // field tags require full JSON+YAML names + TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty" yaml:"token_endpoint_auth_method,omitempty"` + // UserInfo contains configuration for fetching user information (optional). // When nil, the provider does not support UserInfo fetching. UserInfo *UserInfoConfig `json:"userinfo,omitempty" yaml:"userinfo,omitempty"` @@ -307,7 +318,14 @@ func newBaseOAuth2Provider(config *OAuth2Config, hostForClient string) (*BaseOAu return nil, fmt.Errorf("failed to create HTTP client: %w", err) } - // Create the oauth2.Config for use with golang.org/x/oauth2 library + // AuthStyle is derived from the negotiated token_endpoint_auth_method + // rather than hardcoded — see authStyleFromMethod. + authStyle, err := authStyleFromMethod(config.TokenEndpointAuthMethod) + if err != nil { + return nil, err + } + + // Create the oauth2.Config for use with golang.org/x/oauth2 library. oauth2Cfg := &oauth2.Config{ ClientID: config.ClientID, ClientSecret: config.ClientSecret, @@ -316,7 +334,7 @@ func newBaseOAuth2Provider(config *OAuth2Config, hostForClient string) (*BaseOAu Endpoint: oauth2.Endpoint{ AuthURL: config.AuthorizationEndpoint, TokenURL: config.TokenEndpoint, - AuthStyle: oauth2.AuthStyleInParams, // Send client credentials in POST body + AuthStyle: authStyle, }, } @@ -327,6 +345,47 @@ func newBaseOAuth2Provider(config *OAuth2Config, hostForClient string) (*BaseOAu }, nil } +// authStyleFromMethod maps an RFC 7591 token_endpoint_auth_method to the +// oauth2.AuthStyle the golang.org/x/oauth2 library uses when presenting client +// credentials at the token endpoint: +// +// - client_secret_basic → AuthStyleInHeader (HTTP Basic Authorization header) +// - client_secret_post → AuthStyleInParams (credentials in the POST body) +// - none / "" (unset) → AuthStyleInParams (historical default, no secret) +// +// This is the single home for the auth-method rationale. For DCR-registered +// clients the method is copied from the negotiated +// dcr.Resolution.TokenEndpointAuthMethod so the exchange respects whatever the +// upstream advertised — some authorization servers (e.g. Ory Hydra) default +// confidential clients to client_secret_basic and reject client_secret_post +// outright (issue #5865). The same AuthStyle governs the refresh path, which +// shares this oauth2.Config. +// +// AuthStyleAutoDetect is deliberately not used for the default: its Basic-auth +// probe can consume a single-use authorization code against strict servers that +// reject Basic auth, the same failure pkg/auth/oauth/flow.go sidesteps by +// pinning AuthStyleInParams for public clients. +// +// An unrecognised, non-empty method (e.g. private_key_jwt or tls_client_auth, +// which the DCR resolver can negotiate but this client cannot fulfil) returns +// an error rather than silently degrading to POST-body credentials, per the +// "fail loudly on unrecognised enum values" convention in +// .claude/rules/go-style.md. Silently falling back would send an unusable +// credential and reproduce the opaque invalid_client symptom class issue #5865 +// was written to eliminate. +func authStyleFromMethod(method string) (oauth2.AuthStyle, error) { + switch method { + case oauthproto.TokenEndpointAuthMethodClientSecretBasic: + return oauth2.AuthStyleInHeader, nil + case oauthproto.TokenEndpointAuthMethodClientSecretPost, + oauthproto.TokenEndpointAuthMethodNone, + "": + return oauth2.AuthStyleInParams, nil + default: + return oauth2.AuthStyleAutoDetect, fmt.Errorf("unsupported token_endpoint_auth_method %q", method) + } +} + // NewOAuth2Provider creates a new pure OAuth 2.0 provider. // Use this for providers that don't support OIDC discovery. // The config must include AuthorizationEndpoint and TokenEndpoint. @@ -830,12 +889,20 @@ func formatOAuth2Error(err error, prefix string) error { // allowPrivateIPs widens only the private-IP gate: it permits connections to // RFC-1918/link-local addresses (e.g. in-cluster providers) without enabling // the HTTP scheme for non-localhost hosts. +// +// The host-scoped guard policy lives in networking.NewHostScopedClientBuilder +// so this provider path and the DCR resolver share one implementation. +// +// Unlike the DCR resolver's guarded client, this one deliberately leaves +// keep-alive enabled: the returned client is stored on BaseOAuth2Provider and +// reused for many token-refresh/userinfo calls over the provider's lifetime +// against a single operator-configured host, so disabling keep-alive here +// would pay a fresh TCP+TLS handshake on every call. This trades a narrower +// window — a DNS change for that fixed host between connection reuses is not +// re-checked mid-lifetime — for avoiding that cost on a hot path; the DCR +// resolver's per-request-host, low-frequency calls don't have the same +// trade-off, hence the difference. If this provider's threat model changes +// (e.g. it starts dialing caller-varying hosts), revisit this decision. func newHTTPClientForHost(host string, allowPrivateIPs, insecureAllowHTTP bool) (*http.Client, error) { - allowInsecure := networking.IsLocalhost(host) || - insecureAllowHTTP || - strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true") - return networking.NewHttpClientBuilder(). - WithInsecureAllowHTTP(allowInsecure). - WithPrivateIPs(allowInsecure || allowPrivateIPs). - Build() + return networking.NewHostScopedClientBuilder(host, allowPrivateIPs, insecureAllowHTTP).Build() } diff --git a/pkg/authserver/upstream/oauth2_test.go b/pkg/authserver/upstream/oauth2_test.go index 187020cd5d..d00b63afa2 100644 --- a/pkg/authserver/upstream/oauth2_test.go +++ b/pkg/authserver/upstream/oauth2_test.go @@ -31,6 +31,9 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/oauth2" + + "github.com/stacklok/toolhive/pkg/oauthproto" ) // testTokenResponse is a test helper to produce token responses. @@ -219,6 +222,229 @@ func TestNewOAuth2Provider(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "redirect_uri is required") }) + + t.Run("unsupported token endpoint auth method returns error", func(t *testing.T) { + t.Parallel() + + config := &OAuth2Config{ + CommonOAuthConfig: CommonOAuthConfig{ + ClientID: "test-client", + ClientSecret: "test-secret", + RedirectURI: "http://localhost:8080/callback", + }, + AuthorizationEndpoint: mock.URL + "/authorize", + TokenEndpoint: mock.URL + "/token", + TokenEndpointAuthMethod: "private_key_jwt", + } + + _, err := NewOAuth2Provider(config) + require.Error(t, err) + assert.Contains(t, err.Error(), "private_key_jwt", + "construction must fail loudly on a method this client cannot fulfil") + }) +} + +// TestAuthStyleFromMethod pins the mapping from RFC 7591 +// token_endpoint_auth_method values to oauth2.AuthStyle. Recognised methods +// (basic/post/none) and the unset case resolve to a definite style; an +// unrecognised, non-empty method must be rejected (fail loudly) rather than +// silently degrading to the POST-body default. +func TestAuthStyleFromMethod(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + method string + want oauth2.AuthStyle + wantErr bool + }{ + {"client_secret_basic maps to header", oauthproto.TokenEndpointAuthMethodClientSecretBasic, oauth2.AuthStyleInHeader, false}, + {"client_secret_post maps to params", oauthproto.TokenEndpointAuthMethodClientSecretPost, oauth2.AuthStyleInParams, false}, + {"none maps to params", oauthproto.TokenEndpointAuthMethodNone, oauth2.AuthStyleInParams, false}, + {"unset maps to params", "", oauth2.AuthStyleInParams, false}, + {"private_key_jwt is rejected", "private_key_jwt", oauth2.AuthStyleAutoDetect, true}, + {"unknown method is rejected", "totally_bogus", oauth2.AuthStyleAutoDetect, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := authStyleFromMethod(tt.method) + if tt.wantErr { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.method, + "error must name the offending method for operator diagnosis") + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +// TestNewOAuth2Provider_TokenEndpointAuthMethod verifies that the negotiated +// token_endpoint_auth_method controls how client credentials are presented at +// the token endpoint. This is the regression guard for issue #5865: a +// DCR-registered client whose upstream negotiated client_secret_basic must send +// its credentials in the HTTP Basic Authorization header, not the POST body, or +// strict authorization servers (e.g. Ory Hydra) reject the exchange with +// invalid_client. +func TestNewOAuth2Provider_TokenEndpointAuthMethod(t *testing.T) { + t.Parallel() + + const ( + clientID = "dcr-client" + clientSecret = "dcr-secret" + ) + + tests := []struct { + name string + authMethod string + wantAuthStyle oauth2.AuthStyle + wantBasicAuth bool // credentials expected in the Authorization header + }{ + { + name: "client_secret_basic sends credentials in Basic header", + authMethod: oauthproto.TokenEndpointAuthMethodClientSecretBasic, + wantAuthStyle: oauth2.AuthStyleInHeader, + wantBasicAuth: true, + }, + { + name: "client_secret_post sends credentials in body", + authMethod: oauthproto.TokenEndpointAuthMethodClientSecretPost, + wantAuthStyle: oauth2.AuthStyleInParams, + wantBasicAuth: false, + }, + { + name: "unset defaults to credentials in body", + authMethod: "", + wantAuthStyle: oauth2.AuthStyleInParams, + wantBasicAuth: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + mock := newMockOAuth2Server() + t.Cleanup(mock.Close) + + var ( + gotUser, gotPass string + gotBasicAuth bool + gotBodySecret string + ) + mock.tokenHandler = func(w http.ResponseWriter, r *http.Request) { + gotUser, gotPass, gotBasicAuth = r.BasicAuth() + if err := r.ParseForm(); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + gotBodySecret = r.PostForm.Get("client_secret") + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(testTokenResponse{ + AccessToken: "access-token", + TokenType: "Bearer", + }); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } + } + + config := &OAuth2Config{ + CommonOAuthConfig: CommonOAuthConfig{ + ClientID: clientID, + ClientSecret: clientSecret, + RedirectURI: "http://localhost:8080/callback", + }, + AuthorizationEndpoint: mock.URL + "/authorize", + TokenEndpoint: mock.URL + "/token", + TokenEndpointAuthMethod: tt.authMethod, + } + + provider, err := NewOAuth2Provider(config) + require.NoError(t, err) + assert.Equal(t, tt.wantAuthStyle, provider.oauth2Config.Endpoint.AuthStyle, + "oauth2 endpoint auth style must reflect the negotiated method") + + _, err = provider.exchangeCodeForTokens(context.Background(), "auth-code", "") + require.NoError(t, err) + + if tt.wantBasicAuth { + require.True(t, gotBasicAuth, "credentials must be sent in the Basic Authorization header") + assert.Equal(t, clientID, gotUser) + assert.Equal(t, clientSecret, gotPass) + assert.Empty(t, gotBodySecret, "client_secret must not also appear in the POST body") + } else { + assert.False(t, gotBasicAuth, "no Basic Authorization header expected") + assert.Equal(t, clientSecret, gotBodySecret, "client_secret must be sent in the POST body") + } + }) + } +} + +// TestBaseOAuth2Provider_RefreshTokens_TokenEndpointAuthMethod verifies that +// the negotiated auth method also governs the refresh path. RefreshTokens +// shares the same oauth2.Config.Endpoint.AuthStyle as code exchange, so a +// client_secret_basic client must present its credentials in the Basic +// Authorization header on refresh too — this pins the refresh half of the +// contract documented on authStyleFromMethod so a future change that rebuilds +// the config or bypasses AuthStyle in RefreshTokens does not ship unnoticed. +func TestBaseOAuth2Provider_RefreshTokens_TokenEndpointAuthMethod(t *testing.T) { + t.Parallel() + + const ( + clientID = "dcr-client" + clientSecret = "dcr-secret" + ) + + mock := newMockOAuth2Server() + t.Cleanup(mock.Close) + + var ( + gotUser, gotPass string + gotBasicAuth bool + gotBodySecret string + ) + mock.tokenHandler = func(w http.ResponseWriter, r *http.Request) { + gotUser, gotPass, gotBasicAuth = r.BasicAuth() + if err := r.ParseForm(); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + gotBodySecret = r.PostForm.Get("client_secret") + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(testTokenResponse{ + AccessToken: "refreshed-access-token", + TokenType: "Bearer", + }); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } + } + + config := &OAuth2Config{ + CommonOAuthConfig: CommonOAuthConfig{ + ClientID: clientID, + ClientSecret: clientSecret, + RedirectURI: "http://localhost:8080/callback", + }, + AuthorizationEndpoint: mock.URL + "/authorize", + TokenEndpoint: mock.URL + "/token", + TokenEndpointAuthMethod: oauthproto.TokenEndpointAuthMethodClientSecretBasic, + } + + provider, err := NewOAuth2Provider(config) + require.NoError(t, err) + + _, err = provider.RefreshTokens(context.Background(), "some-refresh-token", "") + require.NoError(t, err) + + require.True(t, gotBasicAuth, "refresh must send credentials in the Basic Authorization header") + assert.Equal(t, clientID, gotUser) + assert.Equal(t, clientSecret, gotPass) + assert.Empty(t, gotBodySecret, "client_secret must not also appear in the refresh POST body") } func TestBaseOAuth2Provider_Type(t *testing.T) { diff --git a/pkg/authz/annotation_cache.go b/pkg/authz/annotation_cache.go index 1637e96b29..bda00e21e1 100644 --- a/pkg/authz/annotation_cache.go +++ b/pkg/authz/annotation_cache.go @@ -6,8 +6,7 @@ package authz import ( "sync" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" ) diff --git a/pkg/authz/annotation_cache_test.go b/pkg/authz/annotation_cache_test.go index d4a4483d5c..a8c9930358 100644 --- a/pkg/authz/annotation_cache_test.go +++ b/pkg/authz/annotation_cache_test.go @@ -7,10 +7,10 @@ import ( "sync" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" ) diff --git a/pkg/authz/authorizers/cedar/core.go b/pkg/authz/authorizers/cedar/core.go index 011345b1d4..d6198e71fe 100644 --- a/pkg/authz/authorizers/cedar/core.go +++ b/pkg/authz/authorizers/cedar/core.go @@ -10,6 +10,7 @@ import ( "errors" "fmt" "log/slog" + "slices" "strings" "sync" "time" @@ -149,6 +150,16 @@ var ( // ClientIDContextKey is the key used to store client ID in the context. type ClientIDContextKey struct{} +// claimPrefix namespaces every JWT claim in the Cedar context/principal +// attributes. claimSetPrefix namespaces the synthetic multi-valued-claim Sets. +// They are deliberately disjoint: because every JWT claim is emitted under +// claimPrefix, a token can never produce a claimSetPrefix key, so the synthetic +// Sets cannot be shadowed or spoofed by claim content. +const ( + claimPrefix = "claim_" + claimSetPrefix = "claimset_" +) + // Authorizer authorizes MCP operations using Cedar policies. type Authorizer struct { // Cedar policy set @@ -160,7 +171,8 @@ type Authorizer struct { // Mutex for thread safety mu sync.RWMutex // primaryUpstreamProvider names the upstream IDP provider whose access token - // should be used as the source of JWT claims for Cedar evaluation. + // is the source of JWT claims for Cedar evaluation, aside from the narrow + // user-profile supplement described in resolveClaims. // When empty, claims from the token on the original client request are used, // which may be a ToolHive-issued token or any other bearer token. primaryUpstreamProvider string @@ -179,6 +191,20 @@ type Authorizer struct { // claimKeyLog rate-limits the diagnostic log of resolved JWT claim keys // so it emits at most once per 30 seconds instead of once per authorization check. claimKeyLog *syncutil.AtMost + // supplementLog rate-limits the warning that the upstream access token lacked + // profile claims. It is deliberately separate from claimKeyLog: a shared + // timer would let whichever fires first suppress the other for the rest of + // the window, hiding the claim-key dump on exactly the deployments that need it. + supplementLog *syncutil.AtMost + // missingIDTokenLog rate-limits the warning that the primary provider has no + // usable stored id_token, so profile claims cannot be supplemented at all. + // Separate from supplementLog because the two are mutually exclusive per + // request and describe opposite conditions. + missingIDTokenLog *syncutil.AtMost + // multiValuedClaims lists JWT claim names normalized to a canonical unpadded + // space-delimited string, plus a companion Cedar Set, before Cedar evaluation. + // See ConfigOptions.MultiValuedClaims. + multiValuedClaims []string } // ConfigOptions represents the Cedar-specific authorization configuration options. @@ -190,9 +216,17 @@ type ConfigOptions struct { EntitiesJSON string `json:"entities_json" yaml:"entities_json"` // PrimaryUpstreamProvider names the upstream IDP provider whose access - // token should be used as the source of JWT claims for Cedar evaluation. - // When empty, claims from the ToolHive-issued token are used (current behaviour). + // token is the primary source of JWT claims for Cedar evaluation. + // When empty, claims from the ToolHive-issued token are used. // Must match an entry in identity.UpstreamTokens (e.g. "default", "github"). + // + // The profile claims in profileClaimsFromIDToken fall back to the SAME + // provider's id_token when its access token omits them (many OIDC providers + // assert `email` only in the id_token), so `principal has claim_email` keeps + // meaning "this upstream asserted an email". Every other claim, including all + // group/role/scope claims, comes from the upstream access token or not at all, + // and the ToolHive-issued token the client presented is never a claim source on + // this path. See resolveClaims for the full contract. PrimaryUpstreamProvider string `json:"primary_upstream_provider,omitempty" yaml:"primary_upstream_provider,omitempty"` // GroupClaimName is the JWT claim key that contains group membership for the @@ -214,6 +248,43 @@ type ConfigOptions struct { // names (e.g. "Platform::Group") are not yet supported and are rejected at // construction. See issue #5072. GroupEntityType string `json:"group_entity_type,omitempty" yaml:"group_entity_type,omitempty"` + + // MultiValuedClaims lists JWT claim NAMES (bare, e.g. "scp" or "scope") whose + // value is exposed to Cedar in two normalized forms, regardless of whether the + // IdP emits the claim as a space-delimited string (Entra `scp`, Keycloak `scope`) + // or a JSON array (Okta `scp`): + // + // 1. "claim_" (both principal attribute and context) is a space-delimited + // string: an array is joined with single spaces; a string value is passed + // through VERBATIM — its spacing and order are preserved, not + // re-canonicalized. Use this with `like`/`==` string policies. + // 2. "claimset_" (bare key, both principal attribute and context) is a Cedar + // Set of the claim's elements. Use this with `.contains`/`.containsAll`/ + // `.containsAny` for exact-element matching that cannot be fooled by a + // substring (e.g. "Mail.Read" will not match "Mail.ReadWrite"). The "claimset_" + // prefix is reserved and never collides with a JWT claim, since all JWT + // claims are surfaced under the "claim_" prefix. + // + // Only for claims whose ELEMENTS are space-free (OAuth scopes, per RFC 6749 §3.3). + // Do NOT list group/role claims or any claim whose elements can contain spaces + // (e.g. group display names) — element boundaries would become ambiguous. + // + // Backward compatibility: opt-in. When empty (default), "claim_" is + // byte-identical to today and no "claimset_" key is added. Listing a claim + // leaves OTHER claims untouched. + // + // Migration hazard: listing a claim that an IdP emits as an ARRAY changes its + // "claim_" form from a Cedar Set to a String. A pre-existing policy that + // tested it with `like`/substring — which errored on the Set and so failed + // closed (deny) — will then match against the joined string and can + // substring-match (e.g. `context.claim_scp like "*Mail.Read*"` also matches + // "Mail.ReadWrite"), an unintended grant. Audit existing `like` policies on a + // claim before listing it, and prefer "claimset_" with `.contains`/ + // `.containsAll`/`.containsAny` for exact-element membership. + // + // Note: this unifies claim SHAPE, not NAME — `scp` and `scope` remain distinct + // claim names, so a policy still references one specific name. + MultiValuedClaims []string `json:"multi_valued_claims,omitempty" yaml:"multi_valued_claims,omitempty"` } // validateGroupEntityType validates a GroupEntityType value. Empty string is @@ -269,6 +340,9 @@ func NewCedarAuthorizer(options ConfigOptions, serverName string) (authorizers.A roleClaimName: options.RoleClaimName, serverName: serverName, claimKeyLog: syncutil.NewAtMost(30 * time.Second), + supplementLog: syncutil.NewAtMost(30 * time.Second), + missingIDTokenLog: syncutil.NewAtMost(30 * time.Second), + multiValuedClaims: options.MultiValuedClaims, } // Load policies @@ -479,53 +553,213 @@ func (a *Authorizer) IsAuthorized( } // resolveClaims determines which JWT claims to use for Cedar policy evaluation. -// When primaryUpstreamProvider is set, claims are extracted from the upstream -// IDP token stored in the identity. Otherwise, claims from the token on the -// original client request are used, which may be a ToolHive-issued token or -// any other bearer token. +// +// When primaryUpstreamProvider is empty (the default), the claims of the token on +// the original client request are used as-is. That may be a ToolHive-issued token +// or any other bearer token. +// +// When primaryUpstreamProvider is set (the embedded auth server path), the claim +// source is that provider's upstream credentials — and only that provider's. Its +// access token is primary; the profile claims listed in profileClaimsFromIDToken +// fall back to the same provider's id_token when the access token does not carry +// them, because many OIDC providers assert `email` only in the id_token and the +// access token alone was therefore a deny-all trap (#5916). +// +// Both tokens belong to the same upstream provider and session: the auth +// middleware splits a single credential bundle into Identity.UpstreamTokens and +// Identity.UpstreamIDTokens (see TokenValidator.Middleware in pkg/auth/token.go), +// so a given key denotes the same provider and session in both maps. That is what +// makes a supplemented claim carry the same provenance as one read from the access +// token: `principal has claim_email` means "this upstream asserted an email". (The +// two maps' key SETS may differ — a provider can have an access token but no +// id_token, which is why the fallback is conditional. That is not a case of one key +// meaning two different identities.) +// +// The ToolHive-issued token the client presented is deliberately NOT a claim +// source here, even though it mirrors a name/email — in a multi-upstream chain +// those belong to the first configured upstream, which need not be the pinned +// provider, so using them could attribute one IdP's email to another. +// +// The supplement is restricted to profileClaimsFromIDToken rather than merging +// the two token bodies. Read this before widening it: +// +// - An id_token's registered claims (`iss`, `aud`, `exp`, `nonce`, `at_hash`, +// `azp`) describe that token, not the user. Merging them would overwrite the +// access token's own `aud`/`exp` with values that mean something different. +// - Authorization-bearing claims (`groups`, `roles`, `scope`) are left alone for +// now. Taking them from this provider's id_token would be provenance-correct — +// unlike taking them from the ToolHive-issued token, which must never happen — +// but it would newly grant requests that deny today, so it belongs in its own +// change rather than riding along with a deny-all fix. +// - `sub` is excluded because it becomes the Cedar principal entity ID rather +// than an attribute, and Cedar has no `has`-style guard in the principal +// position. An access token without `sub` violates RFC 9068; failing closed +// with ErrMissingPrincipal beats silently choosing a principal. +// +// An access-token value always wins, so a claim the access token does assert can +// never be shadowed. A claim absent from both tokens stays absent, so `has`-guarded +// policies still fail closed. +// +// The id_token is used without checking `exp`, deliberately: it is read as a record +// of what the upstream asserted at login, not presented as a credential, and +// rejecting it once expired would flip a policy from permit to deny mid-session. +// See Identity.UpstreamIDTokens for that contract and its two consumers. func (a *Authorizer) resolveClaims(identity *auth.Identity) (jwt.MapClaims, error) { - if a.primaryUpstreamProvider != "" { - // Embedded auth server path: use the upstream IDP token's claims. - upstreamToken, tokenFound := identity.UpstreamTokens[a.primaryUpstreamProvider] - if !tokenFound || upstreamToken == "" { - return nil, authorizers.NewUpstreamAuthenticationRequiredError(a.primaryUpstreamProvider, "token_unavailable") + requestClaims := jwt.MapClaims(identity.Claims) + + if a.primaryUpstreamProvider == "" { + // Default path: use claims from the original client request's token. + a.logClaimKeys("token", requestClaims) + return requestClaims, nil + } + + // Embedded auth server path: the upstream IDP token is the primary claim source. + upstreamToken, tokenFound := identity.UpstreamTokens[a.primaryUpstreamProvider] + if !tokenFound || upstreamToken == "" { + // The upstream token must be present if the authorizer is configured to use it. + // Missing token means the session has no upstream credential; deny. + return nil, fmt.Errorf("upstream token for provider %q not found in identity", + a.primaryUpstreamProvider) + } + + upstreamClaims, err := parseUpstreamJWTClaims(upstreamToken) + if err != nil { + // Distinguish "not JWT-shaped" (opaque OAuth 2.0 access token — + // Google's ya29.*, GitHub's gho_*, etc.) from "JWT-shaped but + // malformed/tampered" (a JWT with three segments that fails to + // parse). Only fall back for the former; preserve the deny for + // the latter so a tampered upstream JWT cannot bypass policy. + // + // Policies that reference upstream-only claims (groups, hd, custom + // namespaced claims) see those attributes as absent on this branch and + // must be authored defensively (`principal has claim_groups && ...`). + if !looksLikeJWT(upstreamToken) { + // The Warn shares claimKeyLog with logClaimKeys below so a busy + // Google/GitHub deployment does not emit one line per tool call. + a.claimKeyLog.Do(func() { + slog.Warn("upstream token is not a JWT; falling back to request-token claims for Cedar evaluation", + "provider", a.primaryUpstreamProvider) + }) + a.logClaimKeys("token-fallback", requestClaims) + return requestClaims, nil } - parsedClaims, err := parseUpstreamJWTClaims(upstreamToken) - if err != nil { - // Distinguish "not JWT-shaped" (opaque OAuth 2.0 access token — - // Google's ya29.*, GitHub's gho_*, etc.) from "JWT-shaped but - // malformed/tampered" (a JWT with three segments that fails to - // parse). Only fall back for the former; preserve the deny for - // the latter so a tampered upstream JWT cannot bypass policy. - // - // The embedded auth server already mirrors the upstream OIDC - // sub/email/name claims into its issued AS token (see - // pkg/authserver/server/session/session.go). For opaque-token - // providers, falling back to identity.Claims preserves identity - // for policies referencing standard OIDC claims; policies that - // reference upstream-only claims (groups, hd, custom namespaced - // claims) will see those attributes as absent and must be - // authored defensively (`has(claim_groups) && ...`). - if !looksLikeJWT(upstreamToken) { - // The Warn shares claimKeyLog with logClaimKeys below so a busy - // Google/GitHub deployment does not emit one line per tool call. - a.claimKeyLog.Do(func() { - slog.Warn("upstream token is not a JWT; falling back to request-token claims for Cedar evaluation", - "provider", a.primaryUpstreamProvider) - }) - a.logClaimKeys("token-fallback", jwt.MapClaims(identity.Claims)) - return jwt.MapClaims(identity.Claims), nil - } - return nil, fmt.Errorf("failed to parse upstream token for provider %q: %w", - a.primaryUpstreamProvider, err) + return nil, fmt.Errorf("failed to parse upstream token for provider %q: %w", + a.primaryUpstreamProvider, err) + } + + merged := a.supplementFromIDToken(identity, upstreamClaims) + a.logClaimKeys("upstream", merged) + return merged, nil +} + +// supplementFromIDToken returns upstreamClaims with the profileClaimsFromIDToken +// it lacks filled in from the primary provider's stored id_token, logging what it +// did. upstreamClaims is not mutated. +// +// A provider with no stored id_token (an OAuth 2.0 upstream that was never asked +// for `openid`, so nothing was captured at login) yields no supplement: the claims +// stay absent and `has`-guarded policies deny. OIDC upstreams always have one — +// the provider rejects a login without it (upstream/oidc.go) and a refresh that +// omits a rotated id_token carries the original forward (upstreamtoken/service.go), +// so it does not vanish mid-session. +func (a *Authorizer) supplementFromIDToken(identity *auth.Identity, upstreamClaims jwt.MapClaims) jwt.MapClaims { + idToken := identity.UpstreamIDTokens[a.primaryUpstreamProvider] // nil map safe in Go + if idToken == "" { + a.missingIDTokenLog.Do(func() { + slog.Warn("no upstream ID token stored for provider; policies referencing profile claims "+ + "the access token omits will deny", + "provider", a.primaryUpstreamProvider, + "profile_claims", profileClaimsFromIDToken) + }) + return upstreamClaims + } + + idTokenClaims, err := parseUpstreamJWTClaims(idToken) + if err != nil { + // Do not fail the request: the access token parsed, so policies that only + // reference its claims must keep working. An unparsable id_token means no + // supplement, which denies exactly the policies that needed it. + a.missingIDTokenLog.Do(func() { + slog.Warn("upstream ID token for provider is not a parsable JWT; not supplementing profile claims", + "provider", a.primaryUpstreamProvider, + "error", err) + }) + return upstreamClaims + } + + merged, filled := supplementProfileClaims(upstreamClaims, idTokenClaims) + if len(filled) > 0 { + // Rate-limited on its own timer, not claimKeyLog: sharing one would let + // this Warn consume the window and permanently starve the claim-key dump + // in resolveClaims, which is what you want when debugging a denying policy. + // + // Claim keys only — never claim values, which hold user PII. + a.supplementLog.Do(func() { + slog.Warn("upstream access token lacks profile claims; using the upstream ID token's values "+ + "for Cedar evaluation", + "provider", a.primaryUpstreamProvider, + "claims", filled) + }) + } + return merged +} + +// OIDC Core 1.0 §5.1 standard claim names, declared here rather than borrowed +// from another package. What this code reads is an upstream provider's id_token, +// so these are wire names fixed by the OIDC spec — not names ToolHive chooses. +// Anchoring them to a ToolHive constant would invert the dependency: renaming that +// constant would silently change which claim Cedar reads out of a third party's +// token. Two literals also do not justify pulling an authorization-server +// dependency tree into the authorizer. +const ( + oidcNameClaim = "name" + oidcEmailClaim = "email" +) + +// profileClaimsFromIDToken are the OIDC Core §5.1 profile claims that may be read +// from the primary provider's id_token when its access token omits them. They +// coincide with the two claims the embedded auth server extracts from that same +// id_token: pkg/authserver/upstream/oidc.go reads them by their OIDC names, and the +// callback handler then passes them to session.New +// (pkg/authserver/server/handlers/callback.go). Both sides key off the wire names +// independently — a shared reliance on the spec, not a coupling. If the auth server +// ever mirrors more, adding it here is a separate, deliberate decision. +// +// Widening this list widens what can satisfy a policy — see resolveClaims first. +var profileClaimsFromIDToken = []string{oidcNameClaim, oidcEmailClaim} + +// supplementProfileClaims returns a fresh claim set holding every access-token +// claim plus, for each name in profileClaimsFromIDToken that the access token does +// not carry, the id_token's value for that name. It also returns the sorted list of +// names actually supplied by the id_token, for logging. +// +// Neither input is mutated and the result aliases neither, so the caller may hand +// it to code that adds synthetic keys. Absent claims are never fabricated: a name +// missing from both tokens is missing from the result, which keeps `has`-guarded +// policies failing closed. +func supplementProfileClaims(accessTokenClaims, idTokenClaims jwt.MapClaims) (merged jwt.MapClaims, filled []string) { + merged = make(jwt.MapClaims, len(accessTokenClaims)+len(profileClaimsFromIDToken)) + for k, v := range accessTokenClaims { + merged[k] = v + } + + for _, name := range profileClaimsFromIDToken { + if _, ok := merged[name]; ok { + // The access token asserted this claim; it always wins. + continue } - a.logClaimKeys("upstream", parsedClaims) - return parsedClaims, nil + v, ok := idTokenClaims[name] + if !ok { + continue + } + merged[name] = v + filled = append(filled, name) } - // Default path: use claims from the original client request's token. - claims := jwt.MapClaims(identity.Claims) - a.logClaimKeys("token", claims) - return claims, nil + + // Sorted for a canonical, easily-greppable log line — profileClaimsFromIDToken + // is ordered for readability (name, email), not alphabetically. + slices.Sort(filled) + return merged, filled } // looksLikeJWT returns true when the token has the three-segment shape of a @@ -584,12 +818,134 @@ func extractClientIDFromClaims(claims jwt.MapClaims) (string, bool) { func preprocessClaims(claims jwt.MapClaims) map[string]interface{} { preprocessed := make(map[string]interface{}) for k, v := range claims { - claimKey := fmt.Sprintf("claim_%s", k) + claimKey := claimPrefix + k preprocessed[claimKey] = v } return preprocessed } +// normalizeMultiValuedClaims returns claims with every claim named in names +// and present in claims rewritten to a canonical unpadded space-delimited +// string. This lets a single Cedar `like`/`==` policy match a claim regardless +// of whether the IdP emitted it as a space-delimited string or a JSON array. +// See ConfigOptions.MultiValuedClaims for the full rationale. +// +// When names is empty the input is returned directly (no copy); otherwise a +// shallow copy is returned and the input is never mutated. +// +// - array ([]interface{} or []string): elements joined with single spaces. +// Non-string elements are stringified with fmt.Sprint rather than +// dropped; nested array/object elements are skipped with a slog.Debug +// since they have no scalar string form. +// - string: passthrough, completely unchanged (not trimmed). +// - empty array, or an array with no usable elements: "" (empty string — +// present but empty). +// - claim absent from claims, or names empty: left untouched (no key +// fabricated). +// +// Only top-level claim keys are matched; there is no dot-notation support. +// claims is not mutated. +func normalizeMultiValuedClaims(claims jwt.MapClaims, names []string) jwt.MapClaims { + if len(names) == 0 { + // Nothing to normalize. Return the input directly — the sole caller + // (preprocessClaims) only reads it, so no defensive copy is needed and + // the opt-in-empty path pays no allocation. + return claims + } + + out := make(jwt.MapClaims, len(claims)) + for k, v := range claims { + out[k] = v + } + + for _, name := range names { + v, ok := claims[name] + if !ok { + continue + } + + switch val := v.(type) { + case string: + // Verbatim passthrough: out[name] already equals val (out is a copy + // of claims), so this is intentionally a no-op. The explicit branch + // documents the string contract and, importantly, keeps a string + // value out of the default "unrecognized type" log below. This is the + // one place the string handling diverges from addMultiValuedClaimSets, + // which collapses whitespace via strings.Fields for the Set form. + out[name] = val + case []interface{}: + out[name] = strings.Join(multiValuedTokens(name, val), " ") + case []string: + out[name] = strings.Join(val, " ") + default: + slog.Debug("multi-valued claim has unrecognized type, leaving unchanged", + "claim", name, "type", fmt.Sprintf("%T", v)) + } + } + + return out +} + +// addMultiValuedClaimSets adds a bare "claimset_" key to processed for every +// claim named in names and present in raw, holding the claim's elements as a +// []string (converted to a Cedar Set by convertToCedarValueAtDepth). This is +// the companion to normalizeMultiValuedClaims's "claim_" string form: +// where "claim_" supports substring-style `like` matching, "claimset_" +// supports exact-element `.contains`/`.containsAll`/`.containsAny` matching. +// The "claimset_" prefix is reserved and cannot collide with a JWT claim, since all +// JWT claims are "claim_"-prefixed by preprocessClaims. +// +// processed must already have the "claim_" prefix applied (i.e. this must run +// after preprocessClaims) so the added key stays bare. raw is not mutated. +func addMultiValuedClaimSets(processed map[string]interface{}, raw jwt.MapClaims, names []string) { + for _, name := range names { + v, ok := raw[name] + if !ok { + continue + } + + // The array path shares multiValuedTokens with normalizeMultiValuedClaims + // (string-only elements) so the Set and the joined string stay in sync. + // The string path deliberately differs: normalize passes the string + // through verbatim, whereas here strings.Fields splits it into exact + // elements (collapsing irregular whitespace) — the two cannot share one + // coercion for that reason. + switch val := v.(type) { + case string: + processed[claimSetPrefix+name] = strings.Fields(val) + case []interface{}: + processed[claimSetPrefix+name] = multiValuedTokens(name, val) + case []string: + processed[claimSetPrefix+name] = slices.Clone(val) + default: + slog.Debug("multi-valued claim has unrecognized type, omitting claimset set", + "claim", name, "type", fmt.Sprintf("%T", v)) + } + } +} + +// multiValuedTokens extracts the string elements of a []interface{} claim value +// for normalizeMultiValuedClaims and addMultiValuedClaimSets. Only string +// elements are kept: OAuth scope tokens are strings (RFC 6749 §3.3), so a +// non-string element (number, bool, null, or a nested array/object) is +// malformed and is skipped with a slog.Debug rather than stringified. Skipping +// (a) avoids spurious tokens like "" or "1e+06" polluting both claim +// forms, and (b) keeps the array shape consistent with the string shape, whose +// strings.Fields path likewise yields only real string tokens. +func multiValuedTokens(claimName string, elems []interface{}) []string { + tokens := make([]string, 0, len(elems)) + for _, elem := range elems { + s, ok := elem.(string) + if !ok { + slog.Debug("multi-valued claim element is not a string, skipping", + "claim", claimName, "type", fmt.Sprintf("%T", elem)) + continue + } + tokens = append(tokens, s) + } + return tokens +} + // preprocessArguments adds an "arg_" prefix to all argument keys. // For complex types, it just notes their presence with an "_present" suffix. func preprocessArguments(arguments map[string]interface{}) map[string]interface{} { @@ -881,8 +1237,12 @@ func (a *Authorizer) AuthorizeWithJWTClaims( extractGroups(resolvedClaims, a.roleClaimName)..., )) - // Preprocess claims and arguments - processedClaims := preprocessClaims(resolvedClaims) + // Preprocess claims and arguments. Multi-valued claim normalization runs + // after group/role extraction (which needs the raw claim shapes) and + // before the "claim_" prefix is applied. + normalizedClaims := normalizeMultiValuedClaims(resolvedClaims, a.multiValuedClaims) + processedClaims := preprocessClaims(normalizedClaims) + addMultiValuedClaimSets(processedClaims, resolvedClaims, a.multiValuedClaims) processedArgs := preprocessArguments(arguments) // Authorize based on the feature and operation diff --git a/pkg/authz/authorizers/cedar/core_test.go b/pkg/authz/authorizers/cedar/core_test.go index 42005a9d39..936e50fc90 100644 --- a/pkg/authz/authorizers/cedar/core_test.go +++ b/pkg/authz/authorizers/cedar/core_test.go @@ -8,6 +8,8 @@ import ( "context" "encoding/json" "log/slog" + "maps" + "sync" "testing" cedar "github.com/cedar-policy/cedar-go" @@ -1402,19 +1404,7 @@ func TestAuthorizeWithJWTClaims_UpstreamProvider(t *testing.T) { UpstreamTokens: map[string]string{}, }, wantErr: true, - errContains: "upstream authentication required", - }, - { - name: "upstream_token_missing_reports_token_unavailable", - identity: &auth.Identity{ - PrincipalInfo: auth.PrincipalInfo{ - Subject: "thv-user", - Claims: map[string]any{"sub": "thv-user"}, - }, - UpstreamTokens: map[string]string{}, - }, - wantErr: true, - errContains: "token_unavailable", + errContains: "upstream token for provider", }, { name: "upstream_token_opaque_falls_back_to_request_claims_denied", @@ -1479,7 +1469,7 @@ func TestAuthorizeWithJWTClaims_UpstreamProvider(t *testing.T) { UpstreamTokens: nil, }, wantErr: true, - errContains: "upstream authentication required", + errContains: "upstream token for provider", }, { name: "upstream_token_has_no_sub_claim", @@ -1495,6 +1485,11 @@ func TestAuthorizeWithJWTClaims_UpstreamProvider(t *testing.T) { }), }, }, + // `sub` is excluded from mirroredIdentityClaims, so the request + // token's subject must NOT stand in: the AS-issued subject is an + // internal user ID, and supplementing it would silently retarget the + // Cedar principal instead of failing closed. An upstream access token + // with no `sub` violates RFC 9068 and stays an error. wantErr: true, errContains: "missing principal", }, @@ -1528,6 +1523,497 @@ func TestAuthorizeWithJWTClaims_UpstreamProvider(t *testing.T) { } } +// TestSupplementProfileClaims covers which id_token claims may stand in for a +// missing access-token claim, and which may never do so. +func TestSupplementProfileClaims(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + accessTokenClaims jwt.MapClaims + idTokenClaims jwt.MapClaims + wantMerged jwt.MapClaims + wantFilled []string + }{ + { + name: "access_token_wins_on_shared_profile_claims", + accessTokenClaims: jwt.MapClaims{"sub": "okta|alice", "email": "at@example.com", "name": "AT Alice"}, + idTokenClaims: jwt.MapClaims{"sub": "okta|alice", "email": "idt@example.com", "name": "IDT Alice"}, + wantMerged: jwt.MapClaims{"sub": "okta|alice", "email": "at@example.com", "name": "AT Alice"}, + wantFilled: nil, + }, + { + // The #5916 shape: provider asserts email in the id_token only. + name: "id_token_fills_missing_profile_claims", + accessTokenClaims: jwt.MapClaims{"sub": "okta|alice", "iss": "https://idp.example.com"}, + idTokenClaims: jwt.MapClaims{"sub": "okta|alice", "email": "alice@example.com", "name": "Alice"}, + wantMerged: jwt.MapClaims{ + "sub": "okta|alice", + "iss": "https://idp.example.com", + "email": "alice@example.com", + "name": "Alice", + }, + // Sorted, so the logged claim list is stable across runs. + wantFilled: []string{"email", "name"}, + }, + { + // Taking these from this provider's id_token would be provenance-correct, + // but it would newly grant requests that deny today, so it is out of scope + // here and pinned as unchanged behaviour. + name: "authorization_bearing_claims_are_not_filled", + accessTokenClaims: jwt.MapClaims{"sub": "okta|alice"}, + idTokenClaims: jwt.MapClaims{ + "groups": []interface{}{"platform-eng"}, + "roles": []interface{}{"admin"}, + "scope": "tools:write", + "hd": "example.com", + }, + wantMerged: jwt.MapClaims{"sub": "okta|alice"}, + wantFilled: nil, + }, + { + // An id_token's registered claims describe that token. Merging them would + // overwrite the access token's own aud/exp with different-meaning values. + name: "id_token_registered_claims_are_never_filled_or_overwritten", + accessTokenClaims: jwt.MapClaims{ + "sub": "okta|alice", "aud": "api://resource", "exp": 2000000000, + }, + idTokenClaims: jwt.MapClaims{ + "aud": "toolhive-as-client", "exp": 1000000000, + "nonce": "n-abc", "at_hash": "h-abc", "azp": "client-x", + }, + wantMerged: jwt.MapClaims{ + "sub": "okta|alice", "aud": "api://resource", "exp": 2000000000, + }, + wantFilled: nil, + }, + { + // `sub` becomes the Cedar principal entity ID, where Cedar offers no + // `has`-style guard, so it fails closed rather than being chosen for us. + name: "subject_is_never_filled_from_the_id_token", + accessTokenClaims: jwt.MapClaims{"iss": "https://idp.example.com"}, + idTokenClaims: jwt.MapClaims{"sub": "okta|alice", "email": "alice@example.com"}, + wantMerged: jwt.MapClaims{ + "iss": "https://idp.example.com", + "email": "alice@example.com", + }, + wantFilled: []string{"email"}, + }, + { + // Never fabricate: a claim absent from both tokens stays absent so + // `has`-guarded policies keep failing closed. + name: "claim_absent_from_both_stays_absent", + accessTokenClaims: jwt.MapClaims{"sub": "okta|alice"}, + idTokenClaims: jwt.MapClaims{"sub": "okta|alice"}, + wantMerged: jwt.MapClaims{"sub": "okta|alice"}, + wantFilled: nil, + }, + { + // A key the access token maps to nil is still "asserted by the access + // token": presence, not truthiness, decides precedence. + name: "explicit_nil_access_token_value_still_wins", + accessTokenClaims: jwt.MapClaims{"email": nil}, + idTokenClaims: jwt.MapClaims{"email": "alice@example.com"}, + wantMerged: jwt.MapClaims{"email": nil}, + wantFilled: nil, + }, + { + name: "nil_inputs_yield_empty_result", + accessTokenClaims: nil, + idTokenClaims: nil, + wantMerged: jwt.MapClaims{}, + wantFilled: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Snapshot the inputs so the no-mutation contract can be asserted. + accessBefore := maps.Clone(tt.accessTokenClaims) + idTokenBefore := maps.Clone(tt.idTokenClaims) + + merged, filled := supplementProfileClaims(tt.accessTokenClaims, tt.idTokenClaims) + + assert.Equal(t, tt.wantMerged, merged) + assert.Equal(t, tt.wantFilled, filled) + assert.Equal(t, accessBefore, tt.accessTokenClaims, "access token claims must not be mutated") + assert.Equal(t, idTokenBefore, tt.idTokenClaims, "id token claims must not be mutated") + + // The result must not alias either input, since the caller adds + // synthetic keys to the claim set it receives. + merged["injected-by-caller"] = true + assert.NotContains(t, tt.accessTokenClaims, "injected-by-caller") + assert.NotContains(t, tt.idTokenClaims, "injected-by-caller") + }) + } +} + +// TestResolveClaimsUpstreamSupplement pins the claim-source contract of +// resolveClaims on the embedded-auth-server path: the pinned provider's access +// token wins, its id_token supplies only the profile claims the access token +// omits, and the ToolHive-issued token is not a claim source at all. +// +// The email-filling case is the #5916 reproducer — an upstream IdP whose access +// token is a well-formed JWT carrying no `email` (the provider asserts it in the +// id_token only), which used to make every `claim_email` policy deny. +func TestResolveClaimsUpstreamSupplement(t *testing.T) { + t.Parallel() + + const providerName = "hub" + + // What the embedded auth server mirrors into the token it issues. In a + // multi-upstream chain these values come from the FIRST configured upstream, + // which need not be the pinned provider, so they must never reach Cedar on + // this path. Distinctive values make a leak obvious. + asIssuedClaims := map[string]any{ + "sub": "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42", + "email": "leaked-from-as-token@example.com", + "name": "Leaked From AS Token", + "client_id": "vscode", + } + + tests := []struct { + name string + provider string + requestClaims map[string]any + upstreamToken string + idToken string + wantClaims jwt.MapClaims + }{ + { + name: "access_token_without_email_falls_back_to_id_token_email", + provider: providerName, + upstreamToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "iss": "https://hub.example.com", + // No email: the provider asserts it in the id_token only. + }), + idToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@example.com", + "name": "Alice", + "aud": "toolhive-as-client", + "nonce": "n-abc", + }), + requestClaims: asIssuedClaims, + wantClaims: jwt.MapClaims{ + "sub": "hub|alice", + "iss": "https://hub.example.com", + "email": "alice@example.com", + "name": "Alice", + // The id_token's own aud/nonce are absent: they describe that + // token, not the user. + }, + }, + { + name: "access_token_claims_win_over_id_token_claims", + provider: providerName, + upstreamToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@access-token.example", + "groups": []interface{}{"engineering"}, + }), + idToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@id-token.example", + "groups": []interface{}{"admins"}, + }), + requestClaims: asIssuedClaims, + wantClaims: jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@access-token.example", + "groups": []interface{}{"engineering"}, + }, + }, + { + // The multi-upstream misattribution guard. With no id_token stored for + // the pinned provider, `email` stays absent and the policy denies — the + // ToolHive-issued token's mirrored email must NOT stand in, because in a + // chain it belongs to the first configured upstream, not this one. + name: "no_id_token_leaves_profile_claims_absent", + provider: providerName, + upstreamToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "iss": "https://hub.example.com", + }), + idToken: "", + requestClaims: asIssuedClaims, + wantClaims: jwt.MapClaims{ + "sub": "hub|alice", + "iss": "https://hub.example.com", + }, + }, + { + // An id_token is read as a record of what the upstream asserted at + // login, not presented as a credential, so `exp` is deliberately not + // enforced. Enforcing it would make a policy permit early in a session + // and silently deny later, since sessions outlive id_tokens by days. + name: "expired_id_token_is_still_used", + provider: providerName, + upstreamToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + }), + idToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@example.com", + "exp": 1000000000, // long past + }), + requestClaims: asIssuedClaims, + wantClaims: jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@example.com", + }, + }, + { + // A malformed id_token must not fail the request: the access token + // parsed, so policies referencing only its claims keep working. + name: "unparsable_id_token_degrades_to_no_supplement", + provider: providerName, + upstreamToken: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + }), + idToken: "not-base64.not-base64.not-base64", + requestClaims: asIssuedClaims, + wantClaims: jwt.MapClaims{"sub": "hub|alice"}, + }, + { + name: "no_provider_configured_uses_request_claims_verbatim", + provider: "", + requestClaims: map[string]any{"sub": "7f3c1e64-uuid", "email": "alice@thv.example"}, + // Present but irrelevant: without a configured provider neither upstream + // token is consulted. + upstreamToken: makeUnsignedJWT(jwt.MapClaims{"sub": "hub|alice"}), + idToken: makeUnsignedJWT(jwt.MapClaims{"sub": "hub|alice"}), + wantClaims: jwt.MapClaims{"sub": "7f3c1e64-uuid", "email": "alice@thv.example"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{`permit(principal, action, resource);`}, + EntitiesJSON: `[]`, + PrimaryUpstreamProvider: tt.provider, + }, "") + require.NoError(t, err) + + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{Subject: "7f3c1e64-uuid", Claims: tt.requestClaims}, + UpstreamTokens: map[string]string{providerName: tt.upstreamToken}, + } + if tt.idToken != "" { + identity.UpstreamIDTokens = map[string]string{providerName: tt.idToken} + } + claimsBefore := maps.Clone(tt.requestClaims) + + resolved, err := authorizer.(*Authorizer).resolveClaims(identity) + require.NoError(t, err) + + assert.Equal(t, tt.wantClaims, resolved) + // Identity MUST NOT be modified after it is placed in the request + // context, since it is read concurrently. + assert.Equal(t, claimsBefore, tt.requestClaims, "identity claims must not be mutated") + }) + } +} + +// TestAuthorizeWithJWTClaims_UpstreamJWTMissingIdentityClaim is the end-to-end +// authorizer-level reproducer for #5916: the exact policy shape from the report +// (a defensively-authored `has`-guarded domain gate plus a tool permit) evaluated +// against an upstream access token that is a well-formed JWT without `email`. +// Before the fix this denied every request for every user. +func TestAuthorizeWithJWTClaims_UpstreamJWTMissingIdentityClaim(t *testing.T) { + t.Parallel() + + const providerName = "hub" + + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{ + `forbid(principal, action, resource) + unless { principal has claim_email && principal.claim_email like "*@example.com" };`, + `permit(principal, action == Action::"call_tool", resource);`, + }, + EntitiesJSON: `[]`, + PrimaryUpstreamProvider: providerName, + }, "") + require.NoError(t, err) + + // A JWT-shaped upstream access token that parses cleanly but carries no + // identity claims beyond `sub`. looksLikeJWT is true here, so the #5147 + // opaque-token fallback does not apply — the id_token supplement is what + // recovers `email`. + upstreamToken := makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|alice", + "iss": "https://hub.example.com", + }) + + // The ToolHive-issued token mirrors an email too, but it must never be the + // source on this path; it is outside the gated domain so a leak would show up + // as an unexpected deny in the permitting cases. + asIssuedClaims := map[string]any{ + "sub": "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42", + "email": "leaked-from-as-token@wrong-provider.example", + } + + tests := []struct { + name string + idTokenClaims jwt.MapClaims + wantAuthorize bool + }{ + { + // The upstream asserts the email in its id_token, so the domain gate + // can be satisfied even though the access token omits it. + name: "id_token_email_in_gated_domain_permits", + idTokenClaims: jwt.MapClaims{"sub": "hub|alice", "email": "alice@example.com"}, + wantAuthorize: true, + }, + { + // The supplement supplies the claim; it does not weaken the gate. + name: "id_token_email_outside_gated_domain_denies", + idTokenClaims: jwt.MapClaims{"sub": "hub|alice", "email": "alice@evil.example"}, + wantAuthorize: false, + }, + { + // Neither upstream token carries `email`: the `has` guard fails and the + // forbid applies. The supplement never fabricates a claim, and the + // ToolHive-issued token's mirrored email does not stand in. + name: "email_absent_from_both_upstream_tokens_denies", + idTokenClaims: jwt.MapClaims{"sub": "hub|alice"}, + wantAuthorize: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42", + Claims: asIssuedClaims, + }, + UpstreamTokens: map[string]string{providerName: upstreamToken}, + UpstreamIDTokens: map[string]string{providerName: makeUnsignedJWT(tt.idTokenClaims)}, + } + ctx := auth.WithIdentity(context.Background(), identity) + + authorized, err := authorizer.AuthorizeWithJWTClaims( + ctx, + authorizers.MCPFeatureTool, + authorizers.MCPOperationCall, + "deploy", + nil, + ) + require.NoError(t, err) + assert.Equal(t, tt.wantAuthorize, authorized) + }) + } +} + +// TestAuthorizeWithJWTClaims_PrincipalStaysUpstreamSourced pins the property that +// keeps `sub` out of mirroredIdentityClaims: the Cedar principal entity ID is +// derived from the upstream subject, never from the AS-issued token's internal +// user ID. Unlike an attribute, the principal has no `has`-style guard available +// in Cedar, so a supplemented subject would silently retarget every principal-keyed +// rule instead of failing closed. +func TestAuthorizeWithJWTClaims_PrincipalStaysUpstreamSourced(t *testing.T) { + t.Parallel() + + const providerName = "hub" + + // A banned-user rule keyed on the upstream subject, plus a broad permit that + // would otherwise let the request through. If the principal were ever built + // from the AS-issued subject, the forbid would stop matching and the permit + // would apply — the exact silent-widening this guards against. + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{ + `forbid(principal == Client::"hub|banned-alice", action, resource);`, + `permit(principal, action == Action::"call_tool", resource);`, + }, + EntitiesJSON: `[]`, + PrimaryUpstreamProvider: providerName, + }, "") + require.NoError(t, err) + + // The AS-issued subject: a UserResolver UUID, unrelated to the upstream one. + const asSubject = "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42" + + tests := []struct { + name string + upstreamClaims jwt.MapClaims + wantAuthorize bool + wantErr bool + }{ + { + name: "banned_upstream_subject_is_forbidden", + upstreamClaims: jwt.MapClaims{"sub": "hub|banned-alice"}, + wantAuthorize: false, + }, + { + name: "other_upstream_subject_is_permitted", + upstreamClaims: jwt.MapClaims{"sub": "hub|bob"}, + wantAuthorize: true, + }, + { + // The sharp case. If `sub` were added to mirroredIdentityClaims, the + // AS subject would stand in, the principal would become + // Client::"7f3c1e64-..." and the broad permit would apply — the + // request would be ALLOWED instead of erroring. Fail closed instead. + name: "missing_upstream_subject_fails_closed_not_to_the_as_subject", + upstreamClaims: jwt.MapClaims{"iss": "https://hub.example.com"}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: asSubject, + Claims: map[string]any{"sub": asSubject, "email": "alice@example.com"}, + }, + UpstreamTokens: map[string]string{ + providerName: makeUnsignedJWT(tt.upstreamClaims), + }, + // The id_token carries a subject too. It is the correct upstream + // subject, but `sub` is still excluded from the supplement so a + // missing access-token `sub` fails closed rather than having a + // principal chosen for it. + UpstreamIDTokens: map[string]string{ + providerName: makeUnsignedJWT(jwt.MapClaims{ + "sub": "hub|from-id-token", + "email": "alice@example.com", + }), + }, + } + ctx := auth.WithIdentity(context.Background(), identity) + + authorized, err := authorizer.AuthorizeWithJWTClaims( + ctx, + authorizers.MCPFeatureTool, + authorizers.MCPOperationCall, + "deploy", + nil, + ) + + if tt.wantErr { + require.ErrorIs(t, err, ErrMissingPrincipal) + assert.False(t, authorized) + return + } + + require.NoError(t, err) + assert.Equal(t, tt.wantAuthorize, authorized) + }) + } +} + // TestAuthorizeWithJWTClaims_GroupMembership verifies that Cedar policies using // "principal in THVGroup::..." are enforced when groups are present in the claims. func TestAuthorizeWithJWTClaims_GroupMembership(t *testing.T) { @@ -2583,6 +3069,205 @@ func TestAuthorizeWithJWTClaims_BackwardCompat(t *testing.T) { } } +// TestAuthorizeWithJWTClaims_MultiValuedClaimSets is the end-to-end proof that +// both normalized surfaces work for both IdP shapes (array and space-delimited +// string): the "claimset_" Cedar Set for exact-element `.contains`/ +// `.containsAll` matching, and (in the containsAll subtest) that a missing +// token correctly denies. It also proves `.contains` is exact-element — a +// "Mail.ReadWrite" token does not satisfy a "Mail.Read" check. +func TestAuthorizeWithJWTClaims_MultiValuedClaimSets(t *testing.T) { + t.Parallel() + + newAuthorizerCtx := func(t *testing.T, scp interface{}) context.Context { + t.Helper() + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user1", + Claims: map[string]any{ + "sub": "user1", + "scp": scp, + }, + }, + } + return auth.WithIdentity(context.Background(), identity) + } + + t.Run("contains_exact_element_no_substring_bleed", func(t *testing.T) { + t.Parallel() + + policy := ` + permit(principal, action, resource) + when { context.claimset_scp.contains("Mail.Read") }; + ` + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{policy}, + EntitiesJSON: `[]`, + MultiValuedClaims: []string{"scp"}, + }, "") + require.NoError(t, err) + + tests := []struct { + name string + scp interface{} + wantAuth bool + }{ + { + name: "array_shaped_scp_permitted", + scp: []interface{}{"User.Read.All", "Mail.Read"}, + wantAuth: true, + }, + { + name: "string_shaped_scp_permitted", + scp: "User.Read.All Mail.Read", + wantAuth: true, + }, + { + name: "substring_scope_not_matched", + scp: []interface{}{"Mail.ReadWrite"}, + wantAuth: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + authorized, err := authorizer.AuthorizeWithJWTClaims( + newAuthorizerCtx(t, tt.scp), + authorizers.MCPFeatureTool, + authorizers.MCPOperationCall, + "any-tool", + nil, + ) + require.NoError(t, err) + assert.Equal(t, tt.wantAuth, authorized) + }) + } + }) + + t.Run("contains_all_tokens", func(t *testing.T) { + t.Parallel() + + policy := ` + permit(principal, action, resource) + when { context.claimset_scp.containsAll(["User.Read.All", "Mail.Read"]) }; + ` + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{policy}, + EntitiesJSON: `[]`, + MultiValuedClaims: []string{"scp"}, + }, "") + require.NoError(t, err) + + tests := []struct { + name string + scp interface{} + wantAuth bool + }{ + { + name: "array_shaped_scp_has_both_tokens_permitted", + scp: []interface{}{"User.Read.All", "Mail.Read", "Extra.Scope"}, + wantAuth: true, + }, + { + name: "string_shaped_scp_has_both_tokens_permitted", + scp: "User.Read.All Mail.Read", + wantAuth: true, + }, + { + name: "missing_one_token_denied", + scp: []interface{}{"User.Read.All"}, + wantAuth: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + authorized, err := authorizer.AuthorizeWithJWTClaims( + newAuthorizerCtx(t, tt.scp), + authorizers.MCPFeatureTool, + authorizers.MCPOperationCall, + "any-tool", + nil, + ) + require.NoError(t, err) + assert.Equal(t, tt.wantAuth, authorized) + }) + } + }) + + // principal_claimset_set proves the claimset_ Set also appears as a principal + // entity attribute, not just in the context — mirroring how claim_ + // is available on both surfaces. + t.Run("principal_claimset_set", func(t *testing.T) { + t.Parallel() + + policy := ` + permit(principal, action, resource) + when { principal.claimset_scp.contains("Mail.Read") }; + ` + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{policy}, + EntitiesJSON: `[]`, + MultiValuedClaims: []string{"scp"}, + }, "") + require.NoError(t, err) + + authorized, err := authorizer.AuthorizeWithJWTClaims( + newAuthorizerCtx(t, []interface{}{"User.Read.All", "Mail.Read"}), + authorizers.MCPFeatureTool, + authorizers.MCPOperationCall, + "any-tool", + nil, + ) + require.NoError(t, err) + assert.True(t, authorized) + }) +} + +// TestAuthorizeWithJWTClaims_MultiValuedClaimsEmptyIsBackwardCompat verifies that +// when MultiValuedClaims is unset, a whole-string equality policy on the claim +// still matches — i.e. normalization is fully opt-in and byte-identical to the +// pre-#5828 behavior when no claim names are listed. +func TestAuthorizeWithJWTClaims_MultiValuedClaimsEmptyIsBackwardCompat(t *testing.T) { + t.Parallel() + + policy := ` + permit(principal, action, resource) + when { context.claim_scp == "User.Read.All Mail.Read" }; + ` + + authorizer, err := NewCedarAuthorizer(ConfigOptions{ + Policies: []string{policy}, + EntitiesJSON: `[]`, + // MultiValuedClaims intentionally left unset. + }, "") + require.NoError(t, err) + + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "user1", + Claims: map[string]any{ + "sub": "user1", + "scp": "User.Read.All Mail.Read", + }, + }, + } + ctx := auth.WithIdentity(context.Background(), identity) + + authorized, err := authorizer.AuthorizeWithJWTClaims( + ctx, + authorizers.MCPFeatureTool, + authorizers.MCPOperationCall, + "any-tool", + nil, + ) + require.NoError(t, err) + assert.True(t, authorized, "unpadded whole-string equality must still match when MultiValuedClaims is unset") +} + // TestParseCedarEntityID tests the parseCedarEntityID helper function. func TestParseCedarEntityID(t *testing.T) { t.Parallel() @@ -2780,6 +3465,201 @@ func TestPreprocessClaims(t *testing.T) { } } +// TestNormalizeMultiValuedClaims tests the normalizeMultiValuedClaims helper, +// which rewrites listed claims to a canonical unpadded space-delimited string +// so a single Cedar `like`/`==` policy matches regardless of whether the IdP +// emitted the claim as a space-delimited string or a JSON array. +func TestNormalizeMultiValuedClaims(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + claims jwt.MapClaims + names []string + want jwt.MapClaims + }{ + { + name: "array_of_strings_joined_unpadded", + claims: jwt.MapClaims{"scp": []interface{}{"a", "b"}}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": "a b"}, + }, + { + name: "string_array_joined_unpadded", + claims: jwt.MapClaims{"scp": []string{"a", "b"}}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": "a b"}, + }, + { + name: "single_element_array_unpadded", + claims: jwt.MapClaims{"scp": []interface{}{"a"}}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": "a"}, + }, + { + name: "string_passthrough_unchanged", + claims: jwt.MapClaims{"scp": "a b"}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": "a b"}, + }, + { + name: "string_with_edge_and_internal_spaces_unchanged", + claims: jwt.MapClaims{"scp": " a b "}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": " a b "}, + }, + { + name: "empty_array_becomes_empty_string", + claims: jwt.MapClaims{"scp": []interface{}{}}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": ""}, + }, + { + name: "absent_claim_stays_absent", + claims: jwt.MapClaims{"sub": "user1"}, + names: []string{"scp"}, + want: jwt.MapClaims{"sub": "user1"}, + }, + { + name: "claim_not_in_names_untouched", + claims: jwt.MapClaims{"scp": []interface{}{"a", "b"}, "sub": "user1"}, + names: []string{"scope"}, + want: jwt.MapClaims{"scp": []interface{}{"a", "b"}, "sub": "user1"}, + }, + { + name: "non_string_elements_skipped", + claims: jwt.MapClaims{"scp": []interface{}{"a", 1, true}}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": "a"}, + }, + { + name: "nested_element_skipped", + claims: jwt.MapClaims{"scp": []interface{}{"a", []interface{}{"nested"}}}, + names: []string{"scp"}, + want: jwt.MapClaims{"scp": "a"}, + }, + { + name: "empty_names_list_leaves_contents_identical", + claims: jwt.MapClaims{"scp": []interface{}{"a", "b"}, "sub": "user1"}, + names: nil, + want: jwt.MapClaims{"scp": []interface{}{"a", "b"}, "sub": "user1"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := normalizeMultiValuedClaims(tt.claims, tt.names) + assert.Equal(t, tt.want, got) + }) + } +} + +// TestAddMultiValuedClaimSets tests the addMultiValuedClaimSets helper, which +// injects a bare "claimset_" key holding the claim's elements as a []string +// (converted to a Cedar Set downstream) — the companion to +// normalizeMultiValuedClaims's "claim_" string form. +func TestAddMultiValuedClaimSets(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + processed map[string]interface{} + raw jwt.MapClaims + names []string + want map[string]interface{} + }{ + { + name: "array_claim_becomes_set", + processed: map[string]interface{}{"claim_scp": "a b"}, + raw: jwt.MapClaims{"scp": []interface{}{"a", "b"}}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": "a b", "claimset_scp": []string{"a", "b"}}, + }, + { + name: "string_claim_split_into_set", + processed: map[string]interface{}{"claim_scp": "a b"}, + raw: jwt.MapClaims{"scp": "a b"}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": "a b", "claimset_scp": []string{"a", "b"}}, + }, + { + // Companion to normalizeMultiValuedClaims's verbatim passthrough: + // claim_scp keeps irregular spacing untouched, while claimset_scp + // collapses it via strings.Fields into exact elements. + name: "multi_space_string_collapsed_into_set", + processed: map[string]interface{}{"claim_scp": " a b "}, + raw: jwt.MapClaims{"scp": " a b "}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": " a b ", "claimset_scp": []string{"a", "b"}}, + }, + { + name: "empty_array_becomes_empty_set", + processed: map[string]interface{}{"claim_scp": ""}, + raw: jwt.MapClaims{"scp": []interface{}{}}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": "", "claimset_scp": []string{}}, + }, + { + name: "empty_string_becomes_empty_set", + processed: map[string]interface{}{"claim_scp": ""}, + raw: jwt.MapClaims{"scp": ""}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": "", "claimset_scp": []string{}}, + }, + { + name: "absent_claim_no_key_added", + processed: map[string]interface{}{"claim_sub": "user1"}, + raw: jwt.MapClaims{"sub": "user1"}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_sub": "user1"}, + }, + { + name: "claim_not_in_names_no_key_added", + processed: map[string]interface{}{"claim_scp": "a b"}, + raw: jwt.MapClaims{"scp": []interface{}{"a", "b"}}, + names: []string{"scope"}, + want: map[string]interface{}{"claim_scp": "a b"}, + }, + { + name: "non_string_elements_skipped", + processed: map[string]interface{}{"claim_scp": "a"}, + raw: jwt.MapClaims{"scp": []interface{}{"a", 1, true}}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": "a", "claimset_scp": []string{"a"}}, + }, + { + name: "nested_element_skipped", + processed: map[string]interface{}{"claim_scp": "a"}, + raw: jwt.MapClaims{"scp": []interface{}{"a", []interface{}{"nested"}}}, + names: []string{"scp"}, + want: map[string]interface{}{"claim_scp": "a", "claimset_scp": []string{"a"}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + addMultiValuedClaimSets(tt.processed, tt.raw, tt.names) + assert.Equal(t, tt.want, tt.processed) + + // The injected key must be bare "claimset_", never + // "claim_claimset_" — it must not be reachable via the + // claim_-prefix path. + for _, name := range tt.names { + _, hadRaw := tt.raw[name] + if !hadRaw { + continue + } + _, hasBareKey := tt.processed["claimset_"+name] + _, hasWrongKey := tt.processed["claim_claimset_"+name] + assert.True(t, hasBareKey, "expected bare claimset_%s key", name) + assert.False(t, hasWrongKey, "must not add claim_claimset_%s", name) + } + }) + } +} + // TestPreprocessArguments tests the preprocessArguments helper. func TestPreprocessArguments(t *testing.T) { t.Parallel() @@ -2997,6 +3877,51 @@ func TestConfigOptionsRoleClaimNameJSON(t *testing.T) { } } +// TestConfigOptionsMultiValuedClaimsJSON verifies JSON marshal/unmarshal of the +// MultiValuedClaims field, including backward compatibility when the field is absent. +func TestConfigOptionsMultiValuedClaimsJSON(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + jsonInput string + want []string + wantOmitOnMar bool // when true, marshal output must NOT contain "multi_valued_claims" + }{ + { + name: "present", + jsonInput: `{"policies":["permit(principal,action,resource);"],"multi_valued_claims":["scp","scope"]}`, + want: []string{"scp", "scope"}, + }, + { + name: "absent_gives_nil", + jsonInput: `{"policies":["permit(principal,action,resource);"]}`, + want: nil, + wantOmitOnMar: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var opts ConfigOptions + err := json.Unmarshal([]byte(tt.jsonInput), &opts) + require.NoError(t, err) + assert.Equal(t, tt.want, opts.MultiValuedClaims) + + marshalled, err := json.Marshal(opts) + require.NoError(t, err) + if tt.wantOmitOnMar { + assert.NotContains(t, string(marshalled), "multi_valued_claims", + "empty MultiValuedClaims must be omitted from JSON output") + } else { + assert.Contains(t, string(marshalled), "multi_valued_claims") + } + }) + } +} + // TestValidateGroupEntityType exercises the private validateGroupEntityType helper // directly. Each case names an input, states whether it should succeed, and — for // error cases — a substring that the error message must contain so operators can @@ -3208,14 +4133,42 @@ func TestNewCedarAuthorizerGroupEntityTypeValidation(t *testing.T) { } } -// captureSlogWarn redirects slog's default logger to a bytes.Buffer for the -// duration of f, then restores the original default. Returns the captured -// output. This helper exists because slog.SetDefault is a process-global -// side effect — tests that use it must NOT run in parallel. +// syncBuffer is a concurrency-safe io.Writer over a bytes.Buffer. +// +// captureSlogWarn needs one because slog.SetDefault is process-global: while the +// capturing handler is installed, ANY goroutine in the test binary can emit a +// record into it. Keeping the capturing test itself non-parallel does not prevent +// that — other tests already running in parallel still log — so the buffer must +// be safe on its own. A plain bytes.Buffer here is a data race that only shows up +// once some concurrently-reachable code path starts logging. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +// captureSlogWarn redirects slog's default logger to a buffer for the duration of +// f, then restores the original default, and returns the captured output. +// +// Because slog.SetDefault is process-global, the returned output may also contain +// records emitted by other tests running concurrently. Callers must therefore +// assert on the specific record they expect (or its specific absence) and never +// on the output being empty overall. func captureSlogWarn(t *testing.T, f func()) string { t.Helper() - var buf bytes.Buffer + var buf syncBuffer handler := slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn}) orig := slog.Default() slog.SetDefault(slog.New(handler)) @@ -3279,14 +4232,19 @@ func TestStaleTHVGroupWarning(t *testing.T) { require.NoError(t, err) }) + // Keyed on this specific warning rather than on the buffer being + // empty: slog.Default is process-global, so concurrently-running + // tests can also land records in the captured output. + const staleWarn = "Cedar entities_json contains THVGroup entities" + if tt.wantWarn { - require.NotEmpty(t, output, "expected a warn log") + require.Contains(t, output, staleWarn, "expected the stale-THVGroup warn log") for _, want := range tt.wantContains { assert.Contains(t, output, want, "warn log must mention %q", want) } } else { - assert.Empty(t, output, "no warning expected") + assert.NotContains(t, output, staleWarn, "no stale-THVGroup warning expected") } }) } diff --git a/pkg/authz/authorizers/cedar/record_test.go b/pkg/authz/authorizers/cedar/record_test.go index 30f3b91bbd..80a3a589c9 100644 --- a/pkg/authz/authorizers/cedar/record_test.go +++ b/pkg/authz/authorizers/cedar/record_test.go @@ -172,9 +172,7 @@ func TestConvertMapToCedarRecord(t *testing.T) { { name: "Nested map converts to Cedar record", input: map[string]interface{}{ - "act": map[string]interface{}{ - "sub": "spiffe://toolhive.dev/ns/agents/sa/devops-agent", - }, + "act": map[string]interface{}{"sub": "spiffe://toolhive.dev/ns/agents/sa/devops-agent"}, }, expected: map[string]cedar.Value{ "act": cedar.NewRecord(cedar.RecordMap{ diff --git a/pkg/authz/authorizers/errors.go b/pkg/authz/authorizers/errors.go deleted file mode 100644 index 357f10f52b..0000000000 --- a/pkg/authz/authorizers/errors.go +++ /dev/null @@ -1,54 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package authorizers - -import ( - "errors" - "fmt" -) - -var errUpstreamAuthenticationRequired = errors.New("upstream authentication required") - -// UpstreamAuthenticationRequiredError indicates that authorization depends on -// an upstream identity provider token that is no longer usable and must be -// re-established via a fresh login. -type UpstreamAuthenticationRequiredError struct { - Provider string - Reason string -} - -func (e *UpstreamAuthenticationRequiredError) Error() string { - if e == nil { - return errUpstreamAuthenticationRequired.Error() - } - if e.Provider == "" && e.Reason == "" { - return errUpstreamAuthenticationRequired.Error() - } - if e.Reason == "" { - return fmt.Sprintf("%s for provider %q", errUpstreamAuthenticationRequired, e.Provider) - } - if e.Provider == "" { - return fmt.Sprintf("%s: %s", errUpstreamAuthenticationRequired, e.Reason) - } - return fmt.Sprintf("%s for provider %q: %s", errUpstreamAuthenticationRequired, e.Provider, e.Reason) -} - -func (*UpstreamAuthenticationRequiredError) Unwrap() error { - return errUpstreamAuthenticationRequired -} - -// NewUpstreamAuthenticationRequiredError creates a typed authorization error -// that callers can translate into a 401 re-authentication challenge. -func NewUpstreamAuthenticationRequiredError(provider, reason string) error { - return &UpstreamAuthenticationRequiredError{ - Provider: provider, - Reason: reason, - } -} - -// IsUpstreamAuthenticationRequired reports whether err indicates an upstream -// provider session that now requires re-authentication. -func IsUpstreamAuthenticationRequired(err error) bool { - return errors.Is(err, errUpstreamAuthenticationRequired) -} diff --git a/pkg/authz/errors.go b/pkg/authz/errors.go deleted file mode 100644 index 3bb7385923..0000000000 --- a/pkg/authz/errors.go +++ /dev/null @@ -1,58 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package authz - -import ( - "errors" - "net/http" - - "github.com/stacklok/toolhive/pkg/authz/authorizers" -) - -const upstreamAuthRequiredDescription = "upstream token is no longer valid; re-authentication required" - -var errForbiddenAuthorization = errors.New("forbidden authorization request") - -type forbiddenAuthorizationError struct { - message string -} - -func (e *forbiddenAuthorizationError) Error() string { - if e == nil || e.message == "" { - return "Unauthorized" - } - return e.message -} - -func (*forbiddenAuthorizationError) Unwrap() error { - return errForbiddenAuthorization -} - -func newForbiddenAuthorizationError(message string) error { - return &forbiddenAuthorizationError{message: message} -} - -func isForbiddenAuthorizationError(err error) bool { - return errors.Is(err, errForbiddenAuthorization) -} - -func bearerInvalidTokenChallenge(description string) string { - return `Bearer error="invalid_token", error_description="` + description + `"` -} - -func classifyAuthorizationError(err error) (httpStatus int, rpcCode int, message string, wwwAuthenticate string) { - if authorizers.IsUpstreamAuthenticationRequired(err) { - return http.StatusUnauthorized, - http.StatusUnauthorized, - upstreamAuthRequiredDescription, - bearerInvalidTokenChallenge(upstreamAuthRequiredDescription) - } - if isForbiddenAuthorizationError(err) { - return http.StatusForbidden, http.StatusForbidden, err.Error(), "" - } - if err != nil { - return http.StatusInternalServerError, http.StatusInternalServerError, err.Error(), "" - } - return http.StatusForbidden, http.StatusForbidden, "Unauthorized", "" -} diff --git a/pkg/authz/integration_test.go b/pkg/authz/integration_test.go index 79ca874820..b8cd12473a 100644 --- a/pkg/authz/integration_test.go +++ b/pkg/authz/integration_test.go @@ -11,11 +11,11 @@ import ( "testing" "github.com/golang-jwt/jwt/v5" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" mcpparser "github.com/stacklok/toolhive/pkg/mcp" @@ -1187,3 +1187,172 @@ func TestIntegrationPrimaryUpstreamProviderClaimAttributeAccess(t *testing.T) { }) } } + +// TestIntegrationUpstreamJWTWithoutEmailClaim is the end-to-end reproducer for +// stacklok/toolhive#5916, driven through the full authorization middleware with +// the exact policy pair from the report. +// +// When the embedded auth server is active, the runner force-injects the first +// upstream provider as Cedar's PrimaryUpstreamProvider, so Cedar sources claims +// from the upstream IdP's access token. Providers that issue JWT-shaped access +// tokens without `email` (JetBrains Hub is a public example) used to make every +// request deny: the token parsed cleanly, so the opaque-token fallback added in +// #5147 did not apply, and `principal has claim_email` was false for every user. +// Opaque-token upstreams (Google, GitHub) masked the bug because their tokens hit +// that fallback and saw the AS-issued token's mirrored `email`. +// +// The provider does assert `email` — in its id_token, which the auth server stores +// alongside the access token for the same provider. The fix reads the profile +// claims the access token omits from that id_token, keeping the claim's provenance +// with the pinned upstream. +func TestIntegrationUpstreamJWTWithoutEmailClaim(t *testing.T) { + t.Parallel() + + const providerName = "hub" + + authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ + // The report's policies verbatim: a defensively `has`-guarded domain + // gate plus a tool permit. + Policies: []string{ + `forbid(principal, action, resource) + unless { principal has claim_email && principal.claim_email like "*@example.com" };`, + `permit(principal, action == Action::"call_tool", resource);`, + }, + EntitiesJSON: "[]", + PrimaryUpstreamProvider: providerName, + }, "") + require.NoError(t, err) + + // Claims of the token the client actually presented, as minted by the embedded + // auth server: an internal ToolHive user ID (a UserResolver UUID, deliberately + // NOT the upstream subject) plus the profile it mirrored from the FIRST + // configured upstream. In a multi-upstream chain that is not necessarily the + // pinned provider, so these must never reach Cedar on this path — the + // distinctive email makes a leak visible as an unexpected deny. + asIssuedClaims := jwt.MapClaims{ + "sub": "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42", + "email": "leaked-from-as-token@wrong-provider.example", + "name": "Leaked From AS Token", + "iss": "https://thv-as.example.com/", + "aud": "toolhive", + "client_id": "vscode", + } + + // The pinned provider's id_token, stored beside its access token at login. + // This is where this provider asserts the email. + upstreamIDToken := makeUnsignedJWT(t, jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@example.com", + "name": "Alice", + }) + + tests := []struct { + name string + upstreamToken string + requestClaims jwt.MapClaims + expectAllowed bool + }{ + { + // The reported case: parses as a JWT, carries no `email`. + name: "jwt_access_token_without_email_uses_id_token_email", + upstreamToken: makeUnsignedJWT(t, jwt.MapClaims{ + "sub": "hub|alice", + "iss": "https://hub.example.com", + }), + expectAllowed: true, + }, + { + // The access token's own `email` still wins where it asserts one, and + // it is outside the gated domain here, so the gate must reject. + name: "access_token_email_wins_over_id_token_email", + upstreamToken: makeUnsignedJWT(t, jwt.MapClaims{ + "sub": "hub|alice", + "email": "alice@contractor.example", + }), + expectAllowed: false, + }, + { + // Regression guard for the #5147 path, which this change leaves intact: + // an opaque access token has no claims to read, so that branch still + // evaluates the ToolHive-issued token's claims. Those carry the + // upstream profile only because the auth server mirrors the FIRST + // configured upstream, so this case supplies a single-upstream + // AS token rather than the leak sentinel the JWT cases use. + // + // That leaves the two branches inconsistent about claim provenance + // (and about the principal: this branch's `sub` is ToolHive's internal + // user ID, not the upstream subject). Unifying them on the id_token is + // a follow-up — it changes behaviour on a shipped path. + name: "opaque_upstream_token_still_falls_back", + upstreamToken: "ya29.opaque-google-style-token", + requestClaims: jwt.MapClaims{ + "sub": "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42", + "email": "alice@example.com", + "name": "Alice", + }, + expectAllowed: true, + }, + { + // A JWT-shaped but unparsable upstream token must stay a hard deny: + // silently degrading a tampered token to other claims would let it + // bypass policy. + name: "tampered_upstream_jwt_still_denies", + upstreamToken: "not-base64.not-base64.not-base64", + expectAllowed: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + params, err := json.Marshal(map[string]interface{}{ + "name": "deploy", + "arguments": map[string]interface{}{}, + }) + require.NoError(t, err) + callReq, err := jsonrpc2.NewCall(jsonrpc2.Int64ID(1), string(mcp.MethodToolsCall), json.RawMessage(params)) + require.NoError(t, err) + reqJSON, err := jsonrpc2.EncodeMessage(callReq) + require.NoError(t, err) + + httpReq, err := http.NewRequest(http.MethodPost, "/messages", bytes.NewBuffer(reqJSON)) + require.NoError(t, err) + httpReq.Header.Set("Content-Type", "application/json") + + requestClaims := tt.requestClaims + if requestClaims == nil { + requestClaims = asIssuedClaims + } + + identity := &auth.Identity{ + PrincipalInfo: auth.PrincipalInfo{ + Subject: "7f3c1e64-9b2a-4d51-8e77-1c0a5f3b9d42", + Claims: requestClaims, + }, + UpstreamTokens: map[string]string{providerName: tt.upstreamToken}, + UpstreamIDTokens: map[string]string{providerName: upstreamIDToken}, + } + httpReq = httpReq.WithContext(auth.WithIdentity(httpReq.Context(), identity)) + + rr := httptest.NewRecorder() + var handlerCalled bool + mockHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + handlerCalled = true + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"result": "ok"}`)) + }) + + middleware := mcpparser.ParsingMiddleware(Middleware(authorizer, mockHandler, nil)) + middleware.ServeHTTP(rr, httpReq) + + if tt.expectAllowed { + assert.Equal(t, http.StatusOK, rr.Code) + assert.True(t, handlerCalled) + } else { + assert.Equal(t, http.StatusForbidden, rr.Code) + assert.False(t, handlerCalled) + } + }) + } +} diff --git a/pkg/authz/middleware.go b/pkg/authz/middleware.go index 9be0aa9436..bd29e2c37a 100644 --- a/pkg/authz/middleware.go +++ b/pkg/authz/middleware.go @@ -19,6 +19,8 @@ import ( "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/ssecommon" "github.com/stacklok/toolhive/pkg/transport/types" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" + "github.com/stacklok/toolhive/pkg/vmcp/schema" "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" ) @@ -54,6 +56,34 @@ var MCPMethodToFeatureOperation = map[string]featureOperation{ "features/list": {Feature: "", Operation: authorizers.MCPOperationList}, // Capability discovery "roots/list": {Feature: "", Operation: ""}, // Root directory discovery + // server/discover, Modern's (2026-07-28) replacement for initialize+capability + // negotiation, is always-allowed on THIS path (the single-server pkg/runner HTTP + // authz Middleware -- vMCP's Modern dispatcher never consults this map at all, it + // re-homes admission through core.Check*/core.List* directly). The always-allowed + // choice rests on initialize parity, not on any per-request filtering this map + // enforces: DiscoverResult carries the exact same Capabilities *ServerCapabilities + // (+ Instructions) shape InitializeResult does, and "initialize" above has always + // been always-allowed in this map. discover therefore adds no new exposure class -- + // note ServerCapabilities.Experimental/.Extensions (arbitrary backend-authored maps) + // and Instructions (free text) are already freeform fields a backend can populate on + // the always-allowed initialize response today, so "no descriptors" is a property of + // how vMCP's dispatcher happens to build the value, not a guarantee this wire shape + // makes on its own. Classifying it as MCPOperationList instead would be safe too -- + // response_filter.go hardcodes an exact 4-method filter list (tools/list, + // prompts/list, resources/list, find_tool), so server/discover would just pass + // through unfiltered -- but always-allowed is simpler and equally safe here. + "server/discover": {Feature: "", Operation: ""}, + + // Subscriptions - always allowed for now. This method carries no single resource + // identifier the parser extracts (params are a notification-type filter with an + // optional resourceSubscriptions array), so routing it through Cedar with an empty + // ResourceID would risk matching a broad allow rule. Notification delivery and + // per-resource authorization of resourceSubscriptions URIs are future work. + // + // TODO(#5755): when subscription notification delivery is implemented, replace this + // always-allowed entry with real per-resource authorization of resourceSubscriptions URIs. + "subscriptions/listen": {Feature: "", Operation: ""}, + // Logging and client preferences - always allowed "logging/setLevel": {Feature: "", Operation: ""}, // Client preference for server logging @@ -82,10 +112,12 @@ var MCPMethodToFeatureOperation = map[string]featureOperation{ } // shouldSkipInitialAuthorization checks if the request should skip authorization -// before reading the request body. +// before reading the request body. Content-Type is deliberately NOT consulted +// here: the middleware body refuses non-JSON POSTs with an explicit early +// return before this function is reached. func shouldSkipInitialAuthorization(r *http.Request) bool { - // Skip authorization for non-POST requests and non-JSON content types - if r.Method != http.MethodPost || !strings.HasPrefix(r.Header.Get("Content-Type"), "application/json") { + // Skip authorization for non-POST requests + if r.Method != http.MethodPost { return true } @@ -108,33 +140,39 @@ func shouldSkipSubsequentAuthorization(method string) bool { return false } -// handleUnauthorized handles unauthorized requests. +// handleUnauthorized handles unauthorized requests. The client always sees the fixed +// "Unauthorized" message -- err (an authorizer failure) can carry policy detail that +// security.md forbids returning to callers, so it is logged server-side instead. +// Cedar's evaluation context can embed decoded JWT claim values (see the claim-keys- +// only rule in authorizers/cedar/core.go), so err must not be surfaced to the client, +// nor copied into additional log lines or fields beyond the single Warn below. func handleUnauthorized(w http.ResponseWriter, msgID interface{}, err error) { + if err != nil { + slog.Warn("authorization denied", "error", err) + } + // Create a JSON-RPC error response - id, idErr := mcp.ConvertToJSONRPC2ID(msgID) - if idErr != nil { + id, convErr := mcp.ConvertToJSONRPC2ID(msgID) + if convErr != nil { id = jsonrpc2.ID{} // Use empty ID if conversion fails } - httpStatus, rpcCode, errorMsg, wwwAuthenticate := classifyAuthorizationError(err) - errorResponse := &jsonrpc2.Response{ ID: id, - Error: jsonrpc2.NewError(int64(rpcCode), errorMsg), + Error: jsonrpc2.NewError(mcp.JSONRPCCodeDenied, "Unauthorized"), } - // Set the response headers - w.Header().Set("Content-Type", "application/json") - if wwwAuthenticate != "" { - w.Header().Set("WWW-Authenticate", wwwAuthenticate) - } - w.WriteHeader(httpStatus) + // The helper encodes before writing any header, so a marshal failure never + // leaves a half-written response (e.g. a 403 header followed by a second + // 500 write). Nothing to do on a write error for a denial body. + _ = mcp.WriteJSONRPCError(w, http.StatusForbidden, errorResponse) +} - // Write the error response - if err := json.NewEncoder(w).Encode(errorResponse); err != nil { - // If we can't encode the error response, log it and return a simple error - http.Error(w, "Internal server error", http.StatusInternalServerError) - } +// rejectInvalidMCPRequest writes the 400 response for requests that arrive +// without a parsed MCP message: non-JSON POSTs refused by the middleware +// (the load-bearing security refusal) and malformed JSON POSTs. +func rejectInvalidMCPRequest(w http.ResponseWriter) { + http.Error(w, "Invalid or malformed MCP request", http.StatusBadRequest) } // Middleware creates an HTTP middleware that authorizes MCP requests. @@ -159,6 +197,22 @@ func Middleware(a authorizers.Authorizer, next http.Handler, passThroughTools ma annotationCache := NewAnnotationCache() return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Non-JSON POSTs are rejected deliberately and must never be passed + // through. Such a request is not parsed as MCP, so message-level + // authorization cannot run, but the proxy still forwards the body + // verbatim and MCP backends parse JSON-RPC without checking + // Content-Type. This early return is load-bearing for security: it is + // the only point that keeps a JSON-RPC body smuggled under text/plain + // from reaching the backend un-authorized. The marker lets the outer + // audit middleware record the refusal as a denial, not a 400 failure. + if r.Method == http.MethodPost && !mcp.RequestHasJSONContentType(r) { + if marker, ok := mcp.AuthzDenialMarkerFromContext(r.Context()); ok { + marker.Denied = true + } + rejectInvalidMCPRequest(w) + return + } + // Check if we should skip authorization before checking parsed data if shouldSkipInitialAuthorization(r) { next.ServeHTTP(w, r) @@ -168,15 +222,17 @@ func Middleware(a authorizers.Authorizer, next http.Handler, passThroughTools ma // Get parsed MCP request from context (set by parsing middleware) parsedRequest := mcp.GetParsedMCPRequest(r.Context()) if parsedRequest == nil { - // No parsed MCP request available for a request that should have been parsed - // This indicates either a malformed request or missing parsing middleware - http.Error(w, "Invalid or malformed MCP request", http.StatusBadRequest) + // Non-JSON POSTs are already rejected by the early return above, + // so a nil parsed request here means a malformed JSON body or a + // missing parsing middleware. This branch is now only a + // belt-and-braces fallback behind the content-type refusal. + rejectInvalidMCPRequest(w) return } // Streamable HTTP allows clients to POST JSON-RPC responses back to the server. // These are not authorization subjects and must reach the transport handler. - if !parsedRequest.IsRequest { + if parsedRequest.IsClientResponse() { next.ServeHTTP(w, r) return } @@ -190,13 +246,13 @@ func Middleware(a authorizers.Authorizer, next http.Handler, passThroughTools ma // Get the feature and operation from the method featureOp, ok := MCPMethodToFeatureOperation[parsedRequest.Method] if !ok { - // Unknown method - deny by default for security - // Methods must be explicitly added to MCPMethodToFeatureOperation to be allowed - handleUnauthorized(w, parsedRequest.ID, - newForbiddenAuthorizationError( - fmt.Sprintf("unknown MCP method: %s (not configured for authorization)", parsedRequest.Method), - ), - ) + // Unknown method - deny by default for security. Methods must be + // explicitly added to MCPMethodToFeatureOperation to be allowed. + // This is expected traffic (e.g. a newer-spec client), not an + // operational failure, so it's Debug rather than the Warn + // handleUnauthorized logs for an actual authorizer error. + slog.Debug("MCP method denied by default", "method", parsedRequest.Method) + handleUnauthorized(w, parsedRequest.ID, nil) return } @@ -267,7 +323,9 @@ func authorizeAndServe( // It always fully handles the request (authorization, unauthorized response, or serving). // // For pass-through meta-tools (find_tool, call_tool): -// - call_tool: authorizes the real inner tool name from arguments["tool_name"]. +// - call_tool: authorizes the real inner tool name, decoded from the request +// arguments exactly as dispatch decodes them, so the two cannot disagree about +// which tool a request names. Arguments that do not decode are denied. // - find_tool (and other pass-through tools without a tool_name): allowed through // as a discovery operation with no policy check. // @@ -283,12 +341,28 @@ func handleToolsCall( next http.Handler, ) { if _, isPassThrough := passThroughTools[parsedRequest.ResourceID]; isPassThrough { - if toolName, ok := parsedRequest.Arguments[optimizerdec.CallToolArgToolName].(string); ok && toolName != "" { + // Decode with the same call the two call_tool dispatch sites use rather than + // indexing the arguments map. encoding/json matches struct fields + // case-insensitively, so a map index on "tool_name" misses a request carrying + // "Tool_Name" that dispatch resolves and runs. Going through CallToolInput + // also applies the nested-tool_name hoist, so the name authorized here is the + // one that will execute. + input, err := schema.Translate[optimizer.CallToolInput](parsedRequest.Arguments) + if err != nil { + // The arguments are not a decodable call_tool payload, so the tool they + // target cannot be established. Deny rather than pass through; dispatch + // decodes the same map with the same call and rejects it too, so no + // legitimate invocation is lost. + slog.Warn("denying pass-through tool call with undecodable arguments", + "tool", parsedRequest.ResourceID, "error", err) + handleUnauthorized(w, parsedRequest.ID, nil) + return + } + if input.ToolName != "" { // call_tool: authorize the real backend tool name. - innerArgs, _ := parsedRequest.Arguments[optimizerdec.CallToolArgParameters].(map[string]interface{}) authorizeAndServe(w, r, a, annotationCache, featureOp.Feature, featureOp.Operation, - parsedRequest.ID, toolName, innerArgs, next) + parsedRequest.ID, input.ToolName, input.Parameters, next) return } // find_tool: allow through but filter the tools list in the response so diff --git a/pkg/authz/middleware_test.go b/pkg/authz/middleware_test.go index 8fd1d64fe6..e3bfdf6034 100644 --- a/pkg/authz/middleware_test.go +++ b/pkg/authz/middleware_test.go @@ -11,15 +11,16 @@ import ( "net/http" "net/http/httptest" "os" + "strings" "testing" "github.com/golang-jwt/jwt/v5" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" @@ -330,6 +331,28 @@ func TestMiddleware(t *testing.T) { expectStatus: http.StatusForbidden, expectAuthorized: false, }, + { + name: "Server discover is always allowed", + method: "server/discover", + params: map[string]interface{}{}, + claims: jwt.MapClaims{ + "sub": "user123", + "name": "John Doe", + }, + expectStatus: http.StatusOK, + expectAuthorized: true, + }, + { + name: "Subscriptions listen is always allowed", + method: "subscriptions/listen", + params: map[string]interface{}{}, + claims: jwt.MapClaims{ + "sub": "user123", + "name": "John Doe", + }, + expectStatus: http.StatusOK, + expectAuthorized: true, + }, { name: "Sampling createMessage is denied by default (security-sensitive)", method: "sampling/createMessage", @@ -434,6 +457,26 @@ func TestMiddleware(t *testing.T) { } } +// TestSubscriptionsListenIsAllowlistedPendingDelivery guards a deliberate, temporary +// exception: subscriptions/listen is always-allowed only because notification delivery +// for it is not yet implemented, so it exposes no data. When delivery lands, this entry +// must become a real Feature/Operation with per-resource authorization of +// resourceSubscriptions URIs (see TODO(#5755) in MCPMethodToFeatureOperation) — this test +// should fail at that point as a reminder to update it deliberately. +func TestSubscriptionsListenIsAllowlistedPendingDelivery(t *testing.T) { + t.Parallel() + require.Equal(t, featureOperation{}, MCPMethodToFeatureOperation["subscriptions/listen"]) +} + +// TestServerDiscoverIsAllowlisted guards the now-safe allow-listing of server/discover: +// its Modern envelope is post-admission capability flags (booleans), never per-resource +// descriptors, so unlike tools/list or prompts/list there is nothing here for +// ResponseFilteringWriter to filter -- always-allowed is correct, not a bypass. +func TestServerDiscoverIsAllowlisted(t *testing.T) { + t.Parallel() + require.Equal(t, featureOperation{}, MCPMethodToFeatureOperation["server/discover"]) +} + // TestMiddlewareWithGETRequest tests that the middleware doesn't panic with GET requests. func TestMiddlewareWithGETRequest(t *testing.T) { t.Parallel() @@ -471,6 +514,103 @@ func TestMiddlewareWithGETRequest(t *testing.T) { assert.Equal(t, http.StatusOK, rr.Code, "Response status code should be OK") } +func TestMiddlewareRejectsNonJSONPost(t *testing.T) { + t.Parallel() + // Even a permissive policy must not see this request: a non-JSON POST is + // never parsed as MCP, so message-level authorization cannot run, while the + // proxy would still forward the body verbatim to a backend that parses + // JSON-RPC without checking Content-Type. + authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ + Policies: []string{ + `permit(principal, action, resource);`, + }, + EntitiesJSON: `[]`, + }, "") + require.NoError(t, err, "Failed to create Cedar authorizer") + + var handlerCalled bool + handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + handlerCalled = true + w.WriteHeader(http.StatusOK) + }) + + middleware := mcpparser.ParsingMiddleware(Middleware(authorizer, handler, nil)) + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"weather","arguments":{}}}` + req, err := http.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + require.NoError(t, err, "Failed to create HTTP request") + req.Header.Set("Content-Type", "text/plain") + + rr := httptest.NewRecorder() + middleware.ServeHTTP(rr, req) + + assert.False(t, handlerCalled, "handler must not be reached by a non-JSON POST carrying JSON-RPC") + assert.Equal(t, http.StatusBadRequest, rr.Code, "non-JSON POST should be rejected") +} + +// TestMiddlewareNonJSONPostVariants pins the content-type handling around the +// explicit non-JSON refusal: which declarations still authorize and reach the +// handler, and which get the 400. +func TestMiddlewareNonJSONPostVariants(t *testing.T) { + t.Parallel() + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"weather","arguments":{}}}` + + cases := []struct { + name string + path string + contentType string // empty means the header is not set at all + wantStatus int + wantHandled bool + }{ + {"charset variant still authorizes", "/mcp", "application/json; charset=utf-8", http.StatusOK, true}, + {"uppercase media type still authorizes", "/mcp", "Application/JSON", http.StatusOK, true}, + {"missing Content-Type is rejected", "/mcp", "", http.StatusBadRequest, false}, + {"non-JSON POST to SSE path is rejected", "/sse", "text/plain", http.StatusBadRequest, false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ + Policies: []string{ + `permit(principal, action, resource);`, + }, + EntitiesJSON: `[]`, + }, "") + require.NoError(t, err, "Failed to create Cedar authorizer") + + var handlerCalled bool + handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + handlerCalled = true + w.WriteHeader(http.StatusOK) + }) + middleware := mcpparser.ParsingMiddleware(Middleware(authorizer, handler, nil)) + + req, err := http.NewRequest(http.MethodPost, tc.path, strings.NewReader(body)) + require.NoError(t, err, "Failed to create HTTP request") + if tc.contentType != "" { + req.Header.Set("Content-Type", tc.contentType) + } + + // Attach an identity so the permissive policy has a principal to + // authorize, matching the other authorized-call tests. + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "test-user", + Claims: jwt.MapClaims{"sub": "test-user"}, + }} + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + rr := httptest.NewRecorder() + middleware.ServeHTTP(rr, req) + + assert.Equal(t, tc.wantHandled, handlerCalled, "handler reached") + assert.Equal(t, tc.wantStatus, rr.Code, "response status") + }) + } +} + func TestMiddlewareAllowsJSONRPCResponsesThrough(t *testing.T) { t.Parallel() @@ -1061,15 +1201,21 @@ func TestMiddlewareToolsCallTestkit(t *testing.T) { // TestMiddlewareOptimizerMetaTools tests the optimizer meta-tool interception logic. // When a tool is in the passThroughTools set, the middleware handles it specially: -// - call_tool (has "tool_name" in arguments): authorize the inner backend tool -// - find_tool (no "tool_name" in arguments): allow through as a discovery operation +// - call_tool (arguments decode to a tool_name): authorize the inner backend tool +// - find_tool (arguments name no tool): allow through as a discovery operation +// +// The arguments are decoded exactly as dispatch decodes them, so the name authorized +// here and the name that runs are always the same one. func TestMiddlewareOptimizerMetaTools(t *testing.T) { t.Parallel() - // Cedar policy that only permits "allowed_backend" — not "call_tool" or "find_tool". + // Cedar policies permitting "allowed_backend" unconditionally and "args_backend" + // only when the inner arguments reach the authorizer. Neither permits "call_tool" + // or "find_tool" themselves. authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ Policies: []string{ `permit(principal, action == Action::"call_tool", resource == Tool::"allowed_backend");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"args_backend") when { context.arg_query == "x" };`, }, EntitiesJSON: `[]`, }, "") @@ -1130,6 +1276,71 @@ func TestMiddlewareOptimizerMetaTools(t *testing.T) { expectStatus: http.StatusForbidden, expectHandlerHit: false, }, + { + // Dispatch hoists a nested name and runs the tool, so authorization must + // see the same target rather than treating the request as naming no tool. + name: "call_tool with nested unauthorized inner tool is blocked", + toolName: "call_tool", + arguments: map[string]interface{}{ + "parameters": map[string]interface{}{ + "tool_name": "forbidden_backend", + "query": "x", + }, + }, + expectStatus: http.StatusForbidden, + expectHandlerHit: false, + }, + { + name: "call_tool with nested authorized inner tool passes through", + toolName: "call_tool", + arguments: map[string]interface{}{ + "parameters": map[string]interface{}{ + "tool_name": "allowed_backend", + "query": "x", + }, + }, + expectStatus: http.StatusOK, + expectHandlerHit: true, + }, + { + // encoding/json matches struct fields case-insensitively, so dispatch + // resolves and runs this tool. A map index on "tool_name" does not see it + // and would wave the request through with no policy check at all. + name: "call_tool with a case-variant tool_name key is blocked", + toolName: "call_tool", + arguments: map[string]interface{}{ + "Tool_Name": "forbidden_backend", + "parameters": map[string]interface{}{}, + }, + expectStatus: http.StatusForbidden, + expectHandlerHit: false, + }, + { + // The policy for args_backend only permits the call when the inner + // arguments reach the authorizer. A map index on "parameters" drops them + // here, so the tool would be denied a call the policy allows. + name: "call_tool with a case-variant parameters key still carries the inner arguments", + toolName: "call_tool", + arguments: map[string]interface{}{ + "tool_name": "args_backend", + "PARAMETERS": map[string]interface{}{"query": "x"}, + }, + expectStatus: http.StatusOK, + expectHandlerHit: true, + }, + { + // The target cannot be established, so the middleware refuses rather than + // passing through. Dispatch decodes the same arguments with the same call + // and rejects them too, so no legitimate invocation is lost. + name: "call_tool with undecodable arguments is denied", + toolName: "call_tool", + arguments: map[string]interface{}{ + "tool_name": "allowed_backend", + "parameters": "not an object", + }, + expectStatus: http.StatusForbidden, + expectHandlerHit: false, + }, { name: "find_tool request reaches handler (response filtering applied separately)", toolName: "find_tool", @@ -1226,6 +1437,126 @@ func TestMiddlewareOptimizerCallToolJSONRoundTrip(t *testing.T) { assert.True(t, handlerCalled, "handler should be called for authorized call_tool") } +// TestHandleUnauthorizedIsConformantJSONRPC drives the real middleware chain (parsing + +// authorization) for a denied request and asserts the 403 body is a spec-conformant +// JSON-RPC error response: lowercase "jsonrpc"/"id"/"error" keys and, for notifications, +// no "id" key at all. assert.JSONEq is case-sensitive, so it alone would catch "ID"/"Error" +// substituted for "id"/"error" and would fail on an unexpected "id" key — the id-presence +// check is asserted separately anyway, as a belt-and-braces, self-documenting assertion. +func TestHandleUnauthorizedIsConformantJSONRPC(t *testing.T) { + t.Parallel() + + // The policy content is irrelevant: tasks/list is unknown to + // MCPMethodToFeatureOperation and denied before any policy is consulted. If a + // future PR adds authorization for tasks/* and this test starts failing, swap + // in any other method still absent from that map rather than reading it as an + // envelope regression. + authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ + Policies: []string{`permit(principal, action == Action::"call_tool", resource == Tool::"weather");`}, + EntitiesJSON: `[]`, + }, "") + require.NoError(t, err) + + handler := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("next handler must not be called for a denied request") + }) + middleware := mcpparser.ParsingMiddleware(Middleware(authorizer, handler, nil)) + + const wantErr = `"error":{"code":403,"message":"Unauthorized"}` + + testCases := []struct { + name string + reqIDJSON string // raw "id" field in the request; "" omits it (notification) + wantBody string + }{ + { + name: "int64 id", + reqIDJSON: `"id":42,`, + wantBody: `{"jsonrpc":"2.0","id":42,` + wantErr + `}`, + }, + { + name: "string id", + reqIDJSON: `"id":"abc",`, + wantBody: `{"jsonrpc":"2.0","id":"abc",` + wantErr + `}`, + }, + { + name: "nil id (notification) omits id entirely", + reqIDJSON: "", + wantBody: `{"jsonrpc":"2.0",` + wantErr + `}`, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + reqBody := `{"jsonrpc":"2.0",` + tc.reqIDJSON + `"method":"tasks/list"}` + req, err := http.NewRequest(http.MethodPost, "/messages", bytes.NewBufferString(reqBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "test-user", + Claims: jwt.MapClaims{"sub": "test-user"}, + }} + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + rr := httptest.NewRecorder() + middleware.ServeHTTP(rr, req) + + require.Equal(t, http.StatusForbidden, rr.Code) + assert.Equal(t, "application/json", rr.Header().Get("Content-Type")) + assert.JSONEq(t, tc.wantBody, rr.Body.String()) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &decoded)) + _, hasID := decoded["id"] + assert.Equal(t, tc.reqIDJSON != "", hasID, "id key presence mismatch") + _, hasResult := decoded["result"] + assert.False(t, hasResult, "denied response must not contain a result key") + }) + } +} + +// TestHandleUnauthorizedScrubsAuthorizerError drives the real middleware chain +// (parsing + authorization) for a denial caused solely by the authorizer erroring +// while reporting allowed=true (as opposed to a clean policy deny), and asserts the +// 403 body carries the fixed "Unauthorized" message rather than the authorizer's +// error text. allowed is deliberately true here: authorizeAndServe's condition is +// `err != nil || !authorized`, so with allowed=false the test would still pass via +// the !authorized disjunct even if err were dropped somewhere -- this fixture pins +// the err-only branch specifically. +func TestHandleUnauthorizedScrubsAuthorizerError(t *testing.T) { + t.Parallel() + + sensitiveErr := errors.New("cedar entity store lookup failed: secret-backend-url unreachable") + stub := &stubAuthorizer{allowed: true, err: sensitiveErr} + + handler := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("next handler must not be called for a denied request") + }) + middleware := mcpparser.ParsingMiddleware(Middleware(stub, handler, nil)) + + reqBody := `{"jsonrpc":"2.0","id":1,"method":"resources/read","params":{"uri":"data"}}` + req, err := http.NewRequest(http.MethodPost, "/messages", bytes.NewBufferString(reqBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "test-user", + Claims: jwt.MapClaims{"sub": "test-user"}, + }} + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + rr := httptest.NewRecorder() + middleware.ServeHTTP(rr, req) + + require.Equal(t, http.StatusForbidden, rr.Code) + assert.JSONEq(t, `{"jsonrpc":"2.0","id":1,"error":{"code":403,"message":"Unauthorized"}}`, rr.Body.String()) + assert.NotContains(t, rr.Body.String(), "secret-backend-url", + "authorizer error text must not leak to the client") +} + // TestConvertToJSONRPC2ID tests the convertToJSONRPC2ID function with various ID types func TestConvertToJSONRPC2ID(t *testing.T) { t.Parallel() @@ -1313,7 +1644,6 @@ func TestAuthorizeAndServe(t *testing.T) { expectHandlerHit bool expectStatus int expectAnnotations bool // whether annotations should be in handler context - expectChallenge string }{ { name: "authorized with cache miss — next called, no annotations", @@ -1336,19 +1666,23 @@ func TestAuthorizeAndServe(t *testing.T) { expectStatus: http.StatusForbidden, }, { - name: "authorizer error — 500, next not called", + name: "authorizer error — 403, next not called", allowed: false, authErr: errors.New("policy evaluation failed"), expectHandlerHit: false, - expectStatus: http.StatusInternalServerError, + expectStatus: http.StatusForbidden, }, { - name: "upstream auth required error — 401, next not called", - allowed: false, - authErr: authorizers.NewUpstreamAuthenticationRequiredError("cognito", "refresh_failed"), + // allowed=true here discriminates fail-closed-on-error from + // fail-closed-on-deny: with allowed=false above, the 403 could come + // from either disjunct of `err != nil || !authorized`, so it alone + // doesn't pin that an authorizer error fails closed even when it + // also happens to report allowed. + name: "authorizer error with allowed=true — still 403, next not called", + allowed: true, + authErr: errors.New("policy evaluation failed"), expectHandlerHit: false, - expectStatus: http.StatusUnauthorized, - expectChallenge: `Bearer error="invalid_token", error_description="upstream token is no longer valid; re-authentication required"`, + expectStatus: http.StatusForbidden, }, } @@ -1381,7 +1715,6 @@ func TestAuthorizeAndServe(t *testing.T) { assert.Equal(t, tc.expectHandlerHit, handlerCalled) assert.Equal(t, tc.expectStatus, rr.Code) - assert.Equal(t, tc.expectChallenge, rr.Header().Get("WWW-Authenticate")) if tc.expectAnnotations { ann := authorizers.ToolAnnotationsFromContext(ctxInHandler) require.NotNil(t, ann) diff --git a/pkg/authz/response_filter.go b/pkg/authz/response_filter.go index 736b80b17c..027a38261a 100644 --- a/pkg/authz/response_filter.go +++ b/pkg/authz/response_filter.go @@ -13,10 +13,11 @@ import ( "net/http" "strings" - "github.com/mark3labs/mcp-go/mcp" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" ) @@ -69,8 +70,15 @@ func (rfw *ResponseFilteringWriter) WriteHeader(statusCode int) { // FlushAndFilter processes the captured response and applies filtering if needed. // Returns an error if filtering or writing fails. func (rfw *ResponseFilteringWriter) FlushAndFilter() error { - // If it's not a successful response, just pass it through - if rfw.statusCode != http.StatusOK && rfw.statusCode != http.StatusAccepted { + // Only successful responses can deliver a list result to a client, so + // non-2xx responses pass through unfiltered: an error body isn't a list, + // and rewriting it would only hurt debuggability. This deliberately + // covers the whole 2xx range, not just 200/202: fetch-based MCP clients + // (including the reference TypeScript transport) gate on response.ok, + // which accepts 200-299, so a backend answering tools/list with e.g. 201 + // could otherwise smuggle an unfiltered list past the filter. A 204 has + // no body and is passed through by the empty-response check below. + if rfw.statusCode < http.StatusOK || rfw.statusCode >= http.StatusMultipleChoices { rfw.ResponseWriter.WriteHeader(rfw.statusCode) _, err := rfw.ResponseWriter.Write(rfw.buffer.Bytes()) //nolint:gosec // G705 - JSON-RPC response, not rendered as HTML return err @@ -92,7 +100,13 @@ func (rfw *ResponseFilteringWriter) FlushAndFilter() error { return err } - mimeType := strings.Split(rfw.ResponseWriter.Header().Get("Content-Type"), ";")[0] + // Media type names are case-insensitive and parameters may be preceded by + // whitespace (RFC 9110 section 8.3.1), so normalize before matching. + // Without this, "Application/JSON" or "application/json ; charset=utf-8" + // -- both valid labels for a JSON body that clients happily parse -- would + // fall through to the default branch and skip filtering entirely. + contentType := rfw.ResponseWriter.Header().Get("Content-Type") + mimeType := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) switch mimeType { case "application/json": @@ -107,6 +121,28 @@ func (rfw *ResponseFilteringWriter) FlushAndFilter() error { rfw.ResponseWriter.Header().Del("Content-Length") return rfw.processSSEResponse(rawResponse) default: + // A successful response to a method whose result must be filtered, + // yet labeled with neither MCP-supported media type. That could be an + // accident, or a backend deliberately mislabeling the response to + // smuggle an unfiltered list past the filter -- the same + // disguised-result concern as #5257, handled for the tool filter in + // pkg/mcp/tool_filter.go's processUnrecognizedMimeType. Sniff the + // body under both supported shapes and, when it carries a JSON-RPC + // result, process it exactly as if it had been labeled correctly + // (which filters clean frames and fails closed on dirty ones). Only a + // body carrying no result under either shape passes through. + if carriesResult(rawResponse) { + slog.Warn("response with unrecognized media type carries a JSON-RPC result; filtering as JSON", + "method", rfw.method, "contentType", contentType) + rfw.ResponseWriter.Header().Del("Content-Length") + return rfw.processJSONResponse(rawResponse) + } + if sseCarriesResult(rawResponse) { + slog.Warn("response with unrecognized media type carries a JSON-RPC result; filtering as SSE", + "method", rfw.method, "contentType", contentType) + rfw.ResponseWriter.Header().Del("Content-Length") + return rfw.processSSEResponse(rawResponse) + } rfw.ResponseWriter.WriteHeader(rfw.statusCode) _, err := rfw.ResponseWriter.Write(rawResponse) return err @@ -140,20 +176,22 @@ func (rfw *ResponseFilteringWriter) processJSONResponse(rawResponse []byte) erro if carriesResult(rawResponse) { slog.Warn("JSON response carried a result outside a clean Response frame; dropping as a protocol violation", "method", rfw.method) - return rfw.writeErrorResponse(jsonrpc2.ID{}, - fmt.Errorf("dropped a frame carrying a result outside a clean Response")) + // A batch is replaced wholesale by this single error envelope + // even if only one of many elements carried the smuggled + // result: JSON-RPC does define batch responses, but we don't + // attempt a partial one here, and mcpparser.ParsingMiddleware + // already rejects batch requests outright (see #5745), so a + // batch response on this path is already off-spec traffic. + // Fail-closed is the right side to err on regardless. + return rfw.writeErrorResponse(rfw.requestID(), + errors.New("dropped a frame carrying a result outside a clean Response")) } rfw.ResponseWriter.WriteHeader(rfw.statusCode) _, werr := rfw.ResponseWriter.Write(rawResponse) return werr } - filteredResponse, err := rfw.filterListResponse(response) - if err != nil { - return rfw.writeErrorResponse(response.ID, err) - } - - filteredData, err := jsonrpc2.EncodeMessage(filteredResponse) + filteredData, err := rfw.filterAndEncode(response) if err != nil { return rfw.writeErrorResponse(response.ID, err) } @@ -163,58 +201,100 @@ func (rfw *ResponseFilteringWriter) processJSONResponse(rawResponse []byte) erro return err } +// filterAndEncode filters a Response and encodes the result for the wire. +func (rfw *ResponseFilteringWriter) filterAndEncode(response *jsonrpc2.Response) ([]byte, error) { + filteredResponse, err := rfw.filterListResponse(response) + if err != nil { + return nil, err + } + return jsonrpc2.EncodeMessage(filteredResponse) +} + +// sseLine pairs one scanned SSE line with the terminator that followed it in +// the raw body: "\n", or nil for a final line with none. Lines are split on +// LF only (see processSSEResponse), so a line that was actually terminated by +// "\r\n" keeps its trailing "\r" as the last byte of text -- text+term always +// reproduces exactly the bytes that were scanned. +type sseLine struct { + text []byte + term []byte +} + +// processSSEResponse rewrites a text/event-stream body per SSE framing rules +// rather than treating it as an arbitrary sequence of lines: within one +// event, all `data:` field values concatenate (joined by LF, per spec) into a +// single payload, and a blank line is structural — it dispatches the event +// rather than being emittable content in its own right. +// +// Lines are split on LF only, matching the client that actually consumes +// these streams (github.com/modelcontextprotocol/go-sdk's mcp/event.go, +// which ReadBytes('\n')s and never treats a lone CR as a line terminator) -- +// not the WHATWG EventSource grammar, which additionally splits on a bare CR. +// The two disagree on a body containing an interior "\r" (e.g. inside a +// data: payload): WHATWG would cut the line there, but the client reads it as +// one line and parses through. Following the WHATWG grammar here decoded a +// SUPERSET of what the client would treat as one line, letting the shorter, +// wrongly-split pieces slip past filtering undetected -- a real regression, +// not a hypothetical one. Splitting on LF only means we parse a superset of +// what a stricter reader would call one line, so we decode and filter *more* +// than before, never less: worst case for a client that does honor a lone CR +// (e.g. a browser EventSource) is a payload it mis-frames identically to how +// it would mis-frame the backend's raw bytes directly, since output stays +// byte-preserving -- denial, never disclosure. +// +// Every line is re-emitted with the same terminator it was scanned with, so +// the output is structurally identical to the input (minus a leading UTF-8 +// BOM, stripped below); only `data:` payloads are rewritten. func (rfw *ResponseFilteringWriter) processSSEResponse(rawResponse []byte) error { // Note: this routine is adapted from the one in pkg/mcp/tool_filter.go. // I don't see an obvious way to factor out the commonalities, so I'm // duplicating it here, but we should refactor response parsing // respecting mime types to a common routine. - var linesep []byte - if bytes.Contains(rawResponse, []byte("\r\n")) { - linesep = []byte("\r\n") - } else if bytes.Contains(rawResponse, []byte("\n")) { - linesep = []byte("\n") - } else if bytes.Contains(rawResponse, []byte("\r")) { - linesep = []byte("\r") - } else { - return fmt.Errorf("unsupported separator: %s", string(rawResponse)) - } - - var linesepTotal, linesepCount int - linesepTotal = bytes.Count(rawResponse, linesep) - lines := bytes.Split(rawResponse, linesep) - for _, line := range lines { - if len(line) == 0 { - continue - } - var written bool - if data, ok := bytes.CutPrefix(line, []byte("data:")); ok { - w, stop, err := rfw.writeSSEDataLine(data) - if stop { - return err - } - written = w + // A client strips a leading BOM per the WHATWG UTF-8 decode algorithm + // before parsing lines, so strip it here too: otherwise the first line's + // "data:" prefix wouldn't match and the event would pass through + // unfiltered. + rawResponse = bytes.TrimPrefix(rawResponse, []byte("\xEF\xBB\xBF")) + + var outputLines []sseLine + var event []sseLine + for len(rawResponse) > 0 { + idx := bytes.IndexByte(rawResponse, '\n') + var line, term []byte + if idx == -1 { + line, term, rawResponse = rawResponse, nil, nil + } else { + line, term, rawResponse = rawResponse[:idx], rawResponse[idx:idx+1], rawResponse[idx+1:] } - if !written { - _, err := rfw.ResponseWriter.Write(line) - if err != nil { - return fmt.Errorf("%w: %w", errBug, err) - } + // A line is structurally blank once its trailing CR (left over from + // a "\r\n" terminator the LF-only split doesn't consume) is ignored; + // without this a CRLF-terminated blank line would scan as the + // single byte "\r" and never be recognised as the event separator. + // The blank line's own bytes (that "\r", if present) are preserved + // in the output below rather than dropped, so structure stays + // byte-identical to the input. + if len(bytes.TrimRight(line, "\r")) == 0 { + outputLines = append(outputLines, rfw.resolveSSEEvent(event)...) + outputLines = append(outputLines, sseLine{text: line, term: term}) + event = nil + continue } + event = append(event, sseLine{text: line, term: term}) + } + // A trailing event with no closing blank line: the go-sdk client (see + // mcp/event.go's yieldEvent, called once more after its read loop hits + // EOF) does still dispatch whatever is buffered at that point, so this + // is filtered like any other event, not dropped. We write these bytes + // (filtered) without fabricating a terminator the backend didn't send. + outputLines = append(outputLines, rfw.resolveSSEEvent(event)...) - _, err := rfw.ResponseWriter.Write(linesep) - if err != nil { + for _, l := range outputLines { + if _, err := rfw.ResponseWriter.Write(l.text); err != nil { return fmt.Errorf("%w: %w", errBug, err) } - linesepCount++ - } - - // This ensures we don't send too few line separators, which might break - // SSE parsing. - if linesepCount < linesepTotal { - _, err := rfw.ResponseWriter.Write(linesep) - if err != nil { + if _, err := rfw.ResponseWriter.Write(l.term); err != nil { return fmt.Errorf("%w: %w", errBug, err) } } @@ -222,53 +302,198 @@ func (rfw *ResponseFilteringWriter) processSSEResponse(rawResponse []byte) error return nil } -// writeSSEDataLine classifies a single SSE `data:` payload and writes the -// filtered response when it is a clean Response. written reports whether the -// caller should skip the raw passthrough for this line. stop reports that -// processSSEResponse should return err immediately (an error response has -// already been written on the filter/encode error paths). -func (rfw *ResponseFilteringWriter) writeSSEDataLine(data []byte) (written, stop bool, err error) { +// resolveSSEEvent assembles the `data:` payload of one SSE event (the lines +// between two blank lines, exclusive) and returns the lines to emit in its +// place. An event with no `data:` field (e.g. a bare "event: ping") or whose +// assembled payload passes filterSSEEventData unchanged is returned as-is. +func (rfw *ResponseFilteringWriter) resolveSSEEvent(event []sseLine) []sseLine { + if len(event) == 0 { + return nil + } + + dataValues := eventDataValues(event) + if len(dataValues) == 0 { + return event + } + + // No whole-buffer trim needed here: the per-value trim in eventDataValues + // already leaves nothing but JSON-significant bytes at either edge of the join. + assembled := bytes.Join(dataValues, []byte("\n")) + if len(assembled) == 0 { + // A client never dispatches an empty data buffer, so there's nothing + // to classify or fail closed on. + return event + } + + replacement, failedClosed := rfw.filterSSEEventData(assembled) + if replacement == nil { + replacement, failedClosed = rfw.probeValuesForSmuggledResult(dataValues) + } + if replacement == nil { + return event + } + return rebuildEventWithPayload(event, replacement, failedClosed) +} + +// eventDataValues extracts one event's `data:` field values, in order, each +// trimmed exactly as the client trims it. +func eventDataValues(event []sseLine) [][]byte { + var dataValues [][]byte + for _, l := range event { + // A line ending in "\r\n" keeps that "\r" as the trailing byte of + // l.text (see sseLine), so trim it before matching the "data:" + // prefix -- mirroring the client, which strips trailing "\r\n" from + // every line before inspecting it (mcp/event.go). + trimmed := bytes.TrimRight(l.text, "\r") + // A bare "data" line with no colon, or one with a space before the + // colon ("data :foo"), doesn't match the "data:" prefix below, so both + // are excluded from assembly here, unlike strict WHATWG grammar (which + // treats a bare "data" as a field with an empty value). That's safe + // for two reasons: (1) the only effect on a spec-compliant client's + // assembled buffer is one extra LF, which is JSON-insignificant + // whitespace outside a string literal and can never turn JSON we'd + // reject into JSON a client would accept; (2) for + // "data :foo", bytes.Cut on the raw line yields before == "data " (not + // "data"), so the go-sdk's own line parser ignores it too, exactly + // like us and strict WHATWG. + data, ok := bytes.CutPrefix(trimmed, []byte("data:")) + if !ok { + continue + } + // Per the WHATWG EventSource grammar, exactly one leading space after + // "data:" is part of the field delimiter, not the payload. Stripping + // more (or a differently-positioned space) would corrupt a payload + // split mid-string-literal across lines. + data, _ = bytes.CutPrefix(data, []byte(" ")) + // TrimSpace each value independently, before joining, not on the + // assembled buffer afterward. Go's JSON scanner treats only SP, TAB, + // CR and LF as whitespace, but unicode.IsSpace also covers U+000B, + // U+000C, U+0085, U+00A0, U+1680 and U+3000 -- and the go-sdk client + // trims each data value with exactly this function before joining + // (mcp/event.go), not the whole buffer once at the end. A + // whole-buffer trim only catches one of those bytes sitting at the + // very front or back of the assembled payload; one sitting at an + // interior boundary between two data: lines of the SAME event + // survived undetected, because it was never at either edge of the + // joined buffer. Trimming here catches it regardless of position. + data = bytes.TrimSpace(data) + dataValues = append(dataValues, data) + } + return dataValues +} + +// probeValuesForSmuggledResult is the fallback for an event whose assembled +// payload was undecodable as a whole. filterSSEEventData's carriesResult scan +// stops at the first undecodable JSON value in the joined payload, so a +// well-formed result-bearing value sitting right after a malformed one within +// the SAME event would otherwise ride through as "undecodable, nothing to +// filter" -- reopening the split-payload bypass this rewrite exists to close +// (#5257), just one line lower than before. Probing each value independently +// only ever turns a pass-through into a fail-closed error envelope, never the +// reverse, so it cannot itself reintroduce a bypass. A nil replacement means +// no value carried a result and the event genuinely needs no filtering. +func (rfw *ResponseFilteringWriter) probeValuesForSmuggledResult(dataValues [][]byte) (replacement []byte, failedClosed bool) { + for _, v := range dataValues { + if carriesResult(v) { + slog.Warn("SSE event's assembled payload was undecodable but an individual data value carries a result; failing closed", + "method", rfw.method) + return rfw.errorResponseBody(rfw.requestID(), + errors.New("dropped a frame carrying a result outside a clean Response")), true + } + } + return nil, false +} + +// rebuildEventWithPayload re-emits one event's lines with replacement standing +// in for the first `data:` line's value, dropping the later `data:` lines whose +// content is already folded into it. +func rebuildEventWithPayload(event []sseLine, replacement []byte, failedClosed bool) []sseLine { + out := make([]sseLine, 0, len(event)) + replaced := false + for _, l := range event { + trimmed := bytes.TrimRight(l.text, "\r") + // The go-sdk streamable client only dispatches unnamed ("message") + // events, so a named event's own "event:" field would silently + // swallow the fail-closed envelope substituted below, reproducing + // the #6037 hang. Drop it: the substituted payload is either the + // error envelope itself or a freshly filtered result, neither of + // which the original event name describes any more. + if failedClosed && bytes.HasPrefix(trimmed, []byte("event:")) { + continue + } + if !bytes.HasPrefix(trimmed, []byte("data:")) { + out = append(out, l) + continue + } + if replaced { + // Subsequent data: lines are already folded into the assembled + // payload; only the first line carries it. + continue + } + // A "\r\n"-terminated line keeps that "\r" as l.text's trailing + // byte (see sseLine), not as part of l.term. Since the replacement + // text discards l.text entirely, re-attach any such trailing "\r" to + // the terminator here -- otherwise a CRLF-terminated data line would + // silently downgrade to bare LF once filtered, while an untouched + // (pass-through) line keeps its CRLF because it reuses l.text verbatim. + term := l.term + if cr := l.text[len(trimmed):]; len(cr) > 0 { + term = append(append([]byte{}, cr...), term...) + } + // jsonrpc2.EncodeMessage and json.Marshal both escape newlines, so + // replacement can never contain a raw line separator and is always + // safe to emit as a single data: line. + out = append(out, sseLine{text: append([]byte("data: "), replacement...), term: term}) + replaced = true + } + return out +} + +// filterSSEEventData classifies an event's assembled data payload and returns +// the payload to emit in its place, plus whether that payload is a fail-closed +// error envelope (as opposed to a successfully filtered result). A nil +// replacement means the event's lines pass through unchanged, and failedClosed +// is meaningless in that case. +func (rfw *ResponseFilteringWriter) filterSSEEventData(data []byte) (replacement []byte, failedClosed bool) { message, decodeErr := jsonrpc2.DecodeMessage(data) response, isResponse := message.(*jsonrpc2.Response) switch { case isResponse: - filteredResponse, ferr := rfw.filterListResponse(response) + filteredData, ferr := rfw.filterAndEncode(response) if ferr != nil { - return false, true, rfw.writeErrorResponse(response.ID, ferr) - } - filteredData, ferr := jsonrpc2.EncodeMessage(filteredResponse) - if ferr != nil { - return false, true, rfw.writeErrorResponse(response.ID, ferr) - } - // The caller writes linesep after this line, so do not append a - // terminator here. A hardcoded "\n" desyncs \r\n streams. - if _, werr := rfw.ResponseWriter.Write([]byte("data: " + string(filteredData))); werr != nil { - return false, true, fmt.Errorf("%w: %w", errBug, werr) + // errorResponseBody logs ferr itself, so don't repeat it here. + slog.Warn("emitting a JSON-RPC error in place of an SSE frame that failed to filter", + "method", rfw.method) + return rfw.errorResponseBody(response.ID, ferr), true } - return true, false, nil + return filteredData, false case carriesResult(data): // The frame is not a clean Response but still carries a result. This // covers a non-Response type (a request/notification frame smuggling a // result), a decode error (missing or invalid jsonrpc tag, a response // frame with no or a non-scalar id), and a batch array. All are - // upstream-controlled shapes that would otherwise fall through and leak - // the very list this filter scrubs (#5257). Fail closed and drop the line. - slog.Warn("SSE data line carried a result outside a clean Response frame; dropping as a protocol violation", + // upstream-controlled shapes that would otherwise leak the very list + // this filter scrubs (#5257). Fail closed with an explicit error + // envelope rather than dropping the event silently: an event + // dispatched with an empty data buffer isn't delivered to the + // client's handler at all, leaving it hanging on its own timeout + // (#6037). + slog.Warn("SSE event carried a result outside a clean Response frame; failing closed as a protocol violation", "method", rfw.method) - return true, false, nil + return rfw.errorResponseBody(rfw.requestID(), + errors.New("dropped a frame carrying a result outside a clean Response")), true case decodeErr != nil: - // Genuinely undecodable and no smuggled result. Pass this line through - // unfiltered (this line only). - slog.Warn("SSE data line could not be decoded as JSON-RPC; passing through unfiltered", + // Genuinely undecodable and no smuggled result. Pass through unfiltered. + slog.Warn("SSE event data could not be decoded as JSON-RPC; passing through unfiltered", "method", rfw.method, "error", decodeErr) - return false, false, nil + return nil, false default: - // Genuine non-Response frame (e.g. an interleaved notifications/* message) - // with no result payload. Routine SSE traffic, so log at Debug to keep the - // suspicious branches above from being buried. Pass through this line only. - slog.Debug("SSE data line was not a JSON-RPC Response; passing through unfiltered", + // Genuine non-Response frame (e.g. an interleaved notifications/* + // message) with no result payload. Routine SSE traffic, so log at + // Debug to keep the suspicious branches above from being buried. + slog.Debug("SSE event data was not a JSON-RPC Response; passing through unfiltered", "method", rfw.method) - return false, false, nil + return nil, false } } @@ -288,26 +513,55 @@ func requiresResponseFiltering(method string) bool { // result, a shape DecodeMessage rejects, or a batch array) but still carries a // result must not pass through: it would leak the list the filter exists to // scrub. See issue #5257. +// +// The payload may hold more than one top-level JSON value: an SSE event's +// assembled data can concatenate several messages (e.g. a notification and a +// response, one per data: line) into something DecodeMessage rejects as a +// single value. json.Decoder reads one JSON value at a time and stops at the +// first one it can't decode, rather than Unmarshal-ing the whole payload as +// one document, so a result-bearing value after the first is still caught +// instead of making the whole payload look undecodable. Stopping there +// (rather than trying to resync past the bad value) is safe: a payload the +// decoder can't parse value-by-value isn't a single parseable JSON document +// either, so no strict client can read past that point. func carriesResult(data []byte) bool { - trimmed := bytes.TrimSpace(data) + dec := json.NewDecoder(bytes.NewReader(data)) + for { + var value json.RawMessage + if dec.Decode(&value) != nil { + return false // EOF, or malformed JSON with nothing left to check + } + if valueCarriesResult(value) { + return true + } + } +} + +// valueCarriesResult applies carriesResult's check to a single JSON value. +func valueCarriesResult(value json.RawMessage) bool { + trimmed := bytes.TrimSpace(value) if len(trimmed) == 0 { return false } switch trimmed[0] { case '{': - var probe struct { - Result json.RawMessage `json:"result"` - } - return json.Unmarshal(trimmed, &probe) == nil && probe.Result != nil + return objectCarriesResult(trimmed) case '[': - var batch []struct { - Result json.RawMessage `json:"result"` - } + // A single level, not recursive: a JSON-RPC batch is a flat array of + // message objects, so that's the whole legitimate surface. Recursing + // through nested arrays is O(d^2) on encoding/json's own nesting cap + // (10000), an easy amplification handle for something no client + // actually flattens. + var batch []json.RawMessage if json.Unmarshal(trimmed, &batch) != nil { return false } - for i := range batch { - if batch[i].Result != nil { + for _, el := range batch { + elTrimmed := bytes.TrimSpace(el) + if len(elTrimmed) == 0 || elTrimmed[0] != '{' { + continue + } + if objectCarriesResult(elTrimmed) { return true } } @@ -315,10 +569,54 @@ func carriesResult(data []byte) bool { return false } +// sseCarriesResult reports whether rawResponse contains an SSE "data:" line +// whose payload carries a JSON-RPC result. It is a lightweight detector, used +// only to decide whether a 2xx body with an unrecognized media type needs the +// full SSE processing path (which applies its own event-based filtering and +// fail-closed rules); it mirrors sniffSSEToolsList in pkg/mcp/tool_filter.go. +func sseCarriesResult(rawResponse []byte) bool { + normalized := bytes.ReplaceAll(rawResponse, []byte("\r\n"), []byte("\n")) + normalized = bytes.ReplaceAll(normalized, []byte("\r"), []byte("\n")) + for _, line := range bytes.Split(normalized, []byte("\n")) { + if data, ok := bytes.CutPrefix(line, []byte("data:")); ok && carriesResult(data) { + return true + } + } + return false +} + +// objectCarriesResult reports whether a trimmed JSON object value has a +// "result" key present, including an explicit `"result":null` — RawMessage's +// UnmarshalJSON runs even for a JSON null, so probe.Result is non-nil (the +// 4 bytes "null") whenever the key exists at all. Treating a null result as +// carrying one is deliberate: it's the fail-closed direction. +func objectCarriesResult(trimmedObject json.RawMessage) bool { + var probe struct { + Result json.RawMessage `json:"result"` + } + return json.Unmarshal(trimmedObject, &probe) == nil && probe.Result != nil +} + // filterListResponse filters the list response based on authorization policies func (rfw *ResponseFilteringWriter) filterListResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) { if response.Error != nil { - // If there's an error in the response, don't filter + // A Response carrying both error and result is a protocol violation: + // DecodeMessage populates both fields as given, and some clients + // (e.g. the reference TS SDK's zod schema) silently strip the + // unexpected "error" key and treat it as a successful result. Fail + // closed rather than passing the smuggled list through under cover + // of the error field. See #5257. + // A literal `"result":null` is exempted: DecodeMessage sets Result to + // the non-nil 4 bytes "null" whenever the key is present at all (see + // objectCarriesResult), so without this a legitimate upstream error + // that happens to carry an explicit null result would lose its own + // code (e.g. -32601) to our generic internal-error envelope for no + // security benefit -- a null result can never carry a list, so + // failing closed on it buys nothing. + if response.Result != nil && !bytes.Equal(bytes.TrimSpace(response.Result), []byte("null")) { + return nil, errors.New("response carried both error and result") + } + // If there's an error and no result (or an explicit null result), don't filter return response, nil } @@ -377,10 +675,7 @@ func (rfw *ResponseFilteringWriter) filterToolsResponse(response *jsonrpc2.Respo // filterToolsByPolicy checks each tool against the caller's Cedar policies // (injecting annotations into context for when-clause evaluation) and returns // only tools the caller is authorized to call. - policyFiltered, err := filterToolsByPolicy(rfw.request.Context(), rfw.authorizer, regular) - if err != nil { - return nil, err - } + policyFiltered := filterToolsByPolicy(rfw.request.Context(), rfw.authorizer, regular) filteredTools := make([]mcp.Tool, 0, len(passThrough)+len(policyFiltered)) filteredTools = append(filteredTools, passThrough...) filteredTools = append(filteredTools, policyFiltered...) @@ -408,10 +703,6 @@ func (rfw *ResponseFilteringWriter) filterToolsResponse(response *jsonrpc2.Respo // filterPromptsResponse filters prompts based on get_prompt authorization func (rfw *ResponseFilteringWriter) filterPromptsResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) { - if rfw.authorizer == nil { - return response, nil - } - // Parse the result as a ListPromptsResult var listResult mcp.ListPromptsResult if err := json.Unmarshal(response.Result, &listResult); err != nil { @@ -432,7 +723,9 @@ func (rfw *ResponseFilteringWriter) filterPromptsResponse(response *jsonrpc2.Res nil, // No arguments for the authorization check ) if err != nil { - return nil, err + slog.Warn("Authorization check failed for prompt, skipping", + "prompt", prompt.Name, "error", err) + continue } if authorized { @@ -471,10 +764,6 @@ func (rfw *ResponseFilteringWriter) filterPromptsResponse(response *jsonrpc2.Res // filterResourcesResponse filters resources based on read_resource authorization func (rfw *ResponseFilteringWriter) filterResourcesResponse(response *jsonrpc2.Response) (*jsonrpc2.Response, error) { - if rfw.authorizer == nil { - return response, nil - } - // Parse the result as a ListResourcesResult var listResult mcp.ListResourcesResult if err := json.Unmarshal(response.Result, &listResult); err != nil { @@ -495,7 +784,9 @@ func (rfw *ResponseFilteringWriter) filterResourcesResponse(response *jsonrpc2.R nil, // No arguments for the authorization check ) if err != nil { - return nil, err + slog.Warn("Authorization check failed for resource, skipping", + "resource", resource.URI, "error", err) + continue } if authorized { @@ -532,32 +823,60 @@ func (rfw *ResponseFilteringWriter) filterResourcesResponse(response *jsonrpc2.R return filteredResponse, nil } -// writeErrorResponse writes an error response -func (rfw *ResponseFilteringWriter) writeErrorResponse(id jsonrpc2.ID, err error) error { - httpStatus, rpcCode, errorMsg, wwwAuthenticate := classifyAuthorizationError(err) - +// errorResponseBody logs the full filtering error server-side and encodes a +// JSON-RPC error response carrying a deliberately generic client-visible +// message: err can originate in policy evaluation and name tools or resources, +// so security.md forbids returning it verbatim (#6066). Both transports build +// their body here, so neither can leak it. id is echoed so the client can +// correlate the error with its request. +func (rfw *ResponseFilteringWriter) errorResponseBody(id jsonrpc2.ID, err error) []byte { + slog.Error("error filtering response", "method", rfw.method, "error", err) errorResponse := &jsonrpc2.Response{ ID: id, - Error: jsonrpc2.NewError(int64(rpcCode), errorMsg), + Error: jsonrpc2.NewError(mcpparser.CodeInternalError, "internal error"), } - errorData, marshalErr := json.Marshal(errorResponse) - if marshalErr != nil { - // If we can't even marshal the error, write a simple error - rfw.ResponseWriter.WriteHeader(http.StatusInternalServerError) - _, writeErr := rfw.ResponseWriter.Write([]byte(`{"error": "Internal server error"}`)) - return writeErr + body, encErr := jsonrpc2.EncodeMessage(errorResponse) + if encErr != nil { + // Unreachable in practice: errorResponse is always a well-formed + // jsonrpc2.Response built from a valid ID and message above. Fall back + // to a hardcoded valid JSON-RPC error body rather than writing nothing. + slog.Error("failed to encode JSON-RPC error response", "error", encErr) + return fmt.Appendf(nil, `{"jsonrpc":"2.0","error":{"code":%d,"message":"internal error"}}`, + mcpparser.CodeInternalError) } + return body +} - rfw.ResponseWriter.Header().Set("Content-Type", "application/json") - if wwwAuthenticate != "" { - rfw.ResponseWriter.Header().Set("WWW-Authenticate", wwwAuthenticate) - } - rfw.ResponseWriter.WriteHeader(httpStatus) - _, writeErr := rfw.ResponseWriter.Write(errorData) +// writeErrorResponse writes an error response for the application/json path. +// The SSE path must NOT use this: WriteHeader is a no-op once Flush() has +// committed the headers, which is the case for any real SSE response. On the +// SSE path a filtering failure instead becomes the replacement payload of the +// event's data: line (see filterSSEEventData), written through +// processSSEResponse's normal line writer alongside every other line — no +// separate write path or WriteHeader call is needed there. +func (rfw *ResponseFilteringWriter) writeErrorResponse(id jsonrpc2.ID, err error) error { + rfw.ResponseWriter.WriteHeader(http.StatusInternalServerError) + _, writeErr := rfw.ResponseWriter.Write(rfw.errorResponseBody(id, err)) return writeErr } +// requestID recovers the JSON-RPC id of the original request so an error +// response written mid-filtering can still be correlated by the client. +// Falls back to an empty jsonrpc2.ID (which EncodeMessage omits from the +// wire entirely) if the request was never parsed or its id doesn't convert. +func (rfw *ResponseFilteringWriter) requestID() jsonrpc2.ID { + parsed := mcpparser.GetParsedMCPRequest(rfw.request.Context()) + if parsed == nil { + return jsonrpc2.ID{} + } + id, err := mcpparser.ConvertToJSONRPC2ID(parsed.ID) + if err != nil { + return jsonrpc2.ID{} + } + return id +} + // filterFindToolResponse filters the tools list embedded in a find_tool tools/call // response. The response is a CallToolResult whose first text content item contains // a JSON-encoded optimizer.FindToolOutput. Only tools the caller is authorized to @@ -601,11 +920,7 @@ func (rfw *ResponseFilteringWriter) filterFindToolResponse(response *jsonrpc2.Re // that annotations are available even for tools that Cedar will deny. rfw.annotationCache.SetFromToolsList(output.Tools) - filteredTools, err := filterToolsByPolicy(rfw.request.Context(), rfw.authorizer, output.Tools) - if err != nil { - return nil, err - } - output.Tools = filteredTools + output.Tools = filterToolsByPolicy(rfw.request.Context(), rfw.authorizer, output.Tools) filteredText, err := json.Marshal(output) if err != nil { diff --git a/pkg/authz/response_filter_test.go b/pkg/authz/response_filter_test.go index 13425a7d7c..f09cc8c3d7 100644 --- a/pkg/authz/response_filter_test.go +++ b/pkg/authz/response_filter_test.go @@ -4,6 +4,7 @@ package authz import ( + "bytes" "context" "encoding/json" "errors" @@ -17,11 +18,11 @@ import ( "testing" "github.com/golang-jwt/jwt/v5" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" @@ -486,13 +487,14 @@ func TestResponseFilteringWriter_ErrorResponse(t *testing.T) { }, "") require.NoError(t, err, "Failed to create Cedar authorizer") - // Create an error response + // Create an error response in wire format (json.Marshal on the struct would + // produce Go field names, not a JSON-RPC frame). jsonrpcResponse := &jsonrpc2.Response{ ID: jsonrpc2.Int64ID(1), Error: jsonrpc2.NewError(404, "Not found"), } - responseBytes, err := json.Marshal(jsonrpcResponse) + responseBytes, err := jsonrpc2.EncodeMessage(jsonrpcResponse) require.NoError(t, err, "Failed to marshal JSON-RPC response") // Create an HTTP request @@ -504,6 +506,7 @@ func TestResponseFilteringWriter_ErrorResponse(t *testing.T) { // Create the response filtering writer filteringWriter := NewResponseFilteringWriter(rr, authorizer, req, "tools/list", nil, nil) + filteringWriter.ResponseWriter.Header().Set("Content-Type", "application/json") // Write the response data _, err = filteringWriter.Write(responseBytes) @@ -517,87 +520,6 @@ func TestResponseFilteringWriter_ErrorResponse(t *testing.T) { assert.Equal(t, responseBytes, rr.Body.Bytes(), "Error response should pass through unchanged") } -func TestResponseFilteringWriter_ListAuthorizationErrorsReturnExplicitResponses(t *testing.T) { - t.Parallel() - - listResult := mcp.ListToolsResult{ - Tools: []mcp.Tool{ - {Name: "weather", Description: "Get weather information"}, - }, - } - resultJSON, err := json.Marshal(listResult) - require.NoError(t, err) - - jsonrpcResponse := &jsonrpc2.Response{ - ID: jsonrpc2.Int64ID(1), - Result: json.RawMessage(resultJSON), - } - responseBytes, err := jsonrpc2.EncodeMessage(jsonrpcResponse) - require.NoError(t, err) - - tests := []struct { - name string - authErr error - wantStatus int - wantCode int64 - wantMessage string - wantChallenge string - }{ - { - name: "upstream auth required becomes 401", - authErr: authorizers.NewUpstreamAuthenticationRequiredError("cognito", "refresh_failed"), - wantStatus: http.StatusUnauthorized, - wantCode: http.StatusUnauthorized, - wantMessage: upstreamAuthRequiredDescription, - wantChallenge: `Bearer error="invalid_token", error_description="upstream token is no longer valid; re-authentication required"`, - }, - { - name: "generic authz failure becomes 500", - authErr: errors.New("policy evaluation failed"), - wantStatus: http.StatusInternalServerError, - wantCode: http.StatusInternalServerError, - wantMessage: "policy evaluation failed", - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - req, err := http.NewRequest(http.MethodPost, "/messages", nil) - require.NoError(t, err) - - rr := httptest.NewRecorder() - filteringWriter := NewResponseFilteringWriter( - rr, - &stubAuthorizer{allowed: false, err: tc.authErr}, - req, - string(mcp.MethodToolsList), - nil, - nil, - ) - filteringWriter.ResponseWriter.Header().Set("Content-Type", "application/json") - - _, err = filteringWriter.Write(responseBytes) - require.NoError(t, err) - require.NoError(t, filteringWriter.FlushAndFilter()) - - assert.Equal(t, tc.wantStatus, rr.Code) - assert.Equal(t, tc.wantChallenge, rr.Header().Get("WWW-Authenticate")) - - var rpcResp struct { - Error struct { - Code int64 `json:"code"` - Message string `json:"message"` - } `json:"error"` - } - require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &rpcResp)) - assert.Equal(t, tc.wantCode, rpcResp.Error.Code) - assert.Equal(t, tc.wantMessage, rpcResp.Error.Message) - }) - } -} - // TestResponseFilteringWriter_ContentLengthMismatch reproduces a bug where // httputil.ReverseProxy copies the backend's Content-Length header to the // underlying ResponseWriter via Header() (which ResponseFilteringWriter does @@ -947,33 +869,170 @@ func TestOptimizerPassThroughToolsInResponseFilter(t *testing.T) { }) } -// TestResponseFilteringWriter_SSE_PerLineFallthrough is a regression test for -// issue #5257: when an SSE upstream interleaves a non-Response data line (e.g. -// an MCP notification) or an undecodable data line with a real list response, -// the filter previously wrote the entire raw upstream payload and returned, -// leaking the unfiltered list past Cedar. It must instead pass only the -// offending line through and continue filtering the rest of the stream. -// -// The same code path runs for every method covered by -// requiresResponseFiltering, so each of tools/list, prompts/list, and -// resources/list is exercised below. -func TestResponseFilteringWriter_SSE_PerLineFallthrough(t *testing.T) { +// TestCarriesResult pins carriesResult's (and, through it, valueCarriesResult's) +// classification of the shapes the SSE and JSON filters route through it: no +// clean single Response, but the caller needs to know whether a "result" key +// is hiding somewhere in the payload regardless. +func TestCarriesResult(t *testing.T) { t.Parallel() - authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ - Policies: []string{ - `permit(principal, action == Action::"call_tool", resource == Tool::"weather");`, - `permit(principal, action == Action::"get_prompt", resource == Prompt::"greeting");`, - `permit(principal, action == Action::"read_resource", resource == Resource::"data");`, + const notification = `{"jsonrpc":"2.0","method":"notifications/message","params":{}}` + const result = `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}` + + testCases := []struct { + name string + data string + want bool + }{ + { + name: "concatenated notification then result", + // The exact shape an SSE event's assembled data takes when two + // data: lines (a notification, then a result-bearing frame) + // concatenate into one payload: DecodeMessage rejects the whole + // thing as a single value, but the embedded result must still be + // caught. + data: notification + "\n" + result, + want: true, + }, + { + name: "heterogeneous batch array", + // A batch whose first element isn't an object: Unmarshal-ing the + // whole array into a single probe struct would fail outright, + // missing the result-bearing element further down. + data: `[1,` + result + `]`, + want: true, + }, + { + name: "undecodable prefix before a result stops at the first value", + // Deliberate: json.Decoder stops at the first value it can't + // decode ("garbage") and never reaches the result-bearing value + // after it. This is safe rather than a missed catch, because no + // strict client can read past that point either. Don't "fix" + // this into scanning past bad input looking for a result. + data: "garbage\n" + result, + want: false, }, + { + name: "plain notification carries no result", + data: notification, + want: false, + }, + { + name: "empty input", + data: "", + want: false, + }, + { + name: "nested array is not recursed into", + // Deliberate: a JSON-RPC batch is a flat array of message + // objects, so that's the only legitimate batch surface. Recursing + // into nested arrays is an O(d^2) amplification handle on + // encoding/json's own nesting cap for something no client + // actually flattens. Don't "fix" this into recursive descent. + data: `[[` + result + `]]`, + want: false, + }, + { + name: "explicit null result still counts as carrying one", + // Deliberate, the fail-closed direction: RawMessage's + // UnmarshalJSON runs even for a JSON null, so a "result" key + // present with a null value is still treated as carrying a + // result. Don't "fix" this into skipping nulls. + data: `{"jsonrpc":"2.0","id":1,"result":null}`, + want: true, + }, + { + name: "top-level scalar carries no result", + data: `"just a string"`, + want: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.want, carriesResult([]byte(tc.data))) + }) + } +} + +// newWeatherOnlyAuthorizer builds the Cedar authorizer most SSE regression +// tests below share: a single policy permitting call_tool on "weather" plus +// any extraPolicies the caller needs for other resource types. +func newWeatherOnlyAuthorizer(t *testing.T, extraPolicies ...string) authorizers.Authorizer { + t.Helper() + policies := append([]string{ + `permit(principal, action == Action::"call_tool", resource == Tool::"weather");`, + }, extraPolicies...) + authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ + Policies: policies, EntitiesJSON: `[]`, }, "") require.NoError(t, err) + return authorizer +} + +// newUser1Request builds a POST /messages request carrying the "user1" +// identity most SSE regression tests below share. +func newUser1Request(t *testing.T) *http.Request { + t.Helper() + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user1", + Claims: map[string]interface{}{"sub": "user1"}, + }} + req, err := http.NewRequest(http.MethodPost, "/messages", nil) + require.NoError(t, err) + return req.WithContext(auth.WithIdentity(req.Context(), identity)) +} +// newParsedUser1Request is newUser1Request plus mcpparser.ParsingMiddleware, +// for tests whose error envelope must correlate against the request's real +// JSON-RPC id: requestID() reads the parsed request back out of context, so +// a plain request (no parsed MCP request in context) would always produce a +// zero, id-less envelope regardless of what the fix is supposed to do. +func newParsedUser1Request(t *testing.T, reqBody string) *http.Request { + t.Helper() identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ Subject: "user1", Claims: map[string]interface{}{"sub": "user1"}, }} + req, err := http.NewRequest(http.MethodPost, "/messages", strings.NewReader(reqBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + var parsedReq *http.Request + mcpparser.ParsingMiddleware(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + parsedReq = r + })).ServeHTTP(httptest.NewRecorder(), req) + require.NotNil(t, parsedReq) + // ParsingMiddleware calls next even when parsing fails, so a nil parsedReq + // alone wouldn't catch a typo'd reqBody that fails to parse — that would + // silently degrade every caller to the zero-ID path. A notification (no + // "id" in reqBody) still parses to a non-nil ParsedMCPRequest with a nil + // ID, so this assertion holds for every caller, not just id-bearing ones. + require.NotNil(t, mcpparser.GetParsedMCPRequest(parsedReq.Context()), + "reqBody must have parsed into a ParsedMCPRequest, or requestID() silently falls back to the zero ID") + return parsedReq +} + +// TestResponseFilteringWriter_SSE_PerEventFallthrough is a regression test for +// issue #5257: when an SSE upstream interleaves a non-Response event (e.g. an +// MCP notification) or an undecodable event with a real list response, the +// filter previously wrote the entire raw upstream payload and returned, +// leaking the unfiltered list past Cedar. It must instead pass only the +// offending event through and continue filtering the rest of the stream. +// +// The same code path runs for every method covered by +// requiresResponseFiltering, so each of tools/list, prompts/list, and +// resources/list is exercised below. +func TestResponseFilteringWriter_SSE_PerEventFallthrough(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t, + `permit(principal, action == Action::"get_prompt", resource == Prompt::"greeting");`, + `permit(principal, action == Action::"read_resource", resource == Resource::"data");`, + ) // encodeListResponse marshals a list result type into a JSON-RPC Response // data line. @@ -1093,16 +1152,21 @@ func TestResponseFilteringWriter_SSE_PerLineFallthrough(t *testing.T) { t.Run(mc.name+"/"+plc.name, func(t *testing.T) { t.Parallel() - req, err := http.NewRequest(http.MethodPost, "/messages", nil) - require.NoError(t, err) - req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + req := newUser1Request(t) rr := httptest.NewRecorder() rfw := NewResponseFilteringWriter(rr, authorizer, req, mc.method, nil, nil) rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") - body := strings.Join([]string{plc.line, mc.respLine, ""}, "\n") - _, err = rfw.Write([]byte(body)) + // A blank line separates plc.line and mc.respLine into two + // distinct SSE events, matching how a real MCP server frames + // SSE (one JSON-RPC message per event): otherwise, with no + // blank line between them, they'd be two data: fields of the + // SAME event and their values would concatenate, which is not + // what this test means to model. The trailing "" pair + // terminates the second event. + body := strings.Join([]string{plc.line, "", mc.respLine, "", ""}, "\n") + _, err := rfw.Write([]byte(body)) require.NoError(t, err) require.NoError(t, rfw.FlushAndFilter()) @@ -1139,6 +1203,13 @@ func TestResponseFilteringWriter_SSE_PerLineFallthrough(t *testing.T) { // must not appear in the wire output. assert.NotContains(t, out, `"`+mc.unauthorizedName+`"`, "unfiltered list payload leaked into SSE output") + + // The blank-line event terminator must survive the rewrite: + // processSSEResponse's structural invariant is that every + // line the input had, including blank lines, is written + // exactly once in its original position. + assert.True(t, strings.HasSuffix(out, "\n\n"), + "output must end with a blank-line SSE terminator") }) } } @@ -1147,23 +1218,21 @@ func TestResponseFilteringWriter_SSE_PerLineFallthrough(t *testing.T) { // TestResponseFilteringWriter_SSE_DisguisedResponseFrame is a regression test // for a residual of issue #5257: a frame that carries both a method field (so // jsonrpc2.DecodeMessage classifies it as a request/notification rather than a -// Response) and a result field smuggling a real list. Such a frame must be -// dropped, not passed through, or the unfiltered list leaks past Cedar. +// Response) and a result field smuggling a real list. Such a frame must fail +// closed with an error envelope in its place — not pass the smuggled list +// through, and not silently drop the event's data (which per WHATWG SSE +// semantics leaves the client waiting on a dispatch that never carries a +// usable payload; see #6037). +// +// The envelope must also be one the client can actually correlate: it's +// decoded and checked for a non-nil Error carrying the *request's* id, not +// just a substring match, because an id-less envelope (the zero jsonrpc2.ID, +// which EncodeMessage omits from the wire) reproduces the #6037 hang even +// though an "error" key is present on the wire. func TestResponseFilteringWriter_SSE_DisguisedResponseFrame(t *testing.T) { t.Parallel() - authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ - Policies: []string{ - `permit(principal, action == Action::"call_tool", resource == Tool::"weather");`, - }, - EntitiesJSON: `[]`, - }, "") - require.NoError(t, err) - - identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ - Subject: "user1", - Claims: map[string]interface{}{"sub": "user1"}, - }} + authorizer := newWeatherOnlyAuthorizer(t) // A genuine, filterable tools/list response on a later line. realResultJSON, err := json.Marshal(mcp.ListToolsResult{ @@ -1179,7 +1248,7 @@ func TestResponseFilteringWriter_SSE_DisguisedResponseFrame(t *testing.T) { // Each frame smuggles a tools/list result outside a clean Response, via a // different upstream-controlled shape. DecodeMessage either classifies them // as non-Response or rejects them outright, so the pre-fix code passed them - // through unfiltered. All must be dropped, not leaked. + // through unfiltered. All must fail closed into an error envelope instead. frames := []struct { name string payload string @@ -1189,6 +1258,10 @@ func TestResponseFilteringWriter_SSE_DisguisedResponseFrame(t *testing.T) { {"no id", `data: {"jsonrpc":"2.0","result":{"tools":[{"name":"admin_tool"}]}}`}, {"non-scalar id", `data: {"jsonrpc":"2.0","id":{"x":1},"result":{"tools":[{"name":"admin_tool"}]}}`}, {"batch array", `data: [{"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"admin_tool"}]}}]`}, + // A batch array whose first element isn't an object: Unmarshal-ing + // the whole array into []struct{Result} would previously fail + // outright, missing the result-bearing element further down. + {"heterogeneous batch array", `data: [1,{"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"admin_tool"}]}}]`}, } for _, f := range frames { @@ -1196,16 +1269,20 @@ func TestResponseFilteringWriter_SSE_DisguisedResponseFrame(t *testing.T) { t.Run(f.name, func(t *testing.T) { t.Parallel() - req, err := http.NewRequest(http.MethodPost, "/messages", nil) - require.NoError(t, err) - req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + // Run the real MCP parsing middleware so the request context + // carries a ParsedMCPRequest with a real id, matching production: + // requestID() reads it back out to correlate the error envelope. + parsedReq := newParsedUser1Request(t, `{"jsonrpc":"2.0","id":99,"method":"tools/list"}`) rr := httptest.NewRecorder() - rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw := NewResponseFilteringWriter(rr, authorizer, parsedReq, string(mcp.MethodToolsList), nil, nil) rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") - body := strings.Join([]string{f.payload, "data: " + string(realResp), ""}, "\n") - _, err = rfw.Write([]byte(body)) + // A blank line separates the disguised frame from the genuine + // response into two distinct SSE events, matching how a real MCP + // server frames SSE (one JSON-RPC message per event). + body := strings.Join([]string{f.payload, "", "data: " + string(realResp), "", ""}, "\n") + _, err := rfw.Write([]byte(body)) require.NoError(t, err) require.NoError(t, rfw.FlushAndFilter()) @@ -1213,46 +1290,1224 @@ func TestResponseFilteringWriter_SSE_DisguisedResponseFrame(t *testing.T) { assert.NotContains(t, out, "admin_tool", "smuggled result leaked past the filter") assert.Contains(t, out, "weather", - "genuine list response on a later line must still be delivered") + "genuine list response in a later event must still be delivered") + + // The disguised frame's event is first on the wire; decode it and + // require an error envelope the client can correlate. TrimSuffix + // guards a CRLF body leaving a trailing \r on the line. + firstLine := strings.TrimSuffix(strings.SplitN(out, "\n", 2)[0], "\r") + payload := strings.TrimPrefix(firstLine, "data: ") + msg, err := jsonrpc2.DecodeMessage([]byte(payload)) + require.NoError(t, err, "the fail-closed envelope must itself be valid JSON-RPC") + resp, ok := msg.(*jsonrpc2.Response) + require.True(t, ok, "the fail-closed envelope must be a clean Response") + require.NotNil(t, resp.Error, "the disguised frame's event must fail closed into an error envelope, not vanish silently") + assert.Equal(t, jsonrpc2.Int64ID(99), resp.ID, + "the error envelope must carry the request's id so the client can correlate it, or the #6037 hang reproduces") }) } } -// TestResponseFilteringWriter_JSON_DisguisedResponseFrame is the application/json -// counterpart: the disguised-result bypass is transport-independent, so the JSON -// path must also fail closed on a smuggled result rather than write it through. -func TestResponseFilteringWriter_JSON_DisguisedResponseFrame(t *testing.T) { +// TestResponseFilteringWriter_SSE_FailClosedDropsEventName is a regression +// test for #6037: the go-sdk streamable client only dispatches unnamed +// ("message"-typed) events to its MCP request handler, so substituting a +// fail-closed error envelope in place of a NAMED event's data (while leaving +// its "event:" field untouched) meant the envelope was never delivered and +// the caller hung on its own timeout -- the very hang the #6037 envelope was +// meant to prevent. This asserts what parseSSEStream (the go-sdk client +// model, not the WHATWG grammar) actually dispatches, since that's what +// determines delivery. +func TestResponseFilteringWriter_SSE_FailClosedDropsEventName(t *testing.T) { t.Parallel() - authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ - Policies: []string{ - `permit(principal, action == Action::"call_tool", resource == Tool::"weather");`, - }, - EntitiesJSON: `[]`, - }, "") + authorizer := newWeatherOnlyAuthorizer(t) + parsedReq := newParsedUser1Request(t, `{"jsonrpc":"2.0","id":99,"method":"tools/list"}`) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, parsedReq, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + // A method+result frame that fails closed (see + // TestResponseFilteringWriter_SSE_DisguisedResponseFrame), framed under a + // custom event name the way a real backend might tag a tool-list push. + body := "event: toolsUpdate\n" + + `data: {"jsonrpc":"2.0","method":"notifications/initialized","id":1,` + + `"result":{"tools":[{"name":"admin_tool"}]}}` + "\n\n" + _, err := rfw.Write([]byte(body)) require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) - identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ - Subject: "user1", - Claims: map[string]interface{}{"sub": "user1"}, - }} + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", "smuggled result leaked past the filter") + assert.NotContains(t, out, "event:", + "the offending event's event: field must be dropped, or the go-sdk client never dispatches the fail-closed envelope") + + events := parseSSEStream(t, rr.Body.Bytes()) + require.Len(t, events, 1, "body: %q", out) + assert.Empty(t, events[0].eventType, + "the go-sdk client only dispatches unnamed/'message' events to its MCP handler; a named event here reproduces the #6037 hang") + + msg, err := jsonrpc2.DecodeMessage([]byte(events[0].data)) + require.NoError(t, err, "the dispatched envelope must itself be valid JSON-RPC") + resp, ok := msg.(*jsonrpc2.Response) + require.True(t, ok, "the dispatched envelope must be a clean Response") + require.NotNil(t, resp.Error, "the event must fail closed into an error envelope") + assert.Equal(t, jsonrpc2.Int64ID(99), resp.ID, + "the error envelope must carry the request's id so the client can correlate it") +} - // A batch array smuggling a tools/list result. DecodeMessage rejects the - // array, so the pre-fix JSON path wrote it through raw. - smuggled := `[{"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"admin_tool"}]}}]` +// TestResponseFilteringWriter_SSE_ConcatenatedEventBypass is a regression test +// for a #5257-class leak the event-based rewrite itself introduced: two +// data: fields in ONE event (a notification followed by a result-bearing +// frame) assemble, per WHATWG semantics, into "{notif}\n{result}". That +// payload fails jsonrpc2.DecodeMessage (trailing data after the first JSON +// value), and carriesResult must still catch the embedded result — Unmarshal +// on the whole concatenated payload rejects it the same way DecodeMessage +// does, which would otherwise route it to the "pass through unfiltered" +// branch and leak the unauthorized tool straight through. +func TestResponseFilteringWriter_SSE_ConcatenatedEventBypass(t *testing.T) { + t.Parallel() - req, err := http.NewRequest(http.MethodPost, "/messages", nil) + authorizer := newWeatherOnlyAuthorizer(t) + + notificationLine := `data: {"jsonrpc":"2.0","method":"notifications/message","params":{}}` + resultLine := `data: {"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"admin_tool"}]}}` + + // Run the real MCP parsing middleware so the request context carries a + // ParsedMCPRequest with a real id, matching production: requestID() reads + // it back out to correlate the error envelope. Without this, requestID() + // falls back to the zero ID, the envelope carries no "id", and this test + // would accept an envelope no client could actually correlate. + req := newParsedUser1Request(t, `{"jsonrpc":"2.0","id":99,"method":"tools/list"}`) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + // No blank line between the two data: lines: one event, two fields. + body := strings.Join([]string{notificationLine, resultLine, "", ""}, "\n") + _, err := rfw.Write([]byte(body)) require.NoError(t, err) - req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a result concatenated after a notification within one event must not bypass the filter") + + // NotContains alone would also pass if the event were dropped silently + // (#6037's hang); require the fail-closed envelope actually went out and + // carries the request's id so the client can correlate it. + firstLine := strings.TrimSuffix(strings.SplitN(out, "\n", 2)[0], "\r") + payload := strings.TrimPrefix(firstLine, "data: ") + msg, err := jsonrpc2.DecodeMessage([]byte(payload)) + require.NoError(t, err, "the fail-closed envelope must itself be valid JSON-RPC") + resp, ok := msg.(*jsonrpc2.Response) + require.True(t, ok, "the fail-closed envelope must be a clean Response") + require.NotNil(t, resp.Error, "the concatenated event must fail closed into an error envelope, not vanish silently") + assert.Equal(t, jsonrpc2.Int64ID(99), resp.ID, + "the error envelope must carry the request's id so the client can correlate it, or the #6037 hang reproduces") +} + +// TestResponseFilteringWriter_SSE_MixedSeparatorsBypass is a regression test +// for a #5257-class leak: SSE permits CR, LF, and CRLF mixed within one +// stream, but sniffing a single stream-wide separator (the pre-fix approach) +// disagrees with how a spec-compliant client actually splits the body into +// lines and events. Here the body sniffs as CRLF, so a conformant client's +// LF-terminated lines collapse into one giant "line" for the sniffing filter, +// whose data: prefix strip leaves an undecodable payload that falls through +// unfiltered — while the client itself correctly sees three separate events, +// the third a clean Response carrying the unauthorized tool. +func TestResponseFilteringWriter_SSE_MixedSeparatorsBypass(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + resultLine := "data: " + `{"jsonrpc":"2.0","id":1,"result":` + string(resultJSON) + "}" + + body := "data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/message\",\"params\":{}}\r\n" + + "\r\n" + + "data: x\n" + + "\n" + + resultLine + "\n" + + "\n" + _, err = rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "mixed line-terminator conventions within one SSE body must not bypass the filter") + // NotContains alone would also pass if the event were dropped rather than + // filtered (#6037's hang); prove it was actually filtered, not eaten. + assert.Contains(t, out, "weather", "the authorized tool must survive filtering") +} + +// TestResponseFilteringWriter_SSE_LeadingBOMBypass is a regression test for a +// #5257-class leak: a client strips a leading UTF-8 BOM per the WHATWG decode +// algorithm before parsing lines, but the filter compared the raw first line +// against "data:". The BOM byte prefix made that comparison fail, so the +// event (and the unauthorized tool inside it) passed through unfiltered. +func TestResponseFilteringWriter_SSE_LeadingBOMBypass(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + req := newUser1Request(t) rr := httptest.NewRecorder() rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) - rfw.ResponseWriter.Header().Set("Content-Type", "application/json") + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + body := "\xEF\xBB\xBF" + `data: {"jsonrpc":"2.0","id":1,"result":` + string(resultJSON) + "}\n\n" - _, err = rfw.Write([]byte(smuggled)) + _, err = rfw.Write([]byte(body)) require.NoError(t, err) require.NoError(t, rfw.FlushAndFilter()) - assert.NotContains(t, rr.Body.String(), "admin_tool", - "smuggled result on the application/json transport leaked past the filter") + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a leading UTF-8 BOM must not bypass the filter") + assert.Contains(t, out, "weather", "the authorized tool must survive filtering") + assert.False(t, strings.HasPrefix(out, "\xEF\xBB\xBF"), + "the leading BOM must be dropped from the output, not stripped-for-matching then re-emitted") +} + +// TestResponseFilteringWriter_SSE_ErrorAndResultBypass is a regression test +// for a #5257-class leak: jsonrpc2.DecodeMessage and EncodeMessage both +// populate/re-emit "error" and "result" together on one Response, and +// filterListResponse's `response.Error != nil` check returned the whole +// response, list intact, as soon as an error field was present. The reference +// TS SDK's zod schema silently strips the unknown "error" key rather than +// rejecting the message, so it would accept this as a successful response +// carrying the full unfiltered list. +// +// The error-only row distinguishes "failed closed" from "passed through +// unfiltered": both leave no admin_tool on the wire, so a substring check +// alone can't tell a fix from a regression that fails closed on every +// legitimate upstream error. Only the error *code* does — mcpparser.CodeInternalError +// is our envelope, anything else is the upstream's own error surviving untouched. +func TestResponseFilteringWriter_SSE_ErrorAndResultBypass(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + testCases := []struct { + name string + body string + wantErrorCode float64 + }{ + { + name: "error and result together fails closed", + body: `data: {"jsonrpc":"2.0","id":1,"error":{"code":1,"message":"x"},` + + `"result":{"tools":[{"name":"admin_tool"}]}}` + "\n\n", + wantErrorCode: float64(mcpparser.CodeInternalError), // our envelope, not the upstream's code 1 + }, + { + name: "error alone passes through unfiltered", + body: `data: {"jsonrpc":"2.0","id":1,"error":{"code":1,"message":"x"}}` + "\n\n", + wantErrorCode: 1, // the upstream's own error, untouched + }, + { + // A literal `"result":null` is exempted from the both-error- + // and-result fail-closed rule (see filterListResponse): a null + // result can never carry a list, so failing closed on it only + // destroys the upstream's real error code for no security + // benefit. Regression for that exemption. + name: "error with explicit null result passes through unfiltered", + body: `data: {"jsonrpc":"2.0","id":1,"error":{"code":404,"message":"not found"},"result":null}` + "\n\n", + wantErrorCode: 404, // the upstream's own error, untouched + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + _, err := rfw.Write([]byte(tc.body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a Response carrying both error and result must not pass the list through under cover of the error field") + + firstLine := strings.TrimSuffix(strings.SplitN(out, "\n", 2)[0], "\r") + var decoded struct { + Error struct { + Code float64 `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal([]byte(strings.TrimPrefix(firstLine, "data: ")), &decoded)) + assert.Equal(t, tc.wantErrorCode, decoded.Error.Code, + "the error code distinguishes fail-closed (mcpparser.CodeInternalError, our envelope) from pass-through (the upstream's own code)") + }) + } +} + +// TestResponseFilteringWriter_JSON_DisguisedResponseFrame is the application/json +// counterpart: the disguised-result bypass is transport-independent, so the JSON +// path must also fail closed on a smuggled result rather than write it through. +// +// It also pins the writeErrorResponse envelope: the error body written on this +// path must be conformant JSON-RPC (lowercase "jsonrpc"/"id"/"error" keys, no +// "id" key for a notification) and must carry the *original request's* id, not +// an empty one, so the client can still correlate the denial. +func TestResponseFilteringWriter_JSON_DisguisedResponseFrame(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + // A batch array smuggling a tools/list result. DecodeMessage rejects the + // array, so the pre-fix JSON path wrote it through raw. + const smuggled = `[{"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"admin_tool"}]}}]` + + // writeErrorResponse logs the carriesResult error server-side and sends the + // client a fixed generic message alongside the standard JSON-RPC Internal + // Error code (mcpparser.CodeInternalError), not the HTTP status. + wantErr := `"error":{"code":` + strconv.FormatInt(mcpparser.CodeInternalError, 10) + `,"message":"internal error"}` + + testCases := []struct { + name string + reqIDJSON string // raw "id" field on the *incoming* request; "" omits it (notification) + wantIDKey string // expected "id" fragment in the error envelope; "" if id must be absent + }{ + { + name: "int id is recovered onto the error envelope", + reqIDJSON: `"id":42,`, + wantIDKey: `"id":42,`, + }, + { + name: "string id is recovered onto the error envelope", + reqIDJSON: `"id":"abc",`, + wantIDKey: `"id":"abc",`, + }, + { + name: "notification (no id) omits id entirely", + reqIDJSON: "", + wantIDKey: "", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Run the real MCP parsing middleware so the request context carries a + // ParsedMCPRequest, matching how the request would look at this point + // in the real middleware chain. writeErrorResponse's id-recovery path + // reads the request id back out of that context. + reqBody := `{"jsonrpc":"2.0",` + tc.reqIDJSON + `"method":"tools/list"}` + parsedReq := newParsedUser1Request(t, reqBody) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, parsedReq, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "application/json") + + _, err := rfw.Write([]byte(smuggled)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + assert.Equal(t, http.StatusInternalServerError, rr.Code) + + body := rr.Body.String() + assert.NotContains(t, body, "admin_tool", + "smuggled result on the application/json transport leaked past the filter") + + // assert.JSONEq is key-case-sensitive, so this single assertion catches + // "Error"/"ID" substituted for "error"/"id" (the historical bug), a + // missing "jsonrpc" tag, and a wrong or missing id all at once. + wantBody := `{"jsonrpc":"2.0",` + tc.wantIDKey + wantErr + `}` + assert.JSONEq(t, wantBody, body) + + var decoded map[string]json.RawMessage + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &decoded)) + _, hasID := decoded["id"] + assert.Equal(t, tc.wantIDKey != "", hasID, "id key presence mismatch") + _, hasResult := decoded["result"] + assert.False(t, hasResult, "error response must not contain a result key") + }) + } +} + +// TestErrorResponseBody pins the wire shape of the JSON-RPC error envelope +// errorResponseBody produces for a filter/encode failure: a "2.0" jsonrpc +// tag, the standard JSON-RPC internal-error code (mcpparser.CodeInternalError) +// with a fixed generic client-visible message that never echoes the wrapped +// error's text (#6066 — that error can originate in policy evaluation and +// name tools or resources), and an id that round-trips for a real request id +// but is entirely ABSENT (not present and null) for the zero-value id. +// Absence matters because MCP types the error response id as optional +// (id?: RequestId), and the reference TypeScript SDK's strict schema admits +// undefined but not null — a null id would make that client throw inside +// its transport. +func TestErrorResponseBody(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + // A distinctive sentinel, not a generic message: its absence from the + // encoded body is what actually pins the #6066 scrub, rather than + // passing by the coincidence of a short or common wrapped error string. + wrapped := errors.New("leaky-tool-name-sentinel") + + testCases := []struct { + name string + id jsonrpc2.ID + wantID any // nil means the id key must be absent + }{ + {name: "int64 id round-trips", id: jsonrpc2.Int64ID(7), wantID: float64(7)}, + {name: "string id round-trips", id: jsonrpc2.StringID("abc"), wantID: "abc"}, + {name: "zero-value id key is absent, not null", id: jsonrpc2.ID{}}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := newUser1Request(t) + rfw := NewResponseFilteringWriter(httptest.NewRecorder(), authorizer, req, string(mcp.MethodToolsList), nil, nil) + body := rfw.errorResponseBody(tc.id, wrapped) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(body, &decoded)) + + assert.Equal(t, "2.0", decoded["jsonrpc"]) + + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "error field must decode as an object") + assert.Equal(t, float64(mcpparser.CodeInternalError), errObj["code"]) + assert.Equal(t, "internal error", errObj["message"], + "the client-visible message must be the fixed generic string, not the wrapped error") + assert.NotContains(t, string(body), wrapped.Error(), + "the wrapped error's text must never reach the wire (#6066)") + + idValue, hasID := decoded["id"] + require.Equal(t, tc.wantID != nil, hasID, "id key presence mismatch") + if tc.wantID != nil { + assert.Equal(t, tc.wantID, idValue) + } + }) + } +} + +// TestResponseFilteringWriter_SSE_SplitPayloadAcrossDataLines is a regression +// test for the #6037 rewrite: SSE data fields belonging to the same event +// concatenate (per WHATWG semantics) before the event is decoded, so a +// JSON-RPC frame split across two data: lines within one event must still be +// reconstructed and filtered, not smuggled past the filter because each half +// independently fails to decode. +// +// The split point is chosen right after a comma (a whitespace-legal position +// in JSON grammar), so each half is independently invalid JSON on its own, +// but the two halves, reassembled with the LF the SSE spec mandates between +// concatenated data values, form the original valid JSON-RPC frame again. +func TestResponseFilteringWriter_SSE_SplitPayloadAcrossDataLines(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + encoded, err := jsonrpc2.EncodeMessage(&jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(1), + Result: json.RawMessage(resultJSON), + }) + require.NoError(t, err) + + splitAt := bytes.Index(encoded, []byte(`,"result"`)) + 1 + require.Greater(t, splitAt, 0, "test fixture assumption broken: no comma before \"result\" in the encoded frame") + firstHalf, secondHalf := encoded[:splitAt], encoded[splitAt:] + + // Each half is independently undecodable and carries no result of its + // own, so the pre-rewrite per-line filter passed both through unfiltered. + _, err = jsonrpc2.DecodeMessage(firstHalf) + require.Error(t, err, "test fixture assumption broken: first half must not decode on its own") + require.False(t, carriesResult(firstHalf), "test fixture assumption broken: first half must not independently carry a result") + _, err = jsonrpc2.DecodeMessage(secondHalf) + require.Error(t, err, "test fixture assumption broken: second half must not decode on its own") + require.False(t, carriesResult(secondHalf), "test fixture assumption broken: second half must not independently carry a result") + + body := strings.Join([]string{"data: " + string(firstHalf), "data: " + string(secondHalf), "", ""}, "\n") + + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + _, err = rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a JSON-RPC frame split across two data: lines within one event must still be filtered") + assert.Contains(t, out, "weather", "the authorized tool must survive filtering") +} + +// TestResponseFilteringWriter_SSE_MultiEventBody is a regression test for the +// #6037 rewrite: a body carrying two complete SSE events (each terminated by +// its own blank line) must deliver both events, with the blank line between +// them preserved, and the second event's result still filtered. The +// pre-rewrite per-line implementation dropped blank lines outright and +// reconciled the deficit with at most one substitute separator, collapsing +// any two-event body into something that looks like one malformed event. +func TestResponseFilteringWriter_SSE_MultiEventBody(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + notificationLine := `data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"warming up"}}` + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + encoded, err := jsonrpc2.EncodeMessage(&jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(1), + Result: json.RawMessage(resultJSON), + }) + require.NoError(t, err) + respLine := "data: " + string(encoded) + + // Two complete events, each terminated by its own blank line. + body := strings.Join([]string{notificationLine, "", respLine, "", ""}, "\n") + + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + _, err = rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + + // Both events must survive as separate events: the notification verbatim, + // separated by its own blank line from the (filtered) list response. + assert.Contains(t, out, notificationLine+"\n\n", + "the notification event must be delivered with its terminating blank line intact") + assert.NotContains(t, out, "admin_tool", "the second event's result must still be filtered") + assert.Contains(t, out, "weather", "the authorized tool must survive filtering") +} + +// TestResponseFilteringWriter_SSE_StructuralRoundTrip verifies the core +// structural invariant of the #6037 rewrite: a body needing no filtering +// (here, an interleaved notification whose event carries no result) is +// reproduced byte-for-byte, regardless of how the body's separators trail. +func TestResponseFilteringWriter_SSE_StructuralRoundTrip(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + // The notification event carries no result, so it never reaches Cedar + // evaluation; the policy content is irrelevant to this test. + const dataLine = `data: {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"hi"}}` + + testCases := []struct { + name string + body string + }{ + {name: "ends with a blank line", body: "event: message\n" + dataLine + "\n\n"}, + {name: "ends with a single separator, no blank line", body: "event: message\n" + dataLine + "\n"}, + {name: "no separator anywhere in the body", body: dataLine}, + {name: "CRLF body ending with a blank line", body: "event: message\r\n" + dataLine + "\r\n\r\n"}, + // One event terminated by CRLF then a lone-LF blank line, followed by + // a second event terminated by a lone CR then a CRLF blank line: all + // three SSE-legal terminators appear in one body, which per WHATWG an + // upstream may do freely. + {name: "mixed CRLF, LF and CR within one body", body: dataLine + "\r\n" + "\n" + "event: ping" + "\r" + "\r\n"}, + {name: "consecutive blank lines", body: dataLine + "\n\n\n"}, + {name: "blank lines only", body: "\n\n\n"}, + // Guards the idx+1 < len(rawResponse) bounds check: a lone trailing + // CR with nothing after it must not index out of range. + {name: "ends with a lone CR", body: dataLine + "\r"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + _, err := rfw.Write([]byte(tc.body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + assert.Equal(t, tc.body, rr.Body.String(), + "a body needing no filtering must be reproduced byte-for-byte") + }) + } +} + +// TestResponseFilteringWriter_SSE_FilteredEventTerminators covers a hole +// TestResponseFilteringWriter_SSE_StructuralRoundTrip leaves open: every row +// there asserts byte-identity for a body that needs no filtering, so nothing +// ever reaches resolveSSEEvent's replacement branch. Swapping that branch's +// `term: l.term` for a hardcoded []byte("\n") would still pass the whole +// suite. These cases exercise a body that IS filtered, so they assert the +// (changed) output rather than a reproduced input and don't belong under the +// byte-identity test's name. +func TestResponseFilteringWriter_SSE_FilteredEventTerminators(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + encoded, err := jsonrpc2.EncodeMessage(&jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(1), + Result: json.RawMessage(resultJSON), + }) + require.NoError(t, err) + + t.Run("CRLF body: filtered data line and blank line stay CRLF-terminated", func(t *testing.T) { + t.Parallel() + + req := newUser1Request(t) + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + body := "data: " + string(encoded) + "\r\n\r\n" + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", "the filtered event must not leak the denied tool") + assert.Contains(t, out, "weather", "the authorized tool must survive filtering") + require.True(t, strings.HasSuffix(out, "\r\n\r\n"), + "the replacement data line and the following blank line must both stay CRLF-terminated: %q", out) + + // Every LF in the output must be immediately preceded by a CR: a bare + // LF proves the replacement line (or the blank line after it) fell + // back to a hardcoded "\n" instead of the CRLF the input actually used. + for i := 0; i < len(out); i++ { + if out[i] == '\n' { + require.True(t, i > 0 && out[i-1] == '\r', "found a bare LF at byte %d in %q", i, out) + } + } + }) + + t.Run("trailing filtered event with no closing blank line", func(t *testing.T) { + t.Parallel() + + req := newUser1Request(t) + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + // No trailing terminator at all: the backend closed the stream + // mid-event, same fixture as the unterminated-body case documented + // in processSSEResponse. + body := "data: " + string(encoded) + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", "the filtered event must not leak the denied tool") + assert.Contains(t, out, "weather", "the authorized tool must survive filtering") + assert.False(t, strings.HasSuffix(out, "\n") || strings.HasSuffix(out, "\r"), + "an unterminated input must not gain a fabricated terminator: %q", out) + }) +} + +// TestResponseFilteringWriter_FilterBypassAttempts verifies that list results a +// backend tries to slip past the filter -- via media-type casing or whitespace +// (RFC 9110 makes both legal), a non-200 2xx status (fetch-based clients accept +// any 2xx as ok), or an unrecognized Content-Type on a result-carrying body -- +// are still filtered, while responses with nothing to filter pass through. +func TestResponseFilteringWriter_FilterBypassAttempts(t *testing.T) { + t.Parallel() + + // The policy permits only "weather"; "admin_tool" must never reach the client. + authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ + Policies: []string{ + `permit(principal, action == Action::"call_tool", resource == Tool::"weather");`, + }, + EntitiesJSON: `[]`, + }, "") + require.NoError(t, err) + + resultData, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Admin operations"}, + }, + }) + require.NoError(t, err) + listBody, err := jsonrpc2.EncodeMessage(&jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(1), + Result: json.RawMessage(resultData), + }) + require.NoError(t, err) + sseListBody := []byte("event: message\ndata: " + string(listBody) + "\n\n") + + testCases := []struct { + name string + contentType string // "" leaves the Content-Type header unset + statusCode int + body []byte + // wantFiltered: the denied tool is stripped from the client-visible + // body; otherwise the body passes through byte-for-byte. + wantFiltered bool + }{ + { + name: "uppercase media type is normalized and filtered", + contentType: "Application/JSON", + statusCode: http.StatusOK, + body: listBody, + wantFiltered: true, + }, + { + name: "whitespace before media type parameters is normalized and filtered", + contentType: "application/json ; charset=utf-8", + statusCode: http.StatusOK, + body: listBody, + wantFiltered: true, + }, + { + name: "2xx status other than 200 and 202 is filtered", + contentType: "application/json", + statusCode: http.StatusCreated, + body: listBody, + wantFiltered: true, + }, + { + name: "unrecognized media type carrying a JSON result is filtered", + contentType: "text/plain", + statusCode: http.StatusOK, + body: listBody, + wantFiltered: true, + }, + { + name: "missing media type carrying an SSE result is filtered", + contentType: "", + statusCode: http.StatusOK, + body: sseListBody, + wantFiltered: true, + }, + { + name: "unrecognized media type with no JSON-RPC result passes through", + contentType: "text/plain", + statusCode: http.StatusOK, + body: []byte("plain text, nothing to filter"), + wantFiltered: false, + }, + { + name: "non-2xx error response passes through unfiltered", + contentType: "application/json", + statusCode: http.StatusInternalServerError, + body: listBody, + wantFiltered: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req, err := http.NewRequest(http.MethodPost, "/messages", nil) + require.NoError(t, err) + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: "user123", + Claims: jwt.MapClaims{"sub": "user123"}, + }} + req = req.WithContext(auth.WithIdentity(req.Context(), identity)) + + rr := httptest.NewRecorder() + fw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + if tc.contentType != "" { + fw.ResponseWriter.Header().Set("Content-Type", tc.contentType) + } + fw.WriteHeader(tc.statusCode) + _, err = fw.Write(tc.body) + require.NoError(t, err) + require.NoError(t, fw.FlushAndFilter()) + + assert.Equal(t, tc.statusCode, rr.Code, "original status code should be preserved") + if tc.wantFiltered { + body := rr.Body.String() + assert.Contains(t, body, "weather", "permitted tool should remain in the response") + assert.NotContains(t, body, "admin_tool", "denied tool leaked to the client") + } else { + assert.Equal(t, tc.body, rr.Body.Bytes(), "response should pass through unchanged") + } + }) + } +} + +// TestResponseFilteringWriter_SSE_LeadingNonJSONWhitespaceStillFiltered is a +// regression test for the assembled-payload TrimSpace fix: Go's JSON scanner +// only treats SP/TAB/CR/LF as whitespace, but unicode.IsSpace (and the +// go-sdk client's own buffer trim, mcp/event.go) also covers U+000B, U+000C, +// U+0085, U+00A0, U+1680 and U+3000. Any of those left on the front of the +// assembled payload made DecodeMessage and carriesResult both reject it while +// the client parsed it fine, so the event fell to the undecodable-passthrough +// branch and the unfiltered list went out (#5257). +// +// NotContains on admin_tool alone would also be satisfied by the event +// failing closed, which would silently degrade every legitimate response +// that happens to carry leading whitespace -- so each case also requires the +// permitted tool to survive, proving the event was filtered, not dropped. +func TestResponseFilteringWriter_SSE_LeadingNonJSONWhitespaceStillFiltered(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + encoded, err := jsonrpc2.EncodeMessage(&jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(1), + Result: json.RawMessage(resultJSON), + }) + require.NoError(t, err) + + testCases := []struct { + name string + prefix string + }{ + {name: "no prefix (baseline)", prefix: ""}, + {name: "U+000B line tabulation", prefix: "\u000B"}, + {name: "U+000C form feed", prefix: "\u000C"}, + {name: "U+0085 next line", prefix: "\u0085"}, + {name: "U+00A0 no-break space", prefix: "\u00A0"}, + {name: "U+1680 ogham space mark", prefix: "\u1680"}, + {name: "U+3000 ideographic space", prefix: "\u3000"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := newUser1Request(t) + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + body := "data: " + tc.prefix + string(encoded) + "\n\n" + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a leading non-JSON-whitespace byte before the payload must not bypass the filter") + assert.Contains(t, out, "weather", + "the permitted tool must survive filtering, proving the event was filtered rather than dropped") + }) + } +} + +// TestResponseFilteringWriter_SSE_PerValueProbeCatchesConcatenatedGarbage is a +// regression test for the per-value probe added to resolveSSEEvent: one +// malformed data: value anywhere in an event used to give up filtering for +// the WHOLE event. Concretely, "garbage" followed (no blank line, so same +// event) by a genuine result-bearing frame assembles into +// "garbage\n{...result...}"; jsonrpc2.DecodeMessage stops at "garbage" and +// carriesResult on the whole joined payload does too, so both lines went out +// unfiltered. The fix probes each data: value independently before +// conceding the event needs no filtering. +// +// NotContains(admin_tool) alone would also be satisfied by the event being +// dropped silently (the #6037 hang), so this also requires the fail-closed +// error envelope to actually go out, correlated to the request's real id via +// newParsedUser1Request -- an unparsed request would let requestID() fall +// back to the zero ID and this test would accept an envelope no client could +// actually correlate. +func TestResponseFilteringWriter_SSE_PerValueProbeCatchesConcatenatedGarbage(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + resultJSON, err := json.Marshal(mcp.ListToolsResult{ + Tools: []mcp.Tool{ + {Name: "weather", Description: "Get weather information"}, + {Name: "admin_tool", Description: "Sensitive admin operations"}, + }, + }) + require.NoError(t, err) + encoded, err := jsonrpc2.EncodeMessage(&jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(1), + Result: json.RawMessage(resultJSON), + }) + require.NoError(t, err) + + testCases := []struct { + name string + garbage string + }{ + {name: "plain garbage prefix", garbage: "garbage"}, + {name: "comment-like prefix", garbage: "//comment"}, + {name: "bare NaN prefix", garbage: "NaN"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + joined := tc.garbage + "\n" + string(encoded) + _, decErr := jsonrpc2.DecodeMessage([]byte(joined)) + require.Error(t, decErr, + "test fixture assumption broken: assembled payload must not decode as a single Response") + require.False(t, carriesResult([]byte(joined)), + "test fixture assumption broken: the whole-payload scan must not itself catch the smuggled result "+ + "-- otherwise this test would not exercise the per-value probe at all") + + parsedReq := newParsedUser1Request(t, `{"jsonrpc":"2.0","id":99,"method":"tools/list"}`) + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, parsedReq, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + // One event, two data: lines with no blank line between them, so + // WHATWG assembly joins them with LF into one payload the + // decoder gives up on at the first byte. + body := "data: " + tc.garbage + "\n" + "data: " + string(encoded) + "\n\n" + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a result-bearing data: line following a malformed one in the same event must not bypass the filter") + + firstLine := strings.TrimSuffix(strings.SplitN(out, "\n", 2)[0], "\r") + payload := strings.TrimPrefix(firstLine, "data: ") + msg, decErr := jsonrpc2.DecodeMessage([]byte(payload)) + require.NoError(t, decErr, "the fail-closed envelope must itself be valid JSON-RPC") + resp, ok := msg.(*jsonrpc2.Response) + require.True(t, ok, "the fail-closed envelope must be a clean Response") + require.NotNil(t, resp.Error, "the event must fail closed into an error envelope, not vanish silently") + assert.Equal(t, jsonrpc2.Int64ID(99), resp.ID, + "the error envelope must carry the request's id so the client can correlate it, or the #6037 hang reproduces") + }) + } +} + +// TestResponseFilteringWriter_SSE_InteriorBareCRNotALineBreak is a regression +// test for a client-divergence bypass: the client that actually consumes +// these streams (github.com/modelcontextprotocol/go-sdk's mcp/event.go) +// splits lines on LF only, so a bare "\r" sitting inside a single data: +// payload is insignificant JSON whitespace to it, not a line break. Splitting +// on "\r" as well (as strict WHATWG grammar does) cuts the line there +// instead, leaving a "data:"-less second half that contributes nothing to +// the assembled payload -- so the first half (truncated, undecodable) is +// classified as carrying no result and re-emitted verbatim, leaking +// admin_tool. NotContains(admin_tool) alone would also pass if the whole +// event were merely dropped, so this also requires weather to survive, +// proving the event was filtered rather than eaten. +func TestResponseFilteringWriter_SSE_InteriorBareCRNotALineBreak(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + // A bare "\r" sits between "id":1, and "result", inside what a + // non-conformant-to-WHATWG (but client-accurate) reader treats as ONE + // data: line. "\r" is legal JSON whitespace between a comma and the next + // key, so the client parses this payload intact. + body := "data: {\"jsonrpc\":\"2.0\",\"id\":1,\r\"result\":" + + `{"tools":[{"name":"weather"},{"name":"admin_tool"}]}}` + "\n\n" + + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a bare CR inside a single data: payload must not split the line and truncate it past the filter") + assert.Contains(t, out, "weather", + "the authorized tool must survive filtering, proving the event was filtered rather than dropped") +} + +// TestResponseFilteringWriter_SSE_InteriorExoticWhitespaceAcrossDataLines is a +// regression test for a client-divergence bypass distinct from the leading- +// whitespace case above: the go-sdk client TrimSpaces each data: value +// independently before joining them with LF (mcp/event.go), but a single +// whole-buffer trim (applied once, after joining) only reaches the very +// front and back of the assembled payload. An exotic unicode.IsSpace byte +// (one of U+000B, U+000C, U+0085, U+00A0, U+1680, U+3000 -- whitespace to the +// client's trim but not to Go's own JSON scanner) sitting at the boundary +// BETWEEN two data: lines of the same event survives a whole-buffer trim, +// because it's never at either edge of the joined result. That leaves an +// illegal byte mid-token for our decoder while the client's own per-value +// trim removes it before joining, so the client parses fine while ours +// rejects the payload as undecodable and (pre-fix) re-emitted it raw. +func TestResponseFilteringWriter_SSE_InteriorExoticWhitespaceAcrossDataLines(t *testing.T) { + t.Parallel() + + exoticBytes := []struct { + name string + suffix string + }{ + {name: "U+000B line tabulation", suffix: "\u000B"}, + {name: "U+000C form feed", suffix: "\u000C"}, + {name: "U+0085 next line", suffix: "\u0085"}, + {name: "U+00A0 no-break space", suffix: "\u00A0"}, + {name: "U+1680 ogham space mark", suffix: "\u1680"}, + {name: "U+3000 ideographic space", suffix: "\u3000"}, + } + + for _, tc := range exoticBytes { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + req := newUser1Request(t) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + // Two data: lines, no blank line between them, so they belong to + // the SAME event and their values concatenate. The exotic byte + // trails the first value, sitting exactly at the interior join + // boundary once assembled -- never at the front or back of the + // whole buffer. + body := `data: {"jsonrpc":"2.0","id":1,` + tc.suffix + "\n" + + `data: "result":{"tools":[{"name":"weather"},{"name":"admin_tool"}]}}` + "\n\n" + + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "exotic whitespace at an interior data: value boundary must not bypass the filter") + assert.Contains(t, out, "weather", + "the authorized tool must survive filtering, proving the event was filtered rather than dropped") + }) + } +} + +// TestResponseFilteringWriter_SSE_FailClosedDoesNotLeakSubsequentEvents pins +// the #6037 fail-closed property that neither the event-based rewrite nor the +// TrimSpace fix's tests cover: after a mid-stream event fails to filter and +// is replaced with an error envelope, the loop must keep filtering every +// event that follows, not fall back to passing the remainder of the stream +// through raw. failingFrame is a Response carrying both "error" and "result" +// (see TestResponseFilteringWriter_SSE_ErrorAndResultBypass), a real +// reachable trigger for the fail-closed branch on this path. +func TestResponseFilteringWriter_SSE_FailClosedDoesNotLeakSubsequentEvents(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + + failingFrame := `{"jsonrpc":"2.0","id":1,"error":{"code":1,"message":"x"},` + + `"result":{"tools":[{"name":"admin_tool"}]}}` + + // The second event carries an authorized tool alongside the denied one on + // purpose. A NotContains check alone would also be satisfied by the loop + // TRUNCATING the stream after the failure — writing nothing downstream — so + // it would pass under exactly the regression this test exists to catch. + // Asserting the authorized tool survives is what proves the loop kept + // filtering rather than giving up. Do not reduce this to one assertion. + stream := "data: " + failingFrame + "\n\n" + + "data: " + `{"jsonrpc":"2.0","id":9,"result":{"tools":[{"name":"weather"},{"name":"admin_tool"}]}}` + "\n\n" + + req := newUser1Request(t) + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, req, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + _, err := rfw.Write([]byte(stream)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + out := rr.Body.String() + assert.NotContains(t, out, "admin_tool", + "a smuggled-result event after a mid-stream filter failure must still be filtered, not passed through raw") + assert.Contains(t, out, "weather", + "the authorized tool proves the loop kept filtering after the failure instead of truncating the stream") + assert.Contains(t, out, `"error"`, + "the failing event must still emit its correlatable error envelope") +} + +// sseEvent is one event dispatched by parseSSEStream. +type sseEvent struct { + eventType string + data string +} + +// parseSSEStream implements the WHATWG EventSource parsing grammar strictly, +// so tests assert actual client-visible delivery rather than merely +// conformant-looking bytes: after #6066 the error payload can look correct +// on the wire while never being dispatched to a client at all (e.g. a bare +// JSON object with no "data:" prefix). Only what this dispatches counts as +// delivered. +// +// Per the grammar: lines end at CRLF, LF, or CR; a leading ':' marks a +// comment line with no effect; a "name:value" line strips exactly one +// leading space from the value; a line with no colon is a field with an +// empty value; only "data" and "event" have any effect, every other field +// name is ignored outright; a blank line dispatches the event only if the +// data buffer is non-empty (and is reset either way); any data left +// buffered at EOF is discarded, never dispatched. +func parseSSEStream(t *testing.T, body []byte) []sseEvent { + t.Helper() + + var events []sseEvent + var eventType string + var dataLines []string + + dispatch := func() { + if len(dataLines) > 0 { + events = append(events, sseEvent{eventType: eventType, data: strings.Join(dataLines, "\n")}) + } + eventType = "" + dataLines = nil + } + + for len(body) > 0 { + idx := bytes.IndexAny(body, "\r\n") + var line []byte + switch { + case idx == -1: + line, body = body, nil + case body[idx] == '\r' && idx+1 < len(body) && body[idx+1] == '\n': + line, body = body[:idx], body[idx+2:] + default: + line, body = body[:idx], body[idx+1:] + } + + if len(line) == 0 { + dispatch() + continue + } + if line[0] == ':' { + continue + } + + field, value, hasColon := bytes.Cut(line, []byte(":")) + if !hasColon { + field, value = line, nil + } + value = bytes.TrimPrefix(value, []byte(" ")) + + switch string(field) { + case "data": + dataLines = append(dataLines, string(value)) + case "event": + eventType = string(value) + } + // Every other field name is ignored outright. + } + // Pending data at EOF is never dispatched. + + return events +} + +// TestParseSSEStream_DeliversOnlyDispatchedEvents guards the SSE error path +// with the strict grammar parser rather than a raw byte/substring check: it +// asserts the fail-closed envelope is actually dispatched as a single +// default-typed event, which is what an MCP client (default/"message" event +// handler only) will actually see. +func TestParseSSEStream_DeliversOnlyDispatchedEvents(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + parsedReq := newParsedUser1Request(t, `{"jsonrpc":"2.0","id":99,"method":"tools/list"}`) + + rr := httptest.NewRecorder() + rfw := NewResponseFilteringWriter(rr, authorizer, parsedReq, string(mcp.MethodToolsList), nil, nil) + rfw.ResponseWriter.Header().Set("Content-Type", "text/event-stream") + + body := `data: {"jsonrpc":"2.0","method":"notifications/initialized","id":1,` + + `"result":{"tools":[{"name":"admin_tool"}]}}` + "\n\n" + _, err := rfw.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, rfw.FlushAndFilter()) + + wantErrorJSON := string(rfw.errorResponseBody(jsonrpc2.Int64ID(99), + errors.New("dropped a frame carrying a result outside a clean Response"))) + + events := parseSSEStream(t, rr.Body.Bytes()) + require.Len(t, events, 1, "body: %q", rr.Body.String()) + assert.Empty(t, events[0].eventType, "MCP clients only process default/'message' events") + assert.JSONEq(t, wantErrorJSON, events[0].data) +} + +// TestParseSSEStream_RejectsBareJSONWithoutDataPrefix guards the guard: a +// bare JSON-RPC envelope with no "data:" prefix -- exactly what the pre-#6066 +// code wrote directly to the response body -- must never be recognized as a +// dispatched SSE event, with or without a trailing newline. If +// parseSSEStream ever loosened enough to accept bare JSON, this fails first. +func TestParseSSEStream_RejectsBareJSONWithoutDataPrefix(t *testing.T) { + t.Parallel() + + authorizer := newWeatherOnlyAuthorizer(t) + req := newUser1Request(t) + rfw := NewResponseFilteringWriter(httptest.NewRecorder(), authorizer, req, string(mcp.MethodToolsList), nil, nil) + + bareBody := rfw.errorResponseBody(jsonrpc2.Int64ID(1), errors.New("x")) + + testCases := []struct { + name string + body []byte + }{ + {name: "no trailing newline", body: bareBody}, + {name: "with trailing newline", body: append(append([]byte{}, bareBody...), '\n')}, + } + + for _, tc := range testCases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + events := parseSSEStream(t, tc.body) + assert.Empty(t, events, "a bare JSON envelope with no data: prefix must never dispatch as an SSE event") + }) + } } diff --git a/pkg/authz/tool_filter.go b/pkg/authz/tool_filter.go index 0ac87c4a7b..44c736714e 100644 --- a/pkg/authz/tool_filter.go +++ b/pkg/authz/tool_filter.go @@ -7,8 +7,7 @@ import ( "context" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" ) @@ -16,11 +15,9 @@ import ( // For each tool, it checks whether the caller (identified by JWT claims in ctx) // is authorized to call that tool. Only authorized tools are returned. // If authorizer is nil, all tools are returned unmodified. -// Returns an error when policy evaluation itself fails; denials still filter -// normally without failing the whole list operation. -func filterToolsByPolicy(ctx context.Context, a authorizers.Authorizer, tools []mcp.Tool) ([]mcp.Tool, error) { +func filterToolsByPolicy(ctx context.Context, a authorizers.Authorizer, tools []mcp.Tool) []mcp.Tool { if a == nil { - return tools, nil + return tools } // Note: instantiating the list ensures that no null value is sent over the wire. @@ -39,7 +36,9 @@ func filterToolsByPolicy(ctx context.Context, a authorizers.Authorizer, tools [] authorized, err := authorizeToolCall(toolCtx, a, tool.Name, nil) if err != nil { - return nil, err + slog.Warn("Authorization check failed for tool, skipping", + "tool", tool.Name, "error", err) + continue } if authorized { @@ -55,7 +54,7 @@ func filterToolsByPolicy(ctx context.Context, a authorizers.Authorizer, tools [] "total", len(tools), "allowed", len(filtered), "denied", denied) } - return filtered, nil + return filtered } // authorizeToolCall checks whether the caller is authorized to call a specific tool diff --git a/pkg/authz/tool_filter_test.go b/pkg/authz/tool_filter_test.go index c6c1bfef28..882cb9fd93 100644 --- a/pkg/authz/tool_filter_test.go +++ b/pkg/authz/tool_filter_test.go @@ -8,10 +8,10 @@ import ( "errors" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" @@ -75,7 +75,6 @@ func TestFilterToolsByPolicy(t *testing.T) { authorizer authorizers.Authorizer tools []mcp.Tool wantNames []string - wantErr error }{ { name: "nil authorizer returns all tools unchanged", @@ -107,7 +106,7 @@ func TestFilterToolsByPolicy(t *testing.T) { wantNames: []string{"allowed", "allowed2"}, }, { - name: "returns error when policy evaluation fails", + name: "skips tools where authorizer returns error", authorizer: &mockAuthorizer{results: map[string]mockResult{ "good": {authorized: true}, "error": {err: errAuth}, @@ -116,7 +115,7 @@ func TestFilterToolsByPolicy(t *testing.T) { makeTool("good", nil), makeTool("error", nil), }, - wantErr: errAuth, + wantNames: []string{"good"}, }, { name: "tools with annotations are still filtered", @@ -136,14 +135,7 @@ func TestFilterToolsByPolicy(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() - got, err := filterToolsByPolicy(context.Background(), tc.authorizer, tc.tools) - if tc.wantErr != nil { - require.Error(t, err) - assert.ErrorIs(t, err, tc.wantErr) - assert.Nil(t, got) - return - } - require.NoError(t, err) + got := filterToolsByPolicy(context.Background(), tc.authorizer, tc.tools) gotNames := make([]string, len(got)) for i, tool := range got { @@ -162,8 +154,7 @@ func TestFilterToolsByPolicy_CallsAuthorizerCorrectly(t *testing.T) { }} tools := []mcp.Tool{makeTool("tool1", nil)} - _, err := filterToolsByPolicy(context.Background(), mock, tools) - require.NoError(t, err) + filterToolsByPolicy(context.Background(), mock, tools) require.Len(t, mock.calls, 1) assert.Equal(t, authorizers.MCPFeatureTool, mock.calls[0].feature) @@ -202,8 +193,7 @@ func TestFilterToolsByPolicy_WithCedarAuthorizer(t *testing.T) { {Name: "translator", Description: "Translate text"}, } - got, err := filterToolsByPolicy(cedarCtx(t), cedarAuth, tools) - require.NoError(t, err) + got := filterToolsByPolicy(cedarCtx(t), cedarAuth, tools) require.Len(t, got, 1) assert.Equal(t, "weather", got[0].Name) @@ -218,8 +208,7 @@ func TestFilterToolsByPolicy_WithCedarAuthorizer(t *testing.T) { {Name: "translator", Description: "Translate text"}, } - got, err := filterToolsByPolicy(cedarCtx(t), cedarAuth, tools) - require.NoError(t, err) + got := filterToolsByPolicy(cedarCtx(t), cedarAuth, tools) assert.Empty(t, got) }) diff --git a/pkg/config/registry.go b/pkg/config/registry.go index cb0f1cb2cc..a19d6c9c37 100644 --- a/pkg/config/registry.go +++ b/pkg/config/registry.go @@ -177,9 +177,9 @@ func isValidRegistryJSON(client *http.Client, url string) error { } // parseAndValidateUpstreamRegistry parses data as an UpstreamRegistry and -// verifies that it contains at least one server or skill. If the data looks -// like the legacy ToolHive registry format, returns an error containing the -// migration hint message. +// verifies that it contains at least one server, skill, or plugin. If the data +// looks like the legacy ToolHive registry format, returns an error containing +// the migration hint message. func parseAndValidateUpstreamRegistry(data []byte) error { if legacyhint.Looks(data) { return errors.New(legacyhint.MigrationMessage) @@ -188,8 +188,8 @@ func parseAndValidateUpstreamRegistry(data []byte) error { if err := json.Unmarshal(data, &upstream); err != nil { return fmt.Errorf("invalid upstream registry format: %w", err) } - if len(upstream.Data.Servers) == 0 && len(upstream.Data.Skills) == 0 { - return errors.New("upstream registry contains no servers or skills") + if len(upstream.Data.Servers) == 0 && len(upstream.Data.Skills) == 0 && len(upstream.Data.Plugins) == 0 { + return errors.New("upstream registry contains no servers, skills, or plugins") } return nil } diff --git a/pkg/config/registry_test.go b/pkg/config/registry_test.go index d9b4476a18..ec739914e3 100644 --- a/pkg/config/registry_test.go +++ b/pkg/config/registry_test.go @@ -311,7 +311,7 @@ func TestValidateRegistryFileStructure_UpstreamFormat(t *testing.T) { "data": {"servers": []} }`, expectError: true, - errMsgContains: "no servers or skills", + errMsgContains: "no servers, skills, or plugins", }, { name: "legacy format with top-level servers returns migration hint", diff --git a/pkg/container/docker/client.go b/pkg/container/docker/client.go index e85142682e..99e038747f 100644 --- a/pkg/container/docker/client.go +++ b/pkg/container/docker/client.go @@ -12,6 +12,7 @@ import ( "fmt" "io" "log/slog" + "maps" "net/netip" "os" "path/filepath" @@ -99,16 +100,6 @@ type deployOps interface { networkName string, endpointsConfig map[string]*network.EndpointSettings, ) (string, string, error) - createEgressSquidContainer( - ctx context.Context, - containerName string, - squidContainerName string, - attachStdio bool, - exposedPorts map[string]struct{}, - endpointsConfig map[string]*network.EndpointSettings, - perm *permissions.NetworkPermissions, - allowDockerGateway bool, - ) (string, error) createMcpContainer( ctx context.Context, name string, @@ -124,14 +115,7 @@ type deployOps interface { portBindings map[string][]runtime.PortBinding, isolateNetwork bool, ) error - createIngressContainer( - ctx context.Context, - containerName string, - upstreamPort int, - attachStdio bool, - externalEndpointsConfig map[string]*network.EndpointSettings, - networkPermissions *permissions.NetworkPermissions, - ) (int, error) + getContainerNetworkIP(ctx context.Context, containerName, networkName string) (string, error) } // Client implements the Deployer interface for Docker (and compatible runtimes) @@ -142,6 +126,7 @@ type Client struct { api dockerAPI imageManager images.ImageManager ops deployOps + proxy networkProxy // selected at construction time via newNetworkProxy } // NewClient creates a new container client @@ -163,26 +148,13 @@ func NewClient(ctx context.Context) (*Client, error) { // Default ops implementation uses the real client methods. c.ops = c - return c, nil -} + proxy, err := newNetworkProxy(c) + if err != nil { + return nil, fmt.Errorf("failed to initialize network proxy: %w", err) + } + c.proxy = proxy -// createEgressSquidContainer wraps the package-level createEgressSquidContainer to satisfy deployOps. -func (c *Client) createEgressSquidContainer( - ctx context.Context, - containerName string, - squidContainerName string, - attachStdio bool, - exposedPorts map[string]struct{}, - endpointsConfig map[string]*network.EndpointSettings, - perm *permissions.NetworkPermissions, - allowDockerGateway bool, -) (string, error) { - gatewayIP := c.getDockerBridgeGatewayIP(ctx) - return createEgressSquidContainer( - ctx, c, containerName, squidContainerName, - attachStdio, exposedPorts, endpointsConfig, perm, - allowDockerGateway, gatewayIP, - ) + return c, nil } // DeployWorkload creates and starts a workload. @@ -222,8 +194,25 @@ func (c *Client) DeployWorkload( networkName: {}, } + // Network isolation is only enforceable in bridge mode. For non-bridge + // modes (host/none/custom) the sidecars have no route out and the workload + // bypasses them, so isolation must be dropped. See issue #5775. + effectiveIsolation := isolateNetwork && networking.IsBridgeMode(permissionConfig.NetworkMode) + if isolateNetwork && !effectiveIsolation { + if permissionConfig.NetworkMode == "none" { + // "none" is already maximally confined; isolation is merely redundant. + //nolint:gosec // G706: network mode from permission config + slog.Debug(networking.NetworkIsolationNoneRedundantMsg, "network_mode", permissionConfig.NetworkMode) + } else { + // host (or other non-bridge) is less restrictive than isolation, so + // dropping it is a real reduction in confinement. + //nolint:gosec // G706: network mode from permission config + slog.Warn(networking.NetworkIsolationHostDroppedMsg, "network_mode", permissionConfig.NetworkMode) + } + } + // Only create external networks and add endpoints if we're not using a custom network mode like "host" or "none" - if permissionConfig.NetworkMode == "" || permissionConfig.NetworkMode == "bridge" || permissionConfig.NetworkMode == "default" { + if networking.IsBridgeMode(permissionConfig.NetworkMode) { // Add toolhive-external to endpoints config for default networking modes externalEndpointsConfig["toolhive-external"] = &network.EndpointSettings{} @@ -236,8 +225,24 @@ func (c *Client) DeployWorkload( slog.Debug("skipping external network creation for custom network mode", "network_mode", permissionConfig.NetworkMode) } + // For non-stdio isolated workloads, extract the upstream port so the ingress + // proxy can be configured. Done here so a bad exposed-ports config fails fast. + var upstreamPort int + if transportType != "stdio" && effectiveIsolation { + upstreamPort, err = extractFirstPort(options) + if err != nil { + return 0, err // extractFirstPort already wraps the error with context. + } + } + networkIsolation := false - if isolateNetwork { + // pspec and egress are populated in the isolation block and reused by + // SetupIngress after the MCP container is created. + var ( + pspec proxySpec + egress egressResult + ) + if effectiveIsolation { networkIsolation = true internalNetworkLabels := map[string]string{} @@ -257,24 +262,24 @@ func (c *Client) DeployWorkload( return 0, fmt.Errorf("failed to create dns container: %w", err) } - // create egress container - egressContainerName := fmt.Sprintf("%s-egress", name) - allowDockerGateway := options != nil && options.AllowDockerGateway - _, err = c.ops.createEgressSquidContainer( - ctx, - name, - egressContainerName, - attachStdio, - nil, - externalEndpointsConfig, - permissionProfile.Network, - allowDockerGateway, - ) - if err != nil { - return 0, fmt.Errorf("failed to create egress container: %w", err) + pspec = proxySpec{ + WorkloadName: name, + Permissions: permissionProfile.Network, + AllowDockerGateway: options != nil && options.AllowDockerGateway, + GatewayIP: c.getDockerBridgeGatewayIP(ctx), + TransportType: transportType, + UpstreamPort: upstreamPort, + AttachStdio: attachStdio, + Endpoints: externalEndpointsConfig, } - envVars = addEgressEnvVars(envVars, egressContainerName) + // SetupEgress runs before createMcpContainer so its env vars can be + // injected into the workload and the egress proxy has a head start. + egress, err = c.proxy.SetupEgress(ctx, pspec) + if err != nil { + return 0, fmt.Errorf("failed to set up egress proxy: %w", err) + } + envVars = mergeEnvVars(envVars, egress.EnvVars) } else { networkName = "" } @@ -303,7 +308,7 @@ func (c *Client) DeployWorkload( additionalDNS, options.ExposedPorts, newPortBindings, - isolateNetwork, + effectiveIsolation, ) if err != nil { return 0, fmt.Errorf("failed to create mcp container: %w", err) @@ -318,12 +323,21 @@ func (c *Client) DeployWorkload( if err != nil { return 0, err // extractFirstPort already wraps the error with context. } - if isolateNetwork { - // just extract the first exposed port - hostPort, err = c.ops.createIngressContainer(ctx, name, firstPortInt, attachStdio, externalEndpointsConfig, - permissionProfile.Network) + if effectiveIsolation { + // Point the ingress reverse proxy at the MCP container's IP rather than + // its name. Ordering ingress after createMcpContainer is not enough: under + // concurrent startup the container's DNS record can still lag, and the + // squid cache_peer caches the failed lookup and never recovers within the + // readiness window (see #6063). An IP has no lookup, so a transient + // connection failure is retried instead of latching the peer dead. + upstreamIP, ipErr := c.ops.getContainerNetworkIP(ctx, name, networkName) + if ipErr != nil { + return 0, fmt.Errorf("failed to resolve MCP container IP for ingress: %w", ipErr) + } + pspec.UpstreamHost = upstreamIP + hostPort, err = c.proxy.SetupIngress(ctx, pspec, egress) if err != nil { - return 0, fmt.Errorf("failed to create ingress container: %w", err) + return 0, fmt.Errorf("failed to set up ingress proxy: %w", err) } } @@ -1179,18 +1193,49 @@ func (c *Client) createNetwork( _, err = c.client.NetworkCreate(ctx, name, networkCreate) if err != nil { + // The existence check above is not atomic with the create: another + // process (e.g. a second `thv run` sharing the "toolhive-external" + // network) can create the same network in between. Docker rejects a + // duplicate name with a conflict error, so treat that as success — the + // network is already in the state we wanted. Without this, workloads + // that start concurrently race here and all but one fail to launch. + if errdefs.IsConflict(err) { + return nil + } return err } return nil } +// getContainerNetworkIP returns a container's IP address on the named network. +// Used to give the ingress proxy a DNS-free upstream address (see proxySpec.UpstreamHost). +func (c *Client) getContainerNetworkIP(ctx context.Context, containerName, networkName string) (string, error) { + inspectResult, err := c.client.ContainerInspect(ctx, containerName, mobyclient.ContainerInspectOptions{}) + if err != nil { + return "", fmt.Errorf("failed to inspect container %s: %w", containerName, err) + } + netSettings, ok := inspectResult.Container.NetworkSettings.Networks[networkName] + if !ok { + return "", fmt.Errorf("network %s not found in container %s network settings", networkName, containerName) + } + // EndpointSettings.IPAddress is a netip.Addr; render it to string form. + if !netSettings.IPAddress.IsValid() { + return "", fmt.Errorf("container %s has no IP address on network %s", containerName, networkName) + } + return netSettings.IPAddress.String(), nil +} + // getDockerBridgeGatewayIP returns the gateway IP of the Docker default bridge // network by inspecting it at runtime. This handles platform differences: // Linux Docker Engine uses 172.17.0.1 by default, while Docker Desktop on macOS // uses 192.168.65.1 and Colima typically uses 192.168.5.1 or similar. Querying // the daemon directly is more accurate than hardcoding platform-specific IPs. -// Falls back to dockerDefaultBridgeGatewayIP on any error. +// Falls back to dockerDefaultBridgeGatewayIP on any error or when the underlying +// Docker client is unavailable. func (c *Client) getDockerBridgeGatewayIP(ctx context.Context) string { + if c.client == nil { + return dockerDefaultBridgeGatewayIP + } nr, err := c.client.NetworkInspect(ctx, "bridge", mobyclient.NetworkInspectOptions{}) if err != nil { slog.Debug("failed to inspect bridge network, using default gateway IP", "error", err) @@ -1413,7 +1458,34 @@ func (c *Client) createContainer( EndpointsConfig: endpointsConfig, } - // Create the container + id, err := c.createAndStartContainer(ctx, containerName, config, hostConfig, networkConfig) + if err != nil && errdefs.IsNotFound(err) { + if _, sharesExternalNetwork := endpointsConfig["toolhive-external"]; sharesExternalNetwork { + // "toolhive-external" is shared, concurrency-created infrastructure + // reused by every workload; deleteNetworks removes it opportunistically + // once it looks like no workload needs it anymore. Under concurrent + // `thv` processes that teardown can race this container's own creation, + // so the network can vanish between our create-time existence check and + // this container actually attaching to it, surfacing as a not-found + // error from either Create or Start. Recreate it and retry once instead + // of failing the whole workload for what is a transient condition. + if recreateErr := c.createExternalNetworks(ctx); recreateErr == nil { + id, err = c.createAndStartContainer(ctx, containerName, config, hostConfig, networkConfig) + } + } + } + return id, err +} + +// createAndStartContainer creates and starts a container, wrapping any +// Docker API error with contextual information. +func (c *Client) createAndStartContainer( + ctx context.Context, + containerName string, + config *container.Config, + hostConfig *container.HostConfig, + networkConfig *network.NetworkingConfig, +) (string, error) { resp, err := c.api.ContainerCreate(ctx, mobyclient.ContainerCreateOptions{ Config: config, HostConfig: hostConfig, @@ -1432,6 +1504,12 @@ func (c *Client) createContainer( // Start the container _, err = c.api.ContainerStart(ctx, resp.ID, mobyclient.ContainerStartOptions{}) if err != nil { + // Docker leaves the container behind in a "created" state on a failed + // start (e.g. its network vanished underneath it). Clean it up so a + // caller-driven retry under the same name doesn't collide with it. + if _, rmErr := c.api.ContainerRemove(ctx, resp.ID, mobyclient.ContainerRemoveOptions{Force: true}); rmErr != nil { + slog.Debug("failed to remove container after failed start", "id", resp.ID, "error", rmErr) + } return "", NewContainerError(err, resp.ID, fmt.Sprintf("failed to start container: %v", err)) } @@ -1568,7 +1646,7 @@ func (c *Client) createMcpContainer( internalEndpointsConfig := map[string]*network.EndpointSettings{} // Check if we have a custom network mode (e.g., "host", "none", etc.) - if permissionConfig.NetworkMode != "" && permissionConfig.NetworkMode != "bridge" && permissionConfig.NetworkMode != "default" { + if !networking.IsBridgeMode(permissionConfig.NetworkMode) { // For custom network modes like "host", "none", etc., don't add any endpoint configurations // The NetworkMode in hostConfig will handle the networking //nolint:gosec // G706: network mode from permission config @@ -1593,26 +1671,48 @@ func (c *Client) createMcpContainer( } -// addEgressEnvVars adds environment variables for egress proxy configuration. +// addEgressEnvVars returns a new map containing all entries from envVars plus +// the HTTP_PROXY/HTTPS_PROXY/NO_PROXY variables for the given egress container. +// The caller's map is never mutated. func addEgressEnvVars(envVars map[string]string, egressContainerName string) map[string]string { egressHost := fmt.Sprintf("http://%s:3128", egressContainerName) - if envVars == nil { - envVars = make(map[string]string) - } - envVars["HTTP_PROXY"] = egressHost - envVars["HTTPS_PROXY"] = egressHost - envVars["http_proxy"] = egressHost - envVars["https_proxy"] = egressHost - envVars["NO_PROXY"] = "localhost,127.0.0.1,::1" - envVars["no_proxy"] = "localhost,127.0.0.1,::1" - return envVars + result := maps.Clone(envVars) + if result == nil { + result = make(map[string]string) + } + result["HTTP_PROXY"] = egressHost + result["HTTPS_PROXY"] = egressHost + result["http_proxy"] = egressHost + result["https_proxy"] = egressHost + result["NO_PROXY"] = "localhost,127.0.0.1,::1" + result["no_proxy"] = "localhost,127.0.0.1,::1" + return result +} + +// mergeEnvVars returns a new map containing all entries from base with all +// entries from extra added (extra values overwrite base on conflict). Neither +// input map is mutated. +func mergeEnvVars(base, extra map[string]string) map[string]string { + result := make(map[string]string, len(base)+len(extra)) + for k, v := range base { + result[k] = v + } + for k, v := range extra { + result[k] = v + } + return result } -func (c *Client) createIngressContainer(ctx context.Context, containerName string, upstreamPort int, attachStdio bool, - externalEndpointsConfig map[string]*network.EndpointSettings, networkPermissions *permissions.NetworkPermissions) (int, error) { - squidPort, err := networking.FindOrUsePort(upstreamPort + 1) +// setupIngressContainer creates the ingress Squid reverse-proxy container for +// the workload and returns the host-side port it is bound on. +func (c *Client) setupIngressContainer(ctx context.Context, containerName, upstreamHost string, upstreamPort int, + attachStdio bool, externalEndpointsConfig map[string]*network.EndpointSettings, + networkPermissions *permissions.NetworkPermissions) (int, error) { + // A random port avoids every same-image workload racing to bind the same + // preferred port when starting concurrently (see #6063). + squidPort, err := networking.FindOrUsePort(0) if err != nil { - return 0, fmt.Errorf("failed to find or use port %d: %w", squidPort, err) + return 0, fmt.Errorf("failed to find an available port: %w", err) } squidExposedPorts := map[string]struct{}{ fmt.Sprintf("%d/tcp", squidPort): {}, @@ -1630,6 +1730,7 @@ func (c *Client) createIngressContainer(ctx context.Context, containerName strin ctx, c, containerName, + upstreamHost, ingressContainerName, attachStdio, upstreamPort, @@ -1677,6 +1778,8 @@ func (c *Client) createExternalNetworks(ctx context.Context) error { func generatePortBindings(labels map[string]string, portBindings map[string][]runtime.PortBinding) (map[string][]runtime.PortBinding, int, error) { + // Clone portBindings so we never mutate the caller's map. + portBindings = maps.Clone(portBindings) var hostPort int // check if we need to map to a random port of not if _, ok := labels[ToolhiveAuxiliaryWorkloadLabel]; ok && labels[ToolhiveAuxiliaryWorkloadLabel] == LabelValueTrue { diff --git a/pkg/container/docker/client_deploy_test.go b/pkg/container/docker/client_deploy_test.go index 3b1a98f223..d61abb3ca8 100644 --- a/pkg/container/docker/client_deploy_test.go +++ b/pkg/container/docker/client_deploy_test.go @@ -5,6 +5,7 @@ package docker import ( "context" + "fmt" "testing" "github.com/moby/moby/api/types/network" @@ -31,12 +32,10 @@ type fakeDeployOps struct { dnsID string dnsIP string - egressCalled bool - egressID string - egressAllowDockerGW bool - - ingressCalled bool - ingressPort int + // mcpIP is returned by getContainerNetworkIP as the MCP container's IP that + // the ingress proxy targets. errMcpIP injects a lookup failure. + mcpIP string + errMcpIP error mcpCalled bool mcpName string @@ -52,12 +51,13 @@ type fakeDeployOps struct { mcpPortBindings map[string][]runtime.PortBinding mcpIsolate bool + // callOrder tracks the sequence of operation calls for ordering assertions. + callOrder *[]string + // error injection errExternalNetworks error errCreateNetwork error errDNS error - errEgress error - errIngress error errMcp error } @@ -80,10 +80,8 @@ func (f *fakeDeployOps) createDnsContainer(_ context.Context, _ string, _ bool, return f.dnsID, f.dnsIP, f.errDNS } -func (f *fakeDeployOps) createEgressSquidContainer(_ context.Context, _ string, _ string, _ bool, _ map[string]struct{}, _ map[string]*network.EndpointSettings, _ *permissions.NetworkPermissions, allowDockerGateway bool) (string, error) { - f.egressCalled = true - f.egressAllowDockerGW = allowDockerGateway - return f.egressID, f.errEgress +func (f *fakeDeployOps) getContainerNetworkIP(_ context.Context, _, _ string) (string, error) { + return f.mcpIP, f.errMcpIP } func (f *fakeDeployOps) createMcpContainer( @@ -101,6 +99,9 @@ func (f *fakeDeployOps) createMcpContainer( portBindings map[string][]runtime.PortBinding, isolateNetwork bool, ) error { + if f.callOrder != nil { + *f.callOrder = append(*f.callOrder, "createMcpContainer") + } f.mcpCalled = true f.mcpName = name f.mcpNetworkName = networkName @@ -117,22 +118,58 @@ func (f *fakeDeployOps) createMcpContainer( return f.errMcp } -func (f *fakeDeployOps) createIngressContainer(_ context.Context, _ string, _ int, _ bool, _ map[string]*network.EndpointSettings, _ *permissions.NetworkPermissions) (int, error) { +// fakeNetworkProxy implements networkProxy for testing DeployWorkload without real proxy containers. +type fakeNetworkProxy struct { + setupCalled bool // SetupEgress was called + ingressCalled bool // SetupIngress was called + capturedSpec proxySpec + ingressPort int // returned by SetupIngress + err error + // callOrder tracks cross-component ordering when shared with fakeDeployOps. + callOrder *[]string +} + +func (f *fakeNetworkProxy) SetupEgress(_ context.Context, spec proxySpec) (egressResult, error) { + if f.callOrder != nil { + *f.callOrder = append(*f.callOrder, "SetupEgress") + } + f.setupCalled = true + f.capturedSpec = spec + if f.err != nil { + return egressResult{}, f.err + } + // Return realistic env vars based on spec so MCP container env var assertions remain meaningful. + egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName) + return egressResult{EnvVars: addEgressEnvVars(nil, egressContainerName)}, nil +} + +func (f *fakeNetworkProxy) SetupIngress(_ context.Context, spec proxySpec, _ egressResult) (int, error) { + if f.callOrder != nil { + *f.callOrder = append(*f.callOrder, "SetupIngress") + } f.ingressCalled = true - if f.errIngress != nil { - return 0, f.errIngress + f.capturedSpec = spec + if f.err != nil { + return 0, f.err } return f.ingressPort, nil } -// newClientWithOps creates a minimal client with the provided ops and a fake dockerAPI. -func newClientWithOps(ops deployOps) *Client { +// newClientWithOpsAndProxy creates a minimal client with the provided ops, proxy, and a fake dockerAPI. +func newClientWithOpsAndProxy(ops deployOps, proxy networkProxy) *Client { return &Client{ - api: opsToFakeDockerAPI(), - ops: ops, + api: opsToFakeDockerAPI(), + ops: ops, + proxy: proxy, } } +// newClientWithOps creates a minimal client with the provided ops and a zero fakeNetworkProxy. +// Use this for tests that do not exercise isolated-network paths. +func newClientWithOps(ops deployOps) *Client { + return newClientWithOpsAndProxy(ops, &fakeNetworkProxy{}) +} + // opsToFakeDockerAPI returns a fake dockerAPI that won't be used by DeployWorkload tests directly. func opsToFakeDockerAPI() dockerAPI { return &fakeDockerAPI{} @@ -142,10 +179,10 @@ func TestDeployWorkload_Stdio_IsolatedNetwork_SkipsIngressAndSetsEgressEnv(t *te t.Parallel() fops := &fakeDeployOps{ - dnsIP: "172.18.0.10", - ingressPort: 18080, // should be ignored for stdio + dnsIP: "172.18.0.10", } - c := newClientWithOps(fops) + fproxy := &fakeNetworkProxy{} + c := newClientWithOpsAndProxy(fops, fproxy) opts := runtime.NewDeployWorkloadOptions() opts.AttachStdio = true @@ -179,8 +216,13 @@ func TestDeployWorkload_Stdio_IsolatedNetwork_SkipsIngressAndSetsEgressEnv(t *te require.Len(t, fops.createNetworkCalls, 1) assert.True(t, fops.createNetworkCalls[0].internal) assert.True(t, fops.dnsCalled) - assert.True(t, fops.egressCalled) - assert.False(t, fops.ingressCalled) + + // Proxy must be invoked for isolated-network deployment + assert.True(t, fproxy.setupCalled) + // stdio transport: no ingress needed + assert.Equal(t, "stdio", fproxy.capturedSpec.TransportType) + // AllowDockerGateway was not set + assert.False(t, fproxy.capturedSpec.AllowDockerGateway) // MCP container created with egress env vars present require.True(t, fops.mcpCalled) @@ -194,24 +236,25 @@ func TestDeployWorkload_Stdio_IsolatedNetwork_SkipsIngressAndSetsEgressEnv(t *te // SELinux labeling should be disabled assert.Contains(t, fops.mcpPermissionCfg.SecurityOpt, "label:disable", "expected SELinux labeling to be disabled") - - // TODO: Test for disabled SELinux labeling in the rest of workload containers } func TestDeployWorkload_SSE_IsolatedNetwork_ReturnsIngressPortAndPassesDNS(t *testing.T) { t.Parallel() fops := &fakeDeployOps{ - dnsIP: "172.18.0.20", + dnsIP: "172.18.0.20", + mcpIP: "172.18.0.21", + } + fproxy := &fakeNetworkProxy{ ingressPort: 18081, } - c := newClientWithOps(fops) + c := newClientWithOpsAndProxy(fops, fproxy) opts := runtime.NewDeployWorkloadOptions() opts.ExposedPorts = map[string]struct{}{"8080/tcp": {}} opts.PortBindings = map[string][]runtime.PortBinding{ "8080/tcp": { - {HostIP: "127.0.0.1", HostPort: ""}, // random/non-deterministic is fine; will be overridden by ingress + {HostIP: "127.0.0.1", HostPort: ""}, }, } @@ -231,9 +274,14 @@ func TestDeployWorkload_SSE_IsolatedNetwork_ReturnsIngressPortAndPassesDNS(t *te ) require.NoError(t, err) - // For non-stdio with network isolation, returned port comes from ingress proxy + // For non-stdio with network isolation, returned port comes from ingress proxy. assert.Equal(t, 18081, hostPort) - assert.True(t, fops.ingressCalled) + assert.True(t, fproxy.setupCalled) + // The upstream port should be the first exposed port (8080). + assert.Equal(t, 8080, fproxy.capturedSpec.UpstreamPort) + // The ingress proxy must target the MCP container's IP (not its name) so it + // has no DNS dependency to latch on under concurrent startup (see #6063). + assert.Equal(t, "172.18.0.21", fproxy.capturedSpec.UpstreamHost) require.True(t, fops.mcpCalled) assert.Equal(t, "172.18.0.20", fops.mcpAdditionalDNS, "additionalDNS passed to MCP container should come from DNS container IP") } @@ -270,10 +318,8 @@ func TestDeployWorkload_NoIsolation_ReturnsPortFromBindingsAndSkipsAuxContainers ) require.NoError(t, err) - // Should not create internal network, DNS, egress, or ingress + // Should not create internal network, DNS, or invoke the proxy assert.False(t, fops.dnsCalled) - assert.False(t, fops.egressCalled) - assert.False(t, fops.ingressCalled) assert.Empty(t, fops.createNetworkCalls, "internal network should not be created when isolation is disabled") // MCP should be created on default network (empty name) @@ -288,7 +334,8 @@ func TestDeployWorkload_AllowDockerGateway_ForwardedToEgress(t *testing.T) { t.Parallel() fops := &fakeDeployOps{dnsIP: "172.18.0.10"} - c := newClientWithOps(fops) + fproxy := &fakeNetworkProxy{} + c := newClientWithOpsAndProxy(fops, fproxy) opts := runtime.NewDeployWorkloadOptions() opts.AttachStdio = true @@ -304,12 +351,12 @@ func TestDeployWorkload_AllowDockerGateway_ForwardedToEgress(t *testing.T) { &permissions.Profile{}, "stdio", opts, - true, // isolateNetwork required for egress container to be created + true, // isolateNetwork required for proxy to be invoked ) require.NoError(t, err) - require.True(t, fops.egressCalled, "egress container must be created when isolateNetwork=true") - assert.True(t, fops.egressAllowDockerGW, "AllowDockerGateway must be forwarded to createEgressSquidContainer") + require.True(t, fproxy.setupCalled, "proxy must be set up when isolateNetwork=true") + assert.True(t, fproxy.capturedSpec.AllowDockerGateway, "AllowDockerGateway must be forwarded to SetupProxies") } // TestDeployWorkload_AllowDockerGateway_DefaultsToNotForwarded guards against a @@ -320,7 +367,8 @@ func TestDeployWorkload_AllowDockerGateway_DefaultsToNotForwarded(t *testing.T) t.Parallel() fops := &fakeDeployOps{dnsIP: "172.18.0.10"} - c := newClientWithOps(fops) + fproxy := &fakeNetworkProxy{} + c := newClientWithOpsAndProxy(fops, fproxy) opts := runtime.NewDeployWorkloadOptions() opts.AttachStdio = true @@ -336,16 +384,90 @@ func TestDeployWorkload_AllowDockerGateway_DefaultsToNotForwarded(t *testing.T) &permissions.Profile{}, "stdio", opts, - true, // isolateNetwork required for egress container to be created + true, // isolateNetwork required for the proxy to be set up ) require.NoError(t, err) - require.True(t, fops.egressCalled, "egress container must be created when isolateNetwork=true") - assert.False(t, fops.egressAllowDockerGW, + require.True(t, fproxy.setupCalled, "network proxy must be set up when isolateNetwork=true") + assert.False(t, fproxy.capturedSpec.AllowDockerGateway, "AllowDockerGateway must default to false so the gateway deny rules stay in place") assert.True(t, fops.dnsCalled, "DNS container must be created on the isolation path") } +// TestDeployWorkload_NonBridgeNetwork_DropsIsolation verifies that when +// isolateNetwork=true but the network mode is non-bridge (host/none), the +// isolation sidecars are never created, no proxy env vars are injected, and the +// network-isolation label records the effective (false) value. Building the +// sidecars on the internal network for a host/none workload would silently break +// all outbound traffic. See issue #5775. +func TestDeployWorkload_NonBridgeNetwork_DropsIsolation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + networkMode string + }{ + {"host mode", "host"}, + {"none mode", "none"}, + {"custom mode", "container:foo"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + fops := &fakeDeployOps{dnsIP: "172.18.0.10"} + fproxy := &fakeNetworkProxy{ingressPort: 18081} + c := newClientWithOpsAndProxy(fops, fproxy) + + opts := runtime.NewDeployWorkloadOptions() + opts.ExposedPorts = map[string]struct{}{"8080/tcp": {}} + opts.PortBindings = map[string][]runtime.PortBinding{ + "8080/tcp": {{HostIP: "127.0.0.1", HostPort: "12345"}}, + } + + labels := map[string]string{} + env := map[string]string{"EXISTING": "1"} + + _, err := c.DeployWorkload( + t.Context(), + "ghcr.io/example/mcp:latest", + "app", + []string{"serve"}, + env, + labels, + &permissions.Profile{ + Network: &permissions.NetworkPermissions{Mode: tt.networkMode}, + }, + "sse", + opts, + true, // isolateNetwork requested, but must be dropped for non-bridge modes + ) + require.NoError(t, err) + + // No isolation sidecars should be created. + assert.False(t, fops.dnsCalled, "DNS container must not be created for non-bridge modes") + assert.False(t, fproxy.setupCalled, "network proxy must not be set up for non-bridge modes") + assert.Empty(t, fops.createNetworkCalls, "internal network must not be created for non-bridge modes") + + // The external network is a bridge-only construct. + assert.False(t, fops.externalNetworksCalled, "external network must not be created for non-bridge modes") + + // No proxy env vars should be injected. + require.True(t, fops.mcpCalled) + assert.NotContains(t, fops.mcpEnvVars, "HTTP_PROXY", "HTTP_PROXY must not be injected without isolation") + assert.NotContains(t, fops.mcpEnvVars, "HTTPS_PROXY", "HTTPS_PROXY must not be injected without isolation") + + // createMcpContainer must receive the effective (false) isolation value. + assert.False(t, fops.mcpIsolate, "createMcpContainer must receive effective isolation=false") + + // The network-isolation label must record the effective value. + assert.False(t, lb.HasNetworkIsolation(labels), + "network isolation label must record effective (false) value for non-bridge modes") + }) + } +} + func TestDeployWorkload_UnsupportedTransport_PropagatesError(t *testing.T) { t.Parallel() @@ -375,3 +497,46 @@ func TestDeployWorkload_UnsupportedTransport_PropagatesError(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "unsupported transport type") } + +// TestDeployWorkload_Isolated_EgressBeforeMcpIngressAfter locks in the ordering +// contract: egress is provisioned BEFORE the MCP container (so proxy env vars are +// injected), and ingress is provisioned AFTER it. Creating the ingress before the +// MCP container caused the squid reverse proxy to cache a negative DNS lookup for +// the not-yet-existent upstream and never recover, leaving the workload stuck. +func TestDeployWorkload_Isolated_EgressBeforeMcpIngressAfter(t *testing.T) { + t.Parallel() + + callOrder := make([]string, 0, 3) + fops := &fakeDeployOps{ + dnsIP: "172.18.0.10", + callOrder: &callOrder, + } + fproxy := &fakeNetworkProxy{ + ingressPort: 18081, + callOrder: &callOrder, + } + c := newClientWithOpsAndProxy(fops, fproxy) + + opts := runtime.NewDeployWorkloadOptions() + opts.ExposedPorts = map[string]struct{}{"8080/tcp": {}} + opts.PortBindings = map[string][]runtime.PortBinding{ + "8080/tcp": {{HostIP: "127.0.0.1", HostPort: ""}}, + } + + _, err := c.DeployWorkload( + t.Context(), + "ghcr.io/example/mcp:latest", + "app", + []string{"serve"}, + map[string]string{}, + map[string]string{}, + &permissions.Profile{}, + "sse", + opts, + true, + ) + require.NoError(t, err) + + require.Equal(t, []string{"SetupEgress", "createMcpContainer", "SetupIngress"}, callOrder, + "egress must be set up before the MCP container and ingress after it") +} diff --git a/pkg/container/docker/client_network_test.go b/pkg/container/docker/client_network_test.go new file mode 100644 index 0000000000..0c990aecbd --- /dev/null +++ b/pkg/container/docker/client_network_test.go @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package docker + +import ( + "context" + "fmt" + "os" + "sync" + "testing" + + mobyclient "github.com/moby/moby/client" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/container/docker/sdk" +) + +// TestCreateNetwork_ConcurrentSameName reproduces the shared-network startup +// race: multiple workloads (e.g. several `thv run` sharing "toolhive-external") +// call createNetwork for the same name at once. The existence check and the +// create are not atomic, so without idempotent conflict handling all but one +// caller fails with "network ... already exists" and its workload never +// launches. This test drives that exact contention against a real daemon and +// asserts every caller succeeds with exactly one network created. +// +// Skips when Docker is unavailable, matching the real-daemon tests in this +// package (see envoy_test.go). +func TestCreateNetwork_ConcurrentSameName(t *testing.T) { + t.Parallel() + + ctx := context.Background() + mc, _, _, err := sdk.NewDockerClient(ctx) + if err != nil { + t.Skipf("docker not available; skipping network create race test: %v", err) + } + c := &Client{client: mc} + + // Unique per test process so parallel runs / leftovers never collide. + name := fmt.Sprintf("toolhive-test-netrace-%d", os.Getpid()) + labels := map[string]string{"toolhive": "true"} + t.Cleanup(func() { removeNetworksByName(context.Background(), mc, name) }) + removeNetworksByName(ctx, mc, name) // start clean + + const workers = 8 + var wg sync.WaitGroup + start := make(chan struct{}) + errs := make([]error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start // release all goroutines together to maximize contention + errs[i] = c.createNetwork(ctx, name, labels, false) + }(i) + } + close(start) + wg.Wait() + + for i, e := range errs { + require.NoErrorf(t, e, "worker %d: concurrent createNetwork must treat a lost race as success", i) + } + + require.Equal(t, 1, countNetworksByName(ctx, mc, name), + "exactly one shared network should exist after the concurrent race") +} + +func countNetworksByName(ctx context.Context, mc *mobyclient.Client, name string) int { + networks, err := mc.NetworkList(ctx, mobyclient.NetworkListOptions{ + Filters: mobyclient.Filters{}.Add("name", name), + }) + if err != nil { + return -1 + } + n := 0 + for _, item := range networks.Items { + if item.Name == name { + n++ + } + } + return n +} + +func removeNetworksByName(ctx context.Context, mc *mobyclient.Client, name string) { + networks, err := mc.NetworkList(ctx, mobyclient.NetworkListOptions{ + Filters: mobyclient.Filters{}.Add("name", name), + }) + if err != nil { + return + } + for _, item := range networks.Items { + if item.Name == name { + _, _ = mc.NetworkRemove(ctx, item.ID, mobyclient.NetworkRemoveOptions{}) + } + } +} diff --git a/pkg/container/docker/client_stop_test.go b/pkg/container/docker/client_stop_test.go index 41a42dd46b..4bc00df183 100644 --- a/pkg/container/docker/client_stop_test.go +++ b/pkg/container/docker/client_stop_test.go @@ -138,3 +138,69 @@ func TestStopWorkload_NotFound_ReturnsNil(t *testing.T) { // StopWorkload should treat a not-found workload as success require.NoError(t, err) } + +// TestStopProxyContainer_HandlesEnvoyAndSquidTopologies exercises the building +// block of the shared cleanup path. StopWorkload iterates the fixed proxy +// suffixes (-egress, -ingress, -dns) and calls stopProxyContainer for each; this +// test verifies stopProxyContainer's per-container contract directly: stop a +// container that exists, and silently tolerate one that does not. That tolerance +// is what lets the same suffix iteration clean up both the 3-container Squid +// workload and the 2-container Envoy workload (which has no -ingress container) +// without backend-specific logic. See #5902. +func TestStopProxyContainer_HandlesEnvoyAndSquidTopologies(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + containerName string + present bool // whether the named proxy container exists + wantStopped bool + }{ + { + name: "present container is stopped (Squid egress/ingress/dns, Envoy egress/dns)", + containerName: "app-egress", + present: true, + wantStopped: true, + }, + { + name: "absent ingress is tolerated (Envoy has no -ingress container)", + containerName: "app-ingress", + present: false, + wantStopped: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var stoppedID string + api := &fakeDockerAPI{ + listFunc: func(_ context.Context, _ mobyclient.ContainerListOptions) ([]container.Summary, error) { + if !tt.present { + return []container.Summary{}, nil + } + return []container.Summary{ + {ID: "cid-" + tt.containerName, Names: []string{"/" + tt.containerName}}, + }, nil + }, + stopFunc: func(_ context.Context, id string, _ mobyclient.ContainerStopOptions) error { + stoppedID = id + return nil + }, + } + c := &Client{api: api} + + // Must not panic or error regardless of whether the container exists. + c.stopProxyContainer(t.Context(), tt.containerName, 30) + + if tt.wantStopped { + assert.Equal(t, "cid-"+tt.containerName, stoppedID, + "expected the existing proxy container to be stopped") + } else { + assert.Empty(t, stoppedID, + "expected no ContainerStop call for an absent proxy container") + } + }) + } +} diff --git a/pkg/container/docker/envoy.go b/pkg/container/docker/envoy.go new file mode 100644 index 0000000000..143e48d832 --- /dev/null +++ b/pkg/container/docker/envoy.go @@ -0,0 +1,985 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package docker + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "os" + "regexp" + "strconv" + "strings" + + "github.com/moby/moby/api/types/container" + "k8s.io/utils/ptr" + + "github.com/stacklok/toolhive/pkg/container/runtime" + lb "github.com/stacklok/toolhive/pkg/labels" + "github.com/stacklok/toolhive/pkg/networking" +) + +const ( + // defaultEnvoyImage is pinned by tag+digest for supply-chain integrity. + // Override with TOOLHIVE_ENVOY_IMAGE (accepts any docker pull reference). + defaultEnvoyImage = "envoyproxy/envoy-distroless:v1.32.3@sha256:375aab0d80b3c0e1b42a776b4cb1743ed79012032051d2da19cbc93ea884fb81" + + // Protobuf type URLs required by Envoy's protobuf-JSON bootstrap format. + // Every typed_config field must carry an @type URL or Envoy will reject the + // config on startup. + typeHCM = "type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager" + typeHTTPRBAC = "type.googleapis.com/envoy.extensions.filters.http.rbac.v3.RBAC" + typeDFPFilter = "type.googleapis.com/envoy.extensions.filters.http.dynamic_forward_proxy.v3.FilterConfig" + typeRouter = "type.googleapis.com/envoy.extensions.filters.http.router.v3.Router" + typeDFPCluster = "type.googleapis.com/envoy.extensions.clusters.dynamic_forward_proxy.v3.ClusterConfig" + + // dfpCacheName is the shared DNS cache config name used by both the DFP + // HTTP filter and the DFP cluster extension. + dfpCacheName = "dynamic_forward_proxy_cache_config" + + // dfpClusterName is the cluster referenced by the egress route config. + dfpClusterName = "dynamic_forward_proxy_cluster" + + // ingressClusterName is the cluster referenced by the ingress route config. + ingressClusterName = "ingress_upstream" +) + +func getEnvoyImage() string { + if img := os.Getenv("TOOLHIVE_ENVOY_IMAGE"); img != "" { + return img + } + return defaultEnvoyImage +} + +// ── Bootstrap ──────────────────────────────────────────────────────────────── + +// envoyBootstrap is the top-level Envoy bootstrap configuration. +// The admin interface is intentionally omitted: Envoy does not start an admin +// server when the field is absent, eliminating the attack surface at runtime. +type envoyBootstrap struct { + StaticResources envoyStaticResources `json:"static_resources"` +} + +// envoyStaticResources holds the static listeners and clusters. +type envoyStaticResources struct { + Listeners []envoyListener `json:"listeners,omitempty"` + Clusters []envoyCluster `json:"clusters,omitempty"` +} + +// ── Address types ──────────────────────────────────────────────────────────── + +// envoyAddress wraps a socket address for Envoy config. +type envoyAddress struct { + SocketAddress envoySocketAddress `json:"socket_address"` +} + +// envoySocketAddress is an IP + port pair used throughout Envoy config. +type envoySocketAddress struct { + Address string `json:"address"` + PortValue int `json:"port_value"` +} + +// ── Listener / filter chain ────────────────────────────────────────────────── + +// envoyListener is an Envoy listener binding on a socket address. +type envoyListener struct { + Name string `json:"name"` + Address envoyAddress `json:"address"` + FilterChains []envoyFilterChain `json:"filter_chains"` +} + +// envoyFilterChain is a sequence of network-level filters applied to matching +// connections. +type envoyFilterChain struct { + Filters []envoyNetworkFilter `json:"filters"` +} + +// envoyNetworkFilter is a named network filter whose typed config is an +// arbitrary protobuf-JSON object (e.g. envoyHCM). The any type lets us embed +// concrete structs directly so that JSON serialisation includes @type. +type envoyNetworkFilter struct { + Name string `json:"name"` + TypedConfig any `json:"typed_config"` +} + +// ── HttpConnectionManager ─────────────────────────────────────────────────── + +// envoyHCM is the typed config for +// envoy.filters.network.http_connection_manager. +type envoyHCM struct { + Type string `json:"@type"` + StatPrefix string `json:"stat_prefix"` + AccessLog []envoyAccessLog `json:"access_log,omitempty"` + // StreamIdleTimeout caps how long a stream may be idle. Envoy defaults it to + // 5m, which would reap a sparse-but-open SSE stream; "0s" disables it. + StreamIdleTimeout string `json:"stream_idle_timeout,omitempty"` + UpgradeConfigs []envoyUpgradeConfig `json:"upgrade_configs,omitempty"` + HTTPFilters []envoyHTTPFilter `json:"http_filters"` + RouteConfig *envoyRouteConfig `json:"route_config,omitempty"` +} + +// envoyAccessLog configures request access logging for an HCM. +type envoyAccessLog struct { + Name string `json:"name"` + TypedConfig any `json:"typed_config"` +} + +// envoyStdoutAccessLog is the typed config for envoy.access_loggers.stdout. +type envoyStdoutAccessLog struct { + Type string `json:"@type"` +} + +const typeStdoutAccessLog = "type.googleapis.com/envoy.extensions.access_loggers.stream.v3.StdoutAccessLog" + +func stdoutAccessLog() []envoyAccessLog { + return []envoyAccessLog{{ + Name: "envoy.access_loggers.stdout", + TypedConfig: &envoyStdoutAccessLog{Type: typeStdoutAccessLog}, + }} +} + +// envoyUpgradeConfig enables HTTP upgrade protocols (e.g. CONNECT) in HCM. +type envoyUpgradeConfig struct { + UpgradeType string `json:"upgrade_type"` + ConnectConfig *envoyConnectConfig `json:"connect_config,omitempty"` +} + +// envoyHTTPFilter is a named HTTP-layer filter embedded inside HCM. +type envoyHTTPFilter struct { + Name string `json:"name"` + TypedConfig any `json:"typed_config"` +} + +// ── RBAC ───────────────────────────────────────────────────────────────────── + +// envoyHTTPRBAC is the typed config for envoy.filters.http.rbac. +type envoyHTTPRBAC struct { + Type string `json:"@type"` + Rules envoyRBACRules `json:"rules"` +} + +// envoyRBACRules holds the RBAC action and policy map. +// CRITICAL: Policies must NOT have omitempty. An empty map serializes as {} and +// is intentional — an absent field would silently turn deny-all into allow-all. +type envoyRBACRules struct { + Action string `json:"action"` + Policies map[string]envoyRBACPolicy `json:"policies"` +} + +// envoyRBACPolicy pairs permissions with principals for a single RBAC policy. +type envoyRBACPolicy struct { + Permissions []envoyPermission `json:"permissions"` + Principals []envoyPrincipal `json:"principals"` +} + +// envoyPermission matches a request by various criteria. Exactly one field +// should be set per permission entry. +type envoyPermission struct { + Any *bool `json:"any,omitempty"` + Header *envoyHeaderMatcher `json:"header,omitempty"` + OrRules *envoyOrRules `json:"or_rules,omitempty"` +} + +// envoyOrRules composes multiple permissions with logical OR. +type envoyOrRules struct { + Rules []envoyPermission `json:"rules"` +} + +// envoyHeaderMatcher matches an HTTP header value. +type envoyHeaderMatcher struct { + Name string `json:"name"` + StringMatch *envoyStringMatch `json:"string_match,omitempty"` +} + +// envoyStringMatch matches a string by exact value, prefix, suffix, or regex. +type envoyStringMatch struct { + Exact string `json:"exact,omitempty"` + Prefix string `json:"prefix,omitempty"` + Suffix string `json:"suffix,omitempty"` + SafeRegex *envoySafeRegex `json:"safe_regex,omitempty"` +} + +// envoySafeRegex is an RE2 regex matcher. Envoy anchors the pattern fully +// (implicit ^ and $), so the pattern must match the entire input. +type envoySafeRegex struct { + Regex string `json:"regex"` +} + +// envoyPrincipal matches a downstream principal. Any:true is a wildcard. +type envoyPrincipal struct { + Any bool `json:"any"` +} + +// ── DFP filter ─────────────────────────────────────────────────────────────── + +// envoyDFPFilter is the typed config for +// envoy.filters.http.dynamic_forward_proxy. +type envoyDFPFilter struct { + Type string `json:"@type"` + DNSCacheConfig envoyDNSCache `json:"dns_cache_config"` +} + +// envoyDNSCache is the shared DNS cache config referenced by both the DFP +// HTTP filter and the DFP cluster extension. +type envoyDNSCache struct { + Name string `json:"name"` + DNSLookupFamily string `json:"dns_lookup_family"` +} + +// ── Router ──────────────────────────────────────────────────────────────────── + +// envoyRouter is the typed config for envoy.filters.http.router. +type envoyRouter struct { + Type string `json:"@type"` +} + +// ── Route config ───────────────────────────────────────────────────────────── + +// envoyRouteConfig holds the list of virtual hosts for a listener. +type envoyRouteConfig struct { + VirtualHosts []envoyVirtualHost `json:"virtual_hosts"` +} + +// envoyVirtualHost matches requests by domain and dispatches to a cluster. +type envoyVirtualHost struct { + Name string `json:"name"` + Domains []string `json:"domains"` + Routes []envoyRoute `json:"routes"` +} + +// envoyRoute matches a request prefix and forwards it to a cluster. +type envoyRoute struct { + Match envoyRouteMatch `json:"match"` + Route *envoyRouteAction `json:"route,omitempty"` +} + +// envoyRouteMatch matches incoming requests by URI prefix or CONNECT method. +// Exactly one of Prefix or ConnectMatcher must be set. +type envoyRouteMatch struct { + Prefix string `json:"prefix,omitempty"` + ConnectMatcher *envoyConnectMatcher `json:"connect_matcher,omitempty"` +} + +// envoyConnectMatcher matches HTTP CONNECT requests (used for HTTPS tunneling). +type envoyConnectMatcher struct{} + +// envoyRouteAction forwards matched requests to an upstream cluster. +type envoyRouteAction struct { + Cluster string `json:"cluster"` + // Timeout is the route's total request→response cap. Envoy defaults it to + // 15s, which truncates long-lived MCP streams (SSE, streamable-http); set + // "0s" to disable it for proxied traffic. + Timeout string `json:"timeout,omitempty"` + UpgradeConfigs []envoyUpgradeConfig `json:"upgrade_configs,omitempty"` +} + +// envoyConnectConfig is the per-route CONNECT tunnel configuration. +type envoyConnectConfig struct{} + +// ── Cluster ─────────────────────────────────────────────────────────────────── + +// envoyCluster is an Envoy upstream cluster definition. +type envoyCluster struct { + Name string `json:"name"` + ConnectTimeout string `json:"connect_timeout"` + LbPolicy string `json:"lb_policy,omitempty"` + Type string `json:"type,omitempty"` + // DnsLookupFamily restricts DNS resolution to IPv4 only. Without this, + // STRICT_DNS clusters attempt AAAA lookups first, which adds latency in + // environments where IPv6 is unavailable or times out. + DnsLookupFamily string `json:"dns_lookup_family,omitempty"` + ClusterType *envoyClusterType `json:"cluster_type,omitempty"` + LoadAssignment *envoyLoadAssignment `json:"load_assignment,omitempty"` +} + +// envoyClusterType is the custom cluster discovery extension (e.g. DFP). +type envoyClusterType struct { + Name string `json:"name"` + TypedConfig any `json:"typed_config"` +} + +// envoyDFPClusterConfig is the typed config for +// envoy.clusters.dynamic_forward_proxy. +type envoyDFPClusterConfig struct { + Type string `json:"@type"` + DNSCacheConfig envoyDNSCache `json:"dns_cache_config"` +} + +// envoyLoadAssignment is an EDS-style static load assignment. +type envoyLoadAssignment struct { + ClusterName string `json:"cluster_name"` + Endpoints []envoyEndpoint `json:"endpoints"` +} + +// envoyEndpoint is a group of LB endpoints. +type envoyEndpoint struct { + LBEndpoints []envoyLBEndpoint `json:"lb_endpoints"` +} + +// envoyLBEndpoint is a single upstream endpoint. +type envoyLBEndpoint struct { + Endpoint envoyEndpointAddress `json:"endpoint"` +} + +// envoyEndpointAddress wraps an address for an LB endpoint. +type envoyEndpointAddress struct { + Address envoyAddress `json:"address"` +} + +// ── Builder functions ───────────────────────────────────────────────────────── + +// buildEgressListener builds the Envoy listener config for outbound traffic. +// +// The resulting HCM HTTP filter chain is: +// 1. (when !spec.AllowDockerGateway) RBAC DENY on gateway IP and hostnames +// 2. RBAC ALLOW allowlist — empty policies map is Envoy's deny-all +// 3. dynamic_forward_proxy HTTP filter +// 4. router +// +// CONNECT upgrades are enabled so that HTTPS CONNECT tunnels pass through. +func buildEgressListener(spec proxySpec) envoyListener { + httpFilters := buildEgressHTTPFilters(spec) + + hcm := &envoyHCM{ + Type: typeHCM, + StatPrefix: "egress_http", + AccessLog: stdoutAccessLog(), + StreamIdleTimeout: "0s", // don't reap sparse long-lived outbound streams + UpgradeConfigs: []envoyUpgradeConfig{{UpgradeType: "CONNECT"}}, + HTTPFilters: httpFilters, + RouteConfig: &envoyRouteConfig{ + VirtualHosts: []envoyVirtualHost{ + { + Name: "local_service", + Domains: []string{"*"}, + Routes: []envoyRoute{ + // CONNECT match must come first: handles HTTPS tunneling. + // Without this, CONNECT requests don't match prefix "/" and get 404. + { + Match: envoyRouteMatch{ConnectMatcher: &envoyConnectMatcher{}}, + Route: &envoyRouteAction{ + Cluster: dfpClusterName, + UpgradeConfigs: []envoyUpgradeConfig{ + {UpgradeType: "CONNECT", ConnectConfig: &envoyConnectConfig{}}, + }, + }, + }, + // Prefix match handles plain HTTP requests. Timeout "0s" + // disables Envoy's 15s default so long-lived plain-HTTP + // outbound streams aren't truncated (CONNECT tunnels are + // unaffected by the default, but disabling is harmless there). + { + Match: envoyRouteMatch{Prefix: "/"}, + Route: &envoyRouteAction{Cluster: dfpClusterName, Timeout: "0s"}, + }, + }, + }, + }, + }, + } + + return envoyListener{ + Name: fmt.Sprintf("%s-egress", spec.WorkloadName), + Address: envoyAddress{ + SocketAddress: envoySocketAddress{ + Address: "0.0.0.0", + PortValue: 3128, + }, + }, + FilterChains: []envoyFilterChain{ + { + Filters: []envoyNetworkFilter{ + { + Name: "envoy.filters.network.http_connection_manager", + TypedConfig: hcm, + }, + }, + }, + }, + } +} + +// buildEgressHTTPFilters constructs the ordered list of HTTP filters for the +// egress HCM. Gateway deny rules (when !spec.AllowDockerGateway) are placed +// first so they are evaluated before the allowlist. +func buildEgressHTTPFilters(spec proxySpec) []envoyHTTPFilter { + var filters []envoyHTTPFilter + + if !spec.AllowDockerGateway { + filters = append(filters, envoyHTTPFilter{ + Name: "envoy.filters.http.rbac", + TypedConfig: buildGatewayDenyRBAC(spec.GatewayIP), + }) + } + + filters = append(filters, + envoyHTTPFilter{ + Name: "envoy.filters.http.rbac", + TypedConfig: buildAllowlistRBAC(spec), + }, + envoyHTTPFilter{ + Name: "envoy.filters.http.dynamic_forward_proxy", + TypedConfig: &envoyDFPFilter{ + Type: typeDFPFilter, + DNSCacheConfig: envoyDNSCache{ + Name: dfpCacheName, + DNSLookupFamily: "V4_ONLY", + }, + }, + }, + envoyHTTPFilter{ + Name: "envoy.filters.http.router", + TypedConfig: &envoyRouter{Type: typeRouter}, + }, + ) + + return filters +} + +// buildGatewayDenyRBAC builds an RBAC DENY filter that blocks the Docker bridge +// gateway — both the resolved gateway IP and the Docker-internal hostnames — +// matched on the :authority header. This filter must precede the allowlist. +// +// Matching is on :authority (not destination_ip): for a forward proxy the RBAC +// destination_ip resolves to Envoy's own listener socket, not the proxied +// target, so an L3 CIDR rule never matches the upstream. The forward-proxy +// target is carried in :authority for both plain HTTP and HTTPS CONNECT, so +// authority matching — including the gateway IP literal — is what actually +// blocks it (mirroring Squid's `dst`/`dstdomain` denies). +func buildGatewayDenyRBAC(gatewayIP string) *envoyHTTPRBAC { + return &envoyHTTPRBAC{ + Type: typeHTTPRBAC, + Rules: envoyRBACRules{ + Action: "DENY", + Policies: map[string]envoyRBACPolicy{ + "gateway": { + Permissions: gatewayAuthorityPermissions(gatewayIP), + Principals: []envoyPrincipal{{Any: true}}, + }, + }, + }, + } +} + +// gatewayAuthorityPermissions returns RBAC permissions that match the Docker +// gateway on the :authority header — the resolved gateway IP literal plus the +// two Docker-internal hostnames. Each is matched with hostMatchRegex, so the +// match is anchored, case-insensitive (HOST.DOCKER.INTERNAL cannot bypass it), +// and tolerant of an optional :port. Shared by the deny path (gateway blocked) +// and the allow path (--allow-docker-gateway). +func gatewayAuthorityPermissions(gatewayIP string) []envoyPermission { + patterns := []string{ + hostMatchRegex(gatewayIP), + hostMatchRegex(dockerGatewayHostname), + hostMatchRegex(dockerAltGatewayHostname), + } + rules := make([]envoyPermission, 0, len(patterns)) + for _, re := range patterns { + rules = append(rules, envoyPermission{ + Header: &envoyHeaderMatcher{ + Name: ":authority", + StringMatch: &envoyStringMatch{SafeRegex: &envoySafeRegex{Regex: re}}, + }, + }) + } + return []envoyPermission{{OrRules: &envoyOrRules{Rules: rules}}} +} + +// buildAllowlistRBAC builds an RBAC ALLOW filter encoding the outbound +// allowlist. An empty Policies map is Envoy's deny-all under ALLOW action. +func buildAllowlistRBAC(spec proxySpec) *envoyHTTPRBAC { + return &envoyHTTPRBAC{ + Type: typeHTTPRBAC, + Rules: envoyRBACRules{ + Action: "ALLOW", + Policies: buildAllowlistPolicies(spec), + }, + } +} + +// buildAllowlistPolicies returns the policy map for the egress RBAC ALLOW +// filter. An empty map (deny-all) is returned when: +// - spec.Permissions is nil +// - spec.Permissions.Outbound is nil +// +// InsecureAllowAll produces a single wildcard policy. Each AllowHost entry +// becomes a policy matching the :authority header via hostMatchRegex (Squid +// dstdomain semantics). When AllowPort is also set, each host policy AND-s a +// port-suffix permission so that host AND port must both match — mirroring +// Squid's "allowed_ports AND allowed_dsts" combination. +// +// Known divergence from Squid: plain-HTTP requests omit the port from +// :authority (e.g. "example.com" not "example.com:80"). A port-suffix matcher +// for ":80" will not match the portless authority, so AllowPort:[80] with plain +// HTTP is more restrictive than Squid's allowed_ports ACL. This is fail-closed +// (over-blocks) rather than a bypass. +// +// When spec.AllowDockerGateway is set, an explicit ALLOW policy for the gateway +// is added: omitting the deny filter alone is not enough, because the allowlist +// is default-deny, so the gateway would still be blocked under a non-permissive +// profile. This mirrors Squid, which emits explicit allow rules for the flag. +func buildAllowlistPolicies(spec proxySpec) map[string]envoyRBACPolicy { + policies := make(map[string]envoyRBACPolicy) + + // --allow-docker-gateway: grant the gateway explicitly (see doc comment). + // Harmless under InsecureAllowAll (allow-all already covers it). + if spec.AllowDockerGateway { + policies["gateway"] = envoyRBACPolicy{ + Permissions: gatewayAuthorityPermissions(spec.GatewayIP), + Principals: []envoyPrincipal{{Any: true}}, + } + } + + if spec.Permissions == nil || spec.Permissions.Outbound == nil { + return policies + } + out := spec.Permissions.Outbound + if out.InsecureAllowAll { + policies["allow-all"] = envoyRBACPolicy{ + Permissions: []envoyPermission{{Any: ptr.To(true)}}, + Principals: []envoyPrincipal{{Any: true}}, + } + return policies + } + if len(out.AllowHost) == 0 && len(out.AllowPort) > 0 { + // AllowPort only — any host, but only on the listed ports. + // Use a single regex so both the explicit-port ("host:80") and the + // bare-hostname ("host", implying port 80) forms are handled correctly. + policies["any-host-allowed-ports"] = envoyRBACPolicy{ + Permissions: []envoyPermission{{ + Header: &envoyHeaderMatcher{ + Name: ":authority", + StringMatch: &envoyStringMatch{SafeRegex: &envoySafeRegex{Regex: anyHostPortRegex(out.AllowPort)}}, + }, + }}, + Principals: []envoyPrincipal{{Any: true}}, + } + return policies + } + + for _, host := range out.AllowHost { + var regex string + if len(out.AllowPort) > 0 { + // AllowHost + AllowPort: build a single regex that encodes both the + // host pattern and the allowed ports. Using a combined regex rather + // than separate AND-d permissions means plain-HTTP requests (where + // :authority omits the default port 80) are handled correctly — a + // bare "example.com" is permitted when port 80 is in the list. + regex = hostPortMatchRegex(host, out.AllowPort) + } else { + // AllowHost only — any port permitted. + regex = hostMatchRegex(host) + } + policies[host] = envoyRBACPolicy{ + Permissions: []envoyPermission{{ + Header: &envoyHeaderMatcher{ + Name: ":authority", + StringMatch: &envoyStringMatch{SafeRegex: &envoySafeRegex{Regex: regex}}, + }, + }}, + Principals: []envoyPrincipal{{Any: true}}, + } + } + return policies +} + +// portGroupRegex builds the ":(?:port1|port2|…)" alternation for a list of +// allowed ports. Returns (group, includes80) where includes80 indicates whether +// port 80 is present (caller uses this to decide whether the port suffix is +// optional). Shared by hostPortMatchRegex and anyHostPortRegex. +func portGroupRegex(ports []int) (group string, includes80 bool) { + alts := make([]string, 0, len(ports)) + for _, p := range ports { + alts = append(alts, strconv.Itoa(p)) + if p == 80 { + includes80 = true + } + } + return ":(?:" + strings.Join(alts, "|") + ")", includes80 +} + +// hostBasePattern returns the escaped, case-insensitive host portion of an +// :authority regex — without any trailing port group. It mirrors hostMatchRegex +// semantics: leading dot / *. = subdomain wildcard, optional trailing dot. +// Shared by hostPortMatchRegex and hostMatchRegex. +func hostBasePattern(host string) string { + subdomains := false + switch { + case strings.HasPrefix(host, "*."): + host = host[2:] + subdomains = true + case strings.HasPrefix(host, "."): + host = host[1:] + subdomains = true + } + p := regexp.QuoteMeta(host) + if subdomains { + p = `(.*\.)?` + p + } + return p +} + +// hostPortMatchRegex builds a single anchored RE2 pattern that encodes both a +// host allowlist entry and a port allowlist, so that AllowHost+AllowPort can be +// expressed as one safe_regex permission rather than AND-ing two separate matchers. +// +// Port 80 is treated as the HTTP default: when 80 is in the ports list the port +// suffix is optional, so a bare "example.com" authority (plain HTTP) is permitted +// alongside "example.com:80". For all other ports the port suffix is required. +// +// The host portion mirrors hostMatchRegex semantics (leading dot / *. = subdomain +// wildcard, case-insensitive, optional trailing dot, no port in the base pattern). +func hostPortMatchRegex(host string, ports []int) string { + portGroup, includes80 := portGroupRegex(ports) + hostPattern := hostBasePattern(host) + if includes80 { + // Port is optional: bare "example.com" counts as port 80. + return `(?i)` + hostPattern + `\.?(` + portGroup + `)?` + } + // Port is required: bare hostname is not accepted. + return `(?i)` + hostPattern + `\.?` + portGroup +} + +// anyHostPortRegex builds a regex matching any hostname on the given ports. +// When port 80 is in the list the port suffix is optional (bare hostname implies +// port 80 for plain HTTP). For other ports the port suffix is required. +// +// [^:]+ matches any hostname without an explicit port; IPv6 bracket addresses +// (which contain colons) are outside the supported scope for this proxy backend. +func anyHostPortRegex(ports []int) string { + portGroup, includes80 := portGroupRegex(ports) + if includes80 { + return `(?i)[^:]+(` + portGroup + `)?` + } + return `(?i)[^:]+` + portGroup +} + +// hostMatchRegex builds an anchored, case-insensitive RE2 pattern that matches +// an AllowHost entry against the HTTP :authority header. It tolerates an +// optional ":port" suffix because HTTPS CONNECT authorities include the port +// (e.g. "example.com:443"), and it mirrors Squid's dstdomain semantics: +// +// - ".example.com" or "*.example.com" — matches the apex "example.com" AND +// any subdomain ("www.example.com", "a.b.example.com"). +// - "example.com" — matches that host exactly (no subdomains). +// +// The domain is regex-escaped so dots are literal, and the pattern is anchored +// by Envoy so it cannot match "example.com.attacker.com". +func hostMatchRegex(host string) string { + pattern := hostBasePattern(host) + // (?i) case-insensitive (hostnames are); \.? tolerates a trailing-dot FQDN + // ("host.docker.internal." resolves identically and would otherwise bypass a + // deny); (:[0-9]+)? optional port. + return `(?i)` + pattern + `\.?(:[0-9]+)?` +} + +// buildEgressCluster returns the dynamic_forward_proxy cluster required by the +// egress listener's route config. +func buildEgressCluster() envoyCluster { + return envoyCluster{ + Name: dfpClusterName, + ConnectTimeout: "10s", + LbPolicy: "CLUSTER_PROVIDED", + ClusterType: &envoyClusterType{ + Name: "envoy.clusters.dynamic_forward_proxy", + TypedConfig: &envoyDFPClusterConfig{ + Type: typeDFPCluster, + DNSCacheConfig: envoyDNSCache{ + Name: dfpCacheName, + DNSLookupFamily: "V4_ONLY", + }, + }, + }, + } +} + +// buildIngressListener builds the Envoy listener config for inbound (ingress) +// traffic. It binds on 0.0.0.0:hostPort inside the container so that Docker's +// port forwarding (host:127.0.0.1: → container:) can reach +// it. The host-side HostIP:"127.0.0.1" in the port binding provides the +// localhost-only restriction; the container-side address must be 0.0.0.0 or +// Docker's bridge forwarding cannot deliver traffic to the listener. +// +// The virtual host domain is always "*". Inbound host filtering is enforced +// by the egress RBAC `:authority` matcher (hostMatchRegex), not the ingress +// virtual host list — the transparent proxy sends "127.0.0.1:" as the +// Host header (port included), which would not match bare hostnames in a +// domain list. +func buildIngressListener(spec proxySpec, hostPort int) envoyListener { + domains := ingressDomains(spec) + + hcm := &envoyHCM{ + Type: typeHCM, + StatPrefix: "ingress_http", + AccessLog: stdoutAccessLog(), + // Disable the idle timer so a sparse-but-open SSE stream isn't reaped. + StreamIdleTimeout: "0s", + HTTPFilters: []envoyHTTPFilter{ + { + Name: "envoy.filters.http.router", + TypedConfig: &envoyRouter{Type: typeRouter}, + }, + }, + RouteConfig: &envoyRouteConfig{ + VirtualHosts: []envoyVirtualHost{ + { + Name: "ingress_service", + Domains: domains, + Routes: []envoyRoute{ + { + Match: envoyRouteMatch{Prefix: "/"}, + // Timeout "0s" disables Envoy's 15s default total-response + // cap, which would otherwise truncate SSE / streamable-http + // MCP streams. This is the primary transport path. + Route: &envoyRouteAction{Cluster: ingressClusterName, Timeout: "0s"}, + }, + }, + }, + }, + }, + } + + return envoyListener{ + Name: fmt.Sprintf("%s-ingress", spec.WorkloadName), + Address: envoyAddress{ + SocketAddress: envoySocketAddress{ + Address: "0.0.0.0", // must be 0.0.0.0; Docker port forwarding targets the bridge IP, not container loopback + PortValue: hostPort, + }, + }, + FilterChains: []envoyFilterChain{ + { + Filters: []envoyNetworkFilter{ + { + Name: "envoy.filters.network.http_connection_manager", + TypedConfig: hcm, + }, + }, + }, + }, + } +} + +// ingressDomains returns the virtual host domain list for the ingress listener. +// Always returns a wildcard: the Inbound.AllowHost list contains bare hostnames +// (e.g. "localhost", "127.0.0.1") but the transparent proxy sends a Host header +// with a port suffix ("127.0.0.1:22354"), which Envoy would not match against +// a bare hostname. The inbound access restriction is instead enforced by the +// host-side port binding to 127.0.0.1, which already limits the ingress to +// local connections only. +func ingressDomains(_ proxySpec) []string { + return []string{"*"} +} + +// buildIngressCluster returns the STRICT_DNS upstream cluster for the ingress +// listener, pointing at spec.WorkloadName:spec.UpstreamPort. +func buildIngressCluster(spec proxySpec) envoyCluster { + // Prefer the resolved upstream IP (no DNS dependency); fall back to the + // workload name. See proxySpec.UpstreamHost. + upstreamHost := spec.UpstreamHost + if upstreamHost == "" { + upstreamHost = spec.WorkloadName + } + return envoyCluster{ + Name: ingressClusterName, + ConnectTimeout: "10s", + Type: "STRICT_DNS", + DnsLookupFamily: "V4_ONLY", + LoadAssignment: &envoyLoadAssignment{ + ClusterName: ingressClusterName, + Endpoints: []envoyEndpoint{ + { + LBEndpoints: []envoyLBEndpoint{ + { + Endpoint: envoyEndpointAddress{ + Address: envoyAddress{ + SocketAddress: envoySocketAddress{ + Address: upstreamHost, + PortValue: spec.UpstreamPort, + }, + }, + }, + }, + }, + }, + }, + }, + } +} + +// writeEnvoyBootstrap marshals b to JSON and writes it to a temporary file at +// mode 0600, returning the path. On success the file must outlive this call — +// the Envoy container bind-mounts it read-only — so it is not removed here; the +// caller removes it if setup fails afterward. On any error during writing, the +// partially-created file is removed before returning. +func writeEnvoyBootstrap(b envoyBootstrap) (string, error) { + data, err := json.Marshal(b) + if err != nil { + return "", fmt.Errorf("failed to marshal envoy bootstrap: %w", err) + } + tmpFile, err := os.CreateTemp("", "envoy-bootstrap-*.json") + if err != nil { + return "", fmt.Errorf("failed to create envoy bootstrap temp file: %w", err) + } + created := tmpFile.Name() + defer func() { + if cerr := tmpFile.Close(); cerr != nil { + slog.Warn("failed to close envoy bootstrap temp file", "error", cerr) + } + }() + if _, err := tmpFile.Write(data); err != nil { + _ = os.Remove(created) + return "", fmt.Errorf("failed to write envoy bootstrap: %w", err) + } + // 0o644: world-readable so the Envoy distroless container (UID 101) can + // read the bind-mounted file. On Linux Docker Engine, strict POSIX + // permissions apply — 0o600 prevents the container user from reading the + // file, causing Envoy to crash-loop. The bootstrap contains no secrets + // (only network topology: hostnames, ports, RBAC rules). + if err := tmpFile.Chmod(0o644); err != nil { + _ = os.Remove(created) + return "", fmt.Errorf("failed to set envoy bootstrap file permissions: %w", err) + } + return created, nil +} + +// ── envoyProxy ──────────────────────────────────────────────────────────────── + +// envoyProxy implements networkProxy using Envoy as the proxy backend. +// It creates a single Envoy container that handles both egress (forward proxy +// on :3128) and ingress (reverse proxy) as separate listeners, reducing aux +// container count from 3 (Squid: egress + ingress + dns) to 2 (Envoy: combined +// + dns). +type envoyProxy struct { + client *Client +} + +// SetupEgress implements networkProxy for the Envoy backend. +// +// Only the egress proxy env vars are returned here — no container is created +// yet. Container creation is deferred to SetupIngress (which runs after +// createMcpContainer) so that the MCP container's hostname is resolvable the +// moment the Envoy STRICT_DNS ingress cluster first probes it. Creating Envoy +// before the MCP container caused the ingress cluster to cache a negative DNS +// response on Linux Docker Engine, preventing the server from ever becoming +// ready within the readiness window (see #5922). +func (*envoyProxy) SetupEgress(_ context.Context, spec proxySpec) (egressResult, error) { + egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName) + return egressResult{EnvVars: addEgressEnvVars(nil, egressContainerName)}, nil +} + +// SetupIngress implements networkProxy for the Envoy backend. +// +// Creates the Envoy container after the MCP container exists, with both the +// egress forward-proxy listener and (for non-stdio transports) the ingress +// reverse-proxy listener. Running after createMcpContainer ensures the +// STRICT_DNS upstream cluster resolves the MCP hostname on the first probe, +// avoiding the Linux Docker Engine readiness failure described in #5922. +func (e *envoyProxy) SetupIngress(ctx context.Context, spec proxySpec, _ egressResult) (int, error) { + // The container is named -egress (not -proxy or similar) to + // preserve parity with the Squid backend: addEgressEnvVars and the suffix- + // based cleanup loop both key off this name. + egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName) + + bootstrap := envoyBootstrap{ + StaticResources: envoyStaticResources{ + Listeners: []envoyListener{buildEgressListener(spec)}, + Clusters: []envoyCluster{buildEgressCluster()}, + }, + } + + var ingressPort int + if spec.TransportType != "stdio" && spec.UpstreamPort > 0 { + // A random port avoids every same-image workload racing to bind the + // same preferred port when starting concurrently (see #6063). + port, err := networking.FindOrUsePort(0) + if err != nil { + return 0, fmt.Errorf("failed to find ingress port: %w", err) + } + ingressPort = port + bootstrap.StaticResources.Listeners = append( + bootstrap.StaticResources.Listeners, + buildIngressListener(spec, ingressPort), + ) + bootstrap.StaticResources.Clusters = append( + bootstrap.StaticResources.Clusters, + buildIngressCluster(spec), + ) + } + + configPath, err := writeEnvoyBootstrap(bootstrap) + if err != nil { + return 0, err + } + success := false + defer func() { + if !success { + _ = os.Remove(configPath) + } + }() + + envoyImage := getEnvoyImage() + //nolint:gosec // G706: envoy image name from config + slog.Debug("setting up envoy container", "name", egressContainerName, "image", envoyImage) + + if err := e.client.imageManager.PullImage(ctx, envoyImage); err != nil { + exists, inspectErr := e.client.imageManager.ImageExists(ctx, envoyImage) + if inspectErr != nil || !exists { + return 0, fmt.Errorf("failed to pull envoy image: %w", err) + } + //nolint:gosec // G706: envoy image name from config + slog.Debug("envoy image exists locally, continuing despite pull failure", "image", envoyImage) + } + + envoyLabels := map[string]string{} + lb.AddStandardLabels(envoyLabels, egressContainerName, egressContainerName, "stdio", 80) + envoyLabels[ToolhiveAuxiliaryWorkloadLabel] = LabelValueTrue + + containerConfig := &container.Config{ + Image: envoyImage, + Cmd: []string{"-c", "/etc/envoy/envoy.json"}, + Labels: envoyLabels, + } + + mounts := []runtime.Mount{ + {Source: configPath, Target: "/etc/envoy/envoy.json", ReadOnly: true}, + } + + var exposedPorts map[string]struct{} + var portBindings map[string][]runtime.PortBinding + if ingressPort > 0 { + portStr := strconv.Itoa(ingressPort) + portKey := portStr + "/tcp" + exposedPorts = map[string]struct{}{portKey: {}} + portBindings = map[string][]runtime.PortBinding{ + portKey: {{HostIP: "127.0.0.1", HostPort: portStr}}, + } + } + + hostConfig := &container.HostConfig{ + Mounts: convertMounts(mounts), + NetworkMode: container.NetworkMode("bridge"), + SecurityOpt: []string{"label:disable"}, + CapDrop: []string{"ALL"}, + RestartPolicy: container.RestartPolicy{ + Name: "unless-stopped", + }, + } + if portBindings != nil { + if err := setupPortBindings(hostConfig, portBindings); err != nil { + return 0, fmt.Errorf("failed to setup port bindings: %w", err) + } + } + if err := setupExposedPorts(containerConfig, exposedPorts); err != nil { + return 0, fmt.Errorf("failed to setup exposed ports: %w", err) + } + + if _, err := e.client.createContainer(ctx, egressContainerName, containerConfig, hostConfig, spec.Endpoints); err != nil { + return 0, fmt.Errorf("failed to create envoy container: %w", err) + } + + success = true + return ingressPort, nil +} diff --git a/pkg/container/docker/envoy_test.go b/pkg/container/docker/envoy_test.go new file mode 100644 index 0000000000..a45c8b0bef --- /dev/null +++ b/pkg/container/docker/envoy_test.go @@ -0,0 +1,1234 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package docker + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "regexp" + "strings" + "sync" + "testing" + "time" + + "github.com/moby/moby/api/types/container" + "github.com/moby/moby/api/types/network" + mobyclient "github.com/moby/moby/client" + v1 "github.com/opencontainers/image-spec/specs-go/v1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/permissions" +) + +// Compile-time assertion: envoyProxy must satisfy networkProxy. +var _ networkProxy = (*envoyProxy)(nil) + +// findRBACFilter walks the HTTP filters inside the first HCM in a listener's +// first filter chain and returns the first RBAC filter with action == "ALLOW". +// It returns nil if no matching filter is found. +func findRBACFilter(listener envoyListener) *envoyHTTPRBAC { + if len(listener.FilterChains) == 0 { + return nil + } + fc := listener.FilterChains[0] + if len(fc.Filters) == 0 { + return nil + } + hcm, ok := fc.Filters[0].TypedConfig.(*envoyHCM) + if !ok { + return nil + } + for _, f := range hcm.HTTPFilters { + rbac, ok := f.TypedConfig.(*envoyHTTPRBAC) + if ok && rbac.Rules.Action == "ALLOW" { + return rbac + } + } + return nil +} + +// collectAuthorityRegexes walks every policy/permission in an RBAC filter +// (including or_rules) and returns all :authority safe_regex patterns. +func collectAuthorityRegexes(rbac *envoyHTTPRBAC) []string { + var out []string + var walk func(perms []envoyPermission) + walk = func(perms []envoyPermission) { + for _, p := range perms { + if p.Header != nil && p.Header.Name == ":authority" && + p.Header.StringMatch != nil && p.Header.StringMatch.SafeRegex != nil { + out = append(out, p.Header.StringMatch.SafeRegex.Regex) + } + if p.OrRules != nil { + walk(p.OrRules.Rules) + } + } + } + for _, pol := range rbac.Rules.Policies { + walk(pol.Permissions) + } + return out +} + +// authorityMatched reports whether any of the given RE2 patterns matches the +// authority, replicating Envoy's full-string anchoring. +func authorityMatched(patterns []string, authority string) bool { + for _, p := range patterns { + if regexp.MustCompile("^(?:" + p + ")$").MatchString(authority) { + return true + } + } + return false +} + +// listenerAllowsAuthority returns true when the ALLOW RBAC filter in the given +// listener would permit a request with the given :authority value. +func listenerAllowsAuthority(listener envoyListener, authority string) bool { + rbac := findRBACFilter(listener) + if rbac == nil { + return false + } + for _, policy := range rbac.Rules.Policies { + allMatch := true + for _, perm := range policy.Permissions { + if !permMatchesAuthority(perm, authority) { + allMatch = false + break + } + } + if allMatch { + return true + } + } + return false +} + +// findDenyRBACFilter returns the first RBAC filter with action == "DENY" from +// the HCM inside the listener's first filter chain. +func findDenyRBACFilter(listener envoyListener) *envoyHTTPRBAC { + if len(listener.FilterChains) == 0 { + return nil + } + fc := listener.FilterChains[0] + if len(fc.Filters) == 0 { + return nil + } + hcm, ok := fc.Filters[0].TypedConfig.(*envoyHCM) + if !ok { + return nil + } + for _, f := range hcm.HTTPFilters { + rbac, ok := f.TypedConfig.(*envoyHTTPRBAC) + if ok && rbac.Rules.Action == "DENY" { + return rbac + } + } + return nil +} + +// TestBuildEgressListener_AllowlistAndDefaultDeny exercises the RBAC policy +// generation logic of buildEgressListener across the main permission scenarios. +func TestBuildEgressListener_AllowlistAndDefaultDeny(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + spec proxySpec + wantRBACPresent bool + wantRBACAction string // "ALLOW" or "DENY" + wantPolicies []string + wantGatewayDenyAbsent bool + wantGatewayDenyL3 bool // CIDR deny on GatewayIP + wantGatewayDenyL7 bool // authority deny on host.docker.internal + }{ + { + name: "nil permissions, InsecureAllowAll=false produces deny-all RBAC", + spec: proxySpec{ + WorkloadName: "myserver", + Permissions: nil, + AllowDockerGateway: false, + GatewayIP: dockerDefaultBridgeGatewayIP, + }, + wantRBACPresent: true, + wantRBACAction: "ALLOW", + wantPolicies: nil, // no policies → Envoy deny-all + wantGatewayDenyL3: true, + wantGatewayDenyL7: true, + }, + { + name: "InsecureAllowAll=true produces allow-all RBAC policy", + spec: proxySpec{ + WorkloadName: "myserver", + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + InsecureAllowAll: true, + }, + }, + AllowDockerGateway: false, + GatewayIP: dockerDefaultBridgeGatewayIP, + }, + wantRBACPresent: true, + wantRBACAction: "ALLOW", + wantPolicies: []string{"allow-all"}, + wantGatewayDenyL3: true, + wantGatewayDenyL7: true, + }, + { + name: "AllowHost list produces per-host RBAC policies", + spec: proxySpec{ + WorkloadName: "myserver", + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com", "api.example.com"}, + }, + }, + AllowDockerGateway: false, + GatewayIP: dockerDefaultBridgeGatewayIP, + }, + wantRBACPresent: true, + wantRBACAction: "ALLOW", + wantPolicies: []string{"example.com", "api.example.com"}, + wantGatewayDenyL3: true, + wantGatewayDenyL7: true, + }, + { + name: "AllowDockerGateway=true omits gateway deny rules", + spec: proxySpec{ + WorkloadName: "myserver", + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + InsecureAllowAll: true, + }, + }, + AllowDockerGateway: true, + GatewayIP: dockerDefaultBridgeGatewayIP, + }, + wantRBACPresent: true, + wantRBACAction: "ALLOW", + wantPolicies: []string{"allow-all"}, + wantGatewayDenyAbsent: true, + }, + { + name: "wildcard AllowHost produces correct authority pattern", + spec: proxySpec{ + WorkloadName: "myserver", + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"*.example.com"}, + }, + }, + AllowDockerGateway: false, + GatewayIP: dockerDefaultBridgeGatewayIP, + }, + wantRBACPresent: true, + wantRBACAction: "ALLOW", + wantPolicies: []string{"*.example.com"}, + wantGatewayDenyL3: true, + wantGatewayDenyL7: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + listener := buildEgressListener(tt.spec) + + // The listener must have at least one filter chain. + require.NotEmpty(t, listener.FilterChains, "expected at least one filter chain") + + if tt.wantRBACPresent { + rbac := findRBACFilter(listener) + require.NotNil(t, rbac, "expected RBAC filter to be present in the listener") + assert.Equal(t, tt.wantRBACAction, rbac.Rules.Action, + "RBAC action mismatch") + + if tt.wantPolicies == nil { + assert.Empty(t, rbac.Rules.Policies, + "expected empty policy set for deny-all") + } else { + for _, policyName := range tt.wantPolicies { + _, ok := rbac.Rules.Policies[policyName] + assert.True(t, ok, "expected RBAC policy %q to be present", policyName) + } + } + } + + if tt.wantGatewayDenyAbsent { + // --allow-docker-gateway: no DENY filter, and the ALLOW filter must + // explicitly permit the gateway (omitting the deny alone is not + // enough under a default-deny allowlist — see #2). + assert.Nil(t, findDenyRBACFilter(listener), + "no DENY RBAC filter should be present when AllowDockerGateway=true") + allow := findRBACFilter(listener) + require.NotNil(t, allow) + pats := collectAuthorityRegexes(allow) + assert.True(t, authorityMatched(pats, dockerDefaultBridgeGatewayIP), + "ALLOW filter must permit the gateway IP when AllowDockerGateway=true") + assert.True(t, authorityMatched(pats, dockerGatewayHostname), + "ALLOW filter must permit host.docker.internal when AllowDockerGateway=true") + } + + // wantGatewayDenyL3/L7: the DENY filter must block the gateway by + // :authority — the IP literal (with/without port) AND the hostnames, + // case-insensitively — without a destination_ip rule (inert for a + // forward proxy). See #1/#4. + if tt.wantGatewayDenyL3 || tt.wantGatewayDenyL7 { + deny := findDenyRBACFilter(listener) + require.NotNil(t, deny, "expected a DENY RBAC filter for the gateway") + pats := collectAuthorityRegexes(deny) + + assert.True(t, authorityMatched(pats, dockerDefaultBridgeGatewayIP), + "deny must match the gateway IP") + assert.True(t, authorityMatched(pats, dockerDefaultBridgeGatewayIP+":8080"), + "deny must match the gateway IP with a port (raw-IP bypass)") + assert.True(t, authorityMatched(pats, dockerGatewayHostname), + "deny must match host.docker.internal") + assert.True(t, authorityMatched(pats, dockerAltGatewayHostname), + "deny must match gateway.docker.internal") + assert.True(t, authorityMatched(pats, "HOST.DOCKER.INTERNAL"), + "deny must be case-insensitive (uppercase must not bypass)") + assert.True(t, authorityMatched(pats, dockerGatewayHostname+"."), + "deny must match a trailing-dot FQDN (host.docker.internal.)") + assert.True(t, authorityMatched(pats, dockerDefaultBridgeGatewayIP+"."), + "deny must match a trailing-dot gateway IP (172.17.0.1.)") + assert.False(t, authorityMatched(pats, dockerGatewayHostname+".attacker.com"), + "deny must be anchored (lookalike suffix must not match)") + + // The dead destination_ip L3 rule must be gone. + raw, err := json.Marshal(listener) + require.NoError(t, err) + assert.NotContains(t, string(raw), "destination_ip", + "destination_ip is inert for a forward proxy and must not be used") + } + }) + } +} + +// TestHostMatchRegex verifies that AllowHost entries translate to :authority +// patterns matching Squid's dstdomain semantics: a leading dot (or "*.") matches +// the apex and all subdomains, a bare host matches exactly, and every form +// tolerates an optional ":port" (HTTPS CONNECT authorities carry the port). +// Crucially, no form may match a lookalike parent domain like +// "example.com.attacker.com". +func TestHostMatchRegex(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + host string + matches []string + noMatches []string + }{ + { + name: "exact host, no subdomains", + host: "example.com", + // A trailing-dot FQDN resolves identically and must match, or it becomes + // a bypass on the deny path. + matches: []string{"example.com", "example.com:443", "EXAMPLE.COM", "example.com.", "example.com.:443"}, + noMatches: []string{"www.example.com", "notexample.com", "example.com.attacker.com"}, + }, + { + name: "leading dot matches apex and subdomains", + host: ".example.com", + matches: []string{"example.com", "www.example.com", "a.b.example.com", "www.example.com:8080"}, + noMatches: []string{"notexample.com", "example.com.attacker.com", "evil.com"}, + }, + { + name: "asterisk-dot is an alias for leading dot", + host: "*.example.com", + matches: []string{"example.com", "api.example.com", "api.example.com:443"}, + noMatches: []string{"attacker-example.com", "example.com.attacker.com"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Envoy anchors safe_regex fully; replicate that with ^...$ here. + re := regexp.MustCompile("^(?:" + hostMatchRegex(tt.host) + ")$") + for _, m := range tt.matches { + assert.True(t, re.MatchString(m), "expected %q to match AllowHost %q", m, tt.host) + } + for _, nm := range tt.noMatches { + assert.False(t, re.MatchString(nm), "expected %q NOT to match AllowHost %q", nm, tt.host) + } + }) + } +} + +// TestBuildEgressListener_EmptyAllowHostDenyAll is a mandatory regression guard: +// buildEgressListener with an empty AllowHost and InsecureAllowAll=false must +// produce a listener where the RBAC filter is present with action=ALLOW and +// zero policies. This is Envoy's deny-all behavior. The test guards against a +// bug that silently omits the RBAC filter and produces allow-all. +// +// Note: a JSON round-trip assertion is intentionally omitted here. The +// envoyNetworkFilter.TypedConfig field is typed as any, so concrete pointer +// types (*envoyHTTPRBAC, etc.) do not survive JSON round-trip — the unmarshaled +// value becomes map[string]any. Behavioral correctness is verified directly on +// the Go struct returned by buildEgressListener. +func TestBuildEgressListener_EmptyAllowHostDenyAll(t *testing.T) { + t.Parallel() + + spec := proxySpec{ + WorkloadName: "guard-test", + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + InsecureAllowAll: false, + AllowHost: []string{}, + }, + }, + AllowDockerGateway: false, + GatewayIP: dockerDefaultBridgeGatewayIP, + } + + listener := buildEgressListener(spec) + require.NotEmpty(t, listener.FilterChains) + + rbac := findRBACFilter(listener) + require.NotNil(t, rbac, + "RBAC filter must be present — its absence would silently allow all traffic") + assert.Equal(t, "ALLOW", rbac.Rules.Action, + "action must be ALLOW; an empty policy set under ALLOW is Envoy's deny-all") + assert.Empty(t, rbac.Rules.Policies, + "policy set must be empty to achieve deny-all semantics") + + // Verify the config serialises to valid JSON. + raw, err := json.Marshal(listener) + require.NoError(t, err) + assert.NotEmpty(t, raw, "serialized listener must not be empty") +} + +// TestBuildIngressListener_PortAndHostGating verifies that buildIngressListener +// wires the upstream port, host-port binding, and virtual-host domain gating +// correctly for several input scenarios. +func TestBuildIngressListener_PortAndHostGating(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + spec proxySpec + hostPort int + wantUpstreamRef string // substring that must appear in the listener JSON + wantHostPortBound int // the listener's bind port + wantDomains []string + }{ + { + name: "sse transport binds hostPort and routes to upstream", + spec: proxySpec{ + WorkloadName: "myserver", + UpstreamPort: 8080, + TransportType: "sse", + Permissions: nil, + }, + hostPort: 18080, + wantUpstreamRef: "myserver", + wantHostPortBound: 18080, + }, + { + // Inbound.AllowHost is intentionally ignored for Envoy's ingress virtual + // host domains. The transparent proxy rewrites Host to include the port + // ("127.0.0.1:19090"), which would not match a bare hostname list. The + // inbound access restriction is enforced by the 127.0.0.1 port binding. + name: "inbound AllowHost is ignored — ingress always uses wildcard domain", + spec: proxySpec{ + WorkloadName: "svc", + UpstreamPort: 9090, + TransportType: "streamable-http", + Permissions: &permissions.NetworkPermissions{ + Inbound: &permissions.InboundNetworkPermissions{ + AllowHost: []string{"app.example.com"}, + }, + }, + }, + hostPort: 19090, + wantUpstreamRef: "svc", + wantHostPortBound: 19090, + wantDomains: []string{`"*"`}, + }, + { + // No inbound AllowHost → wildcard virtual host. This is safe because + // the ingress listener is bound to 127.0.0.1 on the host, so only the + // local proxy runner can reach it regardless of the vhost domain. + name: "nil permissions defaults to wildcard virtual host", + spec: proxySpec{ + WorkloadName: "tool", + UpstreamPort: 7070, + TransportType: "sse", + Permissions: nil, + }, + hostPort: 17070, + wantUpstreamRef: "tool", + wantHostPortBound: 17070, + wantDomains: []string{`"*"`}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + listener := buildIngressListener(tt.spec, tt.hostPort) + + // Serialize to JSON for substring assertions — this is simpler + // than deeply navigating the typed structs and more resilient to + // internal refactors that rename unexported fields. + raw, err := json.Marshal(listener) + require.NoError(t, err) + s := string(raw) + + assert.Contains(t, s, tt.wantUpstreamRef, + "listener config must reference upstream %s", tt.wantUpstreamRef) + + if tt.wantDomains != nil { + for _, domain := range tt.wantDomains { + assert.Contains(t, s, domain, + "listener config must contain domain restriction %q", domain) + } + } + + // The listener must not be bound on port 0. + assert.NotContains(t, s, `"port_value":0`, + "listener must not bind on port 0") + }) + } +} + +// TestBuildIngressCluster_UpstreamAddress verifies that buildIngressCluster +// produces a STRICT_DNS cluster pointing at the correct workload address. +func TestBuildIngressCluster_UpstreamAddress(t *testing.T) { + t.Parallel() + + spec := proxySpec{ + WorkloadName: "myserver", + UpstreamPort: 8080, + } + + cluster := buildIngressCluster(spec) + + assert.Equal(t, ingressClusterName, cluster.Name) + assert.Equal(t, "STRICT_DNS", cluster.Type) + // V4_ONLY prevents slow AAAA lookup timeouts when IPv6 is unavailable. + // A regression dropping this would fail silently (extra latency only). + assert.Equal(t, "V4_ONLY", cluster.DnsLookupFamily, + "ingress STRICT_DNS cluster must use V4_ONLY to avoid IPv6 lookup timeouts") + require.NotNil(t, cluster.LoadAssignment) + require.NotEmpty(t, cluster.LoadAssignment.Endpoints) + require.NotEmpty(t, cluster.LoadAssignment.Endpoints[0].LBEndpoints) + + ep := cluster.LoadAssignment.Endpoints[0].LBEndpoints[0] + assert.Equal(t, "myserver", ep.Endpoint.Address.SocketAddress.Address) + assert.Equal(t, 8080, ep.Endpoint.Address.SocketAddress.PortValue) +} + +// TestBuildEgressCluster_DFPConfig verifies that buildEgressCluster produces a +// dynamic_forward_proxy cluster with the expected configuration. +func TestBuildEgressCluster_DFPConfig(t *testing.T) { + t.Parallel() + + cluster := buildEgressCluster() + + assert.Equal(t, dfpClusterName, cluster.Name) + assert.Equal(t, "CLUSTER_PROVIDED", cluster.LbPolicy) + require.NotNil(t, cluster.ClusterType) + assert.Equal(t, "envoy.clusters.dynamic_forward_proxy", cluster.ClusterType.Name) + + dfp, ok := cluster.ClusterType.TypedConfig.(*envoyDFPClusterConfig) + require.True(t, ok, "ClusterType.TypedConfig must be *envoyDFPClusterConfig") + assert.Equal(t, typeDFPCluster, dfp.Type) + assert.Equal(t, dfpCacheName, dfp.DNSCacheConfig.Name) + assert.Equal(t, "V4_ONLY", dfp.DNSCacheConfig.DNSLookupFamily) +} + +// TestWriteEnvoyBootstrap_FileMode verifies that writeEnvoyBootstrap writes a +// valid JSON bootstrap file at mode 0600. +func TestWriteEnvoyBootstrap_FileMode(t *testing.T) { + t.Parallel() + + b := envoyBootstrap{ + StaticResources: envoyStaticResources{}, + } + + path, err := writeEnvoyBootstrap(b) + require.NoError(t, err) + require.NotEmpty(t, path, "returned path must be non-empty") + + t.Cleanup(func() { _ = os.Remove(path) }) + + // File must exist and be readable. + info, err := os.Stat(path) + require.NoError(t, err) + + // Mode must be 0644: world-readable so the Envoy distroless container + // (UID 101) can read the bind-mounted file on Linux Docker Engine, where + // strict POSIX permissions prevent a different UID from reading 0600 files. + // The bootstrap contains no secrets, so world-readable is safe. + assert.Equal(t, os.FileMode(0o644), info.Mode().Perm(), + "bootstrap file must be written at mode 0644") + + // File must contain valid JSON that deserializes back into envoyBootstrap. + data, err := os.ReadFile(path) + require.NoError(t, err) + require.NotEmpty(t, data, "bootstrap file must not be empty") + + var roundTripped envoyBootstrap + require.NoError(t, json.Unmarshal(data, &roundTripped), + "bootstrap file must contain valid JSON") +} + +// TestEnvoyAdmin_Absent asserts that the admin interface is entirely absent from +// the generated bootstrap. Omitting the admin block causes Envoy to skip the +// admin server, removing the attack surface without affecting proxy behaviour. +func TestEnvoyAdmin_Absent(t *testing.T) { + t.Parallel() + + b := envoyBootstrap{ + StaticResources: envoyStaticResources{}, + } + + path, err := writeEnvoyBootstrap(b) + require.NoError(t, err) + require.NotEmpty(t, path) + t.Cleanup(func() { _ = os.Remove(path) }) + + data, err := os.ReadFile(path) + require.NoError(t, err) + + var m map[string]any + require.NoError(t, json.Unmarshal(data, &m), "bootstrap must be valid JSON") + + _, hasAdmin := m["admin"] + assert.False(t, hasAdmin, "bootstrap must not contain an admin block") +} + +// TestGetEnvoyImage verifies that getEnvoyImage returns the default image when +// the override env var is unset and the override when it is set. +// NOTE: t.Setenv is used so t.Parallel() is intentionally omitted here — env +// mutations are global state and are incompatible with parallel execution. +func TestGetEnvoyImage(t *testing.T) { + tests := []struct { + name string + envValue string + wantImage string + wantEnvoy bool // true: assert the result contains "envoy" + }{ + { + name: "empty env returns non-empty default containing envoy", + envValue: "", + wantEnvoy: true, + }, + { + name: "custom image override is returned verbatim", + envValue: "my-custom-envoy:latest", + wantImage: "my-custom-envoy:latest", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("TOOLHIVE_ENVOY_IMAGE", tt.envValue) + + got := getEnvoyImage() + require.NotEmpty(t, got, "getEnvoyImage must never return an empty string") + + if tt.wantEnvoy { + assert.Contains(t, got, "envoy", + "default image must contain 'envoy' to be identifiable") + } + + if tt.wantImage != "" { + assert.Equal(t, tt.wantImage, got) + } + }) + } +} + +// TestEgressListenerHCMTypeURLs verifies that the egress listener serializes +// with correct protobuf @type URLs so Envoy can parse it. +func TestEgressListenerHCMTypeURLs(t *testing.T) { + t.Parallel() + + spec := proxySpec{ + WorkloadName: "myserver", + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + }, + }, + AllowDockerGateway: false, + GatewayIP: dockerDefaultBridgeGatewayIP, + } + + listener := buildEgressListener(spec) + raw, err := json.Marshal(listener) + require.NoError(t, err) + s := string(raw) + + assert.Contains(t, s, typeHCM, "@type for HCM must be present in serialized JSON") + assert.Contains(t, s, typeHTTPRBAC, "@type for RBAC must be present in serialized JSON") + assert.Contains(t, s, typeDFPFilter, "@type for DFP filter must be present in serialized JSON") + assert.Contains(t, s, typeRouter, "@type for router must be present in serialized JSON") +} + +// TestEgressClusterTypeURL verifies that the egress cluster serializes with the +// correct DFP cluster @type URL. +func TestEgressClusterTypeURL(t *testing.T) { + t.Parallel() + + cluster := buildEgressCluster() + raw, err := json.Marshal(cluster) + require.NoError(t, err) + s := string(raw) + + assert.Contains(t, s, typeDFPCluster, + "@type for DFP cluster config must be present in serialized JSON") +} + +// TestEnvoyBootstrap_ValidatesAgainstRealEnvoy asserts that the generated +// bootstrap is accepted by a real Envoy (`--mode validate`), across the deny, +// allow-gateway, and allowlist permutations. Hand-rolled protobuf-JSON can pass +// Go-level unit tests yet be rejected by Envoy (wrong @type, bad matcher shape), +// so this closes that gap. Skips when docker is unavailable. +func TestEnvoyBootstrap_ValidatesAgainstRealEnvoy(t *testing.T) { + t.Parallel() + + if _, err := exec.LookPath("docker"); err != nil { + t.Skip("docker not available; skipping real-Envoy validation") + } + + cases := []struct { + name string + spec proxySpec + }{ + { + name: "allowlist + gateway deny + ingress", + spec: proxySpec{ + WorkloadName: "demo", TransportType: "streamable-http", UpstreamPort: 8080, + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{AllowHost: []string{"example.com", ".github.com"}}, + Inbound: &permissions.InboundNetworkPermissions{AllowHost: []string{"app.internal"}}, + }, + }, + }, + { + name: "allow-docker-gateway + insecure-allow-all + stdio (egress only)", + spec: proxySpec{ + WorkloadName: "demo", TransportType: "stdio", AllowDockerGateway: true, + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{Outbound: &permissions.OutboundNetworkPermissions{InsecureAllowAll: true}}, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + listeners := []envoyListener{buildEgressListener(tc.spec)} + clusters := []envoyCluster{buildEgressCluster()} + if tc.spec.TransportType != "stdio" { + listeners = append(listeners, buildIngressListener(tc.spec, 18080)) + clusters = append(clusters, buildIngressCluster(tc.spec)) + } + b := envoyBootstrap{ + StaticResources: envoyStaticResources{Listeners: listeners, Clusters: clusters}, + } + cfg, err := json.Marshal(b) + require.NoError(t, err) + + // Pass the config inline via --config-yaml (JSON is valid YAML) rather + // than bind-mounting a file: the bootstrap is written 0600 and envoy + // runs as nonroot, so a bind-mounted file is unreadable to the + // container on native-Linux CI. Inline config sidesteps that entirely. + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + //nolint:gosec // G204: fixed args; cfg is test-generated JSON + out, err := exec.CommandContext(ctx, "docker", "run", "--rm", + defaultEnvoyImage, "--mode", "validate", "--config-yaml", string(cfg)).CombinedOutput() + require.NoError(t, err, "envoy rejected the generated config:\n%s", out) + assert.Contains(t, string(out), "configuration '' OK") + }) + } +} + +// TestEnvoyProxy_SetupOrchestration covers the container-orchestration branch +// logic of SetupEgress/SetupIngress without a live daemon. SetupEgress must NOT +// create any container — it only returns env vars. SetupIngress creates the +// Envoy container (named -egress) and returns the ingress port (0 for +// stdio). +func TestEnvoyProxy_SetupOrchestration(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + transport string + upstreamPort int + wantIngress bool + }{ + {name: "non-stdio reserves an ingress port", transport: "streamable-http", upstreamPort: 8080, wantIngress: true}, + {name: "stdio reserves no ingress port", transport: "stdio", upstreamPort: 0, wantIngress: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var createdName string + api := &fakeDockerAPI{ + createFunc: func(_ context.Context, _ *container.Config, _ *container.HostConfig, _ *network.NetworkingConfig, _ *v1.Platform, name string) (container.CreateResponse, error) { + createdName = name + return container.CreateResponse{ID: "cid-envoy"}, nil + }, + startFunc: func(_ context.Context, _ string, _ mobyclient.ContainerStartOptions) error { return nil }, + } + c := &Client{ + api: api, + imageManager: &fakeImageManager{availableImages: map[string]struct{}{defaultEnvoyImage: {}}}, + } + e := &envoyProxy{client: c} + + spec := proxySpec{ + WorkloadName: "app", + TransportType: tt.transport, + UpstreamPort: tt.upstreamPort, + GatewayIP: dockerDefaultBridgeGatewayIP, + Endpoints: map[string]*network.EndpointSettings{}, + } + + // SetupEgress must NOT create the container — it only returns env vars. + // Container creation is deferred to SetupIngress so the MCP hostname + // resolves on first probe (see #5922). + egress, err := e.SetupEgress(t.Context(), spec) + require.NoError(t, err) + assert.Empty(t, createdName, "SetupEgress must not create the container") + assert.Equal(t, "http://app-egress:3128", egress.EnvVars["HTTP_PROXY"]) + + // SetupIngress creates the container and returns the ingress port. + ingressPort, err := e.SetupIngress(t.Context(), spec, egress) + require.NoError(t, err) + assert.Equal(t, "app-egress", createdName, "SetupIngress must create the -egress container") + if tt.wantIngress { + assert.Positive(t, ingressPort, "non-stdio must reserve an ingress port") + } else { + assert.Zero(t, ingressPort, "stdio must not reserve an ingress port") + } + }) + } +} + +// TestEnvoyProxy_SetupIngress_ConcurrentSameUpstreamPort guards against a +// deterministic ingress port collision: two workloads built from the same +// image share the same fixed UpstreamPort, so deriving the ingress port from +// UpstreamPort+1 made every concurrently-starting instance request the +// identical host port. The fix picks a random port instead, so concurrent +// SetupIngress calls sharing UpstreamPort must land on different ports. +func TestEnvoyProxy_SetupIngress_ConcurrentSameUpstreamPort(t *testing.T) { + t.Parallel() + + const upstreamPort = 8080 + const workloadCount = 4 + + ports := make([]int, workloadCount) + errs := make([]error, workloadCount) + + var wg sync.WaitGroup + for i := range workloadCount { + wg.Add(1) + go func(idx int) { + defer wg.Done() + api := &fakeDockerAPI{ + createFunc: func(_ context.Context, _ *container.Config, _ *container.HostConfig, + _ *network.NetworkingConfig, _ *v1.Platform, _ string) (container.CreateResponse, error) { + return container.CreateResponse{ID: "cid-envoy"}, nil + }, + startFunc: func(_ context.Context, _ string, _ mobyclient.ContainerStartOptions) error { return nil }, + } + e := &envoyProxy{client: &Client{ + api: api, + imageManager: &fakeImageManager{availableImages: map[string]struct{}{defaultEnvoyImage: {}}}, + }} + spec := proxySpec{ + WorkloadName: fmt.Sprintf("app-%d", idx), + TransportType: "streamable-http", + UpstreamPort: upstreamPort, + GatewayIP: dockerDefaultBridgeGatewayIP, + Endpoints: map[string]*network.EndpointSettings{}, + } + port, err := e.SetupIngress(t.Context(), spec, egressResult{}) + ports[idx] = port + errs[idx] = err + }(i) + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent SetupIngress calls") + } + + seen := make(map[int]bool, workloadCount) + for i, port := range ports { + require.NoError(t, errs[i]) + assert.NotEqual(t, upstreamPort+1, port, "ingress port must not be derived from UpstreamPort+1") + assert.False(t, seen[port], "port %d was reserved by more than one concurrent workload", port) + seen[port] = true + } +} + +// TestListenerTimeoutsDisabledForStreaming guards against Envoy's 15s default +// RouteAction.timeout (a hard total-response cap) and 5m stream_idle_timeout +// truncating long-lived MCP streams (SSE / streamable-http). Both the ingress +// reverse-proxy route and the egress plain-HTTP route must disable the route +// timeout, and both HCMs must disable the idle timeout. +func TestListenerTimeoutsDisabledForStreaming(t *testing.T) { + t.Parallel() + + spec := proxySpec{ + WorkloadName: "app", + TransportType: "streamable-http", + UpstreamPort: 8080, + GatewayIP: dockerDefaultBridgeGatewayIP, + } + + hcmOf := func(l envoyListener) *envoyHCM { + hcm, ok := l.FilterChains[0].Filters[0].TypedConfig.(*envoyHCM) + require.True(t, ok, "first network filter must be the HCM") + return hcm + } + // routeTimeouts returns the Timeout of every route action in an HCM. + routeTimeouts := func(hcm *envoyHCM) []string { + var out []string + for _, vh := range hcm.RouteConfig.VirtualHosts { + for _, r := range vh.Routes { + if r.Route != nil { + out = append(out, r.Route.Timeout) + } + } + } + return out + } + + ingress := hcmOf(buildIngressListener(spec, 18080)) + assert.Equal(t, "0s", ingress.StreamIdleTimeout, "ingress HCM must disable the idle timeout") + for _, tmo := range routeTimeouts(ingress) { + assert.Equal(t, "0s", tmo, "ingress route must disable the 15s total-response cap") + } + + egress := hcmOf(buildEgressListener(spec)) + assert.Equal(t, "0s", egress.StreamIdleTimeout, "egress HCM must disable the idle timeout") + // The plain-HTTP egress route must disable the cap; the CONNECT route is + // unaffected by the default (no end-of-stream while tunneling) so it need not. + var sawPlainHTTP bool + for _, vh := range egress.RouteConfig.VirtualHosts { + for _, r := range vh.Routes { + if r.Match.Prefix == "/" && r.Route != nil { + sawPlainHTTP = true + assert.Equal(t, "0s", r.Route.Timeout, "egress plain-HTTP route must disable the 15s cap") + } + } + } + assert.True(t, sawPlainHTTP, "egress must have a plain-HTTP route") +} + +// TestBuildAllowlistPolicies_AllowPort verifies that AllowPort entries are +// translated into combined host+port regex matchers, matching Squid's +// allowed_ports ACL behaviour including plain-HTTP bare-hostname support. +// Tests cover: +// - AllowPort alone (any host on listed ports, bare hostname allowed when port 80 in list) +// - AllowHost + AllowPort (both must match; bare hostname allowed when port 80 in list) +// - AllowHost without AllowPort (any port, original behaviour) +// - InsecureAllowAll ignores port constraints +func TestBuildAllowlistPolicies_AllowPort(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + out *permissions.OutboundNetworkPermissions + // authority strings that must be allowed / denied + allowed []string + denied []string + }{ + { + name: "AllowPort only — any host on listed ports; bare hostname allowed when port 80 listed", + out: &permissions.OutboundNetworkPermissions{AllowPort: []int{80, 443}}, + // Bare "example.com" = plain HTTP = port 80 → allowed + allowed: []string{"example.com:80", "anything.io:443", "example.com"}, + denied: []string{"example.com:8080", "example.com:22"}, + }, + { + name: "AllowHost + AllowPort — host AND port must match", + out: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + AllowPort: []int{443}, + }, + allowed: []string{"example.com:443"}, + denied: []string{ + "example.com:80", // right host, wrong port + "other.com:443", // right port, wrong host + "other.com:80", // neither + }, + }, + { + name: "AllowHost without AllowPort — any port allowed", + out: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + }, + allowed: []string{"example.com:443", "example.com:8080", "example.com:22"}, + denied: []string{"other.com:443"}, + }, + { + name: "InsecureAllowAll ignores AllowPort", + out: &permissions.OutboundNetworkPermissions{InsecureAllowAll: true, AllowPort: []int{443}}, + allowed: []string{"example.com:80", "anything.io:9999"}, + denied: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{Outbound: tt.out}, + } + listener := buildEgressListener(spec) + for _, a := range tt.allowed { + assert.True(t, listenerAllowsAuthority(listener, a), + "authority %q should be allowed", a) + } + for _, a := range tt.denied { + assert.False(t, listenerAllowsAuthority(listener, a), + "authority %q should be denied", a) + } + }) + } +} + +// permMatchesAuthority tests a single envoyPermission against an authority +// string, replicating Envoy's full-string anchoring for safe_regex and +// exact/suffix/prefix string matches. +func permMatchesAuthority(perm envoyPermission, authority string) bool { + if perm.Any != nil && *perm.Any { + return true + } + if perm.Header != nil && perm.Header.Name == ":authority" && perm.Header.StringMatch != nil { + sm := perm.Header.StringMatch + switch { + case sm.SafeRegex != nil: + return regexp.MustCompile("^(?:" + sm.SafeRegex.Regex + ")$").MatchString(authority) + case sm.Exact != "": + return authority == sm.Exact + case sm.Suffix != "": + return strings.HasSuffix(authority, sm.Suffix) + case sm.Prefix != "": + return strings.HasPrefix(authority, sm.Prefix) + } + } + if perm.OrRules != nil { + for _, r := range perm.OrRules.Rules { + if permMatchesAuthority(r, authority) { + return true + } + } + } + return false +} + +// TestBuildAllowlistPolicies_AllowPort_Structure guards against a wrong +// permMatchesAuthority re-implementation masking a config regression: it checks +// the actual RBAC policy shape (single SafeRegex combining host + port) rather +// than just whether the helper says the authority is allowed or denied (F6). +func TestBuildAllowlistPolicies_AllowPort_Structure(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + AllowPort: []int{80, 443}, + }, + }, + } + listener := buildEgressListener(spec) + rbac := findRBACFilter(listener) + require.NotNil(t, rbac) + policy, ok := rbac.Rules.Policies["example.com"] + require.True(t, ok, "expected policy keyed 'example.com'") + // Must have exactly 1 permission: a single SafeRegex encoding both host and port. + require.Len(t, policy.Permissions, 1, + "AllowHost+AllowPort policy must have exactly 1 combined safe_regex permission") + perm := policy.Permissions[0] + require.NotNil(t, perm.Header, "permission must be a header matcher") + require.NotNil(t, perm.Header.StringMatch, "permission must have a string match") + require.NotNil(t, perm.Header.StringMatch.SafeRegex, "permission must use safe_regex") + // The regex must actually match the expected authorities — including bare + // hostname (plain HTTP) because port 80 is listed. Evaluating the full regex + // here rather than inspecting its text guards against a reimplementation that + // produces a different but equivalent string (N2: the previous assertion of + // assert.Contains(re,"?") was vacuous because "?" appears in non-capturing groups). + re := perm.Header.StringMatch.SafeRegex.Regex + compile := func(r string) *regexp.Regexp { return regexp.MustCompile("^(?:" + r + ")$") } + assert.True(t, compile(re).MatchString("example.com"), "bare hostname must match (plain HTTP, port 80 implied)") + assert.True(t, compile(re).MatchString("example.com:80"), "explicit port 80 must match") + assert.True(t, compile(re).MatchString("example.com:443"), "explicit port 443 must match") + assert.False(t, compile(re).MatchString("example.com:8080"), "non-listed port must not match") + assert.False(t, compile(re).MatchString("other.com:443"), "wrong host must not match") +} + +// TestBuildAllowlistPolicies_AdditionalCases covers edge cases not in the main +// AllowPort table (F7): multi-port OR shape, port-less authority divergence from +// Squid, and gateway-deny ordering under AllowPort-only (F8). +func TestBuildAllowlistPolicies_AdditionalCases(t *testing.T) { + t.Parallel() + + // F7a: AllowHost + multi-port — authority with a non-listed port is denied. + t.Run("AllowHost+multi-port denies non-listed port", func(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + AllowPort: []int{80, 443}, + }, + }, + } + l := buildEgressListener(spec) + assert.True(t, listenerAllowsAuthority(l, "example.com:443"), "port 443 must be allowed") + assert.True(t, listenerAllowsAuthority(l, "example.com:80"), "port 80 must be allowed") + assert.False(t, listenerAllowsAuthority(l, "example.com:8080"), "port 8080 must be denied") + assert.False(t, listenerAllowsAuthority(l, "other.com:443"), "wrong host must be denied") + }) + + // F7b: Port-less authority (plain HTTP default port) — Squid parity. + // Plain HTTP omits the port from :authority ("example.com" not "example.com:80"). + // When port 80 is in AllowPort the port suffix is made optional so both the + // bare form and the explicit ":80" form are accepted, matching Squid behaviour. + t.Run("port-less authority is allowed when port 80 is in AllowPort (Squid parity)", func(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + AllowPort: []int{80}, + }, + }, + } + l := buildEgressListener(spec) + // Bare hostname (plain HTTP) must be allowed when port 80 is listed. + assert.True(t, listenerAllowsAuthority(l, "example.com"), + "port-less authority must be allowed when port 80 is in AllowPort (plain HTTP parity with Squid)") + assert.True(t, listenerAllowsAuthority(l, "example.com:80"), + "explicit port 80 must be allowed") + assert.False(t, listenerAllowsAuthority(l, "example.com:443"), + "port 443 must not be allowed when only port 80 is listed") + // Bare hostname must be denied when port 80 is NOT in the list. + spec443 := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"example.com"}, + AllowPort: []int{443}, + }, + }, + } + l443 := buildEgressListener(spec443) + assert.False(t, listenerAllowsAuthority(l443, "example.com"), + "port-less authority must be denied when port 80 is NOT in AllowPort") + assert.True(t, listenerAllowsAuthority(l443, "example.com:443"), + "explicit port 443 must be allowed") + }) + + // F8: Gateway-deny DENY filter is evaluated before the ALLOW filter even + // when AllowPort-only is configured (deny-before-allow ordering). + t.Run("gateway-deny precedes AllowPort-only allow", func(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowPort: []int{443}, + }, + }, + } + l := buildEgressListener(spec) + // The gateway deny filter must precede the allow filter. + deny := findDenyRBACFilter(l) + require.NotNil(t, deny, "DENY filter must be present") + hcm := l.FilterChains[0].Filters[0].TypedConfig.(*envoyHCM) + var denyIdx, allowIdx int + for i, f := range hcm.HTTPFilters { + if rbac, ok := f.TypedConfig.(*envoyHTTPRBAC); ok { + if rbac.Rules.Action == "DENY" { + denyIdx = i + } else { + allowIdx = i + } + } + } + assert.Less(t, denyIdx, allowIdx, "DENY filter must appear before ALLOW filter") + // Gateway IP must be denied even though port 443 is in AllowPort. + gatewayPats := collectAuthorityRegexes(deny) + assert.True(t, authorityMatched(gatewayPats, dockerDefaultBridgeGatewayIP+":443"), + "gateway IP on port 443 must be caught by the DENY filter") + }) + + // N3: AllowPort-only with port 443 (no port 80) — bare hostname denied, + // explicit :443 allowed. Exercises the includes80=false branch of anyHostPortRegex. + t.Run("AllowPort-only port-443 — bare hostname denied, explicit port allowed", func(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowPort: []int{443}, + }, + }, + } + l := buildEgressListener(spec) + assert.True(t, listenerAllowsAuthority(l, "anything.io:443"), "port 443 must be allowed") + assert.False(t, listenerAllowsAuthority(l, "anything.io"), "bare hostname must be denied (port 80 not listed)") + assert.False(t, listenerAllowsAuthority(l, "anything.io:80"), "port 80 must be denied") + }) + + // N4: Wildcard host + AllowPort — the subdomain (.*\.)? path in hostPortMatchRegex. + t.Run("wildcard AllowHost + AllowPort matches subdomains on listed ports only", func(t *testing.T) { + t.Parallel() + spec := proxySpec{ + WorkloadName: "app", + GatewayIP: dockerDefaultBridgeGatewayIP, + Permissions: &permissions.NetworkPermissions{ + Outbound: &permissions.OutboundNetworkPermissions{ + AllowHost: []string{"*.example.com"}, + AllowPort: []int{443}, + }, + }, + } + l := buildEgressListener(spec) + assert.True(t, listenerAllowsAuthority(l, "api.example.com:443"), "subdomain on allowed port must match") + assert.True(t, listenerAllowsAuthority(l, "example.com:443"), "apex on allowed port must match") + assert.False(t, listenerAllowsAuthority(l, "api.example.com:80"), "subdomain on disallowed port must not match") + assert.False(t, listenerAllowsAuthority(l, "api.example.com"), "bare subdomain must not match (port 80 not listed)") + assert.False(t, listenerAllowsAuthority(l, "other.com:443"), "non-matching host must not match") + }) +} diff --git a/pkg/container/docker/networkproxy.go b/pkg/container/docker/networkproxy.go new file mode 100644 index 0000000000..7b27a74a34 --- /dev/null +++ b/pkg/container/docker/networkproxy.go @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package docker + +import ( + "context" + "fmt" + "os" + + "github.com/moby/moby/api/types/network" + + "github.com/stacklok/toolhive-core/permissions" +) + +// networkProxy is the single enforcement point for outbound allowlisting, +// docker-gateway blocking, and inbound reverse-proxying for isolated workloads. +// +// What it enables: all egress and ingress proxy containers are created through +// this interface, ensuring a consistent policy-enforcement seam that can be +// swapped at startup via TOOLHIVE_NETWORK_PROXY. +// +// What it does NOT solve: non-cooperative egress is contained by the +// Internal:true network blackhole created by createNetwork, not by this proxy. +// True L3/L4 traffic interception is a separate Phase 2 concern and is not +// addressed here. +type networkProxy interface { + // SetupEgress provisions egress enforcement BEFORE the MCP container is + // created and returns the environment variables to inject into the workload + // (HTTP_PROXY etc.). + // + // Squid creates its egress container here. The Envoy backend creates no + // container in SetupEgress — it returns only env vars — and defers all + // container creation (both egress and ingress listeners, one container) to + // SetupIngress so the STRICT_DNS ingress cluster can resolve the MCP hostname + // on its first probe. + // + // It must run before createMcpContainer so the returned env vars land in the + // workload's environment. + SetupEgress(ctx context.Context, spec proxySpec) (egressResult, error) + + // SetupIngress creates the ingress proxy AFTER the MCP container exists and + // returns the host-side ingress port (0 for stdio / UpstreamPort==0). + // + // Squid creates its ingress container here. The Envoy backend creates its + // single dual-listener container here (both egress forward-proxy and ingress + // reverse-proxy listeners). Running after createMcpContainer ensures the MCP + // container's hostname resolves on first probe, avoiding a cached negative DNS + // lookup that would leave the ingress permanently unable to reach the upstream. + // + // It must run after createMcpContainer. + SetupIngress(ctx context.Context, spec proxySpec, egress egressResult) (int, error) +} + +// proxySpec contains all the parameters needed to set up proxy containers for +// an isolated workload. +type proxySpec struct { + // WorkloadName is the base name of the MCP container (e.g. "myserver"). + WorkloadName string + // UpstreamHost is the address the ingress reverse proxy connects to for the + // MCP container. It is the container's resolved IP on the internal network + // (not its name) so the ingress proxy has no DNS dependency: under + // concurrent startup the container's DNS record can lag, and a hostname-based + // Squid cache_peer caches the failed lookup and never recovers within the + // readiness window (see #6063). Empty falls back to WorkloadName. + UpstreamHost string + // Permissions holds the network permission profile from the workload's + // permission profile, governing what outbound traffic is allowed. + Permissions *permissions.NetworkPermissions + // AllowDockerGateway, when true, skips the docker-gateway deny rules in the + // egress proxy configuration. + AllowDockerGateway bool + // GatewayIP is the Docker bridge gateway IP resolved at runtime. + GatewayIP string + // TransportType is the MCP transport in use (e.g. "stdio", "sse", + // "streamable-http"). A value of "stdio" suppresses ingress proxy creation. + TransportType string + // UpstreamPort is the container port the MCP server listens on. Ignored + // when TransportType is "stdio" or the value is 0. + UpstreamPort int + // AttachStdio controls whether the proxy containers attach stdio streams. + AttachStdio bool + // Endpoints is the set of network endpoints the proxy containers should + // join, keyed by network name. + Endpoints map[string]*network.EndpointSettings +} + +// egressResult is the output of a successful SetupEgress call. It is passed to +// SetupIngress so backends can carry any state set up during egress forward +// without holding per-workload state on the shared proxy. +type egressResult struct { + // EnvVars contains environment variables that must be merged into the MCP + // container's environment (e.g. HTTP_PROXY, HTTPS_PROXY). + EnvVars map[string]string +} + +// newNetworkProxy reads the TOOLHIVE_NETWORK_PROXY environment variable and +// returns the corresponding networkProxy implementation. An empty value or +// "squid" selects the default squid-based proxy. Any other value returns an +// error so that misconfiguration is caught at startup. +func newNetworkProxy(c *Client) (networkProxy, error) { + val := os.Getenv("TOOLHIVE_NETWORK_PROXY") + switch val { + case "", "squid": + return &squidProxy{client: c}, nil + case "envoy": + return &envoyProxy{client: c}, nil + default: + return nil, fmt.Errorf("unknown TOOLHIVE_NETWORK_PROXY value %q: supported values are \"squid\" (default), \"envoy\"", val) + } +} + +// Compile-time assertion that squidProxy satisfies networkProxy. +var _ networkProxy = (*squidProxy)(nil) + +// Compile-time assertion that envoyProxy satisfies networkProxy. +var _ networkProxy = (*envoyProxy)(nil) diff --git a/pkg/container/docker/networkproxy_test.go b/pkg/container/docker/networkproxy_test.go new file mode 100644 index 0000000000..46c577dc35 --- /dev/null +++ b/pkg/container/docker/networkproxy_test.go @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package docker + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewNetworkProxy(t *testing.T) { + // t.Setenv is used in subtests so the outer test must NOT call t.Parallel(). + // The subtests run sequentially to avoid concurrent environment mutation. + + tests := []struct { + name string + envValue string + wantSquid bool + wantEnvoy bool + wantErr bool + errContains []string + }{ + { + name: "empty env returns squidProxy", + envValue: "", + wantSquid: true, + wantErr: false, + }, + { + name: "squid returns squidProxy", + envValue: "squid", + wantSquid: true, + wantErr: false, + }, + { + // After Stage 1 this case becomes a success path — envoy returns + // *envoyProxy, no error. Flip wantErr to false and add wantEnvoy + // assertion once envoy.go exists. + name: "envoy returns envoyProxy", + envValue: "envoy", + wantSquid: false, + wantEnvoy: true, + wantErr: false, + }, + { + name: "bogus value returns error", + envValue: "bogus", + wantSquid: false, + wantErr: true, + errContains: []string{"bogus"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // t.Setenv is incompatible with t.Parallel() — env mutations are + // global state; subtests run sequentially within this parent test. + t.Setenv("TOOLHIVE_NETWORK_PROXY", tt.envValue) + + c := &Client{} + proxy, err := newNetworkProxy(c) + + if tt.wantErr { + require.Error(t, err) + for _, substr := range tt.errContains { + assert.Contains(t, err.Error(), substr) + } + return + } + + require.NoError(t, err) + require.NotNil(t, proxy) + + if tt.wantSquid { + _, ok := proxy.(*squidProxy) + assert.True(t, ok, "expected proxy to be *squidProxy, got %T", proxy) + } + + if tt.wantEnvoy { + _, ok := proxy.(*envoyProxy) + assert.True(t, ok, "expected proxy to be *envoyProxy, got %T", proxy) + } + }) + } +} diff --git a/pkg/container/docker/squid.go b/pkg/container/docker/squid.go index 24a4f08a81..f784e3d802 100644 --- a/pkg/container/docker/squid.go +++ b/pkg/container/docker/squid.go @@ -42,6 +42,7 @@ func createIngressSquidContainer( ctx context.Context, c *Client, containerName string, + upstreamHost string, squidContainerName string, attachStdio bool, upstreamPort int, @@ -51,7 +52,7 @@ func createIngressSquidContainer( portBindings map[string][]runtime.PortBinding, networkPermissions *permissions.NetworkPermissions, ) (string, error) { - squidConfPath, err := createTempIngressSquidConf(containerName, upstreamPort, squidPort, networkPermissions) + squidConfPath, err := createTempIngressSquidConf(containerName, upstreamHost, upstreamPort, squidPort, networkPermissions) if err != nil { return "", fmt.Errorf("failed to create temporary squid.conf: %w", err) } @@ -353,8 +354,56 @@ func getSquidImage() string { return defaultSquidImage } +// squidProxy is the default networkProxy implementation. It creates egress and +// ingress Squid-based proxy containers for isolated workloads. +type squidProxy struct { + client *Client +} + +// SetupEgress creates the egress Squid container before the MCP container is +// created and returns the proxy env vars to inject into the workload. +func (s *squidProxy) SetupEgress(ctx context.Context, spec proxySpec) (egressResult, error) { + egressContainerName := fmt.Sprintf("%s-egress", spec.WorkloadName) + _, err := createEgressSquidContainer( + ctx, s.client, spec.WorkloadName, egressContainerName, + spec.AttachStdio, nil, spec.Endpoints, spec.Permissions, + spec.AllowDockerGateway, spec.GatewayIP, + ) + if err != nil { + return egressResult{}, fmt.Errorf("failed to create egress container: %w", err) + } + // ingressPort stays 0: squid creates and binds the ingress container later, + // in SetupIngress, once the MCP container's hostname resolves. + return egressResult{EnvVars: addEgressEnvVars(nil, egressContainerName)}, nil +} + +// SetupIngress creates the ingress Squid container after the MCP container +// exists. Creating it here (rather than before the MCP container) ensures the +// cache_peer hostname resolves on first probe; a Squid ingress created against a +// not-yet-existent upstream caches the negative DNS lookup and never recovers +// within the workload's readiness window. +func (s *squidProxy) SetupIngress(ctx context.Context, spec proxySpec, _ egressResult) (int, error) { + if spec.TransportType == "stdio" || spec.UpstreamPort == 0 { + return 0, nil + } + // Prefer the resolved upstream IP so the cache_peer has no DNS dependency; + // fall back to the workload name when unset (see proxySpec.UpstreamHost). + upstreamHost := spec.UpstreamHost + if upstreamHost == "" { + upstreamHost = spec.WorkloadName + } + ingressPort, err := s.client.setupIngressContainer( + ctx, spec.WorkloadName, upstreamHost, spec.UpstreamPort, spec.AttachStdio, spec.Endpoints, spec.Permissions, + ) + if err != nil { + return 0, err + } + return ingressPort, nil +} + func createTempIngressSquidConf( serverHostname string, + upstreamHost string, upstreamPort int, squidPort int, networkPermissions *permissions.NetworkPermissions, @@ -363,7 +412,7 @@ func createTempIngressSquidConf( writeCommonConfig(&sb, serverHostname, proxyIngress) - writeIngressProxyConfig(&sb, serverHostname, upstreamPort, squidPort, networkPermissions) + writeIngressProxyConfig(&sb, serverHostname, upstreamHost, upstreamPort, squidPort, networkPermissions) sb.WriteString("http_access deny all\n") tmpFile, err := os.CreateTemp("", "squid-*.conf") @@ -392,17 +441,26 @@ func createTempIngressSquidConf( func writeIngressProxyConfig( sb *strings.Builder, serverHostname string, + upstreamHost string, upstreamPort int, squidPort int, networkPermissions *permissions.NetworkPermissions, ) { portNum := strconv.Itoa(upstreamPort) squidPortNum := strconv.Itoa(squidPort) + // cache_peer targets upstreamHost (the MCP container's IP), not its name, so + // the peer has no DNS lookup to cache-and-latch on a cold-start miss (#6063). + // defaultsite keeps the name for the origin Host header. standby=2 keeps warm + // idle connections open to the upstream so the first request after a cold + // start (notably a long-lived GET SSE stream that a server-initiated request + // rides on) is forwarded without paying inline TCP connect latency. Without + // it, the cold first GET races behind a later POST that reuses a warmed path, + // reordering server->client streams. sb.WriteString( "\n# Reverse proxy setup for port " + portNum + "\n" + "http_port 0.0.0.0:" + squidPortNum + " accel defaultsite=" + serverHostname + "\n" + - "cache_peer " + serverHostname + " parent " + portNum + " 0 no-query originserver name=origin_" + - portNum + " connect-timeout=5 connect-fail-limit=5\n") + "cache_peer " + upstreamHost + " parent " + portNum + " 0 no-query originserver name=origin_" + + portNum + " connect-timeout=5 connect-fail-limit=5 standby=2\n") // Check if inbound network permissions are configured if networkPermissions != nil && networkPermissions.Inbound != nil && len(networkPermissions.Inbound.AllowHost) > 0 { diff --git a/pkg/container/docker/squid_test.go b/pkg/container/docker/squid_test.go index 483dbfaec3..0e99d630f2 100644 --- a/pkg/container/docker/squid_test.go +++ b/pkg/container/docker/squid_test.go @@ -193,7 +193,7 @@ func TestCreateTempEgressSquidConf_WithACLs(t *testing.T) { func TestCreateTempIngressSquidConf_Basics(t *testing.T) { t.Parallel() - fp, err := createTempIngressSquidConf("svc-example", 8080, 18080, nil) + fp, err := createTempIngressSquidConf("svc-example", "10.89.0.7", 8080, 18080, nil) require.NoError(t, err) t.Cleanup(func() { _ = os.Remove(fp) }) @@ -203,8 +203,13 @@ func TestCreateTempIngressSquidConf_Basics(t *testing.T) { assert.Contains(t, s, "visible_hostname svc-example-ingress") assert.Contains(t, s, "\n# Reverse proxy setup for port 8080\n") + // defaultsite keeps the name; cache_peer targets the resolved upstream IP so + // the peer has no DNS lookup to latch on (see #6063). assert.Contains(t, s, "http_port 0.0.0.0:18080 accel defaultsite=svc-example") - assert.Contains(t, s, "cache_peer svc-example parent 8080 0 no-query originserver name=origin_8080") + assert.Contains(t, s, "cache_peer 10.89.0.7 parent 8080 0 no-query originserver name=origin_8080") + // standby=2 pre-warms upstream connections so a cold first GET SSE stream is + // not reordered behind a later POST (fixes the sampling conformance flake). + assert.Contains(t, s, "connect-timeout=5 connect-fail-limit=5 standby=2") assert.Contains(t, s, "acl site_8080 dstdomain svc-example") assert.Contains(t, s, "http_access allow site_8080") assert.True(t, strings.HasSuffix(strings.TrimSpace(s), "http_access deny all")) @@ -223,7 +228,7 @@ func TestCreateTempIngressSquidConf_WithOverrideHosts(t *testing.T) { }, } - fp, err := createTempIngressSquidConf("svc-example", 8080, 18080, networkPermissions) + fp, err := createTempIngressSquidConf("svc-example", "10.89.0.7", 8080, 18080, networkPermissions) require.NoError(t, err) t.Cleanup(func() { _ = os.Remove(fp) }) @@ -234,7 +239,7 @@ func TestCreateTempIngressSquidConf_WithOverrideHosts(t *testing.T) { assert.Contains(t, s, "visible_hostname svc-example-ingress") assert.Contains(t, s, "\n# Reverse proxy setup for port 8080\n") assert.Contains(t, s, "http_port 0.0.0.0:18080 accel defaultsite=svc-example") - assert.Contains(t, s, "cache_peer svc-example parent 8080 0 no-query originserver name=origin_8080") + assert.Contains(t, s, "cache_peer 10.89.0.7 parent 8080 0 no-query originserver name=origin_8080") // Test that override mode is used - no default ACLs assert.NotContains(t, s, "acl site_8080 dstdomain svc-example") @@ -266,7 +271,7 @@ func TestCreateTempIngressSquidConf_EmptyInboundHosts(t *testing.T) { }, } - fp, err := createTempIngressSquidConf("svc-example", 8080, 18080, networkPermissions) + fp, err := createTempIngressSquidConf("svc-example", "10.89.0.7", 8080, 18080, networkPermissions) require.NoError(t, err) t.Cleanup(func() { _ = os.Remove(fp) }) @@ -496,7 +501,7 @@ func TestTempFilesWrittenToSystemTempDir(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = os.Remove(fp1) }) - fp2, err := createTempIngressSquidConf("s2", 8081, 18081, nil) + fp2, err := createTempIngressSquidConf("s2", "10.89.0.8", 8081, 18081, nil) require.NoError(t, err) t.Cleanup(func() { _ = os.Remove(fp2) }) diff --git a/pkg/container/images/keychain.go b/pkg/container/images/keychain.go index e0a4471168..c194c8c2db 100644 --- a/pkg/container/images/keychain.go +++ b/pkg/container/images/keychain.go @@ -4,73 +4,14 @@ package images import ( - "os" - "strings" - "github.com/google/go-containerregistry/pkg/authn" -) - -// envKeychain implements a keychain that reads credentials from environment variables -type envKeychain struct{} - -// Resolve implements the authn.Keychain interface -func (*envKeychain) Resolve(target authn.Resource) (authn.Authenticator, error) { - registry := target.RegistryStr() - - // Try registry-specific environment variables first - // Format: REGISTRY__USERNAME/PASSWORD, i.e., REGISTRY_DOCKER_IO_USERNAME - normalizedRegistry := strings.ToUpper(strings.ReplaceAll(registry, ".", "_")) - normalizedRegistry = strings.ReplaceAll(normalizedRegistry, "-", "_") - - username := os.Getenv("REGISTRY_" + normalizedRegistry + "_USERNAME") - password := os.Getenv("REGISTRY_" + normalizedRegistry + "_PASSWORD") - - // If registry-specific vars not found, try generic one REGISTRY_USERNAME/PASSWORD - if username == "" || password == "" { - username = os.Getenv("REGISTRY_USERNAME") - password = os.Getenv("REGISTRY_PASSWORD") - } - - if username != "" && password != "" { - return &authn.Basic{ - Username: username, - Password: password, - }, nil - } - return authn.Anonymous, nil -} - -// compositeKeychain combines multiple keychains and tries them in order -type compositeKeychain struct { - keychains []authn.Keychain -} - -// Resolve implements the authn.Keychain interface -func (c *compositeKeychain) Resolve(target authn.Resource) (authn.Authenticator, error) { - for _, keychain := range c.keychains { - auth, err := keychain.Resolve(target) - if err != nil { - continue - } - - // Check if we got actual credentials (not anonymous) - if auth != authn.Anonymous { - return auth, nil - } - } - - // If no keychain provided credentials, return anonymous - return authn.Anonymous, nil -} + coreimages "github.com/stacklok/toolhive-core/container/images" +) // NewCompositeKeychain creates a keychain that tries environment variables first, -// then falls back to the default keychain +// then falls back to the default keychain. +// Deprecated: use github.com/stacklok/toolhive-core/container/images.NewCompositeKeychain. func NewCompositeKeychain() authn.Keychain { - return &compositeKeychain{ - keychains: []authn.Keychain{ - &envKeychain{}, // Try environment variables first - authn.DefaultKeychain, // Then try default keychain (Docker config, etc.) - }, - } + return coreimages.NewCompositeKeychain() } diff --git a/pkg/container/images/registry.go b/pkg/container/images/registry.go index 1acf52219c..acb1f6ee4a 100644 --- a/pkg/container/images/registry.go +++ b/pkg/container/images/registry.go @@ -71,12 +71,21 @@ func (r *RegistryImageManager) ImageExists(_ context.Context, imageName string) return true, nil } -// PullImage pulls an image from a registry and saves it to the local daemon +// PullImage pulls an image from a registry and saves it to the local daemon. +// +// For digest-pinned references (tag@digest or digest-only), the Docker daemon's +// native ImagePull is used. The daemon registers the OCI manifest digest so that +// subsequent container creates using a tag@digest reference can resolve the image. +// go-containerregistry's daemon.Write stores images under a tag only, without +// registering the manifest digest; Docker's tag@digest cross-check at +// container-create time then fails with "No such image". +// +// For plain tag references, go-containerregistry is used for cross-platform +// support and custom keychain auth. func (r *RegistryImageManager) PullImage(ctx context.Context, imageName string) error { //nolint:gosec // G706: image name from user/config input slog.Info("pulling image", "image", imageName) - // Parse the image reference ref, err := name.ParseReference(imageName) if err != nil { return fmt.Errorf("failed to parse image reference %q: %w", imageName, err) @@ -91,29 +100,24 @@ func (r *RegistryImageManager) PullImage(ctx context.Context, imageName string) return fmt.Errorf("unsupported image reference type %T", ref) } - // Configure remote options remoteOpts := []remote.Option{ remote.WithAuthFromKeychain(r.keychain), remote.WithContext(ctx), } - if r.platform != nil { remoteOpts = append(remoteOpts, remote.WithPlatform(*r.platform)) } - // Pull the image from the registry img, err := remote.Image(ref, remoteOpts...) if err != nil { return fmt.Errorf("failed to pull image from registry: %w", err) } - // Save the image to the local daemon response, err := daemon.Write(tag, img, daemon.WithClient(r.dockerClient)) if err != nil { return fmt.Errorf("failed to write image to daemon: %w", err) } - // Display success message if _, err := fmt.Fprintf(os.Stdout, "Successfully pulled %s\n", imageName); err != nil { slog.Debug("failed to write success message", "error", err) } @@ -123,7 +127,11 @@ func (r *RegistryImageManager) PullImage(ctx context.Context, imageName string) return nil } +// pullDigestViaDockerDaemon pulls a digest-pinned image (tag@digest or +// digest-only form) using the Docker daemon client so the OCI manifest digest +// is stored and Docker can resolve the reference at container-create time. func (r *RegistryImageManager) pullDigestViaDockerDaemon(ctx context.Context, imageName string, ref name.Reference) error { + // Resolve auth from the keychain and encode for the daemon API. var registryAuth string if auth, err := r.keychain.Resolve(ref.Context().Registry); err == nil && auth != authn.Anonymous { if authCfg, err := auth.Authorization(); err == nil { @@ -133,9 +141,6 @@ func (r *RegistryImageManager) pullDigestViaDockerDaemon(ctx context.Context, im } } - if r.dockerClient == nil { - return fmt.Errorf("failed to pull image %q: docker daemon client unavailable", imageName) - } resp, err := r.dockerClient.ImagePull(ctx, imageName, mobyclient.ImagePullOptions{ RegistryAuth: registryAuth, }) diff --git a/pkg/container/templates/runtime_config.go b/pkg/container/templates/runtime_config.go index 4a71c800ed..be5a59c7c7 100644 --- a/pkg/container/templates/runtime_config.go +++ b/pkg/container/templates/runtime_config.go @@ -20,6 +20,15 @@ const maxPackageNameLength = 128 // dots, underscores, plus signs, or hyphens. var packageNamePattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._+\-]*$`) +// buildWithPattern matches a safe subset of PEP 508 requirement specifiers for +// BuildWith entries: a package name (optionally with extras) followed by version +// specifiers, e.g. "mcp<2", "mcp>=1.27,<2", "pkg[extra]==1.2.*", "foo~=1.4". +// The allowlist deliberately excludes quotes, backticks, dollar signs, +// semicolons, parentheses, and backslashes: entries are interpolated into a +// single-quoted shell word inside a Dockerfile RUN instruction, so anything +// that could close the quote or expand in shell context is rejected. +var buildWithPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._+\[\],<>=!~* -]*$`) + // envKeyPattern matches valid environment variable names for RuntimeEnv. // Must start with an uppercase letter, followed by uppercase letters, numbers, or underscores. var envKeyPattern = regexp.MustCompile(`^[A-Z][A-Z0-9_]*$`) @@ -54,6 +63,14 @@ type RuntimeConfig struct { // Examples for Debian: ["git", "build-essential"] AdditionalPackages []string `json:"additional_packages,omitempty" yaml:"additional_packages,omitempty"` + // BuildWith lists build-time dependency constraints, interpreted per + // package ecosystem. For uvx:// builds these are PEP 508 requirement + // specifiers passed to `uv tool install --with`, used to constrain + // transitive dependencies the package itself leaves unbounded + // (e.g. "mcp<2"). Ecosystems without constraint support (npx://, go://) + // reject a non-empty BuildWith at build time. + BuildWith []string `json:"build_with,omitempty" yaml:"build_with,omitempty"` + // RuntimeEnv contains environment variables to inject into the Dockerfile's // final runtime stage. Unlike BuildEnv (pkg/container/templates.TemplateData.BuildEnv), // which only affects the builder stage, these variables are baked into the @@ -100,6 +117,22 @@ func (rc *RuntimeConfig) Validate() error { } } + // Validate each BuildWith entry against a strict allowlist so specifiers + // cannot escape the single-quoted --with argument in the uvx Dockerfile. + for _, spec := range rc.BuildWith { + if len(spec) > maxPackageNameLength { + errs = append(errs, fmt.Errorf( + "build_with specifier %q exceeds maximum length of %d characters", + spec, maxPackageNameLength, + )) + } else if !buildWithPattern.MatchString(spec) { + errs = append(errs, fmt.Errorf( + "invalid build_with specifier %q: must match %s", + spec, buildWithPattern.String(), + )) + } + } + // Validate each RuntimeEnv entry to ensure keys and values are safe to // interpolate into a Dockerfile ENV instruction. for key, value := range rc.RuntimeEnv { diff --git a/pkg/container/templates/runtime_config_test.go b/pkg/container/templates/runtime_config_test.go index 7f29640870..b69a87a7be 100644 --- a/pkg/container/templates/runtime_config_test.go +++ b/pkg/container/templates/runtime_config_test.go @@ -457,3 +457,75 @@ func TestRuntimeConfigValidate_MultipleErrorsWithRuntimeEnv(t *testing.T) { assert.Contains(t, err.Error(), "is reserved and cannot be overridden") assert.Contains(t, err.Error(), "contains potentially dangerous characters") } + +func TestRuntimeConfigValidate_ValidBuildWith(t *testing.T) { + t.Parallel() + + valid := []string{ + "mcp<2", + "mcp>=1.27,<2", + "mcp==1.29.0", + "package[extra]==1.2.*", + "foo~=1.4", + "Foo_bar!=2.0", + "pkg >= 1.0, < 3", + } + for _, spec := range valid { + rc := &RuntimeConfig{BuildWith: []string{spec}} + assert.NoError(t, rc.Validate(), "specifier %q should be valid", spec) + } +} + +func TestRuntimeConfigValidate_InvalidBuildWith(t *testing.T) { + t.Parallel() + + invalid := []string{ + "", // empty + "mcp<2'; rm -rf /", // single quote escapes the --with argument + "mcp<2\" || true", // double quote + "mcp<2`id`", // backtick command substitution + "mcp<2$(id)", // dollar command substitution + "mcp<2;id", // command separator + "mcp<2|id", // pipe + "mcp<2&id", // background + "mcp<2\\", // backslash + "mcp<2\ninject", // newline breaks out of the RUN line + "-e evil", // leading dash could become a flag + strings.Repeat("a", 129), // over length bound + "pkg; python_version<'3.8'", // env markers need quotes, deliberately unsupported + } + for _, spec := range invalid { + rc := &RuntimeConfig{BuildWith: []string{spec}} + assert.Error(t, rc.Validate(), "specifier %q should be rejected", spec) + } +} + +func TestUVXTemplateRendersBuildWith(t *testing.T) { + t.Parallel() + + rc := GetDefaultRuntimeConfig(TransportTypeUVX) + rc.BuildWith = []string{"mcp<2", "other>=1,<4"} + data := TemplateData{ + MCPPackage: "arxiv-mcp-server", + RuntimeConfig: &rc, + } + dockerfile, err := GetDockerfileTemplate(TransportTypeUVX, data) + require.NoError(t, err) + assert.Contains(t, dockerfile, `uv tool install --with 'mcp<2' --with 'other>=1,<4' "$package_spec"`, + "each BuildWith specifier must be passed as a single-quoted --with argument") +} + +func TestUVXTemplateWithoutBuildWithIsUnchanged(t *testing.T) { + t.Parallel() + + rc := GetDefaultRuntimeConfig(TransportTypeUVX) + data := TemplateData{ + MCPPackage: "arxiv-mcp-server", + RuntimeConfig: &rc, + } + dockerfile, err := GetDockerfileTemplate(TransportTypeUVX, data) + require.NoError(t, err) + assert.Contains(t, dockerfile, `uv tool install "$package_spec"`, + "no --with arguments should appear when BuildWith is empty") + assert.NotContains(t, dockerfile, "--with") +} diff --git a/pkg/container/templates/uvx.tmpl b/pkg/container/templates/uvx.tmpl index 0a290ffd46..07b456b9be 100644 --- a/pkg/container/templates/uvx.tmpl +++ b/pkg/container/templates/uvx.tmpl @@ -77,7 +77,7 @@ ENV UV_TOOL_DIR=/opt/uv-tools \ RUN package="{{.MCPPackage}}"; \ # Replace @ with == for uv tool install (Python uses == for version pinning) package_spec=$(echo "$package" | sed 's/@/==/'); \ - uv tool install "$package_spec" && \ + uv tool install {{range .RuntimeConfig.BuildWith}}--with '{{.}}' {{end}}"$package_spec" && \ # List installed executables for debugging ls -la /opt/uv-tools/bin/ {{end}} diff --git a/pkg/git/client.go b/pkg/git/client.go index 2ff35b87ee..b3240ab07f 100644 --- a/pkg/git/client.go +++ b/pkg/git/client.go @@ -7,6 +7,7 @@ import ( "context" "errors" "fmt" + "io" "log/slog" "path/filepath" "strings" @@ -16,6 +17,7 @@ import ( "github.com/go-git/go-git/v5" "github.com/go-git/go-git/v5/plumbing" "github.com/go-git/go-git/v5/plumbing/cache" + "github.com/go-git/go-git/v5/plumbing/object" "github.com/go-git/go-git/v5/plumbing/transport" "github.com/go-git/go-git/v5/storage/filesystem" ) @@ -37,13 +39,28 @@ type Client interface { // GetFileContent retrieves the content of a file from the repository GetFileContent(repoInfo *RepositoryInfo, path string) ([]byte, error) - // HeadCommitHash returns the commit hash of the HEAD reference. - HeadCommitHash(repoInfo *RepositoryInfo) (string, error) + // HeadCommit returns the hash and signature of the HEAD commit. + HeadCommit(repoInfo *RepositoryInfo) (HeadCommit, error) // Cleanup removes local repository directory Cleanup(ctx context.Context, repoInfo *RepositoryInfo) error } +// HeadCommit describes the HEAD commit of a cloned repository. +type HeadCommit struct { + // Hash is the commit hash. + Hash string + // Signature is the armored signature attached to the commit, empty + // when the commit is unsigned. The signature is UNVERIFIED — it is + // whatever bytes the commit carries; callers must cryptographically + // verify it before treating it as provenance. + Signature string + // Payload is the encoded commit object without its signature — the + // exact bytes the signature signs. Verifiers check Signature over + // Payload. + Payload []byte +} + // DefaultGitClient implements Client using go-git type DefaultGitClient struct { // auth is the optional authentication method for cloning @@ -251,16 +268,43 @@ func (*DefaultGitClient) updateRepositoryInfo(repoInfo *RepositoryInfo) error { return nil } -// HeadCommitHash returns the commit hash of the HEAD reference. -func (*DefaultGitClient) HeadCommitHash(repoInfo *RepositoryInfo) (string, error) { +// HeadCommit returns the hash and (unverified) signature of the HEAD +// commit in a single lookup, so both describe the same commit. +func (*DefaultGitClient) HeadCommit(repoInfo *RepositoryInfo) (HeadCommit, error) { if repoInfo == nil || repoInfo.Repository == nil { - return "", ErrNilRepository + return HeadCommit{}, ErrNilRepository } ref, err := repoInfo.Repository.Head() if err != nil { - return "", fmt.Errorf("failed to get HEAD reference: %w", err) + return HeadCommit{}, fmt.Errorf("failed to get HEAD reference: %w", err) + } + commit, err := repoInfo.Repository.CommitObject(ref.Hash()) + if err != nil { + return HeadCommit{}, fmt.Errorf("failed to read HEAD commit: %w", err) + } + payload, err := commitPayload(commit) + if err != nil { + return HeadCommit{}, fmt.Errorf("failed to encode HEAD commit payload: %w", err) } + return HeadCommit{ + Hash: ref.Hash().String(), + Signature: commit.PGPSignature, + Payload: payload, + }, nil +} - return ref.Hash().String(), nil +// commitPayload returns the encoded commit object without its signature — +// the bytes a commit signature is computed over. +func commitPayload(commit *object.Commit) ([]byte, error) { + obj := &plumbing.MemoryObject{} + if err := commit.EncodeWithoutSignature(obj); err != nil { + return nil, err + } + r, err := obj.Reader() + if err != nil { + return nil, err + } + defer r.Close() //nolint:errcheck // in-memory reader + return io.ReadAll(r) } diff --git a/pkg/git/client_test.go b/pkg/git/client_test.go index 4898caad20..ac8fb81742 100644 --- a/pkg/git/client_test.go +++ b/pkg/git/client_test.go @@ -104,7 +104,7 @@ func TestDefaultGitClient_GetFileContent_PathTraversal(t *testing.T) { } } -func TestDefaultGitClient_HeadCommitHash_NilInputs(t *testing.T) { +func TestDefaultGitClient_HeadCommit_NilInputs(t *testing.T) { t.Parallel() client := NewDefaultGitClient() @@ -119,14 +119,14 @@ func TestDefaultGitClient_HeadCommitHash_NilInputs(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - hash, err := client.HeadCommitHash(tt.repoInfo) + head, err := client.HeadCommit(tt.repoInfo) require.ErrorIs(t, err, ErrNilRepository) - assert.Empty(t, hash) + assert.Empty(t, head) }) } } -func TestDefaultGitClient_HeadCommitHash_Valid(t *testing.T) { +func TestDefaultGitClient_HeadCommit_UnsignedCommit(t *testing.T) { t.Parallel() client := NewDefaultGitClient() @@ -135,10 +135,12 @@ func TestDefaultGitClient_HeadCommitHash_Valid(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = client.Cleanup(t.Context(), repoInfo) }) - hash, err := client.HeadCommitHash(repoInfo) + head, err := client.HeadCommit(repoInfo) require.NoError(t, err) - assert.Len(t, hash, 40, "commit hash should be 40 hex chars") - assert.True(t, isAllHex(hash), "commit hash should be all hex") + assert.Len(t, head.Hash, 40, "commit hash should be 40 hex chars") + assert.True(t, isAllHex(head.Hash), "commit hash should be all hex") + assert.Empty(t, head.Signature, "an unsigned commit must yield an empty signature, not an error") + assert.Contains(t, string(head.Payload), "tree ", "the payload is present even for unsigned commits") } // isAllHex checks if s is a non-empty lowercase hex string. diff --git a/pkg/git/integration_test.go b/pkg/git/integration_test.go index 907f589dd8..7ab998f0a2 100644 --- a/pkg/git/integration_test.go +++ b/pkg/git/integration_test.go @@ -43,6 +43,60 @@ func initTestRepo(t *testing.T, files map[string]string) string { return dir } +// signTestRepoHead rewrites the HEAD commit of the repo at dir with the given +// armored signature attached, mimicking what gitsign/PGP tooling produces. +func signTestRepoHead(t *testing.T, dir, signature string) { + t.Helper() + + repo, err := gogit.PlainOpen(dir) + require.NoError(t, err) + head, err := repo.Head() + require.NoError(t, err) + commit, err := repo.CommitObject(head.Hash()) + require.NoError(t, err) + + commit.PGPSignature = signature + obj := repo.Storer.NewEncodedObject() + require.NoError(t, commit.Encode(obj)) + hash, err := repo.Storer.SetEncodedObject(obj) + require.NoError(t, err) + require.NoError(t, repo.Storer.SetReference(plumbing.NewHashReference(head.Name(), hash))) +} + +// repoInfoHeadHash returns the HEAD hash of a cloned repo directly via +// go-git, for cross-checking HeadCommit. +func repoInfoHeadHash(t *testing.T, repoInfo *RepositoryInfo) string { + t.Helper() + ref, err := repoInfo.Repository.Head() + require.NoError(t, err) + return ref.Hash().String() +} + +// TestDefaultGitClient_HeadCommit_SignedCommit verifies the armored +// signature stored on a commit survives clone and is returned verbatim. +func TestDefaultGitClient_HeadCommit_SignedCommit(t *testing.T) { + t.Parallel() + + const armoredSig = "-----BEGIN SIGNED MESSAGE-----\nMIIC(test-signature)\n-----END SIGNED MESSAGE-----\n" + repoDir := initTestRepo(t, map[string]string{"test.txt": "content"}) + signTestRepoHead(t, repoDir, armoredSig) + + client := NewDefaultGitClient() + repoInfo, err := client.Clone(t.Context(), &CloneConfig{URL: repoDir}) + require.NoError(t, err) + t.Cleanup(func() { _ = client.Cleanup(t.Context(), repoInfo) }) + + head, err := client.HeadCommit(repoInfo) + require.NoError(t, err) + assert.Equal(t, armoredSig, head.Signature) + assert.Equal(t, repoInfoHeadHash(t, repoInfo), head.Hash, + "hash and signature must describe the same commit") + assert.NotContains(t, string(head.Payload), "gpgsig", + "the payload must be the commit without its signature header") + assert.Contains(t, string(head.Payload), "tree ", + "the payload must be the encoded commit object") +} + // TestDefaultGitClient_FullWorkflow exercises the complete clone → read → cleanup lifecycle // against a local git repository to verify end-to-end correctness. func TestDefaultGitClient_FullWorkflow(t *testing.T) { diff --git a/pkg/llm/proxy/proxy.go b/pkg/llm/proxy/proxy.go index 5f96bf7960..086f43b717 100644 --- a/pkg/llm/proxy/proxy.go +++ b/pkg/llm/proxy/proxy.go @@ -21,12 +21,21 @@ import ( "github.com/stacklok/toolhive/pkg/networking" ) -// tokenFetchTimeout bounds the per-request token fetch. It is deliberately -// generous: the first request after a lazy setup ("thv llm setup --lazy") drives -// the interactive OIDC browser login through the token source, and a human needs -// time to complete it. Once a token is cached, subsequent requests are served -// from cache and return well within this bound, so the only request that ever -// approaches it is the initial login. +// tokenFetchTimeout bounds the per-request token fetch. The bound is achievable +// because the fetch is rooted in the proxy-lifetime context passed to Start, not +// in the inbound request's context: the first request after a lazy setup +// ("thv llm setup --lazy") drives the interactive OIDC browser login through the +// token source, and a human needs time to complete it. An impatient client +// disconnecting no longer cancels the login, but Ctrl+C (cancelling Start's ctx) +// still does. Once a token is cached, subsequent requests are served from cache +// and return well within this bound, so the only request that ever approaches it +// is the initial login. +// +// The token source serializes Token() internally, so requests that arrive while +// that initial login is in flight queue behind it and are then served from the +// resulting cached token — one browser flow, however many concurrent requests. +// Their own clients may time out and disconnect first; that is fine, since the +// login they were waiting on completes regardless and their retry hits the cache. const tokenFetchTimeout = 3 * time.Minute // TokenSource obtains fresh OIDC access tokens for the LLM gateway. @@ -155,8 +164,10 @@ func (p *Proxy) Addr() string { } // handler returns an http.Handler that injects a fresh OIDC token and proxies -// the request to the upstream gateway. -func (p *Proxy) handler() http.Handler { +// the request to the upstream gateway. The token fetch is bounded by baseCtx +// (the proxy lifetime context passed to Start), not the inbound request's +// context, so a client disconnect cannot abort an in-flight interactive login. +func (p *Proxy) handler(baseCtx context.Context) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Guard against DNS-rebinding attacks: reject requests whose Host header // does not resolve to a loopback address. A loopback-only bind prevents @@ -167,7 +178,7 @@ func (p *Proxy) handler() http.Handler { return } - tokenCtx, cancel := context.WithTimeout(r.Context(), tokenFetchTimeout) + tokenCtx, cancel := context.WithTimeout(baseCtx, tokenFetchTimeout) defer cancel() token, err := p.tokenSource.Token(tokenCtx) if err != nil { @@ -206,7 +217,7 @@ func (p *Proxy) handler() http.Handler { // cancelled, then performs a graceful shutdown with a 5-second timeout. func (p *Proxy) Start(ctx context.Context) error { p.server = &http.Server{ - Handler: p.handler(), + Handler: p.handler(ctx), ReadHeaderTimeout: 30 * time.Second, } diff --git a/pkg/llm/proxy/proxy_test.go b/pkg/llm/proxy/proxy_test.go index 11516f6638..c805c004e8 100644 --- a/pkg/llm/proxy/proxy_test.go +++ b/pkg/llm/proxy/proxy_test.go @@ -26,9 +26,18 @@ import ( type stubTokenSource struct { token string err error + // block, if non-nil, is invoked from Token. Tests use it to observe the + // ctx the proxy passed in and to control timing at the observable boundary + // (when the fetch resolves). It receives the token-fetch ctx and the token + // the stub will return once it unblocks. This is test-only instrumentation; + // production structs are unchanged. + block func(ctx context.Context, token string) } -func (s *stubTokenSource) Token(_ context.Context) (string, error) { +func (s *stubTokenSource) Token(ctx context.Context) (string, error) { + if s.block != nil { + s.block(ctx, s.token) + } return s.token, s.err } @@ -119,7 +128,7 @@ func TestHandler_InjectsToken(t *testing.T) { req := loopbackRequest("/v1/models") req.Header.Set("Authorization", "Bearer old-token") w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) mu.Lock() defer mu.Unlock() @@ -145,7 +154,7 @@ func TestHandler_StripsIncomingAuthorization(t *testing.T) { req := loopbackRequest("/v1/models") req.Header.Set("Authorization", "Bearer user-supplied-token") w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) mu.Lock() defer mu.Unlock() @@ -168,7 +177,7 @@ func TestHandler_RejectsDNSRebindingHost(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) req.Host = host w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) assert.Equal(t, http.StatusForbidden, w.Code, "host %q should be rejected", host) } @@ -177,7 +186,7 @@ func TestHandler_RejectsDNSRebindingHost(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) req.Host = host w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) assert.NotEqual(t, http.StatusForbidden, w.Code, "host %q should be allowed", host) } } @@ -194,7 +203,7 @@ func TestHandler_Returns502OnTokenError(t *testing.T) { req := loopbackRequest("/v1/chat/completions") w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) assert.Equal(t, http.StatusBadGateway, w.Code) assert.Equal(t, "application/json", w.Header().Get("Content-Type")) @@ -215,7 +224,7 @@ func TestHandler_Returns401WithActionableMessageOnErrTokenRequired(t *testing.T) req := loopbackRequest("/v1/chat/completions") w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) assert.Equal(t, http.StatusUnauthorized, w.Code) assert.Equal(t, "application/json", w.Header().Get("Content-Type")) @@ -364,7 +373,7 @@ func TestWithTLSSkipVerify(t *testing.T) { req := loopbackRequest("/v1/models") w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) // Certificate verification failure surfaces as 502 Bad Gateway. assert.Equal(t, http.StatusBadGateway, w.Code) @@ -382,7 +391,7 @@ func TestWithTLSSkipVerify(t *testing.T) { req := loopbackRequest("/v1/models") w := httptest.NewRecorder() - p.handler().ServeHTTP(w, req) + p.handler(context.Background()).ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) }) @@ -409,3 +418,145 @@ func TestProxy_PassesThroughErrorResponses(t *testing.T) { }) } } + +// newBlockingTokenProxy builds a Proxy whose token source blocks inside Token +// until the returned release channel is closed (or the fetch ctx is cancelled, +// whichever is first). It also hands the fetch ctx to sawCtx, exactly once, +// so the caller can inspect the context the proxy rooted the fetch in. +// +// The gateway returns 200 so that, once the token is delivered, the handler +// completes normally; tests that only care about the fetch phase never reach it. +func newBlockingTokenProxy(t *testing.T) (p *Proxy, release chan struct{}, sawCtx chan context.Context) { + t.Helper() + + gateway := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(gateway.Close) + + cfg := &llm.Config{ + GatewayURL: gateway.URL, + Proxy: llm.ProxyConfig{ListenPort: freePort(t)}, + } + release = make(chan struct{}) + sawCtx = make(chan context.Context, 1) + src := &stubTokenSource{ + token: "fresh-token", + block: func(ctx context.Context, _ string) { + select { + case sawCtx <- ctx: + default: + } + // Block until the test releases us or the proxy cancels the fetch. + select { + case <-release: + case <-ctx.Done(): + } + }, + } + p, err := New(cfg, src) + require.NoError(t, err) + p.transport = gateway.Client().Transport + t.Cleanup(func() { _ = p.listener.Close() }) + return p, release, sawCtx +} + +// waitForFetchCtx waits for the stubbed Token to report the fetch ctx, failing +// fast on timeout so a miswire never hangs the suite. +func waitForFetchCtx(t *testing.T, sawCtx <-chan context.Context) context.Context { + t.Helper() + select { + case ctx := <-sawCtx: + return ctx + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for token fetch to start") + return nil + } +} + +// TestHandler_TokenFetchSurvivesClientDisconnect pins the Step 1 contract that a +// client disconnecting mid-request does NOT cancel the token fetch: the fetch is +// rooted in the proxy-lifetime baseCtx, not the inbound request's context. +func TestHandler_TokenFetchSurvivesClientDisconnect(t *testing.T) { + t.Parallel() + p, release, sawCtx := newBlockingTokenProxy(t) + + baseCtx, baseCancel := context.WithCancel(context.Background()) + t.Cleanup(baseCancel) + + // Inbound request carries its own cancellable context, as a real + // disconnected client would. Host is loopback to pass the DNS-rebinding guard. + reqCtx, reqCancel := context.WithCancel(context.Background()) + t.Cleanup(reqCancel) + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("")) + req = req.WithContext(reqCtx) + req.Host = "127.0.0.1" + + handlerDone := make(chan struct{}) + go func() { + w := httptest.NewRecorder() + p.handler(baseCtx).ServeHTTP(w, req) + close(handlerDone) + }() + + fetchCtx := waitForFetchCtx(t, sawCtx) + + // Disconnect the client while the fetch is still in flight. + reqCancel() + + // Give the cancellation a moment to propagate, then assert the fetch ctx is + // untouched: under the Step 1 contract the fetch is rooted in baseCtx, which + // is still alive, so the fetch must keep running. + select { + case <-fetchCtx.Done(): + t.Fatal("fetch ctx was cancelled by request ctx — fetch must be rooted in baseCtx, not r.Context()") + case <-time.After(100 * time.Millisecond): + } + + // Releasing the stub delivers the token and the handler finishes normally. + close(release) + select { + case <-handlerDone: + case <-time.After(2 * time.Second): + t.Fatal("handler did not return after token was delivered") + } +} + +// TestHandler_TokenFetchCancelsWithStartContext pins the other half of the +// Step 1 contract: cancelling Start's ctx (baseCtx) — e.g. via Ctrl+C — DOES +// cancel the in-flight token fetch. +func TestHandler_TokenFetchCancelsWithStartContext(t *testing.T) { + t.Parallel() + p, _, sawCtx := newBlockingTokenProxy(t) + + baseCtx, baseCancel := context.WithCancel(context.Background()) + t.Cleanup(baseCancel) + + req := loopbackRequest("/v1/models") + handlerDone := make(chan struct{}) + go func() { + w := httptest.NewRecorder() + p.handler(baseCtx).ServeHTTP(w, req) + close(handlerDone) + }() + + fetchCtx := waitForFetchCtx(t, sawCtx) + + // Cancel Start's lifetime context, exactly as Ctrl+C would. + baseCancel() + + // The fetch ctx is derived from baseCtx, so it must be cancelled now. + select { + case <-fetchCtx.Done(): + case <-time.After(2 * time.Second): + t.Fatal("fetch ctx was not cancelled when baseCtx was cancelled") + } + + // And the handler must unwind: the stub also returns on ctx.Done(), so the + // token is delivered and the handler completes via the 200 proxy path. + select { + case <-handlerDone: + case <-time.After(2 * time.Second): + t.Fatal("handler did not return after baseCtx cancellation") + } +} diff --git a/pkg/llm/tokensource.go b/pkg/llm/tokensource.go index e9e3780eda..a95a2e6e57 100644 --- a/pkg/llm/tokensource.go +++ b/pkg/llm/tokensource.go @@ -39,7 +39,10 @@ type TokenSource = tokensource.OAuthTokenSource // secretsProvider may be nil if the secrets store is unavailable. // tokenRefUpdater is called after login/refresh to persist the token reference // into config — pass nil to skip config persistence (useful in tests). -// Set interactive to false for non-interactive callers such as thv llm token. +// interactive controls whether a genuine cache miss may launch the browser +// OIDC flow. thv llm token passes true so a prior "thv llm setup --lazy" signs +// the user in transparently on first use; cached tokens are served without a +// browser prompt either way. // When skipBrowser is true, an interactive login prints the authorization URL // instead of opening a browser (headless/SSH/CI use); it has no effect unless // interactive is also true. diff --git a/pkg/mcp/batch.go b/pkg/mcp/batch.go new file mode 100644 index 0000000000..ce9866d1e2 --- /dev/null +++ b/pkg/mcp/batch.go @@ -0,0 +1,71 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "bytes" + "net/http" +) + +// batchUnsupported is the shared error rendered by both batch-rejection helpers. +var batchUnsupported = &BatchUnsupportedError{} + +// BatchUnsupportedError indicates a client sent a JSON-RPC batch (a top-level +// JSON array). Batching was removed from MCP in the 2025-06-18 revision, and +// ToolHive serves only 2025-11-25 and 2026-07-28, so a batch can only originate +// from an unsupported pre-2025-06-18 client. +// +// Batches are rejected outright rather than executed: the parser and the +// downstream authz, audit, and tool-filter middleware each inspect a single +// parsed request per call, so a batch that reached the backend would smuggle +// every nested call past those controls (see #5745). +type BatchUnsupportedError struct{} + +func (*BatchUnsupportedError) Error() string { + return "JSON-RPC batch requests are not supported (batching was removed in MCP revision 2025-06-18)" +} + +// Code implements CodedError. +func (*BatchUnsupportedError) Code() int64 { return CodeInvalidRequest } + +// Data implements CodedError. +func (*BatchUnsupportedError) Data() map[string]any { return nil } + +// IsBatchRequest reports whether body is a JSON-RPC batch, i.e. a top-level JSON +// array. Only the first non-whitespace byte is inspected, so a syntactically +// malformed array (e.g. "[...") still classifies as a batch and is rejected +// rather than silently mis-parsed as a single message. +// +// It trims the same Unicode whitespace as encoding/json (via bytes.TrimSpace) so +// callers that gate execution on this function and callers that later +// json-decode the body agree on what counts as a leading '[' — a narrower trim +// would let a batch prefixed with, say, a vertical-tab slip past detection yet +// still decode as an array. Every path that could forward a batch to a backend +// must gate on this one function so the two never drift (see #5745). +func IsBatchRequest(body []byte) bool { + trimmed := bytes.TrimSpace(body) + return len(trimmed) > 0 && trimmed[0] == '[' +} + +// WriteBatchUnsupportedError writes an HTTP 400 response carrying a JSON-RPC +// "Invalid Request" error for a rejected batch. A batch has no single request +// id to echo, so the "id" key is omitted entirely (schema/2025-11-25 types the +// error response's id as optional, not nullable; see session.HasJSONRPCID). +// Use this where an http.ResponseWriter is available (the streamable proxy, +// ParsingMiddleware). +func WriteBatchUnsupportedError(w http.ResponseWriter) { + body := classificationErrorBody(nil, batchUnsupported) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client + _, _ = w.Write(body) +} + +// BatchUnsupportedResponse builds an *http.Response carrying the batch-rejection +// JSON-RPC error, for proxies that intercept at the RoundTripper layer (the +// transparent proxy) where no http.ResponseWriter is available. Mirrors +// ClassificationErrorResponse; the JSON-RPC id is omitted, not null. +func BatchUnsupportedResponse(req *http.Request) *http.Response { + return jsonRPCErrorResponse(req, nil, batchUnsupported) +} diff --git a/pkg/mcp/batch_test.go b/pkg/mcp/batch_test.go new file mode 100644 index 0000000000..5ea28ff0ea --- /dev/null +++ b/pkg/mcp/batch_test.go @@ -0,0 +1,83 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIsBatchRequest(t *testing.T) { + t.Parallel() + tests := []struct { + name string + body string + want bool + }{ + {name: "single request object", body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`, want: false}, + {name: "batch array", body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"}]`, want: true}, + {name: "empty array", body: `[]`, want: true}, + {name: "array with leading whitespace", body: " \n\t [1]", want: true}, + // Vertical tab / form feed are Unicode whitespace that json ignores; + // bytes.TrimSpace must strip them so detection matches decoding (#5745). + {name: "array with leading vertical tab", body: "\v\f[1]", want: true}, + {name: "object with leading whitespace", body: " \n\t {}", want: false}, + {name: "empty body", body: ``, want: false}, + {name: "whitespace only", body: " \n\t", want: false}, + {name: "malformed unterminated array", body: `[{"a":1}`, want: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, IsBatchRequest([]byte(tt.body))) + }) + } +} + +func TestBatchUnsupportedError(t *testing.T) { + t.Parallel() + err := &BatchUnsupportedError{} + // Must satisfy CodedError so the shared JSON-RPC error writer picks up its code. + var coded CodedError = err + assert.Equal(t, CodeInvalidRequest, coded.Code()) + assert.Nil(t, coded.Data()) + assert.NotEmpty(t, coded.Error()) +} + +func TestWriteBatchUnsupportedError(t *testing.T) { + t.Parallel() + w := httptest.NewRecorder() + + WriteBatchUnsupportedError(w) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + + var resp struct { + JSONRPC string `json:"jsonrpc"` + Error struct { + Code int64 `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.Equal(t, "2.0", resp.JSONRPC) + assert.Equal(t, CodeInvalidRequest, resp.Error.Code) + assert.NotEmpty(t, resp.Error.Message) + + // A batch has no single request id to echo. A tagged-struct decode + // (above) can't tell an absent "id" key from a present null -- both + // decode to the zero value -- so the omission is checked via a map + // decode instead. + var asMap map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &asMap)) + _, hasID := asMap["id"] + assert.False(t, hasID, `"id" key must be omitted, not present as null`) +} diff --git a/pkg/mcp/classification_response.go b/pkg/mcp/classification_response.go new file mode 100644 index 0000000000..724ae942ef --- /dev/null +++ b/pkg/mcp/classification_response.go @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + + "github.com/stacklok/toolhive/pkg/transport/session" +) + +// WriteClassificationError writes an HTTP 400 response with a JSON-RPC error +// body for an mcp.ClassifyRevision failure. Use this with http.ResponseWriter +// in the streamable HTTP proxy. +func WriteClassificationError(w http.ResponseWriter, requestID any, err error) { + body := classificationErrorBody(requestID, err) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client + _, _ = w.Write(body) +} + +// ClassificationErrorResponse constructs an *http.Response with HTTP 400 and +// a JSON-RPC error body for an mcp.ClassifyRevision failure. Use this in +// httputil.ReverseProxy.ModifyResponse/RoundTrip (transparent proxy) where no +// http.ResponseWriter is available. +// +// req is attached as resp.Request, matching session.NotFoundResponse: a +// RoundTripper-produced *http.Response is generally expected to carry the +// request it answers, and httputil.ReverseProxy relies on that field. +func ClassificationErrorResponse(req *http.Request, requestID any, err error) *http.Response { + return jsonRPCErrorResponse(req, requestID, err) +} + +// jsonRPCErrorResponse builds an HTTP 400 *http.Response carrying the JSON-RPC +// error rendered from a CodedError. It backs both ClassificationErrorResponse +// and BatchUnsupportedResponse so the RoundTripper-layer error shape lives in +// one place. +func jsonRPCErrorResponse(req *http.Request, requestID any, err error) *http.Response { + body := classificationErrorBody(requestID, err) + hdr := make(http.Header) + hdr.Set("Content-Type", "application/json") + return &http.Response{ + StatusCode: http.StatusBadRequest, + Status: fmt.Sprintf("%d %s", http.StatusBadRequest, http.StatusText(http.StatusBadRequest)), + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: hdr, + ContentLength: int64(len(body)), + Body: io.NopCloser(bytes.NewReader(body)), + Request: req, + } +} + +// classificationErrorBody renders an mcp.ClassifyRevision error as a JSON-RPC +// error body, modeled on session.NotFoundBody: the body is marshaled first +// (with a hand-crafted fallback on marshal failure) so callers only write +// headers/status once a valid body is ready. +// +// It uses the error's Code(), Error() message, and Data() (when non-empty) if +// the error implements CodedError, falling back to the standard JSON-RPC +// Invalid Params code otherwise -- a fallback that is currently unreachable, +// since every error ClassifyRevision returns implements CodedError. +// +// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error +// response as `id?: RequestId` where RequestId = string | number, so an +// absent id is encoded by omitting the "id" key, never as null. See +// session.HasJSONRPCID. +func classificationErrorBody(requestID any, err error) []byte { + code := CodeInvalidParams + var coded CodedError + var data map[string]any + if errors.As(err, &coded) { + code = coded.Code() + data = coded.Data() + } + + errBody := map[string]any{ + "code": code, + "message": err.Error(), + } + if len(data) > 0 { + errBody["data"] = data + } + resp := map[string]any{ + "jsonrpc": "2.0", + "error": errBody, + } + if session.HasJSONRPCID(requestID) { + resp["id"] = requestID + } + + body, marshalErr := json.Marshal(resp) + if marshalErr != nil { + // This should never happen with simple map types, but return a + // hand-crafted fallback to guarantee a valid JSON-RPC error. + return []byte(`{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid params"}}`) + } + return body +} diff --git a/pkg/mcp/classification_response_test.go b/pkg/mcp/classification_response_test.go new file mode 100644 index 0000000000..efc7948681 --- /dev/null +++ b/pkg/mcp/classification_response_test.go @@ -0,0 +1,217 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestClassificationError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + requestID any + err error + wantCode int64 + wantData bool + wantIDKey bool + wantIDEcho string // substring the body must contain, verbatim, when wantIDKey + }{ + { + name: "header mismatch", + requestID: "req-1", + err: &HeaderMismatchError{Header: "2026-07-28", Body: "2025-11-25"}, + wantCode: CodeHeaderMismatch, + wantData: true, + wantIDKey: true, + wantIDEcho: `"id":"req-1"`, + }, + { + name: "unsupported version", + requestID: float64(42), + err: &UnsupportedVersionError{Requested: "1999-01-01", Supported: []string{MCPVersionModern}}, + wantCode: CodeUnsupportedProtocolVersion, + wantData: true, + wantIDKey: true, + wantIDEcho: `"id":42`, + }, + { + name: "missing client capability", + requestID: "req-3", + err: &MissingClientCapabilityError{RequiredCapabilities: map[string]any{}}, + wantCode: CodeMissingClientCapability, + wantData: true, + wantIDKey: true, + wantIDEcho: `"id":"req-3"`, + }, + { + name: "missing modern metadata", + requestID: "req-4", + err: &MissingModernMetadataError{}, + wantCode: CodeInvalidParams, + wantData: false, // Data() returns an empty (non-nil) map, so len(data) == 0 + wantIDKey: true, + wantIDEcho: `"id":"req-4"`, + }, + { + // A nil requestID means the caller couldn't determine an id at + // all (e.g. the request failed to parse). MCP narrows base + // JSON-RPC to omit the "id" key here rather than emit + // "id":null: see session.HasJSONRPCID. + name: "nil request id omits the id key", + requestID: nil, + err: &HeaderMismatchError{Header: "a", Body: "b"}, + wantCode: CodeHeaderMismatch, + wantData: true, + wantIDKey: false, + }, + { + // The transparent proxy threads the incoming id through as raw + // bytes; literal null must be treated the same as a missing id. + name: "raw null request id omits the id key", + requestID: json.RawMessage(`null`), + err: &HeaderMismatchError{Header: "a", Body: "b"}, + wantCode: CodeHeaderMismatch, + wantData: true, + wantIDKey: false, + }, + { + name: "nil RawMessage omits the id key", + requestID: json.RawMessage(nil), + err: &HeaderMismatchError{Header: "a", Body: "b"}, + wantCode: CodeHeaderMismatch, + wantData: true, + wantIDKey: false, + }, + { + name: "zero request id is still present", + requestID: float64(0), + err: &HeaderMismatchError{Header: "a", Body: "b"}, + wantCode: CodeHeaderMismatch, + wantData: true, + wantIDKey: true, + wantIDEcho: `"id":0`, + }, + { + name: "plain non-coded error falls back to invalid params", + requestID: "req-6", + err: errors.New("boom"), + wantCode: CodeInvalidParams, + wantData: false, + wantIDKey: true, + wantIDEcho: `"id":"req-6"`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + WriteClassificationError(rec, tt.requestID, tt.err) + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + resp := ClassificationErrorResponse(req, tt.requestID, tt.err) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("WriteClassificationError: status = %d, want %d", rec.Code, http.StatusBadRequest) + } + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("ClassificationErrorResponse: status = %d, want %d", resp.StatusCode, http.StatusBadRequest) + } + if ct := rec.Header().Get("Content-Type"); ct != "application/json" { + t.Fatalf("WriteClassificationError: Content-Type = %q, want application/json", ct) + } + if ct := resp.Header.Get("Content-Type"); ct != "application/json" { + t.Fatalf("ClassificationErrorResponse: Content-Type = %q, want application/json", ct) + } + if resp.Request != req { + t.Fatalf("ClassificationErrorResponse: Request not attached") + } + if resp.Body == nil { + t.Fatalf("ClassificationErrorResponse: Body is nil") + } + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("reading response body: %v", err) + } + if resp.ContentLength != int64(len(respBody)) { + t.Fatalf("ContentLength = %d, want %d", resp.ContentLength, len(respBody)) + } + + wireBody := rec.Body.Bytes() + if string(wireBody) != string(respBody) { + t.Fatalf("WriteClassificationError and ClassificationErrorResponse bodies differ:\n%s\nvs\n%s", wireBody, respBody) + } + + var decoded struct { + JSONRPC string `json:"jsonrpc"` + Error struct { + Code int64 `json:"code"` + Message string `json:"message"` + Data map[string]any `json:"data"` + } `json:"error"` + } + if err := json.Unmarshal(wireBody, &decoded); err != nil { + t.Fatalf("unmarshaling response body: %v", err) + } + + if decoded.JSONRPC != "2.0" { + t.Errorf("jsonrpc = %q, want \"2.0\"", decoded.JSONRPC) + } + if decoded.Error.Code != tt.wantCode { + t.Errorf("code = %d, want %d", decoded.Error.Code, tt.wantCode) + } + if decoded.Error.Message != tt.err.Error() { + t.Errorf("message = %q, want %q", decoded.Error.Message, tt.err.Error()) + } + if tt.wantData && len(decoded.Error.Data) == 0 { + t.Errorf("expected non-empty data, got %v", decoded.Error.Data) + } + if !tt.wantData && len(decoded.Error.Data) != 0 { + t.Errorf("expected no data, got %v", decoded.Error.Data) + } + + // Decoding into map[string]any (rather than a tagged struct) is + // the only way to tell an absent "id" key apart from a present + // null: a struct field would read as the zero value either way. + var asMap map[string]any + require.NoError(t, json.Unmarshal(wireBody, &asMap)) + _, hasID := asMap["id"] + assert.Equal(t, tt.wantIDKey, hasID, `"id" key presence`) + if tt.wantIDKey { + require.NotEmpty(t, tt.wantIDEcho, "table entry must set wantIDEcho when wantIDKey is true") + // A map decode would coerce a numeric id to float64, losing + // the verbatim-echo guarantee #5945 pinned -- so check the + // raw wire bytes instead. + assert.Contains(t, string(wireBody), tt.wantIDEcho) + } + }) + } +} + +// TestClassificationErrorBodyMarshalFallback verifies the defensive fallback: +// if the response fails to marshal (unreachable in production, where requestID +// is always a json.RawMessage), classificationErrorBody still returns a valid +// JSON-RPC error rather than an empty body. A channel is not JSON-marshalable, +// so it forces json.Marshal to fail. +func TestClassificationErrorBodyMarshalFallback(t *testing.T) { + t.Parallel() + + const want = `{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid params"}}` + got := classificationErrorBody(make(chan int), errors.New("boom")) + if string(got) != want { + t.Fatalf("fallback body = %s, want %s", got, want) + } +} diff --git a/pkg/mcp/client/call.go b/pkg/mcp/client/call.go index 84e050a696..d51c191b73 100644 --- a/pkg/mcp/client/call.go +++ b/pkg/mcp/client/call.go @@ -8,7 +8,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // CallTool connects to the MCP server at serverURL, performs the initialize diff --git a/pkg/mcp/client/client.go b/pkg/mcp/client/client.go index 34db99dd3e..0ae9c77baf 100644 --- a/pkg/mcp/client/client.go +++ b/pkg/mcp/client/client.go @@ -16,9 +16,8 @@ import ( "net/url" "strings" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" - + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/transport/ssecommon" "github.com/stacklok/toolhive/pkg/transport/streamable" "github.com/stacklok/toolhive/pkg/transport/types" diff --git a/pkg/mcp/errors.go b/pkg/mcp/errors.go index 1cc1aacf2d..b3990b2b52 100644 --- a/pkg/mcp/errors.go +++ b/pkg/mcp/errors.go @@ -3,6 +3,48 @@ package mcp +import "net/http" + +// JSONRPCCodeDenied is the application-space JSON-RPC error code ToolHive uses +// when a policy denies a call, alongside HTTP 403. Both denial paths reference +// this constant — the single-server authorization middleware +// (pkg/authz.handleUnauthorized) and the vMCP Serve-path call gate — so a denial +// is represented by one code across `thv run` and vMCP, by reference rather than +// by two copies of the literal. It is deliberately outside the reserved +// -32768..-32000 JSON-RPC range so it never collides with an SDK-generated code. +// It is also deliberately untyped: it is consumed both as an int +// (server.Denial.Code, map[string]any envelopes) and as an int64 +// (jsonrpc2.NewError, JSONRPCCodeForStatus). Typing it would break +// assert.Equal sites on one side. +const JSONRPCCodeDenied = 403 + +// JSONRPCCodeForStatus maps an HTTP status code to the JSON-RPC error code +// that belongs in a response body's error.code field. The HTTP status and the +// JSON-RPC code are independent dimensions: ToolHive's error paths pick an +// HTTP status per failure class, and each class has exactly one correct +// JSON-RPC code. This function is the single place that mapping lives, so +// callers never pass an HTTP status straight through as the JSON-RPC code. +// +// http.StatusUnprocessableEntity maps to JSONRPCCodeDenied because ToolHive's +// mutating-webhook path relays a webhook's HTTP 422 always-deny response +// (webhook.IsAlwaysDenyError) as HTTP 422 downstream — the validating path +// relays the same condition as 403. So a 422 arriving here is semantically a +// denial. +// +// Any status with no explicit case below maps to CodeInternalError, so an +// unmapped status fails safe to a standard code instead of leaking a raw HTTP +// status into the JSON-RPC code field. +func JSONRPCCodeForStatus(status int) int64 { + switch status { + case http.StatusForbidden, http.StatusUnprocessableEntity: + return JSONRPCCodeDenied + case http.StatusRequestEntityTooLarge: + return CodeInvalidRequest + default: + return CodeInternalError + } +} + // CodedError is implemented by domain errors that should surface a stable error // code and optional structured data in an MCP tool result. // diff --git a/pkg/mcp/errors_test.go b/pkg/mcp/errors_test.go new file mode 100644 index 0000000000..411eb6a0e2 --- /dev/null +++ b/pkg/mcp/errors_test.go @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestJSONRPCCodeForStatus(t *testing.T) { + t.Parallel() + tests := []struct { + name string + status int + want int64 + }{ + {name: "forbidden maps to denied", status: http.StatusForbidden, want: JSONRPCCodeDenied}, + {name: "request entity too large maps to invalid request", status: http.StatusRequestEntityTooLarge, want: CodeInvalidRequest}, + {name: "unprocessable entity maps to denied", status: http.StatusUnprocessableEntity, want: JSONRPCCodeDenied}, + {name: "internal server error falls through to internal error", status: http.StatusInternalServerError, want: CodeInternalError}, + {name: "unmapped status defaults to internal error", status: http.StatusBadGateway, want: CodeInternalError}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, JSONRPCCodeForStatus(tt.status)) + }) + } +} diff --git a/pkg/mcp/jsonrpc_write.go b/pkg/mcp/jsonrpc_write.go new file mode 100644 index 0000000000..fce87708a3 --- /dev/null +++ b/pkg/mcp/jsonrpc_write.go @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "log/slog" + "net/http" + + "golang.org/x/exp/jsonrpc2" +) + +// EncodeJSONRPCError returns the wire encoding of a jsonrpc2 error response. +// +// This is the single sanctioned way to serialize a jsonrpc2 envelope in this +// codebase. NEVER json.Marshal (or json.NewEncoder().Encode) a jsonrpc2 type: +// jsonrpc2.Response carries no json tags and jsonrpc2.ID's only field is +// unexported, so reflection produces Go-capitalized keys, drops the mandatory +// "jsonrpc":"2.0" tag, and renders every id — valid or not — as "{}" (#5950). +// jsonrpc2.EncodeMessage marshals through the library's tagged wire struct: +// lowercase keys, the version tag stamped unconditionally, and an id that is +// omitted when absent (never null, which MCP cannot express and the reference +// TypeScript SDK client rejects, #6038) while a legitimate id of 0 survives. +// +// EncodeMessage cannot fail for a Response built from an ID and a +// jsonrpc2.NewError; if it somehow does, this logs and falls back to a +// hardcoded, conformant Internal Error body rather than returning nothing — +// the original code and id are not trustworthy at that point. +func EncodeJSONRPCError(resp *jsonrpc2.Response) []byte { + body, err := jsonrpc2.EncodeMessage(resp) + if err != nil { + slog.Error("failed to encode JSON-RPC error response", "error", err) + return []byte(`{"jsonrpc":"2.0","error":{"code":-32603,"message":"Internal error"}}`) + } + return body +} + +// WriteJSONRPCError writes resp as an HTTP response: the JSON-RPC error body +// under Content-Type application/json with the given HTTP status. The body is +// encoded before any header is written, so an encode failure never leaves a +// half-written response (see EncodeJSONRPCError for the fallback). Returns +// the body write error; callers with no recovery path may discard it. +// +// Do not use this to write into an SSE stream — a bare JSON object in a +// text/event-stream parses as an unrecognized field and is silently +// discarded (#6037); SSE paths must frame the body as an event instead. +func WriteJSONRPCError(w http.ResponseWriter, httpStatus int, resp *jsonrpc2.Response) error { + body := EncodeJSONRPCError(resp) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(httpStatus) + _, err := w.Write(body) + return err +} diff --git a/pkg/mcp/jsonrpc_write_test.go b/pkg/mcp/jsonrpc_write_test.go new file mode 100644 index 0000000000..338cb24de0 --- /dev/null +++ b/pkg/mcp/jsonrpc_write_test.go @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" +) + +// TestEncodeJSONRPCError pins the wire properties the helper exists to +// guarantee: lowercase keys with the version tag stamped (the #5950 bug was +// reflection over the untagged jsonrpc2.Response), an absent id omitted +// rather than rendered null (#6038), and a legitimate id of 0 surviving — +// the omitempty on the library's wire struct tests IsNil, not zero-ness. +func TestEncodeJSONRPCError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + id jsonrpc2.ID + want string + }{ + { + name: "int64 id", + id: jsonrpc2.Int64ID(42), + want: `{"jsonrpc":"2.0","id":42,"error":{"code":403,"message":"denied"}}`, + }, + { + name: "string id", + id: jsonrpc2.StringID("abc"), + want: `{"jsonrpc":"2.0","id":"abc","error":{"code":403,"message":"denied"}}`, + }, + { + name: "zero id survives", + id: jsonrpc2.Int64ID(0), + want: `{"jsonrpc":"2.0","id":0,"error":{"code":403,"message":"denied"}}`, + }, + { + name: "empty id omits the key", + id: jsonrpc2.ID{}, + want: `{"jsonrpc":"2.0","error":{"code":403,"message":"denied"}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + body := EncodeJSONRPCError(&jsonrpc2.Response{ + ID: tt.id, + Error: jsonrpc2.NewError(JSONRPCCodeDenied, "denied"), + }) + + // assert.JSONEq is key-case-sensitive, so this single assertion + // catches "Error"/"ID" substituted for "error"/"id" and a missing + // "jsonrpc" tag; key presence is asserted separately because a + // decoded nil is indistinguishable from a decoded null. + assert.JSONEq(t, tt.want, string(body)) + + var decoded map[string]json.RawMessage + require.NoError(t, json.Unmarshal(body, &decoded)) + _, hasID := decoded["id"] + assert.Equal(t, tt.id != jsonrpc2.ID{}, hasID, "id key presence mismatch") + }) + } +} + +func TestWriteJSONRPCError(t *testing.T) { + t.Parallel() + + rr := httptest.NewRecorder() + err := WriteJSONRPCError(rr, http.StatusForbidden, &jsonrpc2.Response{ + ID: jsonrpc2.Int64ID(7), + Error: jsonrpc2.NewError(JSONRPCCodeDenied, "denied"), + }) + require.NoError(t, err) + + assert.Equal(t, http.StatusForbidden, rr.Code) + assert.Equal(t, "application/json", rr.Header().Get("Content-Type")) + assert.JSONEq(t, `{"jsonrpc":"2.0","id":7,"error":{"code":403,"message":"denied"}}`, rr.Body.String()) +} diff --git a/pkg/mcp/parser.go b/pkg/mcp/parser.go index 5016df2d8e..9eddf5047d 100644 --- a/pkg/mcp/parser.go +++ b/pkg/mcp/parser.go @@ -8,7 +8,9 @@ import ( "bytes" "context" "encoding/json" + "errors" "io" + "mime" "net/http" "strconv" "strings" @@ -41,9 +43,27 @@ type ParsedMCPRequest struct { // Meta contains the _meta field from the request params for protocol-level metadata // such as progress tokens, trace IDs, or custom namespaced metadata Meta map[string]interface{} + // MCPMethodHeader is the value of the Modern (stateless MCP) "Mcp-Method" + // request header, if present. Mandatory on every Modern POST per the draft spec. + MCPMethodHeader string + // MCPNameHeader is the raw, as-received value of the Modern (stateless MCP) + // "Mcp-Name" request header, if present. Required for tools/call, + // resources/read, prompts/get. Stored undecoded: the spec allows the header + // value to be sentinel-encoded (=?base64?...?=), and a caller comparing it to + // the plain body name/uri during validation must decode the header first. + MCPNameHeader string + // ClientInfo is the client implementation info surfaced via _meta for Modern + // (stateless) requests, sourced from _meta["io.modelcontextprotocol/clientInfo"]. + ClientInfo map[string]interface{} + // ProtocolVersion is the per-request protocol version surfaced via _meta for + // Modern (stateless) requests, sourced from _meta["io.modelcontextprotocol/protocolVersion"]. + ProtocolVersion string // IsRequest indicates if this is a JSON-RPC request (vs response or notification) IsRequest bool - // IsBatch indicates if this is a batch request + // IsBatch indicates if this is a batch request. It is always false: JSON-RPC + // batches are rejected by ParsingMiddleware before a ParsedMCPRequest is ever + // constructed (batching was removed in MCP 2025-06-18; see batch.go). The + // field is retained for wire/telemetry compatibility. IsBatch bool } @@ -61,10 +81,10 @@ type ParsedMCPRequest struct { // Example usage: // // middlewares := []types.Middleware{ -// authMiddleware, // Authentication first +// auditMiddleware, // Audit wraps the chain; parsed data flows back via ParsedRequestHolder +// authMiddleware, // Authentication // mcp.ParsingMiddleware, // MCP parsing after auth // authzMiddleware, // Authorization uses parsed data -// auditMiddleware, // Audit uses parsed data // } func ParsingMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -92,9 +112,28 @@ func ParsingMiddleware(next http.Handler) http.Handler { // Restore the request body for downstream handlers r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + // Reject JSON-RPC batches before any downstream middleware or the + // transport proxy's batch executor can act on them. Batching was + // removed in MCP revision 2025-06-18; ToolHive serves only 2025-11-25 + // and 2026-07-28. Authz, audit, and tool filtering each inspect a + // single parsed request per call, so a batch reaching the backend would + // smuggle every nested call past those controls (see #5745). + if IsBatchRequest(bodyBytes) { + WriteBatchUnsupportedError(w) + return + } + // Parse the MCP request and store in context parsedRequest := parseMCPRequest(bodyBytes) if parsedRequest != nil { + parsedRequest.MCPMethodHeader = r.Header.Get("Mcp-Method") + parsedRequest.MCPNameHeader = r.Header.Get("Mcp-Name") + // Publish to a ParsedRequestHolder if one is present, so middleware + // wrapping the parser (e.g. audit) can observe the parsed request + // even though the derived context only flows downstream. + if holder, ok := ParsedRequestHolderFromContext(r.Context()); ok { + holder.Parsed = parsedRequest + } ctx := context.WithValue(r.Context(), MCPRequestContextKey, parsedRequest) r = r.WithContext(ctx) } @@ -104,6 +143,121 @@ func ParsingMiddleware(next http.Handler) http.Handler { }) } +// RepublishParsedMCPRequest refreshes the cached parse after middleware has +// rewritten the request body, returning the request to pass downstream. +// ParsingMiddleware deliberately parses only once, so any middleware that +// replaces r.Body MUST call this or downstream consumers (authorization, +// audit, telemetry) will decide on the bytes that arrived rather than the +// bytes the backend executes. +// +// On error, the returned request is nil and must not be passed downstream: +// the caller is responsible for terminating the request (e.g. writing an +// error response) instead of proceeding with a stale or absent parse. +// +// This only refreshes consumers that read the parse from the request context or +// from a ParsedRequestHolder. Middleware that inspects the raw body from OUTSIDE +// ParsingMiddleware — the tool-call filter and the rate limiter — has already +// decided against the pre-rewrite body and is not corrected by republishing. +// +// The caller must also refresh r.ContentLength when it replaces r.Body, or the +// reverse proxy will reject the forwarded request. +func RepublishParsedMCPRequest(r *http.Request, body []byte) (*http.Request, error) { + // Batch-reject before parsing, using the same guard ParsingMiddleware uses, + // so a mutation can never smuggle a batch past authz/audit by rewriting a + // single request into an array (see IsBatchRequest's doc comment). + if IsBatchRequest(body) { + return nil, &BatchUnsupportedError{} + } + + parsed := parseMCPRequest(body) + if parsed == nil { + return nil, errors.New("republished body is not a valid JSON-RPC request") + } + parsed.MCPMethodHeader = r.Header.Get("Mcp-Method") + parsed.MCPNameHeader = r.Header.Get("Mcp-Name") + + if holder, ok := ParsedRequestHolderFromContext(r.Context()); ok { + holder.Parsed = parsed + } + + return r.WithContext(context.WithValue(r.Context(), MCPRequestContextKey, parsed)), nil +} + +// parsedRequestHolderContextKey is the context key for ParsedRequestHolder. +type parsedRequestHolderContextKey struct{} + +// ParsedRequestHolder is a mutable carrier that lets middleware running +// OUTSIDE the parsing middleware observe the parsed MCP request. Context +// values only flow downstream, so a wrapper such as the audit middleware +// cannot read the parsed request that the parser attaches for inner handlers. +// The wrapper injects an empty holder via WithParsedRequestHolder before +// calling the inner chain; ParsingMiddleware fills it; the wrapper reads it +// after the inner chain returns. +// +// The holder is written and read by the single request goroutine (writes +// happen-before the wrapper's post-ServeHTTP read), so no synchronization is +// needed. +type ParsedRequestHolder struct { + Parsed *ParsedMCPRequest +} + +// WithParsedRequestHolder returns a new context carrying the given holder. +func WithParsedRequestHolder(ctx context.Context, holder *ParsedRequestHolder) context.Context { + return context.WithValue(ctx, parsedRequestHolderContextKey{}, holder) +} + +// ParsedRequestHolderFromContext retrieves the ParsedRequestHolder from the +// context. Returns (nil, false) if no holder is present. +func ParsedRequestHolderFromContext(ctx context.Context) (*ParsedRequestHolder, bool) { + holder, ok := ctx.Value(parsedRequestHolderContextKey{}).(*ParsedRequestHolder) + return holder, ok && holder != nil +} + +// authzDenialMarkerContextKey is the context key for AuthzDenialMarker. +type authzDenialMarkerContextKey struct{} + +// AuthzDenialMarker is a mutable carrier that lets the authorization +// middleware (pkg/authz), which runs INSIDE the audit middleware, flag a +// request it refused before message-level authorization could run, such as a +// non-JSON POST carrying a smuggled JSON-RPC body. It follows the same +// propagation pattern as ParsedRequestHolder: the audit wrapper injects an +// empty marker via WithAuthzDenialMarker, the inner middleware fills it, and +// the wrapper reads it back after the inner chain returns so the refusal is +// audited as a denial rather than a generic 400 failure. The type lives in +// this package because pkg/authz already depends on pkg/audit transitively, +// so the carrier cannot live in either of those two packages. +// +// Like ParsedRequestHolder, the marker is written and read by the single +// request goroutine, so no synchronization is needed. +type AuthzDenialMarker struct { + Denied bool +} + +// WithAuthzDenialMarker returns a new context carrying the given marker. +func WithAuthzDenialMarker(ctx context.Context, marker *AuthzDenialMarker) context.Context { + return context.WithValue(ctx, authzDenialMarkerContextKey{}, marker) +} + +// AuthzDenialMarkerFromContext retrieves the AuthzDenialMarker from the +// context. Returns (nil, false) if no marker is present. +func AuthzDenialMarkerFromContext(ctx context.Context) (*AuthzDenialMarker, bool) { + marker, ok := ctx.Value(authzDenialMarkerContextKey{}).(*AuthzDenialMarker) + return marker, ok && marker != nil +} + +// RequestHasJSONContentType reports whether r declares an application/json +// Content-Type. Media types are case-insensitive per RFC 9110 8.3.1, so the +// header is parsed with mime.ParseMediaType and compared with EqualFold; the +// case-sensitive prefix match this replaces also wrongly accepted longer +// types such as "application/jsonx". +func RequestHasJSONContentType(r *http.Request) bool { + mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type")) + if err != nil { + return false + } + return strings.EqualFold(mediaType, "application/json") +} + // GetParsedMCPRequest retrieves the parsed MCP request from the request context. // Returns nil if no parsed request is available. func GetParsedMCPRequest(ctx context.Context) *ParsedMCPRequest { @@ -120,8 +274,7 @@ func shouldParseMCPRequest(r *http.Request) bool { return false } - contentType := r.Header.Get("Content-Type") - if !strings.HasPrefix(contentType, "application/json") { + if !RequestHasJSONContentType(r) { return false } @@ -155,6 +308,7 @@ func parseMCPRequest(bodyBytes []byte) *ParsedMCPRequest { case *jsonrpc2.Request: // Extract resource ID, arguments, and meta based on the method resourceID, arguments, meta := extractResourceAndArguments(msg.Method, msg.Params) + clientInfo, protocolVersion := extractModernMeta(meta) // Determine the ID - will be nil for notifications var id interface{} @@ -163,16 +317,26 @@ func parseMCPRequest(bodyBytes []byte) *ParsedMCPRequest { } return &ParsedMCPRequest{ - Method: msg.Method, - ID: id, - Params: msg.Params, - ResourceID: resourceID, - Arguments: arguments, - Meta: meta, - IsRequest: true, - IsBatch: false, // TODO: Add batch request support if needed + Method: msg.Method, + ID: id, + Params: msg.Params, + ResourceID: resourceID, + Arguments: arguments, + Meta: meta, + ClientInfo: clientInfo, + ProtocolVersion: protocolVersion, + IsRequest: true, + // Batches are rejected in ParsingMiddleware before reaching here, so a + // parsed request is always a single, non-batch message. + IsBatch: false, } case *jsonrpc2.Response: + // A client reply to a server-initiated request: no "method" member at all and + // a "result" or "error". Anything else (a request whose method was removed or + // duplicated) is not parsed, so it can never skip authorization as a reply. + if !isClientResponse(bodyBytes) { + return nil + } var id interface{} if msg.ID.IsValid() { id = msg.ID.Raw() @@ -188,6 +352,25 @@ func parseMCPRequest(bodyBytes []byte) *ParsedMCPRequest { } } +// IsClientResponse reports whether p is a client reply to a server-initiated +// request (see isClientResponse); such a message is not an authorization subject. +func (p *ParsedMCPRequest) IsClientResponse() bool { + return p != nil && !p.IsRequest && p.Method == "" +} + +// isClientResponse reports whether body is a JSON object with a "result" or "error" +// member and no "method" member. +func isClientResponse(body []byte) bool { + var members map[string]json.RawMessage + if err := json.Unmarshal(body, &members); err != nil { + return false + } + _, hasMethod := members["method"] + _, hasResult := members["result"] + _, hasError := members["error"] + return !hasMethod && (hasResult || hasError) +} + // extractResourceAndArguments extracts the resource ID, arguments, and _meta field from the JSON-RPC params // based on the MCP method type. // methodHandler defines a function type for handling specific MCP methods @@ -230,6 +413,7 @@ var staticResourceIDs = map[string]string{ "notifications/resources/list_changed": "resources", "notifications/resources/updated": "resources", "notifications/tools/list_changed": "tools", + "server/discover": "discover", } func extractResourceAndArguments(method string, params json.RawMessage) (string, map[string]interface{}, map[string]interface{}) { @@ -242,18 +426,23 @@ func extractResourceAndArguments(method string, params json.RawMessage) (string, return getStaticResourceID(method), nil, nil } - // Extract _meta field if present - var meta map[string]interface{} - if metaRaw, ok := paramsMap["_meta"]; ok { - if metaMap, ok := metaRaw.(map[string]interface{}); ok { - meta = metaMap - } - } + meta := metaFromParamsMap(paramsMap) resourceID, arguments := processMethodWithHandler(method, paramsMap) return resourceID, arguments, meta } +// extractModernMeta surfaces the Modern (stateless MCP) clientInfo and +// protocolVersion fields from a parsed _meta map, if present. It delegates to +// the reserved-key helpers in revision.go so the guarded type assertions live +// in one place; a wrong-shaped value is treated as absent rather than causing +// an error. +func extractModernMeta(meta map[string]interface{}) (clientInfo map[string]interface{}, protocolVersion string) { + clientInfo, _ = objectMetaValue(meta, metaKeyClientInfo) + protocolVersion, _ = stringMetaValue(meta, metaKeyProtocolVersion) + return clientInfo, protocolVersion +} + // getStaticResourceID returns the static resource ID for methods that don't need parameter parsing func getStaticResourceID(method string) string { if resourceID, exists := staticResourceIDs[method]; exists { @@ -506,3 +695,23 @@ func GetMCPMeta(ctx context.Context) map[string]interface{} { } return nil } + +// GetMCPClientInfo is a convenience function to get the Modern (stateless MCP) +// per-request clientInfo from the context. +// Returns nil if no parsed request is available or clientInfo is not present. +func GetMCPClientInfo(ctx context.Context) map[string]interface{} { + if parsed := GetParsedMCPRequest(ctx); parsed != nil { + return parsed.ClientInfo + } + return nil +} + +// GetMCPProtocolVersion is a convenience function to get the Modern (stateless +// MCP) per-request protocol version from the context. +// Returns "" if no parsed request is available or protocolVersion is not present. +func GetMCPProtocolVersion(ctx context.Context) string { + if parsed := GetParsedMCPRequest(ctx); parsed != nil { + return parsed.ProtocolVersion + } + return "" +} diff --git a/pkg/mcp/parser_integration_test.go b/pkg/mcp/parser_integration_test.go index 398dba5f20..2e585d0830 100644 --- a/pkg/mcp/parser_integration_test.go +++ b/pkg/mcp/parser_integration_test.go @@ -10,11 +10,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" ) // TestParsingMiddlewareWithRealMCPClients tests the parsing middleware with real MCP clients and servers @@ -143,17 +144,33 @@ func TestParsingMiddlewareWithRealMCPClients(t *testing.T) { assert.NotEmpty(t, readResult.Contents) // Verify that all requests were parsed by the middleware - assert.GreaterOrEqual(t, len(parsedRequests), 4, "Expected at least 4 parsed requests (initialize, list tools, call tool, list resources, read resource)") - - // Verify specific parsed requests - foundInitialize := false + assert.GreaterOrEqual(t, len(parsedRequests), 4, "Expected at least 4 parsed requests (session establishment, list tools, call tool, list resources, read resource)") + + // Verify specific parsed requests. mcpcompat's client is Modern-first + // (SEP-2575) on every transport, so both arms send server/discover + // first. The SSE server transport advertises 2026-07-28 unconditionally + // (no ProtocolVersionSupporter gate), so discover succeeds and + // "initialize" is never sent there. The streamable-HTTP server here is + // stateful (no WithStateless), which gates 2026-07-28 on being + // stateless, so its discover negotiates down and it additionally sends + // "initialize". + foundSessionEstablished := false foundToolCall := false foundResourceRead := false + var methodsSeen []string for _, parsed := range parsedRequests { + methodsSeen = append(methodsSeen, parsed.Method) switch parsed.Method { case "initialize": - foundInitialize = true + foundSessionEstablished = true assert.Equal(t, "test-client", parsed.ResourceID) + case "server/discover": + if tc.transport == "sse" { + // parser.go maps server/discover to the static ResourceID + // "discover" (not a client-name ResourceID). + foundSessionEstablished = true + assert.Equal(t, "discover", parsed.ResourceID) + } case "tools/call": foundToolCall = true assert.Equal(t, "test_tool", parsed.ResourceID) @@ -162,7 +179,11 @@ func TestParsingMiddlewareWithRealMCPClients(t *testing.T) { assert.Equal(t, "test://resource", parsed.ResourceID) } } - assert.True(t, foundInitialize, "Initialize request should have been parsed") + assert.True(t, foundSessionEstablished, "session establishment request (initialize or server/discover) should have been parsed") + if tc.transport == "sse" { + assert.Contains(t, methodsSeen, "server/discover", "SSE arm should negotiate via server/discover") + assert.NotContains(t, methodsSeen, "initialize", "SSE arm is Modern-first and should never fall back to initialize here") + } assert.True(t, foundToolCall, "Tool call request should have been parsed") assert.True(t, foundResourceRead, "Resource read request should have been parsed") }) diff --git a/pkg/mcp/parser_test.go b/pkg/mcp/parser_test.go index ec653c5635..239cedb0a0 100644 --- a/pkg/mcp/parser_test.go +++ b/pkg/mcp/parser_test.go @@ -100,6 +100,18 @@ func TestParsingMiddleware(t *testing.T) { expectedIsReq: false, expectedID: int64(1), }, + { + name: "server/discover request", + method: "POST", + path: "/messages", + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":10,"method":"server/discover","params":{}}`, + expectParsed: true, + expectedMethod: "server/discover", + expectedID: int64(10), + expectedResID: "discover", + expectedIsReq: true, + }, { name: "GET request - not parsed", method: "GET", @@ -247,6 +259,138 @@ func TestParsingMiddleware(t *testing.T) { } } +func TestParsingMiddlewareRejectsBatch(t *testing.T) { + t.Parallel() + tests := []struct { + name string + body string + }{ + { + name: "batch of requests", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},` + + `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"danger"}}]`, + }, + { + name: "batch with leading whitespace", + body: " \n\t[{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/list\"}]", + }, + { + name: "empty batch", + body: `[]`, + }, + { + name: "malformed batch (unterminated array)", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/call"`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + // The next handler must never run for a batch: reaching it means the + // batch bypassed rejection and would hit the backend uninspected. + nextCalled := false + testHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + }) + + middleware := ParsingMiddleware(testHandler) + + req := httptest.NewRequest("POST", "/messages", bytes.NewBufferString(tt.body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + middleware.ServeHTTP(w, req) + + assert.False(t, nextCalled, "batch must be rejected before the next handler") + assert.Equal(t, http.StatusBadRequest, w.Code) + + var resp struct { + JSONRPC string `json:"jsonrpc"` + Error struct { + Code int64 `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.Equal(t, "2.0", resp.JSONRPC) + assert.Equal(t, CodeInvalidRequest, resp.Error.Code) + assert.NotEmpty(t, resp.Error.Message) + + // A batch has no single request id to echo. A tagged-struct + // decode (above) can't tell an absent "id" key apart from a + // present null -- both decode to the zero value -- so the + // omission is checked via a map decode instead. + var asMap map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &asMap)) + _, hasID := asMap["id"] + assert.False(t, hasID, `"id" key must be omitted, not present as null`) + }) + } +} + +func TestParsingMiddlewareModernHeaders(t *testing.T) { + t.Parallel() + tests := []struct { + name string + mcpMethodHeader string + mcpNameHeader string + expectedMCPMethod string + expectedMCPName string + }{ + { + name: "both headers set", + mcpMethodHeader: "tools/call", + mcpNameHeader: "some-tool", + expectedMCPMethod: "tools/call", + expectedMCPName: "some-tool", + }, + { + name: "neither header set", + expectedMCPMethod: "", + expectedMCPName: "", + }, + { + name: "sentinel-encoded Mcp-Name stored undecoded", + mcpMethodHeader: "tools/call", + mcpNameHeader: "=?base64?dG9vbA==?=", + expectedMCPMethod: "tools/call", + expectedMCPName: "=?base64?dG9vbA==?=", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + var capturedCtx context.Context + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedCtx = r.Context() + w.WriteHeader(http.StatusOK) + }) + + middleware := ParsingMiddleware(testHandler) + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"weather"}}` + req := httptest.NewRequest("POST", "/messages", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + if tt.mcpMethodHeader != "" { + req.Header.Set("Mcp-Method", tt.mcpMethodHeader) + } + if tt.mcpNameHeader != "" { + req.Header.Set("Mcp-Name", tt.mcpNameHeader) + } + w := httptest.NewRecorder() + + middleware.ServeHTTP(w, req) + + parsed := GetParsedMCPRequest(capturedCtx) + require.NotNil(t, parsed) + assert.Equal(t, tt.expectedMCPMethod, parsed.MCPMethodHeader) + assert.Equal(t, tt.expectedMCPName, parsed.MCPNameHeader) + }) + } +} + func TestExtractResourceAndArguments(t *testing.T) { t.Parallel() tests := []struct { @@ -1074,6 +1218,124 @@ func TestMetaFieldParsing(t *testing.T) { } } +func TestModernMetaParsing(t *testing.T) { + t.Parallel() + tests := []struct { + name string + body string + expectedClientInfo map[string]interface{} + expectedProtocolVersion string + }{ + { + name: "clientInfo and protocolVersion present", + body: `{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "weather", + "_meta": { + "io.modelcontextprotocol/clientInfo": {"name": "test-client", "version": "1.0"}, + "io.modelcontextprotocol/protocolVersion": "2026-07-28" + } + } + }`, + expectedClientInfo: map[string]interface{}{ + "name": "test-client", + "version": "1.0", + }, + expectedProtocolVersion: "2026-07-28", + }, + { + name: "_meta absent", + body: `{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "weather" + } + }`, + expectedClientInfo: nil, + expectedProtocolVersion: "", + }, + { + name: "_meta present without modern keys", + body: `{ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "weather", + "_meta": { + "progressToken": "abc123" + } + } + }`, + expectedClientInfo: nil, + expectedProtocolVersion: "", + }, + { + name: "protocolVersion wrong type", + body: `{ + "jsonrpc": "2.0", + "id": 4, + "method": "tools/call", + "params": { + "name": "weather", + "_meta": { + "io.modelcontextprotocol/protocolVersion": 12345 + } + } + }`, + expectedClientInfo: nil, + expectedProtocolVersion: "", + }, + { + name: "clientInfo wrong type", + body: `{ + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": { + "name": "weather", + "_meta": { + "io.modelcontextprotocol/clientInfo": "not-an-object" + } + } + }`, + expectedClientInfo: nil, + expectedProtocolVersion: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + var capturedCtx context.Context + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedCtx = r.Context() + w.WriteHeader(http.StatusOK) + }) + + middleware := ParsingMiddleware(testHandler) + req := httptest.NewRequest("POST", "/messages", bytes.NewBufferString(tt.body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + middleware.ServeHTTP(w, req) + + parsed := GetParsedMCPRequest(capturedCtx) + require.NotNil(t, parsed) + assert.Equal(t, tt.expectedClientInfo, parsed.ClientInfo) + assert.Equal(t, tt.expectedProtocolVersion, parsed.ProtocolVersion) + + assert.Equal(t, tt.expectedClientInfo, GetMCPClientInfo(capturedCtx)) + assert.Equal(t, tt.expectedProtocolVersion, GetMCPProtocolVersion(capturedCtx)) + }) + } +} + func TestMetaFieldInvalidTypes(t *testing.T) { t.Parallel() tests := []struct { @@ -1360,14 +1622,9 @@ func TestParsingMiddlewareErrorHandling(t *testing.T) { body: bytes.NewBufferString(`{"invalid json`), expectParsed: false, }, - { - name: "JSON array instead of object", - method: "POST", - path: "/messages", - contentType: "application/json", - body: bytes.NewBufferString(`[{"jsonrpc":"2.0"}]`), - expectParsed: false, - }, + // A top-level JSON array is a JSON-RPC batch; ParsingMiddleware rejects + // it outright (see TestParsingMiddlewareRejectsBatch) rather than passing + // it through, so it is deliberately not exercised here. } for _, tt := range tests { @@ -1431,6 +1688,11 @@ func TestExtractResourceAndArgumentsNilParams(t *testing.T) { method: "notifications/initialized", expectedResourceID: "initialized", }, + { + name: "server/discover", + method: "server/discover", + expectedResourceID: "discover", + }, } for _, tt := range tests { @@ -1444,32 +1706,6 @@ func TestExtractResourceAndArgumentsNilParams(t *testing.T) { } } -func TestParsingMiddlewareWithBatchRequests(t *testing.T) { - t.Parallel() - // Test batch JSON-RPC requests (currently not supported but should not crash) - batchBody := `[ - {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"tool1"}}, - {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"tool2"}} - ]` - - var capturedCtx context.Context - testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - capturedCtx = r.Context() - w.WriteHeader(http.StatusOK) - }) - - middleware := ParsingMiddleware(testHandler) - req := httptest.NewRequest("POST", "/messages", bytes.NewBufferString(batchBody)) - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() - - middleware.ServeHTTP(w, req) - - // Batch requests should not be parsed (not supported yet) - parsed := GetParsedMCPRequest(capturedCtx) - assert.Nil(t, parsed) -} - func TestConvenienceFunctionsWithNilContext(t *testing.T) { t.Parallel() // Test convenience functions with nil parsed request @@ -1676,3 +1912,142 @@ func TestParsingMiddlewareIntegration(t *testing.T) { }) } } + +func TestRepublishParsedMCPRequest(t *testing.T) { + t.Parallel() + + oldBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"old-tool","arguments":{"a":1}}}` + newBody := `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"new-tool","arguments":{"b":2}}}` + + tests := []struct { + name string + body []byte + expectErr bool + expectBatch bool + expectNilReq bool + checkResponse func(t *testing.T, req *http.Request, err error) + }{ + { + name: "happy re-parse reflects new body", + body: []byte(newBody), + checkResponse: func(t *testing.T, req *http.Request, err error) { + t.Helper() + require.NoError(t, err) + parsed := GetParsedMCPRequest(req.Context()) + require.NotNil(t, parsed) + assert.Equal(t, "tools/call", parsed.Method) + assert.Equal(t, "new-tool", parsed.ResourceID) + assert.Equal(t, map[string]interface{}{"b": float64(2)}, parsed.Arguments) + assert.NotNil(t, parsed.Params) + assert.Equal(t, int64(2), parsed.ID) + assert.False(t, parsed.IsBatch) + }, + }, + { + name: "batch body rejected before parsing", + body: []byte(`[{"jsonrpc":"2.0","method":"tools/call","id":1}]`), + expectErr: true, + expectBatch: true, + }, + { + name: "whitespace-prefixed batch still detected", + body: []byte(" \t[{\"jsonrpc\":\"2.0\",\"method\":\"tools/call\",\"id\":1}]"), + expectErr: true, + expectBatch: true, + }, + { + name: "valid JSON that is not a request", + body: []byte(`{"jsonrpc":"2.0","result":{}}`), + expectErr: true, + }, + { + name: "empty object is not a request", + body: []byte(`{}`), + expectErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/messages", bytes.NewBufferString(oldBody)) + req.Header.Set("Content-Type", "application/json") + var capturedOldCtx context.Context + testHandler := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + capturedOldCtx = r.Context() + }) + ParsingMiddleware(testHandler).ServeHTTP(httptest.NewRecorder(), req) + require.NotNil(t, GetParsedMCPRequest(capturedOldCtx), "precondition: old body must have parsed") + req = req.WithContext(capturedOldCtx) + + republished, err := RepublishParsedMCPRequest(req, tt.body) + + if tt.expectErr { + require.Error(t, err) + assert.Nil(t, republished) + if tt.expectBatch { + var batchErr *BatchUnsupportedError + assert.ErrorAs(t, err, &batchErr) + } + return + } + + require.NoError(t, err) + require.NotNil(t, republished) + if tt.checkResponse != nil { + tt.checkResponse(t, republished, err) + } + + // The original request's context must still yield the OLD parse: + // RepublishParsedMCPRequest must not mutate the caller's request. + oldParsed := GetParsedMCPRequest(req.Context()) + require.NotNil(t, oldParsed) + assert.Equal(t, "old-tool", oldParsed.ResourceID) + }) + } +} + +func TestRepublishParsedMCPRequestHeaders(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/messages", bytes.NewBufferString("")) + req.Header.Set("Mcp-Method", "tools/call") + req.Header.Set("Mcp-Name", "some-tool") + + body := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"weather"}}`) + republished, err := RepublishParsedMCPRequest(req, body) + require.NoError(t, err) + + parsed := GetParsedMCPRequest(republished.Context()) + require.NotNil(t, parsed) + assert.Equal(t, "tools/call", parsed.MCPMethodHeader) + assert.Equal(t, "some-tool", parsed.MCPNameHeader) +} + +func TestRepublishParsedMCPRequestHolder(t *testing.T) { + t.Parallel() + body := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"weather"}}`) + + t.Run("holder present is refreshed", func(t *testing.T) { + t.Parallel() + req := httptest.NewRequest(http.MethodPost, "/messages", bytes.NewBufferString("")) + holder := &ParsedRequestHolder{} + req = req.WithContext(WithParsedRequestHolder(req.Context(), holder)) + + republished, err := RepublishParsedMCPRequest(req, body) + require.NoError(t, err) + require.NotNil(t, republished) + require.NotNil(t, holder.Parsed) + assert.Equal(t, "weather", holder.Parsed.ResourceID) + }) + + t.Run("holder absent does not panic", func(t *testing.T) { + t.Parallel() + req := httptest.NewRequest(http.MethodPost, "/messages", bytes.NewBufferString("")) + + republished, err := RepublishParsedMCPRequest(req, body) + require.NoError(t, err) + require.NotNil(t, republished) + }) +} diff --git a/pkg/mcp/revision.go b/pkg/mcp/revision.go new file mode 100644 index 0000000000..2f2e7e58f9 --- /dev/null +++ b/pkg/mcp/revision.go @@ -0,0 +1,685 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "maps" + "strings" +) + +// Revision identifies which MCP protocol era a request belongs to. +type Revision int + +const ( + // RevisionLegacy is the current 2025-11-25 MCP revision: session-based, + // initialize handshake, Mcp-Session-Id. + RevisionLegacy Revision = iota + // RevisionModern is the 2026-07-28 MCP revision: stateless, no initialize, + // protocol metadata carried per-request in _meta. + RevisionModern +) + +// MCPVersionModern is the single Modern (stateless) protocol version this +// build understands. +const MCPVersionModern = "2026-07-28" + +// MCPVersionLegacy is the single Legacy (session-based) protocol version this +// build understands. It is what RevisionLegacy names on the wire and the one +// version mcpcompat's initialize handshake negotiates. +const MCPVersionLegacy = "2025-11-25" + +// String returns the wire protocol-version string for the revision: +// MCPVersionModern for RevisionModern, MCPVersionLegacy otherwise. +func (r Revision) String() string { + if r == RevisionModern { + return MCPVersionModern + } + return MCPVersionLegacy +} + +// metaKeyProtocolVersion is the reserved _meta key that carries the per-request +// protocol version on Modern (stateless) MCP requests, per the draft MCP schema's +// RequestMetaObject. +const metaKeyProtocolVersion = "io.modelcontextprotocol/protocolVersion" + +// metaKeyClientInfo is the reserved _meta key that carries the per-request client +// implementation info on Modern (stateless) MCP requests, per the draft MCP +// schema's RequestMetaObject. +const metaKeyClientInfo = "io.modelcontextprotocol/clientInfo" + +// metaKeyClientCapabilities is the reserved _meta key that carries the per-request +// client capabilities on Modern (stateless) MCP requests, per the draft MCP +// schema's RequestMetaObject. +const metaKeyClientCapabilities = "io.modelcontextprotocol/clientCapabilities" + +// MetaKeyLogLevel carries the per-request minimum log level on Modern +// (2026-07-28) requests (draft schema RequestMetaObject; go-sdk protocol.go). +// It is a reserved per-hop key that must be stripped before a Legacy backend +// hop, but — unlike the other reserved keys — its mere presence is NOT a claim +// of the Modern revision: go-sdk's validateRequestMeta gates Modern-ness purely +// on protocolVersion, and SEP-2577 already deprecates logLevel. It therefore +// belongs in the strip set (ReservedMetaPrefix, minus passthroughMetaKeys) but +// not the ingress signal set (modernSignalMetaKeys). +// +// Exported so the vMCP Modern client can overlay it onto the request _meta it +// mints (the Modern replacement for the removed logging/setLevel RPC). +const MetaKeyLogLevel = "io.modelcontextprotocol/logLevel" + +// ReservedMetaPrefix is the _meta key namespace the MCP spec reserves for the +// protocol's own use. StripReservedMeta removes every key carrying it except +// the passthroughMetaKeys below. +// +// The trailing slash is load-bearing. The registry's +// "io.modelcontextprotocol.registry/publisher-provided" (docs/registry/schema.md) +// is ALSO spec-reserved — per the 2025-11-25 _meta key-name rules any prefix +// whose second label is "modelcontextprotocol" or "mcp" is reserved — but it is +// server-descriptor provenance, not per-hop control, and never rides a +// tools/call result. The slash is what keeps this predicate off it. +// +// Matching one literal prefix deliberately does not implement that full +// second-label rule: "dev.mcp/", "org.modelcontextprotocol.api/" and friends are +// equally reserved and sail through. Nothing in this repo emits them and no +// threat model calls for parsing the general rule. +const ReservedMetaPrefix = "io.modelcontextprotocol/" + +// passthroughMetaKeys are the reserved keys StripReservedMeta must NOT remove. +// +// The reserved namespace holds two different kinds of key, and only one of them +// is a gateway's to terminate: +// +// - per-hop protocol control (protocolVersion, clientInfo, clientCapabilities, +// logLevel, serverInfo, subscriptionId) — scoped to a single MCP connection. +// vMCP is its client's peer and its backends' peer on two different hops, so +// it must strip these and mint its own. +// - end-to-end semantic payload — meaningful to the ORIGINAL endpoints and +// merely carried by intermediaries. Stripping these breaks the feature. +// +// Both entries below are the second kind, from the 2025-11-25 task facility: +// "related-task" is a MUST on every task-related request AND response +// (tasks/result in particular carries the task id nowhere else), and +// "model-immediate-response" rides CreateTaskResult._meta. vMCP does not route +// tasks/* today, so these are inert — but this strip is what would silently +// break them later, and silent broken correlation is expensive to debug. +// +// Extend this by category, not by string: ask whether the key names something +// about THIS hop (strip) or about the two endpoints (pass through). +// A set, not a lookup table: struct{} values make membership the only question +// this map can answer. A map[string]bool would let a future "key": false entry +// read as "in the set" while silently stripping. +var passthroughMetaKeys = map[string]struct{}{ + "io.modelcontextprotocol/related-task": {}, + "io.modelcontextprotocol/model-immediate-response": {}, +} + +// modernSignalMetaKeys is the INGRESS/detection set consumed by hasModernSignal: +// the reserved keys a Legacy client never sets, whose presence — independent of +// whether the value is well-formed — is itself a claim of the Modern revision +// and must not be silently downgraded to Legacy. Only a malformed/absent +// protocolVersion alongside one of them turns into a rejection, never a +// downgrade. +// +// It is deliberately narrower than what StripReservedMeta removes: logLevel is +// excluded so a request carrying only logLevel — which go-sdk's +// validateRequestMeta accepts (gating purely on protocolVersion) and which +// SEP-2577 deprecates — is classified Legacy rather than misdetected as Modern +// and then rejected. "Strip this key on the way out" and "this key means the +// request is Modern" are separate decisions; conflating them would make a +// request rejectable merely for carrying a strippable-but-not-signalling key. +var modernSignalMetaKeys = []string{ + metaKeyProtocolVersion, + metaKeyClientInfo, + metaKeyClientCapabilities, +} + +// StripReservedMeta returns a copy of meta with every ReservedMetaPrefix key +// removed except the passthroughMetaKeys, leaving all other caller-supplied keys +// (progressToken, trace context, backend-custom fields) untouched. The input is +// never mutated (maps.Clone). +// +// Use it on BOTH hops, in both directions, wherever a _meta map crosses vMCP: +// +// - Legacy backend egress (request): a downstream Modern caller's +// _meta.protocolVersion is only valid on a stateless Modern hop. go-sdk v1.7 +// rejects the request outright (HTTP 400: "protocol version ... is only +// supported on stateless HTTP servers") for ANY _meta.protocolVersion on a +// stateful streamable-HTTP server, regardless of its value. +// - Modern backend egress (request): vMCP overlays its own authoritative +// values afterwards — see ModernRequestMeta / mergeModernMeta. +// - client egress (response, and server->client requests): a backend must not +// speak for vMCP. Only serverInfo is even schema-legal on a result +// (ResultMetaObject); the request-only keys arriving on one are a backend +// fabricating the client's own identity. The spec says nothing about what a +// gateway should forward, so this is derived rather than quoted: serverInfo +// identifies "the server software producing the response", vMCP is that +// software, therefore vMCP's value is the correct one. +// +// TRIPWIRE: if vMCP ever relays resource subscriptions it must MINT its own +// io.modelcontextprotocol/subscriptionId from its own downstream listen request +// — never forward a backend's, which names a request id on the vMCP<->backend +// connection and collides across backends. That belongs in the notification +// relay (pkg/vmcp/client/forwarding.go), not here; this helper only sees +// requests and results. See pkg/transport/proxy/streamable/dispatcher_streams.go +// for the same concern in the streamable proxy. +// +// Returns nil whenever the result would be empty -- for nil or empty input, and +// also for a map whose keys were ALL reserved. Callers can therefore treat nil +// as "no _meta to send" without a second length check (mergeModernMeta's +// caller-tolerant convention, and what lets conversion.ToMCPMeta omit _meta +// entirely rather than emit an empty object). A map with keys left over is +// returned as a copy (maps.Clone), never the original. +func StripReservedMeta(meta map[string]any) map[string]any { + if len(meta) == 0 { + return nil + } + stripped := maps.Clone(meta) + maps.DeleteFunc(stripped, func(k string, _ any) bool { + _, passthrough := passthroughMetaKeys[k] + return strings.HasPrefix(k, ReservedMetaPrefix) && !passthrough + }) + if len(stripped) == 0 { + return nil + } + return stripped +} + +// ModernRequestMeta builds the reserved _meta object every Modern (2026-07-28) +// request must carry: protocolVersion, clientInfo, and (empty) clientCapabilities. +// It is the single source of truth for the client side of the reserved +// io.modelcontextprotocol/* keys, kept consistent with what ClassifyRevision and +// ValidateHeaderConsistency require of the server side — protocolVersion equal to +// MCPVersionModern and clientCapabilities present as a JSON object. +func ModernRequestMeta(clientName, clientVersion string) map[string]any { + return map[string]any{ + metaKeyProtocolVersion: MCPVersionModern, + metaKeyClientInfo: map[string]any{"name": clientName, "version": clientVersion}, + metaKeyClientCapabilities: map[string]any{}, + } +} + +// The following JSON-RPC error codes are defined by the draft MCP spec +// (schema/draft/schema.ts) for the stateless "Modern" revision, except for the +// last three, which are standard JSON-RPC 2.0 codes rather than draft-spec- +// specific ones. They are declared as local literals, matching this repo's +// existing convention for JSON-RPC codes (e.g. streamable_proxy.go's -32603). +// The only SDK with equivalents, modelcontextprotocol/go-sdk, is reachable +// solely as a transitive dependency (via toolhive-core's mcpcompat) and is +// pinned to an older draft snapshot (2026-06-30, CodeHeaderMismatch = -32001) +// with no equivalents at all for the other two codes below, so importing it +// would risk wiring in stale values. Revisit once the go-sdk dependency is +// bumped to a revision that matches MCPVersionModern. +// +// These are exported so other packages (e.g. the HTTP layer) can reference +// the same wire values instead of hardcoding or redeclaring them. +const ( + // CodeHeaderMismatch signals a mismatch between the MCP-Protocol-Version + // header and the _meta protocol version (schema.ts HeaderMismatchError). + // Also reused by RequestHeaderMismatchError for the Mcp-Method/Mcp-Name + // headers: keep both error types' Code() in sync with this constant. + CodeHeaderMismatch int64 = -32020 + // CodeMissingClientCapability signals that _meta is missing a client + // capability required for the request (schema.ts MissingRequiredClientCapabilityError). + CodeMissingClientCapability int64 = -32021 + // CodeUnsupportedProtocolVersion signals that the _meta protocol version + // is not one this server supports (schema.ts UnsupportedProtocolVersionError). + CodeUnsupportedProtocolVersion int64 = -32022 + // CodeInvalidParams is the standard JSON-RPC Invalid Params code, used + // as a fallback when the draft spec defines no dedicated code for the failure. + CodeInvalidParams int64 = -32602 + // CodeInvalidRequest is the standard JSON-RPC Invalid Request code. ToolHive + // uses it to reject a message whose shape is not a valid single request for + // the negotiated protocol version — currently a JSON-RPC batch, which was + // removed from MCP in the 2025-06-18 revision (see batch.go) — and also for + // a request body over the configured size limit (HTTP 413, see + // JSONRPCCodeForStatus). + CodeInvalidRequest int64 = -32600 + // CodeInternalError is the standard JSON-RPC Internal Error code, used for + // server-side failures that are not attributable to the request's shape. + CodeInternalError int64 = -32603 +) + +// HeaderMismatchError indicates the MCP-Protocol-Version HTTP header did not match +// the io.modelcontextprotocol/protocolVersion carried in the request's _meta field. +// Neither value wins: this is a hard rejection of the request, not a signal to +// proceed with either value. +type HeaderMismatchError struct { + // Header is the MCP-Protocol-Version header value. + Header string + // Body is the _meta protocol version value. + Body string +} + +func (e *HeaderMismatchError) Error() string { + return fmt.Sprintf("MCP-Protocol-Version header %q does not match _meta protocol version %q", e.Header, e.Body) +} + +// Code implements CodedError. +func (*HeaderMismatchError) Code() int64 { return CodeHeaderMismatch } + +// Data implements CodedError. +func (e *HeaderMismatchError) Data() map[string]any { + return map[string]any{"header": e.Header, "body": e.Body} +} + +// requestHeaderMismatchReason distinguishes why a Mcp-Method/Mcp-Name header +// check failed, so Error() can describe the actual problem instead of always +// reporting a value "mismatch" — the wire code (CodeHeaderMismatch, -32020) +// is the same in every case; only the human-readable message differs. +type requestHeaderMismatchReason int + +const ( + // headerMismatchReasonValue is the default: the header and body both + // carry a value, but they disagree. + headerMismatchReasonValue requestHeaderMismatchReason = iota + // headerMismatchReasonMissing means the header was required but absent. + headerMismatchReasonMissing + // headerMismatchReasonMalformed means the header value itself could not + // be decoded (e.g. invalid base64 in a Mcp-Name sentinel). + headerMismatchReasonMalformed +) + +// RequestHeaderMismatchError indicates a Modern (2026-07-28) request header +// contradicted, was missing from, or was malformed relative to the value +// carried in the request body. This is the generalized sibling of +// HeaderMismatchError for the Mcp-Method/Mcp-Name headers: HeaderMismatchError +// is specific to MCP-Protocol-Version, while this type covers any other +// header/body consistency check and names which header failed. +type RequestHeaderMismatchError struct { + // Header is the name of the header that failed the check (e.g. "Mcp-Method", "Mcp-Name"). + Header string + // HeaderValue is the (decoded, where applicable) value carried by the header. + HeaderValue string + // BodyValue is the value the header was compared against in the request body. + BodyValue string + + // reason is unexported: it only shapes Error()'s message, not Data()'s + // three-key shape or Code()'s wire value. + reason requestHeaderMismatchReason +} + +func (e *RequestHeaderMismatchError) Error() string { + switch e.reason { + case headerMismatchReasonValue: + return fmt.Sprintf("%s header %q does not match request body value %q", e.Header, e.HeaderValue, e.BodyValue) + case headerMismatchReasonMissing: + return fmt.Sprintf("%s header is missing (required for this request)", e.Header) + case headerMismatchReasonMalformed: + return fmt.Sprintf("%s header value %q is malformed", e.Header, e.HeaderValue) + default: + return fmt.Sprintf("%s header %q does not match request body value %q", e.Header, e.HeaderValue, e.BodyValue) + } +} + +// Code implements CodedError. +func (*RequestHeaderMismatchError) Code() int64 { return CodeHeaderMismatch } + +// Data implements CodedError. +func (e *RequestHeaderMismatchError) Data() map[string]any { + return map[string]any{"header": e.Header, "headerValue": e.HeaderValue, "bodyValue": e.BodyValue} +} + +// UnsupportedVersionError indicates the _meta protocol version named a Modern +// revision this server does not support. +type UnsupportedVersionError struct { + // Requested is the _meta protocol version the client asked for. + Requested string + // Supported lists the protocol versions this server supports. + Supported []string +} + +func (e *UnsupportedVersionError) Error() string { + return fmt.Sprintf("unsupported MCP protocol version %q (supported: %v)", e.Requested, e.Supported) +} + +// Code implements CodedError. +func (*UnsupportedVersionError) Code() int64 { return CodeUnsupportedProtocolVersion } + +// Data implements CodedError. +func (e *UnsupportedVersionError) Data() map[string]any { + return map[string]any{"supported": e.Supported, "requested": e.Requested} +} + +// MissingClientCapabilityError indicates a Modern request's _meta is missing +// clientCapabilities (clientInfo is optional per the draft schema and is not +// checked here). +// +// The draft types MissingRequiredClientCapabilityError.data.requiredCapabilities +// as a ClientCapabilities object, not a list of names. The classifier cannot +// compute per-method required capabilities (that check is deferred to the +// caller/handler layer), so RequiredCapabilities is populated best-effort and +// may be an empty object. +type MissingClientCapabilityError struct { + // RequiredCapabilities is the ClientCapabilities object the request was + // missing, if known. + RequiredCapabilities map[string]any +} + +func (e *MissingClientCapabilityError) Error() string { + return fmt.Sprintf("request _meta is missing required client capabilities: %v", e.RequiredCapabilities) +} + +// Code implements CodedError. +func (*MissingClientCapabilityError) Code() int64 { return CodeMissingClientCapability } + +// Data implements CodedError. +func (e *MissingClientCapabilityError) Data() map[string]any { + return map[string]any{"requiredCapabilities": e.RequiredCapabilities} +} + +// MissingModernMetadataError indicates the request carried a Modern signal +// via one of the reserved io.modelcontextprotocol/* _meta keys — with no +// MCP-Protocol-Version header present at all — but _meta carried no valid +// protocolVersion. (When a header IS present, a missing or invalid body +// protocolVersion is a header/body mismatch instead; see HeaderMismatchError.) +// The draft spec defines no dedicated error code for this reserved-key-only +// case, so it falls back to the standard JSON-RPC Invalid Params code. +type MissingModernMetadataError struct{} + +func (*MissingModernMetadataError) Error() string { + return "request carries a Modern signal but no valid io.modelcontextprotocol/protocolVersion in _meta" +} + +// Code implements CodedError. +func (*MissingModernMetadataError) Code() int64 { return CodeInvalidParams } + +// Data implements CodedError. +func (*MissingModernMetadataError) Data() map[string]any { return map[string]any{} } + +// ClassifyRevision determines whether a single MCP request is Legacy or Modern, +// and whether it is valid. +// +// When err != nil, the caller MUST reject the request; the returned Revision is +// then purely informational — it records that the request claimed Modern, not +// that classification succeeded. +// +// method == "initialize" always classifies Legacy immediately, unconditionally — +// Modern never sends initialize (it is the Legacy session-start marker by +// definition) — which also guards against a spoofed Modern _meta on a Legacy +// call. +// +// Otherwise, a request "signals" Modern if the MCP-Protocol-Version header is +// exactly MCPVersionModern, OR _meta carries ANY of the reserved +// io.modelcontextprotocol/* keys (protocolVersion, clientInfo, +// clientCapabilities) — presence of the key is the signal, independent of +// whether its value is well-formed, since a Legacy client never sets these +// keys at all. A request with no signal anywhere classifies Legacy, the safe +// default. A request with a signal is never silently downgraded: it either +// classifies Modern with a nil error, or Modern with an error the caller must +// reject on. +// +// A non-empty MCP-Protocol-Version header that names some OTHER version (not +// MCPVersionModern) and carries no reserved _meta key is, by design, not a +// Modern signal: the request body's _meta is authoritative for the protocol +// version, and an unrecognized header value alone classifies Legacy rather +// than erroring. +// +// Given a Modern signal, checks run in this order: +// 1. meta[metaKeyProtocolVersion] must be a non-empty string. If it is not, +// and protoHeader is non-empty, the body has nothing valid to match against +// the header: this is a header/body mismatch (*HeaderMismatchError). If +// protoHeader is empty — the signal came only from a reserved _meta key, +// so there is no header to mismatch against — the request is malformed +// (*MissingModernMetadataError). +// 2. that string must equal MCPVersionModern, or it names an unsupported +// version (*UnsupportedVersionError). +// 3. if protoHeader is non-empty it must equal the body version, or the two +// conflict (*HeaderMismatchError, a hard rejection — neither value wins). +// 4. _meta must carry clientCapabilities (clientInfo is optional per the +// draft schema), or the client capabilities are missing +// (*MissingClientCapabilityError, with per-method requirements deferred to +// the caller). +// +// A request that passes all four classifies Modern with a nil error. +func ClassifyRevision(method string, meta map[string]any, protoHeader string) (Revision, error) { + if method == "initialize" { + return RevisionLegacy, nil + } + + if !hasModernSignal(meta, protoHeader) { + return RevisionLegacy, nil + } + + bodyVersion, hasBodyVersion := stringMetaValue(meta, metaKeyProtocolVersion) + if !hasBodyVersion { + if protoHeader != "" { + return RevisionModern, &HeaderMismatchError{Header: protoHeader, Body: ""} + } + // TODO: this always returns -32602, but ClassifyRevision is transport-agnostic + // and cannot tell "stdio, no header concept" (where -32602 is correct) apart from + // "HTTP, header omitted" (where the draft's Server Validation rules make a missing + // MCP-Protocol-Version header itself a -32020 HeaderMismatch condition). Revisit + // once the classifier is wired into the HTTP request path and can be given + // transport context. + return RevisionModern, &MissingModernMetadataError{} + } + + if bodyVersion != MCPVersionModern { + return RevisionModern, &UnsupportedVersionError{Requested: bodyVersion, Supported: []string{MCPVersionModern}} + } + + if protoHeader != "" && protoHeader != bodyVersion { + return RevisionModern, &HeaderMismatchError{Header: protoHeader, Body: bodyVersion} + } + + if !hasObjectMetaValue(meta, metaKeyClientCapabilities) { + return RevisionModern, &MissingClientCapabilityError{RequiredCapabilities: map[string]any{}} + } + + return RevisionModern, nil +} + +// nameRequiredMethods are the Modern (2026-07-28) methods for which the draft +// spec's "Server Validation" section requires a Mcp-Name header naming the +// body's tool/resource/prompt identifier. Methods outside this set (e.g. +// tools/list) have no per-request name to check, so Mcp-Name is never +// required for them — though if a client sends one anyway, it is still +// validated for consistency. +var nameRequiredMethods = map[string]bool{ + "tools/call": true, + "resources/read": true, + "prompts/get": true, +} + +// IsNameRequiredMethod reports whether the Modern (2026-07-28) method requires +// an Mcp-Name request header naming the target tool/resource/prompt. It shares +// nameRequiredMethods with ValidateHeaderConsistency so a Modern client sets the +// header for exactly the methods the server validates it on. +func IsNameRequiredMethod(method string) bool { + return nameRequiredMethods[method] +} + +// ValidateHeaderConsistency enforces the Modern (2026-07-28) Mcp-Method and +// Mcp-Name request headers against the corresponding parsed request body +// fields (Method and ResourceID). +// +// Caller contract: only call this for a request ClassifyRevision has already +// classified Modern. Mcp-Method/Mcp-Name are HTTP-only fields — populated +// solely by ParsingMiddleware reading real HTTP headers (see parser.go) — +// so there is no stdio/HTTP transport ambiguity here to excuse a Legacy +// request from this check; the caller must simply not invoke this function +// for Legacy traffic in the first place. +// +// Mcp-Method is required on every Modern request: a missing or empty header +// is itself a rejection (*RequestHeaderMismatchError), not a silent no-op, +// and a present-but-different value is a mismatch. +// +// Mcp-Name is required only for the methods in nameRequiredMethods +// (tools/call, resources/read, prompts/get) — for any other method, an +// absent Mcp-Name header is fine. When present, it is decoded via +// decodeSentinelName before comparison against ResourceID, since the draft +// spec allows it to be sentinel-encoded; a decode failure or mismatch is a +// rejection. +func ValidateHeaderConsistency(parsed *ParsedMCPRequest) error { + if parsed.MCPMethodHeader == "" { + return &RequestHeaderMismatchError{Header: "Mcp-Method", reason: headerMismatchReasonMissing} + } + if parsed.MCPMethodHeader != parsed.Method { + return &RequestHeaderMismatchError{Header: "Mcp-Method", HeaderValue: parsed.MCPMethodHeader, BodyValue: parsed.Method} + } + + if parsed.MCPNameHeader == "" { + if nameRequiredMethods[parsed.Method] { + return &RequestHeaderMismatchError{Header: "Mcp-Name", reason: headerMismatchReasonMissing} + } + return nil + } + + decoded, err := decodeSentinelName(parsed.MCPNameHeader) + if err != nil { + return &RequestHeaderMismatchError{ + Header: "Mcp-Name", + HeaderValue: parsed.MCPNameHeader, + BodyValue: parsed.ResourceID, + reason: headerMismatchReasonMalformed, + } + } + if decoded != parsed.ResourceID { + return &RequestHeaderMismatchError{Header: "Mcp-Name", HeaderValue: decoded, BodyValue: parsed.ResourceID} + } + + return nil +} + +// ExtractMeta pulls the "_meta" object out of raw JSON-RPC request params, for +// use with ClassifyRevision. It is deliberately tolerant: absent params, params +// that don't decode as a JSON object, or a "_meta" value that isn't itself an +// object all yield a nil map rather than an error. Only a well-formed object +// "_meta" is returned. +func ExtractMeta(params json.RawMessage) map[string]any { + if len(params) == 0 { + return nil + } + var paramsMap map[string]any + if err := json.Unmarshal(params, ¶msMap); err != nil { + return nil + } + return metaFromParamsMap(paramsMap) +} + +// metaFromParamsMap reports the "_meta" value of an already-decoded JSON-RPC +// params map, if it decodes as a JSON object. A missing key or a wrong-typed +// value (e.g. a string or number) both yield nil. +func metaFromParamsMap(paramsMap map[string]any) map[string]any { + meta, ok := paramsMap["_meta"].(map[string]any) + if !ok { + return nil + } + return meta +} + +// hasModernSignal reports whether the request signals the Modern revision: +// either the header exactly names MCPVersionModern, or _meta carries any of the +// modernSignalMetaKeys (regardless of whether their values are well-formed). +// It reads modernSignalMetaKeys, NOT everything StripReservedMeta removes — a +// key can be reserved-for-stripping without signalling Modern (logLevel). +func hasModernSignal(meta map[string]any, protoHeader string) bool { + if protoHeader == MCPVersionModern { + return true + } + for _, key := range modernSignalMetaKeys { + if _, ok := meta[key]; ok { + return true + } + } + return false +} + +// stringMetaValue reports the non-empty string value of meta[key], if present. +func stringMetaValue(meta map[string]any, key string) (string, bool) { + raw, ok := meta[key] + if !ok { + return "", false + } + s, ok := raw.(string) + if !ok || s == "" { + return "", false + } + return s, true +} + +// hasObjectMetaValue reports whether meta[key] is present and decodes as a JSON +// object; see objectMetaValue for the shape this covers. +func hasObjectMetaValue(meta map[string]any, key string) bool { + _, ok := objectMetaValue(meta, key) + return ok +} + +// objectMetaValue reports the value of meta[key] if it decodes as a JSON +// object (map[string]any) — the shape clientInfo and clientCapabilities take +// in _meta. A missing key or a wrong-typed value (e.g. a string or number) +// both count as "not present". +func objectMetaValue(meta map[string]any, key string) (map[string]any, bool) { + raw, ok := meta[key] + if !ok { + return nil, false + } + obj, ok := raw.(map[string]any) + if !ok { + return nil, false + } + return obj, true +} + +// sentinelPrefix and sentinelSuffix wrap a base64-encoded Mcp-Name header +// value per the draft MCP spec. This is NOT RFC 2047 encoded-word syntax +// (there is no encoding-letter field); both markers are literal and +// case-sensitive. +const ( + sentinelPrefix = "=?base64?" + sentinelSuffix = "?=" +) + +// decodeSentinelName decodes a Mcp-Name header value that may be wrapped in +// the draft spec's base64 sentinel format (=?base64??=). A value +// that isn't wrapped in the sentinel markers is returned unchanged, since it +// is already the plain name/uri. A wrapped value whose payload fails to +// base64-decode is reported as an error. +func decodeSentinelName(v string) (string, error) { + if !strings.HasPrefix(v, sentinelPrefix) || !strings.HasSuffix(v, sentinelSuffix) { + return v, nil + } + + payload := strings.TrimSuffix(strings.TrimPrefix(v, sentinelPrefix), sentinelSuffix) + decoded, err := base64.StdEncoding.DecodeString(payload) + if err != nil { + return "", fmt.Errorf("decoding base64 sentinel Mcp-Name payload: %w", err) + } + return string(decoded), nil +} + +// EncodeSentinelName encodes v into the draft spec's base64 sentinel format +// (=?base64??=) when required — the mirror of decodeSentinelName. +// Encoding is required when EITHER: +// - v is not safely representable as a plain ASCII header value: any byte +// falls outside printable ASCII 0x21-0x7E, which also covers leading/ +// trailing whitespace and CR/LF (both fall below 0x21 and would +// otherwise make net/http reject the request outright), or +// - v already matches the sentinel pattern (has both sentinelPrefix and +// sentinelSuffix), which must be escaped so the server doesn't mistake a +// literal name for an encoded payload. +// +// Otherwise v is returned unchanged. The result always round-trips through +// decodeSentinelName back to v. +func EncodeSentinelName(v string) string { + if !needsSentinelEncoding(v) { + return v + } + return sentinelPrefix + base64.StdEncoding.EncodeToString([]byte(v)) + sentinelSuffix +} + +// needsSentinelEncoding reports whether v requires sentinel encoding; see +// EncodeSentinelName for the two conditions checked. +func needsSentinelEncoding(v string) bool { + if strings.HasPrefix(v, sentinelPrefix) && strings.HasSuffix(v, sentinelSuffix) { + return true + } + for i := 0; i < len(v); i++ { + if v[i] < 0x21 || v[i] > 0x7E { + return true + } + } + return false +} diff --git a/pkg/mcp/revision_test.go b/pkg/mcp/revision_test.go new file mode 100644 index 0000000000..a70c50839b --- /dev/null +++ b/pkg/mcp/revision_test.go @@ -0,0 +1,763 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/base64" + "encoding/json" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var _ CodedError = (*HeaderMismatchError)(nil) +var _ CodedError = (*RequestHeaderMismatchError)(nil) +var _ CodedError = (*UnsupportedVersionError)(nil) +var _ CodedError = (*MissingClientCapabilityError)(nil) +var _ CodedError = (*MissingModernMetadataError)(nil) + +func validModernMeta() map[string]any { + return map[string]any{ + metaKeyProtocolVersion: MCPVersionModern, + metaKeyClientInfo: map[string]any{"name": "test-client", "version": "1.0.0"}, + metaKeyClientCapabilities: map[string]any{}, + } +} + +func TestClassifyRevision(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + method string + meta map[string]any + protoHeader string + expectedRev Revision + checkErr func(t *testing.T, err error) + }{ + { + name: "modern: valid meta, no header", + method: "tools/call", + meta: validModernMeta(), + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "modern: valid meta, matching header", + method: "tools/call", + meta: validModernMeta(), + protoHeader: MCPVersionModern, + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "modern: mismatched header", + method: "tools/call", + meta: validModernMeta(), + protoHeader: "2025-11-25", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *HeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, CodeHeaderMismatch, mismatchErr.Code()) + assert.Equal(t, "2025-11-25", mismatchErr.Header) + assert.Equal(t, MCPVersionModern, mismatchErr.Body) + assert.Equal(t, map[string]any{"header": "2025-11-25", "body": MCPVersionModern}, mismatchErr.Data()) + }, + }, + { + name: "modern: unsupported future body version", + method: "tools/call", + meta: map[string]any{ + metaKeyProtocolVersion: "2099-01-01", + }, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var unsupportedErr *UnsupportedVersionError + require.ErrorAs(t, err, &unsupportedErr) + assert.Equal(t, CodeUnsupportedProtocolVersion, unsupportedErr.Code()) + data := unsupportedErr.Data() + assert.Equal(t, "2099-01-01", data["requested"]) + assert.Equal(t, []string{MCPVersionModern}, data["supported"]) + }, + }, + { + name: "modern header but meta absent entirely is a header mismatch", + method: "tools/call", + meta: nil, + protoHeader: MCPVersionModern, + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *HeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, CodeHeaderMismatch, mismatchErr.Code()) + assert.Equal(t, MCPVersionModern, mismatchErr.Header) + assert.Empty(t, mismatchErr.Body) + }, + }, + { + name: "modern header but meta missing protocol version key is a header mismatch", + method: "tools/call", + meta: map[string]any{"other": "value"}, + protoHeader: MCPVersionModern, + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + var mismatchErr *HeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + }, + }, + { + name: "modern header but body version wrong-typed is a header mismatch", + method: "tools/call", + meta: map[string]any{metaKeyProtocolVersion: 42}, + protoHeader: MCPVersionModern, + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + var mismatchErr *HeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + }, + }, + { + name: "modern header but body version empty string is a header mismatch", + method: "tools/call", + meta: map[string]any{metaKeyProtocolVersion: ""}, + protoHeader: MCPVersionModern, + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + var mismatchErr *HeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + }, + }, + { + name: "modern: clientInfo omitted is valid", + method: "tools/call", + meta: map[string]any{ + metaKeyProtocolVersion: MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "modern: missing clientCapabilities", + method: "tools/call", + meta: map[string]any{ + metaKeyProtocolVersion: MCPVersionModern, + metaKeyClientInfo: map[string]any{"name": "test-client"}, + }, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var missingCapErr *MissingClientCapabilityError + require.ErrorAs(t, err, &missingCapErr) + assert.Equal(t, CodeMissingClientCapability, missingCapErr.Code()) + }, + }, + { + name: "legacy: absent meta", + method: "tools/call", + meta: nil, + protoHeader: "", + expectedRev: RevisionLegacy, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "legacy: meta missing protocol version key", + method: "tools/call", + meta: map[string]any{"other": "value"}, + protoHeader: "", + expectedRev: RevisionLegacy, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "legacy: unrecognized header version, no reserved meta key", + method: "tools/call", + meta: map[string]any{"other": "value"}, + protoHeader: "2099-01-01", + expectedRev: RevisionLegacy, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + // logLevel is a reserved key that must be STRIPPED on egress but is + // NOT a Modern signal (go-sdk's validateRequestMeta gates purely on + // protocolVersion; SEP-2577 deprecates logLevel). Guards the split + // between what StripReservedMeta removes and modernSignalMetaKeys: if + // hasModernSignal ever iterated the strip set again, this request + // would be misdetected Modern and rejected instead of classified + // Legacy. + name: "legacy: logLevel reserved key alone is not a Modern signal", + method: "tools/call", + meta: map[string]any{MetaKeyLogLevel: "debug"}, + protoHeader: "", + expectedRev: RevisionLegacy, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "modern signal: reserved protocolVersion key wrong-typed", + method: "tools/call", + meta: map[string]any{metaKeyProtocolVersion: 42}, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var missingMetaErr *MissingModernMetadataError + require.ErrorAs(t, err, &missingMetaErr) + assert.Equal(t, CodeInvalidParams, missingMetaErr.Code()) + }, + }, + { + name: "modern signal: reserved protocolVersion key empty string", + method: "tools/call", + meta: map[string]any{metaKeyProtocolVersion: ""}, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var missingMetaErr *MissingModernMetadataError + require.ErrorAs(t, err, &missingMetaErr) + assert.Equal(t, CodeInvalidParams, missingMetaErr.Code()) + }, + }, + { + name: "modern signal via clientCapabilities key, no protocolVersion", + method: "tools/call", + meta: map[string]any{metaKeyClientCapabilities: map[string]any{}}, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var missingMetaErr *MissingModernMetadataError + require.ErrorAs(t, err, &missingMetaErr) + assert.Equal(t, CodeInvalidParams, missingMetaErr.Code()) + }, + }, + { + name: "modern signal via clientInfo key, broken protocolVersion", + method: "tools/call", + meta: map[string]any{ + metaKeyClientInfo: map[string]any{"name": "test-client"}, + metaKeyProtocolVersion: "", + }, + protoHeader: "", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var missingMetaErr *MissingModernMetadataError + require.ErrorAs(t, err, &missingMetaErr) + assert.Equal(t, CodeInvalidParams, missingMetaErr.Code()) + }, + }, + { + name: "modern signal via reserved key with non-modern header is a header mismatch", + method: "tools/call", + meta: map[string]any{ + metaKeyClientCapabilities: map[string]any{}, + }, + protoHeader: "2025-11-25", + expectedRev: RevisionModern, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *HeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, CodeHeaderMismatch, mismatchErr.Code()) + assert.Equal(t, "2025-11-25", mismatchErr.Header) + assert.Empty(t, mismatchErr.Body) + }, + }, + { + name: "legacy: initialize with nil meta", + method: "initialize", + meta: nil, + protoHeader: "", + expectedRev: RevisionLegacy, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "legacy: initialize wins over spoofed modern meta and header", + method: "initialize", + meta: validModernMeta(), + protoHeader: MCPVersionModern, + expectedRev: RevisionLegacy, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + rev, err := ClassifyRevision(tt.method, tt.meta, tt.protoHeader) + + assert.Equal(t, tt.expectedRev, rev) + tt.checkErr(t, err) + }) + } +} + +func TestExtractMeta(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + params json.RawMessage + want map[string]any + }{ + { + name: "well-formed object _meta is returned", + params: json.RawMessage(`{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28"}}`), + want: map[string]any{"io.modelcontextprotocol/protocolVersion": "2026-07-28"}, + }, + { + name: "empty object _meta is returned", + params: json.RawMessage(`{"_meta":{}}`), + want: map[string]any{}, + }, + { + name: "nil params yields nil", + params: nil, + want: nil, + }, + { + name: "empty params yields nil", + params: json.RawMessage(``), + want: nil, + }, + { + name: "params without _meta yields nil", + params: json.RawMessage(`{"other":"value"}`), + want: nil, + }, + { + name: "params as JSON array yields nil", + params: json.RawMessage(`["_meta"]`), + want: nil, + }, + { + name: "params as JSON scalar yields nil", + params: json.RawMessage(`42`), + want: nil, + }, + { + name: "malformed params bytes yield nil", + params: json.RawMessage(`{not json`), + want: nil, + }, + { + name: "_meta as string yields nil", + params: json.RawMessage(`{"_meta":"not-an-object"}`), + want: nil, + }, + { + name: "_meta as number yields nil", + params: json.RawMessage(`{"_meta":42}`), + want: nil, + }, + { + name: "_meta as array yields nil", + params: json.RawMessage(`{"_meta":[1,2,3]}`), + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, ExtractMeta(tt.params)) + }) + } +} + +func sentinelEncode(name string) string { + return "=?base64?" + base64.StdEncoding.EncodeToString([]byte(name)) + "?=" +} + +func TestValidateHeaderConsistency(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + parsed *ParsedMCPRequest + checkErr func(t *testing.T, err error) + }{ + { + name: "missing Mcp-Method header is rejected", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *RequestHeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, CodeHeaderMismatch, mismatchErr.Code()) + assert.Equal(t, "Mcp-Method", mismatchErr.Header) + }, + }, + { + name: "method not in the name-required set needs no Mcp-Name", + parsed: &ParsedMCPRequest{ + Method: "tools/list", + MCPMethodHeader: "tools/list", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "method in the name-required set missing Mcp-Name is rejected", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *RequestHeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, CodeHeaderMismatch, mismatchErr.Code()) + assert.Equal(t, "Mcp-Name", mismatchErr.Header) + }, + }, + { + name: "Mcp-Method matches body method", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + MCPNameHeader: "my-tool", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "Mcp-Method mismatches body method", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "resources/read", + MCPNameHeader: "my-tool", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *RequestHeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, int64(-32020), mismatchErr.Code()) + assert.Equal(t, "Mcp-Method", mismatchErr.Header) + assert.Equal(t, "resources/read", mismatchErr.HeaderValue) + assert.Equal(t, "tools/call", mismatchErr.BodyValue) + }, + }, + { + name: "Mcp-Name plain string matches ResourceID", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + MCPNameHeader: "my-tool", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "Mcp-Name plain string mismatches ResourceID", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + MCPNameHeader: "other-tool", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *RequestHeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, CodeHeaderMismatch, mismatchErr.Code()) + assert.Equal(t, "Mcp-Name", mismatchErr.Header) + assert.Equal(t, "other-tool", mismatchErr.HeaderValue) + assert.Equal(t, "my-tool", mismatchErr.BodyValue) + }, + }, + { + name: "Mcp-Name sentinel-encoded decodes to matching ResourceID", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + MCPNameHeader: sentinelEncode("my-tool"), + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.NoError(t, err) + }, + }, + { + name: "Mcp-Name sentinel-encoded decodes to mismatching value", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + MCPNameHeader: sentinelEncode("other-tool"), + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *RequestHeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, "Mcp-Name", mismatchErr.Header) + assert.Equal(t, "other-tool", mismatchErr.HeaderValue) + assert.Equal(t, "my-tool", mismatchErr.BodyValue) + }, + }, + { + name: "Mcp-Name sentinel wrapper with invalid base64 payload", + parsed: &ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "my-tool", + MCPMethodHeader: "tools/call", + MCPNameHeader: "=?base64?not-valid-base64!!?=", + }, + checkErr: func(t *testing.T, err error) { + t.Helper() + require.Error(t, err) + var mismatchErr *RequestHeaderMismatchError + require.ErrorAs(t, err, &mismatchErr) + assert.Equal(t, "Mcp-Name", mismatchErr.Header) + assert.Equal(t, "=?base64?not-valid-base64!!?=", mismatchErr.HeaderValue) + assert.Equal(t, "my-tool", mismatchErr.BodyValue) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := ValidateHeaderConsistency(tt.parsed) + tt.checkErr(t, err) + }) + } +} + +func TestDecodeSentinelName(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + want string + wantErr bool + }{ + { + name: "non-sentinel value passes through unchanged", + input: "my-tool", + want: "my-tool", + }, + { + name: "empty value passes through unchanged", + input: "", + want: "", + }, + { + name: "valid sentinel decodes the base64 payload", + input: sentinelEncode("my-tool"), + want: "my-tool", + }, + { + name: "sentinel wrapper with invalid base64 payload errors", + input: "=?base64?not-valid-base64!!?=", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := decodeSentinelName(tt.input) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +// TestEncodeSentinelName pins EncodeSentinelName as the exact mirror of +// decodeSentinelName: every case must round-trip back to the original value. +func TestEncodeSentinelName(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + wantSame bool // true if the value must pass through unchanged + }{ + {name: "plain ASCII name unchanged", input: "my-tool", wantSame: true}, + {name: "URI with colon and slash unchanged", input: "file:///tmp/foo.txt", wantSame: true}, + {name: "accented character encoded", input: "café-résumé", wantSame: false}, + {name: "CJK character encoded", input: "工具", wantSame: false}, + {name: "CR in value encoded", input: "bad\rname", wantSame: false}, + {name: "LF in value encoded", input: "bad\nname", wantSame: false}, + {name: "value already shaped like a sentinel is escaped", input: sentinelEncode("my-tool"), wantSame: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := EncodeSentinelName(tt.input) + if tt.wantSame { + assert.Equal(t, tt.input, got) + } else { + assert.NotEqual(t, tt.input, got) + } + + decoded, err := decodeSentinelName(got) + require.NoError(t, err) + assert.Equal(t, tt.input, decoded, "must round-trip through decodeSentinelName") + }) + } +} + +// TestStripReservedMeta pins the copy-before-mutate contract and the two halves +// of the strip predicate: every ReservedMetaPrefix key goes EXCEPT the +// passthroughMetaKeys, and nothing outside the prefix is touched. +func TestStripReservedMeta(t *testing.T) { + t.Parallel() + + t.Run("nil input returns nil", func(t *testing.T) { + t.Parallel() + assert.Nil(t, StripReservedMeta(nil)) + }) + + t.Run("empty input returns nil", func(t *testing.T) { + t.Parallel() + assert.Nil(t, StripReservedMeta(map[string]any{})) + }) + + t.Run("removes reserved keys, preserves the rest", func(t *testing.T) { + t.Parallel() + in := map[string]any{ + // Request-side reserved keys. + metaKeyProtocolVersion: MCPVersionModern, + metaKeyClientInfo: map[string]any{"name": "x"}, + metaKeyClientCapabilities: map[string]any{}, + MetaKeyLogLevel: "debug", + // Response/notification-side reserved keys: a backend must not be able + // to speak for vMCP on the way back to the client (#5986). + "io.modelcontextprotocol/serverInfo": map[string]any{"name": "attacker"}, + "io.modelcontextprotocol/subscriptionId": "sub-1", + // An unknown future reserved key must go too -- that is the whole point + // of matching a namespace rather than a fixed list. + "io.modelcontextprotocol/futureThing": "whatever", + // Not reserved: ordinary caller/backend metadata. + "progressToken": "tok-1", + "traceparent": "00-abc-def-01", + "custom": 42, + // Reserved-adjacent, must NOT match: the second label differs, so the + // trailing slash in ReservedMetaPrefix is what saves it. This is + // registry provenance metadata, not per-hop control. + "io.modelcontextprotocol.registry/publisher-provided": map[string]any{"x": 1}, + } + got := StripReservedMeta(in) + + for k := range in { + if strings.HasPrefix(k, ReservedMetaPrefix) { + assert.NotContains(t, got, k, "reserved key %q must be stripped", k) + } + } + assert.Equal(t, "tok-1", got["progressToken"]) + assert.Equal(t, "00-abc-def-01", got["traceparent"]) + assert.Equal(t, 42, got["custom"]) + assert.Contains(t, got, "io.modelcontextprotocol.registry/publisher-provided", + "the registry namespace is not this predicate's to strip") + }) + + t.Run("passthrough keys survive despite the reserved prefix", func(t *testing.T) { + t.Parallel() + // related-task is a 2025-11-25 MUST on task-related requests AND responses + // (tasks/result carries the task id nowhere else), so the strip must not + // eat it even though it sits under the reserved prefix. + in := map[string]any{ + metaKeyProtocolVersion: MCPVersionModern, + "io.modelcontextprotocol/related-task": map[string]any{"taskId": "t-1"}, + "io.modelcontextprotocol/model-immediate-response": true, + } + got := StripReservedMeta(in) + assert.NotContains(t, got, metaKeyProtocolVersion) + assert.Equal(t, map[string]any{"taskId": "t-1"}, got["io.modelcontextprotocol/related-task"]) + assert.Equal(t, true, got["io.modelcontextprotocol/model-immediate-response"]) + }) + + t.Run("does not mutate the caller's map", func(t *testing.T) { + t.Parallel() + in := map[string]any{metaKeyProtocolVersion: MCPVersionModern, "custom": 1} + _ = StripReservedMeta(in) + assert.Contains(t, in, metaKeyProtocolVersion, "caller's map must be untouched") + assert.Len(t, in, 2) + }) + + t.Run("no reserved keys returns a copy, not the original", func(t *testing.T) { + t.Parallel() + in := map[string]any{"custom": 1} + got := StripReservedMeta(in) + require.Equal(t, in, got) + got["custom"] = 2 + assert.Equal(t, 1, in["custom"], "returned value must be a copy") + }) +} diff --git a/pkg/mcp/tool_filter.go b/pkg/mcp/tool_filter.go index e277bbe788..e3446c45cd 100644 --- a/pkg/mcp/tool_filter.go +++ b/pkg/mcp/tool_filter.go @@ -13,9 +13,17 @@ import ( "net/http" "strings" + "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/transport/types" ) +// filteredToolNotFoundMessage is the generic message returned to a client +// that called a tool blocked by the tool filter. It deliberately matches what +// a client would see for a tool that doesn't exist at all: see the NOTE above +// the *toolCallFilter case in NewToolCallMappingMiddleware for why a filtered +// tool must not be distinguishable from a nonexistent one. +const filteredToolNotFoundMessage = "tool not found" + var errToolNameNotFound = errors.New("tool name not found") var errBug = errors.New("there's a bug") var errKeepBuffering = errors.New("keep buffering") @@ -275,9 +283,22 @@ func NewToolCallMappingMiddleware(opts ...ToolMiddlewareOption) (types.Middlewar // Unfortunately, implementing this behaviour is not trivial and requires // session management, as the SSE stream is managed by the proxy in an entirely // different thread of execution. As a consequence, the best thing we can - // do that is still compliant with the spec is to return a 400 Bad Request - // to the client. + // do that is still compliant with the spec is to return a JSON-RPC error + // for this call, over HTTP 200, that looks the same as calling a tool that + // doesn't exist (see filteredToolNotFoundMessage above). + // + // A client that only accepts an event stream (the legacy HTTP+SSE + // transport, whose real response is delivered on a separate stream we + // don't control from this middleware) can't be answered this way: the + // body written here IS the only response this call gets, and it must be + // something other than a stray JSON blob on a connection the client + // expects to be text/event-stream. So that path keeps the 400, same as + // the session-management limitation described above. case *toolCallFilter: + if clientAcceptsJSON(r) { + writeFilteredToolCallError(w, toolCallRequest.ID) + return + } w.WriteHeader(http.StatusBadRequest) return @@ -298,7 +319,7 @@ func NewToolCallMappingMiddleware(opts ...ToolMiddlewareOption) (types.Middlewar r.ContentLength = int64(len(bodyBytes)) // According to the current version of the MCP spec at - // https://modelcontextprotocol.io/specification/2025-06-18/schema#calltoolrequest + // https://modelcontextprotocol.io/specification/2025-11-25/schema#calltoolrequest // this case can only happen if the request is malformed. The proxied MCP // server should be able to process the request, but since we detect it here // we short-circuit returning an error. @@ -320,6 +341,80 @@ func NewToolCallMappingMiddleware(opts ...ToolMiddlewareOption) (types.Middlewar }, nil } +// clientAcceptsJSON reports whether r's Accept header allows an +// application/json response. An empty/absent header is treated as +// accepting JSON, matching the common case of a client that doesn't bother +// to negotiate. Each comma-separated entry's media-type token (i.e. +// everything before any ";param=..." qualifier) is compared case- +// insensitively against "application/json", the type wildcard +// "application/*", and the full wildcard "*/*". Quality values (";q=") are +// not honored: a token present with any q-value counts as accepted. +func clientAcceptsJSON(r *http.Request) bool { + accept := r.Header.Get("Accept") + if accept == "" { + return true + } + + for _, entry := range strings.Split(accept, ",") { + mediaType := strings.TrimSpace(strings.Split(entry, ";")[0]) + if strings.EqualFold(mediaType, "application/json") || + strings.EqualFold(mediaType, "application/*") || + mediaType == "*/*" { + return true + } + } + + return false +} + +// writeFilteredToolCallError writes a JSON-RPC error response, over HTTP 200, +// for a tools/call request blocked by the tool filter. The error message is +// the generic filteredToolNotFoundMessage: see the NOTE in NewToolCallMappingMiddleware +// for why a filtered tool must look the same as one that doesn't exist. +// +// HTTP 200 (not 400) is deliberate: under MCP streamable HTTP, a validly +// received JSON-RPC request that fails at the application level rides back +// in a 200 response carrying a JSON-RPC error object, mirroring how +// WriteClassificationError's peers (e.g. writeRateLimited) shape their body, +// modulo status code. +func writeFilteredToolCallError(w http.ResponseWriter, id any) { + body := filteredToolCallErrorBody(id) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client + _, _ = w.Write(body) +} + +// filteredToolCallErrorBody renders the JSON-RPC error body for a filtered +// tool call, modeled on classificationErrorBody: the body is marshaled first +// (with a hand-crafted fallback on marshal failure) so the caller only writes +// headers/status once a valid body is ready. +// +// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error +// response as `id?: RequestId` where RequestId = string | number, so an +// absent id is encoded by omitting the "id" key, never as null. See +// session.HasJSONRPCID. +func filteredToolCallErrorBody(id any) []byte { + resp := map[string]any{ + "jsonrpc": "2.0", + "error": map[string]any{ + "code": CodeInvalidParams, + "message": filteredToolNotFoundMessage, + }, + } + if session.HasJSONRPCID(id) { + resp["id"] = id + } + + body, err := json.Marshal(resp) + if err != nil { + // This should never happen with simple map types, but return a + // hand-crafted fallback to guarantee a valid JSON-RPC error. + return []byte(`{"jsonrpc":"2.0","error":{"code":-32602,"message":"tool not found"}}`) + } + return body +} + // toolFilterWriter wraps http.ResponseWriter to capture and process SSE responses type toolFilterWriter struct { http.ResponseWriter @@ -602,7 +697,9 @@ func processEventStream( } else if len(buffer) >= 1 && buffer[len(buffer)-1] == '\r' { linesep = []byte("\r") } else { - return fmt.Errorf("unsupported separator: %s", string(buffer)) + // Length only, not the buffer: buffer is the full pre-filter payload + // (e.g. the unfiltered tools/list), and this error is logged upstream. + return fmt.Errorf("unsupported SSE line separator in %d-byte buffer", len(buffer)) } var linesepTotal, linesepCount int diff --git a/pkg/mcp/tool_filter_test.go b/pkg/mcp/tool_filter_test.go index 58de837dc9..8d967437f7 100644 --- a/pkg/mcp/tool_filter_test.go +++ b/pkg/mcp/tool_filter_test.go @@ -7,8 +7,10 @@ import ( "bytes" "encoding/json" "errors" + "fmt" "net/http" "net/http/httptest" + "strings" "testing" "github.com/stretchr/testify/assert" @@ -1875,3 +1877,255 @@ func TestNewListToolsMappingMiddleware_MissingContentTypeCannotSmuggleToolsList( }) } } + +func TestClientAcceptsJSON(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + accept string + want bool + }{ + {name: "application/json", accept: "application/json", want: true}, + {name: "application/json with other alternatives", accept: "application/json, text/event-stream", want: true}, + {name: "text/event-stream only", accept: "text/event-stream", want: false}, + {name: "wildcard", accept: "*/*", want: true}, + {name: "application subtype wildcard", accept: "application/*", want: true}, + {name: "empty header", accept: "", want: true}, + {name: "application/json with quality parameter", accept: "application/json;q=0.9", want: true}, + {name: "uppercase media type is case-insensitive", accept: "APPLICATION/JSON", want: true}, + {name: "application/json-patch+json is not application/json", accept: "application/json-patch+json", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/", nil) + if tt.accept != "" { + req.Header.Set("Accept", tt.accept) + } + + assert.Equal(t, tt.want, clientAcceptsJSON(req)) + }) + } +} + +// TestNewToolCallMappingMiddleware_FilteredTool covers the HTTP-level +// behavior of the *toolCallFilter case in NewToolCallMappingMiddleware: a +// tools/call request for a tool blocked by the filter must be rejected with +// a JSON-RPC error (never the raw, filter-specific "blocked" reason -- see +// the NOTE above the *toolCallFilter case) rather than a bare HTTP 400, +// unless the client can only receive an event stream, in which case the +// legacy 400 fallback still applies. +func TestNewToolCallMappingMiddleware_FilteredTool(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + accept string + setAccept bool + id any + expectStatus int + expectJSONRPC bool + }{ + { + name: "Accept: application/json - numeric id", + accept: "application/json", + setAccept: true, + id: 1, + expectStatus: http.StatusOK, + expectJSONRPC: true, + }, + { + name: "Accept: application/json - string id", + accept: "application/json", + setAccept: true, + id: "req-42", + expectStatus: http.StatusOK, + expectJSONRPC: true, + }, + { + // Zero is a valid JSON-RPC id and must stay present, not be + // mistaken for absent. + name: "Accept: application/json - zero id is still present", + accept: "application/json", + setAccept: true, + id: 0, + expectStatus: http.StatusOK, + expectJSONRPC: true, + }, + { + name: "Accept: application/json, text/event-stream", + accept: "application/json, text/event-stream", + setAccept: true, + id: 1, + expectStatus: http.StatusOK, + expectJSONRPC: true, + }, + { + name: "Accept: text/event-stream only - falls back to 400", + accept: "text/event-stream", + setAccept: true, + id: 1, + expectStatus: http.StatusBadRequest, + }, + { + name: "No Accept header - JSON allowed by default", + setAccept: false, + id: 1, + expectStatus: http.StatusOK, + expectJSONRPC: true, + }, + { + // MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the + // error response id as optional, not nullable, so a request with + // no usable id gets an error body that omits the "id" key + // entirely rather than echoing "id":null (session.HasJSONRPCID). + name: "null id omits the id key", + accept: "application/json", + setAccept: true, + id: nil, + expectStatus: http.StatusOK, + expectJSONRPC: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + middleware, err := NewToolCallMappingMiddleware(WithToolsFilter("allowed_tool")) + require.NoError(t, err) + + nextCalled := false + next := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + nextCalled = true + }) + handler := middleware(next) + + idJSON, err := json.Marshal(tt.id) + require.NoError(t, err) + body := fmt.Sprintf( + `{"jsonrpc":"2.0","id":%s,"method":"tools/call","params":{"name":"blocked_tool"}}`, + idJSON, + ) + + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + if tt.setAccept { + req.Header.Set("Accept", tt.accept) + } + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.False(t, nextCalled, "next handler must not be called for a filtered tool") + assert.Equal(t, tt.expectStatus, recorder.Code) + + if !tt.expectJSONRPC { + return + } + + assert.Equal(t, "application/json", recorder.Header().Get("Content-Type")) + + var response map[string]any + require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response)) + + assert.Equal(t, "2.0", response["jsonrpc"]) + errObj, ok := response["error"].(map[string]any) + require.True(t, ok, "response must carry a JSON-RPC error object") + assert.Equal(t, float64(CodeInvalidParams), errObj["code"]) + assert.Equal(t, "tool not found", errObj["message"]) + assert.NotContains(t, errObj["message"], "filter", + "a filtered tool must look the same as a nonexistent one") + + // Map decode is the only way to tell an absent "id" key apart + // from a present null -- response["id"] reads back as the same + // nil interface value either way. + _, hasID := response["id"] + if tt.id == nil { + assert.False(t, hasID, `"id" key must be omitted for a nil request id, not present as null`) + return + } + assert.True(t, hasID, `"id" key must be present`) + + expectedID, err := json.Marshal(tt.id) + require.NoError(t, err) + actualID, err := json.Marshal(response["id"]) + require.NoError(t, err) + assert.JSONEq(t, string(expectedID), string(actualID)) + }) + } +} + +// TestNewToolCallMappingMiddleware_AllowedToolPassesThrough verifies that a +// tools/call request for an allowed tool is forwarded unmodified. +func TestNewToolCallMappingMiddleware_AllowedToolPassesThrough(t *testing.T) { + t.Parallel() + + middleware, err := NewToolCallMappingMiddleware(WithToolsFilter("allowed_tool")) + require.NoError(t, err) + + nextCalled := false + next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + }) + handler := middleware(next) + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"allowed_tool"}}` + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set("Accept", "application/json") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.True(t, nextCalled, "next handler must be called for an allowed tool") + assert.Equal(t, http.StatusOK, recorder.Code) +} + +// TestNewToolCallMappingMiddleware_BogusRequest verifies that a malformed +// tools/call request (missing/nil params, or params without a name) is still +// rejected with a bare HTTP 400 -- unlike a filtered tool, this is a +// malformed-request case, not a filtering decision, per the current MCP spec. +func TestNewToolCallMappingMiddleware_BogusRequest(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + }{ + { + name: "params missing name", + body: `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"arguments":{}}}`, + }, + { + name: "nil params", + body: `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":null}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + middleware, err := NewToolCallMappingMiddleware(WithToolsFilter("allowed_tool")) + require.NoError(t, err) + + nextCalled := false + next := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + nextCalled = true + }) + handler := middleware(next) + + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(tt.body)) + req.Header.Set("Accept", "application/json") + recorder := httptest.NewRecorder() + + handler.ServeHTTP(recorder, req) + + assert.False(t, nextCalled, "next handler must not be called for a malformed request") + assert.Equal(t, http.StatusBadRequest, recorder.Code) + }) + } +} diff --git a/pkg/mcp/tool_middleware_test.go b/pkg/mcp/tool_middleware_test.go index fc8298c1b7..161870d2c1 100644 --- a/pkg/mcp/tool_middleware_test.go +++ b/pkg/mcp/tool_middleware_test.go @@ -186,12 +186,11 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { t.Parallel() tests := []struct { - name string - serverOpts []testkit.TestMCPServerOption - opts *[]ToolMiddlewareOption - expected any - callToolName string - expectedStatus int + name string + serverOpts []testkit.TestMCPServerOption + opts *[]ToolMiddlewareOption + expected any + callToolName string }{ { name: "No filter, No override - Call Foo", @@ -201,10 +200,9 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { //nolint:goconst testkit.WithTool("Bar", "Bar tool", func() string { return "Bar" }), }, - opts: nil, - expected: "Foo", - callToolName: "Foo", - expectedStatus: http.StatusOK, + opts: nil, + expected: "Foo", + callToolName: "Foo", }, { name: "Filter Foo, No override - Call Foo", @@ -217,9 +215,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { opts: &[]ToolMiddlewareOption{ WithToolsFilter("Foo"), }, - expected: "Foo", - callToolName: "Foo", - expectedStatus: http.StatusOK, + expected: "Foo", + callToolName: "Foo", }, { name: "Filter Foo, No override - Call Bar", @@ -232,9 +229,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { opts: &[]ToolMiddlewareOption{ WithToolsFilter("Foo"), }, - expected: nil, - callToolName: "Bar", - expectedStatus: http.StatusBadRequest, // Bar is filtered out + expected: nil, + callToolName: "Bar", }, { name: "No filter, Override MyFoo -> Foo - Call MyFoo", @@ -247,9 +243,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { opts: &[]ToolMiddlewareOption{ WithToolsOverride("Foo", "MyFoo", "Override description"), }, - expected: "Foo", - callToolName: "MyFoo", - expectedStatus: http.StatusOK, + expected: "Foo", + callToolName: "MyFoo", }, { name: "No filter, Override MyFoo -> Foo - Call Bar", @@ -262,9 +257,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { opts: &[]ToolMiddlewareOption{ WithToolsOverride("Foo", "MyFoo", "Override description"), }, - expected: "Bar", - callToolName: "Bar", - expectedStatus: http.StatusOK, + expected: "Bar", + callToolName: "Bar", }, { name: "Filter MyFoo, Override MyFoo -> Foo - Call MyFoo", @@ -278,9 +272,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { WithToolsFilter("MyFoo"), WithToolsOverride("Foo", "MyFoo", "Override description"), }, - expected: "Foo", - callToolName: "MyFoo", - expectedStatus: http.StatusOK, + expected: "Foo", + callToolName: "MyFoo", }, { name: "Filter MyFoo, Override MyFoo -> Foo - Call Bar", @@ -294,9 +287,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { WithToolsFilter("MyFoo"), WithToolsOverride("Foo", "MyFoo", "Override description"), }, - expected: nil, - callToolName: "Bar", - expectedStatus: http.StatusBadRequest, // Bar is filtered out + expected: nil, + callToolName: "Bar", }, { name: "No filter, Override Bar -> Foo - Call Foo", @@ -309,9 +301,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { opts: &[]ToolMiddlewareOption{ WithToolsOverride("Bar", "Foo", ""), }, - expected: "Bar", - callToolName: "Foo", - expectedStatus: http.StatusOK, + expected: "Bar", + callToolName: "Foo", }, { name: "No filter, Override Bar -> Foo - Call Bar", @@ -324,9 +315,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { opts: &[]ToolMiddlewareOption{ WithToolsOverride("Foo", "Bar", ""), }, - expected: "Foo", - callToolName: "Bar", - expectedStatus: http.StatusOK, + expected: "Foo", + callToolName: "Bar", }, { name: "Filter MyFoo, Override Foo -> MyFoo", @@ -340,9 +330,8 @@ func TestNewToolCallMappingMiddleware_Scenarios(t *testing.T) { WithToolsFilter("Foo"), WithToolsOverride("Foo", "MyFoo", "Override description"), }, - expected: nil, - callToolName: "MyFoo", - expectedStatus: http.StatusBadRequest, + expected: nil, + callToolName: "MyFoo", }, } diff --git a/pkg/mcp/xmcpheader.go b/pkg/mcp/xmcpheader.go new file mode 100644 index 0000000000..4bcd689537 --- /dev/null +++ b/pkg/mcp/xmcpheader.go @@ -0,0 +1,477 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/json" + "errors" + "fmt" + "math" + "slices" + "strconv" + "strings" +) + +// SEP-2243 ("HTTP standardization") lets an MCP server mark individual tool +// parameters for mirroring into HTTP request headers, via an x-mcp-header +// annotation inside the parameter's schema in the tool's inputSchema. A server +// MAY use the annotation; a client MUST honour it, sending each designated +// parameter's value as the header Mcp-Param-{name} on the tools/call request. A +// server that designated a parameter and did not receive its header rejects the +// call with -32020. +// +// This file owns the annotation's vocabulary and its validation. It is +// deliberately the single place that reads x-mcp-header, so the two consumers -- +// ingestion-time validation of a backend's advertised tools, and call-time +// derivation of the outgoing headers -- cannot drift in their reading of the +// constraints. (Compare the "third independent copy" problem called out for the +// reserved _meta keys in #5986.) +const ( + // XMCPHeaderAnnotation is the JSON Schema extension key a server sets on a + // tool parameter to designate it for header mirroring. Its value is the + // header's name suffix, NOT the full header name. + XMCPHeaderAnnotation = "x-mcp-header" + + // ParamHeaderPrefix prefixes every mirrored parameter header: an annotation + // of "Region" is sent as "Mcp-Param-Region". + ParamHeaderPrefix = "Mcp-Param-" +) + +// maxSchemaDepth bounds the inputSchema traversal. Annotations are legal at any +// nesting depth, so the walk has to recurse, and the schema is attacker-supplied +// from vMCP's perspective (a backend can advertise any tool list it likes). A +// depth cap keeps a hostile or accidentally-recursive schema from exhausting the +// stack. 64 is far past any hand-written tool schema. +const maxSchemaDepth = 64 + +// ErrSchemaTooDeep is returned when an inputSchema nests past maxSchemaDepth. +// It is deliberately distinguishable from a constraint violation: the schema may +// be perfectly valid and merely beyond what this walk will inspect, so a caller +// that wants to fail open on depth alone can single it out. +var ErrSchemaTooDeep = errors.New("tool inputSchema exceeds maximum inspection depth") + +// ParamHeader is one x-mcp-header annotation found in a tool's inputSchema. +type ParamHeader struct { + // Path is the property path from the root of inputSchema to the annotated + // parameter, e.g. ["filter", "region"] for a nested property. Length is + // always >= 1. + Path []string + + // Name is the annotation's value -- the suffix appended to + // ParamHeaderPrefix, not the full header name. Use HeaderName for that. + Name string + + // Type is the annotated parameter's declared JSON Schema type, guaranteed by + // validation to be one of "string", "integer", or "boolean". + Type string +} + +// HeaderName is the full HTTP header name this annotation mirrors into. +func (p ParamHeader) HeaderName() string { + return ParamHeaderPrefix + p.Name +} + +// ParamHeaders walks a tool's inputSchema and returns every x-mcp-header +// annotation it declares, in a deterministic order (by path). A schema with no +// annotations -- the overwhelmingly common case -- returns nil, nil. +// +// It returns an error when the schema violates any of SEP-2243's constraints on +// the annotation: +// +// - the value must be a non-empty string; +// - it must be a valid HTTP field-name token (RFC 9110 tchar), which also +// excludes control characters and whitespace; +// - it must be unique, case-insensitively, within the whole inputSchema +// (HTTP field names are case-insensitive, so two spellings would collide on +// the wire); +// - it may only annotate a parameter whose declared type is "string", +// "integer", or "boolean" -- notably NOT "number", which SEP-2243 excludes +// because a float has no canonical wire spelling. +// +// A missing or non-string "type" on an annotated parameter is treated as a +// violation. SEP-2243 permits the annotation only on primitive parameters, and +// an undeclared type is not a declared primitive; equally, mirroring cannot +// serialize a value whose type it does not know. This is the strict reading: +// it rejects the tool rather than guessing a spelling for the header value. +// +// Traversal covers "properties", array "items", and the "oneOf"/"anyOf"/"allOf" +// combinators, so an annotation is found wherever SEP-2243 allows one. Nesting +// past maxSchemaDepth yields ErrSchemaTooDeep. +func ParamHeaders(schema map[string]any) ([]ParamHeader, error) { + if len(schema) == 0 { + return nil, nil + } + // seen maps the case-folded annotation value to the path that first claimed + // it, so a duplicate can name both colliding parameters in its error. + seen := map[string]string{} + var found []ParamHeader + if err := walkSchema(schema, nil, 0, seen, &found); err != nil { + return nil, err + } + if len(found) == 0 { + return nil, nil + } + // Deterministic order so callers (and tests) see a stable sequence + // regardless of Go's map iteration order. + slices.SortFunc(found, func(a, b ParamHeader) int { + return strings.Compare(strings.Join(a.Path, "."), strings.Join(b.Path, ".")) + }) + return found, nil +} + +// ValidateParamHeaders reports whether a tool's inputSchema declares only +// SEP-2243-conformant x-mcp-header annotations. It is ParamHeaders with the +// annotations discarded, for the ingestion-time check where only the verdict +// matters. +func ValidateParamHeaders(schema map[string]any) error { + _, err := ParamHeaders(schema) + return err +} + +// walkSchema recurses one schema node, appending any annotation it finds to +// found. path is the property path to this node ("" at the root, which cannot +// itself be annotated -- an annotation designates a parameter, and the root is +// the parameter object). +func walkSchema(node map[string]any, path []string, depth int, seen map[string]string, found *[]ParamHeader) error { + if depth > maxSchemaDepth { + return fmt.Errorf("%w (%d)", ErrSchemaTooDeep, maxSchemaDepth) + } + + // An annotation on the root node designates no parameter, so it is ignored + // rather than treated as an error: len(path) == 0 only at the root. + if len(path) > 0 { + if raw, ok := node[XMCPHeaderAnnotation]; ok { + hdr, err := parseAnnotation(raw, node, path, seen) + if err != nil { + return err + } + *found = append(*found, hdr) + } + } + + return walkChildren(node, path, depth, seen, found) +} + +// walkChildren recurses into every sub-schema of node that SEP-2243 allows an +// annotation to appear in: object properties, array element schemas, and the +// oneOf/anyOf/allOf combinator branches. Split from walkSchema to keep each +// within the cyclomatic limit. +func walkChildren(node map[string]any, path []string, depth int, seen map[string]string, found *[]ParamHeader) error { + if props, ok := node["properties"].(map[string]any); ok { + for name, sub := range props { + subSchema, ok := sub.(map[string]any) + if !ok { + continue // not a schema object; nothing to inspect + } + if err := walkSchema(subSchema, childPath(path, name), depth+1, seen, found); err != nil { + return err + } + } + } + + // Array element schemas: "items" is a schema object in the JSON Schema + // dialect MCP uses. The element carries no property name of its own, so it + // inherits the array's path with an "[]" marker for legibility in errors. + if items, ok := node["items"].(map[string]any); ok { + if err := walkSchema(items, childPath(path, "[]"), depth+1, seen, found); err != nil { + return err + } + } + + return walkCombinators(node, path, depth, seen, found) +} + +// walkCombinators recurses into the oneOf/anyOf/allOf branches of node. These +// carry real schemas in this repo's backend tool sets (#5976 fixed ingestion +// dropping them), so an annotation inside a branch must be validated like any +// other. +func walkCombinators( + node map[string]any, path []string, depth int, seen map[string]string, found *[]ParamHeader, +) error { + for _, combinator := range []string{"oneOf", "anyOf", "allOf"} { + branches, ok := node[combinator].([]any) + if !ok { + continue + } + for i, branch := range branches { + branchSchema, ok := branch.(map[string]any) + if !ok { + continue + } + marker := fmt.Sprintf("%s[%d]", combinator, i) + if err := walkSchema(branchSchema, childPath(path, marker), depth+1, seen, found); err != nil { + return err + } + } + } + return nil +} + +// maxSafeInteger is JavaScript's Number.MAX_SAFE_INTEGER (2^53 - 1). SEP-2243 +// requires a mirrored integer to sit within this range, since a peer that parses +// the header with JSON-number semantics could not round-trip a larger value. +const maxSafeInteger = 1<<53 - 1 + +// ErrUnmirrorableValue is returned when a designated parameter's value cannot be +// mirrored into a header: a control character (header injection), a non-integral +// or out-of-safe-range integer, or a value whose type contradicts the schema. +var ErrUnmirrorableValue = errors.New("parameter value cannot be mirrored into an HTTP header") + +// MirrorParamHeaders derives the Mcp-Param-* headers to send with a tools/call, +// given the tool's x-mcp-header annotations (from ParamHeaders) and the call's +// arguments. The result maps full header names to values, ready to set on the +// request; it is nil when nothing is to be mirrored. +// +// A designated parameter that is absent from args contributes no header. That is +// deliberate rather than an error: an optional parameter the caller did not +// supply has no value to mirror, and SEP-2243's -32020 covers the server's view +// of a genuinely missing designated value. +// +// Annotations reached through an array element ("[]" in the path) are skipped: an +// array holds many elements and a header holds one value, so there is no +// well-defined single value to send. Combinator segments (oneOf/anyOf/allOf) are +// dropped when resolving, because they are structural to the schema and absent +// from the arguments the schema describes. +// +// It returns an error wrapping ErrUnmirrorableValue when a present value cannot +// be safely rendered. Values originate from the caller (ultimately a model), so +// they are untrusted: a CR, LF, or NUL in a string would let a caller forge +// additional headers on vMCP's outgoing request, and is refused rather than +// silently stripped. +func MirrorParamHeaders(headers []ParamHeader, args map[string]any) (map[string]string, error) { + if len(headers) == 0 || len(args) == 0 { + return nil, nil + } + var out map[string]string + for _, h := range headers { + if slices.Contains(h.Path, "[]") { + continue + } + value, ok := resolveArg(args, h.Path) + if !ok { + continue + } + rendered, err := renderHeaderValue(h, value) + if err != nil { + return nil, err + } + if out == nil { + out = map[string]string{} + } + out[h.HeaderName()] = rendered + } + return out, nil +} + +// ParamHeadersForSchema is ParamHeaders followed by MirrorParamHeaders: it takes +// a tool's inputSchema and a call's arguments and returns the Mcp-Param-* headers +// to send. It exists so the several call sites that mirror headers share one +// reading of the two-step dance rather than each re-deriving it. +// +// A schema error is returned unwrapped from ParamHeaders; a value error wraps +// ErrUnmirrorableValue. Callers distinguish them because the first indicts the +// backend's tool definition and the second the caller's arguments. +func ParamHeadersForSchema(schema map[string]any, args map[string]any) (map[string]string, error) { + annotations, err := ParamHeaders(schema) + if err != nil { + return nil, err + } + if len(annotations) == 0 { + return nil, nil + } + return MirrorParamHeaders(annotations, args) +} + +// resolveArg walks path through args and returns the value at its end. Segments +// naming a combinator are skipped: they exist in the schema's structure, not in +// the data it describes. A missing or non-object intermediate reports false. +func resolveArg(args map[string]any, path []string) (any, bool) { + current := any(args) + for _, seg := range path { + if isCombinatorSegment(seg) { + continue + } + obj, ok := current.(map[string]any) + if !ok { + return nil, false + } + current, ok = obj[seg] + if !ok { + return nil, false + } + } + return current, true +} + +// isCombinatorSegment reports whether a path segment is a combinator marker +// emitted by walkCombinators (e.g. "oneOf[0]") rather than a property name. +func isCombinatorSegment(seg string) bool { + for _, combinator := range []string{"oneOf", "anyOf", "allOf"} { + if strings.HasPrefix(seg, combinator+"[") && strings.HasSuffix(seg, "]") { + return true + } + } + return false +} + +// renderHeaderValue converts a designated parameter's value to its header +// spelling, enforcing the schema's declared type and SEP-2243's integer range. +func renderHeaderValue(h ParamHeader, value any) (string, error) { + where := strings.Join(h.Path, ".") + switch h.Type { + case "string": + s, ok := value.(string) + if !ok { + return "", fmt.Errorf("%w: parameter %q is declared string but got %T", ErrUnmirrorableValue, where, value) + } + if bad, invalid := firstControlChar(s); invalid { + return "", fmt.Errorf( + "%w: parameter %q contains a control character (%q)", ErrUnmirrorableValue, where, bad) + } + return s, nil + case "boolean": + b, ok := value.(bool) + if !ok { + return "", fmt.Errorf("%w: parameter %q is declared boolean but got %T", ErrUnmirrorableValue, where, value) + } + return strconv.FormatBool(b), nil + case "integer": + return renderIntegerHeaderValue(where, value) + default: + // Unreachable: ParamHeaders admits only the three types above. + return "", fmt.Errorf("%w: parameter %q has unsupported type %q", ErrUnmirrorableValue, where, h.Type) + } +} + +// renderIntegerHeaderValue renders an integer-declared parameter. JSON decoding +// yields float64 for every number, so an integer arrives as a float that must be +// checked for integrality and for SEP-2243's safe-integer range; the int/int64 +// cases cover arguments built in Go rather than decoded from JSON. +func renderIntegerHeaderValue(where string, value any) (string, error) { + switch n := value.(type) { + case float64: + if n != math.Trunc(n) { + return "", fmt.Errorf("%w: parameter %q is declared integer but got %v", ErrUnmirrorableValue, where, n) + } + if n > maxSafeInteger || n < -maxSafeInteger { + return "", fmt.Errorf( + "%w: parameter %q value %v is outside the safe integer range", ErrUnmirrorableValue, where, n) + } + return strconv.FormatInt(int64(n), 10), nil + case int: + return renderIntegerHeaderValue(where, float64(n)) + case int64: + return renderIntegerHeaderValue(where, float64(n)) + case json.Number: + i, err := n.Int64() + if err != nil { + return "", fmt.Errorf("%w: parameter %q is declared integer but got %q", ErrUnmirrorableValue, where, n) + } + return renderIntegerHeaderValue(where, float64(i)) + default: + return "", fmt.Errorf( + "%w: parameter %q is declared integer but got %T", ErrUnmirrorableValue, where, value) + } +} + +// firstControlChar returns the first control character in s, reporting true when +// one exists. CR and LF are the header-injection vectors; NUL and the other C0 +// controls are equally illegal in a header value. +func firstControlChar(s string) (string, bool) { + for _, r := range s { + if r < 0x20 || r == 0x7f { + return string(r), true + } + } + return "", false +} + +// childPath returns path with seg appended, always in freshly allocated storage. +// Appending to path directly would let sibling recursions share (and overwrite) +// one backing array, so a path captured deeper in the walk could be rewritten by +// the next sibling. Allocating per child keeps every captured path independent. +func childPath(path []string, seg string) []string { + child := make([]string, len(path)+1) + copy(child, path) + child[len(path)] = seg + return child +} + +// parseAnnotation validates a single x-mcp-header annotation against SEP-2243 +// and records it in seen for the uniqueness check. +func parseAnnotation(raw any, node map[string]any, path []string, seen map[string]string) (ParamHeader, error) { + where := strings.Join(path, ".") + + name, ok := raw.(string) + if !ok { + return ParamHeader{}, fmt.Errorf( + "parameter %q: %s must be a string, got %T", where, XMCPHeaderAnnotation, raw) + } + if name == "" { + return ParamHeader{}, fmt.Errorf("parameter %q: %s must not be empty", where, XMCPHeaderAnnotation) + } + if bad, invalid := firstNonTokenChar(name); invalid { + return ParamHeader{}, fmt.Errorf( + "parameter %q: %s value %q is not a valid HTTP field-name token (offending character %q)", + where, XMCPHeaderAnnotation, name, bad) + } + + // HTTP field names are case-insensitive, so two annotations differing only in + // case would mirror onto the same header and one would silently win. + folded := strings.ToLower(name) + if first, dup := seen[folded]; dup { + return ParamHeader{}, fmt.Errorf( + "%s value %q on parameter %q collides case-insensitively with parameter %q", + XMCPHeaderAnnotation, name, where, first) + } + seen[folded] = where + + declared, ok := node["type"].(string) + if !ok { + return ParamHeader{}, fmt.Errorf( + "parameter %q: %s requires a declared primitive type (string, integer, or boolean)", + where, XMCPHeaderAnnotation) + } + switch declared { + case "string", "integer", "boolean": + default: + return ParamHeader{}, fmt.Errorf( + "parameter %q: %s is not permitted on type %q (only string, integer, or boolean)", + where, XMCPHeaderAnnotation, declared) + } + + // childPath already gives each branch its own storage, so cloning is belt and + // braces -- it keeps the returned Path independent of the walk even if the + // traversal's path handling is later changed. + return ParamHeader{Path: slices.Clone(path), Name: name, Type: declared}, nil +} + +// firstNonTokenChar returns the first byte of s that is not an RFC 9110 tchar, +// reporting true when one exists. Operating on bytes rather than runes is +// correct here: every tchar is ASCII, so any multi-byte rune is invalid and its +// leading byte is a faithful thing to name in the error. +func firstNonTokenChar(s string) (string, bool) { + for i := 0; i < len(s); i++ { + if !isTokenChar(s[i]) { + return string(s[i]), true + } + } + return "", false +} + +// isTokenChar reports whether c is an RFC 9110 tchar, the character set HTTP +// field names are drawn from. +func isTokenChar(c byte) bool { + switch { + case c >= 'a' && c <= 'z', + c >= 'A' && c <= 'Z', + c >= '0' && c <= '9': + return true + } + switch c { + case '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~': + return true + } + return false +} diff --git a/pkg/mcp/xmcpheader_mirror_test.go b/pkg/mcp/xmcpheader_mirror_test.go new file mode 100644 index 0000000000..5c12d42cb6 --- /dev/null +++ b/pkg/mcp/xmcpheader_mirror_test.go @@ -0,0 +1,256 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/json" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParamHeadersForSchema(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + schema map[string]any + args map[string]any + want map[string]string + }{ + { + name: "no annotations mirrors nothing", + schema: objSchema(map[string]any{"query": map[string]any{"type": "string"}}), + args: map[string]any{"query": "select 1"}, + }, + { + // The SEP's worked example: region is designated, query is not. + name: "designated string is mirrored and others are not", + schema: objSchema(map[string]any{ + "region": annotated("string", "Region"), + "query": map[string]any{"type": "string"}, + }), + args: map[string]any{"region": "eu-west1", "query": "select 1"}, + want: map[string]string{"Mcp-Param-Region": "eu-west1"}, + }, + { + name: "boolean is mirrored as true/false", + schema: objSchema(map[string]any{"dry_run": annotated("boolean", "Dry-Run")}), + args: map[string]any{"dry_run": true}, + want: map[string]string{"Mcp-Param-Dry-Run": "true"}, + }, + { + // JSON decoding yields float64 for every number, so the integral + // float path is the one real clients exercise. + name: "integer arriving as a JSON float is mirrored without a decimal point", + schema: objSchema(map[string]any{"attempts": annotated("integer", "Attempts")}), + args: map[string]any{"attempts": float64(3)}, + want: map[string]string{"Mcp-Param-Attempts": "3"}, + }, + { + name: "negative integer is mirrored", + schema: objSchema(map[string]any{"offset": annotated("integer", "Offset")}), + args: map[string]any{"offset": float64(-7)}, + want: map[string]string{"Mcp-Param-Offset": "-7"}, + }, + { + name: "integer arriving as a Go int is mirrored", + schema: objSchema(map[string]any{"attempts": annotated("integer", "Attempts")}), + args: map[string]any{"attempts": 42}, + want: map[string]string{"Mcp-Param-Attempts": "42"}, + }, + { + name: "integer arriving as json.Number is mirrored", + schema: objSchema(map[string]any{"attempts": annotated("integer", "Attempts")}), + args: map[string]any{"attempts": json.Number("42")}, + want: map[string]string{"Mcp-Param-Attempts": "42"}, + }, + { + // An optional designated parameter the caller did not supply has no + // value to mirror. Omitting the header is correct, not an error. + name: "absent designated parameter contributes no header", + schema: objSchema(map[string]any{"region": annotated("string", "Region")}), + args: map[string]any{"query": "select 1"}, + }, + { + name: "nested designated parameter is resolved through the path", + schema: objSchema(map[string]any{ + "filter": objSchema(map[string]any{"region": annotated("string", "Region")}), + }), + args: map[string]any{"filter": map[string]any{"region": "eu-west1"}}, + want: map[string]string{"Mcp-Param-Region": "eu-west1"}, + }, + { + name: "nested path with a missing intermediate contributes no header", + schema: objSchema(map[string]any{ + "filter": objSchema(map[string]any{"region": annotated("string", "Region")}), + }), + args: map[string]any{"other": "x"}, + }, + { + // Combinator segments are structural to the schema and absent from the + // data, so they are skipped when resolving the value's location. + name: "combinator segments are dropped when resolving", + schema: objSchema(map[string]any{ + "target": map[string]any{ + "oneOf": []any{ + objSchema(map[string]any{"zone": annotated("string", "Zone")}), + }, + }, + }), + args: map[string]any{"target": map[string]any{"zone": "a"}}, + want: map[string]string{"Mcp-Param-Zone": "a"}, + }, + { + // An array holds many elements and a header holds one value, so there + // is no well-defined single value to mirror. + name: "annotation behind an array element is skipped", + schema: objSchema(map[string]any{ + "targets": map[string]any{ + "type": "array", + "items": objSchema(map[string]any{"zone": annotated("string", "Zone")}), + }, + }), + args: map[string]any{"targets": []any{map[string]any{"zone": "a"}}}, + }, + { + name: "two designated parameters both mirror", + schema: objSchema(map[string]any{ + "region": annotated("string", "Region"), + "priority": annotated("string", "Priority"), + }), + args: map[string]any{"region": "eu-west1", "priority": "high"}, + want: map[string]string{ + "Mcp-Param-Region": "eu-west1", + "Mcp-Param-Priority": "high", + }, + }, + { + name: "empty args mirrors nothing", + schema: objSchema(map[string]any{"region": annotated("string", "Region")}), + args: map[string]any{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := ParamHeadersForSchema(tt.schema, tt.args) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestParamHeadersForSchema_UnmirrorableValues(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + schema map[string]any + args map[string]any + wantErrPart string + }{ + { + // The header-injection case. The value comes from the caller + // (ultimately a model), so it is untrusted: a CRLF must be refused, + // not silently stripped, or a caller could forge headers on vMCP's + // outgoing request. + name: "CRLF in a string value is refused", + schema: objSchema(map[string]any{"region": annotated("string", "Region")}), + args: map[string]any{"region": "eu\r\nX-Evil: 1"}, + wantErrPart: "control character", + }, + { + name: "bare newline in a string value is refused", + schema: objSchema(map[string]any{"region": annotated("string", "Region")}), + args: map[string]any{"region": "eu\nwest"}, + wantErrPart: "control character", + }, + { + name: "NUL in a string value is refused", + schema: objSchema(map[string]any{"region": annotated("string", "Region")}), + args: map[string]any{"region": "eu\x00west"}, + wantErrPart: "control character", + }, + { + name: "non-integral value for an integer parameter is refused", + schema: objSchema(map[string]any{"attempts": annotated("integer", "Attempts")}), + args: map[string]any{"attempts": 1.5}, + wantErrPart: "declared integer", + }, + { + // SEP-2243 bounds a mirrored integer to JavaScript's safe range. + name: "integer above the safe range is refused", + schema: objSchema(map[string]any{"attempts": annotated("integer", "Attempts")}), + args: map[string]any{"attempts": float64(maxSafeInteger) * 4}, + wantErrPart: "safe integer range", + }, + { + name: "integer below the safe range is refused", + schema: objSchema(map[string]any{"attempts": annotated("integer", "Attempts")}), + args: map[string]any{"attempts": float64(-maxSafeInteger) * 4}, + wantErrPart: "safe integer range", + }, + { + name: "type mismatch against the schema is refused", + schema: objSchema(map[string]any{"region": annotated("string", "Region")}), + args: map[string]any{"region": 42}, + wantErrPart: "declared string", + }, + { + name: "non-boolean for a boolean parameter is refused", + schema: objSchema(map[string]any{"dry_run": annotated("boolean", "Dry-Run")}), + args: map[string]any{"dry_run": "yes"}, + wantErrPart: "declared boolean", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := ParamHeadersForSchema(tt.schema, tt.args) + require.Error(t, err) + assert.Nil(t, got) + assert.Contains(t, err.Error(), tt.wantErrPart) + assert.True(t, errors.Is(err, ErrUnmirrorableValue), + "a bad VALUE must be distinguishable from a bad annotation: %v", err) + }) + } +} + +// TestParamHeadersForSchema_SafeIntegerBoundary pins the inclusive edge of the +// safe-integer range, the off-by-one a reviewer would reasonably doubt. +func TestParamHeadersForSchema_SafeIntegerBoundary(t *testing.T) { + t.Parallel() + + schema := objSchema(map[string]any{"n": annotated("integer", "N")}) + + got, err := ParamHeadersForSchema(schema, map[string]any{"n": float64(maxSafeInteger)}) + require.NoError(t, err, "MAX_SAFE_INTEGER itself is in range") + assert.Equal(t, map[string]string{"Mcp-Param-N": "9007199254740991"}, got) + + got, err = ParamHeadersForSchema(schema, map[string]any{"n": float64(-maxSafeInteger)}) + require.NoError(t, err, "-MAX_SAFE_INTEGER is in range") + assert.Equal(t, map[string]string{"Mcp-Param-N": "-9007199254740991"}, got) +} + +// TestParamHeadersForSchema_InvalidAnnotationIsNotAValueError keeps the two error +// classes separable: callers blame the caller's arguments for one and the +// backend's tool definition for the other. +func TestParamHeadersForSchema_InvalidAnnotationIsNotAValueError(t *testing.T) { + t.Parallel() + + schema := objSchema(map[string]any{"ratio": annotated("number", "Ratio")}) + + got, err := ParamHeadersForSchema(schema, map[string]any{"ratio": 1.5}) + require.Error(t, err) + assert.Nil(t, got) + assert.False(t, errors.Is(err, ErrUnmirrorableValue), + "a malformed ANNOTATION must not masquerade as a bad value: %v", err) +} diff --git a/pkg/mcp/xmcpheader_test.go b/pkg/mcp/xmcpheader_test.go new file mode 100644 index 0000000000..76505c2065 --- /dev/null +++ b/pkg/mcp/xmcpheader_test.go @@ -0,0 +1,342 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "errors" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// objSchema builds an object inputSchema with the given properties, the shape +// every tool inputSchema takes at its root. +func objSchema(props map[string]any) map[string]any { + return map[string]any{"type": "object", "properties": props} +} + +// annotated builds a leaf parameter schema of the given type carrying an +// x-mcp-header annotation. +func annotated(typ, header string) map[string]any { + return map[string]any{"type": typ, XMCPHeaderAnnotation: header} +} + +func TestParamHeaders_Accepted(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + schema map[string]any + want []ParamHeader + }{ + { + name: "nil schema has no annotations", + schema: nil, + }, + { + name: "empty schema has no annotations", + schema: map[string]any{}, + }, + { + // The overwhelmingly common case: SEP-2243 makes the annotation + // optional for servers, so almost every real tool has none. + name: "unannotated schema has no annotations", + schema: objSchema(map[string]any{ + "query": map[string]any{"type": "string"}, + }), + }, + { + // The SEP's own worked example (execute_sql / Region). + name: "single string annotation", + schema: objSchema(map[string]any{ + "region": annotated("string", "Region"), + "query": map[string]any{"type": "string"}, + }), + want: []ParamHeader{{Path: []string{"region"}, Name: "Region", Type: "string"}}, + }, + { + name: "integer and boolean are permitted primitives", + schema: objSchema(map[string]any{ + "attempts": annotated("integer", "Attempts"), + "dry_run": annotated("boolean", "Dry-Run"), + }), + want: []ParamHeader{ + {Path: []string{"attempts"}, Name: "Attempts", Type: "integer"}, + {Path: []string{"dry_run"}, Name: "Dry-Run", Type: "boolean"}, + }, + }, + { + // SEP-2243: "These annotations can be applied to properties at any + // nesting depth." + name: "annotation nested inside an object property", + schema: objSchema(map[string]any{ + "filter": objSchema(map[string]any{ + "region": annotated("string", "Region"), + }), + }), + want: []ParamHeader{{Path: []string{"filter", "region"}, Name: "Region", Type: "string"}}, + }, + { + name: "annotation inside array items", + schema: objSchema(map[string]any{ + "targets": map[string]any{ + "type": "array", + "items": objSchema(map[string]any{"zone": annotated("string", "Zone")}), + }, + }), + want: []ParamHeader{{Path: []string{"targets", "[]", "zone"}, Name: "Zone", Type: "string"}}, + }, + { + // oneOf/anyOf/allOf carry real schemas in this repo's tool sets -- + // #5976 fixed ingestion dropping them -- so the walk must descend them + // or an annotation inside a combinator branch would go unvalidated. + name: "annotation inside a oneOf branch", + schema: objSchema(map[string]any{ + "target": map[string]any{ + "oneOf": []any{ + map[string]any{"type": "string"}, + objSchema(map[string]any{"zone": annotated("string", "Zone")}), + }, + }, + }), + want: []ParamHeader{{Path: []string{"target", "oneOf[1]", "zone"}, Name: "Zone", Type: "string"}}, + }, + { + // Distinct spellings that do not collide case-insensitively are fine. + name: "two distinct annotations", + schema: objSchema(map[string]any{ + "region": annotated("string", "Region"), + "priority": annotated("string", "Priority"), + }), + want: []ParamHeader{ + {Path: []string{"priority"}, Name: "Priority", Type: "string"}, + {Path: []string{"region"}, Name: "Region", Type: "string"}, + }, + }, + { + // An annotation on the root designates no parameter, so it is ignored + // rather than rejected -- the root IS the parameter object. + name: "annotation on the root node is ignored", + schema: map[string]any{ + "type": "object", + XMCPHeaderAnnotation: "Bogus", + "properties": map[string]any{"query": map[string]any{"type": "string"}}, + }, + }, + { + // All tchar punctuation is legal in an HTTP field name. + name: "tchar punctuation is a valid token", + schema: objSchema(map[string]any{ + "weird": annotated("string", "A!#$%&'*+-.^_`|~9z"), + }), + want: []ParamHeader{{Path: []string{"weird"}, Name: "A!#$%&'*+-.^_`|~9z", Type: "string"}}, + }, + { + // A non-schema value where a schema is expected must be skipped, not + // panicked on: the schema is backend-supplied and need not be sane. + name: "non-object property value is skipped", + schema: objSchema(map[string]any{ + "bogus": "not a schema", + "region": annotated("string", "Region"), + }), + want: []ParamHeader{{Path: []string{"region"}, Name: "Region", Type: "string"}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := ParamHeaders(tt.schema) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + // ValidateParamHeaders is ParamHeaders with the result discarded, so + // it must agree on every accepted schema. + assert.NoError(t, ValidateParamHeaders(tt.schema)) + }) + } +} + +func TestParamHeaders_Rejected(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + schema map[string]any + wantErrPart string + }{ + { + name: "empty annotation value", + schema: objSchema(map[string]any{"region": annotated("string", "")}), + wantErrPart: "must not be empty", + }, + { + name: "non-string annotation value", + schema: objSchema(map[string]any{ + "region": map[string]any{"type": "string", XMCPHeaderAnnotation: 42}, + }), + wantErrPart: "must be a string", + }, + { + // A space is not a tchar; allowing it would let a backend inject a + // second header or a request line into the outgoing request. + name: "annotation containing a space", + schema: objSchema(map[string]any{"region": annotated("string", "My Region")}), + wantErrPart: "not a valid HTTP field-name token", + }, + { + // CRLF is the header-injection vector specifically. + name: "annotation containing CRLF", + schema: objSchema(map[string]any{"region": annotated("string", "R\r\nX: y")}), + wantErrPart: "not a valid HTTP field-name token", + }, + { + name: "annotation containing a colon", + schema: objSchema(map[string]any{"region": annotated("string", "R:egion")}), + wantErrPart: "not a valid HTTP field-name token", + }, + { + // SEP-2243 excludes number explicitly: a float has no canonical wire + // spelling, so mirroring it would be lossy. + name: "number is not a permitted type", + schema: objSchema(map[string]any{"ratio": annotated("number", "Ratio")}), + wantErrPart: `not permitted on type "number"`, + }, + { + name: "object is not a permitted type", + schema: objSchema(map[string]any{"blob": annotated("object", "Blob")}), + wantErrPart: `not permitted on type "object"`, + }, + { + name: "array is not a permitted type", + schema: objSchema(map[string]any{"list": annotated("array", "List")}), + wantErrPart: `not permitted on type "array"`, + }, + { + name: "missing type is not a declared primitive", + schema: objSchema(map[string]any{ + "region": map[string]any{XMCPHeaderAnnotation: "Region"}, + }), + wantErrPart: "requires a declared primitive type", + }, + { + // A union type is not a single declared primitive, so mirroring cannot + // know the spelling. + name: "union type is not a declared primitive", + schema: objSchema(map[string]any{ + "region": map[string]any{ + "type": []any{"string", "null"}, + XMCPHeaderAnnotation: "Region", + }, + }), + wantErrPart: "requires a declared primitive type", + }, + { + name: "exact duplicate annotation", + schema: objSchema(map[string]any{ + "a": annotated("string", "Region"), + "b": annotated("string", "Region"), + }), + wantErrPart: "collides case-insensitively", + }, + { + // HTTP field names are case-insensitive, so these two would mirror onto + // the same header and one would silently win. + name: "case-insensitive duplicate annotation", + schema: objSchema(map[string]any{ + "a": annotated("string", "Region"), + "b": annotated("string", "REGION"), + }), + wantErrPart: "collides case-insensitively", + }, + { + // Uniqueness is scoped to the whole inputSchema, not to one object + // level, so a nested collision must also be caught. + name: "duplicate across nesting levels", + schema: objSchema(map[string]any{ + "region": annotated("string", "Region"), + "filter": objSchema(map[string]any{"r": annotated("string", "region")}), + }), + wantErrPart: "collides case-insensitively", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := ParamHeaders(tt.schema) + require.Error(t, err) + assert.Nil(t, got, "a rejected schema must yield no annotations") + assert.Contains(t, err.Error(), tt.wantErrPart) + // ValidateParamHeaders must agree on every rejected schema too. + assert.Error(t, ValidateParamHeaders(tt.schema)) + }) + } +} + +func TestParamHeaders_DepthLimit(t *testing.T) { + t.Parallel() + + // Nest well past maxSchemaDepth. A backend can advertise any tool list it + // likes, so an unbounded recursive walk is a stack-exhaustion vector. + deep := map[string]any{"type": "string", XMCPHeaderAnnotation: "Deep"} + for range maxSchemaDepth + 10 { + deep = objSchema(map[string]any{"next": deep}) + } + + got, err := ParamHeaders(deep) + require.Error(t, err) + assert.Nil(t, got) + assert.True(t, errors.Is(err, ErrSchemaTooDeep), "want ErrSchemaTooDeep, got %v", err) +} + +func TestParamHeaders_ShallowNestingIsNotRejected(t *testing.T) { + t.Parallel() + + // The depth cap must not reject legitimately nested schemas. Build one just + // inside the limit and confirm the annotation is still found, so a future + // change to maxSchemaDepth that makes the walk too strict fails here. + const depth = maxSchemaDepth - 2 + schema := map[string]any{"type": "string", XMCPHeaderAnnotation: "Deep"} + for range depth { + schema = objSchema(map[string]any{"next": schema}) + } + + got, err := ParamHeaders(schema) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, "Deep", got[0].Name) + assert.Len(t, got[0].Path, depth) +} + +func TestParamHeader_HeaderName(t *testing.T) { + t.Parallel() + + // The annotation value is the header's SUFFIX; the wire name carries the + // Mcp-Param- prefix. Conflating the two is the easy mistake for a caller. + p := ParamHeader{Path: []string{"region"}, Name: "Region", Type: "string"} + assert.Equal(t, "Mcp-Param-Region", p.HeaderName()) + assert.True(t, strings.HasPrefix(p.HeaderName(), ParamHeaderPrefix)) +} + +func TestParamHeaders_SiblingPathsAreIndependent(t *testing.T) { + t.Parallel() + + // Guards childPath: appending to the shared path slice would let one branch's + // captured path be overwritten by the next sibling's traversal, silently + // mirroring a value onto the wrong parameter's header. + schema := objSchema(map[string]any{ + "first": objSchema(map[string]any{"alpha": annotated("string", "Alpha")}), + "second": objSchema(map[string]any{"beta": annotated("string", "Beta")}), + }) + + got, err := ParamHeaders(schema) + require.NoError(t, err) + require.Len(t, got, 2) + assert.Equal(t, []string{"first", "alpha"}, got[0].Path) + assert.Equal(t, []string{"second", "beta"}, got[1].Path) +} diff --git a/pkg/networking/http_client.go b/pkg/networking/http_client.go index 86332c50bb..cef216617b 100644 --- a/pkg/networking/http_client.go +++ b/pkg/networking/http_client.go @@ -176,6 +176,31 @@ func NewHttpClientBuilder() *HttpClientBuilder { } } +// NewHostScopedClientBuilder returns an HttpClientBuilder pre-configured with +// the SSRF-guard policy (CWE-918) appropriate for dialing host. By default the +// returned builder blocks plain HTTP and connections to private/loopback/ +// link-local IP ranges. Both gates are relaxed automatically for loopback +// hosts (development/testing) and when INSECURE_DISABLE_URL_VALIDATION is set. +// +// allowPrivateIPs widens only the private-IP gate — for example an in-cluster +// provider reachable solely over an RFC-1918 address — without enabling plain +// HTTP for non-loopback hosts. insecureAllowHTTP additionally permits +// plain-HTTP for non-loopback hosts and must never be set in production. +// +// The returned builder is not yet built: callers may chain further options +// (e.g. WithTimeout, WithDisableKeepAlives) before calling Build. This is the +// single source of truth for the host-scoped guard policy shared by the +// upstream OAuth2/OIDC providers and the DCR resolver so the two paths cannot +// drift. +func NewHostScopedClientBuilder(host string, allowPrivateIPs, insecureAllowHTTP bool) *HttpClientBuilder { + allowInsecure := IsLocalhost(host) || + insecureAllowHTTP || + strings.EqualFold(os.Getenv("INSECURE_DISABLE_URL_VALIDATION"), "true") + return NewHttpClientBuilder(). + WithInsecureAllowHTTP(allowInsecure). + WithPrivateIPs(allowInsecure || allowPrivateIPs) +} + // WithCABundle sets the CA certificate bundle path func (b *HttpClientBuilder) WithCABundle(path string) *HttpClientBuilder { b.caCertPath = path diff --git a/pkg/networking/http_client_test.go b/pkg/networking/http_client_test.go index 7cbc818aa4..fb6f2c85c7 100644 --- a/pkg/networking/http_client_test.go +++ b/pkg/networking/http_client_test.go @@ -33,6 +33,75 @@ func TestNewHttpClientBuilder(t *testing.T) { assert.False(t, builder.allowPrivate) } +func TestNewHostScopedClientBuilder(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + host string + allowPrivateIPs bool + insecureAllowHTTP bool + wantAllowPrivate bool + wantInsecureHTTP bool + }{ + { + name: "external host, guarded by default", + host: "idp.example.com", + }, + { + name: "external host, allowPrivateIPs widens only the private-IP gate", + host: "idp.example.com", + allowPrivateIPs: true, + wantAllowPrivate: true, + wantInsecureHTTP: false, + }, + { + name: "external host, insecureAllowHTTP widens both gates", + host: "idp.example.com", + insecureAllowHTTP: true, + wantAllowPrivate: true, + wantInsecureHTTP: true, + }, + { + name: "loopback host is exempted from both gates regardless of flags", + host: "127.0.0.1:8443", + wantAllowPrivate: true, + wantInsecureHTTP: true, + }, + { + name: "loopback host stays exempted even with allowPrivateIPs also set", + host: "localhost", + allowPrivateIPs: true, + wantAllowPrivate: true, + wantInsecureHTTP: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + builder := NewHostScopedClientBuilder(tt.host, tt.allowPrivateIPs, tt.insecureAllowHTTP) + + assert.Equal(t, tt.wantAllowPrivate, builder.allowPrivate, "allowPrivate mismatch") + assert.Equal(t, tt.wantInsecureHTTP, builder.insecureAllowHTTP, "insecureAllowHTTP mismatch") + }) + } +} + +// TestNewHostScopedClientBuilder_InsecureDisableURLValidationEnvVar pins that +// the env var widens both gates the same way a loopback host does. Kept as a +// standalone test (not a table case) because t.Setenv is incompatible with +// t.Parallel. +func TestNewHostScopedClientBuilder_InsecureDisableURLValidationEnvVar(t *testing.T) { + t.Setenv("INSECURE_DISABLE_URL_VALIDATION", "true") + + builder := NewHostScopedClientBuilder("idp.example.com", false, false) + + assert.True(t, builder.allowPrivate, "env var must widen the private-IP gate") + assert.True(t, builder.insecureAllowHTTP, "env var must widen the HTTP scheme gate") +} + func TestHttpClientBuilder_WithCABundle(t *testing.T) { t.Parallel() diff --git a/pkg/networking/utilities.go b/pkg/networking/utilities.go index 2f55b83a81..bf2c7aeebc 100644 --- a/pkg/networking/utilities.go +++ b/pkg/networking/utilities.go @@ -282,6 +282,30 @@ func IsLoopbackHost(host string) bool { return ip != nil && ip.IsLoopback() } +// Network isolation log messages. Network isolation is only functional in bridge +// mode; host/none/custom network modes bypass the isolation sidecars entirely, so +// isolation must be dropped for them. These constants are shared across the config +// build, config load, and deploy paths so the wording cannot drift. +const ( + // NetworkIsolationHostDroppedMsg is logged at WARN when isolation was requested + // but is dropped because the network mode (host/custom) cannot enforce it. + // Dropping isolation here is a real reduction in confinement. + NetworkIsolationHostDroppedMsg = "network isolation is not enforced with host/custom networking; " + + "disabling isolation for this workload (pass --isolate-network=false to opt out explicitly and silence this warning)" + // NetworkIsolationNoneRedundantMsg is logged at DEBUG when isolation is dropped + // for the "none" network mode, where it is merely redundant (already maximally confined). + NetworkIsolationNoneRedundantMsg = "network isolation is redundant with none networking; " + + "disabling isolation for this workload" +) + +// IsBridgeMode reports whether the given network mode is a bridge-style mode +// (""/"bridge"/"default"). Network isolation is only functional in bridge mode: +// it builds a per-workload internal network plus DNS/egress/ingress sidecars, and +// host/none/custom modes bypass those sidecars entirely. +func IsBridgeMode(mode string) bool { + return mode == "" || mode == "bridge" || mode == "default" +} + // IsURL checks if the input is a valid HTTP or HTTPS URL func IsURL(input string) bool { parsedURL, err := url.Parse(input) diff --git a/pkg/oauthproto/constants.go b/pkg/oauthproto/constants.go index 0846899e8c..76575b3349 100644 --- a/pkg/oauthproto/constants.go +++ b/pkg/oauthproto/constants.go @@ -41,6 +41,14 @@ const ( // TokenEndpointAuthMethodNone indicates no client authentication (public clients). // Typically used with PKCE for native/mobile applications. TokenEndpointAuthMethodNone = "none" + + // TokenEndpointAuthMethodClientSecretBasic sends the client credentials in the + // HTTP Basic Authorization header at the token endpoint (RFC 6749 Section 2.3.1). + TokenEndpointAuthMethodClientSecretBasic = "client_secret_basic" + + // TokenEndpointAuthMethodClientSecretPost sends the client credentials in the + // request body at the token endpoint (RFC 6749 Section 2.3.1). + TokenEndpointAuthMethodClientSecretPost = "client_secret_post" ) // PKCE (Proof Key for Code Exchange) methods as defined by RFC 7636. diff --git a/pkg/plugins/adapters/claudecode_test.go b/pkg/plugins/adapters/claudecode_test.go index 14b8ae27d0..ee16707fd4 100644 --- a/pkg/plugins/adapters/claudecode_test.go +++ b/pkg/plugins/adapters/claudecode_test.go @@ -18,6 +18,17 @@ import ( "github.com/stacklok/toolhive/pkg/plugins" ) +// resolvedTempDir returns a temp directory path with symlinks resolved, so that +// paths under it pass plugin extraction's "refusing to extract through symlinks" +// check (e.g. on macOS, t.TempDir() may return /var/folders/... where /var is a +// symlink to /private/var). +func resolvedTempDir(t *testing.T) string { + t.Helper() + resolved, err := filepath.EvalSymlinks(t.TempDir()) + require.NoError(t, err) + return resolved +} + // makePluginLayer builds a tar.gz layer from the given file entries, suitable // for passing as MaterializeRequest.LayerData. Mirrors the fixture helper in // pkg/skills/installer_test.go. @@ -98,7 +109,7 @@ func requireMarketplacePlugin(t *testing.T, mp claudeMarketplace, name string) { func TestClaudeCodeAdapter_SupportedComponents(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -113,7 +124,7 @@ func TestClaudeCodeAdapter_SupportedComponents(t *testing.T) { func TestClaudeCodeAdapter_MaterializeWritesFiles(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -158,7 +169,7 @@ func TestClaudeCodeAdapter_MaterializeWritesFiles(t *testing.T) { func TestClaudeCodeAdapter_ExecutableBitPreserved(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -193,7 +204,7 @@ func TestClaudeCodeAdapter_ExecutableBitPreserved(t *testing.T) { func TestClaudeCodeAdapter_DroppedComponents(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -230,7 +241,7 @@ func TestClaudeCodeAdapter_DroppedComponents(t *testing.T) { func TestClaudeCodeAdapter_DematerializeRemovesDir(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -272,7 +283,7 @@ func TestClaudeCodeAdapter_DematerializeRemovesDir(t *testing.T) { func TestClaudeCodeAdapter_DematerializeIdempotent(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -288,7 +299,7 @@ func TestClaudeCodeAdapter_DematerializeIdempotent(t *testing.T) { func TestClaudeCodeAdapter_RematerializeOverwrites(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -342,7 +353,7 @@ func TestClaudeCodeAdapter_ScopeSupport(t *testing.T) { func TestClaudeCodeAdapter_MaterializePreservesUnrelatedSettingsKeys(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -391,7 +402,7 @@ func TestClaudeCodeAdapter_MaterializePreservesUnrelatedSettingsKeys(t *testing. func TestClaudeCodeAdapter_DematerializeRemovesMarketplaceEntryWhenNoToolhivePluginsRemain(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -436,7 +447,7 @@ func TestClaudeCodeAdapter_DematerializeRemovesMarketplaceEntryWhenNoToolhivePlu func TestClaudeCodeAdapter_NonLifoUninstallKeepsMarketplacePathValid(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) @@ -504,11 +515,11 @@ func TestClaudeCodeAdapter_NonLifoUninstallKeepsMarketplacePathValid(t *testing. func TestClaudeCodeAdapter_ProjectScopeUsesProjectSettings(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewClaudeCodeAdapter(cm) - projectRoot := t.TempDir() + projectRoot := resolvedTempDir(t) layer := makePluginLayer(t, []ociskills.FileEntry{ {Path: "commands/greet.md", Content: []byte("# greet"), Mode: 0644}, }) diff --git a/pkg/plugins/adapters/codex_test.go b/pkg/plugins/adapters/codex_test.go index 580a0395bc..de5b3f03b0 100644 --- a/pkg/plugins/adapters/codex_test.go +++ b/pkg/plugins/adapters/codex_test.go @@ -56,7 +56,7 @@ func findCodexPlugin(t *testing.T, mp codexMarketplace, name string) codexMarket func TestCodexAdapter_SupportedComponents(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -70,7 +70,7 @@ func TestCodexAdapter_SupportedComponents(t *testing.T) { func TestCodexAdapter_MaterializeExtractsAndRegisters(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -111,7 +111,7 @@ func TestCodexAdapter_MaterializeExtractsAndRegisters(t *testing.T) { func TestCodexAdapter_DematerializeRemovesPluginAndEntry(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -140,7 +140,7 @@ func TestCodexAdapter_DematerializeRemovesPluginAndEntry(t *testing.T) { func TestCodexAdapter_DroppedComponentsAllSix(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -178,7 +178,7 @@ func TestCodexAdapter_DroppedComponentsAllSix(t *testing.T) { func TestCodexAdapter_ProjectScopeUsesProjectMarketplace(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -188,7 +188,7 @@ func TestCodexAdapter_ProjectScopeUsesProjectMarketplace(t *testing.T) { // Project scope writes to /.agents/plugins, a Codex discovery // path, so it does NOT degrade. - projectRoot := t.TempDir() + projectRoot := resolvedTempDir(t) res, err := a.Materialize(context.Background(), plugins.MaterializeRequest{ Name: "proj-plugin", LayerData: layer, @@ -214,7 +214,7 @@ func TestCodexAdapter_ProjectScopeUsesProjectMarketplace(t *testing.T) { func TestCodexAdapter_DematerializeIdempotent(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -240,7 +240,7 @@ func TestCodexAdapter_ScopeSupport(t *testing.T) { // second deletes it. func TestCodexAdapter_MarketplaceSurvivesWhenOtherPluginRemains(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) @@ -270,7 +270,7 @@ func TestCodexAdapter_MarketplaceSurvivesWhenOtherPluginRemains(t *testing.T) { // with its relative source and alpha's dir is intact. func TestCodexAdapter_NonLifoUninstallKeepsMarketplaceValid(t *testing.T) { t.Parallel() - tempHome := t.TempDir() + tempHome := resolvedTempDir(t) cm := newTestClientManager(t, tempHome) a := NewCodexAdapter(cm) diff --git a/pkg/plugins/client/client.go b/pkg/plugins/client/client.go new file mode 100644 index 0000000000..cb75e1c04c --- /dev/null +++ b/pkg/plugins/client/client.go @@ -0,0 +1,402 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package client provides an HTTP client for the ToolHive Plugins API. +package client + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/url" + "strings" + "time" + + "github.com/stacklok/toolhive-core/env" + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/plugins" + "github.com/stacklok/toolhive/pkg/server/discovery" +) + +const ( + pluginsBasePath = "/api/v1beta/plugins" + defaultBaseURL = "http://127.0.0.1:8080" + // defaultTimeout is a backstop against a wedged server, not a budget for + // the operation. Install, build, and push move OCI artifacts of unbounded + // size over the user's connection, so a timeout tight enough to feel + // responsive would abort legitimate work — and because giving up cancels + // the request context, it aborts it server-side too, leaving nothing for + // a retry to reuse. Callers who want to fail fast can set + // TOOLHIVE_API_TIMEOUT or pass WithTimeout. + defaultTimeout = 10 * time.Minute + envAPIURL = "TOOLHIVE_API_URL" + envAPITimeout = "TOOLHIVE_API_TIMEOUT" + maxResponseSize = 1 << 20 // 1 MiB — defensive response cap; consistent with the skills client + maxErrorBodySize = 1 << 16 // 64 KiB — matches auth/token and DCR limits +) + +// ErrServerUnreachable is returned when the client cannot connect to the +// ToolHive API server. The most common cause is that "thv serve" is not +// running. +var ErrServerUnreachable = errors.New("could not reach ToolHive API server — is 'thv serve' running?") + +// ErrRequestTimeout is returned when the server was reached but did not +// respond within the client timeout. This is a distinct condition from +// ErrServerUnreachable: the server is healthy and was working on the request. +// Note that abandoning the request cancels its context server-side, so the +// operation does not continue in the background and retrying restarts it. +var ErrRequestTimeout = errors.New("the ToolHive API server did not respond within the client timeout") + +// Compile-time interface check. +var _ plugins.PluginService = (*Client)(nil) + +// Client is an HTTP client for the ToolHive Plugins API. +type Client struct { + baseURL string + httpClient *http.Client +} + +// Option configures a Client. +type Option func(*Client) + +// WithTimeout sets the HTTP client timeout. +func WithTimeout(d time.Duration) Option { + return func(c *Client) { + c.httpClient.Timeout = d + } +} + +// WithHTTPClient replaces the underlying *http.Client entirely. +// This overrides any previously applied options such as WithTimeout. +func WithHTTPClient(hc *http.Client) Option { + return func(c *Client) { + c.httpClient = hc + } +} + +// NewClient creates a new Plugins API client with the given base URL. +func NewClient(baseURL string, opts ...Option) *Client { + c := &Client{ + baseURL: strings.TrimRight(baseURL, "/"), + httpClient: &http.Client{Timeout: defaultTimeout}, + } + for _, o := range opts { + o(c) + } + return c +} + +// NewDefaultClient creates a Plugins API client by trying, in order: +// 1. The TOOLHIVE_API_URL environment variable (explicit override) +// 2. The server discovery file (auto-detected running server) +// 3. The default URL http://127.0.0.1:8080 +// +// The context is used for the server discovery health check; it is not stored. +func NewDefaultClient(ctx context.Context, opts ...Option) *Client { + return newDefaultClientWithEnv(ctx, &env.OSReader{}, resolveViaDiscovery, opts...) +} + +// discoverFunc resolves a running server's base URL and any transport options +// (e.g. a Unix socket client). It returns an empty base URL when no running +// server is found. resolveViaDiscovery is the production implementation; tests +// inject a stub so the discovery step does not read real local state. +type discoverFunc func(ctx context.Context) (string, []Option) + +// newDefaultClientWithEnv is the testable core of NewDefaultClient. The +// envReader and discover dependencies are injected so each resolution step +// can be exercised in isolation. +func newDefaultClientWithEnv(ctx context.Context, envReader env.Reader, discover discoverFunc, opts ...Option) *Client { + // An env-supplied timeout applies to every resolution path, but ranks + // below caller-supplied options so a WithTimeout argument still wins. + withEnvTimeout := func(base []Option) []Option { + merged := make([]Option, 0, len(base)+len(opts)+1) + merged = append(merged, base...) + if d, ok := timeoutFromEnv(envReader); ok { + merged = append(merged, WithTimeout(d)) + } + return append(merged, opts...) + } + + // 1. Explicit env var override always wins. + if base := envReader.Getenv(envAPIURL); base != "" { + return NewClient(base, withEnvTimeout(nil)...) + } + + // 2. Try server discovery. + if base, httpOpts := discover(ctx); base != "" { + // Discovery opts go first so caller-supplied opts can override them + // (e.g. a caller-provided WithTimeout replaces the discovery default). + return NewClient(base, withEnvTimeout(httpOpts)...) + } + + // 3. Fall back to the default URL. + return NewClient(defaultBaseURL, withEnvTimeout(nil)...) +} + +// timeoutFromEnv reads TOOLHIVE_API_TIMEOUT as a Go duration (e.g. "45s", +// "5m"). An unset, unparsable, or non-positive value is ignored so a typo +// degrades to the default rather than disabling the timeout entirely. +func timeoutFromEnv(envReader env.Reader) (time.Duration, bool) { + raw := strings.TrimSpace(envReader.Getenv(envAPITimeout)) + if raw == "" { + return 0, false + } + d, err := time.ParseDuration(raw) + if err != nil || d <= 0 { + slog.Warn("ignoring invalid "+envAPITimeout, "value", raw, "error", err) + return 0, false + } + return d, true +} + +// resolveViaDiscovery attempts to find a running server via the discovery file. +// It returns the base URL and any additional options (e.g. a Unix socket transport). +// On failure it returns empty values and the caller falls back to the default. +func resolveViaDiscovery(ctx context.Context) (string, []Option) { + result, err := discovery.Discover(ctx) + if err != nil { + slog.Debug("server discovery failed", "error", err) + return "", nil + } + if result.State != discovery.StateRunning { + return "", nil + } + + client, baseURL, err := discovery.HTTPClientForURL(result.Info.URL) + if err != nil { + slog.Debug("invalid URL in discovery file", "url", result.Info.URL, "error", err) + return "", nil + } + client.Timeout = defaultTimeout + + return baseURL, []Option{WithHTTPClient(client)} +} + +// --- PluginService implementation --- + +// List returns all installed plugins matching the given options. +func (c *Client) List(ctx context.Context, opts plugins.ListOptions) ([]plugins.InstalledPlugin, error) { + q := url.Values{} + if opts.Scope != "" { + q.Set("scope", string(opts.Scope)) + } + if opts.ClientApp != "" { + q.Set("client", opts.ClientApp) + } + if opts.ProjectRoot != "" { + q.Set("project_root", opts.ProjectRoot) + } + if opts.Group != "" { + q.Set("group", opts.Group) + } + + var resp listResponse + if err := c.doJSONRequest(ctx, http.MethodGet, "", q, nil, &resp); err != nil { + return nil, err + } + return resp.Plugins, nil +} + +// Install installs a plugin from a remote source. +func (c *Client) Install(ctx context.Context, opts plugins.InstallOptions) (*plugins.InstallResult, error) { + body := installRequest{ + Name: opts.Name, + Version: opts.Version, + Scope: opts.Scope, + ProjectRoot: opts.ProjectRoot, + Clients: opts.Clients, + Force: opts.Force, + Group: opts.Group, + } + + var resp installResponse + if err := c.doJSONRequest(ctx, http.MethodPost, "", nil, body, &resp); err != nil { + return nil, err + } + return &plugins.InstallResult{Plugin: resp.Plugin}, nil +} + +// Uninstall removes an installed plugin. +func (c *Client) Uninstall(ctx context.Context, opts plugins.UninstallOptions) error { + q := url.Values{} + if opts.Scope != "" { + q.Set("scope", string(opts.Scope)) + } + if opts.ProjectRoot != "" { + q.Set("project_root", opts.ProjectRoot) + } + + path := "/" + url.PathEscape(opts.Name) + return c.doJSONRequest(ctx, http.MethodDelete, path, q, nil, nil) +} + +// Info returns detailed information about a plugin. +func (c *Client) Info(ctx context.Context, opts plugins.InfoOptions) (*plugins.PluginInfo, error) { + q := url.Values{} + if opts.Scope != "" { + q.Set("scope", string(opts.Scope)) + } + if opts.ProjectRoot != "" { + q.Set("project_root", opts.ProjectRoot) + } + + path := "/" + url.PathEscape(opts.Name) + var info plugins.PluginInfo + if err := c.doJSONRequest(ctx, http.MethodGet, path, q, nil, &info); err != nil { + return nil, err + } + return &info, nil +} + +// Validate checks whether a plugin definition is valid. +func (c *Client) Validate(ctx context.Context, path string) (*plugins.ValidationResult, error) { + body := validateRequest{Path: path} + + var result plugins.ValidationResult + if err := c.doJSONRequest(ctx, http.MethodPost, "/validate", nil, body, &result); err != nil { + return nil, err + } + return &result, nil +} + +// Build builds a plugin from a local directory into an OCI artifact. +func (c *Client) Build(ctx context.Context, opts plugins.BuildOptions) (*plugins.BuildResult, error) { + body := buildRequest{ + Path: opts.Path, + Tag: opts.Tag, + } + + var result plugins.BuildResult + if err := c.doJSONRequest(ctx, http.MethodPost, "/build", nil, body, &result); err != nil { + return nil, err + } + return &result, nil +} + +// Push pushes a built plugin artifact to a remote registry. +func (c *Client) Push(ctx context.Context, opts plugins.PushOptions) error { + body := pushRequest{Reference: opts.Reference} + return c.doJSONRequest(ctx, http.MethodPost, "/push", nil, body, nil) +} + +// ListBuilds returns all locally-built OCI plugin artifacts in the local store. +func (c *Client) ListBuilds(ctx context.Context) ([]plugins.LocalBuild, error) { + var resp listBuildsResponse + if err := c.doJSONRequest(ctx, http.MethodGet, "/builds", nil, nil, &resp); err != nil { + return nil, err + } + return resp.Builds, nil +} + +// DeleteBuild removes a locally-built OCI plugin artifact from the local store. +func (c *Client) DeleteBuild(ctx context.Context, tag string) error { + return c.doJSONRequest(ctx, http.MethodDelete, "/builds/"+url.PathEscape(tag), nil, nil, nil) +} + +// GetContent retrieves the plugin.json body and file listing from an OCI artifact without installing it. +func (c *Client) GetContent(ctx context.Context, opts plugins.ContentOptions) (*plugins.PluginContent, error) { + q := url.Values{} + q.Set("ref", opts.Reference) + var content plugins.PluginContent + if err := c.doJSONRequest(ctx, http.MethodGet, "/content", q, nil, &content); err != nil { + return nil, err + } + return &content, nil +} + +// --- internal helpers --- + +func (c *Client) buildURL(path string, query url.Values) string { + u := c.baseURL + pluginsBasePath + path + if len(query) > 0 { + u += "?" + query.Encode() + } + return u +} + +// doJSONRequest performs the full HTTP request lifecycle: marshal body, build +// URL, create request with context, set headers, execute, check status, and +// decode response or return *APIError. +func (c *Client) doJSONRequest( + ctx context.Context, + method, path string, + query url.Values, + reqBody any, + result any, +) error { + var bodyReader io.Reader + if reqBody != nil { + data, err := json.Marshal(reqBody) + if err != nil { + return fmt.Errorf("marshaling request body: %w", err) + } + bodyReader = bytes.NewReader(data) + } + + reqURL := c.buildURL(path, query) + + req, err := http.NewRequestWithContext(ctx, method, reqURL, bodyReader) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + + if reqBody != nil { + req.Header.Set("Content-Type", "application/json") + } + req.Header.Set("Accept", "application/json") + + resp, err := c.httpClient.Do(req) // #nosec G704 -- baseURL is a trusted local API server URL + if err != nil { + return fmt.Errorf("%w: %w", classifyTransportError(ctx, err), err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= http.StatusBadRequest { + return handleErrorResponse(resp) + } + + if result != nil { + limited := io.LimitReader(resp.Body, maxResponseSize) + if err := json.NewDecoder(limited).Decode(result); err != nil { + // Drain before close so the connection can be reused. + _, _ = io.Copy(io.Discard, limited) + return fmt.Errorf("decoding response: %w", err) + } + } + + return nil +} + +// classifyTransportError decides which sentinel a failed round-trip belongs +// to. Reporting a timeout as an unreachable server sends users to check +// whether "thv serve" is running when it is running and was mid-operation. +func classifyTransportError(ctx context.Context, err error) error { + // The caller's own context ending takes precedence — neither the server's + // availability nor the client timeout is why the request stopped. + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return ErrRequestTimeout + } + if errors.Is(err, context.DeadlineExceeded) { + return ErrRequestTimeout + } + return ErrServerUnreachable +} + +// handleErrorResponse reads the response body and returns an *httperr.CodedError. +func handleErrorResponse(resp *http.Response) error { + body, err := io.ReadAll(io.LimitReader(resp.Body, maxErrorBodySize)) + if err != nil { + return httperr.New("failed to read error response body", resp.StatusCode) + } + return httperr.New(strings.TrimSpace(string(body)), resp.StatusCode) +} diff --git a/pkg/plugins/client/client_test.go b/pkg/plugins/client/client_test.go new file mode 100644 index 0000000000..e87a11f524 --- /dev/null +++ b/pkg/plugins/client/client_test.go @@ -0,0 +1,1001 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + envmocks "github.com/stacklok/toolhive-core/env/mocks" + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/plugins" +) + +// newTestClient returns a *Client pointed at the given test server. +func newTestClient(t *testing.T, srv *httptest.Server) *Client { + t.Helper() + return NewClient(srv.URL) +} + +func TestList(t *testing.T) { + t.Parallel() + + now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + opts plugins.ListOptions + wantQuery map[string]string + response listResponse + statusCode int + wantErr bool + }{ + { + name: "no filters", + opts: plugins.ListOptions{}, + response: listResponse{Plugins: []plugins.InstalledPlugin{ + { + Metadata: plugins.PluginMetadata{Name: "my-plugin", Version: "1.0.0"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + InstalledAt: now, + }, + }}, + statusCode: http.StatusOK, + }, + { + name: "with all filters", + opts: plugins.ListOptions{ + Scope: plugins.ScopeProject, + ClientApp: "claude-code", + ProjectRoot: "/home/user/proj", + Group: "my-group", + }, + wantQuery: map[string]string{ + "scope": "project", + "client": "claude-code", + "project_root": "/home/user/proj", + "group": "my-group", + }, + response: listResponse{Plugins: []plugins.InstalledPlugin{}}, + statusCode: http.StatusOK, + }, + { + name: "server error", + opts: plugins.ListOptions{}, + statusCode: http.StatusInternalServerError, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodGet, r.Method) + assert.Equal(t, pluginsBasePath, r.URL.Path) + + for k, v := range tt.wantQuery { + assert.Equal(t, v, r.URL.Query().Get(k), "query param %s", k) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "something went wrong", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.List(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response.Plugins, got) + }) + } +} + +func TestInstall(t *testing.T) { + t.Parallel() + + now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + opts plugins.InstallOptions + wantBody installRequest + response installResponse + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success", + opts: plugins.InstallOptions{ + Name: "my-plugin", + Version: "1.0.0", + Scope: plugins.ScopeUser, + Clients: []string{"claude-code"}, + Force: true, + }, + wantBody: installRequest{ + Name: "my-plugin", + Version: "1.0.0", + Scope: plugins.ScopeUser, + Clients: []string{"claude-code"}, + Force: true, + }, + response: installResponse{Plugin: plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin", Version: "1.0.0"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + InstalledAt: now, + }}, + statusCode: http.StatusCreated, + }, + { + name: "bad request", + opts: plugins.InstallOptions{Name: ""}, + statusCode: http.StatusBadRequest, + wantErr: true, + wantCode: http.StatusBadRequest, + }, + { + name: "conflict", + opts: plugins.InstallOptions{Name: "existing-plugin"}, + statusCode: http.StatusConflict, + wantErr: true, + wantCode: http.StatusConflict, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, pluginsBasePath, r.URL.Path) + + if tt.wantBody.Name != "" { + var got installRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&got)) + assert.Equal(t, tt.wantBody, got) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "error", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(tt.statusCode) + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.Install(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response.Plugin, got.Plugin) + }) + } +} + +func TestUninstall(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts plugins.UninstallOptions + wantPath string + wantQuery map[string]string + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success", + opts: plugins.UninstallOptions{Name: "my-plugin"}, + wantPath: pluginsBasePath + "/my-plugin", + statusCode: http.StatusNoContent, + }, + { + name: "with scope and project root", + opts: plugins.UninstallOptions{ + Name: "my-plugin", + Scope: plugins.ScopeProject, + ProjectRoot: "/home/user/proj", + }, + wantPath: pluginsBasePath + "/my-plugin", + wantQuery: map[string]string{ + "scope": "project", + "project_root": "/home/user/proj", + }, + statusCode: http.StatusNoContent, + }, + { + name: "not found", + opts: plugins.UninstallOptions{Name: "missing"}, + wantPath: pluginsBasePath + "/missing", + statusCode: http.StatusNotFound, + wantErr: true, + wantCode: http.StatusNotFound, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodDelete, r.Method) + assert.Equal(t, tt.wantPath, r.URL.Path) + + for k, v := range tt.wantQuery { + assert.Equal(t, v, r.URL.Query().Get(k), "query param %s", k) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "not found", tt.statusCode) + return + } + w.WriteHeader(tt.statusCode) + })) + defer srv.Close() + + c := newTestClient(t, srv) + err := c.Uninstall(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + }) + } +} + +func TestInfo(t *testing.T) { + t.Parallel() + + now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + opts plugins.InfoOptions + wantPath string + response plugins.PluginInfo + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success", + opts: plugins.InfoOptions{Name: "my-plugin"}, + wantPath: pluginsBasePath + "/my-plugin", + response: plugins.PluginInfo{ + Metadata: plugins.PluginMetadata{Name: "my-plugin", Version: "1.0.0"}, + InstalledPlugin: &plugins.InstalledPlugin{ + Metadata: plugins.PluginMetadata{Name: "my-plugin", Version: "1.0.0"}, + Scope: plugins.ScopeUser, + Status: plugins.InstallStatusInstalled, + InstalledAt: now, + }, + }, + statusCode: http.StatusOK, + }, + { + name: "not found", + opts: plugins.InfoOptions{Name: "missing"}, + wantPath: pluginsBasePath + "/missing", + statusCode: http.StatusNotFound, + wantErr: true, + wantCode: http.StatusNotFound, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodGet, r.Method) + assert.Equal(t, tt.wantPath, r.URL.Path) + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "not found", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.Info(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response, *got) + }) + } +} + +func TestValidate(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + path string + wantBody validateRequest + response plugins.ValidationResult + statusCode int + wantErr bool + }{ + { + name: "valid plugin", + path: "/home/user/my-plugin", + wantBody: validateRequest{Path: "/home/user/my-plugin"}, + response: plugins.ValidationResult{Valid: true}, + statusCode: http.StatusOK, + }, + { + name: "invalid plugin", + path: "/home/user/bad-plugin", + wantBody: validateRequest{Path: "/home/user/bad-plugin"}, + response: plugins.ValidationResult{ + Valid: false, + Errors: []string{"missing name field"}, + }, + statusCode: http.StatusOK, + }, + { + name: "bad request", + path: "", + statusCode: http.StatusBadRequest, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, pluginsBasePath+"/validate", r.URL.Path) + + if tt.wantBody.Path != "" { + var got validateRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&got)) + assert.Equal(t, tt.wantBody, got) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "bad request", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.Validate(t.Context(), tt.path) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response, *got) + }) + } +} + +func TestBuild(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts plugins.BuildOptions + wantBody buildRequest + response plugins.BuildResult + statusCode int + wantErr bool + }{ + { + name: "success", + opts: plugins.BuildOptions{Path: "/home/user/my-plugin", Tag: "v1.0.0"}, + wantBody: buildRequest{Path: "/home/user/my-plugin", Tag: "v1.0.0"}, + response: plugins.BuildResult{Reference: "ghcr.io/org/my-plugin:v1.0.0"}, + statusCode: http.StatusOK, + }, + { + name: "bad request", + opts: plugins.BuildOptions{}, + statusCode: http.StatusBadRequest, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, pluginsBasePath+"/build", r.URL.Path) + + if tt.wantBody.Path != "" { + var got buildRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&got)) + assert.Equal(t, tt.wantBody, got) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "bad request", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.Build(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response, *got) + }) + } +} + +func TestPush(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts plugins.PushOptions + wantBody pushRequest + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success", + opts: plugins.PushOptions{Reference: "ghcr.io/org/my-plugin:v1.0.0"}, + wantBody: pushRequest{Reference: "ghcr.io/org/my-plugin:v1.0.0"}, + statusCode: http.StatusNoContent, + }, + { + name: "not found", + opts: plugins.PushOptions{Reference: "ghcr.io/org/missing:v1"}, + statusCode: http.StatusNotFound, + wantErr: true, + wantCode: http.StatusNotFound, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, pluginsBasePath+"/push", r.URL.Path) + + if tt.wantBody.Reference != "" { + var got pushRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&got)) + assert.Equal(t, tt.wantBody, got) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "not found", tt.statusCode) + return + } + w.WriteHeader(tt.statusCode) + })) + defer srv.Close() + + c := newTestClient(t, srv) + err := c.Push(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + }) + } +} + +func TestListBuilds(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + response listBuildsResponse + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success", + response: listBuildsResponse{Builds: []plugins.LocalBuild{ + {Tag: "my-plugin", Digest: "sha256:abc", Name: "my-plugin", Version: "1.0.0"}, + }}, + statusCode: http.StatusOK, + }, + { + name: "empty", + response: listBuildsResponse{Builds: []plugins.LocalBuild{}}, + statusCode: http.StatusOK, + }, + { + name: "server error", + statusCode: http.StatusInternalServerError, + wantErr: true, + wantCode: http.StatusInternalServerError, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodGet, r.Method) + assert.Equal(t, pluginsBasePath+"/builds", r.URL.Path) + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "error", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.ListBuilds(t.Context()) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response.Builds, got) + }) + } +} + +func TestDeleteBuild(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + tag string + wantPath string + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success", + tag: "my-plugin", + wantPath: pluginsBasePath + "/builds/my-plugin", + statusCode: http.StatusNoContent, + }, + { + name: "not found", + tag: "missing", + wantPath: pluginsBasePath + "/builds/missing", + statusCode: http.StatusNotFound, + wantErr: true, + wantCode: http.StatusNotFound, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodDelete, r.Method) + assert.Equal(t, tt.wantPath, r.URL.Path) + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "error", tt.statusCode) + return + } + w.WriteHeader(tt.statusCode) + })) + defer srv.Close() + + c := newTestClient(t, srv) + err := c.DeleteBuild(t.Context(), tt.tag) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + }) + } +} + +func TestGetContent(t *testing.T) { + t.Parallel() + + response := plugins.PluginContent{ + Name: "my-plugin", + Description: "A test plugin", + Version: "1.0.0", + License: "Apache-2.0", + Manifest: `{"name":"my-plugin"}`, + Files: []plugins.PluginFileEntry{{Path: "plugin.json", Size: 42}}, + } + + tests := []struct { + name string + opts plugins.ContentOptions + wantQuery string + response plugins.PluginContent + statusCode int + wantErr bool + wantCode int + }{ + { + name: "success with local tag", + opts: plugins.ContentOptions{Reference: "my-plugin"}, + wantQuery: "my-plugin", + response: response, + statusCode: http.StatusOK, + }, + { + name: "success with OCI reference", + opts: plugins.ContentOptions{Reference: "ghcr.io/org/my-plugin:v1"}, + wantQuery: "ghcr.io/org/my-plugin:v1", + response: response, + statusCode: http.StatusOK, + }, + { + name: "server error propagates", + opts: plugins.ContentOptions{Reference: "missing"}, + statusCode: http.StatusBadRequest, + wantErr: true, + wantCode: http.StatusBadRequest, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodGet, r.Method) + assert.Equal(t, pluginsBasePath+"/content", r.URL.Path) + if tt.wantQuery != "" { + assert.Equal(t, tt.wantQuery, r.URL.Query().Get("ref")) + } + + if tt.statusCode >= http.StatusBadRequest { + http.Error(w, "bad request", tt.statusCode) + return + } + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(tt.response)) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.GetContent(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + return + } + require.NoError(t, err) + assert.Equal(t, tt.response, *got) + }) + } +} + +func TestConnectionError(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + srv.Close() + + c := NewClient(srv.URL) + _, err := c.List(t.Context(), plugins.ListOptions{}) + + require.Error(t, err) + assert.True(t, errors.Is(err, ErrServerUnreachable), "expected ErrServerUnreachable, got: %v", err) +} + +func TestNewDefaultClient(t *testing.T) { + t.Parallel() + + // noDiscovery stubs server discovery to report no running server, isolating + // the test from any real local server (e.g. a running Desktop app). + noDiscovery := func(context.Context) (string, []Option) { return "", nil } + + // failDiscovery fails the test if discovery is consulted, asserting that an + // earlier resolution step short-circuited. + failDiscovery := func(t *testing.T) discoverFunc { + t.Helper() + return func(context.Context) (string, []Option) { + t.Error("discovery should not be called when TOOLHIVE_API_URL is set") + return "", nil + } + } + + t.Run("falls back to default URL when env is empty", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery) + assert.Equal(t, defaultBaseURL, c.baseURL) + }) + + t.Run("uses TOOLHIVE_API_URL from env", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("http://localhost:9999") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() + + c := newDefaultClientWithEnv(t.Context(), mockEnv, failDiscovery(t)) + assert.Equal(t, "http://localhost:9999", c.baseURL) + }) + + t.Run("uses discovered server when env is empty", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() + + discover := func(context.Context) (string, []Option) { + return "http://127.0.0.1:54321", nil + } + c := newDefaultClientWithEnv(t.Context(), mockEnv, discover) + assert.Equal(t, "http://127.0.0.1:54321", c.baseURL) + }) + + t.Run("applies options", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery, WithTimeout(5*time.Second)) + assert.Equal(t, 5*time.Second, c.httpClient.Timeout) + }) +} + +func TestWithHTTPClient(t *testing.T) { + t.Parallel() + + custom := &http.Client{Timeout: 99 * time.Second} + c := NewClient("http://example.com", WithHTTPClient(custom)) + assert.Equal(t, custom, c.httpClient) +} + +func TestURLEncodesPluginNames(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, pluginsBasePath+"/my%20plugin%2Fv2", r.URL.RawPath) + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(plugins.PluginInfo{ + Metadata: plugins.PluginMetadata{Name: "my plugin/v2"}, + })) + })) + defer srv.Close() + + c := newTestClient(t, srv) + got, err := c.Info(t.Context(), plugins.InfoOptions{Name: "my plugin/v2"}) + require.NoError(t, err) + assert.Equal(t, "my plugin/v2", got.Metadata.Name) +} + +func TestHandleErrorResponseReadFailure(t *testing.T) { + t.Parallel() + + resp := &http.Response{ + StatusCode: http.StatusInternalServerError, + Body: io.NopCloser(&failReader{}), + } + err := handleErrorResponse(resp) + + require.Error(t, err) + assert.Equal(t, http.StatusInternalServerError, httperr.Code(err)) + assert.Contains(t, err.Error(), "failed to read error response body") +} + +type failReader struct{} + +func (*failReader) Read([]byte) (int, error) { + return 0, errors.New("simulated read error") +} + +func TestTimeoutFromEnv(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + value string + want time.Duration + ok bool + }{ + {name: "unset falls back to the default", value: "", want: 0, ok: false}, + {name: "a duration is honored", value: "45s", want: 45 * time.Second, ok: true}, + {name: "minutes parse", value: "5m", want: 5 * time.Minute, ok: true}, + {name: "surrounding whitespace is tolerated", value: " 90s ", want: 90 * time.Second, ok: true}, + {name: "a bare number is not a duration", value: "60", want: 0, ok: false}, + {name: "garbage degrades to the default", value: "soon", want: 0, ok: false}, + {name: "zero would disable the timeout, so it is ignored", value: "0s", want: 0, ok: false}, + {name: "negative is ignored", value: "-5s", want: 0, ok: false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPITimeout).Return(tc.value) + + got, ok := timeoutFromEnv(mockEnv) + assert.Equal(t, tc.ok, ok) + assert.Equal(t, tc.want, got) + }) + } +} + +func TestNewDefaultClientTimeoutPrecedence(t *testing.T) { + t.Parallel() + + noDiscovery := func(context.Context) (string, []Option) { return "", nil } + + t.Run("defaults when the env is unset", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("") + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery) + assert.Equal(t, defaultTimeout, c.httpClient.Timeout) + }) + + t.Run("env overrides the default", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("2m") + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery) + assert.Equal(t, 2*time.Minute, c.httpClient.Timeout) + }) + + t.Run("an explicit option outranks the env", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("2m") + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery, WithTimeout(7*time.Second)) + assert.Equal(t, 7*time.Second, c.httpClient.Timeout) + }) +} + +// TestTimeoutIsReportedAsTimeoutNotUnreachable pins the distinction the CLI +// hint depends on: a healthy-but-slow server must not be reported as absent. +func TestTimeoutIsReportedAsTimeoutNotUnreachable(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-time.After(30 * time.Second): + } + })) + t.Cleanup(srv.Close) + + c := NewClient(srv.URL, WithTimeout(50*time.Millisecond)) + _, err := c.List(t.Context(), plugins.ListOptions{}) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrRequestTimeout) + assert.NotErrorIs(t, err, ErrServerUnreachable, + "the server answered the connection; only the response was slow") +} + +func TestUnreachableServerIsStillUnreachable(t *testing.T) { + t.Parallel() + + // Bind and immediately release a port so nothing is listening on it. + l, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := l.Addr().String() + require.NoError(t, l.Close()) + + c := NewClient("http://"+addr, WithTimeout(2*time.Second)) + _, err = c.List(t.Context(), plugins.ListOptions{}) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrServerUnreachable) + assert.NotErrorIs(t, err, ErrRequestTimeout) +} + +// TestCallerCancellationIsNeitherSentinel keeps a user pressing Ctrl-C from +// being reported as a server problem. +func TestCallerCancellationIsNeitherSentinel(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + t.Cleanup(srv.Close) + + ctx, cancel := context.WithCancel(t.Context()) + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + c := NewClient(srv.URL, WithTimeout(30*time.Second)) + _, err := c.List(ctx, plugins.ListOptions{}) + + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + assert.NotErrorIs(t, err, ErrRequestTimeout) + assert.NotErrorIs(t, err, ErrServerUnreachable) +} diff --git a/pkg/plugins/client/dto.go b/pkg/plugins/client/dto.go new file mode 100644 index 0000000000..00f0f29e67 --- /dev/null +++ b/pkg/plugins/client/dto.go @@ -0,0 +1,43 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import "github.com/stacklok/toolhive/pkg/plugins" + +// --- request/response dto (mirror pkg/api/v1/plugins_types.go) --- + +type installRequest struct { + Name string `json:"name"` + Version string `json:"version,omitempty"` + Scope plugins.Scope `json:"scope,omitempty"` + ProjectRoot string `json:"project_root,omitempty"` + Clients []string `json:"clients,omitempty"` + Force bool `json:"force,omitempty"` + Group string `json:"group,omitempty"` +} + +type validateRequest struct { + Path string `json:"path"` +} + +type buildRequest struct { + Path string `json:"path"` + Tag string `json:"tag,omitempty"` +} + +type pushRequest struct { + Reference string `json:"reference"` +} + +type listResponse struct { + Plugins []plugins.InstalledPlugin `json:"plugins"` +} + +type installResponse struct { + Plugin plugins.InstalledPlugin `json:"plugin"` +} + +type listBuildsResponse struct { + Builds []plugins.LocalBuild `json:"builds"` +} diff --git a/pkg/plugins/pluginsvc/content.go b/pkg/plugins/pluginsvc/content.go index bc0fc817ca..1bacf3ebd8 100644 --- a/pkg/plugins/pluginsvc/content.go +++ b/pkg/plugins/pluginsvc/content.go @@ -85,6 +85,10 @@ func (s *service) getContentFromOCI(ctx context.Context, ref string) (*plugins.P ) } + if err := validateOCIRegistryHost(ociRef); err != nil { + return nil, err + } + qualifiedRef := qualifiedOCIRef(ociRef) pullCtx, cancel := context.WithTimeout(ctx, ociPullTimeout) defer cancel() diff --git a/pkg/plugins/pluginsvc/install.go b/pkg/plugins/pluginsvc/install.go index 122707c91e..e660060c96 100644 --- a/pkg/plugins/pluginsvc/install.go +++ b/pkg/plugins/pluginsvc/install.go @@ -116,43 +116,171 @@ func (s *service) installByName( // installFromRegistryLookup resolves a plain plugin name via the registry // lookup and dispatches to the appropriate installer. When no lookup is // configured or it returns no hits, returns a 404 with an install hint. +// +// Resolution mirrors skillsvc.resolveFromRegistry exactly: +// - a lookup error is logged and treated as "not found" (fall back to the +// 404 hint) rather than propagated, and +// - search hits are post-filtered for an exact, case-insensitive name match +// (PluginSearchHit carries Name only, so there is no namespace filter). +// +// When opts.Version is set, exact-name matches are further narrowed to those +// whose Version equals opts.Version; no version match falls through to the 404 +// not-found path (which includes the requested version in the message). +// +// Multiple exact matches are ambiguous and produce a 409 listing candidates. +// Package selection mirrors skillsvc.resolveRegistryPackages: the first OCI +// package (Type == "oci") is chosen; when none exists, a 422 is returned. func (s *service) installFromRegistryLookup( ctx context.Context, opts plugins.InstallOptions, scope plugins.Scope, ) (*plugins.InstallResult, error) { if s.pluginLookup != nil { - hits, err := s.pluginLookup.SearchPlugins(ctx, opts.Name) + // Use the last path segment as the search query (matching + // skillsvc.resolveFromRegistry's splitQualifiedName), since + // SearchPlugins matches on name substring. + _, searchName := splitQualifiedName(opts.Name) + + hits, err := s.pluginLookup.SearchPlugins(ctx, searchName) if err != nil { - return nil, fmt.Errorf("searching registry for plugin %q: %w", opts.Name, err) + slog.Warn("registry plugin lookup failed, falling back to not-found", "name", opts.Name, "error", err) + hits = nil } - if len(hits) > 0 && len(hits[0].Packages) > 0 { - pkg := hits[0].Packages[0] - slog.Info("resolved plugin from registry", "name", opts.Name, "reference", pkg.Reference) - opts.Name = pkg.Reference - ref, _, parseErr := parseOCIReference(pkg.Reference) - if parseErr != nil { - return nil, httperr.WithCode( - fmt.Errorf("registry returned invalid OCI reference %q: %w", pkg.Reference, parseErr), - http.StatusUnprocessableEntity, - ) + + // Filter for exact match. Case-insensitive because registry data + // may not be normalized to lowercase even though local plugin names are. + var matches []PluginSearchHit + for _, hit := range hits { + if !strings.EqualFold(hit.Name, searchName) { + continue } - result, ociErr := s.installFromOCI(ctx, opts, scope, ref) - if ociErr != nil { - return nil, ociErr + matches = append(matches, hit) + } + + // When a version was requested, narrow the exact-name matches to those + // whose Version equals opts.Version (string equality). A hit that does + // not carry a Version cannot satisfy a versioned request. No match + // falls through to the 404 not-found path below. + if opts.Version != "" { + var versioned []PluginSearchHit + for _, m := range matches { + if m.Version == opts.Version { + versioned = append(versioned, m) + } } - return s.installAndRegister(ctx, result, opts.Group, result.Plugin.Metadata.Name, scope, opts.ProjectRoot) + matches = versioned + } + + if len(matches) == 1 { + return s.installFromRegistryHit(ctx, opts, scope, matches[0]) + } + + if len(matches) > 1 { + return nil, ambiguousPluginNameError(opts.Name, matches) } } + if opts.Version != "" { + return nil, httperr.WithCode( + fmt.Errorf("plugin %q version %q not found in local store or registry;"+ + " install by OCI reference:\n thv ai-plugin install ghcr.io//%s:%s", + opts.Name, opts.Version, opts.Name, opts.Version), + http.StatusNotFound, + ) + } return nil, httperr.WithCode( fmt.Errorf("plugin %q not found in local store or registry;"+ - " install by OCI reference:\n thv plugin install ghcr.io//%s:", + " install by OCI reference:\n thv ai-plugin install ghcr.io//%s:", opts.Name, opts.Name), http.StatusNotFound, ) } +// installFromRegistryHit selects the OCI package from a single resolved hit, +// parses its reference, and dispatches to installFromOCI. It guards against a +// malformed catalog package whose Reference is not structurally an OCI +// reference (parseOCIReference returns (nil, false, nil) for such input), +// which would otherwise nil-deref in validateOCIRegistryHost. +func (s *service) installFromRegistryHit( + ctx context.Context, + opts plugins.InstallOptions, + scope plugins.Scope, + hit PluginSearchHit, +) (*plugins.InstallResult, error) { + pkg, pkgErr := selectOCIPluginPackage(opts.Name, hit.Packages) + if pkgErr != nil { + return nil, pkgErr + } + slog.Info("resolved plugin from registry", "name", opts.Name, "reference", pkg.Reference) + opts.Name = pkg.Reference + ref, isOCIRef, parseErr := parseOCIReference(pkg.Reference) + if parseErr != nil { + return nil, httperr.WithCode( + fmt.Errorf("registry returned invalid OCI reference %q: %w", pkg.Reference, parseErr), + http.StatusUnprocessableEntity, + ) + } + if !isOCIRef || ref == nil { + return nil, httperr.WithCode( + fmt.Errorf("registry returned invalid OCI reference %q", pkg.Reference), + http.StatusUnprocessableEntity, + ) + } + result, ociErr := s.installFromOCI(ctx, opts, scope, ref) + if ociErr != nil { + return nil, ociErr + } + return s.installAndRegister(ctx, result, opts.Group, result.Plugin.Metadata.Name, scope, opts.ProjectRoot) +} + +// selectOCIPluginPackage selects the first OCI package from a registry entry's +// package list. Mirror of skillsvc.resolveRegistryPackages' OCI branch: OCI +// packages are preferred, and a missing OCI package yields a 422 error. +func selectOCIPluginPackage(name string, packages []PluginPackage) (PluginPackage, error) { + for _, pkg := range packages { + if pkg.Type == "oci" && pkg.Reference != "" { + return pkg, nil + } + } + return PluginPackage{}, httperr.WithCode( + fmt.Errorf("plugin %q found in registry but has no installable OCI package", name), + http.StatusUnprocessableEntity, + ) +} + +// ambiguousPluginNameError builds a 409 error listing each ambiguous match, +// capped at 5 candidates with an "and N more" suffix. Mirror of +// skillsvc.resolveFromRegistry's ambiguity path. +func ambiguousPluginNameError(name string, matches []PluginSearchHit) error { + const maxCandidates = 5 + candidates := make([]string, 0, len(matches)) + for _, m := range matches { + candidates = append(candidates, m.Name) + } + suffix := "" + if len(candidates) > maxCandidates { + suffix = fmt.Sprintf(" and %d more", len(candidates)-maxCandidates) + candidates = candidates[:maxCandidates] + } + return httperr.WithCode( + fmt.Errorf("ambiguous plugin name %q matches multiple registry entries: %s%s; install by full OCI reference instead", + name, strings.Join(candidates, ", "), suffix), + http.StatusConflict, + ) +} + +// splitQualifiedName splits "namespace/name" into (namespace, name). If the +// input has no "/" it returns ("", name) unchanged. Plugin names validated by +// ValidatePluginName never contain "/", so this is a structural mirror of +// skillsvc.splitQualifiedName that is a no-op for valid plugin names. +func splitQualifiedName(s string) (namespace, name string) { + idx := strings.LastIndex(s, "/") + if idx < 0 { + return "", s + } + return s[:idx], s[idx+1:] +} + // registerPluginInGroup adds the plugin to the requested group when a group // manager is configured. When groupName is empty it defaults to the "default" // group, matching workload behavior. diff --git a/pkg/plugins/pluginsvc/install_git.go b/pkg/plugins/pluginsvc/install_git.go index 3712477001..1048c26f08 100644 --- a/pkg/plugins/pluginsvc/install_git.go +++ b/pkg/plugins/pluginsvc/install_git.go @@ -116,10 +116,11 @@ func (s *service) cloneAndCollectPlugin( } defer func() { _ = client.Cleanup(ctx, repoInfo) }() - commitHash, err := client.HeadCommitHash(repoInfo) + head, err := client.HeadCommit(repoInfo) if err != nil { - return nil, nil, "", fmt.Errorf("getting commit hash: %w", err) + return nil, nil, "", fmt.Errorf("getting HEAD commit: %w", err) } + commitHash := head.Hash // Read the plugin manifest. It lives at /.claude-plugin/plugin.json. // We collect the whole file tree first, then locate the manifest among the diff --git a/pkg/plugins/pluginsvc/install_oci.go b/pkg/plugins/pluginsvc/install_oci.go index 8b2aa03f27..98bcb81e1f 100644 --- a/pkg/plugins/pluginsvc/install_oci.go +++ b/pkg/plugins/pluginsvc/install_oci.go @@ -42,6 +42,10 @@ func (s *service) installFromOCI( ) } + if err := validateOCIRegistryHost(ref); err != nil { + return nil, err + } + ociRef := qualifiedOCIRef(ref) pullCtx, cancel := context.WithTimeout(ctx, ociPullTimeout) diff --git a/pkg/plugins/pluginsvc/install_registry_test.go b/pkg/plugins/pluginsvc/install_registry_test.go new file mode 100644 index 0000000000..be19bc5257 --- /dev/null +++ b/pkg/plugins/pluginsvc/install_registry_test.go @@ -0,0 +1,376 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package pluginsvc + +import ( + "context" + "fmt" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + ociplugins "github.com/stacklok/toolhive-core/oci/plugins" + ocimocks "github.com/stacklok/toolhive-core/oci/plugins/mocks" + "github.com/stacklok/toolhive/pkg/plugins" + plugmocks "github.com/stacklok/toolhive/pkg/plugins/mocks" + "github.com/stacklok/toolhive/pkg/storage" + storemocks "github.com/stacklok/toolhive/pkg/storage/mocks" +) + +// stubLookup is a test helper implementing PluginLookup with canned results. +type stubLookup struct { + hits []PluginSearchHit + err error +} + +func (s *stubLookup) SearchPlugins(_ context.Context, _ string) ([]PluginSearchHit, error) { + return s.hits, s.err +} + +func TestInstallRegistryResolution(t *testing.T) { + t.Parallel() + + t.Run("resolves plain name via lookup and installs from OCI", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + + ociStore, err := ociplugins.NewStore(tempDir(t)) + require.NoError(t, err) + indexDigest := buildTestPlugin(t, ociStore, "my-plugin", "1.0.0") + + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), ociStore, "ghcr.io/org/my-plugin:v1"). + Return(indexDigest, nil) + + store := storemocks.NewMockPluginStore(ctrl) + adapter := plugmocks.NewMockMaterializationAdapter(ctrl) + store.EXPECT().Get(gomock.Any(), "my-plugin", plugins.ScopeUser, "").Return(plugins.InstalledPlugin{}, storage.ErrNotFound) + adapter.EXPECT().Materialize(gomock.Any(), gomock.Any()).Return(&plugins.MaterializeResult{}, nil) + store.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, p plugins.InstalledPlugin) error { + assert.Equal(t, "my-plugin", p.Metadata.Name) + assert.Equal(t, "1.0.0", p.Metadata.Version) + assert.Equal(t, "ghcr.io/org/my-plugin:v1", p.Reference) + return nil + }) + + lookup := &stubLookup{hits: []PluginSearchHit{ + { + Name: "my-plugin", + Description: "test plugin", + Packages: []PluginPackage{{Reference: "ghcr.io/org/my-plugin:v1", Type: "oci"}}, + }, + }} + + svc := newTestService( + WithStore(store), + WithOCIStore(ociStore), + WithRegistryClient(reg), + WithMaterializers(map[string]plugins.MaterializationAdapter{"claude-code": adapter}), + WithPluginLookup(lookup), + ) + result, err := svc.Install(t.Context(), plugins.InstallOptions{ + Name: "my-plugin", + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + assert.Equal(t, "my-plugin", result.Plugin.Metadata.Name) + assert.Equal(t, "1.0.0", result.Plugin.Metadata.Version) + assert.Equal(t, "ghcr.io/org/my-plugin:v1", result.Plugin.Reference) + }) + + t.Run("exact name wins over substring matches", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + + ociStore, err := ociplugins.NewStore(tempDir(t)) + require.NoError(t, err) + indexDigest := buildTestPlugin(t, ociStore, "my-plugin", "1.0.0") + + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), ociStore, "ghcr.io/org/my-plugin:v1"). + Return(indexDigest, nil) + + store := storemocks.NewMockPluginStore(ctrl) + adapter := plugmocks.NewMockMaterializationAdapter(ctrl) + store.EXPECT().Get(gomock.Any(), "my-plugin", plugins.ScopeUser, "").Return(plugins.InstalledPlugin{}, storage.ErrNotFound) + adapter.EXPECT().Materialize(gomock.Any(), gomock.Any()).Return(&plugins.MaterializeResult{}, nil) + store.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, p plugins.InstalledPlugin) error { + assert.Equal(t, "my-plugin", p.Metadata.Name) + return nil + }) + + // Search returns several substring hits, only one with an exact Name + // match. The exact match must be selected even though it is not first. + lookup := &stubLookup{hits: []PluginSearchHit{ + {Name: "my-plugin-extra", Packages: []PluginPackage{{Reference: "ghcr.io/org/other:v1", Type: "oci"}}}, + {Name: "my-plugin", Packages: []PluginPackage{{Reference: "ghcr.io/org/my-plugin:v1", Type: "oci"}}}, + {Name: "my-plugin-tool", Packages: []PluginPackage{{Reference: "ghcr.io/org/other2:v1", Type: "oci"}}}, + }} + + svc := newTestService( + WithStore(store), + WithOCIStore(ociStore), + WithRegistryClient(reg), + WithMaterializers(map[string]plugins.MaterializationAdapter{"claude-code": adapter}), + WithPluginLookup(lookup), + ) + result, err := svc.Install(t.Context(), plugins.InstallOptions{ + Name: "my-plugin", + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + assert.Equal(t, "my-plugin", result.Plugin.Metadata.Name) + assert.Equal(t, "ghcr.io/org/my-plugin:v1", result.Plugin.Reference) + }) + + t.Run("ambiguous exact matches return 409", func(t *testing.T) { + t.Parallel() + + // Two hits with the exact same Name across namespaces — ambiguous. + lookup := &stubLookup{hits: []PluginSearchHit{ + {Name: "my-plugin", Packages: []PluginPackage{{Reference: "ghcr.io/org1/my-plugin:v1", Type: "oci"}}}, + {Name: "my-plugin", Packages: []PluginPackage{{Reference: "ghcr.io/org2/my-plugin:v1", Type: "oci"}}}, + }} + + svc := newTestService( + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{Name: "my-plugin"}) + require.Error(t, err) + assert.Equal(t, http.StatusConflict, httperr.Code(err)) + assert.Contains(t, err.Error(), "ambiguous plugin name") + assert.Contains(t, err.Error(), "my-plugin") + }) + + t.Run("selects oci package over positional git package", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + + ociStore, err := ociplugins.NewStore(tempDir(t)) + require.NoError(t, err) + indexDigest := buildTestPlugin(t, ociStore, "my-plugin", "1.0.0") + + reg := ocimocks.NewMockRegistryClient(ctrl) + // The OCI package is second in the list; the first is a git package. + // Selection must pick the OCI package, not Packages[0]. + reg.EXPECT().Pull(gomock.Any(), ociStore, "ghcr.io/org/my-plugin:v1"). + Return(indexDigest, nil) + + store := storemocks.NewMockPluginStore(ctrl) + adapter := plugmocks.NewMockMaterializationAdapter(ctrl) + store.EXPECT().Get(gomock.Any(), "my-plugin", plugins.ScopeUser, "").Return(plugins.InstalledPlugin{}, storage.ErrNotFound) + adapter.EXPECT().Materialize(gomock.Any(), gomock.Any()).Return(&plugins.MaterializeResult{}, nil) + store.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, p plugins.InstalledPlugin) error { + assert.Equal(t, "ghcr.io/org/my-plugin:v1", p.Reference) + return nil + }) + + lookup := &stubLookup{hits: []PluginSearchHit{ + { + Name: "my-plugin", + Packages: []PluginPackage{ + {Reference: "https://github.com/org/repo", Type: "git"}, + {Reference: "ghcr.io/org/my-plugin:v1", Type: "oci"}, + }, + }, + }} + + svc := newTestService( + WithStore(store), + WithOCIStore(ociStore), + WithRegistryClient(reg), + WithMaterializers(map[string]plugins.MaterializationAdapter{"claude-code": adapter}), + WithPluginLookup(lookup), + ) + result, err := svc.Install(t.Context(), plugins.InstallOptions{ + Name: "my-plugin", + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + assert.Equal(t, "ghcr.io/org/my-plugin:v1", result.Plugin.Reference) + }) + + t.Run("no oci package returns 422", func(t *testing.T) { + t.Parallel() + + // Exact match but only a git package (plugins have no git install flow). + lookup := &stubLookup{hits: []PluginSearchHit{ + {Name: "my-plugin", Packages: []PluginPackage{{Reference: "https://github.com/org/repo", Type: "git"}}}, + }} + + svc := newTestService( + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{Name: "my-plugin"}) + require.Error(t, err) + assert.Equal(t, http.StatusUnprocessableEntity, httperr.Code(err)) + assert.Contains(t, err.Error(), "no installable OCI package") + }) + + t.Run("lookup returns no hits returns 404", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + + store := storemocks.NewMockPluginStore(ctrl) + + lookup := &stubLookup{hits: nil} + + svc := newTestService( + WithStore(store), + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{Name: "nonexistent"}) + require.Error(t, err) + assert.Equal(t, http.StatusNotFound, httperr.Code(err)) + assert.Contains(t, err.Error(), "not found in local store or registry") + }) + + t.Run("lookup search error falls back to 404 not-found", func(t *testing.T) { + t.Parallel() + + lookup := &stubLookup{err: fmt.Errorf("registry timeout")} + + svc := newTestService( + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{Name: "some-plugin"}) + require.Error(t, err) + // A lookup error is logged and treated as not-found (mirroring + // skillsvc.resolveFromRegistry), not propagated. + assert.Equal(t, http.StatusNotFound, httperr.Code(err)) + assert.Contains(t, err.Error(), "not found in local store or registry") + }) + + // Regression: a malformed catalog package typed "oci" but whose Reference + // has no '/', ':', or '@' must not reach installFromOCI with a nil ref. + // parseOCIReference returns (nil, false, nil) for such input; previously + // the isOCI bool was discarded and the nil ref caused a panic in + // validateOCIRegistryHost via ref.Context(). Must return 422 instead. + t.Run("malformed oci package reference returns 422 no panic", func(t *testing.T) { + t.Parallel() + + lookup := &stubLookup{hits: []PluginSearchHit{ + { + Name: "my-plugin", + Version: "1.0.0", + Packages: []PluginPackage{ + {Reference: "foo", Type: "oci"}, + }, + }, + }} + + svc := newTestService( + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{Name: "my-plugin"}) + require.Error(t, err) + assert.Equal(t, http.StatusUnprocessableEntity, httperr.Code(err)) + assert.Contains(t, err.Error(), "invalid OCI reference") + assert.Contains(t, err.Error(), "foo") + }) + + // Regression: when opts.Version is set, only an exact-name hit whose + // Version matches is installed. + t.Run("version requested with matching hit version installs", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + + ociStore, err := ociplugins.NewStore(tempDir(t)) + require.NoError(t, err) + indexDigest := buildTestPlugin(t, ociStore, "my-plugin", "1.0.0") + + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), ociStore, "ghcr.io/org/my-plugin:1.0.0"). + Return(indexDigest, nil) + + store := storemocks.NewMockPluginStore(ctrl) + adapter := plugmocks.NewMockMaterializationAdapter(ctrl) + store.EXPECT().Get(gomock.Any(), "my-plugin", plugins.ScopeUser, "").Return(plugins.InstalledPlugin{}, storage.ErrNotFound) + adapter.EXPECT().Materialize(gomock.Any(), gomock.Any()).Return(&plugins.MaterializeResult{}, nil) + store.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, p plugins.InstalledPlugin) error { + assert.Equal(t, "my-plugin", p.Metadata.Name) + assert.Equal(t, "ghcr.io/org/my-plugin:1.0.0", p.Reference) + return nil + }) + + lookup := &stubLookup{hits: []PluginSearchHit{ + { + Name: "my-plugin", + Version: "1.0.0", + Packages: []PluginPackage{ + {Reference: "ghcr.io/org/my-plugin:1.0.0", Type: "oci"}, + }, + }, + }} + + svc := newTestService( + WithStore(store), + WithOCIStore(ociStore), + WithRegistryClient(reg), + WithMaterializers(map[string]plugins.MaterializationAdapter{"claude-code": adapter}), + WithPluginLookup(lookup), + ) + result, err := svc.Install(t.Context(), plugins.InstallOptions{ + Name: "my-plugin", + Version: "1.0.0", + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + assert.Equal(t, "my-plugin", result.Plugin.Metadata.Name) + assert.Equal(t, "ghcr.io/org/my-plugin:1.0.0", result.Plugin.Reference) + }) + + // Regression: a version request with only non-matching hit versions must + // fall through to the 404 path, and the message must mention the version. + t.Run("version requested with non-matching hit version returns 404 mentioning version", func(t *testing.T) { + t.Parallel() + + lookup := &stubLookup{hits: []PluginSearchHit{ + { + Name: "my-plugin", + Version: "2.0.0", + Packages: []PluginPackage{ + {Reference: "ghcr.io/org/my-plugin:2.0.0", Type: "oci"}, + }, + }, + }} + + svc := newTestService( + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{ + Name: "my-plugin", + Version: "1.0.0", + }) + require.Error(t, err) + assert.Equal(t, http.StatusNotFound, httperr.Code(err)) + assert.Contains(t, err.Error(), "not found in local store or registry") + assert.Contains(t, err.Error(), "1.0.0", "404 message should mention the requested version") + }) + + // Regression: the install hint text uses the renamed command + // "thv ai-plugin install" (was "thv plugin install"). + t.Run("not found hint text references thv ai-plugin install", func(t *testing.T) { + t.Parallel() + + lookup := &stubLookup{hits: nil} + + svc := newTestService( + WithPluginLookup(lookup), + ) + _, err := svc.Install(t.Context(), plugins.InstallOptions{Name: "nonexistent"}) + require.Error(t, err) + assert.Equal(t, http.StatusNotFound, httperr.Code(err)) + assert.Contains(t, err.Error(), "thv ai-plugin install") + assert.NotContains(t, err.Error(), "thv plugin install") + }) +} diff --git a/pkg/plugins/pluginsvc/oci.go b/pkg/plugins/pluginsvc/oci.go index 4bbb6c06a5..a70bb3dacb 100644 --- a/pkg/plugins/pluginsvc/oci.go +++ b/pkg/plugins/pluginsvc/oci.go @@ -8,7 +8,9 @@ import ( "encoding/json" "errors" "fmt" + "net" "net/http" + "os" "strings" nameref "github.com/google/go-containerregistry/pkg/name" @@ -17,6 +19,7 @@ import ( "github.com/stacklok/toolhive-core/httperr" ociplugins "github.com/stacklok/toolhive-core/oci/plugins" + "github.com/stacklok/toolhive/pkg/networking" ) // maxConfigSize bounds the OCI image config blob read into memory during @@ -59,6 +62,46 @@ func parseOCIReference(name string) (nameref.Reference, bool, error) { return ref, true, nil } +// validateOCIRegistryHost rejects OCI references whose registry host is +// localhost or a private/loopback IP, preventing SSRF via attacker-controlled +// references on the install and content endpoints. Mirrors the git install +// path's validateHost (pkg/skills/gitresolver). In dev mode (TOOLHIVE_DEV=true) +// the check is relaxed so E2E tests can pull from a local registry. +// +// NOTE: This check only validates literal IPs and known localhost strings. +// Hostnames that DNS-resolve to private IPs (DNS rebinding) are NOT caught. +func validateOCIRegistryHost(ref nameref.Reference) error { + host := ref.Context().RegistryStr() + + // Strip port if present. + hostname := host + if h, _, err := net.SplitHostPort(host); err == nil { + hostname = h + } + + // Dev-mode bypass mirrors gitresolver.isDevMode() (TOOLHIVE_DEV=true); + // that helper is unexported, so the same env-var check is replicated here. + if strings.EqualFold(os.Getenv("TOOLHIVE_DEV"), "true") { + return nil + } + + if networking.IsLocalhost(hostname) { + return httperr.WithCode( + fmt.Errorf("registry host %q is not allowed: localhost is rejected for SSRF prevention", host), + http.StatusBadRequest, + ) + } + + if ip := net.ParseIP(hostname); ip != nil && networking.IsPrivateIP(ip) { + return httperr.WithCode( + fmt.Errorf("registry host %q is not allowed: private/loopback IPs are rejected for SSRF prevention", host), + http.StatusBadRequest, + ) + } + + return nil +} + // isPluginArtifact reports whether the OCI descriptor at digest d carries // ArtifactType == ArtifactTypePlugin. It inspects the top-level index or // manifest without descending into layers, so it is cheap to call. diff --git a/pkg/plugins/pluginsvc/oci_ssrf_test.go b/pkg/plugins/pluginsvc/oci_ssrf_test.go new file mode 100644 index 0000000000..c24472d309 --- /dev/null +++ b/pkg/plugins/pluginsvc/oci_ssrf_test.go @@ -0,0 +1,146 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package pluginsvc + +import ( + "net/http" + "testing" + + nameref "github.com/google/go-containerregistry/pkg/name" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + ociplugins "github.com/stacklok/toolhive-core/oci/plugins" + ocimocks "github.com/stacklok/toolhive-core/oci/plugins/mocks" + "github.com/stacklok/toolhive/pkg/plugins" + plugmocks "github.com/stacklok/toolhive/pkg/plugins/mocks" +) + +func TestValidateOCIRegistryHost(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + ref string + wantCode int + wantErr string + }{ + { + name: "localhost registry rejected", + ref: "localhost:5000/org/plugin:v1", + wantCode: http.StatusBadRequest, + wantErr: "localhost is rejected for SSRF prevention", + }, + { + name: "loopback IP registry rejected", + ref: "127.0.0.1:5000/org/plugin:v1", + wantCode: http.StatusBadRequest, + wantErr: "rejected for SSRF prevention", + }, + { + name: "private IP registry rejected", + ref: "10.0.0.5/org/plugin:v1", + wantCode: http.StatusBadRequest, + wantErr: "private/loopback IPs are rejected for SSRF prevention", + }, + { + name: "public registry allowed", + ref: "ghcr.io/org/plugin:v1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ref, err := nameref.ParseReference(tt.ref) + require.NoError(t, err) + + err = validateOCIRegistryHost(ref) + if tt.wantCode == 0 { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Equal(t, tt.wantCode, httperr.Code(err)) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +// Dev mode relaxes the SSRF check so E2E tests can use a local registry. +// Not parallel: mutates the process environment. +func TestValidateOCIRegistryHost_DevMode(t *testing.T) { + t.Setenv("TOOLHIVE_DEV", "true") + + ref, err := nameref.ParseReference("localhost:5000/org/plugin:v1") + require.NoError(t, err) + assert.NoError(t, validateOCIRegistryHost(ref)) +} + +func TestInstallOCIRejectsPrivateRegistryHost(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + ref string + }{ + {name: "loopback registry", ref: "127.0.0.1:5000/org/plugin:v1"}, + {name: "localhost registry", ref: "localhost:5000/org/plugin:v1"}, + {name: "private IP registry", ref: "10.0.0.5/org/plugin:v1"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + // No Pull expectation: gomock fails the test if Pull is called. + reg := ocimocks.NewMockRegistryClient(ctrl) + adapter := plugmocks.NewMockMaterializationAdapter(ctrl) + ociStore, err := ociplugins.NewStore(t.TempDir()) + require.NoError(t, err) + + svc := newTestService( + WithRegistryClient(reg), + WithOCIStore(ociStore), + WithMaterializers(map[string]plugins.MaterializationAdapter{"claude-code": adapter}), + ) + _, err = svc.Install(t.Context(), plugins.InstallOptions{Name: tt.ref}) + require.Error(t, err) + assert.Equal(t, http.StatusBadRequest, httperr.Code(err)) + assert.Contains(t, err.Error(), "SSRF prevention") + }) + } +} + +func TestGetContentRejectsPrivateRegistryHost(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + ref string + }{ + {name: "loopback registry", ref: "127.0.0.1:5000/org/plugin:v1"}, + {name: "localhost registry", ref: "localhost:5000/org/plugin:v1"}, + {name: "private IP registry", ref: "10.0.0.5/org/plugin:v1"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + // No Pull expectation: gomock fails the test if Pull is called. + reg := ocimocks.NewMockRegistryClient(ctrl) + ociStore, err := ociplugins.NewStore(t.TempDir()) + require.NoError(t, err) + + svc := New(WithOCIStore(ociStore), WithRegistryClient(reg)) + _, err = svc.GetContent(t.Context(), plugins.ContentOptions{Reference: tt.ref}) + require.Error(t, err) + assert.Equal(t, http.StatusBadRequest, httperr.Code(err)) + assert.Contains(t, err.Error(), "SSRF prevention") + }) + } +} diff --git a/pkg/plugins/pluginsvc/service.go b/pkg/plugins/pluginsvc/service.go index 23ae33b67b..ec97f72d9e 100644 --- a/pkg/plugins/pluginsvc/service.go +++ b/pkg/plugins/pluginsvc/service.go @@ -92,6 +92,13 @@ func WithClientManager(cm *client.ClientManager) Option { type PluginSearchHit struct { // Name is the plugin name (kebab-case). Name string `json:"name"` + // Namespace is the registry namespace the hit was published under + // (reverse-DNS, e.g. "io.github.user"). Carried so callers can disambiguate + // same-named plugins across namespaces. + Namespace string `json:"namespace,omitempty"` + // Version is the plugin version string reported by the registry. Used by + // the install flow to honor an explicit --version request. + Version string `json:"version,omitempty"` // Description is a human-readable description of the plugin. Description string `json:"description,omitempty"` // Packages lists the OCI packages that publish this plugin. @@ -104,6 +111,10 @@ type PluginPackage struct { Reference string `json:"reference"` // Type is the package type (e.g. "oci"). Type string `json:"type,omitempty"` + // Digest is the OCI digest of the package artifact, when reported by the + // registry. Carried as part of the trust tuple (reference + digest) so + // callers can pin installs to a verified artifact. + Digest string `json:"digest,omitempty"` } // PluginLookup resolves a plain plugin name against a registry/index. This seam diff --git a/pkg/ratelimit/errors.go b/pkg/ratelimit/errors.go index 1f00743c67..85fff2c6ba 100644 --- a/pkg/ratelimit/errors.go +++ b/pkg/ratelimit/errors.go @@ -11,9 +11,17 @@ import ( ) const ( - // CodeRateLimited is the JSON-RPC error code for rate-limited requests. - // Per RFC THV-0057: implementation-defined code in the -32000 to -32099 range. - CodeRateLimited int64 = -32029 + // CodeRateLimited is the JSON-RPC error code for rate-limited requests, + // mirroring HTTP 429 (Too Many Requests) the same way mcp.JSONRPCCodeDenied + // mirrors 403. It sits outside the JSON-RPC reserved range + // (-32768..-32000) so a single value is conformant on every MCP revision: + // 2026-07-28 reserves -32020..-32099 exclusively for spec-defined codes + // and directs application codes outside the reserved range, while + // 2025-11-25 and earlier inherit plain JSON-RPC 2.0, which leaves the + // space outside that range application-defined. Replaces -32029, chosen + // by RFC THV-0057 when -32000..-32099 was uniformly implementation-defined + // — a premise the 2026-07-28 partition invalidated (#6101). + CodeRateLimited int64 = 429 // MessageRateLimited is the error message for rate-limited requests. MessageRateLimited = "Rate limit exceeded" diff --git a/pkg/ratelimit/middleware.go b/pkg/ratelimit/middleware.go index c1a842ab4f..e1a64af465 100644 --- a/pkg/ratelimit/middleware.go +++ b/pkg/ratelimit/middleware.go @@ -20,6 +20,7 @@ import ( v1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/transport/types" ) @@ -160,7 +161,13 @@ func writeRateLimited(w http.ResponseWriter, requestID any, retryAfter time.Dura _, _ = w.Write(rateLimitedBody(requestID, retryAfter)) } -// rateLimitedBody returns the JSON-encoded body for a rate-limited JSON-RPC error. +// rateLimitedBody returns the JSON-encoded body for a rate-limited JSON-RPC +// error. +// +// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error +// response as `id?: RequestId` where RequestId = string | number, so an +// absent id is encoded by omitting the "id" key, never as null. See +// session.HasJSONRPCID. func rateLimitedBody(requestID any, retryAfter time.Duration) []byte { retrySeconds := math.Ceil(retryAfter.Seconds()) resp := map[string]any{ @@ -172,12 +179,14 @@ func rateLimitedBody(requestID any, retryAfter time.Duration) []byte { "retryAfterSeconds": retrySeconds, }, }, - "id": requestID, + } + if session.HasJSONRPCID(requestID) { + resp["id"] = requestID } data, err := json.Marshal(resp) if err != nil { return []byte(fmt.Sprintf( - `{"jsonrpc":"2.0","error":{"code":-32029,"message":"Rate limit exceeded","data":{"retryAfterSeconds":%.0f}},"id":null}`, + `{"jsonrpc":"2.0","error":{"code":429,"message":"Rate limit exceeded","data":{"retryAfterSeconds":%.0f}}}`, retrySeconds, )) } diff --git a/pkg/ratelimit/middleware_test.go b/pkg/ratelimit/middleware_test.go index 512ff382ec..c0727290f6 100644 --- a/pkg/ratelimit/middleware_test.go +++ b/pkg/ratelimit/middleware_test.go @@ -108,13 +108,105 @@ func TestRateLimitHandler_ToolCallRejected(t *testing.T) { var resp map[string]any require.NoError(t, json.Unmarshal(body, &resp)) errObj := resp["error"].(map[string]any) - assert.Equal(t, float64(-32029), errObj["code"]) + assert.Equal(t, float64(429), errObj["code"]) assert.Equal(t, "Rate limit exceeded", errObj["message"]) data := errObj["data"].(map[string]any) assert.Equal(t, float64(5), data["retryAfterSeconds"]) assert.Equal(t, float64(42), resp["id"]) } +// TestRateLimitHandler_NotificationRejected is the live regression case: +// rateLimitHandler has no notification guard (unlike dispatchModern), so a +// well-formed tools/call notification (no id) that gets rate-limited reaches +// writeRateLimited with a nil requestID. Before the fix this emitted +// "id":null at HTTP 429; it must now omit the key entirely. +func TestRateLimitHandler_NotificationRejected(t *testing.T) { + t.Parallel() + + limiter := &dummyLimiter{decision: &Decision{Allowed: false, RetryAfter: 5 * time.Second}} + handler := rateLimitHandler(limiter)(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("next handler should not be called when rate limited") + })) + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + req = withParsedMCPRequest(req, "tools/call", "echo", nil) // no id: a notification + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + assert.Equal(t, http.StatusTooManyRequests, w.Code) + + body, err := io.ReadAll(w.Body) + require.NoError(t, err) + + var resp map[string]any + require.NoError(t, json.Unmarshal(body, &resp)) + _, hasID := resp["id"] + assert.False(t, hasID, `"id" key must be absent for a notification, never "id":null`) +} + +// TestRateLimitedBody covers the "id" key's presence/absence per +// session.HasJSONRPCID: MCP narrows base JSON-RPC 2.0 so an absent id is +// encoded by omitting the key entirely, never as "id":null (schema/2025-11-25 +// types the error response id as optional string|number). +func TestRateLimitedBody(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + requestID any + wantIDKey bool + wantIDEcho string // substring the body must contain, verbatim, when wantIDKey + }{ + {name: "nil id is omitted", requestID: nil, wantIDKey: false}, + {name: "null RawMessage id is omitted", requestID: json.RawMessage("null"), wantIDKey: false}, + {name: "nil RawMessage id is omitted", requestID: json.RawMessage(nil), wantIDKey: false}, + // int64 matches the real wire type: parsed.ID comes from jsonrpc2.ID.Raw(), + // which yields nil | int64 | string, never float64. + {name: "zero id is still present", requestID: int64(0), wantIDKey: true, wantIDEcho: `"id":0`}, + {name: "present id is present", requestID: int64(42), wantIDKey: true, wantIDEcho: `"id":42`}, + {name: "string id is present", requestID: "req-1", wantIDKey: true, wantIDEcho: `"id":"req-1"`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + body := rateLimitedBody(tt.requestID, 5*time.Second) + + // Decoding into map[string]any (rather than a tagged struct) is + // the only way to tell an absent "id" key apart from a present + // null: a struct field would read as the zero value either way. + var asMap map[string]any + require.NoError(t, json.Unmarshal(body, &asMap)) + _, hasID := asMap["id"] + assert.Equal(t, tt.wantIDKey, hasID, `"id" key presence`) + if tt.wantIDKey { + // A map decode would coerce a numeric id to float64, losing + // the verbatim-echo guarantee -- so check the raw wire bytes + // instead. + assert.Contains(t, string(body), tt.wantIDEcho) + } + }) + } +} + +// TestRateLimitedBodyMarshalFallback verifies the defensive fallback: if the +// response fails to marshal (unreachable in production, where requestID is +// always a simple value), rateLimitedBody still returns a valid JSON-RPC +// error with no "id" key (never "id":null) and the retryAfterSeconds value +// correctly interpolated. A channel is not JSON-marshalable and, being +// non-nil, satisfies session.HasJSONRPCID -- forcing the map's marshal to +// fail after "id" was already added. +func TestRateLimitedBodyMarshalFallback(t *testing.T) { + t.Parallel() + + const want = `{"jsonrpc":"2.0","error":{"code":429,"message":"Rate limit exceeded","data":{"retryAfterSeconds":5}}}` + got := rateLimitedBody(make(chan int), 5*time.Second) + assert.Equal(t, want, string(got)) + assert.NotContains(t, string(got), `"id"`) +} + func TestRateLimitHandler_RedisErrorFailOpen(t *testing.T) { t.Parallel() diff --git a/pkg/recovery/recovery.go b/pkg/recovery/recovery.go index 230812038f..7e54b133b0 100644 --- a/pkg/recovery/recovery.go +++ b/pkg/recovery/recovery.go @@ -1,19 +1,18 @@ // SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. // SPDX-License-Identifier: Apache-2.0 -// Package recovery provides panic recovery middleware for HTTP handlers. +// Package recovery adapts toolhive-core's panic recovery middleware to +// ToolHive's middleware factory, wiring in ToolHive's observability: +// OTel span error recording and Sentry issue reporting. package recovery import ( - "errors" - "fmt" - "log/slog" "net/http" - "runtime/debug" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/trace" + corerecovery "github.com/stacklok/toolhive-core/recovery" sentrypkg "github.com/stacklok/toolhive/pkg/sentry" "github.com/stacklok/toolhive/pkg/transport/types" ) @@ -21,37 +20,32 @@ import ( // MiddlewareType is the type constant for recovery middleware const MiddlewareType = "recovery" +// panicHandler wires core's recovery hooks to ToolHive's OTel spans and +// Sentry reporting. +type panicHandler struct{} + +// RecordError records a sanitized error on the request's span. The generic +// message is deliberate: panic values may embed credentials or internal +// state that must not reach external telemetry backends. Full details are +// in the log and in Sentry. +func (panicHandler) RecordError(r *http.Request, err error) { + span := trace.SpanFromContext(r.Context()) + span.RecordError(err) + span.SetStatus(codes.Error, "panic recovered") +} + +// ReportPanic reports the raw panic value to Sentry. The Sentry span +// processor only creates transactions; reporting explicitly makes panics +// also appear as Issues in the Sentry Issues tab. +func (panicHandler) ReportPanic(r *http.Request, v any) { + sentrypkg.RecoverPanic(r, v) +} + // Middleware is an HTTP middleware that recovers from panics. // When a panic occurs, it logs the error and returns // a 500 Internal Server Error response to the client. func Middleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - defer func() { - if rec := recover(); rec != nil { - // Re-panic http.ErrAbortHandler so Go's HTTP server can - // handle it as designed (silently close the connection). - // ReverseProxy panics with this sentinel when a streaming - // response breaks mid-copy; catching it would log noisy - // stack traces and corrupt the already-in-flight response. - if isErrAbortHandler(rec) { - panic(http.ErrAbortHandler) - } - stack := debug.Stack() - slog.Error(fmt.Sprintf("Panic recovered: %v\nStack trace:\n%s", rec, stack)) - span := trace.SpanFromContext(r.Context()) - // Use a generic message on the span to avoid sending potentially - // sensitive panic values (which may embed credentials or internal - // state) to external telemetry backends. Full details are in the log. - span.RecordError(errors.New("panic recovered")) - span.SetStatus(codes.Error, "panic recovered") - // Sentry span processor only creates transactions; call RecoverPanic - // explicitly so panics also appear as Issues in the Sentry Issues tab. - sentrypkg.RecoverPanic(r, rec) - http.Error(w, "Internal Server Error", http.StatusInternalServerError) - } - }() - next.ServeHTTP(w, r) - }) + return corerecovery.Middleware(next, corerecovery.WithPanicHandler(panicHandler{})) } // FactoryMiddleware wraps recovery middleware functionality for the factory pattern. @@ -75,25 +69,3 @@ func CreateMiddleware(_ *types.MiddlewareConfig, runner types.MiddlewareRunner) runner.AddMiddleware(MiddlewareType, recoveryMw) return nil } - -// isErrAbortHandler reports whether rec is the net/http abort-handler sentinel -// or an error wrapping it (see errors.Is). httputil.ReverseProxy uses this -// panic to stop copying a streaming response when the backend or client drops -// the connection. -// -// We must not treat it as a normal panic: logging it as ERROR and calling -// http.Error would run after headers may already be sent (SSE), which produces -// "superfluous response.WriteHeader" and corrupts the response. -func isErrAbortHandler(rec any) bool { - if rec == nil { - return false - } - if rec == http.ErrAbortHandler { - return true - } - err, ok := rec.(error) - if !ok { - return false - } - return errors.Is(err, http.ErrAbortHandler) -} diff --git a/pkg/registry/api/plugins_client.go b/pkg/registry/api/plugins_client.go new file mode 100644 index 0000000000..dac96b11e8 --- /dev/null +++ b/pkg/registry/api/plugins_client.go @@ -0,0 +1,270 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package api + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "net/url" + "strings" + + thvregistry "github.com/stacklok/toolhive-core/registry/types" + "github.com/stacklok/toolhive/pkg/registry/auth" + "github.com/stacklok/toolhive/pkg/versions" +) + +const pluginsBasePath = "/v0.1/x/dev.toolhive/plugins" + +// PluginsListOptions contains options for listing plugins. +type PluginsListOptions struct { + // Search is an optional search query to filter plugins. + Search string + // Limit is the maximum number of plugins per page (default: 100). + Limit int + // Cursor is the pagination cursor for fetching the next page. + Cursor string +} + +// PluginsListResult contains a page of plugins and pagination info. +type PluginsListResult struct { + Plugins []*thvregistry.Plugin + NextCursor string +} + +// PluginsClient provides access to the ToolHive Plugins extension API. +type PluginsClient interface { + // GetPlugin retrieves a plugin by namespace and name (latest version). + GetPlugin(ctx context.Context, namespace, name string) (*thvregistry.Plugin, error) + // GetPluginVersion retrieves a specific version of a plugin. + GetPluginVersion(ctx context.Context, namespace, name, version string) (*thvregistry.Plugin, error) + // ListPlugins retrieves plugins with optional filtering and pagination. + ListPlugins(ctx context.Context, opts *PluginsListOptions) (*PluginsListResult, error) + // SearchPlugins searches for plugins matching the query (auto-paginates). + SearchPlugins(ctx context.Context, query string) (*PluginsListResult, error) + // ListPluginVersions lists all versions of a specific plugin. + ListPluginVersions(ctx context.Context, namespace, name string) (*PluginsListResult, error) +} + +// NewPluginsClient creates a new ToolHive Plugins extension API client. +// If tokenSource is non-nil, the HTTP client transport will be wrapped to inject +// Bearer tokens into all requests. +func NewPluginsClient(baseURL string, allowPrivateIp bool, tokenSource auth.TokenSource) (PluginsClient, error) { + httpClient, err := buildHTTPClient(allowPrivateIp, tokenSource) + if err != nil { + return nil, err + } + + // Ensure base URL doesn't have trailing slash + baseURL = strings.TrimRight(baseURL, "/") + + return &mcpPluginsClient{ + baseURL: baseURL, + httpClient: httpClient, + userAgent: versions.GetUserAgent(), + }, nil +} + +// GetPlugin retrieves a plugin by namespace and name (latest version). +func (c *mcpPluginsClient) GetPlugin(ctx context.Context, namespace, name string) (*thvregistry.Plugin, error) { + endpoint, err := url.JoinPath(c.baseURL, pluginsBasePath, url.PathEscape(namespace), url.PathEscape(name)) + if err != nil { + return nil, fmt.Errorf("failed to build plugins URL: %w", err) + } + + var plugin thvregistry.Plugin + if err := c.doPluginsGet(ctx, endpoint, &plugin); err != nil { + return nil, err + } + return &plugin, nil +} + +// GetPluginVersion retrieves a specific version of a plugin. +func (c *mcpPluginsClient) GetPluginVersion(ctx context.Context, namespace, name, version string) (*thvregistry.Plugin, error) { + endpoint, err := url.JoinPath(c.baseURL, pluginsBasePath, + url.PathEscape(namespace), url.PathEscape(name), + "versions", url.PathEscape(version)) + if err != nil { + return nil, fmt.Errorf("failed to build plugins URL: %w", err) + } + + var plugin thvregistry.Plugin + if err := c.doPluginsGet(ctx, endpoint, &plugin); err != nil { + return nil, err + } + return &plugin, nil +} + +// ListPlugins retrieves plugins with optional filtering and pagination. +// It auto-paginates through all available pages, concatenating results. +func (c *mcpPluginsClient) ListPlugins(ctx context.Context, opts *PluginsListOptions) (*PluginsListResult, error) { + if opts == nil { + opts = &PluginsListOptions{} + } + if opts.Limit == 0 { + opts.Limit = 100 + } + + var allPlugins []*thvregistry.Plugin + cursor := opts.Cursor + + // Pagination loop - continue until no more cursors + for { + page, nextCursor, err := c.fetchPluginsPage(ctx, cursor, opts) + if err != nil { + return nil, err + } + + allPlugins = append(allPlugins, page...) + + // Check if we have more pages + if nextCursor == "" { + break + } + + cursor = nextCursor + + // Safety limit: prevent infinite loops + if len(allPlugins) > 10000 { + return nil, fmt.Errorf("exceeded maximum plugins limit (10000)") + } + } + + return &PluginsListResult{ + Plugins: allPlugins, + }, nil +} + +// SearchPlugins searches for plugins matching the query. +// It auto-paginates through all available pages, concatenating results — +// mirroring ListPlugins. This prevents wrong-publisher installs when +// same-named plugins across namespaces span pages. +func (c *mcpPluginsClient) SearchPlugins(ctx context.Context, query string) (*PluginsListResult, error) { + opts := &PluginsListOptions{Search: query, Limit: 100} + + var allPlugins []*thvregistry.Plugin + cursor := "" + + for { + page, nextCursor, err := c.fetchPluginsPage(ctx, cursor, opts) + if err != nil { + return nil, err + } + + allPlugins = append(allPlugins, page...) + + if nextCursor == "" { + break + } + + cursor = nextCursor + + // Safety limit: prevent infinite loops (mirrors ListPlugins). + if len(allPlugins) > 10000 { + return nil, fmt.Errorf("exceeded maximum plugins limit (10000)") + } + } + + return &PluginsListResult{ + Plugins: allPlugins, + }, nil +} + +// ListPluginVersions lists all versions of a specific plugin. +func (c *mcpPluginsClient) ListPluginVersions(ctx context.Context, namespace, name string) (*PluginsListResult, error) { + endpoint, err := url.JoinPath(c.baseURL, pluginsBasePath, url.PathEscape(namespace), url.PathEscape(name), "versions") + if err != nil { + return nil, fmt.Errorf("failed to build plugins URL: %w", err) + } + + var listResp pluginsListResponse + if err := c.doPluginsGet(ctx, endpoint, &listResp); err != nil { + return nil, err + } + + return &PluginsListResult{ + Plugins: listResp.Plugins, + NextCursor: listResp.Metadata.NextCursor, + }, nil +} + +// mcpPluginsClient implements the PluginsClient interface. +type mcpPluginsClient struct { + baseURL string + httpClient *http.Client + userAgent string +} + +// pluginsListResponse is the wire format for list/search responses. +type pluginsListResponse struct { + Plugins []*thvregistry.Plugin `json:"plugins"` + Metadata struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + } `json:"metadata"` +} + +// doPluginsGet performs an HTTP GET request and decodes the JSON response into dest. +func (c *mcpPluginsClient) doPluginsGet(ctx context.Context, endpoint string, dest any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("User-Agent", c.userAgent) + + resp, err := c.httpClient.Do(req) //nolint:gosec // G704: URL from configured registry + if err != nil { + return fmt.Errorf("failed to execute request: %w", err) + } + defer func() { + if err := resp.Body.Close(); err != nil { + slog.Debug("failed to close response body", "error", err) + } + }() + + if resp.StatusCode != http.StatusOK { + return newRegistryHTTPError(resp) + } + + if err := json.NewDecoder(resp.Body).Decode(dest); err != nil { + return fmt.Errorf("failed to decode response: %w", err) + } + return nil +} + +// fetchPluginsPage fetches a single page of plugins. +func (c *mcpPluginsClient) fetchPluginsPage( + ctx context.Context, cursor string, opts *PluginsListOptions, +) ([]*thvregistry.Plugin, string, error) { + params := url.Values{} + if cursor != "" { + params.Add("cursor", cursor) + } + if opts.Limit > 0 { + params.Add("limit", fmt.Sprintf("%d", opts.Limit)) + } + if opts.Search != "" { + params.Add("search", opts.Search) + } + + basePath, err := url.JoinPath(c.baseURL, pluginsBasePath) + if err != nil { + return nil, "", fmt.Errorf("failed to build plugins URL: %w", err) + } + endpoint := func() string { + if len(params) > 0 { + return basePath + "?" + params.Encode() + } + return basePath + }() + + var listResp pluginsListResponse + if err := c.doPluginsGet(ctx, endpoint, &listResp); err != nil { + return nil, "", err + } + + return listResp.Plugins, listResp.Metadata.NextCursor, nil +} diff --git a/pkg/registry/api/plugins_client_test.go b/pkg/registry/api/plugins_client_test.go new file mode 100644 index 0000000000..30887842df --- /dev/null +++ b/pkg/registry/api/plugins_client_test.go @@ -0,0 +1,627 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package api + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + + thvregistry "github.com/stacklok/toolhive-core/registry/types" +) + +func newTestPluginsClient(t *testing.T, server *httptest.Server) PluginsClient { + t.Helper() + client, err := NewPluginsClient(server.URL, true, nil) + require.NoError(t, err) + return client +} + +func TestPluginsClient_GetPlugin(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + namespace string + pluginName string + handler http.HandlerFunc + wantPlugin *thvregistry.Plugin + wantErr bool + }{ + { + name: "success", + namespace: "io.github.user", + pluginName: "my-plugin", + handler: func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/v0.1/x/dev.toolhive/plugins/io.github.user/my-plugin", r.URL.Path) + require.Equal(t, http.MethodGet, r.Method) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(thvregistry.Plugin{ + Namespace: "io.github.user", + Name: "my-plugin", + Version: "1.0.0", + Description: "A test plugin", + }) + require.NoError(t, err) + }, + wantPlugin: &thvregistry.Plugin{ + Namespace: "io.github.user", + Name: "my-plugin", + Version: "1.0.0", + Description: "A test plugin", + }, + }, + { + name: "not found", + namespace: "io.github.user", + pluginName: "nonexistent", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte("plugin not found")) + }, + wantErr: true, + }, + { + name: "server error", + namespace: "io.github.user", + pluginName: "my-plugin", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte("internal error")) + }, + wantErr: true, + }, + { + name: "path escaping", + namespace: "io.github.user/special", + pluginName: "my plugin", + handler: func(w http.ResponseWriter, r *http.Request) { + // Verify that the path components are properly escaped + require.Equal(t, "/v0.1/x/dev.toolhive/plugins/io.github.user%2Fspecial/my%20plugin", r.URL.RawPath) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(thvregistry.Plugin{ + Namespace: "io.github.user/special", + Name: "my plugin", + Version: "1.0.0", + }) + require.NoError(t, err) + }, + wantPlugin: &thvregistry.Plugin{ + Namespace: "io.github.user/special", + Name: "my plugin", + Version: "1.0.0", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(tt.handler) + defer server.Close() + + client := newTestPluginsClient(t, server) + plugin, err := client.GetPlugin(t.Context(), tt.namespace, tt.pluginName) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, tt.wantPlugin, plugin) + }) + } +} + +func TestPluginsClient_GetPluginVersion(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + namespace string + pluginName string + version string + handler http.HandlerFunc + wantPlugin *thvregistry.Plugin + wantErr bool + }{ + { + name: "success", + namespace: "io.github.user", + pluginName: "my-plugin", + version: "2.0.0", + handler: func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/v0.1/x/dev.toolhive/plugins/io.github.user/my-plugin/versions/2.0.0", r.URL.Path) + require.Equal(t, http.MethodGet, r.Method) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(thvregistry.Plugin{ + Namespace: "io.github.user", + Name: "my-plugin", + Version: "2.0.0", + Description: "Version 2", + }) + require.NoError(t, err) + }, + wantPlugin: &thvregistry.Plugin{ + Namespace: "io.github.user", + Name: "my-plugin", + Version: "2.0.0", + Description: "Version 2", + }, + }, + { + name: "version not found", + namespace: "io.github.user", + pluginName: "my-plugin", + version: "99.0.0", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte("version not found")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(tt.handler) + defer server.Close() + + client := newTestPluginsClient(t, server) + plugin, err := client.GetPluginVersion(t.Context(), tt.namespace, tt.pluginName, tt.version) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, tt.wantPlugin, plugin) + }) + } +} + +func TestPluginsClient_ListPlugins(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts *PluginsListOptions + handler http.HandlerFunc + wantCount int + wantErr bool + wantPlugins []*thvregistry.Plugin + }{ + { + name: "single page", + opts: nil, + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.a", Name: "plugin-1", Version: "1.0.0"}, + {Namespace: "io.github.b", Name: "plugin-2", Version: "1.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 2, NextCursor: ""}, + }) + require.NoError(t, err) + }, + wantCount: 2, + wantPlugins: []*thvregistry.Plugin{ + {Namespace: "io.github.a", Name: "plugin-1", Version: "1.0.0"}, + {Namespace: "io.github.b", Name: "plugin-2", Version: "1.0.0"}, + }, + }, + { + name: "pagination across multiple pages", + opts: &PluginsListOptions{Limit: 1}, + handler: func() http.HandlerFunc { + callCount := 0 + return func(w http.ResponseWriter, r *http.Request) { + callCount++ + w.Header().Set("Content-Type", "application/json") + + cursor := r.URL.Query().Get("cursor") + var resp pluginsListResponse + + switch { + case cursor == "" && callCount == 1: + resp = pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.a", Name: "plugin-1", Version: "1.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 1, NextCursor: "page2"}, + } + case cursor == "page2": + resp = pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.b", Name: "plugin-2", Version: "1.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 1, NextCursor: ""}, + } + default: + w.WriteHeader(http.StatusBadRequest) + return + } + + err := json.NewEncoder(w).Encode(resp) + require.NoError(t, err) + } + }(), + wantCount: 2, + wantPlugins: []*thvregistry.Plugin{ + {Namespace: "io.github.a", Name: "plugin-1", Version: "1.0.0"}, + {Namespace: "io.github.b", Name: "plugin-2", Version: "1.0.0"}, + }, + }, + { + name: "empty result", + opts: nil, + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(pluginsListResponse{ + Plugins: []*thvregistry.Plugin{}, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 0, NextCursor: ""}, + }) + require.NoError(t, err) + }, + wantCount: 0, + wantPlugins: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(tt.handler) + defer server.Close() + + client := newTestPluginsClient(t, server) + result, err := client.ListPlugins(t.Context(), tt.opts) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Len(t, result.Plugins, tt.wantCount) + if tt.wantPlugins != nil { + require.Equal(t, tt.wantPlugins, result.Plugins) + } + }) + } +} + +func TestPluginsClient_SearchPlugins(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + query string + handler http.HandlerFunc + wantCount int + wantErr bool + }{ + { + name: "success with results", + query: "kubernetes", + handler: func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "kubernetes", r.URL.Query().Get("search")) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.user", Name: "k8s-plugin", Version: "1.0.0", Description: "Kubernetes plugin"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 1, NextCursor: ""}, + }) + require.NoError(t, err) + }, + wantCount: 1, + }, + { + name: "empty result", + query: "nonexistent", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(pluginsListResponse{ + Plugins: []*thvregistry.Plugin{}, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 0, NextCursor: ""}, + }) + require.NoError(t, err) + }, + wantCount: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(tt.handler) + defer server.Close() + + client := newTestPluginsClient(t, server) + result, err := client.SearchPlugins(t.Context(), tt.query) + + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Len(t, result.Plugins, tt.wantCount) + }) + } +} + +// TestPluginsClient_SearchPluginsPagination verifies that SearchPlugins +// auto-paginates through all available pages, mirroring ListPlugins. Same-named +// plugins across namespaces can span pages; a single-page search would miss the +// later pages and enable wrong-publisher installs. +func TestPluginsClient_SearchPluginsPagination(t *testing.T) { + t.Parallel() + + callCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "kube", r.URL.Query().Get("search")) + w.Header().Set("Content-Type", "application/json") + + callCount++ + cursor := r.URL.Query().Get("cursor") + var resp pluginsListResponse + + switch { + case cursor == "" && callCount == 1: + resp = pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.a", Name: "k8s-plugin", Version: "1.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 1, NextCursor: "page2"}, + } + case cursor == "page2": + resp = pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.b", Name: "k8s-plugin", Version: "2.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 1, NextCursor: ""}, + } + default: + w.WriteHeader(http.StatusBadRequest) + return + } + + err := json.NewEncoder(w).Encode(resp) + require.NoError(t, err) + })) + defer server.Close() + + client := newTestPluginsClient(t, server) + result, err := client.SearchPlugins(t.Context(), "kube") + require.NoError(t, err) + require.Len(t, result.Plugins, 2, "both pages should be concatenated") + require.Equal(t, "io.github.a", result.Plugins[0].Namespace) + require.Equal(t, "io.github.b", result.Plugins[1].Namespace) + require.Equal(t, 2, callCount, "both pages should have been fetched") +} + +func TestPluginsClient_ListPluginVersions(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/v0.1/x/dev.toolhive/plugins/io.github.user/my-plugin/versions", r.URL.Path) + require.Equal(t, http.MethodGet, r.Method) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.user", Name: "my-plugin", Version: "1.0.0"}, + {Namespace: "io.github.user", Name: "my-plugin", Version: "2.0.0"}, + {Namespace: "io.github.user", Name: "my-plugin", Version: "3.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 3, NextCursor: ""}, + }) + require.NoError(t, err) + })) + defer server.Close() + + client := newTestPluginsClient(t, server) + result, err := client.ListPluginVersions(t.Context(), "io.github.user", "my-plugin") + require.NoError(t, err) + require.Len(t, result.Plugins, 3) + require.Equal(t, "1.0.0", result.Plugins[0].Version) + require.Equal(t, "2.0.0", result.Plugins[1].Version) + require.Equal(t, "3.0.0", result.Plugins[2].Version) +} + +func TestPluginsClient_ErrorHandling(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + statusCode int + body string + wantErrIs error + }{ + { + name: "401 unauthorized", + statusCode: http.StatusUnauthorized, + body: "unauthorized", + wantErrIs: ErrRegistryUnauthorized, + }, + { + name: "403 forbidden", + statusCode: http.StatusForbidden, + body: "forbidden", + wantErrIs: ErrRegistryUnauthorized, + }, + { + name: "404 not found", + statusCode: http.StatusNotFound, + body: "not found", + wantErrIs: nil, + }, + { + name: "500 server error does not unwrap to unauthorized", + statusCode: http.StatusInternalServerError, + body: "internal server error", + wantErrIs: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.body)) + })) + defer server.Close() + + client := newTestPluginsClient(t, server) + _, err := client.GetPlugin(t.Context(), "io.github.user", "my-plugin") + require.Error(t, err) + + var httpErr *RegistryHTTPError + require.True(t, errors.As(err, &httpErr), "expected *RegistryHTTPError, got %T", err) + require.Equal(t, tt.statusCode, httpErr.StatusCode) + require.Contains(t, httpErr.Body, tt.body) + + if tt.wantErrIs != nil { + require.True(t, errors.Is(err, tt.wantErrIs), + "expected errors.Is(%v, %v) to be true", err, tt.wantErrIs) + } else { + require.False(t, errors.Is(err, ErrRegistryUnauthorized), + "expected errors.Is(%v, ErrRegistryUnauthorized) to be false", err) + } + }) + } +} + +func TestPluginsClient_MalformedJSON(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{invalid json`)) + })) + defer server.Close() + + client := newTestPluginsClient(t, server) + _, err := client.GetPlugin(t.Context(), "io.github.user", "my-plugin") + require.Error(t, err) + require.Contains(t, err.Error(), "failed to decode response") +} + +func TestPluginsClient_TrailingSlashInBaseURL(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // The path should not have a double slash + require.NotContains(t, r.URL.Path, "//") + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(thvregistry.Plugin{ + Namespace: "io.github.user", + Name: "my-plugin", + Version: "1.0.0", + }) + require.NoError(t, err) + })) + defer server.Close() + + // Create client with trailing slash + client, err := NewPluginsClient(server.URL+"/", true, nil) + require.NoError(t, err) + + plugin, err := client.GetPlugin(t.Context(), "io.github.user", "my-plugin") + require.NoError(t, err) + require.Equal(t, "io.github.user", plugin.Namespace) +} + +func TestPluginsClient_ListPluginsWithSearch(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "test-query", r.URL.Query().Get("search")) + require.Equal(t, "50", r.URL.Query().Get("limit")) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(pluginsListResponse{ + Plugins: []*thvregistry.Plugin{ + {Namespace: "io.github.user", Name: "test-plugin", Version: "1.0.0"}, + }, + Metadata: struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + }{Count: 1, NextCursor: ""}, + }) + require.NoError(t, err) + })) + defer server.Close() + + client := newTestPluginsClient(t, server) + result, err := client.ListPlugins(t.Context(), &PluginsListOptions{ + Search: "test-query", + Limit: 50, + }) + require.NoError(t, err) + require.Len(t, result.Plugins, 1) + require.Equal(t, "test-plugin", result.Plugins[0].Name) +} + +func TestPluginsClient_RegistryHTTPErrorUnwrap(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + statusCode int + wantErrIs error + }{ + {name: "401 wraps unauthorized", statusCode: http.StatusUnauthorized, wantErrIs: ErrRegistryUnauthorized}, + {name: "403 wraps unauthorized", statusCode: http.StatusForbidden, wantErrIs: ErrRegistryUnauthorized}, + {name: "404 unwraps to nil", statusCode: http.StatusNotFound, wantErrIs: nil}, + {name: "500 unwraps to nil", statusCode: http.StatusInternalServerError, wantErrIs: nil}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := &RegistryHTTPError{ + StatusCode: tt.statusCode, + Body: "test body", + URL: "http://example.com/test", + } + require.Equal(t, tt.wantErrIs, err.Unwrap()) + require.Contains(t, err.Error(), fmt.Sprintf("status %d", tt.statusCode)) + }) + } +} diff --git a/pkg/registry/auth/login.go b/pkg/registry/auth/login.go index 3bb4447c09..b952d78470 100644 --- a/pkg/registry/auth/login.go +++ b/pkg/registry/auth/login.go @@ -290,7 +290,10 @@ func ConfigureOAuth( discoveryCtx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() - if _, err := oauth.DiscoverOIDCEndpoints(discoveryCtx, issuer); err != nil { + // blockPrivateIPs=false preserves this call path's existing behavior: + // issuer is an operator-supplied flag here, not remote-server-derived + // discovery input, unlike the CLI DCR flow's use of DiscoverOIDCEndpoints. + if _, err := oauth.DiscoverOIDCEndpoints(discoveryCtx, issuer, false); err != nil { return nil, fmt.Errorf("OIDC discovery failed for issuer %s: %w", issuer, err) } diff --git a/pkg/registry/convert_test.go b/pkg/registry/convert_test.go index 6ed75215fb..fd1a08751d 100644 --- a/pkg/registry/convert_test.go +++ b/pkg/registry/convert_test.go @@ -332,7 +332,7 @@ func TestConvertJSON_OutputPassesSchemaValidation(t *testing.T) { func TestConvertJSON_RoundTripEmbeddedCatalog(t *testing.T) { t.Parallel() - original, _, err := parseRegistryData(catalog.Upstream()) + original, _, _, err := parseRegistryData(catalog.Upstream()) require.NoError(t, err) roundTripped, err := converters.NewUpstreamRegistryFromToolhiveRegistry(original) diff --git a/pkg/registry/mocks/mock_provider.go b/pkg/registry/mocks/mock_provider.go index 5fc97deadf..57bc3d47ad 100644 --- a/pkg/registry/mocks/mock_provider.go +++ b/pkg/registry/mocks/mock_provider.go @@ -40,6 +40,21 @@ func (m *MockProvider) EXPECT() *MockProviderMockRecorder { return m.recorder } +// GetPlugin mocks base method. +func (m *MockProvider) GetPlugin(namespace, name string) (*registry.Plugin, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetPlugin", namespace, name) + ret0, _ := ret[0].(*registry.Plugin) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetPlugin indicates an expected call of GetPlugin. +func (mr *MockProviderMockRecorder) GetPlugin(namespace, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPlugin", reflect.TypeOf((*MockProvider)(nil).GetPlugin), namespace, name) +} + // GetRegistry mocks base method. func (m *MockProvider) GetRegistry() (*registry.Registry, error) { m.ctrl.T.Helper() @@ -85,6 +100,21 @@ func (mr *MockProviderMockRecorder) GetSkill(namespace, name any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSkill", reflect.TypeOf((*MockProvider)(nil).GetSkill), namespace, name) } +// ListAvailablePlugins mocks base method. +func (m *MockProvider) ListAvailablePlugins() ([]registry.Plugin, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAvailablePlugins") + ret0, _ := ret[0].([]registry.Plugin) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAvailablePlugins indicates an expected call of ListAvailablePlugins. +func (mr *MockProviderMockRecorder) ListAvailablePlugins() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAvailablePlugins", reflect.TypeOf((*MockProvider)(nil).ListAvailablePlugins)) +} + // ListAvailableSkills mocks base method. func (m *MockProvider) ListAvailableSkills() ([]registry.Skill, error) { m.ctrl.T.Helper() @@ -115,6 +145,21 @@ func (mr *MockProviderMockRecorder) ListServers() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListServers", reflect.TypeOf((*MockProvider)(nil).ListServers)) } +// SearchPlugins mocks base method. +func (m *MockProvider) SearchPlugins(query string) ([]registry.Plugin, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SearchPlugins", query) + ret0, _ := ret[0].([]registry.Plugin) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// SearchPlugins indicates an expected call of SearchPlugins. +func (mr *MockProviderMockRecorder) SearchPlugins(query any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SearchPlugins", reflect.TypeOf((*MockProvider)(nil).SearchPlugins), query) +} + // SearchServers mocks base method. func (m *MockProvider) SearchServers(query string) ([]registry.ServerMetadata, error) { m.ctrl.T.Helper() diff --git a/pkg/registry/provider.go b/pkg/registry/provider.go index 1cffb4ca61..6194416dde 100644 --- a/pkg/registry/provider.go +++ b/pkg/registry/provider.go @@ -29,4 +29,13 @@ type Provider interface { // SearchSkills searches for skills matching the query SearchSkills(query string) ([]types.Skill, error) + + // ListAvailablePlugins returns plugins discovered from the registry data + ListAvailablePlugins() ([]types.Plugin, error) + + // GetPlugin returns a specific plugin by namespace and name + GetPlugin(namespace, name string) (*types.Plugin, error) + + // SearchPlugins searches for plugins matching the query + SearchPlugins(query string) ([]types.Plugin, error) } diff --git a/pkg/registry/provider_api.go b/pkg/registry/provider_api.go index f8d6b6270e..ee4355eaea 100644 --- a/pkg/registry/provider_api.go +++ b/pkg/registry/provider_api.go @@ -26,6 +26,7 @@ type APIRegistryProvider struct { client api.Client tokenSource auth.TokenSource skillsClient api.SkillsClient + pluginsClient api.PluginsClient } // NewAPIRegistryProvider creates a new API registry provider. @@ -40,12 +41,16 @@ func NewAPIRegistryProvider(apiURL string, allowPrivateIp bool, tokenSource auth // Create skills client (best-effort — skills API may not be available) skillsClient, _ := api.NewSkillsClient(apiURL, allowPrivateIp, tokenSource) + // Create plugins client (best-effort — plugins API may not be available) + pluginsClient, _ := api.NewPluginsClient(apiURL, allowPrivateIp, tokenSource) + p := &APIRegistryProvider{ apiURL: apiURL, allowPrivateIp: allowPrivateIp, client: client, tokenSource: tokenSource, skillsClient: skillsClient, + pluginsClient: pluginsClient, } // Initialize the base provider with the GetRegistry function @@ -227,6 +232,56 @@ func (p *APIRegistryProvider) SearchSkills(query string) ([]types.Skill, error) return skills, nil } +// ListAvailablePlugins returns all plugins from the API, with auto-pagination. +func (p *APIRegistryProvider) ListAvailablePlugins() ([]types.Plugin, error) { + if p.pluginsClient == nil { + return nil, nil + } + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + result, err := p.pluginsClient.ListPlugins(ctx, nil) + if err != nil { + return nil, err + } + plugins := make([]types.Plugin, 0, len(result.Plugins)) + for _, pl := range result.Plugins { + if pl != nil { + plugins = append(plugins, *pl) + } + } + return plugins, nil +} + +// GetPlugin returns a specific plugin by namespace and name from the API. +func (p *APIRegistryProvider) GetPlugin(namespace, name string) (*types.Plugin, error) { + if p.pluginsClient == nil { + return nil, nil + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + return p.pluginsClient.GetPlugin(ctx, namespace, name) +} + +// SearchPlugins searches for plugins matching the query via the API. +func (p *APIRegistryProvider) SearchPlugins(query string) ([]types.Plugin, error) { + if p.pluginsClient == nil { + return nil, nil + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + result, err := p.pluginsClient.SearchPlugins(ctx, query) + if err != nil { + return nil, err + } + plugins := make([]types.Plugin, 0, len(result.Plugins)) + for _, pl := range result.Plugins { + if pl != nil { + plugins = append(plugins, *pl) + } + } + return plugins, nil +} + // ConvertServerJSON converts an MCP Registry API ServerJSON to ToolHive ServerMetadata // Uses converters from converters.go (same package) // Note: Only handles OCI packages and remote servers, skips npm/pypi by design diff --git a/pkg/registry/provider_api_test.go b/pkg/registry/provider_api_test.go new file mode 100644 index 0000000000..9805a1b6b4 --- /dev/null +++ b/pkg/registry/provider_api_test.go @@ -0,0 +1,199 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package registry + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + thvregistry "github.com/stacklok/toolhive-core/registry/types" +) + +// apiRegistryTestServer serves the MCP Registry API endpoints needed to exercise +// APIRegistryProvider plugin methods. It serves the servers list (for the +// constructor's validation probe) plus the ToolHive plugins extension endpoints. +func apiRegistryTestServer(t *testing.T, plugins []*thvregistry.Plugin) *httptest.Server { + t.Helper() + + mux := http.NewServeMux() + + // Servers list endpoint — required by the constructor's validation probe. + mux.HandleFunc("/v0.1/servers", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"servers":[],"metadata":{"next_cursor":""}}`)) + }) + + // List plugins (auto-pagination) — required by ListAvailablePlugins. + mux.HandleFunc("/v0.1/x/dev.toolhive/plugins", func(w http.ResponseWriter, r *http.Request) { + // If there's a search query param, this is the SearchPlugins endpoint. + if r.URL.Query().Get("search") != "" { + searchHandler(w, r, plugins) + return + } + listHandler(w, r, plugins) + }) + + // Get plugin by namespace/name — required by GetPlugin. + mux.HandleFunc("/v0.1/x/dev.toolhive/plugins/", func(w http.ResponseWriter, _ *http.Request) { + // The path is /v0.1/x/dev.toolhive/plugins/{namespace}/{name} + // We serve a fixed plugin for any request (the handler extracts the + // namespace/name from the path but we just return the first plugin + // that matches, or 404). + w.Header().Set("Content-Type", "application/json") + + // Simple match: if plugins has at least one, return the first one. + // The GetPlugin implementation extracts namespace/name from the + // path; we don't need to validate them here — we just need to + // return a valid plugin payload so the caller can assert it. + if len(plugins) > 0 { + require.NoError(t, json.NewEncoder(w).Encode(plugins[0])) + return + } + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte("plugin not found")) + }) + + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +func searchHandler(w http.ResponseWriter, _ *http.Request, plugins []*thvregistry.Plugin) { + w.Header().Set("Content-Type", "application/json") + payload := pluginsListPayload{Plugins: plugins} + payload.Metadata.Count = len(plugins) + _ = json.NewEncoder(w).Encode(payload) +} + +func listHandler(w http.ResponseWriter, _ *http.Request, plugins []*thvregistry.Plugin) { + w.Header().Set("Content-Type", "application/json") + payload := pluginsListPayload{Plugins: plugins} + payload.Metadata.Count = len(plugins) + _ = json.NewEncoder(w).Encode(payload) +} + +func TestAPIRegistryProvider_ListAvailablePlugins(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + { + Namespace: "io.github.stacklok", + Name: "code-reviewer", + Description: "Reviews code for bugs", + Version: "1.0.0", + }, + { + Namespace: "io.github.acme", + Name: "doc-generator", + Description: "Generates documentation", + Version: "0.2.0", + }, + } + + srv := apiRegistryTestServer(t, plugins) + provider, err := NewAPIRegistryProvider(srv.URL, true, nil) + require.NoError(t, err) + require.NotNil(t, provider) + + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 2) + + assert.Equal(t, "code-reviewer", got[0].Name) + assert.Equal(t, "io.github.stacklok", got[0].Namespace) + assert.Equal(t, "Reviews code for bugs", got[0].Description) + assert.Equal(t, "1.0.0", got[0].Version) + + assert.Equal(t, "doc-generator", got[1].Name) + assert.Equal(t, "io.github.acme", got[1].Namespace) + assert.Equal(t, "0.2.0", got[1].Version) +} + +func TestAPIRegistryProvider_GetPlugin(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + { + Namespace: "io.github.stacklok", + Name: "code-reviewer", + Description: "Reviews code for bugs", + Version: "1.0.0", + }, + } + + srv := apiRegistryTestServer(t, plugins) + provider, err := NewAPIRegistryProvider(srv.URL, true, nil) + require.NoError(t, err) + + got, err := provider.GetPlugin("io.github.stacklok", "code-reviewer") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, "code-reviewer", got.Name) + assert.Equal(t, "io.github.stacklok", got.Namespace) + assert.Equal(t, "1.0.0", got.Version) +} + +func TestAPIRegistryProvider_SearchPlugins(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + { + Namespace: "io.github.stacklok", + Name: "code-reviewer", + Description: "Reviews code for bugs", + Version: "1.0.0", + }, + } + + srv := apiRegistryTestServer(t, plugins) + provider, err := NewAPIRegistryProvider(srv.URL, true, nil) + require.NoError(t, err) + + got, err := provider.SearchPlugins("reviewer") + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, "code-reviewer", got[0].Name) + assert.Equal(t, "Reviews code for bugs", got[0].Description) +} + +func TestAPIRegistryProvider_SearchPlugins_Empty(t *testing.T) { + t.Parallel() + + srv := apiRegistryTestServer(t, nil) + provider, err := NewAPIRegistryProvider(srv.URL, true, nil) + require.NoError(t, err) + + got, err := provider.SearchPlugins("nonexistent") + require.NoError(t, err) + assert.Empty(t, got) +} + +func TestAPIRegistryProvider_PluginsNilClient(t *testing.T) { + t.Parallel() + + // Construct an APIRegistryProvider with a nil pluginsClient (simulating + // the case where NewPluginsClient fails, e.g. due to HTTP transport setup). + // The constructor does best-effort plugin client creation, so a nil client + // is possible. The methods must return nil, nil rather than panicking. + p := &APIRegistryProvider{ + pluginsClient: nil, + } + + plugins, err := p.ListAvailablePlugins() + require.NoError(t, err) + assert.Nil(t, plugins) + + got, err := p.GetPlugin("ns", "name") + require.NoError(t, err) + assert.Nil(t, got) + + results, err := p.SearchPlugins("query") + require.NoError(t, err) + assert.Nil(t, results) +} diff --git a/pkg/registry/provider_base.go b/pkg/registry/provider_base.go index 2585b7f188..5619a487f5 100644 --- a/pkg/registry/provider_base.go +++ b/pkg/registry/provider_base.go @@ -136,6 +136,22 @@ func (*BaseProvider) SearchSkills(_ string) ([]types.Skill, error) { return nil, nil } +// ListAvailablePlugins returns an empty slice by default. +// Providers that support plugins (local, remote) override this. +func (*BaseProvider) ListAvailablePlugins() ([]types.Plugin, error) { + return nil, nil +} + +// GetPlugin returns nil for providers that don't support plugins. +func (*BaseProvider) GetPlugin(_, _ string) (*types.Plugin, error) { + return nil, nil +} + +// SearchPlugins returns nil for providers that don't support plugins. +func (*BaseProvider) SearchPlugins(_ string) ([]types.Plugin, error) { + return nil, nil +} + // matchesQuery checks if a server matches the search query func matchesQuery(name, description string, tags []string, query string) bool { // Search in name diff --git a/pkg/registry/provider_cached.go b/pkg/registry/provider_cached.go index e22001be49..c71f7cb096 100644 --- a/pkg/registry/provider_cached.go +++ b/pkg/registry/provider_cached.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "fmt" + "net/http" "os" "path/filepath" "sync" @@ -46,6 +47,12 @@ type CachedAPIRegistryProvider struct { skillsCacheSet bool skillsTime time.Time + // Plugins cache + pluginsMu sync.RWMutex + cachedPlugins []types.Plugin + pluginsCacheSet bool + pluginsTime time.Time + // Cache configuration cacheTTL time.Duration usePersistent bool @@ -149,8 +156,24 @@ func (p *CachedAPIRegistryProvider) refreshCache() (*types.Registry, error) { return registry, nil } -// ForceRefresh forces a cache refresh, ignoring TTL. +// ForceRefresh forces a cache refresh, ignoring TTL. It also invalidates the +// skills and plugins caches so the next ListAvailableSkills/ListAvailablePlugins +// call re-fetches from the API rather than serving stale data until its own +// TTL expires. func (p *CachedAPIRegistryProvider) ForceRefresh() error { + // Invalidate the skills and plugins caches so the next access refetches. + // Without this, ForceRefresh would only refresh the servers cache while + // skills/plugins kept serving stale data for up to cacheTTL after the + // refresh — a confusing inconsistency for callers who expect "force" to + // mean "everything is refreshed". + p.skillsMu.Lock() + p.skillsCacheSet = false + p.skillsMu.Unlock() + + p.pluginsMu.Lock() + p.pluginsCacheSet = false + p.pluginsMu.Unlock() + _, err := p.refreshCache() return err } @@ -447,6 +470,80 @@ func (p *CachedAPIRegistryProvider) ListAvailableSkills() ([]types.Skill, error) return allSkills, nil } +// ListAvailablePlugins returns plugins from the registry API, with caching. +// Creates a PluginsClient on demand and fetches all plugins with auto-pagination. +// +// Error semantics mirror refreshCache's contract for the servers cache: +// - authentication failures (401/403, surfaced as api.RegistryHTTPError that +// unwraps to api.ErrRegistryUnauthorized) are always propagated — stale +// cache must never mask a changed authentication state, or a revoked token +// would silently serve stale data and hide the need to re-auth; +// - other failures (network blip, 5xx) degrade gracefully to stale cache +// when one is present; +// - with no stale cache, the error is returned (never nil,nil), so the v0.1 +// registry route surfaces a real failure instead of an empty 200. +func (p *CachedAPIRegistryProvider) ListAvailablePlugins() ([]types.Plugin, error) { + // Check cache + p.pluginsMu.RLock() + if p.pluginsCacheSet && time.Since(p.pluginsTime) < p.cacheTTL { + plugins := p.cachedPlugins + p.pluginsMu.RUnlock() + return plugins, nil + } + p.pluginsMu.RUnlock() + + // Fetch from API + pluginsClient, err := api.NewPluginsClient(p.apiURL, p.allowPrivateIp, p.tokenSource) + if err != nil { + // Client construction is a local failure (bad URL / transport), not an + // auth-state change: fall back to stale cache when available. + p.pluginsMu.RLock() + defer p.pluginsMu.RUnlock() + if p.pluginsCacheSet { + return p.cachedPlugins, nil + } + return nil, fmt.Errorf("failed to create plugins client: %w", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + // ListPlugins auto-paginates internally, returning all plugins in one call + result, err := pluginsClient.ListPlugins(ctx, nil) + if err != nil { + // Auth failures must propagate — never mask with stale cache. + var httpErr *api.RegistryHTTPError + if errors.As(err, &httpErr) && (httpErr.StatusCode == http.StatusUnauthorized || httpErr.StatusCode == http.StatusForbidden) { + return nil, err + } + // Transient failures: degrade to stale cache if available. + p.pluginsMu.RLock() + defer p.pluginsMu.RUnlock() + if p.pluginsCacheSet { + return p.cachedPlugins, nil + } + // No stale cache: surface the error rather than nil,nil so the + // v0.1 registry route does not answer 200 [] on a real failure. + return nil, err + } + + allPlugins := make([]types.Plugin, 0, len(result.Plugins)) + for _, pl := range result.Plugins { + if pl != nil { + allPlugins = append(allPlugins, *pl) + } + } + + // Update cache + p.pluginsMu.Lock() + p.cachedPlugins = allPlugins + p.pluginsCacheSet = true + p.pluginsTime = time.Now() + p.pluginsMu.Unlock() + + return allPlugins, nil +} + // ConvertServerJSON wraps ConvertServerJSON for cached provider func (*CachedAPIRegistryProvider) ConvertServerJSON(serverJSON *v0.ServerJSON) (types.ServerMetadata, error) { return ConvertServerJSON(serverJSON) diff --git a/pkg/registry/provider_cached_plugins_errors_test.go b/pkg/registry/provider_cached_plugins_errors_test.go new file mode 100644 index 0000000000..6405d8f93f --- /dev/null +++ b/pkg/registry/provider_cached_plugins_errors_test.go @@ -0,0 +1,170 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package registry + +import ( + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + thvregistry "github.com/stacklok/toolhive-core/registry/types" + "github.com/stacklok/toolhive/pkg/registry/api" +) + +// newPluginsStatusServer returns an httptest.Server whose plugins list endpoint +// responds with statusCode when status.Load() is non-zero, and a normal page of +// plugins otherwise. The servers list endpoint always answers an empty list so +// the constructor's validation probe succeeds. +func newPluginsStatusServer(t *testing.T, plugins []*thvregistry.Plugin, status *atomic.Int32) (*httptest.Server, *atomic.Int32) { + t.Helper() + var callCount atomic.Int32 + + mux := http.NewServeMux() + mux.HandleFunc("/v0.1/x/dev.toolhive/plugins", func(w http.ResponseWriter, _ *http.Request) { + callCount.Add(1) + if s := status.Load(); s != 0 { + w.WriteHeader(int(s)) + return + } + w.Header().Set("Content-Type", "application/json") + payload := pluginsListPayload{Plugins: plugins} + payload.Metadata.Count = len(plugins) + require.NoError(t, json.NewEncoder(w).Encode(payload)) + }) + mux.HandleFunc("/v0.1/servers", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"servers":[],"metadata":{"next_cursor":""}}`)) + }) + + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv, &callCount +} + +// TestCachedProvider_PluginsAuthErrorWithWarmCache verifies that a 401 from the +// registry is propagated even when a warm cache could otherwise mask it. A +// revoked token must not silently serve stale data. +func TestCachedProvider_PluginsAuthErrorWithWarmCache(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + {Namespace: "io.github.stacklok", Name: "code-reviewer", Version: "1.0.0"}, + } + var status atomic.Int32 // 0 = healthy + srv, callCount := newPluginsStatusServer(t, plugins, &status) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + // Warm the cache. + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 1) + firstCount := callCount.Load() + + // Flip the registry to 401 and expire the cache. + status.Store(http.StatusUnauthorized) + provider.pluginsMu.Lock() + provider.pluginsTime = provider.pluginsTime.Add(-defaultCacheTTL - 1) + provider.pluginsMu.Unlock() + + got2, err := provider.ListAvailablePlugins() + require.Error(t, err, "401 must propagate, not be masked by stale cache") + require.Nil(t, got2) + require.True(t, + errors.Is(err, api.ErrRegistryUnauthorized), + "expected errors.Is(err, ErrRegistryUnauthorized); got: %v", err) + require.Greater(t, callCount.Load(), firstCount, "expired cache should trigger a re-fetch attempt") +} + +// TestCachedProvider_PluginsTransient500WithWarmCache verifies that a transient +// 500 with a warm cache degrades to stale data (the auth path is the exception, +// not the rule). +func TestCachedProvider_PluginsTransient500WithWarmCache(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + {Namespace: "io.github.stacklok", Name: "code-reviewer", Version: "1.0.0"}, + } + var status atomic.Int32 + srv, callCount := newPluginsStatusServer(t, plugins, &status) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + // Warm the cache. + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 1) + + // Flip to 500 and expire the cache. + status.Store(http.StatusInternalServerError) + provider.pluginsMu.Lock() + provider.pluginsTime = provider.pluginsTime.Add(-defaultCacheTTL - 1) + provider.pluginsMu.Unlock() + + got2, err := provider.ListAvailablePlugins() + require.NoError(t, err, "transient 500 with warm cache should serve stale data") + require.Len(t, got2, 1) + require.Equal(t, "code-reviewer", got2[0].Name) + _ = callCount // re-fetch attempt made (not asserting exact count here) +} + +// TestCachedProvider_PluginsTransient500WithColdCache verifies that a transient +// 500 with no cache returns the error rather than (nil, nil), so the v0.1 +// registry route does not answer 200 [] on a real failure. +func TestCachedProvider_PluginsTransient500WithColdCache(t *testing.T) { + t.Parallel() + + var status atomic.Int32 + status.Store(http.StatusInternalServerError) + srv, _ := newPluginsStatusServer(t, nil, &status) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + got, err := provider.ListAvailablePlugins() + require.Error(t, err, "transient 500 with cold cache must return the error, not nil,nil") + require.Nil(t, got) +} + +// TestCachedProvider_ForceRefreshInvalidatesPluginsCache verifies that +// ForceRefresh invalidates the plugins cache so the next ListAvailablePlugins +// re-fetches from the API instead of serving the pre-refresh cached value. +func TestCachedProvider_ForceRefreshInvalidatesPluginsCache(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + {Namespace: "io.github.stacklok", Name: "code-reviewer", Version: "1.0.0"}, + } + srv, callCount := newPluginsStatusServer(t, plugins, &atomic.Int32{}) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + // Prime the plugins cache. + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 1) + afterFirst := callCount.Load() + + // A second call within TTL must be served from cache (no re-fetch). + got2, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Equal(t, got, got2) + require.Equal(t, afterFirst, callCount.Load(), "second call should be served from cache") + + // ForceRefresh must invalidate the plugins cache so the next call refetches. + require.NoError(t, provider.ForceRefresh()) + + got3, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got3, 1) + require.Greater(t, callCount.Load(), afterFirst, "ForceRefresh should cause the next ListAvailablePlugins to refetch") +} diff --git a/pkg/registry/provider_cached_plugins_test.go b/pkg/registry/provider_cached_plugins_test.go new file mode 100644 index 0000000000..f23cc1ae28 --- /dev/null +++ b/pkg/registry/provider_cached_plugins_test.go @@ -0,0 +1,161 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package registry + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + thvregistry "github.com/stacklok/toolhive-core/registry/types" +) + +// pluginsListPayload is the wire format served by the test registry for the +// plugins list endpoint. Field names mirror pluginsListResponse in +// pkg/registry/api/plugins_client.go so the client decodes them cleanly. +type pluginsListPayload struct { + Plugins []*thvregistry.Plugin `json:"plugins"` + Metadata struct { + Count int `json:"count"` + NextCursor string `json:"nextCursor"` + } `json:"metadata"` +} + +// newPluginsTestServer returns an httptest.Server that serves the plugins list +// endpoint. The server endpoint is /v0.1/x/dev.toolhive/plugins. If fail is +// non-nil and returns true, the server responds 500 to exercise stale-cache +// fallback. callCount tracks how many times the list endpoint was hit. +func newPluginsTestServer(t *testing.T, plugins []*thvregistry.Plugin, fail *atomic.Bool) (*httptest.Server, *atomic.Int32) { + t.Helper() + var callCount atomic.Int32 + + mux := http.NewServeMux() + mux.HandleFunc("/v0.1/x/dev.toolhive/plugins", func(w http.ResponseWriter, _ *http.Request) { + callCount.Add(1) + if fail != nil && fail.Load() { + w.WriteHeader(http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + payload := pluginsListPayload{Plugins: plugins} + payload.Metadata.Count = len(plugins) + require.NoError(t, json.NewEncoder(w).Encode(payload)) + }) + // The constructor's validation probe hits the servers list endpoint; serve + // an empty server list so the probe succeeds and construction completes. + mux.HandleFunc("/v0.1/servers", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"servers":[],"metadata":{"next_cursor":""}}`)) + }) + + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv, &callCount +} + +// TestCachedProvider_PluginsCacheHit verifies that a second ListAvailablePlugins +// call within the TTL is served from cache without hitting the registry. +func TestCachedProvider_PluginsCacheHit(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + {Namespace: "io.github.stacklok", Name: "code-reviewer", Version: "1.0.0"}, + {Namespace: "io.github.acme", Name: "doc-generator", Version: "0.2.0"}, + } + srv, callCount := newPluginsTestServer(t, plugins, nil) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + // First call: cache miss, fetches from API. + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 2) + require.Equal(t, "code-reviewer", got[0].Name) + require.Equal(t, int32(1), callCount.Load()) + + // Second call: cache hit, must not hit the API again. + got2, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Equal(t, got, got2) + require.Equal(t, int32(1), callCount.Load(), "second call should be served from cache") +} + +// TestCachedProvider_PluginsCacheMissFetch verifies that an empty cache +// triggers a fetch from the registry API. +func TestCachedProvider_PluginsCacheMissFetch(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + {Namespace: "io.github.stacklok", Name: "linter", Version: "1.0.0"}, + } + srv, callCount := newPluginsTestServer(t, plugins, nil) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 1) + require.Equal(t, "linter", got[0].Name) + require.Equal(t, int32(1), callCount.Load()) +} + +// TestCachedProvider_PluginsStaleOnFailure verifies that when the registry +// returns a transient error (500) on a cache miss, ListAvailablePlugins +// returns nil (plugins are optional) rather than propagating the error — +// mirroring the skills behavior in provider_cached.go. +func TestCachedProvider_PluginsStaleOnFailure(t *testing.T) { + t.Parallel() + + plugins := []*thvregistry.Plugin{ + {Namespace: "io.github.stacklok", Name: "code-reviewer", Version: "1.0.0"}, + } + var fail atomic.Bool + srv, callCount := newPluginsTestServer(t, plugins, &fail) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + // First call: succeeds and populates the cache. + got, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, got, 1) + firstCount := callCount.Load() + + // Force the registry to fail, then expire the cache so the next call must + // re-fetch. The stale cache should be served instead of the error. + fail.Store(true) + provider.pluginsMu.Lock() + provider.pluginsTime = provider.pluginsTime.Add(-defaultCacheTTL - 1) + provider.pluginsMu.Unlock() + + got2, err := provider.ListAvailablePlugins() + require.NoError(t, err, "transient failure should return stale cache, not error") + require.Len(t, got2, 1, "stale cache should still hold the previously fetched plugin") + require.Equal(t, "code-reviewer", got2[0].Name) + require.Greater(t, callCount.Load(), firstCount, "expired cache should trigger a re-fetch attempt") +} + +// TestCachedProvider_PluginsStaleOnFailureEmptyCache verifies that a transient +// failure with NO cached data returns the error (not nil,nil) so the v0.1 +// registry route surfaces a real failure instead of an empty 200 []. +func TestCachedProvider_PluginsStaleOnFailureEmptyCache(t *testing.T) { + t.Parallel() + + var fail atomic.Bool + fail.Store(true) + srv, _ := newPluginsTestServer(t, nil, &fail) + + provider, err := NewCachedAPIRegistryProvider(srv.URL, true, false, nil) + require.NoError(t, err) + + got, err := provider.ListAvailablePlugins() + require.Error(t, err, "transient failure with empty cache should return the error, not nil,nil") + require.Nil(t, got) +} diff --git a/pkg/registry/provider_local.go b/pkg/registry/provider_local.go index 734292f595..313dc1fc51 100644 --- a/pkg/registry/provider_local.go +++ b/pkg/registry/provider_local.go @@ -16,9 +16,11 @@ import ( // LocalRegistryProvider provides registry data from embedded JSON files or local files type LocalRegistryProvider struct { *BaseProvider - filePath string - skillsMu sync.RWMutex - skills []types.Skill + filePath string + skillsMu sync.RWMutex + skills []types.Skill + pluginsMu sync.RWMutex + plugins []types.Plugin } // NewLocalRegistryProvider creates a new local registry provider @@ -52,11 +54,12 @@ func (p *LocalRegistryProvider) GetRegistry() (*types.Registry, error) { data = catalog.Upstream() } - registry, skills, err := parseRegistryData(data) + registry, skills, plugins, err := parseRegistryData(data) if err != nil { return nil, err } p.setSkills(skills) + p.setPlugins(plugins) // Set name field on each server based on map key for name, server := range registry.Servers { @@ -88,6 +91,12 @@ func (p *LocalRegistryProvider) setSkills(skills []types.Skill) { p.skills = skills } +func (p *LocalRegistryProvider) setPlugins(plugins []types.Plugin) { + p.pluginsMu.Lock() + defer p.pluginsMu.Unlock() + p.plugins = plugins +} + // ListAvailableSkills returns skills discovered from the upstream registry data. // Triggers a registry load if skills haven't been populated yet. func (p *LocalRegistryProvider) ListAvailableSkills() ([]types.Skill, error) { @@ -139,3 +148,55 @@ func (p *LocalRegistryProvider) SearchSkills(query string) ([]types.Skill, error } return results, nil } + +// ListAvailablePlugins returns plugins discovered from the upstream registry data. +// Triggers a registry load if plugins haven't been populated yet. +func (p *LocalRegistryProvider) ListAvailablePlugins() ([]types.Plugin, error) { + p.pluginsMu.RLock() + plugins := p.plugins + p.pluginsMu.RUnlock() + + if plugins == nil { + // Plugins are populated as a side effect of GetRegistry + if _, err := p.GetRegistry(); err != nil { + return nil, err + } + p.pluginsMu.RLock() + plugins = p.plugins + p.pluginsMu.RUnlock() + } + + return plugins, nil +} + +// GetPlugin returns a specific plugin by namespace and name. +func (p *LocalRegistryProvider) GetPlugin(namespace, name string) (*types.Plugin, error) { + plugins, err := p.ListAvailablePlugins() + if err != nil { + return nil, err + } + for i := range plugins { + if plugins[i].Namespace == namespace && plugins[i].Name == name { + return &plugins[i], nil + } + } + return nil, nil +} + +// SearchPlugins searches for plugins matching the query in name, namespace, or description. +func (p *LocalRegistryProvider) SearchPlugins(query string) ([]types.Plugin, error) { + plugins, err := p.ListAvailablePlugins() + if err != nil { + return nil, err + } + query = strings.ToLower(query) + var results []types.Plugin + for _, pl := range plugins { + if strings.Contains(strings.ToLower(pl.Name), query) || + strings.Contains(strings.ToLower(pl.Description), query) || + strings.Contains(strings.ToLower(pl.Namespace), query) { + results = append(results, pl) + } + } + return results, nil +} diff --git a/pkg/registry/provider_plugins_test.go b/pkg/registry/provider_plugins_test.go new file mode 100644 index 0000000000..91d2d80104 --- /dev/null +++ b/pkg/registry/provider_plugins_test.go @@ -0,0 +1,203 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package registry + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// upstreamRegistryWithPlugins is a minimal upstream-format registry payload +// carrying a single server (so validateConnectivity accepts it for the remote +// provider) plus two plugins used to exercise List/Get/Search on both the +// local and remote providers. +const upstreamRegistryWithPlugins = `{ + "$schema": "https://example.com/schema.json", + "version": "1.0.0", + "meta": {"last_updated": "2025-01-01T00:00:00Z"}, + "data": { + "servers": [ + { + "name": "io.example.test-server", + "description": "Test server", + "packages": [ + { + "registryType": "oci", + "identifier": "example/test-server:latest", + "transport": {"type": "stdio"} + } + ] + } + ], + "plugins": [ + { + "namespace": "io.github.stacklok", + "name": "code-reviewer", + "description": "Reviews code for bugs", + "version": "1.0.0" + }, + { + "namespace": "io.github.acme", + "name": "doc-generator", + "description": "Generates documentation", + "version": "0.2.0" + } + ] + } +}` + +// pluginsOnlyRegistry is an upstream-format payload that carries plugins but +// no servers or skills. The remote provider's validateConnectivity guard must +// accept it (regression for the widened empty-data check). +const pluginsOnlyRegistry = `{ + "$schema": "https://example.com/schema.json", + "version": "1.0.0", + "meta": {"last_updated": "2025-01-01T00:00:00Z"}, + "data": { + "servers": [], + "plugins": [ + { + "namespace": "io.github.stacklok", + "name": "code-reviewer", + "description": "Reviews code for bugs", + "version": "1.0.0" + } + ] + } +}` + +func TestParseRegistryData_Plugins(t *testing.T) { + t.Parallel() + + _, _, plugins, err := parseRegistryData([]byte(upstreamRegistryWithPlugins)) + require.NoError(t, err) + require.Len(t, plugins, 2) + assert.Equal(t, "io.github.stacklok", plugins[0].Namespace) + assert.Equal(t, "code-reviewer", plugins[0].Name) + assert.Equal(t, "Reviews code for bugs", plugins[0].Description) +} + +func TestLocalRegistryProvider_Plugins(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + path := filepath.Join(dir, "registry.json") + require.NoError(t, os.WriteFile(path, []byte(upstreamRegistryWithPlugins), 0o600)) + + provider := NewLocalRegistryProvider(path) + + // ListAvailablePlugins yields the plugins from data.plugins. + plugins, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, plugins, 2) + + // GetPlugin returns the matching plugin. + got, err := provider.GetPlugin("io.github.stacklok", "code-reviewer") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, "code-reviewer", got.Name) + assert.Equal(t, "1.0.0", got.Version) + + // GetPlugin returns nil for a missing plugin (no error). + missing, err := provider.GetPlugin("io.github.nope", "absent") + require.NoError(t, err) + assert.Nil(t, missing) + + // SearchPlugins matches by name. + byName, err := provider.SearchPlugins("reviewer") + require.NoError(t, err) + require.Len(t, byName, 1) + assert.Equal(t, "code-reviewer", byName[0].Name) + + // SearchPlugins matches by namespace. + byNs, err := provider.SearchPlugins("acme") + require.NoError(t, err) + require.Len(t, byNs, 1) + assert.Equal(t, "doc-generator", byNs[0].Name) + + // SearchPlugins matches by description. + byDesc, err := provider.SearchPlugins("documentation") + require.NoError(t, err) + require.Len(t, byDesc, 1) + assert.Equal(t, "doc-generator", byDesc[0].Name) + + // SearchPlugins returns nil for a non-matching query. + none, err := provider.SearchPlugins("non-existent-plugin") + require.NoError(t, err) + assert.Nil(t, none) +} + +func TestRemoteRegistryProvider_Plugins(t *testing.T) { + t.Parallel() + + server := createTestServer(upstreamRegistryWithPlugins, 200) + defer server.Close() + + provider, err := NewRemoteRegistryProvider(server.URL, true) + require.NoError(t, err) + require.NotNil(t, provider) + + plugins, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, plugins, 2) + + got, err := provider.GetPlugin("io.github.acme", "doc-generator") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, "doc-generator", got.Name) + + byName, err := provider.SearchPlugins("reviewer") + require.NoError(t, err) + require.Len(t, byName, 1) + assert.Equal(t, "code-reviewer", byName[0].Name) +} + +// TestRemoteRegistryProvider_PluginsOnlyNotRejected confirms the widened +// validateConnectivity guard accepts a plugins-only remote registry. +func TestRemoteRegistryProvider_PluginsOnlyNotRejected(t *testing.T) { + t.Parallel() + + server := createTestServer(pluginsOnlyRegistry, 200) + defer server.Close() + + provider, err := NewRemoteRegistryProvider(server.URL, true) + require.NoError(t, err) + require.NotNil(t, provider) + + plugins, err := provider.ListAvailablePlugins() + require.NoError(t, err) + require.Len(t, plugins, 1) + assert.Equal(t, "code-reviewer", plugins[0].Name) +} + +// TestBaseProvider_PluginsNoOp confirms the BaseProvider no-op defaults return +// nil, nil (providers that don't support plugins inherit these). +func TestBaseProvider_PluginsNoOp(t *testing.T) { + t.Parallel() + + var p *BaseProvider + plugins, err := p.ListAvailablePlugins() + require.NoError(t, err) + assert.Nil(t, plugins) + + got, err := p.GetPlugin("ns", "name") + require.NoError(t, err) + assert.Nil(t, got) + + results, err := p.SearchPlugins("query") + require.NoError(t, err) + assert.Nil(t, results) +} + +// Compile-time assertions that the concrete providers satisfy the interface. +var ( + _ Provider = (*LocalRegistryProvider)(nil) + _ Provider = (*RemoteRegistryProvider)(nil) + _ Provider = (*APIRegistryProvider)(nil) + _ Provider = (*CachedAPIRegistryProvider)(nil) +) diff --git a/pkg/registry/provider_remote.go b/pkg/registry/provider_remote.go index beedf0837b..452a7a6e52 100644 --- a/pkg/registry/provider_remote.go +++ b/pkg/registry/provider_remote.go @@ -25,6 +25,8 @@ type RemoteRegistryProvider struct { allowPrivateIp bool skillsMu sync.RWMutex skills []types.Skill + pluginsMu sync.RWMutex + plugins []types.Plugin } // NewRemoteRegistryProvider creates a new remote registry provider. @@ -89,8 +91,8 @@ func (p *RemoteRegistryProvider) validateConnectivity() error { if err := json.Unmarshal(data, &upstream); err != nil { return fmt.Errorf("registry returned invalid upstream JSON from %s: %w", p.registryURL, err) } - if len(upstream.Data.Servers) == 0 && len(upstream.Data.Skills) == 0 { - return fmt.Errorf("registry at %s returned upstream format with no servers or skills", p.registryURL) + if len(upstream.Data.Servers) == 0 && len(upstream.Data.Skills) == 0 && len(upstream.Data.Plugins) == 0 { + return fmt.Errorf("registry at %s returned upstream format with no servers, skills, or plugins", p.registryURL) } return nil } @@ -129,11 +131,12 @@ func (p *RemoteRegistryProvider) GetRegistry() (*types.Registry, error) { return nil, fmt.Errorf("failed to read registry data from response body: %w", err) } - registry, skills, err := parseRegistryData(data) + registry, skills, plugins, err := parseRegistryData(data) if err != nil { return nil, fmt.Errorf("failed to parse registry data from %s: %w", p.registryURL, err) } p.setSkills(skills) + p.setPlugins(plugins) // Set name field on each server based on map key for name, server := range registry.Servers { @@ -211,8 +214,66 @@ func (p *RemoteRegistryProvider) SearchSkills(query string) ([]types.Skill, erro return results, nil } +// ListAvailablePlugins returns plugins discovered from the remote registry data. +// Triggers a registry load if plugins haven't been populated yet. +func (p *RemoteRegistryProvider) ListAvailablePlugins() ([]types.Plugin, error) { + p.pluginsMu.RLock() + plugins := p.plugins + p.pluginsMu.RUnlock() + + if plugins == nil { + // Plugins are populated as a side effect of GetRegistry + if _, err := p.GetRegistry(); err != nil { + return nil, err + } + p.pluginsMu.RLock() + plugins = p.plugins + p.pluginsMu.RUnlock() + } + + return plugins, nil +} + +// GetPlugin returns a specific plugin by namespace and name. +func (p *RemoteRegistryProvider) GetPlugin(namespace, name string) (*types.Plugin, error) { + plugins, err := p.ListAvailablePlugins() + if err != nil { + return nil, err + } + for i := range plugins { + if plugins[i].Namespace == namespace && plugins[i].Name == name { + return &plugins[i], nil + } + } + return nil, nil +} + +// SearchPlugins searches for plugins matching the query in name, namespace, or description. +func (p *RemoteRegistryProvider) SearchPlugins(query string) ([]types.Plugin, error) { + plugins, err := p.ListAvailablePlugins() + if err != nil { + return nil, err + } + query = strings.ToLower(query) + var results []types.Plugin + for _, pl := range plugins { + if strings.Contains(strings.ToLower(pl.Name), query) || + strings.Contains(strings.ToLower(pl.Description), query) || + strings.Contains(strings.ToLower(pl.Namespace), query) { + results = append(results, pl) + } + } + return results, nil +} + func (p *RemoteRegistryProvider) setSkills(skills []types.Skill) { p.skillsMu.Lock() defer p.skillsMu.Unlock() p.skills = skills } + +func (p *RemoteRegistryProvider) setPlugins(plugins []types.Plugin) { + p.pluginsMu.Lock() + defer p.pluginsMu.Unlock() + p.plugins = plugins +} diff --git a/pkg/registry/provider_test.go b/pkg/registry/provider_test.go index df5e9bebce..28f835e089 100644 --- a/pkg/registry/provider_test.go +++ b/pkg/registry/provider_test.go @@ -223,7 +223,7 @@ func TestRemoteRegistryProvider_ValidateConnectivity(t *testing.T) { responseBody: `{"$schema": "x", "version": "1.0", "meta": {}, "data": {"servers": []}}`, responseStatus: 200, expectError: true, - errorContains: "no servers or skills", + errorContains: "no servers, skills, or plugins", }, { name: "legacy format surfaces migration hint", diff --git a/pkg/registry/schema_validation_test.go b/pkg/registry/schema_validation_test.go index c6e7b8376e..8d33ef70fa 100644 --- a/pkg/registry/schema_validation_test.go +++ b/pkg/registry/schema_validation_test.go @@ -27,7 +27,7 @@ func TestEmbeddedRegistrySchemaValidation(t *testing.T) { func TestValidateEmbeddedRegistryCanLoadData(t *testing.T) { t.Parallel() - registry, skills, err := parseRegistryData(catalog.Upstream()) + registry, skills, _, err := parseRegistryData(catalog.Upstream()) require.NoError(t, err, "Embedded upstream registry should parse successfully") // Verify basic structure @@ -45,7 +45,7 @@ func TestValidateEmbeddedRegistryCanLoadData(t *testing.T) { func TestUpstreamRegistryParsing(t *testing.T) { t.Parallel() - registry, _, err := parseRegistryData(catalog.Upstream()) + registry, _, _, err := parseRegistryData(catalog.Upstream()) require.NoError(t, err) // Verify servers have names set (from conversion) @@ -124,7 +124,7 @@ func TestParseRegistryData_LegacyFormatDetection(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - _, _, err := parseRegistryData([]byte(tt.input)) + _, _, _, err := parseRegistryData([]byte(tt.input)) if tt.wantLegacy { require.Error(t, err) assert.ErrorIs(t, err, errLegacyFormat) @@ -142,7 +142,7 @@ func TestParseRegistryData_LegacyFormatDetection(t *testing.T) { func TestParseRegistryData_MalformedJSON(t *testing.T) { t.Parallel() - _, _, err := parseRegistryData([]byte("not json")) + _, _, _, err := parseRegistryData([]byte("not json")) require.Error(t, err) assert.NotErrorIs(t, err, errLegacyFormat) assert.Contains(t, err.Error(), "failed to parse registry data") diff --git a/pkg/registry/upstream_parser.go b/pkg/registry/upstream_parser.go index c48d8afe22..ffe3325af2 100644 --- a/pkg/registry/upstream_parser.go +++ b/pkg/registry/upstream_parser.go @@ -26,18 +26,19 @@ import ( var errLegacyFormat = &LegacyFormatError{} // parseRegistryData parses raw JSON in the upstream MCP registry format and -// converts it into the internal types.Registry plus any embedded skills. +// converts it into the internal types.Registry plus any embedded skills and +// plugins. // // Returns errLegacyFormat if the input looks like the legacy ToolHive registry // format. -func parseRegistryData(data []byte) (*types.Registry, []types.Skill, error) { +func parseRegistryData(data []byte) (*types.Registry, []types.Skill, []types.Plugin, error) { if !legacyhint.IsUpstream(data) && legacyhint.Looks(data) { - return nil, nil, errLegacyFormat + return nil, nil, nil, errLegacyFormat } var upstream types.UpstreamRegistry if err := json.Unmarshal(data, &upstream); err != nil { - return nil, nil, fmt.Errorf("failed to parse registry data: %w", err) + return nil, nil, nil, fmt.Errorf("failed to parse registry data: %w", err) } // ConvertServersToMetadata expects []*v0.ServerJSON, but UpstreamData.Servers @@ -49,7 +50,7 @@ func parseRegistryData(data []byte) (*types.Registry, []types.Skill, error) { serverMetadata, err := ConvertServersToMetadata(serverPtrs) if err != nil { - return nil, nil, fmt.Errorf("failed to convert servers to metadata: %w", err) + return nil, nil, nil, fmt.Errorf("failed to convert servers to metadata: %w", err) } registry := &types.Registry{ @@ -72,5 +73,5 @@ func parseRegistryData(data []byte) (*types.Registry, []types.Skill, error) { } } - return registry, upstream.Data.Skills, nil + return registry, upstream.Data.Skills, upstream.Data.Plugins, nil } diff --git a/pkg/runner/authz_audit_integration_test.go b/pkg/runner/authz_audit_integration_test.go new file mode 100644 index 0000000000..bb45babc60 --- /dev/null +++ b/pkg/runner/authz_audit_integration_test.go @@ -0,0 +1,233 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package runner + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/authz/authorizers" + "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" + "github.com/stacklok/toolhive/pkg/webhook" + statusesmocks "github.com/stacklok/toolhive/pkg/workloads/statuses/mocks" +) + +// buildRunnerMiddlewareChain populates the middleware configs for runConfig, +// instantiates the middlewares through the runner factories (the same path the +// proxyrunner takes), and wraps final with them in the order the proxies apply +// them (reverse slice order, so index 0 is the outermost wrapper). +func buildRunnerMiddlewareChain(t *testing.T, runConfig *RunConfig, final http.Handler) http.Handler { + t.Helper() + + require.NoError(t, PopulateMiddlewareConfigs(runConfig)) + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockStatusManager := statusesmocks.NewMockStatusManager(ctrl) + + runner := NewRunner(runConfig, mockStatusManager) + for _, mwConfig := range runConfig.MiddlewareConfigs { + factory, ok := runner.supportedMiddleware[mwConfig.Type] + require.True(t, ok, "no factory for middleware type %q", mwConfig.Type) + require.NoError(t, factory(&mwConfig, runner)) + } + require.NotEmpty(t, runner.middlewares) + // The factories acquire resources (the auditor's log file, the usage-metrics + // flush goroutine); release them when the test finishes. + t.Cleanup(func() { + for _, mw := range runner.middlewares { + _ = mw.Close() + } + }) + + handler := final + for i := len(runner.middlewares) - 1; i >= 0; i-- { + handler = runner.middlewares[i].Handler()(handler) + } + return handler +} + +// readAuditEvents parses the newline-delimited JSON audit log at path. +func readAuditEvents(t *testing.T, path string) []map[string]any { + t.Helper() + + data, err := os.ReadFile(path) + require.NoError(t, err) + + var events []map[string]any + for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") { + if line == "" { + continue + } + var event map[string]any + require.NoError(t, json.Unmarshal([]byte(line), &event), "audit log line is not JSON: %s", line) + events = append(events, event) + } + return events +} + +// TestAuthzDecisionIsAudited proves, through the full middleware chain the +// proxyrunner builds via PopulateMiddlewareConfigs, that authorization +// decisions produce audit events: a Cedar-denied tools/call must return HTTP +// 403 AND still emit an audit event with outcome "denied", and an authorized +// call must emit one with outcome "success". This is the regression guard for +// the ordering bug where audit was wired inside authz and denials were never +// audited. +func TestAuthzDecisionIsAudited(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + policy string + wantStatus int + wantOutcome string + wantHandlerHit bool + wantDeniedError bool + }{ + { + name: "denied tool call returns 403 and is audited with outcome denied", + // Permit only an unrelated tool: "target_tool" is default-denied. + policy: `permit(principal, action == Action::"call_tool", resource == Tool::"some_other_tool");`, + wantStatus: http.StatusForbidden, + wantOutcome: "denied", + wantHandlerHit: false, + wantDeniedError: true, + }, + { + name: "allowed tool call returns 200 and is audited with outcome success", + policy: `permit(principal, action, resource);`, + wantStatus: http.StatusOK, + wantOutcome: "success", + wantHandlerHit: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + authzConfig, err := authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{ + Policies: []string{tt.policy}, + EntitiesJSON: "[]", + }, + }) + require.NoError(t, err) + + auditLogPath := filepath.Join(t.TempDir(), "audit.log") + + runConfig := NewRunConfig() + runConfig.Name = "test-server" + runConfig.AuthzConfig = authzConfig + runConfig.AuditConfig = &audit.Config{ + Component: "test-component", + LogFile: auditLogPath, + } + + handlerHit := false + backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + handlerHit = true + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`)) + }) + + handler := buildRunnerMiddlewareChain(t, runConfig, backend) + + reqBody := `{"jsonrpc":"2.0","method":"tools/call","id":1,"params":{"name":"target_tool","arguments":{}}}` + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(reqBody)) + req.Header.Set("Content-Type", "application/json") + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + require.Equal(t, tt.wantStatus, rr.Code, "response body: %s", rr.Body.String()) + assert.Equal(t, tt.wantHandlerHit, handlerHit) + if tt.wantDeniedError { + assert.Contains(t, rr.Body.String(), "Unauthorized", + "denied response must carry the authorization-denied error") + } + + events := readAuditEvents(t, auditLogPath) + require.Len(t, events, 1, "exactly one audit event must be emitted") + event := events[0] + assert.Equal(t, "mcp_tool_call", event["type"], + "the parsed MCP method must drive the audit event type") + assert.Equal(t, tt.wantOutcome, event["outcome"]) + }) + } +} + +// TestWebhookDenialIsAudited proves, through the full middleware chain built +// by PopulateMiddlewareConfigs, that a validating-webhook policy denial (403) +// still produces an audit event with outcome "denied". Webhooks run inside +// the audit middleware, so the rejection must be captured like any other. +func TestWebhookDenialIsAudited(t *testing.T) { + t.Parallel() + + denyServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req webhook.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + resp := webhook.Response{ + Version: webhook.APIVersion, + UID: req.UID, + Allowed: false, + Reason: "policy denied", + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + })) + t.Cleanup(denyServer.Close) + + auditLogPath := filepath.Join(t.TempDir(), "audit.log") + + runConfig := NewRunConfig() + runConfig.Name = "test-server" + runConfig.ValidatingWebhooks = []webhook.Config{{ + Name: "deny-all", + URL: denyServer.URL, + Timeout: webhook.DefaultTimeout, + FailurePolicy: webhook.FailurePolicyFail, + TLSConfig: &webhook.TLSConfig{InsecureSkipVerify: true}, + }} + runConfig.AuditConfig = &audit.Config{ + Component: "test-component", + LogFile: auditLogPath, + } + + handlerHit := false + backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + handlerHit = true + w.WriteHeader(http.StatusOK) + }) + + handler := buildRunnerMiddlewareChain(t, runConfig, backend) + + reqBody := `{"jsonrpc":"2.0","method":"tools/call","id":1,"params":{"name":"target_tool","arguments":{}}}` + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(reqBody)) + req.Header.Set("Content-Type", "application/json") + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + require.Equal(t, http.StatusForbidden, rr.Code, "response body: %s", rr.Body.String()) + assert.False(t, handlerHit, "a webhook-denied request must not reach the backend") + + events := readAuditEvents(t, auditLogPath) + require.Len(t, events, 1, "a webhook policy denial must produce exactly one audit event") + assert.Equal(t, "mcp_tool_call", events[0]["type"]) + assert.Equal(t, "denied", events[0]["outcome"]) +} diff --git a/pkg/runner/config.go b/pkg/runner/config.go index 62a3ffd94f..ec71b70eca 100644 --- a/pkg/runner/config.go +++ b/pkg/runner/config.go @@ -207,6 +207,12 @@ type RunConfig struct { // TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies TrustProxyHeaders bool `json:"trust_proxy_headers,omitempty" yaml:"trust_proxy_headers,omitempty"` + // StrictProtocolValidation enables strict MCP-Protocol-Version validation + // on the streamable HTTP proxy: a request whose header names an unknown + // MCP revision is rejected with HTTP 400. Default false accepts any + // version string (an absent header is always accepted in either mode). + StrictProtocolValidation bool `json:"strict_protocol_validation,omitempty" yaml:"strict_protocol_validation,omitempty"` + // Stateless indicates the server only supports POST (no SSE/GET). // When true, the proxy returns 405 for incoming GET requests and uses a // POST-based health check instead of the default GET probe. @@ -395,9 +401,43 @@ func ReadJSON(r io.Reader) (*RunConfig, error) { // Normalize proxyMode so pre-existing configs always reflect the effective protocol config.NormalizeProxyMode() + // Self-heal legacy on-disk configs that persisted isolate_network=true with a + // non-bridge network mode. Isolation is only enforceable in bridge mode, so + // reflect reality in memory for status/list. See issue #5775. + degradeNetworkIsolation(&config) + return &config, nil } +// degradeNetworkIsolation drops network isolation from a loaded config when its +// resolved network mode cannot enforce it, so status/list reflect reality. +// +// ReadJSON is reached from read-only paths (thv list, status, export, API GETs) +// as well as deploy. The "none" case logs at DEBUG since it's merely redundant +// (already maximally confined) and would otherwise fire on every read of a +// static, already-known condition. The host/custom case stays at WARN even on +// read paths: it means the user's --isolate-network request is being silently +// ignored, which they should keep seeing until they either drop the flag or +// switch network modes. See #5775. +func degradeNetworkIsolation(config *RunConfig) { + if !config.IsolateNetwork { + return + } + mode := "" + if config.PermissionProfile != nil && config.PermissionProfile.Network != nil { + mode = config.PermissionProfile.Network.Mode + } + if networking.IsBridgeMode(mode) { + return + } + config.IsolateNetwork = false + if mode == "none" { + slog.Debug(networking.NetworkIsolationNoneRedundantMsg, "network_mode", mode) + return + } + slog.Warn(networking.NetworkIsolationHostDroppedMsg, "network_mode", mode) +} + // migrateOAuthClientSecret migrates plain text OAuth client secrets to CLI format // This handles the transition from storing plain text secrets to CLI format references func migrateOAuthClientSecret(config *RunConfig) error { diff --git a/pkg/runner/config_builder.go b/pkg/runner/config_builder.go index 2f6602f180..417c76a23f 100644 --- a/pkg/runner/config_builder.go +++ b/pkg/runner/config_builder.go @@ -32,6 +32,7 @@ import ( "github.com/stacklok/toolhive/pkg/ignore" "github.com/stacklok/toolhive/pkg/labels" "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/oauthproto/tokenexchange" "github.com/stacklok/toolhive/pkg/recovery" "github.com/stacklok/toolhive/pkg/telemetry" @@ -64,6 +65,10 @@ type runConfigBuilder struct { registryProxyPort int // Store network mode to apply to permission profile after it's loaded networkMode string + // isolateNetworkExplicit records whether the user explicitly requested network + // isolation (e.g. --isolate-network was passed) rather than it being the default. + // It controls whether an incompatible network mode fails fast or silently degrades. + isolateNetworkExplicit bool // Build context determines which validation and features are enabled buildContext BuildContext } @@ -315,7 +320,14 @@ func WithPermissionProfile(profile *permissions.Profile) RunConfigBuilderOption } } -// WithNetworkIsolation sets network isolation +// WithNetworkIsolation sets network isolation. +// +// Network isolation is only enforceable in bridge mode. When it is requested +// together with a non-bridge network mode (host/none/custom), the build +// reconciles the conflict: without the WithNetworkIsolationExplicit companion +// (i.e. the user did not actively pass --isolate-network), a host/custom +// conflict degrades to isolation=false with a warning rather than failing fast. +// See WithNetworkIsolationExplicit and issue #5775. func WithNetworkIsolation(isolate bool) RunConfigBuilderOption { return func(b *runConfigBuilder) error { b.config.IsolateNetwork = isolate @@ -323,6 +335,19 @@ func WithNetworkIsolation(isolate bool) RunConfigBuilderOption { } } +// WithNetworkIsolationExplicit records whether network isolation was explicitly +// requested by the user (i.e. --isolate-network was actually passed) rather than +// defaulted. It exists to drive the fail-fast-vs-degrade decision when isolation +// conflicts with a non-bridge network mode: when true, a host/custom conflict +// fails the build so the user learns their request cannot be honored; when false +// (defaulted), the conflict degrades to isolation=false with a warning. See #5775. +func WithNetworkIsolationExplicit(explicit bool) RunConfigBuilderOption { + return func(b *runConfigBuilder) error { + b.isolateNetworkExplicit = explicit + return nil + } +} + // WithAllowDockerGateway sets whether to allow outbound connections to Docker gateway addresses func WithAllowDockerGateway(allow bool) RunConfigBuilderOption { return func(b *runConfigBuilder) error { @@ -351,6 +376,16 @@ func WithTrustProxyHeaders(trust bool) RunConfigBuilderOption { } } +// WithStrictProtocolValidation sets whether the streamable HTTP proxy rejects +// requests carrying an unknown/unsupported MCP-Protocol-Version header with +// HTTP 400. Default false accepts any version string. +func WithStrictProtocolValidation(strict bool) RunConfigBuilderOption { + return func(b *runConfigBuilder) error { + b.config.StrictProtocolValidation = strict + return nil + } +} + // WithStateless declares the server is stateless (POST-only, no SSE). func WithStateless(stateless bool) RunConfigBuilderOption { return func(b *runConfigBuilder) error { @@ -631,15 +666,24 @@ func WithMiddlewareFromFlags( var middlewareConfigs []types.MiddlewareConfig // NOTE: order matters here. Specifically, these routines use append - // to add new middleware configs, but once these routines are called, - // inside the proxy, they are applied in reverse order, so the first - // being added here is effectively the last being called at HTTP - // request time. + // to add new middleware configs. The proxy wraps the handler in + // reverse slice order (see applyMiddlewares), so the first entry + // added here is the OUTERMOST wrapper and runs first at HTTP request + // time; the last entry runs closest to the handler. // // We should avoid doing this and a better pattern would be to let the // actual proxy determine the order of application of middlewares, since // the types of middleware are known at compile time. + // Audit middleware (if enabled) goes first so it is the outermost + // wrapper after the body-limit prepended in runner.Run: every request + // that passes the size cap produces an audit event no matter which + // middleware rejects it — authentication (401), webhook denials, and + // authorization (403, outcome "denied") included. Identity and parsed + // MCP data are read back from the inner auth/parser middlewares via + // the holder carriers (auth.IdentityHolder, mcp.ParsedRequestHolder). + middlewareConfigs = addAuditMiddleware(middlewareConfigs, enableAudit, auditConfigPath, serverName, transportType) + // Add tool filter middlewares middlewareConfigs = addToolFilterMiddlewares(middlewareConfigs, toolsFilter, toolsOverride) @@ -655,8 +699,11 @@ func WithMiddlewareFromFlags( // // NOTE: addCoreMiddlewares also injects usage metrics before webhook insertion here, // which differs slightly from PopulateMiddlewareConfigs where usage metrics is added - // after webhooks. This is currently benign because usage metrics does not depend on - // webhook state, and the broader ordering TODO remains to unify these paths. + // after webhooks. Since mutating webhooks republish the parsed request + // (mcp.RepublishParsedMCPRequest), the two paths now disagree when a webhook rewrites + // the JSON-RPC method: this path counts the method as received, the operator path counts + // the method as mutated. Only tool-call counts are affected, and only for a webhook that + // patches "method" itself. The broader ordering TODO remains to unify these paths. // Add Mutating webhooks before Validating webhooks var err error @@ -671,16 +718,19 @@ func WithMiddlewareFromFlags( return err } - // Add optional middlewares + // Add optional middlewares. Audit was added at the top of the chain + // so authorization denials (403) still produce an audit event with + // outcome "denied". middlewareConfigs = addTelemetryMiddleware(middlewareConfigs, telemetryConfig, serverName, transportType) var authzErr error middlewareConfigs, authzErr = addAuthzMiddleware(middlewareConfigs, authzConfigPath, b.config.EmbeddedAuthServerConfig) if authzErr != nil { return authzErr } - middlewareConfigs = addAuditMiddleware(middlewareConfigs, enableAudit, auditConfigPath, serverName, transportType) - // Add recovery middleware (always present, added last to be outermost wrapper) + // Add recovery middleware (always present, added last so it is the + // innermost wrapper, executing closest to the handler — matching + // PopulateMiddlewareConfigs) middlewareConfigs = addRecoveryMiddleware(middlewareConfigs) // Set the populated middleware configs @@ -914,9 +964,10 @@ func addAuditMiddleware( return middlewareConfigs } -// addRecoveryMiddleware adds recovery middleware (always present, added last to be outermost wrapper) -// Middleware is applied in reverse order, so adding last means it executes first -// and catches panics from all other middleware and handlers. +// addRecoveryMiddleware adds recovery middleware (always present, added last). +// The proxy wraps the handler in reverse slice order, so the last entry is the +// INNERMOST wrapper: it catches panics from the handler itself, but not from +// middleware added earlier in the slice (which wrap it). func addRecoveryMiddleware(middlewareConfigs []types.MiddlewareConfig) []types.MiddlewareConfig { recoveryConfig, err := types.NewMiddlewareConfig(recovery.MiddlewareType, nil) if err != nil { @@ -1113,6 +1164,15 @@ func (b *runConfigBuilder) validateConfig(imageMetadata *regtypes.ImageMetadata) slog.Debug("Setting network mode on permission profile", "network_mode", b.networkMode) } + // Reconcile network isolation against the resolved network mode. Isolation is + // only enforceable in bridge mode; for host/none/custom modes the isolation + // sidecars have no route out and the workload bypasses them (issue #5775). + // Key off the resolved profile mode so this also covers modes coming from a + // permission-profile file rather than the --network flag. + if err = b.reconcileNetworkIsolation(); err != nil { + return err + } + // Process volume mounts if err = b.processVolumeMounts(); err != nil { return err @@ -1192,6 +1252,51 @@ func (b *runConfigBuilder) validateConfig(imageMetadata *regtypes.ImageMetadata) return nil } +// reconcileNetworkIsolation drops network isolation when the resolved network +// mode cannot enforce it. In bridge mode isolation is left untouched. For "none" +// (already maximally confined) isolation is silently dropped. For host/other +// non-bridge modes isolation is dropped with a warning, unless the user +// explicitly requested isolation, in which case the build fails fast. +func (b *runConfigBuilder) reconcileNetworkIsolation() error { + c := b.config + if !c.IsolateNetwork { + return nil + } + + // Resolve the effective network mode from the permission profile. + mode := "" + if c.PermissionProfile != nil && c.PermissionProfile.Network != nil { + mode = c.PermissionProfile.Network.Mode + } + if networking.IsBridgeMode(mode) { + return nil + } + + if mode == "none" { + // "none" is already maximally confined; isolation is merely redundant. + c.IsolateNetwork = false + slog.Debug(networking.NetworkIsolationNoneRedundantMsg, "network_mode", mode) + return nil + } + + // host (or other non-bridge) is less restrictive than isolation, so dropping + // it reduces confinement. Fail fast when it was explicitly requested. + // + // The mode may come from --permission-profile rather than --network, so the + // message names the resolved mode instead of assuming a flag the user may + // never have passed. See #5794 review discussion. + if b.isolateNetworkExplicit { + return fmt.Errorf( + "network isolation cannot be enforced with the resolved network mode %q: "+ + "use a bridge network mode to keep enforced isolation, or pass --isolate-network=false to keep %q networking", + mode, mode) + } + + c.IsolateNetwork = false + slog.Warn(networking.NetworkIsolationHostDroppedMsg, "network_mode", mode) + return nil +} + func (b *runConfigBuilder) loadPermissionProfile(imageMetadata *regtypes.ImageMetadata) (*permissions.Profile, error) { // The permission profile object takes precedence over the name or path. if b.config.PermissionProfile != nil { diff --git a/pkg/runner/config_builder_test.go b/pkg/runner/config_builder_test.go index f8e63c0964..0087d97ffd 100644 --- a/pkg/runner/config_builder_test.go +++ b/pkg/runner/config_builder_test.go @@ -16,9 +16,11 @@ import ( "github.com/stacklok/toolhive-core/permissions" regtypes "github.com/stacklok/toolhive-core/registry/types" + "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authserver" "github.com/stacklok/toolhive/pkg/authserver/server/registration" + "github.com/stacklok/toolhive/pkg/authz" appconfig "github.com/stacklok/toolhive/pkg/config" "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/networking" @@ -1540,6 +1542,33 @@ func TestWithSessionTTL(t *testing.T) { } } +// TestWithStrictProtocolValidation verifies the builder option sets +// RunConfig.StrictProtocolValidation, mirroring WithTrustProxyHeaders's +// plumbing (see cmd/thv/app/run_flags.go's --strict-protocol-validation flag). +func TestWithStrictProtocolValidation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + strict bool + }{ + {name: "enabled", strict: true}, + {name: "disabled (default)", strict: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + builder := &runConfigBuilder{config: NewRunConfig()} + err := WithStrictProtocolValidation(tt.strict)(builder) + + require.NoError(t, err) + assert.Equal(t, tt.strict, builder.config.StrictProtocolValidation) + }) + } +} + func TestResolveRegistryServerName(t *testing.T) { t.Parallel() @@ -1583,6 +1612,56 @@ func TestResolveRegistryServerName(t *testing.T) { } } +// TestWithMiddlewareFromFlags_AuditWrapsChain pins the same audit-wraps-chain +// ordering invariant on the CLI flag path that +// TestPopulateMiddlewareConfigs_AuditWrapsChain pins on the operator path: +// audit is the first entry of the built slice (earlier entries wrap later +// ones at request time; runner.Run later prepends only body-limit/origin), so +// rejections from auth (401), webhooks, and authorization (403) all still +// produce audit events. +func TestWithMiddlewareFromFlags_AuditWrapsChain(t *testing.T) { + t.Parallel() + + builder := &runConfigBuilder{config: NewRunConfig()} + opt := WithMiddlewareFromFlags( + nil, // oidcConfig + nil, // tokenExchangeConfig + nil, // toolsFilter + nil, // toolsOverride + nil, // telemetryConfig + writeCedarConfigFile(t), // authzConfigPath + true, // enableAudit + "", // auditConfigPath + "test-server", // serverName + "streamable-http", // transportType + true, // disableUsageMetrics + ) + require.NoError(t, opt(builder)) + + typeIndex := make(map[string]int, len(builder.config.MiddlewareConfigs)) + for i, mw := range builder.config.MiddlewareConfigs { + typeIndex[mw.Type] = i + } + + auditIdx, ok := typeIndex[audit.MiddlewareType] + require.True(t, ok, "audit middleware must be present") + authzIdx, ok := typeIndex[authz.MiddlewareType] + require.True(t, ok, "authz middleware must be present") + authIdx, ok := typeIndex[auth.MiddlewareType] + require.True(t, ok, "auth middleware must be present") + parserIdx, ok := typeIndex[mcp.ParserMiddlewareType] + require.True(t, ok, "MCP parser middleware must be present") + + assert.Equal(t, 0, auditIdx, + "audit must be the outermost built entry so every rejection is audited") + assert.Less(t, auditIdx, authIdx, + "audit must wrap auth so authentication failures (401) are audited") + assert.Less(t, auditIdx, parserIdx, + "audit wraps the parser; parsed MCP data flows back via mcp.ParsedRequestHolder") + assert.Less(t, auditIdx, authzIdx, + "audit must wrap authz so authorization denials (403) are audited") +} + // TestWithAdditionalMiddlewareConfigs verifies the generic injected-middleware // builder option: it appends pre-built configs (across multiple calls and // multiple arguments), preserves order, and skips nil entries. The option is @@ -1639,3 +1718,88 @@ func TestWithAdditionalMiddlewareConfigs(t *testing.T) { assert.Nil(t, b.config.AdditionalMiddlewareConfigs) }) } + +// TestRunConfigBuilder_NetworkIsolationReconciliation verifies that network +// isolation is reconciled against the resolved network mode: enforced in bridge +// mode, dropped (with warn or error) for host, and dropped silently for none. +// See issue #5775. +func TestRunConfigBuilder_NetworkIsolationReconciliation(t *testing.T) { + t.Parallel() + + mockValidator := &mockEnvVarValidator{} + + const ( + isolateDefaultTrue = "default-true" // isolation on, not explicitly requested + isolateExplicitTrue = "explicit-true" // isolation on, explicitly requested + isolateFalse = "false" // isolation off + ) + + tests := []struct { + name string + networkMode string + isolate string + expectError bool + expectIsolation bool + }{ + // Bridge modes: isolation left untouched. + {"empty mode default isolate", "", isolateDefaultTrue, false, true}, + {"empty mode explicit isolate", "", isolateExplicitTrue, false, true}, + {"empty mode isolate off", "", isolateFalse, false, false}, + {"bridge mode default isolate", "bridge", isolateDefaultTrue, false, true}, + {"bridge mode explicit isolate", "bridge", isolateExplicitTrue, false, true}, + {"default mode default isolate", "default", isolateDefaultTrue, false, true}, + + // Host mode: drop + warn when defaulted, error when explicit. + {"host mode default isolate degrades", "host", isolateDefaultTrue, false, false}, + {"host mode explicit isolate errors", "host", isolateExplicitTrue, true, false}, + {"host mode isolate off untouched", "host", isolateFalse, false, false}, + + // None mode: always drop silently, never error. + {"none mode default isolate degrades", "none", isolateDefaultTrue, false, false}, + {"none mode explicit isolate degrades", "none", isolateExplicitTrue, false, false}, + {"none mode isolate off", "none", isolateFalse, false, false}, + + // Custom (non-bridge, non-none) mode behaves like host. + {"custom mode default isolate degrades", "container:foo", isolateDefaultTrue, false, false}, + {"custom mode explicit isolate errors", "container:foo", isolateExplicitTrue, true, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + opts := []RunConfigBuilderOption{ + WithImage("test-image"), + WithName("test-name"), + WithNetworkMode(tt.networkMode), + } + switch tt.isolate { + case isolateDefaultTrue: + opts = append(opts, WithNetworkIsolation(true)) + case isolateExplicitTrue: + opts = append(opts, WithNetworkIsolation(true), WithNetworkIsolationExplicit(true)) + case isolateFalse: + opts = append(opts, WithNetworkIsolation(false)) + } + + config, err := NewRunConfigBuilder(t.Context(), nil, map[string]string{}, mockValidator, opts...) + + if tt.expectError { + require.Error(t, err, "expected build to fail fast") + // The error must name the resolved mode (not assume it came from + // --network, since it may have come from --permission-profile) + // and both ways out. + assert.Contains(t, err.Error(), `"`+tt.networkMode+`"`, + "error should name the resolved network mode") + assert.Contains(t, err.Error(), "pass --isolate-network=false", + "error should offer disabling isolation") + return + } + + require.NoError(t, err) + require.NotNil(t, config) + assert.Equal(t, tt.expectIsolation, config.IsolateNetwork, + "IsolateNetwork should reflect reconciled value") + }) + } +} diff --git a/pkg/runner/config_test.go b/pkg/runner/config_test.go index 6ad4e0cacd..4c7d0fcf9d 100644 --- a/pkg/runner/config_test.go +++ b/pkg/runner/config_test.go @@ -174,6 +174,22 @@ func TestRunConfig_NormalizeProxyMode(t *testing.T) { // Note: This test uses actual port finding logic, so it may fail if ports are in use func TestRunConfig_WithPorts(t *testing.T) { t.Parallel() + // Find available ports dynamically to avoid flaky failures + port1 := networking.FindAvailable() + require.NotZero(t, port1, "should find an available proxy port for SSE") + targetPort1 := networking.FindAvailable() + require.NotZero(t, targetPort1, "should find an available target port for SSE") + + port2 := networking.FindAvailable() + require.NotZero(t, port2, "should find an available proxy port for HTTP") + targetPort2 := networking.FindAvailable() + require.NotZero(t, targetPort2, "should find an available target port for HTTP") + + port3 := networking.FindAvailable() + require.NotZero(t, port3, "should find an available proxy port for Stdio") + targetPort3 := networking.FindAvailable() + require.NotZero(t, targetPort3, "should find an available target port for Stdio") + testCases := []struct { name string config *RunConfig @@ -184,8 +200,8 @@ func TestRunConfig_WithPorts(t *testing.T) { { name: "SSE transport with specific ports", config: &RunConfig{Transport: types.TransportTypeSSE}, - port: 8001, - targetPort: 9001, + port: port1, + targetPort: targetPort1, expectError: false, }, { @@ -198,15 +214,15 @@ func TestRunConfig_WithPorts(t *testing.T) { { name: "Streamable HTTP transport with specific ports", config: &RunConfig{Transport: types.TransportTypeStreamableHTTP}, - port: 8002, - targetPort: 9002, + port: port2, + targetPort: targetPort2, expectError: false, }, { name: "Stdio transport with specific port", config: &RunConfig{Transport: types.TransportTypeStdio}, - port: 8003, - targetPort: 9003, // This should be ignored for stdio + port: port3, + targetPort: targetPort3, // This should be ignored for stdio expectError: false, }, } @@ -2601,3 +2617,53 @@ func TestRunConfig_MCPServerGenerationJSONRoundTrip(t *testing.T) { "decoded missing field should be zero, got %d", got.MCPServerGeneration) }) } + +// TestReadJSON_DegradesNetworkIsolation verifies that ReadJSON self-heals legacy +// on-disk configs that persisted isolate_network=true with a non-bridge network +// mode, so status/list reflect reality. See issue #5775. +func TestReadJSON_DegradesNetworkIsolation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mode string + omitProfile bool // omit permission_profile entirely (nil-guard case) + expectIsolation bool + }{ + {name: "host mode degrades", mode: "host", expectIsolation: false}, + {name: "none mode degrades", mode: "none", expectIsolation: false}, + {name: "bridge mode untouched", mode: "bridge", expectIsolation: true}, + {name: "empty mode untouched", mode: "", expectIsolation: true}, + {name: "no permission profile is nil-safe and untouched", omitProfile: true, expectIsolation: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + profileJSON := fmt.Sprintf(`, + "permission_profile": { + "network": { + "mode": "%s" + } + }`, tt.mode) + if tt.omitProfile { + profileJSON = "" + } + + configJSON := fmt.Sprintf(`{ + "schema_version": "v1", + "name": "degrade-test", + "image": "test:latest", + "transport": "stdio", + "isolate_network": true%s + }`, profileJSON) + + config, err := ReadJSON(strings.NewReader(configJSON)) + require.NoError(t, err) + require.NotNil(t, config) + assert.Equal(t, tt.expectIsolation, config.IsolateNetwork, + "IsolateNetwork should reflect reconciled value for mode %q", tt.mode) + }) + } +} diff --git a/pkg/runner/middleware.go b/pkg/runner/middleware.go index 1debe3fe87..da116f1b18 100644 --- a/pkg/runner/middleware.go +++ b/pkg/runner/middleware.go @@ -76,6 +76,28 @@ func PopulateMiddlewareConfigs(config *RunConfig) error { return err } + // Audit middleware (if enabled). Added directly inside body-limit so it + // wraps the REST of the chain: every request that passes the size cap + // produces an audit event no matter which middleware rejects it. + // Authentication (401) and authorization/webhook denials (403) map to + // outcome "denied"; other rejections such as rate limiting (429) map to + // "failure" (see audit.determineOutcome). Identity and parsed MCP data are + // read back from the inner auth/parser middlewares via the holder + // carriers (see auth.IdentityHolder, mcp.ParsedRequestHolder). + if config.AuditConfig != nil { + auditParams := audit.MiddlewareParams{ + ConfigPath: config.AuditConfigPath, // Keep for backwards compatibility + ConfigData: config.AuditConfig, // Use the loaded config data + Component: config.AuditConfig.Component, + TransportType: config.Transport.String(), // Pass the actual transport type + } + auditConfig, err := types.NewMiddlewareConfig(audit.MiddlewareType, auditParams) + if err != nil { + return fmt.Errorf("failed to create audit middleware config: %w", err) + } + middlewareConfigs = append(middlewareConfigs, *auditConfig) + } + // Authentication middleware (always present) authParams := auth.MiddlewareParams{ OIDCConfig: config.OIDCConfig, @@ -151,7 +173,7 @@ func PopulateMiddlewareConfigs(config *RunConfig) error { // Mutating Webhooks middleware (if configured). // Must run BEFORE validating webhooks: - // MCP Parser -> [Mutating Webhooks] -> [Validating Webhooks] -> Authz -> Audit + // Audit -> ... -> MCP Parser -> [Mutating Webhooks] -> [Validating Webhooks] -> Authz middlewareConfigs, err = addMutatingWebhookMiddleware(middlewareConfigs, config) if err != nil { return err @@ -204,21 +226,6 @@ func PopulateMiddlewareConfigs(config *RunConfig) error { middlewareConfigs = append(middlewareConfigs, *authzConfig) } - // Audit middleware (if enabled) - if config.AuditConfig != nil { - auditParams := audit.MiddlewareParams{ - ConfigPath: config.AuditConfigPath, // Keep for backwards compatibility - ConfigData: config.AuditConfig, // Use the loaded config data - Component: config.AuditConfig.Component, - TransportType: config.Transport.String(), // Pass the actual transport type - } - auditConfig, err := types.NewMiddlewareConfig(audit.MiddlewareType, auditParams) - if err != nil { - return fmt.Errorf("failed to create audit middleware config: %w", err) - } - middlewareConfigs = append(middlewareConfigs, *auditConfig) - } - // AWS STS middleware (if configured) // Placed after audit/authz so that authorization is checked before exchanging // credentials, and close to the backend so SigV4 signing happens as late as diff --git a/pkg/runner/middleware_test.go b/pkg/runner/middleware_test.go index bf51aa04ce..2a29e7c4d6 100644 --- a/pkg/runner/middleware_test.go +++ b/pkg/runner/middleware_test.go @@ -1354,6 +1354,51 @@ func TestPopulateMiddlewareConfigs_FullCoverage(t *testing.T) { assert.True(t, typeIndex[audit.MiddlewareType]) } +// TestPopulateMiddlewareConfigs_AuditWrapsChain pins the ordering invariant +// that the audit middleware sits directly inside body-limit, wrapping the +// REST of the chain. Earlier entries wrap later ones at request time, so this +// placement is what guarantees every request that passes the size cap +// produces an audit event no matter which middleware rejects it — +// authentication 401s, webhook denials, and authorization 403s included. +// Identity and parsed MCP data flow back to audit from the inner auth/parser +// middlewares via the holder carriers, so audit no longer needs to run inside +// them. +func TestPopulateMiddlewareConfigs_AuditWrapsChain(t *testing.T) { + t.Parallel() + + config := &RunConfig{ + AuthzConfig: &authz.Config{}, + AuditConfig: &audit.Config{Component: "test-component"}, + } + + require.NoError(t, PopulateMiddlewareConfigs(config)) + + typeIndex := make(map[string]int, len(config.MiddlewareConfigs)) + for i, mw := range config.MiddlewareConfigs { + typeIndex[mw.Type] = i + } + + auditIdx, ok := typeIndex[audit.MiddlewareType] + require.True(t, ok, "audit middleware must be present") + authzIdx, ok := typeIndex[authz.MiddlewareType] + require.True(t, ok, "authz middleware must be present") + authIdx, ok := typeIndex[auth.MiddlewareType] + require.True(t, ok, "auth middleware must be present") + parserIdx, ok := typeIndex[mcp.ParserMiddlewareType] + require.True(t, ok, "MCP parser middleware must be present") + bodyLimitIdx, ok := typeIndex[bodylimit.MiddlewareType] + require.True(t, ok, "body limit middleware must be present") + + assert.Less(t, bodyLimitIdx, auditIdx, + "body limit must stay outside audit so oversized bodies are rejected before audit buffers them") + assert.Less(t, auditIdx, authIdx, + "audit must wrap auth so authentication failures (401) are audited") + assert.Less(t, auditIdx, parserIdx, + "audit wraps the parser; parsed MCP data flows back via mcp.ParsedRequestHolder") + assert.Less(t, auditIdx, authzIdx, + "audit must wrap authz so authorization denials (403) are audited") +} + // TestPopulateMiddlewareConfigs_StripAuthOrdering pins the ordering invariant // for strip-auth: the auth middleware must precede it in the chain so the // client JWT is fully validated (and the identity stored in the request diff --git a/pkg/runner/protocol.go b/pkg/runner/protocol.go index 6f381f3d22..7f4be18f7b 100644 --- a/pkg/runner/protocol.go +++ b/pkg/runner/protocol.go @@ -149,6 +149,15 @@ func createTemplateData( } templateData.RuntimeConfig = runtimeConfig + // Build-time dependency constraints are interpreted per package + // ecosystem; only the uvx builder currently supports them. Reject + // rather than silently ignore for the others. + if transportType != templates.TransportTypeUVX && len(runtimeConfig.BuildWith) > 0 { + return templateData, fmt.Errorf( + "--build-with is not supported for %s:// builds (only uvx://)", transportType, + ) + } + return templateData, nil } @@ -217,6 +226,9 @@ func mergeRuntimeConfig(transportType templates.TransportType, override *templat merged.RuntimeEnv = mergeEnvMaps(defaults.RuntimeEnv, override.RuntimeEnv) + // BuildWith has no defaults; the override's specifiers are used as-is. + merged.BuildWith = override.BuildWith + return merged } diff --git a/pkg/runner/protocol_test.go b/pkg/runner/protocol_test.go index d9bb55d100..6585af3a8a 100644 --- a/pkg/runner/protocol_test.go +++ b/pkg/runner/protocol_test.go @@ -719,3 +719,41 @@ func TestLoadRuntimeConfig_UsesOverrideBuilderImage(t *testing.T) { assert.Equal(t, customImage, got.BuilderImage) assert.Equal(t, base.AdditionalPackages, got.AdditionalPackages) } + +func TestMergeRuntimeConfigCarriesBuildWith(t *testing.T) { + t.Parallel() + + got := mergeRuntimeConfig(templates.TransportTypeUVX, &templates.RuntimeConfig{ + BuildWith: []string{"mcp<2"}, + }) + if len(got.BuildWith) != 1 || got.BuildWith[0] != "mcp<2" { + t.Errorf("BuildWith = %v, want [mcp<2]", got.BuildWith) + } + + // No override specifiers: merged config must not invent any. + got = mergeRuntimeConfig(templates.TransportTypeUVX, &templates.RuntimeConfig{}) + if len(got.BuildWith) != 0 { + t.Errorf("BuildWith = %v, want empty", got.BuildWith) + } +} + +func TestCreateTemplateDataRejectsBuildWithForNonUVX(t *testing.T) { + t.Parallel() + + for _, transport := range []templates.TransportType{ + templates.TransportTypeNPX, templates.TransportTypeGO, + } { + _, err := createTemplateData(transport, "some-package", "", nil, + &templates.RuntimeConfig{BuildWith: []string{"mcp<2"}}) + if err == nil { + t.Errorf("createTemplateData(%s) with BuildWith should error, got nil", transport) + } + } + + // uvx accepts constraints. + _, err := createTemplateData(templates.TransportTypeUVX, "some-package", "", nil, + &templates.RuntimeConfig{BuildWith: []string{"mcp<2"}}) + if err != nil { + t.Errorf("createTemplateData(uvx) with BuildWith should succeed, got %v", err) + } +} diff --git a/pkg/runner/retriever/retriever.go b/pkg/runner/retriever/retriever.go index 0235236a37..5cc2432e00 100644 --- a/pkg/runner/retriever/retriever.go +++ b/pkg/runner/retriever/retriever.go @@ -358,16 +358,25 @@ func VerifyImage(image string, server *types.ImageMetadata, verifySetting string return verifier.ErrProvenanceServerInformationNotSet } - // Create a new verifier + // Create a new verifier. In warn mode this must not block the run: + // construction reaches out to the sigstore TUF repository, so a CDN + // outage or connection reset would otherwise hard-fail a mode whose + // contract is "warn and continue" — an actual failed verification + // below only warns, so an unreachable verifier must not be stricter. v, err := verifier.New(server.Provenance, images.NewCompositeKeychain()) if err != nil { + if verifySetting == VerifyImageWarn { + slog.Warn("image verification unavailable, continuing without it", "image", image, "reason", err) + return nil + } return err } - // Verify the image passing the provenance info + // Verify the image passing the provenance info. Warn mode downgrades + // every failure — signature mismatches and transient registry errors + // alike — to a warning; only enabled mode fails closed. if err = v.VerifyServer(image, server.Provenance); err != nil { - if (errors.Is(err, verifier.ErrImageNotSigned) || errors.Is(err, verifier.ErrProvenanceMismatch)) && - verifySetting == VerifyImageWarn { + if verifySetting == VerifyImageWarn { slog.Warn("MCP server failed image verification", "image", image, "reason", err) return nil } diff --git a/pkg/runner/retriever/retriever_test.go b/pkg/runner/retriever/retriever_test.go index 92d8eb64f9..3701026fe7 100644 --- a/pkg/runner/retriever/retriever_test.go +++ b/pkg/runner/retriever/retriever_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/container/verifier" regtypes "github.com/stacklok/toolhive-core/registry/types" "github.com/stacklok/toolhive/pkg/runner" ) @@ -134,6 +135,74 @@ func TestHasLatestTag(t *testing.T) { } } +func TestVerifyImage(t *testing.T) { + t.Parallel() + + const testImage = "ghcr.io/example/server:v1.0.0" + + // A SigstoreURL with no embedded TUF root makes verifier.New fail + // deterministically without touching the network — the same "verifier + // cannot be constructed" return path a sigstore TUF CDN outage takes. + unreachableVerifier := ®types.ImageMetadata{ + Provenance: ®types.Provenance{SigstoreURL: "tuf.invalid.example"}, + } + + tests := []struct { + name string + server *regtypes.ImageMetadata + verifySetting string + expectErr string + }{ + { + name: "disabled skips verification", + server: unreachableVerifier, + verifySetting: VerifyImageDisabled, + }, + { + name: "warn without provenance continues", + server: ®types.ImageMetadata{}, + verifySetting: VerifyImageWarn, + }, + { + name: "enabled without provenance fails", + server: ®types.ImageMetadata{}, + verifySetting: VerifyImageEnabled, + expectErr: verifier.ErrProvenanceServerInformationNotSet.Error(), + }, + { + name: "warn with unavailable verifier continues", + server: unreachableVerifier, + verifySetting: VerifyImageWarn, + }, + { + name: "enabled with unavailable verifier fails", + server: unreachableVerifier, + verifySetting: VerifyImageEnabled, + expectErr: "root.json", + }, + { + name: "invalid setting fails", + server: ®types.ImageMetadata{}, + verifySetting: "sometimes", + expectErr: "invalid value for --image-verification", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := VerifyImage(testImage, tt.server, tt.verifySetting) + if tt.expectErr != "" { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.expectErr) + } else { + require.NoError(t, err) + } + }) + } +} + // errorPolicyGate is a test PolicyGate that rejects server creation with a // configurable error. It embeds runner.NoopPolicyGate for forward compatibility. type errorPolicyGate struct { diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index 940abed5b9..a253156376 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -45,6 +45,30 @@ var ErrContainerExitedRestartNeeded = errors.New("container exited, restart need const defaultProxySessionKeyPrefix = "thv:proxy:session:" +// probeProtocolVersion is the MCP protocol revision the readiness probe advertises +// in the JSON-RPC params.protocolVersion of its initialize handshake. +// +// It is deliberately pinned to 2025-11-25 — the current stable revision and the +// newest one that still defines the `initialize` method — rather than tracking +// mcp.LATEST_PROTOCOL_VERSION. The upcoming 2026-07-28 revision (draft; subject to +// change until it ships) removes `initialize` entirely (replaced by server/discover +// + per-request _meta, SEP-2575), so a probe that sent method:"initialize" with a +// 2026-07-28 version string would be internally inconsistent. When ToolHive adopts +// 2026-07-28 (issue #5754) the probe must switch to server/discover — which, like +// all 2026-07-28 Streamable HTTP POSTs, carries the required Mcp-Method/Mcp-Name +// headers — and fall back to this initialize path for older backends. Pinning here +// keeps the probe's method and version consistent regardless of what the SDK later +// declares "latest". +// +// The value is sent only in the request body, NOT as an MCP-Protocol-Version +// header: the spec scopes that header to requests made AFTER initialization +// (carrying the negotiated version), and requires a server to reject an unsupported +// header value with HTTP 400. Sending it on the initialize request itself would let +// a backend that only supports an older revision 400 the probe — a false +// "not-ready" — whereas body-based version negotiation degrades gracefully (the +// server answers HTTP 200 with its own supported version). +const probeProtocolVersion = "2025-11-25" + // Runner is responsible for running an MCP server with the provided configuration type Runner struct { // Config is the configuration for the runner @@ -228,16 +252,17 @@ func (r *Runner) Run(ctx context.Context) error { // Create transport with runtime transportConfig := types.Config{ - Type: r.Config.Transport, - ProxyPort: r.Config.Port, - TargetPort: r.Config.TargetPort, - Host: r.Config.Host, - TargetHost: r.Config.TargetHost, - Deployer: r.Config.Deployer, - Debug: r.Config.Debug, - TrustProxyHeaders: r.Config.TrustProxyHeaders, - EndpointPrefix: r.Config.EndpointPrefix, - SessionTTL: effectiveSessionTTL, + Type: r.Config.Transport, + ProxyPort: r.Config.Port, + TargetPort: r.Config.TargetPort, + Host: r.Config.Host, + TargetHost: r.Config.TargetHost, + Deployer: r.Config.Deployer, + Debug: r.Config.Debug, + TrustProxyHeaders: r.Config.TrustProxyHeaders, + StrictProtocolValidation: r.Config.StrictProtocolValidation, + EndpointPrefix: r.Config.EndpointPrefix, + SessionTTL: effectiveSessionTTL, } // Set proxy mode for stdio transport @@ -802,73 +827,139 @@ func (r *Runner) handleRemoteAuthentication(ctx context.Context) (oauth2.TokenSo // Set up token persister to save tokens across restarts if secretManager != nil { authHandler.SetTokenPersister(func(refreshToken string, expiry time.Time) error { - // Generate a unique secret name for this workload's refresh token - secretName, err := authsecrets.GenerateUniqueSecretNameWithPrefix( + return r.persistRefreshToken(ctx, secretManager, refreshToken, expiry) + }) + + // Set up client credentials persister for DCR (Dynamic Client Registration) + authHandler.SetClientCredentialsPersister(func( + clientID, clientSecret string, + secretExpiry time.Time, + regAccessToken, regClientURI string, + tokenEndpointAuthMethod string, + registeredCallbackPort int, + ) error { + return r.persistClientCredentials( + ctx, secretManager, clientID, clientSecret, + secretExpiry, regAccessToken, regClientURI, tokenEndpointAuthMethod, registeredCallbackPort) + }) + } + + // Perform authentication + tokenSource, err := authHandler.Authenticate(ctx, r.Config.RemoteURL) + if err != nil { + return nil, fmt.Errorf("remote authentication failed: %w", err) + } + + return tokenSource, nil +} + +func (r *Runner) persistRefreshToken( + ctx context.Context, + secretManager secrets.Provider, + refreshToken string, + expiry time.Time, +) error { + secretName, err := authsecrets.GenerateUniqueSecretNameWithPrefix( + r.Config.Name, + "OAUTH_REFRESH_TOKEN_", + secretManager, + ) + if err != nil { + return fmt.Errorf("failed to generate secret name: %w", err) + } + + if err := authsecrets.StoreSecretInManagerWithProvider(ctx, secretName, refreshToken, secretManager); err != nil { + return fmt.Errorf("failed to store refresh token: %w", err) + } + + r.Config.RemoteAuthConfig.CachedRefreshTokenRef = secretName + r.Config.RemoteAuthConfig.CachedTokenExpiry = expiry + + if err := r.Config.SaveState(ctx); err != nil { + return fmt.Errorf("failed to save config with token reference: %w", err) + } + + slog.Debug("Stored OAuth refresh token in secret manager", "secret_name", secretName) + return nil +} + +func (r *Runner) persistClientCredentials( + ctx context.Context, + secretManager secrets.Provider, + clientID, clientSecret string, + secretExpiry time.Time, + regAccessToken, regClientURI string, + tokenEndpointAuthMethod string, + registeredCallbackPort int, +) error { + updatedConfig := *r.Config + updatedRemoteAuthConfig := remote.Config{} + if r.Config.RemoteAuthConfig != nil { + updatedRemoteAuthConfig = *r.Config.RemoteAuthConfig + } + updatedConfig.RemoteAuthConfig = &updatedRemoteAuthConfig + updatedRemoteAuthConfig.CachedClientID = clientID + + if clientSecret != "" { + clientSecretSecretName := updatedRemoteAuthConfig.CachedClientSecretRef + if clientSecretSecretName == "" { + var err error + clientSecretSecretName, err = authsecrets.GenerateUniqueSecretNameWithPrefix( r.Config.Name, - "OAUTH_REFRESH_TOKEN_", + "OAUTH_CLIENT_SECRET_", secretManager, ) if err != nil { - return fmt.Errorf("failed to generate secret name: %w", err) + return fmt.Errorf("failed to generate client secret secret name: %w", err) } + } - // Store the refresh token in the secret manager - if err := authsecrets.StoreSecretInManagerWithProvider(ctx, secretName, refreshToken, secretManager); err != nil { - return fmt.Errorf("failed to store refresh token: %w", err) - } + if err := authsecrets.StoreSecretInManagerWithProvider(ctx, clientSecretSecretName, clientSecret, secretManager); err != nil { + return fmt.Errorf("failed to store client secret: %w", err) + } + updatedRemoteAuthConfig.CachedClientSecretRef = clientSecretSecretName + } - // Store the secret reference (not the actual token) in the config - r.Config.RemoteAuthConfig.CachedRefreshTokenRef = secretName - r.Config.RemoteAuthConfig.CachedTokenExpiry = expiry + updatedRemoteAuthConfig.CachedSecretExpiry = secretExpiry - // Save the updated config to persist the reference - if err := r.Config.SaveState(ctx); err != nil { - return fmt.Errorf("failed to save config with token reference: %w", err) + if regAccessToken != "" { + regTokenSecretName := updatedRemoteAuthConfig.CachedRegTokenRef + if regTokenSecretName == "" { + var err error + regTokenSecretName, err = authsecrets.GenerateUniqueSecretNameWithPrefix(r.Config.Name, "OAUTH_REG_TOKEN_", secretManager) + if err != nil { + return fmt.Errorf("failed to generate registration token secret name: %w", err) } + } - slog.Debug("Stored OAuth refresh token in secret manager", "secret_name", secretName) - return nil - }) - - // Set up client credentials persister for DCR (Dynamic Client Registration) - authHandler.SetClientCredentialsPersister(func(clientID, clientSecret string) error { - // Store client ID directly (it's public information) - r.Config.RemoteAuthConfig.CachedClientID = clientID - - // Only store client secret if it's non-empty (PKCE flows may not have one) - if clientSecret != "" { - clientSecretSecretName, err := authsecrets.GenerateUniqueSecretNameWithPrefix( - r.Config.Name, - "OAUTH_CLIENT_SECRET_", - secretManager, - ) - if err != nil { - return fmt.Errorf("failed to generate client secret secret name: %w", err) - } - - if err := authsecrets.StoreSecretInManagerWithProvider(ctx, clientSecretSecretName, clientSecret, secretManager); err != nil { - return fmt.Errorf("failed to store client secret: %w", err) - } - r.Config.RemoteAuthConfig.CachedClientSecretRef = clientSecretSecretName - } + if err := authsecrets.StoreSecretInManagerWithProvider(ctx, regTokenSecretName, regAccessToken, secretManager); err != nil { + return fmt.Errorf("failed to store registration access token: %w", err) + } + updatedRemoteAuthConfig.CachedRegTokenRef = regTokenSecretName + slog.Debug("Stored DCR registration access token for RFC 7592 operations") + } - // Save the updated config to persist the credentials - if err := r.Config.SaveState(ctx); err != nil { - return fmt.Errorf("failed to save config with client credentials: %w", err) - } + updatedRemoteAuthConfig.CachedRegClientURI = regClientURI + updatedRemoteAuthConfig.CachedTokenEndpointAuthMethod = tokenEndpointAuthMethod + updatedRemoteAuthConfig.CachedDCRCallbackPort = registeredCallbackPort - slog.Debug("Stored DCR client credentials", "client_id", clientID) - return nil - }) + if err := updatedConfig.SaveState(ctx); err != nil { + return fmt.Errorf("failed to save config with client credentials: %w", err) } - // Perform authentication - tokenSource, err := authHandler.Authenticate(ctx, r.Config.RemoteURL) - if err != nil { - return nil, fmt.Errorf("remote authentication failed: %w", err) + // Preserve pointer identity so the auth handler and runner continue to + // observe the same complete config after the durable write succeeds. + if r.Config.RemoteAuthConfig == nil { + r.Config.RemoteAuthConfig = &updatedRemoteAuthConfig + } else { + *r.Config.RemoteAuthConfig = updatedRemoteAuthConfig } - return tokenSource, nil + slog.Debug("Stored DCR client credentials", "client_id", clientID, + "has_expiry", !secretExpiry.IsZero(), + "has_reg_token", regAccessToken != "", + "has_reg_uri", regClientURI != "") + return nil } // Cleanup performs cleanup operations for the runner, including shutting down all middleware. @@ -958,9 +1049,12 @@ func waitForInitializeSuccess( // Format: http://localhost:port/mcp endpoint = serverURL method = "POST" - payload = `{"jsonrpc":"2.0","method":"initialize","id":"toolhive-init-check",` + - `"params":{"protocolVersion":"2024-11-05","capabilities":{},` + - `"clientInfo":{"name":"toolhive","version":"1.0"}}}` + payload = fmt.Sprintf( + `{"jsonrpc":"2.0","method":"initialize","id":"toolhive-init-check",`+ + `"params":{"protocolVersion":%q,"capabilities":{},`+ + `"clientInfo":{"name":"toolhive","version":"1.0"}}}`, + probeProtocolVersion, + ) case "sse": // For SSE, just check if the SSE endpoint is available // We can't easily call initialize without establishing a full SSE connection, @@ -985,6 +1079,15 @@ func waitForInitializeSuccess( delay := 100 * time.Millisecond maxDelay := 2 * time.Second // Cap at 2 seconds between retries + // Per-attempt outcomes below log at DEBUG, which is invisible in a default + // deployment — a workload stuck in `starting` for minutes then shows no + // trace of WHY the probe kept failing. Surface a periodic INFO progress + // line (long-running operation) carrying the last observed outcome. + // Every loop iteration assigns lastObserved before it is read. + var lastObserved string + const progressInterval = 30 * time.Second + nextProgressLog := progressInterval + slog.Info("Waiting for MCP server to be ready", "endpoint", endpoint, "timeout", maxWaitTime) // Create HTTP client with a reasonable timeout for requests @@ -1006,11 +1109,15 @@ func waitForInitializeSuccess( if err != nil { slog.Debug("Failed to create request", "attempt", attempt, "error", err) + lastObserved = fmt.Sprintf("failed to create request: %v", err) } else { if method == "POST" { req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") - req.Header.Set("MCP-Protocol-Version", "2024-11-05") + // No MCP-Protocol-Version header: it is scoped to post-initialize + // requests (carrying the negotiated version) and a server MUST reject + // an unsupported value with HTTP 400. The initialize body's + // protocolVersion negotiates gracefully instead. See probeProtocolVersion. } resp, err := httpClient.Do(req) // #nosec G704 -- endpoint is the local MCP server readiness URL @@ -1027,15 +1134,24 @@ func waitForInitializeSuccess( slog.Debug("Server returned status", //nolint:gosec // G706: status code and attempt are integers "status_code", resp.StatusCode, "attempt", attempt) + lastObserved = fmt.Sprintf("HTTP %d", resp.StatusCode) } else { slog.Debug("Failed to reach endpoint", "attempt", attempt, "error", err) + lastObserved = fmt.Sprintf("unreachable: %v", err) } } // Check if we've exceeded the maximum wait time elapsed := time.Since(startTime) if elapsed >= maxWaitTime { - return fmt.Errorf("initialize not successful after %v (%d attempts)", elapsed, attempt) + return fmt.Errorf("initialize not successful after %v (%d attempts, last observed: %s)", + elapsed, attempt, lastObserved) + } + + if elapsed >= nextProgressLog { + slog.Info("Still waiting for MCP server to be ready", //nolint:gosec // G706: attempt is an integer + "endpoint", endpoint, "elapsed", elapsed.Round(time.Second), "attempt", attempt, "last_observed", lastObserved) + nextProgressLog += progressInterval } // Wait before retrying diff --git a/pkg/runner/runner_test.go b/pkg/runner/runner_test.go index e765ea0d39..62d51e5b66 100644 --- a/pkg/runner/runner_test.go +++ b/pkg/runner/runner_test.go @@ -5,21 +5,27 @@ package runner import ( "context" + "io" "net/http" "net/http/httptest" + "sync" "testing" "time" + "github.com/adrg/xdg" "github.com/alicebob/miniredis/v2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive/pkg/auth/remote" "github.com/stacklok/toolhive/pkg/auth/upstreamtoken" "github.com/stacklok/toolhive/pkg/authserver" authserverrunner "github.com/stacklok/toolhive/pkg/authserver/runner" storagemocks "github.com/stacklok/toolhive/pkg/authserver/storage/mocks" rt "github.com/stacklok/toolhive/pkg/container/runtime" + "github.com/stacklok/toolhive/pkg/secrets" + secretsmocks "github.com/stacklok/toolhive/pkg/secrets/mocks" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/transport/types" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" @@ -266,6 +272,53 @@ func TestWaitForInitializeSuccess(t *testing.T) { assert.NoError(t, err) }) + t.Run("Streamable sends pinned probe protocol version in body only", func(t *testing.T) { + t.Parallel() + + var ( + mu sync.Mutex + capturedBody string + capturedHeader []string + headerPresent bool + ) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + + mu.Lock() + capturedBody = string(body) + capturedHeader, headerPresent = r.Header["Mcp-Protocol-Version"] + mu.Unlock() + + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusMethodNotAllowed) + })) + t.Cleanup(server.Close) + + ctx := context.Background() + err := waitForInitializeSuccess(ctx, server.URL, "streamable-http", false, 5*time.Second) + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + // Assert against an independent literal (not probeProtocolVersion) so that + // changing the pin forces a conscious test update — the const is deliberately + // pinned and must not silently track mcp.LATEST_PROTOCOL_VERSION. + assert.Contains(t, capturedBody, `"protocolVersion":"2025-11-25"`) + assert.NotContains(t, capturedBody, "2024-11-05") + // The probe intentionally does not advertise the 2026-07-28 revision, which + // removes the initialize method the probe relies on. + assert.NotContains(t, capturedBody, "2026-07-28") + // The MCP-Protocol-Version header is intentionally NOT sent on the initialize + // request: it is scoped to post-initialize requests and a server MUST 400 an + // unsupported value, which would falsely mark an older backend as not-ready. + assert.False(t, headerPresent, "initialize probe must not send an MCP-Protocol-Version header, got %q", capturedHeader) + }) + t.Run("SSE success", func(t *testing.T) { t.Parallel() @@ -459,6 +512,188 @@ func TestNewTransportSessionStorage_UsesRedisCredentials(t *testing.T) { assert.True(t, ok) } +//nolint:paralleltest // Mutates process-wide runtime and XDG state settings. +func TestRunner_PersistClientCredentials(t *testing.T) { + // SaveState uses process-wide runtime and XDG state settings. Keep this test + // sequential and restore xdg's cached paths after t.Setenv restores the env. + t.Cleanup(xdg.Reload) + t.Setenv("TOOLHIVE_RUNTIME", "") + t.Setenv("KUBERNETES_SERVICE_HOST", "") + t.Setenv("XDG_STATE_HOME", t.TempDir()) + xdg.Reload() + + ctx := context.Background() + writableCapabilities := secrets.ProviderCapabilities{CanWrite: true} + + t.Run("initial registration creates secret names and persists metadata", func(t *testing.T) { + ctrl := gomock.NewController(t) + secretManager := secretsmocks.NewMockProvider(ctrl) + + const ( + clientSecretName = "OAUTH_CLIENT_SECRET_initial-registration" + regTokenName = "OAUTH_REG_TOKEN_initial-registration" + ) + secretExpiry := time.Date(2026, time.July, 15, 12, 0, 0, 0, time.UTC) + + gomock.InOrder( + secretManager.EXPECT().GetSecret(gomock.Any(), clientSecretName).Return("", assert.AnError), + secretManager.EXPECT().Capabilities().Return(writableCapabilities), + secretManager.EXPECT().SetSecret(ctx, clientSecretName, "initial-client-secret").Return(nil), + secretManager.EXPECT().GetSecret(gomock.Any(), regTokenName).Return("", assert.AnError), + secretManager.EXPECT().Capabilities().Return(writableCapabilities), + secretManager.EXPECT().SetSecret(ctx, regTokenName, "initial-registration-token").Return(nil), + ) + + remoteAuthConfig := &remote.Config{Issuer: "https://issuer.example.com"} + runConfig := NewRunConfig() + runConfig.Name = "initial-registration" + runConfig.BaseName = "initial-registration" + runConfig.RemoteAuthConfig = remoteAuthConfig + runner := &Runner{Config: runConfig} + + err := runner.persistClientCredentials( + ctx, + secretManager, + "initial-client-id", + "initial-client-secret", + secretExpiry, + "initial-registration-token", + "https://issuer.example.com/register/initial-client-id", + "client_secret_basic", + 9876, + ) + require.NoError(t, err) + + expected := remote.Config{ + Issuer: "https://issuer.example.com", + CachedClientID: "initial-client-id", + CachedClientSecretRef: clientSecretName, + CachedSecretExpiry: secretExpiry, + CachedRegTokenRef: regTokenName, + CachedRegClientURI: "https://issuer.example.com/register/initial-client-id", + CachedTokenEndpointAuthMethod: "client_secret_basic", + CachedDCRCallbackPort: 9876, + } + assert.Same(t, remoteAuthConfig, runner.Config.RemoteAuthConfig) + assert.Equal(t, expected, *runner.Config.RemoteAuthConfig) + + persisted, err := LoadState(ctx, runConfig.BaseName) + require.NoError(t, err) + require.NotNil(t, persisted.RemoteAuthConfig) + assert.Equal(t, expected, *persisted.RemoteAuthConfig) + }) + + t.Run("renewal reuses secret names and updates the existing config", func(t *testing.T) { + ctrl := gomock.NewController(t) + secretManager := secretsmocks.NewMockProvider(ctrl) + + const ( + clientSecretName = "OAUTH_CLIENT_SECRET_existing" + regTokenName = "OAUTH_REG_TOKEN_existing" + ) + secretExpiry := time.Date(2027, time.January, 2, 3, 4, 5, 0, time.UTC) + secretManager.EXPECT().Capabilities().Return(writableCapabilities) + secretManager.EXPECT().SetSecret(ctx, clientSecretName, "renewed-client-secret").Return(nil) + secretManager.EXPECT().Capabilities().Return(writableCapabilities) + secretManager.EXPECT().SetSecret(ctx, regTokenName, "renewed-registration-token").Return(nil) + + remoteAuthConfig := &remote.Config{ + Issuer: "https://issuer.example.com", + CachedClientID: "old-client-id", + CachedClientSecretRef: clientSecretName, + CachedSecretExpiry: time.Date(2026, time.January, 2, 3, 4, 5, 0, time.UTC), + CachedRegTokenRef: regTokenName, + CachedRegClientURI: "https://issuer.example.com/register/old-client-id", + CachedTokenEndpointAuthMethod: "none", + CachedDCRCallbackPort: 8080, + } + runConfig := NewRunConfig() + runConfig.Name = "renewal" + runConfig.BaseName = "renewal" + runConfig.RemoteAuthConfig = remoteAuthConfig + runner := &Runner{Config: runConfig} + + err := runner.persistClientCredentials( + ctx, + secretManager, + "renewed-client-id", + "renewed-client-secret", + secretExpiry, + "renewed-registration-token", + "https://issuer.example.com/register/renewed-client-id", + "client_secret_post", + 9090, + ) + require.NoError(t, err) + + expected := remote.Config{ + Issuer: "https://issuer.example.com", + CachedClientID: "renewed-client-id", + CachedClientSecretRef: clientSecretName, + CachedSecretExpiry: secretExpiry, + CachedRegTokenRef: regTokenName, + CachedRegClientURI: "https://issuer.example.com/register/renewed-client-id", + CachedTokenEndpointAuthMethod: "client_secret_post", + CachedDCRCallbackPort: 9090, + } + assert.Same(t, remoteAuthConfig, runner.Config.RemoteAuthConfig) + assert.Equal(t, expected, *runner.Config.RemoteAuthConfig) + + persisted, err := LoadState(ctx, runConfig.BaseName) + require.NoError(t, err) + require.NotNil(t, persisted.RemoteAuthConfig) + assert.Equal(t, expected, *persisted.RemoteAuthConfig) + }) + + t.Run("SaveState failure leaves the existing config untouched", func(t *testing.T) { + ctrl := gomock.NewController(t) + secretManager := secretsmocks.NewMockProvider(ctrl) + + const ( + clientSecretName = "OAUTH_CLIENT_SECRET_existing-failure" + regTokenName = "OAUTH_REG_TOKEN_existing-failure" + ) + secretManager.EXPECT().Capabilities().Return(writableCapabilities) + secretManager.EXPECT().SetSecret(ctx, clientSecretName, "new-client-secret").Return(nil) + secretManager.EXPECT().Capabilities().Return(writableCapabilities) + secretManager.EXPECT().SetSecret(ctx, regTokenName, "new-registration-token").Return(nil) + + original := remote.Config{ + Issuer: "https://issuer.example.com", + CachedClientID: "existing-client-id", + CachedClientSecretRef: clientSecretName, + CachedSecretExpiry: time.Date(2026, time.December, 1, 0, 0, 0, 0, time.UTC), + CachedRegTokenRef: regTokenName, + CachedRegClientURI: "https://issuer.example.com/register/existing-client-id", + CachedTokenEndpointAuthMethod: "client_secret_basic", + CachedDCRCallbackPort: 8080, + } + remoteAuthConfig := &remote.Config{} + *remoteAuthConfig = original + runConfig := NewRunConfig() + runConfig.Name = "failed-renewal" + runConfig.BaseName = "../invalid" + runConfig.RemoteAuthConfig = remoteAuthConfig + runner := &Runner{Config: runConfig} + + err := runner.persistClientCredentials( + ctx, + secretManager, + "new-client-id", + "new-client-secret", + time.Date(2027, time.December, 1, 0, 0, 0, 0, time.UTC), + "new-registration-token", + "https://issuer.example.com/register/new-client-id", + "client_secret_post", + 9090, + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to save config with client credentials") + assert.Same(t, remoteAuthConfig, runner.Config.RemoteAuthConfig) + assert.Equal(t, original, *runner.Config.RemoteAuthConfig) + }) +} + // mockMiddlewareImpl is a mock implementation of the types.Middleware interface type mockMiddlewareImpl struct { handler http.Handler diff --git a/pkg/script/executor.go b/pkg/script/executor.go index 0ebd9d8dbf..74d6f77da9 100644 --- a/pkg/script/executor.go +++ b/pkg/script/executor.go @@ -9,9 +9,9 @@ import ( "fmt" "strings" - "github.com/mark3labs/mcp-go/mcp" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/builtins" "github.com/stacklok/toolhive/pkg/script/internal/conversions" "github.com/stacklok/toolhive/pkg/script/internal/core" diff --git a/pkg/script/internal/builtins/builtins.go b/pkg/script/internal/builtins/builtins.go index 3cd94e220f..a7a2d655bf 100644 --- a/pkg/script/internal/builtins/builtins.go +++ b/pkg/script/internal/builtins/builtins.go @@ -8,9 +8,9 @@ import ( "context" "log/slog" - "github.com/mark3labs/mcp-go/mcp" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/conversions" ) diff --git a/pkg/script/internal/builtins/builtins_test.go b/pkg/script/internal/builtins/builtins_test.go index 2068004c77..07417926e0 100644 --- a/pkg/script/internal/builtins/builtins_test.go +++ b/pkg/script/internal/builtins/builtins_test.go @@ -10,10 +10,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/core" ) diff --git a/pkg/script/internal/builtins/tools.go b/pkg/script/internal/builtins/tools.go index 59766c2d1a..a990653741 100644 --- a/pkg/script/internal/builtins/tools.go +++ b/pkg/script/internal/builtins/tools.go @@ -7,9 +7,9 @@ import ( "context" "fmt" - "github.com/mark3labs/mcp-go/mcp" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/conversions" ) diff --git a/pkg/script/internal/conversions/result.go b/pkg/script/internal/conversions/result.go index 2af2951b8f..80717eb348 100644 --- a/pkg/script/internal/conversions/result.go +++ b/pkg/script/internal/conversions/result.go @@ -8,7 +8,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // ParseToolResult converts an mcp.CallToolResult into a Go value suitable diff --git a/pkg/script/internal/conversions/result_test.go b/pkg/script/internal/conversions/result_test.go index 5f534fe04c..821e06f1bb 100644 --- a/pkg/script/internal/conversions/result_test.go +++ b/pkg/script/internal/conversions/result_test.go @@ -6,8 +6,9 @@ package conversions import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestParseToolResult(t *testing.T) { diff --git a/pkg/script/script.go b/pkg/script/script.go index dc1091310b..7696e652e0 100644 --- a/pkg/script/script.go +++ b/pkg/script/script.go @@ -9,7 +9,7 @@ package script import ( "context" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // DefaultStepLimit is the default maximum number of Starlark execution steps. diff --git a/pkg/script/script_test.go b/pkg/script/script_test.go index a9f65ec50b..f92b761e7f 100644 --- a/pkg/script/script_test.go +++ b/pkg/script/script_test.go @@ -9,8 +9,9 @@ import ( "fmt" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestExecutor(t *testing.T) { diff --git a/pkg/server/discovery/discovery.go b/pkg/server/discovery/discovery.go index b00ecddfc2..349eede0c4 100644 --- a/pkg/server/discovery/discovery.go +++ b/pkg/server/discovery/discovery.go @@ -7,6 +7,7 @@ package discovery import ( + "context" "encoding/json" "errors" "fmt" @@ -45,10 +46,30 @@ type ServerInfo struct { StartedAt time.Time `json:"started_at"` } +// discoveryDirChain returns the directories ToolHive creates for the discovery +// file under base, outermost first. The last element is the discovery +// directory itself. +// +// The intermediate toolhive directory is part of the chain on purpose. A user +// who can delete or rename it can substitute a directory of their own, and +// every restriction applied to the leaf would then apply to a directory nobody +// reads. +func discoveryDirChain(base string) []string { + toolhiveDir := filepath.Join(base, "toolhive") + return []string{toolhiveDir, filepath.Join(toolhiveDir, "server")} +} + +// discoveryServerDir returns the leaf directory that holds server.json. +func discoveryServerDir(base string) string { + chain := discoveryDirChain(base) + return chain[len(chain)-1] +} + // defaultDiscoveryDir returns the default directory for the discovery file // based on the XDG Base Directory Specification. func defaultDiscoveryDir() string { - return filepath.Join(xdg.StateHome, "toolhive", "server") + chain := discoveryDirChain(xdg.StateHome) + return chain[len(chain)-1] } // FilePath returns the full path to the server discovery file @@ -57,10 +78,127 @@ func FilePath() string { return filepath.Join(defaultDiscoveryDir(), "server.json") } +// SecureDirResult reports whether ensureSecureDirIn repaired a previously +// insecure discovery directory chain. +type SecureDirResult struct { + RepairedInsecureChain bool +} + +// EnsureSecureDir creates the discovery directory chain and locks it down. +// Callers must invoke it before they read, lock, or otherwise trust anything +// under the discovery directory, not just before they write. +// +// Startup is the case that matters: pkg/api creates the directory, takes +// server.json.lock inside it, and runs Discover before it ever calls +// WriteServerInfo. If a directory left loose by an earlier run only got +// restricted by the write, a pre-planted server.json would be trusted for a +// whole startup cycle, and a Discover result of StateRunning returns before +// the write is reached at all. +func EnsureSecureDir() error { + _, err := EnsureSecureDirEx() + return err +} + +// EnsureSecureDirEx is like EnsureSecureDir but reports whether the chain was +// repaired from an insecure state. pkg/api uses this under the discovery +// file lock to fail closed on an existing record instead of deleting it. +func EnsureSecureDirEx() (SecureDirResult, error) { + return ensureSecureDirIn(xdg.StateHome) +} + +// ensureSecureDirIn creates the discovery directory chain under base and +// restricts the directories ToolHive owns in it. Windows hardens the full +// chain; other platforms chmod only the server leaf so shared toolhive state +// (runconfigs, toolhive.db) keeps its existing group permissions. +// +// A discovery file that existed while the chain was still loose is not removed +// here: it cannot be classified as legitimate or forged without the startup +// lock and a health check. pkg/api calls ReconcileDiscoveryAfterInsecureUpgrade +// under that lock instead. +func ensureSecureDirIn(base string) (SecureDirResult, error) { + secureDirs := discoveryDirsToSecure(base) + + hadInsecureChain, err := discoveryChainWasInsecure(secureDirs) + if err != nil { + return SecureDirResult{}, err + } + + fullChain := discoveryDirChain(base) + if err := mkdirDiscoveryChain(fullChain); err != nil { + return SecureDirResult{}, err + } + + for _, dir := range secureDirs { + if err := restrictDiscoveryDirPermissions(dir); err != nil { + return SecureDirResult{}, err + } + } + + return SecureDirResult{RepairedInsecureChain: hadInsecureChain}, nil +} + +func discoveryChainWasInsecure(chain []string) (bool, error) { + for _, dir := range chain { + loose, err := discoveryDirPermissionsLoose(dir) + if err != nil { + return false, err + } + if loose { + return true, nil + } + } + return false, nil +} + +// ReconcileDiscoveryAfterInsecureUpgrade runs under the discovery file lock +// after EnsureSecureDirEx repaired an insecure chain. An existing record cannot +// be classified as legitimate or forged from directory permissions alone, so +// startup fails closed when a healthy server answers, removes stale records, or +// leaves an unhealthy record for the caller to overwrite. It never deletes a +// record that might belong to a still-running peer. +func ReconcileDiscoveryAfterInsecureUpgrade(ctx context.Context, repaired bool) error { + if !repaired { + return nil + } + + path := FilePath() + if _, err := os.Stat(path); err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err + } + + result, err := Discover(ctx) + if err != nil { + return err + } + + switch result.State { + case StateRunning: + return fmt.Errorf("another ToolHive server is already running at %s (PID %d)", result.Info.URL, result.Info.PID) + case StateStale: + return CleanupStale() + case StateUnhealthy: + return nil + case StateNotFound: + return fmt.Errorf( + "refusing to start: discovery record %s exists after security upgrade but could not be validated; "+ + "remove it manually or stop the other server", + path, + ) + default: + return fmt.Errorf("refusing to start: unexpected discovery state %q after security upgrade", result.State) + } +} + // WriteServerInfo atomically writes the server discovery file. -// It creates the directory if needed, rejects symlinks at the target path, -// and writes with restricted permissions (0600). +// It creates and restricts the directory chain if needed, rejects symlinks at +// the target path, and writes with restricted permissions (0600). func WriteServerInfo(info *ServerInfo) error { + if err := EnsureSecureDir(); err != nil { + return err + } return writeServerInfoTo(defaultDiscoveryDir(), info) } @@ -78,8 +216,10 @@ func writeServerInfoTo(dir string, info *ServerInfo) error { // Tighten permissions on the directory in case it already existed with // looser permissions. MkdirAll only applies mode to newly-created dirs. - if err := os.Chmod(dir, dirPermissions); err != nil { - return fmt.Errorf("failed to set discovery directory permissions: %w", err) + // On Windows this sets an explicit protected DACL (POSIX modes are + // advisory on NTFS); elsewhere it is os.Chmod(dirPermissions). + if err := restrictDiscoveryDirPermissions(dir); err != nil { + return err } path := filepath.Join(dir, "server.json") @@ -114,6 +254,14 @@ func readServerInfoFrom(dir string) (*ServerInfo, error) { } } + // Reject a file another account owns for the same reason we reject a + // symlink: the URL and nonce inside it decide where clients send their + // API traffic. Tightening the directory stops new writes, but a file + // planted while the directory was still loose keeps its old ACL. + if err := validateDiscoveryFileOwner(path); err != nil { + return nil, err + } + data, err := os.ReadFile(path) // #nosec G304 -- path is constructed from a trusted XDG directory, not user input if err != nil { return nil, err diff --git a/pkg/server/discovery/discovery_test.go b/pkg/server/discovery/discovery_test.go index fadf0ec46a..f3b4a218dc 100644 --- a/pkg/server/discovery/discovery_test.go +++ b/pkg/server/discovery/discovery_test.go @@ -6,6 +6,7 @@ package discovery import ( "os" "path/filepath" + "runtime" "testing" "time" @@ -115,6 +116,12 @@ func TestRemoveServerInfo_NotFound(t *testing.T) { func TestWriteServerInfo_FilePermissions(t *testing.T) { t.Parallel() + // POSIX file modes are advisory on NTFS; Windows reports 0666 for + // newly written files regardless of the mode passed to AtomicWriteFile. + // Directory ACL coverage for Windows lives in permissions_windows_test.go. + if runtime.GOOS == "windows" { + t.Skip("POSIX file modes are not meaningful on Windows") + } dir := t.TempDir() info := &ServerInfo{ @@ -132,6 +139,11 @@ func TestWriteServerInfo_FilePermissions(t *testing.T) { func TestWriteServerInfo_CreatesDirectoryWithCorrectPermissions(t *testing.T) { t.Parallel() + // On Windows, writeServerInfoTo sets an explicit DACL instead of relying + // on os.Chmod; see TestWriteServerInfo_WindowsDACL_NoOtherInteractiveUsers. + if runtime.GOOS == "windows" { + t.Skip("POSIX directory modes are not meaningful on Windows; see DACL tests") + } parent := t.TempDir() dir := filepath.Join(parent, "nested", "server") @@ -148,6 +160,69 @@ func TestWriteServerInfo_CreatesDirectoryWithCorrectPermissions(t *testing.T) { assert.Equal(t, os.FileMode(dirPermissions), fi.Mode().Perm()) } +// TestEnsureSecureDirIn_CreatesAndRestrictsChain asserts the chain contract +// startup depends on: both the intermediate toolhive directory and the server +// leaf exist and are restricted before anything reads or locks inside them. +// The Windows DACL half is asserted in +// TestEnsureSecureDirIn_ProtectsIntermediateToolhiveDir. +func TestEnsureSecureDirIn_CreatesAndRestrictsChain(t *testing.T) { + t.Parallel() + + base := t.TempDir() + toolhiveDir := filepath.Join(base, "toolhive") + require.NoError(t, os.MkdirAll(toolhiveDir, 0750)) + + _, err := ensureSecureDirIn(base) + require.NoError(t, err) + + chain := discoveryDirChain(base) + require.Len(t, chain, 2) + toolhiveDir, serverDir := chain[0], chain[1] + for _, dir := range chain { + fi, err := os.Stat(dir) + require.NoError(t, err) + require.True(t, fi.IsDir(), "%s must be a directory", dir) + } + if runtime.GOOS == "windows" { + return + } + fi, err := os.Stat(toolhiveDir) + require.NoError(t, err) + assert.Equal(t, os.FileMode(0750), fi.Mode().Perm(), + "intermediate toolhive dir must keep its existing mode on POSIX") + fi, err = os.Stat(serverDir) + require.NoError(t, err) + assert.Equal(t, os.FileMode(dirPermissions), fi.Mode().Perm(), "mode of %s", serverDir) +} + +// TestEnsureSecureDirIn_TightensExistingChain covers the upgrade path: a chain +// left behind by an earlier run with loose permissions must be tightened, since +// startup trusts the discovery file inside it. +func TestEnsureSecureDirIn_TightensExistingChain(t *testing.T) { + t.Parallel() + if runtime.GOOS == "windows" { + t.Skip("POSIX directory modes are not meaningful on Windows; see DACL tests") + } + + base := t.TempDir() + chain := discoveryDirChain(base) + require.NoError(t, os.MkdirAll(chain[len(chain)-1], 0755)) + for _, dir := range chain { + require.NoError(t, os.Chmod(dir, 0755)) + } + + _, err := ensureSecureDirIn(base) + require.NoError(t, err) + + toolhiveDir, serverDir := chain[0], chain[1] + fi, err := os.Stat(toolhiveDir) + require.NoError(t, err) + assert.Equal(t, os.FileMode(0755), fi.Mode().Perm(), "intermediate toolhive dir must stay loose on POSIX") + fi, err = os.Stat(serverDir) + require.NoError(t, err) + assert.Equal(t, os.FileMode(dirPermissions), fi.Mode().Perm(), "mode of %s", serverDir) +} + func TestWriteServerInfo_RejectsSymlink(t *testing.T) { t.Parallel() dir := t.TempDir() @@ -195,6 +270,11 @@ func TestReadServerInfo_RejectsSymlink(t *testing.T) { func TestWriteServerInfo_TightensExistingDirPermissions(t *testing.T) { t.Parallel() + // On Windows the equivalent "tighten existing" path is covered by + // TestRestrictDiscoveryDirPermissions_ReplacesExistingLooseACL. + if runtime.GOOS == "windows" { + t.Skip("POSIX directory modes are not meaningful on Windows; see DACL tests") + } // Create a directory with deliberately too-loose permissions. dir := t.TempDir() diff --git a/pkg/server/discovery/permissions_other.go b/pkg/server/discovery/permissions_other.go new file mode 100644 index 0000000000..d01f7fa163 --- /dev/null +++ b/pkg/server/discovery/permissions_other.go @@ -0,0 +1,43 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package discovery + +import ( + "errors" + "fmt" + "os" +) + +// restrictDiscoveryDirPermissions tightens POSIX mode bits on the discovery +// directory. On non-Windows platforms this is the Chmod that previously lived +// inline in writeServerInfoTo. +func restrictDiscoveryDirPermissions(dir string) error { + if err := os.Chmod(dir, dirPermissions); err != nil { + return fmt.Errorf("failed to set discovery directory permissions: %w", err) + } + return nil +} + +// discoveryDirPermissionsLoose reports whether dir exists with looser than +// expected POSIX mode bits. +func discoveryDirPermissionsLoose(dir string) (bool, error) { + fi, err := os.Stat(dir) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("failed to stat discovery directory: %w", err) + } + return fi.Mode().Perm() != dirPermissions, nil +} + +// validateDiscoveryFileOwner is a no-op outside Windows. POSIX mode bits are +// enforced rather than advisory, so a 0700 directory owned by the user already +// keeps other accounts from planting a discovery file, and no existing +// non-Windows behavior changes here. +func validateDiscoveryFileOwner(_ string) error { + return nil +} diff --git a/pkg/server/discovery/permissions_windows.go b/pkg/server/discovery/permissions_windows.go new file mode 100644 index 0000000000..3e5745b2bd --- /dev/null +++ b/pkg/server/discovery/permissions_windows.go @@ -0,0 +1,332 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build windows + +package discovery + +import ( + "errors" + "fmt" + "os" + "unsafe" + + "golang.org/x/sys/windows" +) + +// Ancestor invariant on Windows: ToolHive protects and owner-checks every +// directory it creates for the discovery file (%LOCALAPPDATA%\toolhive and +// %LOCALAPPDATA%\toolhive\server). Above that it relies on the OS default ACL +// of %LOCALAPPDATA%, which grants the user, SYSTEM, and Administrators only. +// If another account does hold delete-child there, it can rename the chain +// aside and pre-create its own; the owner check below is what turns that into +// a startup failure instead of silent trust, because a standard user cannot +// create a directory owned by the ToolHive user, SYSTEM, or Administrators. + +// restrictDiscoveryDirPermissions locks down one directory in the discovery +// chain: it fails closed unless the directory is owned by an account we trust, +// then replaces the DACL with an explicit ACL granting FILE-equivalent +// GenericAll only to the process user and SYSTEM, marked protected so parent +// ACEs cannot re-inherit. os.Chmod is advisory on NTFS and does not strip +// inherited ACEs under %LOCALAPPDATA%, which is the gap that lets another +// interactive user rewrite server.json (for example to point at a named pipe +// of their own). +// +// Inheritance (OICI) is intentional: server.json and any future children pick +// up the same restriction instead of inheriting a looser parent ACL. +func restrictDiscoveryDirPermissions(dir string) error { + trusted, err := trustedOwnerSIDs() + if err != nil { + return err + } + return restrictDiscoveryDir(dir, trusted) +} + +// discoveryDirPermissionsLoose reports whether dir carries an unprotected DACL +// or grants access to accounts outside the process user and SYSTEM. +func discoveryDirPermissionsLoose(dir string) (bool, error) { + if _, err := os.Stat(dir); err != nil { + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("failed to stat discovery directory: %w", err) + } + + sd, err := windows.GetNamedSecurityInfo(dir, windows.SE_FILE_OBJECT, windows.DACL_SECURITY_INFORMATION) + if err != nil { + return false, fmt.Errorf("failed to read discovery directory DACL: %w", err) + } + + control, _, err := sd.Control() + if err != nil { + return false, fmt.Errorf("failed to read discovery directory DACL control: %w", err) + } + if control&windows.SE_DACL_PROTECTED == 0 { + return true, nil + } + + dacl, _, err := sd.DACL() + if err != nil { + return false, fmt.Errorf("failed to read discovery directory DACL entries: %w", err) + } + if dacl == nil { + return true, nil + } + + userSID, err := currentProcessUserSID() + if err != nil { + return false, err + } + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + if err != nil { + return false, err + } + + aces, err := allowACEsFromACL(dacl) + if err != nil { + return false, err + } + for _, ace := range aces { + sid := (*windows.SID)(unsafe.Pointer(&ace.SidStart)) + if userSID.Equals(sid) || systemSID.Equals(sid) { + continue + } + return true, nil + } + return false, nil +} + +func allowACEsFromACL(acl *windows.ACL) ([]*windows.ACCESS_ALLOWED_ACE, error) { + aces := make([]*windows.ACCESS_ALLOWED_ACE, 0, acl.AceCount) + for i := uint16(0); i < acl.AceCount; i++ { + var ace *windows.ACCESS_ALLOWED_ACE + if err := windows.GetAce(acl, uint32(i), &ace); err != nil { + return nil, err + } + if ace.Header.AceType != windows.ACCESS_ALLOWED_ACE_TYPE { + continue + } + aces = append(aces, ace) + } + return aces, nil +} + +// restrictDiscoveryDir is the injectable core of +// restrictDiscoveryDirPermissions. Tests pass a trustedOwners set that +// excludes the real owner to exercise the fail-closed path without a second +// user account. +func restrictDiscoveryDir(dir string, trustedOwners []*windows.SID) error { + // Check ownership before writing the DACL, not after. Owners keep + // WRITE_DAC implicitly, so a DACL we set on a directory somebody else owns + // is cosmetic: that owner can make it permissive again at any time. + if err := checkOwner(dir, "directory", trustedOwners); err != nil { + return err + } + + userSID, err := currentProcessUserSID() + if err != nil { + return fmt.Errorf("failed to resolve current user SID: %w", err) + } + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + if err != nil { + return fmt.Errorf("failed to resolve SYSTEM SID: %w", err) + } + + acl, err := windows.ACLFromEntries([]windows.EXPLICIT_ACCESS{ + { + AccessPermissions: windows.GENERIC_ALL, + AccessMode: windows.GRANT_ACCESS, + Inheritance: windows.OBJECT_INHERIT_ACE | windows.CONTAINER_INHERIT_ACE, + Trustee: windows.TRUSTEE{ + TrusteeForm: windows.TRUSTEE_IS_SID, + TrusteeValue: windows.TrusteeValueFromSID(userSID), + }, + }, + { + AccessPermissions: windows.GENERIC_ALL, + AccessMode: windows.GRANT_ACCESS, + Inheritance: windows.OBJECT_INHERIT_ACE | windows.CONTAINER_INHERIT_ACE, + Trustee: windows.TRUSTEE{ + TrusteeForm: windows.TRUSTEE_IS_SID, + TrusteeType: windows.TRUSTEE_IS_WELL_KNOWN_GROUP, + TrusteeValue: windows.TrusteeValueFromSID(systemSID), + }, + }, + }, nil) + if err != nil { + return fmt.Errorf("failed to build discovery directory DACL: %w", err) + } + + if err := windows.SetNamedSecurityInfo( + dir, + windows.SE_FILE_OBJECT, + windows.DACL_SECURITY_INFORMATION|windows.PROTECTED_DACL_SECURITY_INFORMATION, + nil, + nil, + acl, + nil, + ); err != nil { + return fmt.Errorf("failed to set discovery directory DACL: %w", err) + } + + // Re-check the owner before reporting the lockdown as complete. A + // pre-creating owner that raced the DACL write could have taken the + // directory back in between, and the caller is about to trust the contents. + return checkOwner(dir, "directory", trustedOwners) +} + +// validateDiscoveryFileOwner rejects a discovery file owned by an account +// outside the trusted set. The directory lockdown stops new writes but leaves +// the ACL of a file planted while the directory was loose untouched, so the +// file needs its own check before its URL and nonce are believed. +func validateDiscoveryFileOwner(path string) error { + trusted, err := trustedOwnerSIDs() + if err != nil { + return err + } + return validateFileOwner(path, trusted) +} + +// validateFileOwner is the injectable core of validateDiscoveryFileOwner. +func validateFileOwner(path string, trustedOwners []*windows.SID) error { + if _, err := os.Lstat(path); err != nil { + // A missing file is not a trust decision; the caller's own + // os.ErrNotExist handling covers it. + if errors.Is(err, os.ErrNotExist) { + return nil + } + return fmt.Errorf("failed to stat discovery file: %w", err) + } + return checkOwner(path, "file", trustedOwners) +} + +// checkOwner fails closed unless path is owned by one of trustedOwners. +func checkOwner(path string, kind string, trustedOwners []*windows.SID) error { + sd, err := windows.GetNamedSecurityInfo(path, windows.SE_FILE_OBJECT, windows.OWNER_SECURITY_INFORMATION) + if err != nil { + return fmt.Errorf("failed to read discovery %s owner for %s: %w", kind, path, err) + } + owner, _, err := sd.Owner() + if err != nil { + return fmt.Errorf("failed to resolve discovery %s owner for %s: %w", kind, path, err) + } + for _, trustedOwner := range trustedOwners { + if trustedOwner.Equals(owner) { + return nil + } + } + return fmt.Errorf( + "refusing to trust discovery %s %s: it is owned by %s, not the current user, SYSTEM, or Administrators; "+ + "remove it (or restore its ownership) and retry", + kind, path, owner, + ) +} + +// trustedOwnerSIDs returns the accounts allowed to own the discovery directory +// chain and the discovery file: the process user, SYSTEM, and Administrators. +// +// Administrators is in the set because a directory created by an elevated +// `thv serve` is owned by Administrators rather than by the user, and because +// a local administrator is outside this threat model (they can take ownership +// of anything). It does not weaken the check against the attacker this fix is +// about: a standard user cannot make SYSTEM or Administrators the owner of a +// directory they create, so a pre-created hostile directory still fails. +func trustedOwnerSIDs() ([]*windows.SID, error) { + userSID, err := currentProcessUserSID() + if err != nil { + return nil, fmt.Errorf("failed to resolve current user SID: %w", err) + } + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + if err != nil { + return nil, fmt.Errorf("failed to resolve SYSTEM SID: %w", err) + } + adminsSID, err := windows.CreateWellKnownSid(windows.WinBuiltinAdministratorsSid) + if err != nil { + return nil, fmt.Errorf("failed to resolve Administrators SID: %w", err) + } + return []*windows.SID{userSID, systemSID, adminsSID}, nil +} + +func currentProcessUserSID() (*windows.SID, error) { + token := windows.GetCurrentProcessToken() + tokenUser, err := token.GetTokenUser() + if err != nil { + return nil, err + } + return tokenUser.User.Sid, nil +} + +// ValidateRestrictedDiscoveryDACL reports whether dir carries a protected DACL +// that grants FILE access only to the current process user and SYSTEM. +func ValidateRestrictedDiscoveryDACL(dir string) error { + userSID, err := currentProcessUserSID() + if err != nil { + return err + } + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + if err != nil { + return err + } + everyoneSID, err := windows.CreateWellKnownSid(windows.WinWorldSid) + if err != nil { + return err + } + authUsersSID, err := windows.CreateWellKnownSid(windows.WinAuthenticatedUserSid) + if err != nil { + return err + } + + sd, err := windows.GetNamedSecurityInfo( + dir, + windows.SE_FILE_OBJECT, + windows.DACL_SECURITY_INFORMATION, + ) + if err != nil { + return fmt.Errorf("failed to read discovery directory DACL: %w", err) + } + + control, _, err := sd.Control() + if err != nil { + return fmt.Errorf("failed to read discovery directory DACL control: %w", err) + } + if control&windows.SE_DACL_PROTECTED == 0 { + return fmt.Errorf("DACL of %s must be protected against inheritance", dir) + } + + dacl, _, err := sd.DACL() + if err != nil { + return fmt.Errorf("failed to read discovery directory DACL entries: %w", err) + } + if dacl == nil { + return fmt.Errorf("DACL of %s must not be empty", dir) + } + + aces, err := allowACEsFromACL(dacl) + if err != nil { + return err + } + + var userSeen, systemSeen bool + for _, ace := range aces { + sid := (*windows.SID)(unsafe.Pointer(&ace.SidStart)) + switch { + case userSID.Equals(sid): + userSeen = true + case systemSID.Equals(sid): + systemSeen = true + case everyoneSID.Equals(sid): + return fmt.Errorf("DACL still grants Everyone (%s)", sid) + case authUsersSID.Equals(sid): + return fmt.Errorf("DACL still grants Authenticated Users (%s)", sid) + default: + return fmt.Errorf("unexpected allow ACE for SID %s (want only current user + SYSTEM)", sid) + } + } + if !userSeen { + return fmt.Errorf("DACL of %s must grant the current process user", dir) + } + if !systemSeen { + return fmt.Errorf("DACL of %s must grant SYSTEM", dir) + } + return nil +} diff --git a/pkg/server/discovery/permissions_windows_test.go b/pkg/server/discovery/permissions_windows_test.go new file mode 100644 index 0000000000..71f42517a1 --- /dev/null +++ b/pkg/server/discovery/permissions_windows_test.go @@ -0,0 +1,289 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build windows + +package discovery + +import ( + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + "unsafe" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" +) + +// TestWriteServerInfo_WindowsDACL_NoOtherInteractiveUsers asserts the +// acceptance criterion from #5217: after writeServerInfoTo, the discovery +// directory DACL grants access only to the current user and SYSTEM, and does +// not retain Everyone / Authenticated Users / other interactive-user ACEs +// that MkdirAll would otherwise inherit (and that os.Chmod cannot strip). +func TestWriteServerInfo_WindowsDACL_NoOtherInteractiveUsers(t *testing.T) { + t.Parallel() + + parent := t.TempDir() + dir := filepath.Join(parent, "toolhive", "server") + + // Seed a deliberately loose ACL on the parent so newly created children + // inherit Everyone. This models a shared / misconfigured LOCALAPPDATA + // tree better than relying on whatever TempDir happens to carry. + grantEveryone(t, parent) + + info := &ServerInfo{ + URL: "npipe://thv-api", + PID: 1, + Nonce: "dacl-nonce", + StartedAt: time.Now().UTC(), + } + require.NoError(t, writeServerInfoTo(dir, info)) + + assertDiscoveryDACLRestricted(t, dir) +} + +// TestRestrictDiscoveryDirPermissions_ReplacesExistingLooseACL covers the +// sibling failure mode: the discovery directory already exists with a loose +// DACL (Everyone + inherited ACEs). restrictDiscoveryDirPermissions must +// replace that ACL rather than merge, and block further inheritance. +func TestRestrictDiscoveryDirPermissions_ReplacesExistingLooseACL(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + grantEveryone(t, dir) + + before := discoveryDirSDDL(t, dir) + require.Contains(t, strings.ToUpper(before), "WD", "precondition: Everyone (WD) must be present before restrict") + + require.NoError(t, restrictDiscoveryDirPermissions(dir)) + assertDiscoveryDACLRestricted(t, dir) +} + +// TestRestrictDiscoveryDirPermissions_NewDirectory covers the create path +// (no pre-existing ACL to replace) so MkdirAll + restrict still lands a +// protected owner/SYSTEM-only DACL. +func TestRestrictDiscoveryDirPermissions_NewDirectory(t *testing.T) { + t.Parallel() + + parent := t.TempDir() + grantEveryone(t, parent) + dir := filepath.Join(parent, "fresh-server") + require.NoError(t, os.MkdirAll(dir, dirPermissions)) + + require.NoError(t, restrictDiscoveryDirPermissions(dir)) + assertDiscoveryDACLRestricted(t, dir) +} + +// TestEnsureSecureDirIn_ProtectsIntermediateToolhiveDir covers the ancestor +// replacement path. Restricting only the server leaf leaves the intermediate +// toolhive directory with inherited Modify, which carries DELETE on that +// object: another interactive user can rename toolhive aside, recreate +// toolhive\server with an ACL of their own, and the protected leaf is bypassed +// because nothing reads it any more. Both directories in the chain must end up +// protected. +func TestEnsureSecureDirIn_ProtectsIntermediateToolhiveDir(t *testing.T) { + t.Parallel() + + base := t.TempDir() + grantEveryone(t, base) + + chain := discoveryDirChain(base) + require.Len(t, chain, 2) + toolhiveDir, serverDir := chain[0], chain[1] + + // Model the state an earlier ToolHive left behind: both directories exist + // already, created with plain MkdirAll, so both inherited Everyone. + require.NoError(t, os.MkdirAll(serverDir, dirPermissions)) + require.Contains(t, strings.ToUpper(discoveryDirSDDL(t, toolhiveDir)), "WD", + "precondition: intermediate toolhive dir must inherit Everyone (WD)") + + _, err := ensureSecureDirIn(base) + require.NoError(t, err) + + assertDiscoveryDACLRestricted(t, toolhiveDir) + assertDiscoveryDACLRestricted(t, serverDir) +} + +// TestEnsureSecureDirIn_PreservesDiscoveryFileOnUpgrade ensures repairing an +// insecure chain does not delete an existing server.json. pkg/api decides +// under the startup lock whether the record is safe to keep. +func TestEnsureSecureDirIn_PreservesDiscoveryFileOnUpgrade(t *testing.T) { + t.Parallel() + + base := t.TempDir() + grantEveryone(t, base) + + chain := discoveryDirChain(base) + serverDir := chain[len(chain)-1] + require.NoError(t, os.MkdirAll(serverDir, 0700)) + + planted := []byte(`{"url":"npipe://attacker","pid":1,"nonce":"forged","started_at":"2026-08-03T00:00:00Z"}`) + require.NoError(t, os.WriteFile(filepath.Join(serverDir, "server.json"), planted, 0600)) + require.Contains(t, strings.ToUpper(discoveryDirSDDL(t, serverDir)), "WD", + "precondition: server dir must inherit Everyone (WD)") + + _, err := ensureSecureDirIn(base) + require.NoError(t, err) + + assertDiscoveryDACLRestricted(t, chain[0]) + assertDiscoveryDACLRestricted(t, serverDir) + got, err := os.ReadFile(filepath.Join(serverDir, "server.json")) + require.NoError(t, err) + assert.Equal(t, string(planted), string(got)) +} + +// TestRestrictDiscoveryDir_FailsClosedOnUntrustedOwner covers hostile +// ownership. Replacing the DACL is not enough when another account owns the +// directory: owners keep WRITE_DAC implicitly and can make the ACL permissive +// again, so the lockdown must fail rather than report success. +// +// Creating a directory owned by a different account needs a second user or +// SeRestorePrivilege, neither of which a unit test can assume, so the trusted +// owner set is injected instead: the directory is genuinely owned by the test +// user and the set deliberately does not include them. +func TestRestrictDiscoveryDir_FailsClosedOnUntrustedOwner(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + grantEveryone(t, dir) + + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + require.NoError(t, err) + + err = restrictDiscoveryDir(dir, []*windows.SID{systemSID}) + require.Error(t, err) + assert.Contains(t, err.Error(), "refusing to trust discovery directory") + + // The DACL must be left untouched: a half-applied lockdown on a directory + // somebody else owns would look secure to an operator reading icacls. + assert.Contains(t, strings.ToUpper(discoveryDirSDDL(t, dir)), "WD", + "loose ACE must survive: ownership is checked before the DACL is replaced") +} + +// TestRestrictDiscoveryDir_AcceptsCurrentUserOwner is the positive half of the +// ownership check: the real trusted set must accept a directory the process +// created itself, so the fail-closed path above cannot pass by rejecting +// everything. +func TestRestrictDiscoveryDir_AcceptsCurrentUserOwner(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + trusted, err := trustedOwnerSIDs() + require.NoError(t, err) + + require.NoError(t, restrictDiscoveryDir(dir, trusted)) + assertDiscoveryDACLRestricted(t, dir) +} + +// TestValidateFileOwner_RejectsUntrustedOwner covers the read path. The +// directory lockdown stops new writes but does not rewrite the ACL of a +// server.json planted while the directory was still loose, so the file's own +// owner decides whether its URL and nonce can be trusted. +func TestValidateFileOwner_RejectsUntrustedOwner(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + info := &ServerInfo{ + URL: "npipe://attacker-pipe", + PID: 1, + Nonce: "planted-nonce", + StartedAt: time.Now().UTC(), + } + require.NoError(t, writeServerInfoTo(dir, info)) + path := filepath.Join(dir, "server.json") + + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + require.NoError(t, err) + + err = validateFileOwner(path, []*windows.SID{systemSID}) + require.Error(t, err) + assert.Contains(t, err.Error(), "refusing to trust discovery file") + + // A file the process owns stays readable, and a missing file is not a + // trust decision (readServerInfoFrom must still return os.ErrNotExist). + trusted, err := trustedOwnerSIDs() + require.NoError(t, err) + assert.NoError(t, validateFileOwner(path, trusted)) + assert.NoError(t, validateFileOwner(filepath.Join(dir, "absent.json"), trusted)) + + got, err := readServerInfoFrom(dir) + require.NoError(t, err) + assert.Equal(t, "planted-nonce", got.Nonce) +} + +func grantEveryone(t *testing.T, path string) { + t.Helper() + // icacls grants are the product-path way to introduce a loose ACE; + // quoting keeps PowerShell from expanding (OI)/(CI). + cmd := exec.Command("icacls", path, "/grant", "*S-1-1-0:(OI)(CI)M") + out, err := cmd.CombinedOutput() + require.NoError(t, err, "icacls grant Everyone failed: %s", out) +} + +func assertDiscoveryDACLRestricted(t *testing.T, dir string) { + t.Helper() + + userSID, err := currentProcessUserSID() + require.NoError(t, err) + systemSID, err := windows.CreateWellKnownSid(windows.WinLocalSystemSid) + require.NoError(t, err) + everyoneSID, err := windows.CreateWellKnownSid(windows.WinWorldSid) + require.NoError(t, err) + authUsersSID, err := windows.CreateWellKnownSid(windows.WinAuthenticatedUserSid) + require.NoError(t, err) + + sd, err := windows.GetNamedSecurityInfo( + dir, + windows.SE_FILE_OBJECT, + windows.DACL_SECURITY_INFORMATION, + ) + require.NoError(t, err) + + control, _, err := sd.Control() + require.NoError(t, err) + assert.NotZero(t, control&windows.SE_DACL_PROTECTED, "DACL must be protected against inheritance") + + dacl, _, err := sd.DACL() + require.NoError(t, err) + require.NotNil(t, dacl) + + aces, err := allowACEsFromACL(dacl) + require.NoError(t, err) + + var userSeen, systemSeen bool + for _, ace := range aces { + sid := (*windows.SID)(unsafe.Pointer(&ace.SidStart)) + switch { + case userSID.Equals(sid): + userSeen = true + case systemSID.Equals(sid): + systemSeen = true + case everyoneSID.Equals(sid): + t.Fatalf("DACL still grants Everyone (%s)", sid) + case authUsersSID.Equals(sid): + t.Fatalf("DACL still grants Authenticated Users (%s)", sid) + default: + // Administrators / package SIDs / other interactive users must + // not remain after an explicit replace. Fail closed on anything + // that is not the process user or SYSTEM. + t.Fatalf("unexpected allow ACE for SID %s (want only current user + SYSTEM)", sid) + } + } + assert.True(t, userSeen, "DACL must grant the current process user") + assert.True(t, systemSeen, "DACL must grant SYSTEM") +} + +func discoveryDirSDDL(t *testing.T, dir string) string { + t.Helper() + sd, err := windows.GetNamedSecurityInfo( + dir, + windows.SE_FILE_OBJECT, + windows.DACL_SECURITY_INFORMATION, + ) + require.NoError(t, err) + return sd.String() +} diff --git a/pkg/server/discovery/secure_dir_other.go b/pkg/server/discovery/secure_dir_other.go new file mode 100644 index 0000000000..c0f139fdf4 --- /dev/null +++ b/pkg/server/discovery/secure_dir_other.go @@ -0,0 +1,40 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package discovery + +import ( + "fmt" + "os" +) + +const discoveryParentDirPermissions = 0750 + +// discoveryDirsToSecure returns only the server leaf on POSIX. The shared +// toolhive parent also holds runconfigs and toolhive.db created with 0750; +// chmodding it to 0700 would revoke group traversal for unrelated state. +func discoveryDirsToSecure(base string) []string { + return []string{discoveryServerDir(base)} +} + +// mkdirDiscoveryChain creates the discovery path without forcing 0700 on the +// shared intermediate toolhive directory. Parents are created first with +// discoveryParentDirPermissions so a later MkdirAll on the leaf does not +// apply dirPermissions up the chain. +func mkdirDiscoveryChain(chain []string) error { + if len(chain) == 0 { + return nil + } + for _, dir := range chain[:len(chain)-1] { + if err := os.MkdirAll(dir, discoveryParentDirPermissions); err != nil { + return fmt.Errorf("failed to create discovery directory: %w", err) + } + } + leaf := chain[len(chain)-1] + if err := os.MkdirAll(leaf, dirPermissions); err != nil { + return fmt.Errorf("failed to create discovery directory: %w", err) + } + return nil +} diff --git a/pkg/server/discovery/secure_dir_windows.go b/pkg/server/discovery/secure_dir_windows.go new file mode 100644 index 0000000000..bdf61e996e --- /dev/null +++ b/pkg/server/discovery/secure_dir_windows.go @@ -0,0 +1,26 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build windows + +package discovery + +import ( + "fmt" + "os" +) + +// discoveryDirsToSecure returns every directory in the discovery chain that +// receives an explicit protected DACL on Windows. +func discoveryDirsToSecure(base string) []string { + return discoveryDirChain(base) +} + +func mkdirDiscoveryChain(chain []string) error { + for _, dir := range chain { + if err := os.MkdirAll(dir, dirPermissions); err != nil { + return fmt.Errorf("failed to create discovery directory: %w", err) + } + } + return nil +} diff --git a/pkg/server/discovery/validate_dacl_other.go b/pkg/server/discovery/validate_dacl_other.go new file mode 100644 index 0000000000..29230a33fd --- /dev/null +++ b/pkg/server/discovery/validate_dacl_other.go @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package discovery + +import "fmt" + +// ValidateRestrictedDiscoveryDACL is only meaningful on Windows. +func ValidateRestrictedDiscoveryDACL(dir string) error { + return fmt.Errorf("ValidateRestrictedDiscoveryDACL is only supported on Windows (%s)", dir) +} diff --git a/pkg/skills/client/client.go b/pkg/skills/client/client.go index 4c5c69a606..7fbc5b9a37 100644 --- a/pkg/skills/client/client.go +++ b/pkg/skills/client/client.go @@ -12,6 +12,7 @@ import ( "fmt" "io" "log/slog" + "net" "net/http" "net/url" "strings" @@ -24,10 +25,18 @@ import ( ) const ( - skillsBasePath = "/api/v1beta/skills" - defaultBaseURL = "http://127.0.0.1:8080" - defaultTimeout = 30 * time.Second + skillsBasePath = "/api/v1beta/skills" + defaultBaseURL = "http://127.0.0.1:8080" + // defaultTimeout is a backstop against a wedged server, not a budget for + // the operation. Install, sync, and upgrade pull OCI artifacts of + // unbounded size over the user's connection, so a timeout tight enough to + // feel responsive would abort legitimate work — and because giving up + // cancels the request context, it aborts it server-side too, leaving + // nothing for a retry to reuse. Callers who want to fail fast can set + // TOOLHIVE_API_TIMEOUT or pass WithTimeout. + defaultTimeout = 10 * time.Minute envAPIURL = "TOOLHIVE_API_URL" + envAPITimeout = "TOOLHIVE_API_TIMEOUT" maxResponseSize = 1 << 20 // 1 MiB — matches server-side maxRequestBodySize maxErrorBodySize = 1 << 16 // 64 KiB — matches auth/token and DCR limits ) @@ -37,8 +46,16 @@ const ( // running. var ErrServerUnreachable = errors.New("could not reach ToolHive API server — is 'thv serve' running?") +// ErrRequestTimeout is returned when the server was reached but did not +// respond within the client timeout. This is a distinct condition from +// ErrServerUnreachable: the server is healthy and was working on the request. +// Note that abandoning the request cancels its context server-side, so the +// operation does not continue in the background and retrying restarts it. +var ErrRequestTimeout = errors.New("the ToolHive API server did not respond within the client timeout") + // Compile-time interface check. var _ skills.SkillService = (*Client)(nil) +var _ skills.SkillLockService = (*Client)(nil) // Client is an HTTP client for the ToolHive Skills API. type Client struct { @@ -96,23 +113,47 @@ type discoverFunc func(ctx context.Context) (string, []Option) // envReader and discover dependencies are injected so each resolution step // can be exercised in isolation. func newDefaultClientWithEnv(ctx context.Context, envReader env.Reader, discover discoverFunc, opts ...Option) *Client { + // An env-supplied timeout applies to every resolution path, but ranks + // below caller-supplied options so a WithTimeout argument still wins. + withEnvTimeout := func(base []Option) []Option { + merged := make([]Option, 0, len(base)+len(opts)+1) + merged = append(merged, base...) + if d, ok := timeoutFromEnv(envReader); ok { + merged = append(merged, WithTimeout(d)) + } + return append(merged, opts...) + } + // 1. Explicit env var override always wins. if base := envReader.Getenv(envAPIURL); base != "" { - return NewClient(base, opts...) + return NewClient(base, withEnvTimeout(nil)...) } // 2. Try server discovery. if base, httpOpts := discover(ctx); base != "" { // Discovery opts go first so caller-supplied opts can override them // (e.g. a caller-provided WithTimeout replaces the discovery default). - merged := make([]Option, 0, len(httpOpts)+len(opts)) - merged = append(merged, httpOpts...) - merged = append(merged, opts...) - return NewClient(base, merged...) + return NewClient(base, withEnvTimeout(httpOpts)...) } // 3. Fall back to the default URL. - return NewClient(defaultBaseURL, opts...) + return NewClient(defaultBaseURL, withEnvTimeout(nil)...) +} + +// timeoutFromEnv reads TOOLHIVE_API_TIMEOUT as a Go duration (e.g. "45s", +// "5m"). An unset, unparsable, or non-positive value is ignored so a typo +// degrades to the default rather than disabling the timeout entirely. +func timeoutFromEnv(envReader env.Reader) (time.Duration, bool) { + raw := strings.TrimSpace(envReader.Getenv(envAPITimeout)) + if raw == "" { + return 0, false + } + d, err := time.ParseDuration(raw) + if err != nil || d <= 0 { + slog.Warn("ignoring invalid "+envAPITimeout, "value", raw, "error", err) + return 0, false + } + return d, true } // resolveViaDiscovery attempts to find a running server via the discovery file. @@ -163,13 +204,14 @@ func (c *Client) List(ctx context.Context, opts skills.ListOptions) ([]skills.In // Install installs a skill from a remote source. func (c *Client) Install(ctx context.Context, opts skills.InstallOptions) (*skills.InstallResult, error) { body := installRequest{ - Name: opts.Name, - Version: opts.Version, - Scope: opts.Scope, - ProjectRoot: opts.ProjectRoot, - Clients: opts.Clients, - Force: opts.Force, - Group: opts.Group, + Name: opts.Name, + Version: opts.Version, + Scope: opts.Scope, + ProjectRoot: opts.ProjectRoot, + Clients: opts.Clients, + Force: opts.Force, + Group: opts.Group, + AllowUnsigned: opts.AllowUnsigned, } var resp installResponse @@ -267,6 +309,44 @@ func (c *Client) GetContent(ctx context.Context, opts skills.ContentOptions) (*s return &content, nil } +// Sync restores a project's installed skills to match its lock file. +func (c *Client) Sync(ctx context.Context, opts skills.SyncOptions) (*skills.SyncResult, error) { + body := syncRequest{ + ProjectRoot: opts.ProjectRoot, + Clients: opts.Clients, + Prune: opts.Prune, + Check: opts.Check, + Adopt: opts.Adopt, + AllowUnsigned: opts.AllowUnsigned, + } + + var result skills.SyncResult + if err := c.doJSONRequest(ctx, http.MethodPost, "/sync", nil, body, &result); err != nil { + return nil, err + } + return &result, nil +} + +// Upgrade re-resolves a project's lock entries and installs newer content +// where available. +func (c *Client) Upgrade(ctx context.Context, opts skills.UpgradeOptions) (*skills.UpgradeResult, error) { + body := upgradeRequest{ + ProjectRoot: opts.ProjectRoot, + Names: opts.Names, + Preview: opts.Preview, + FailOnChanges: opts.FailOnChanges, + AllowRefChange: opts.AllowRefChange, + AllowSignerChange: opts.AllowSignerChange, + Clients: opts.Clients, + } + + var result skills.UpgradeResult + if err := c.doJSONRequest(ctx, http.MethodPost, "/upgrade", nil, body, &result); err != nil { + return nil, err + } + return &result, nil +} + // --- internal helpers --- func (c *Client) buildURL(path string, query url.Values) string { @@ -310,7 +390,7 @@ func (c *Client) doJSONRequest( resp, err := c.httpClient.Do(req) // #nosec G704 -- baseURL is a trusted local API server URL if err != nil { - return fmt.Errorf("%w: %w", ErrServerUnreachable, err) + return fmt.Errorf("%w: %w", classifyTransportError(ctx, err), err) } defer func() { _ = resp.Body.Close() }() @@ -328,6 +408,25 @@ func (c *Client) doJSONRequest( return nil } +// classifyTransportError decides which sentinel a failed round-trip belongs +// to. Reporting a timeout as an unreachable server sends users to check +// whether "thv serve" is running when it is running and was mid-operation. +func classifyTransportError(ctx context.Context, err error) error { + // The caller's own context ending takes precedence — neither the server's + // availability nor the client timeout is why the request stopped. + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return ErrRequestTimeout + } + if errors.Is(err, context.DeadlineExceeded) { + return ErrRequestTimeout + } + return ErrServerUnreachable +} + // handleErrorResponse reads the response body and returns an *httperr.CodedError. func handleErrorResponse(resp *http.Response) error { body, err := io.ReadAll(io.LimitReader(resp.Body, maxErrorBodySize)) diff --git a/pkg/skills/client/client_test.go b/pkg/skills/client/client_test.go index db747aed8b..0ba894bf45 100644 --- a/pkg/skills/client/client_test.go +++ b/pkg/skills/client/client_test.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "io" + "net" "net/http" "net/http/httptest" "testing" @@ -656,6 +657,7 @@ func TestNewDefaultClient(t *testing.T) { ctrl := gomock.NewController(t) mockEnv := envmocks.NewMockReader(ctrl) mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery) assert.Equal(t, defaultBaseURL, c.baseURL) @@ -666,6 +668,7 @@ func TestNewDefaultClient(t *testing.T) { ctrl := gomock.NewController(t) mockEnv := envmocks.NewMockReader(ctrl) mockEnv.EXPECT().Getenv(envAPIURL).Return("http://localhost:9999") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() c := newDefaultClientWithEnv(t.Context(), mockEnv, failDiscovery(t)) assert.Equal(t, "http://localhost:9999", c.baseURL) @@ -676,6 +679,7 @@ func TestNewDefaultClient(t *testing.T) { ctrl := gomock.NewController(t) mockEnv := envmocks.NewMockReader(ctrl) mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() discover := func(context.Context) (string, []Option) { return "http://127.0.0.1:54321", nil @@ -689,6 +693,7 @@ func TestNewDefaultClient(t *testing.T) { ctrl := gomock.NewController(t) mockEnv := envmocks.NewMockReader(ctrl) mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("").AnyTimes() c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery, WithTimeout(5*time.Second)) assert.Equal(t, 5*time.Second, c.httpClient.Timeout) @@ -740,3 +745,205 @@ type failReader struct{} func (*failReader) Read([]byte) (int, error) { return 0, errors.New("simulated read error") } + +// TestInstallCarriesAllowUnsigned round-trips the unsigned exception through +// the client's request body — without this, the CLI flag silently never +// reaches the server (every --allow-unsigned install would 403 telling the +// user to pass the flag they passed). +func TestInstallCarriesAllowUnsigned(t *testing.T) { + t.Parallel() + + var got installRequest + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.NoError(t, json.NewDecoder(r.Body).Decode(&got)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + _ = json.NewEncoder(w).Encode(installResponse{}) + })) + t.Cleanup(srv.Close) + + _, err := newTestClient(t, srv).Install(t.Context(), skills.InstallOptions{ + Name: "my-skill", + Scope: skills.ScopeProject, + ProjectRoot: "/tmp/project", + AllowUnsigned: true, + }) + require.NoError(t, err) + assert.True(t, got.AllowUnsigned, "allow_unsigned must reach the server") +} + +// TestSyncCarriesAllowUnsigned mirrors TestInstallCarriesAllowUnsigned for +// the sync/adopt path. +func TestSyncCarriesAllowUnsigned(t *testing.T) { + t.Parallel() + + var got syncRequest + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.NoError(t, json.NewDecoder(r.Body).Decode(&got)) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(skills.SyncResult{}) + })) + t.Cleanup(srv.Close) + + _, err := newTestClient(t, srv).Sync(t.Context(), skills.SyncOptions{ + ProjectRoot: "/tmp/project", + Adopt: true, + AllowUnsigned: true, + }) + require.NoError(t, err) + assert.True(t, got.AllowUnsigned, "allow_unsigned must reach the server") +} + +func TestTimeoutFromEnv(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + value string + want time.Duration + ok bool + }{ + {name: "unset falls back to the default", value: "", want: 0, ok: false}, + {name: "a duration is honored", value: "45s", want: 45 * time.Second, ok: true}, + {name: "minutes parse", value: "5m", want: 5 * time.Minute, ok: true}, + {name: "surrounding whitespace is tolerated", value: " 90s ", want: 90 * time.Second, ok: true}, + {name: "a bare number is not a duration", value: "60", want: 0, ok: false}, + {name: "garbage degrades to the default", value: "soon", want: 0, ok: false}, + {name: "zero would disable the timeout, so it is ignored", value: "0s", want: 0, ok: false}, + {name: "negative is ignored", value: "-5s", want: 0, ok: false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPITimeout).Return(tc.value) + + got, ok := timeoutFromEnv(mockEnv) + assert.Equal(t, tc.ok, ok) + assert.Equal(t, tc.want, got) + }) + } +} + +func TestNewDefaultClientTimeoutPrecedence(t *testing.T) { + t.Parallel() + + noDiscovery := func(context.Context) (string, []Option) { return "", nil } + + t.Run("defaults when the env is unset", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("") + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery) + assert.Equal(t, defaultTimeout, c.httpClient.Timeout) + }) + + t.Run("env overrides the default", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("2m") + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery) + assert.Equal(t, 2*time.Minute, c.httpClient.Timeout) + }) + + t.Run("an explicit option outranks the env", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("2m") + + c := newDefaultClientWithEnv(t.Context(), mockEnv, noDiscovery, WithTimeout(7*time.Second)) + assert.Equal(t, 7*time.Second, c.httpClient.Timeout) + }) + + t.Run("env overrides the timeout discovery installs", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mockEnv := envmocks.NewMockReader(ctrl) + mockEnv.EXPECT().Getenv(envAPIURL).Return("") + mockEnv.EXPECT().Getenv(envAPITimeout).Return("3m") + + discover := func(context.Context) (string, []Option) { + return "http://127.0.0.1:54321", []Option{WithHTTPClient(&http.Client{Timeout: defaultTimeout})} + } + c := newDefaultClientWithEnv(t.Context(), mockEnv, discover) + assert.Equal(t, 3*time.Minute, c.httpClient.Timeout, + "WithHTTPClient from discovery must not shadow an operator-set timeout") + }) +} + +// TestTimeoutIsReportedAsTimeoutNotUnreachable pins the distinction the CLI +// hint depends on: a healthy-but-slow server must not be reported as absent. +func TestTimeoutIsReportedAsTimeoutNotUnreachable(t *testing.T) { + t.Parallel() + + blocked := make(chan struct{}) + t.Cleanup(func() { close(blocked) }) + + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + select { + case <-blocked: + case <-r.Context().Done(): + case <-time.After(30 * time.Second): + } + })) + t.Cleanup(srv.Close) + + c := NewClient(srv.URL, WithTimeout(50*time.Millisecond)) + _, err := c.List(t.Context(), skills.ListOptions{}) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrRequestTimeout) + assert.NotErrorIs(t, err, ErrServerUnreachable, + "the server answered the connection; only the response was slow") +} + +func TestUnreachableServerIsStillUnreachable(t *testing.T) { + t.Parallel() + + // Bind and immediately release a port so nothing is listening on it. + l, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := l.Addr().String() + require.NoError(t, l.Close()) + + c := NewClient("http://"+addr, WithTimeout(2*time.Second)) + _, err = c.List(t.Context(), skills.ListOptions{}) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrServerUnreachable) + assert.NotErrorIs(t, err, ErrRequestTimeout) +} + +// TestCallerCancellationIsNeitherSentinel keeps a user pressing Ctrl-C from +// being reported as a server problem. +func TestCallerCancellationIsNeitherSentinel(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + t.Cleanup(srv.Close) + + ctx, cancel := context.WithCancel(t.Context()) + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + c := NewClient(srv.URL, WithTimeout(30*time.Second)) + _, err := c.List(ctx, skills.ListOptions{}) + + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) + assert.NotErrorIs(t, err, ErrRequestTimeout) + assert.NotErrorIs(t, err, ErrServerUnreachable) +} diff --git a/pkg/skills/client/dto.go b/pkg/skills/client/dto.go index 104f7e2fad..c92dad6417 100644 --- a/pkg/skills/client/dto.go +++ b/pkg/skills/client/dto.go @@ -15,6 +15,9 @@ type installRequest struct { Clients []string `json:"clients,omitempty"` Force bool `json:"force,omitempty"` Group string `json:"group,omitempty"` + // AllowUnsigned mirrors skills.InstallOptions.AllowUnsigned; without it + // here the CLI flag would silently never reach the server. + AllowUnsigned bool `json:"allow_unsigned,omitempty"` } type validateRequest struct { @@ -41,3 +44,24 @@ type installResponse struct { type listBuildsResponse struct { Builds []skills.LocalBuild `json:"builds"` } + +type syncRequest struct { + ProjectRoot string `json:"project_root"` + Clients []string `json:"clients,omitempty"` + Prune bool `json:"prune,omitempty"` + Check bool `json:"check,omitempty"` + Adopt bool `json:"adopt,omitempty"` + // AllowUnsigned mirrors skills.SyncOptions.AllowUnsigned for adoption. + AllowUnsigned bool `json:"allow_unsigned,omitempty"` +} + +type upgradeRequest struct { + ProjectRoot string `json:"project_root"` + Names []string `json:"names,omitempty"` + // AllowSignerChange mirrors skills.UpgradeOptions.AllowSignerChange. + AllowSignerChange bool `json:"allow_signer_change,omitempty"` + Preview bool `json:"preview,omitempty"` + FailOnChanges bool `json:"fail_on_changes,omitempty"` + AllowRefChange bool `json:"allow_ref_change,omitempty"` + Clients []string `json:"clients,omitempty"` +} diff --git a/pkg/skills/feature_gate.go b/pkg/skills/feature_gate.go new file mode 100644 index 0000000000..8efff9d00a --- /dev/null +++ b/pkg/skills/feature_gate.go @@ -0,0 +1,27 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skills + +import ( + "os" + "strings" +) + +// LockFileEnvVar gates the project-level skills lock file feature (RFC +// THV-0080) while it lands across multiple PRs. The feature is inert on +// main until every PR in the stack — lock file, sync, upgrade, and Sigstore +// signing/verification — has merged; this keeps each PR mergeable on its own +// without exposing partial, unsigned-by-default behavior in between. +// +// This is intentionally a plain env var, not persisted config or a CLI flag: +// it exists only for the duration of the rollout and is expected to be +// removed once the feature ships, matching the existing TOOLHIVE_DEV / +// TOOLHIVE_REMOTE_HEALTHCHECKS precedent for staged/dev-only behavior. +const LockFileEnvVar = "TOOLHIVE_SKILLS_LOCK_ENABLED" + +// LockFileFeatureEnabled reports whether the project-level skills lock file +// feature is enabled for this process. +func LockFileFeatureEnabled() bool { + return strings.EqualFold(os.Getenv(LockFileEnvVar), "true") +} diff --git a/pkg/skills/feature_gate_test.go b/pkg/skills/feature_gate_test.go new file mode 100644 index 0000000000..cac0d8e4d9 --- /dev/null +++ b/pkg/skills/feature_gate_test.go @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skills + +import "testing" + +func TestLockFileFeatureEnabled(t *testing.T) { + tests := []struct { + name string + value string + want bool + }{ + {name: "unset defaults to disabled", value: "", want: false}, + {name: "true enables", value: "true", want: true}, + {name: "mixed case true enables", value: "True", want: true}, + {name: "false stays disabled", value: "false", want: false}, + {name: "arbitrary value stays disabled", value: "1", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if tt.value == "" { + t.Setenv(LockFileEnvVar, "") + } else { + t.Setenv(LockFileEnvVar, tt.value) + } + if got := LockFileFeatureEnabled(); got != tt.want { + t.Errorf("LockFileFeatureEnabled() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/pkg/skills/gitresolver/resolver.go b/pkg/skills/gitresolver/resolver.go index aaacd13b33..7cfdcb412d 100644 --- a/pkg/skills/gitresolver/resolver.go +++ b/pkg/skills/gitresolver/resolver.go @@ -45,6 +45,14 @@ type ResolveResult struct { Files []FileEntry // CommitHash is the git commit hash (for digest/upgrade detection) CommitHash string + // CommitSignature is the armored signature attached to the resolved + // commit, empty when the commit is unsigned. It is UNVERIFIED here — + // the install-time verifier checks it cryptographically before it is + // trusted or recorded as provenance. + CommitSignature string + // CommitPayload is the encoded commit object without its signature — + // the bytes CommitSignature signs. + CommitPayload []byte } // FileEntry represents a single file from the cloned repository. @@ -136,10 +144,11 @@ func (r *defaultResolver) Resolve(ctx context.Context, ref *GitReference) (*Reso } defer client.Cleanup(ctx, repoInfo) //nolint:errcheck // best-effort cleanup - // Get commit hash for digest tracking - commitHash, err := client.HeadCommitHash(repoInfo) + // Get the commit hash (for digest tracking) and its signature in one + // lookup, so both describe the same commit. + head, err := client.HeadCommit(repoInfo) if err != nil { - return nil, fmt.Errorf("getting commit hash: %w", err) + return nil, fmt.Errorf("getting HEAD commit: %w", err) } // Read SKILL.md from the skill path @@ -173,9 +182,11 @@ func (r *defaultResolver) Resolve(ctx context.Context, ref *GitReference) (*Reso } return &ResolveResult{ - SkillConfig: parsed, - Files: files, - CommitHash: commitHash, + SkillConfig: parsed, + Files: files, + CommitHash: head.Hash, + CommitSignature: head.Signature, + CommitPayload: head.Payload, }, nil } diff --git a/pkg/skills/gitresolver/resolver_test.go b/pkg/skills/gitresolver/resolver_test.go index 22771a1464..a5018f4b40 100644 --- a/pkg/skills/gitresolver/resolver_test.go +++ b/pkg/skills/gitresolver/resolver_test.go @@ -328,6 +328,38 @@ func TestResolver_Resolve_CommitRef(t *testing.T) { require.NotNil(t, result) assert.Equal(t, "my-skill", result.SkillConfig.Name) assert.Equal(t, commitHash, result.CommitHash) + assert.Empty(t, result.CommitSignature, "an unsigned commit must resolve with no signature") +} + +func TestResolver_Resolve_SignedCommitSignaturePropagates(t *testing.T) { + t.Parallel() + + const armoredSig = "-----BEGIN SIGNED MESSAGE-----\nMIIC(test-signature)\n-----END SIGNED MESSAGE-----\n" + repoDir := createTestRepo(t, "", validSkillMD) + + // Rewrite HEAD with the signature attached, mimicking gitsign output. + repo, err := gogit.PlainOpen(repoDir) + require.NoError(t, err) + head, err := repo.Head() + require.NoError(t, err) + commit, err := repo.CommitObject(head.Hash()) + require.NoError(t, err) + commit.PGPSignature = armoredSig + obj := repo.Storer.NewEncodedObject() + require.NoError(t, commit.Encode(obj)) + signedHash, err := repo.Storer.SetEncodedObject(obj) + require.NoError(t, err) + require.NoError(t, repo.Storer.SetReference(plumbing.NewHashReference(head.Name(), signedHash))) + + resolver := NewResolver(WithGitClient(git.NewDefaultGitClient())) + result, err := resolver.Resolve(t.Context(), &GitReference{URL: repoDir}) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, signedHash.String(), result.CommitHash) + assert.Equal(t, armoredSig, result.CommitSignature) + assert.Contains(t, string(result.CommitPayload), "tree ", + "the signed payload must accompany the signature") + assert.NotContains(t, string(result.CommitPayload), "gpgsig") } func TestResolver_Resolve_MissingSkillMD(t *testing.T) { @@ -393,8 +425,8 @@ func (*blockingCloneClient) GetFileContent(_ *git.RepositoryInfo, _ string) ([]b return nil, fmt.Errorf("not implemented") } -func (*blockingCloneClient) HeadCommitHash(_ *git.RepositoryInfo) (string, error) { - return "", fmt.Errorf("not implemented") +func (*blockingCloneClient) HeadCommit(_ *git.RepositoryInfo) (git.HeadCommit, error) { + return git.HeadCommit{}, fmt.Errorf("not implemented") } func (*blockingCloneClient) Cleanup(_ context.Context, _ *git.RepositoryInfo) error { diff --git a/pkg/skills/lock_service.go b/pkg/skills/lock_service.go new file mode 100644 index 0000000000..e0c37acbb6 --- /dev/null +++ b/pkg/skills/lock_service.go @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skills + +import "context" + +//go:generate mockgen -destination=mocks/mock_lock_service.go -package=mocks -source=lock_service.go SkillLockService + +// SkillLockService defines the interface for operations driven by a +// project's toolhive.lock.yaml. It is separate from [SkillService] because +// Sync and Upgrade operate over the whole lock file rather than a single +// named skill, per RFC THV-0080. +type SkillLockService interface { + // Sync restores the project's installed skills to match its lock file. + Sync(ctx context.Context, opts SyncOptions) (*SyncResult, error) + // Upgrade re-resolves each lock entry's source and installs newer content + // when the resolved digest has changed. + Upgrade(ctx context.Context, opts UpgradeOptions) (*UpgradeResult, error) +} diff --git a/pkg/skills/lockfile/contentdigest.go b/pkg/skills/lockfile/contentdigest.go new file mode 100644 index 0000000000..58a00d8af8 --- /dev/null +++ b/pkg/skills/lockfile/contentdigest.go @@ -0,0 +1,118 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package lockfile + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "slices" + "strings" +) + +// ContentFile is a single file used for contentDigest computation. +type ContentFile struct { + // Path is the relative path within the skill directory, using either + // native or slash separators. + Path string + // Content is the raw file bytes. + Content []byte +} + +// ContentDigest computes a deterministic SHA-256 dirhash over files. +// +// Algorithm (content-only; file modes and timestamps are ignored): +// 1. Normalize each path to slash-separated form and sort files by path. +// 2. For each file, feed path + "\x00" + sha256(content) + "\n" into a +// running SHA-256. +// 3. Return "sha256:" + hex(aggregate). +// +// The algorithm is frozen: lock files in the wild pin its output, so any +// change would report every installed skill as drifted. Golden vectors in +// the package tests guard against accidental changes. +// +// Known limitations (accepted, like go.sum's): the digest covers file +// content only — file modes (e.g. an exec-bit flip) are invisible to it, +// and [ContentDigestFromDir] skips symlinks and other non-regular files +// entirely. Trust in the file set's provenance is the Sigstore +// verification layer's job, not this integrity pin's. +func ContentDigest(files []ContentFile) (string, error) { + if len(files) == 0 { + return "", fmt.Errorf("content digest requires at least one file") + } + + normalized := make([]ContentFile, 0, len(files)) + for _, f := range files { + path := filepath.ToSlash(f.Path) + for strings.HasPrefix(path, "./") { + path = strings.TrimPrefix(path, "./") + } + if path == "" || path == ".." || strings.HasPrefix(path, "../") || strings.Contains(path, "/../") || + strings.HasSuffix(path, "/..") || strings.HasPrefix(path, "/") || hasControlChar(path) { + return "", fmt.Errorf("invalid content file path %q", f.Path) + } + normalized = append(normalized, ContentFile{Path: path, Content: f.Content}) + } + slices.SortFunc(normalized, func(a, b ContentFile) int { return strings.Compare(a.Path, b.Path) }) + + h := sha256.New() + for _, f := range normalized { + fileHash := sha256.Sum256(f.Content) + // Writes to a hash.Hash never fail. + _, _ = io.WriteString(h, f.Path) + _, _ = h.Write([]byte{0}) + _, _ = h.Write(fileHash[:]) + _, _ = h.Write([]byte{'\n'}) + } + return ContentDigestPrefix + hex.EncodeToString(h.Sum(nil)), nil +} + +// hasControlChar reports whether s contains a NUL, newline, or other C0/DEL +// control byte. The serialization format used by [ContentDigest] delimits +// fields with NUL and newline bytes; without this check, a Path containing +// those bytes could reproduce the exact byte stream of a different, +// unrelated multi-file tree and collide with its digest. +func hasControlChar(s string) bool { + for i := 0; i < len(s); i++ { + if s[i] < 0x20 || s[i] == 0x7f { + return true + } + } + return false +} + +// ContentDigestFromDir walks skillDir and computes the content digest from +// the on-disk file set. The walk is confined to skillDir via [os.Root], so +// symlinks cannot escape it. +func ContentDigestFromDir(skillDir string) (string, error) { + root, err := os.OpenRoot(skillDir) + if err != nil { + return "", fmt.Errorf("opening skill directory: %w", err) + } + defer func() { _ = root.Close() }() + + var files []ContentFile + err = fs.WalkDir(root.FS(), ".", func(path string, d fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if d.IsDir() || !d.Type().IsRegular() { + return nil + } + data, err := fs.ReadFile(root.FS(), path) + if err != nil { + return fmt.Errorf("reading %q: %w", path, err) + } + files = append(files, ContentFile{Path: path, Content: data}) + return nil + }) + if err != nil { + return "", err + } + return ContentDigest(files) +} diff --git a/pkg/skills/lockfile/contentdigest_test.go b/pkg/skills/lockfile/contentdigest_test.go new file mode 100644 index 0000000000..a8a4d36861 --- /dev/null +++ b/pkg/skills/lockfile/contentdigest_test.go @@ -0,0 +1,188 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package lockfile + +import ( + "crypto/sha256" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Golden vectors freeze the contentDigest algorithm. If these ever need to +// change, every lock file in the wild reports drift — treat a change here as +// a breaking change to the lock file format, not a routine test update. +func TestContentDigestGoldenVectors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + files []ContentFile + want string + }{ + { + name: "single file", + files: []ContentFile{ + {Path: "a.txt", Content: []byte("hello world")}, + }, + want: "sha256:79d1b89c9733bc5a3c51bd715612a85ab13add6933cb9aa7cc1ae44fd5d181a3", + }, + { + name: "multiple files sorted by path", + files: []ContentFile{ + {Path: "SKILL.md", Content: []byte("# Skill")}, + {Path: "refs/guide.md", Content: []byte("guide")}, + }, + want: "sha256:9b46342aaebede490f12bf67a5ecb4be507c7684ee5e11a0511b54c701f464bd", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := ContentDigest(tt.files) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestContentDigestDeterministicAndOrderIndependent(t *testing.T) { + t.Parallel() + + a := []ContentFile{ + {Path: "SKILL.md", Content: []byte("# Skill")}, + {Path: "refs/guide.md", Content: []byte("guide")}, + } + // Same files, different input order: the digest must not depend on + // caller-supplied ordering. + b := []ContentFile{ + {Path: "refs/guide.md", Content: []byte("guide")}, + {Path: "SKILL.md", Content: []byte("# Skill")}, + } + + d1, err := ContentDigest(a) + require.NoError(t, err) + d2, err := ContentDigest(b) + require.NoError(t, err) + assert.Equal(t, d1, d2) +} + +func TestContentDigestDiffersOnContentChange(t *testing.T) { + t.Parallel() + + d1, err := ContentDigest([]ContentFile{{Path: "a.txt", Content: []byte("v1")}}) + require.NoError(t, err) + d2, err := ContentDigest([]ContentFile{{Path: "a.txt", Content: []byte("v2")}}) + require.NoError(t, err) + assert.NotEqual(t, d1, d2) +} + +func TestContentDigestRejectsEmptyFileSet(t *testing.T) { + t.Parallel() + _, err := ContentDigest(nil) + require.Error(t, err) +} + +func TestContentDigestRejectsTraversalPaths(t *testing.T) { + t.Parallel() + + tests := []string{"..", "../escape.txt", "a/../../escape.txt", "/abs/path.txt"} + for _, p := range tests { + t.Run(p, func(t *testing.T) { + t.Parallel() + _, err := ContentDigest([]ContentFile{{Path: p, Content: []byte("x")}}) + require.Error(t, err) + }) + } +} + +func TestContentDigestNormalizesDotSlashPrefix(t *testing.T) { + t.Parallel() + + withPrefix, err := ContentDigest([]ContentFile{{Path: "./a.txt", Content: []byte("hello")}}) + require.NoError(t, err) + withoutPrefix, err := ContentDigest([]ContentFile{{Path: "a.txt", Content: []byte("hello")}}) + require.NoError(t, err) + assert.Equal(t, withoutPrefix, withPrefix) +} + +func TestContentDigestRejectsControlCharsInPath(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + path string + }{ + {"embedded NUL", "a\x00b.txt"}, + {"embedded newline", "a\nb.txt"}, + {"embedded DEL", "a\x7fb.txt"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := ContentDigest([]ContentFile{{Path: tt.path, Content: []byte("x")}}) + require.Error(t, err) + }) + } +} + +// TestContentDigestRejectsForgedPathCollision demonstrates the attack the +// control-char rejection closes: the serialization format delimits fields +// with NUL and newline bytes, so without rejecting those bytes in Path, a +// single crafted entry could reproduce the exact byte stream of two +// legitimate, unrelated entries and collide with their digest. +func TestContentDigestRejectsForgedPathCollision(t *testing.T) { + t.Parallel() + + legit := []ContentFile{ + {Path: "x", Content: []byte("hello world")}, + {Path: "y", Content: []byte("second file content")}, + } + _, err := ContentDigest(legit) + require.NoError(t, err) + + xHash := sha256.Sum256([]byte("hello world")) + forgedPath := "x\x00" + string(xHash[:]) + "\ny" + _, err = ContentDigest([]ContentFile{{Path: forgedPath, Content: []byte("second file content")}}) + require.Error(t, err, "a path forging another entry's serialized bytes must be rejected, not hashed to a collision") +} + +func TestContentDigestNormalizesRepeatedDotSlashPrefix(t *testing.T) { + t.Parallel() + + repeated, err := ContentDigest([]ContentFile{{Path: "././a.txt", Content: []byte("hello")}}) + require.NoError(t, err) + single, err := ContentDigest([]ContentFile{{Path: "a.txt", Content: []byte("hello")}}) + require.NoError(t, err) + assert.Equal(t, single, repeated) +} + +func TestContentDigestFromDir(t *testing.T) { + t.Parallel() + dir, err := filepath.EvalSymlinks(t.TempDir()) + require.NoError(t, err) + require.NoError(t, os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte("# Skill"), 0o644)) + require.NoError(t, os.Mkdir(filepath.Join(dir, "refs"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "refs", "guide.md"), []byte("guide"), 0o644)) + + got, err := ContentDigestFromDir(dir) + require.NoError(t, err) + + want, err := ContentDigest([]ContentFile{ + {Path: "SKILL.md", Content: []byte("# Skill")}, + {Path: "refs/guide.md", Content: []byte("guide")}, + }) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +func TestContentDigestFromDirRejectsMissingDir(t *testing.T) { + t.Parallel() + _, err := ContentDigestFromDir(filepath.Join(t.TempDir(), "does-not-exist")) + require.Error(t, err) +} diff --git a/pkg/skills/lockfile/lockfile.go b/pkg/skills/lockfile/lockfile.go new file mode 100644 index 0000000000..f1032ec678 --- /dev/null +++ b/pkg/skills/lockfile/lockfile.go @@ -0,0 +1,364 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package lockfile manages the project-level skills lock file +// (toolhive.lock.yaml). The lock file pins the exact name, version, source, +// and digests of every project-scoped skill install so a team can restore +// ("thv skill sync") or refresh ("thv skill upgrade") the pinned state on any +// machine. See RFC THV-0080. +// +// All filesystem access goes through [Root], a capability type that can only +// be constructed from a validated project root. Root confines every read and +// write to that directory using [os.Root] (OS-enforced containment, not just +// string validation), so no function in this package can be made to open a +// path outside the project root regardless of what name is requested. +package lockfile + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" + "slices" + "strings" + + "gopkg.in/yaml.v3" + + "github.com/stacklok/toolhive/pkg/fileutils" + "github.com/stacklok/toolhive/pkg/skills" +) + +// FileName is the name of the project-level skills lock file, written to the +// project root alongside .git. +const FileName = "toolhive.lock.yaml" + +// CurrentVersion is the schema version written to new lock files. Loading a +// lock file with a different version is a hard error, never a silent partial +// parse. +const CurrentVersion = 1 + +// Entry represents a single pinned skill installation in the lock file. +type Entry struct { + // Name is the skill's unique name. + Name string `yaml:"name"` + // Version is the skill's declared version, if any (from SKILL.md frontmatter). + Version string `yaml:"version,omitempty"` + // Source is exactly what the user (or the registry resolver) originally + // requested — a plain registry name, an OCI reference, or a git:// + // reference. It is never rewritten; upgrade re-resolves this value to + // check for newer content. + Source string `yaml:"source"` + // ResolvedReference is the concrete OCI reference or git:// URL that + // Source resolved to at install time. + ResolvedReference string `yaml:"resolvedReference,omitempty"` + // Digest pins the exact content installed: an OCI "sha256:..." manifest + // digest or a full git commit hash. + Digest string `yaml:"digest"` + // ContentDigest is a deterministic SHA-256 dirhash of the materialized + // skill file set, used for on-disk integrity verification. + ContentDigest string `yaml:"contentDigest,omitempty"` + // Provenance records the verified signer identity of the installed + // artifact — the trust-on-first-use anchor later installs, syncs, and + // upgrades are checked against. Nil for entries recorded before + // verification existed and for unsigned entries (see Unsigned); an + // entry never carries both. + Provenance *Provenance `yaml:"provenance,omitempty"` + // Unsigned is true when the artifact carried no signature and the user + // explicitly accepted that with --allow-unsigned. It is a recorded + // trust decision, not an omission: replacing an unsigned entry with a + // signed one clears it, and a signed entry never becomes unsigned + // without the same explicit flag. + Unsigned bool `yaml:"unsigned,omitempty"` + // RequiredBy lists parent skill names for transitively materialized + // dependencies (skills declared via toolhive.requires). + RequiredBy []string `yaml:"requiredBy,omitempty"` + // Explicit is true when the user directly installed this skill; explicit + // entries are exempt from cascade removal when RequiredBy becomes empty. + Explicit bool `yaml:"explicit,omitempty"` + // Extra round-trips fields this binary does not know about, so a + // Load→modify→Save cycle by an older binary preserves rather than strips + // them. It applies per entry: an entry this binary rewrites (reinstall, + // upgrade) is built fresh, so its unknown fields — which described the + // previous install — are intentionally dropped along with it. + Extra map[string]any `yaml:",inline"` +} + +// Provenance is the Sigstore signer identity recorded for a verified entry. +type Provenance struct { + // SignerIdentity is the certificate subject identity: for GitHub + // Actions certificates, the workflow path relative to the repository; + // otherwise the certificate SAN verbatim (a URI, email, or SPIFFE ID). + SignerIdentity string `yaml:"signerIdentity"` + // CertIssuer is the OIDC issuer that authenticated the signer. + CertIssuer string `yaml:"certIssuer"` + // RepositoryURI is the source repository from the Fulcio certificate + // extensions, when present. + RepositoryURI string `yaml:"repositoryUri,omitempty"` + // SigstoreURL is the Sigstore instance the signature chains to. + SigstoreURL string `yaml:"sigstoreUrl,omitempty"` + // Provisional marks provenance whose verification has a documented + // gap — currently git-commit signatures, verified for signature and + // certificate chain but without transparency-log proof of signing + // time. Visible in lock diffs so reduced assurance is never silent; + // removed when the gap closes. + Provisional bool `yaml:"provisional,omitempty"` +} + +// Lockfile is the parsed contents of a project's toolhive.lock.yaml. +type Lockfile struct { + // Version is the lock file schema version. + Version int `yaml:"version"` + // Skills is the set of pinned skill installations, sorted by name for + // stable diffs. + Skills []Entry `yaml:"skills,omitempty"` + // Extra round-trips unknown top-level fields, mirroring Entry.Extra. + Extra map[string]any `yaml:",inline"` +} + +// Root is a validated project root directory. It is the only way to address +// a lock file on disk: constructing one runs full project-root validation +// (absolute, NUL-free, no traversal segments, symlink-canonical, git-rooted). +// Every read and write additionally goes through a freshly opened [os.Root] +// scoped to the validated directory, so access is confined at the OS level. +type Root struct { + dir string +} + +// OpenRoot validates projectRoot and returns a Root for it. The zero Root is +// unusable; all lock file operations require a Root produced here. +func OpenRoot(projectRoot string) (Root, error) { + dir, err := skills.ValidateProjectRoot(projectRoot) + if err != nil { + return Root{}, err + } + return Root{dir: dir}, nil +} + +// Dir returns the validated project root directory. +func (r Root) Dir() string { + return r.dir +} + +// Path returns the absolute path of the lock file inside the project root, +// for display purposes (e.g. error messages, CLI output). It is not used by +// this package to open the file; use [Load] and [Save] for that. +func (r Root) Path() (string, error) { + if r.dir == "" { + return "", errors.New("lockfile: uninitialized Root, use OpenRoot") + } + return filepath.Join(r.dir, FileName), nil +} + +// osRoot opens an OS-level containment handle for r. Every name passed to +// its methods is resolved beneath r.dir; a name that would escape it (via +// "..", an absolute path, or a symlink) is rejected by the OS, not by string +// inspection. +func (r Root) osRoot() (*os.Root, error) { + if r.dir == "" { + return nil, errors.New("lockfile: uninitialized Root, use OpenRoot") + } + return os.OpenRoot(r.dir) +} + +// Load reads and parses the lock file for root. A missing lock file is not +// an error — it returns an empty Lockfile ready to be populated. A lock file +// with an unsupported schema version is a hard error. +func Load(root Root) (*Lockfile, error) { + osRoot, err := root.osRoot() + if err != nil { + return nil, err + } + defer func() { _ = osRoot.Close() }() + + data, err := osRoot.ReadFile(FileName) + if errors.Is(err, fs.ErrNotExist) { + return &Lockfile{Version: CurrentVersion}, nil + } + if err != nil { + return nil, fmt.Errorf("reading lock file: %w", err) + } + + var lf Lockfile + if err := yaml.Unmarshal(data, &lf); err != nil { + return nil, fmt.Errorf("parsing lock file: %w", err) + } + if lf.Version == 0 { + // Tolerate a hand-written lock file that omits the version key. + lf.Version = CurrentVersion + } + sortEntries(lf.Skills) + if err := validateLockfile(&lf); err != nil { + return nil, err + } + return &lf, nil +} + +// Get returns the entry for name, if present. +func (l *Lockfile) Get(name string) (Entry, bool) { + for _, e := range l.Skills { + if e.Name == name { + return e, true + } + } + return Entry{}, false +} + +// Upsert inserts or replaces the entry with a matching name, keeping the +// slice sorted by name for stable diffs. +func (l *Lockfile) Upsert(entry Entry) { + for i := range l.Skills { + if l.Skills[i].Name == entry.Name { + l.Skills[i] = entry + return + } + } + l.Skills = append(l.Skills, entry) + sortEntries(l.Skills) +} + +// Remove deletes the entry with the given name, if present. Reports whether +// an entry was removed. +func (l *Lockfile) Remove(name string) bool { + for i, e := range l.Skills { + if e.Name == name { + l.Skills = slices.Delete(l.Skills, i, i+1) + return true + } + } + return false +} + +// RemoveParentFromRequiredBy removes parent from every entry's RequiredBy +// list and returns the names of entries that lost their last parent and are +// not explicit — the candidates for cascade removal. +func (l *Lockfile) RemoveParentFromRequiredBy(parent string) []string { + var cascadeCandidates []string + for i := range l.Skills { + entry := &l.Skills[i] + if !slices.Contains(entry.RequiredBy, parent) { + continue + } + entry.RequiredBy = slices.DeleteFunc(entry.RequiredBy, func(s string) bool { return s == parent }) + if len(entry.RequiredBy) == 0 { + entry.RequiredBy = nil + if !entry.Explicit { + cascadeCandidates = append(cascadeCandidates, entry.Name) + } + } + } + return cascadeCandidates +} + +// tmpFileName is the temporary file Save writes before renaming it into +// tmpFileName returns a per-call temporary file name for Save's atomic +// write. The random suffix means two concurrent direct Save calls (which, +// unlike Update/UpsertEntry/RemoveEntry, hold no file lock) cannot write +// to, rename, or delete each other's temp file — last rename still wins, +// but neither can promote or destroy a half-written file. +func tmpFileName() (string, error) { + var buf [8]byte + if _, err := rand.Read(buf[:]); err != nil { + return "", fmt.Errorf("generating temp file suffix: %w", err) + } + return ".toolhive.lock." + hex.EncodeToString(buf[:]) + ".tmp", nil +} + +// Save writes the lock file into root atomically (temp file + rename), after +// validating it with the same rules [Load] enforces on read. This prevents a +// caller bug from writing a lock file that every subsequent Load/Update call +// would then hard-fail on, with no recovery path through this package. +// Callers that need read-modify-write atomicity across processes must use +// [UpsertEntry], [RemoveEntry], or [Update] instead of Load+Save — Save alone +// is write-atomic but does nothing to serialize concurrent read-modify-write +// cycles. +func (l *Lockfile) Save(root Root) error { + osRoot, err := root.osRoot() + if err != nil { + return err + } + defer func() { _ = osRoot.Close() }() + + if l.Version == 0 { + l.Version = CurrentVersion + } + sortEntries(l.Skills) + + if err := validateLockfile(l); err != nil { + return fmt.Errorf("refusing to save invalid lock file: %w", err) + } + + data, err := yaml.Marshal(l) + if err != nil { + return fmt.Errorf("marshaling lock file: %w", err) + } + + tmpName, err := tmpFileName() + if err != nil { + return err + } + // The lock file is committed to git and not sensitive; 0o644 matches any + // other source file. + if err := osRoot.WriteFile(tmpName, data, 0o644); err != nil { + return fmt.Errorf("writing lock file: %w", err) + } + if err := osRoot.Rename(tmpName, FileName); err != nil { + _ = osRoot.Remove(tmpName) + return fmt.Errorf("saving lock file: %w", err) + } + return nil +} + +// UpsertEntry loads the lock file, upserts entry, and saves it back, all +// under a single file lock so concurrent installs cannot race on +// read-modify-write. +func UpsertEntry(root Root, entry Entry) error { + return Update(root, func(lf *Lockfile) error { + lf.Upsert(entry) + return nil + }) +} + +// RemoveEntry loads the lock file, removes the named entry if present, and +// saves it back, all under a single file lock. Removing an entry that does +// not exist is a no-op, not an error. +func RemoveEntry(root Root, name string) error { + return Update(root, func(lf *Lockfile) error { + if !lf.Remove(name) { + return errSkipSave + } + return nil + }) +} + +// Update loads the lock file, applies fn, and saves the result, all under a +// single file lock. If fn returns an error the lock file is left unchanged. +func Update(root Root, fn func(*Lockfile) error) error { + path, err := root.Path() + if err != nil { + return err + } + return fileutils.WithFileLock(path, func() error { + lf, err := Load(root) + if err != nil { + return err + } + if err := fn(lf); err != nil { + if errors.Is(err, errSkipSave) { + return nil + } + return err + } + return lf.Save(root) + }) +} + +// errSkipSave signals from an Update callback that nothing changed and the +// save should be skipped without reporting an error. +var errSkipSave = errors.New("lockfile: no changes to save") + +func sortEntries(entries []Entry) { + slices.SortFunc(entries, func(a, b Entry) int { return strings.Compare(a.Name, b.Name) }) +} diff --git a/pkg/skills/lockfile/lockfile_test.go b/pkg/skills/lockfile/lockfile_test.go new file mode 100644 index 0000000000..0f04c779f5 --- /dev/null +++ b/pkg/skills/lockfile/lockfile_test.go @@ -0,0 +1,476 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package lockfile + +import ( + "errors" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// hexDigest returns a valid lowercase hex string of exactly n characters. A +// distinct seed produces a distinct string, so tests needing multiple unique +// fixture digests can pass different seeds. +func hexDigest(n int, seed byte) string { + const alphabet = "0123456789abcdef" + b := make([]byte, n) + for i := range b { + b[i] = alphabet[(i+int(seed))%len(alphabet)] + } + return string(b) +} + +// ociDigest returns a valid fixture value for an Entry.Digest/ContentDigest +// field ("sha256:" + 64 hex characters). +func ociDigest(seed byte) string { + return ContentDigestPrefix + hexDigest(sha256HexLength, seed) +} + +// testRoot creates a resolved temp dir with a .git directory (satisfying +// skills.ValidateProjectRoot) and returns its opened Root. +func testRoot(t *testing.T) Root { + t.Helper() + dir, err := filepath.EvalSymlinks(t.TempDir()) + require.NoError(t, err) + require.NoError(t, os.Mkdir(filepath.Join(dir, ".git"), 0o755)) + root, err := OpenRoot(dir) + require.NoError(t, err) + return root +} + +func TestOpenRootRejectsInvalidProjectRoot(t *testing.T) { + t.Parallel() + + _, err := OpenRoot(filepath.Join(t.TempDir(), "missing")) + require.Error(t, err) +} + +func TestZeroRootPathFails(t *testing.T) { + t.Parallel() + + var r Root + _, err := r.Path() + require.Error(t, err) + assert.Contains(t, err.Error(), "uninitialized Root") +} + +func TestLoadMissingFileReturnsEmptyLockfile(t *testing.T) { + t.Parallel() + root := testRoot(t) + + lf, err := Load(root) + require.NoError(t, err) + assert.Equal(t, CurrentVersion, lf.Version) + assert.Empty(t, lf.Skills) +} + +func TestLoadRejectsMalformedYAML(t *testing.T) { + t.Parallel() + root := testRoot(t) + path, err := root.Path() + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, []byte("not: [valid: yaml"), 0o644)) + + _, err = Load(root) + require.Error(t, err) + assert.Contains(t, err.Error(), "parsing lock file") +} + +func TestLoadRejectsUnsupportedVersion(t *testing.T) { + t.Parallel() + root := testRoot(t) + path, err := root.Path() + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, []byte("version: 99\nskills: []\n"), 0o644)) + + _, err = Load(root) + require.Error(t, err) + assert.ErrorIs(t, err, ErrUnsupportedVersion) +} + +func TestSaveAndLoadRoundTrip(t *testing.T) { + t.Parallel() + root := testRoot(t) + + lf := &Lockfile{Version: CurrentVersion} + entry := Entry{ + Name: "code-review", + Version: "1.0.0", + Source: "code-review", + ResolvedReference: "ghcr.io/org/code-review:1.0.0", + Digest: ociDigest(1), + ContentDigest: ociDigest(2), + Explicit: true, + } + lf.Upsert(entry) + require.NoError(t, lf.Save(root)) + + loaded, err := Load(root) + require.NoError(t, err) + require.Len(t, loaded.Skills, 1) + assert.Equal(t, entry, loaded.Skills[0]) +} + +// TestUnknownFieldsSurviveLoadModifySave freezes the version-skew contract: +// a lock file written by a newer binary (e.g. with Sigstore provenance +// fields added under schema version 1) must round-trip through this +// binary's Load -> modify -> Save without the unknown fields being +// stripped — otherwise an older thv silently removes a teammate's signing +// metadata on its next write. +func TestUnknownFieldsSurviveLoadModifySave(t *testing.T) { + t.Parallel() + root := testRoot(t) + path, err := root.Path() + require.NoError(t, err) + + futureLock := "" + + "version: 1\n" + + "futureTopLevelField: keep-me\n" + + "skills:\n" + + " - name: signed-skill\n" + + " source: ghcr.io/org/signed-skill:1.0.0\n" + + " digest: " + ociDigest(1) + "\n" + + " attestations:\n" + + " - predicateType: https://slsa.dev/provenance/v1\n" + + " digest: " + ociDigest(3) + "\n" + require.NoError(t, os.WriteFile(path, []byte(futureLock), 0o644)) + + // Load, touch an unrelated entry, save — the classic older-binary write. + require.NoError(t, UpsertEntry(root, Entry{ + Name: "other-skill", + Source: "other-skill", + Digest: ociDigest(2), + })) + + data, err := os.ReadFile(path) //nolint:gosec // fixed test path + require.NoError(t, err) + saved := string(data) + assert.Contains(t, saved, "predicateType: https://slsa.dev/provenance/v1", "unknown entry fields must survive a Load->Save cycle") + assert.Contains(t, saved, "futureTopLevelField: keep-me", "unknown top-level fields must survive a Load->Save cycle") + + loaded, err := Load(root) + require.NoError(t, err) + signed, ok := loaded.Get("signed-skill") + require.True(t, ok) + assert.Contains(t, signed.Extra, "attestations") +} + +func TestSaveRejectsInvalidLockfile(t *testing.T) { + t.Parallel() + root := testRoot(t) + + lf := &Lockfile{Version: CurrentVersion} + lf.Upsert(Entry{Name: "code-review", Source: "code-review"}) // missing required Digest + + err := lf.Save(root) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid lock file") + + _, statErr := os.Stat(filepath.Join(root.Dir(), FileName)) + require.True(t, os.IsNotExist(statErr), "Save must not write an invalid lock file to disk") +} + +func TestUpdateRejectsInvalidLockfile(t *testing.T) { + t.Parallel() + root := testRoot(t) + + err := Update(root, func(lf *Lockfile) error { + lf.Upsert(Entry{Name: "code-review", Source: "code-review"}) // missing required Digest + return nil + }) + require.Error(t, err) + + // A caller bug that upserts an invalid entry must not leave every + // subsequent Load/Update permanently broken. + loaded, loadErr := Load(root) + require.NoError(t, loadErr) + assert.Empty(t, loaded.Skills) +} + +func TestLockfileGetUpsertRemove(t *testing.T) { + t.Parallel() + + lf := &Lockfile{Version: CurrentVersion} + + _, ok := lf.Get("missing") + assert.False(t, ok) + + a := Entry{Name: "b-skill", Source: "b-skill", Digest: ociDigest(3)} + c := Entry{Name: "a-skill", Source: "a-skill", Digest: ociDigest(4)} + lf.Upsert(a) + lf.Upsert(c) + + // Entries are kept sorted by name. + require.Len(t, lf.Skills, 2) + assert.Equal(t, "a-skill", lf.Skills[0].Name) + assert.Equal(t, "b-skill", lf.Skills[1].Name) + + got, ok := lf.Get("b-skill") + require.True(t, ok) + assert.Equal(t, a.Digest, got.Digest) + + // Upsert replaces an existing entry rather than duplicating it. + updated := Entry{Name: "b-skill", Source: "b-skill", Digest: ociDigest(5)} + lf.Upsert(updated) + require.Len(t, lf.Skills, 2) + got, ok = lf.Get("b-skill") + require.True(t, ok) + assert.Equal(t, updated.Digest, got.Digest) + + removed := lf.Remove("b-skill") + assert.True(t, removed) + require.Len(t, lf.Skills, 1) + + removedAgain := lf.Remove("b-skill") + assert.False(t, removedAgain) +} + +func TestRemoveParentFromRequiredBy(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + entries []Entry + parent string + wantCascadeCandidates []string + wantRemainingParents map[string][]string + }{ + { + name: "last parent removed and not explicit becomes cascade candidate", + entries: []Entry{ + {Name: "dep", RequiredBy: []string{"parent"}}, + }, + parent: "parent", + wantCascadeCandidates: []string{"dep"}, + wantRemainingParents: map[string][]string{"dep": nil}, + }, + { + name: "explicit entry is never a cascade candidate", + entries: []Entry{ + {Name: "dep", RequiredBy: []string{"parent"}, Explicit: true}, + }, + parent: "parent", + wantCascadeCandidates: nil, + wantRemainingParents: map[string][]string{"dep": nil}, + }, + { + name: "dep with another parent survives", + entries: []Entry{ + {Name: "dep", RequiredBy: []string{"parent", "other-parent"}}, + }, + parent: "parent", + wantCascadeCandidates: nil, + wantRemainingParents: map[string][]string{"dep": {"other-parent"}}, + }, + { + name: "entry without the parent is untouched", + entries: []Entry{ + {Name: "unrelated", RequiredBy: []string{"someone-else"}}, + }, + parent: "parent", + wantCascadeCandidates: nil, + wantRemainingParents: map[string][]string{"unrelated": {"someone-else"}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + lf := &Lockfile{Version: CurrentVersion, Skills: tt.entries} + got := lf.RemoveParentFromRequiredBy(tt.parent) + assert.Equal(t, tt.wantCascadeCandidates, got) + for name, want := range tt.wantRemainingParents { + entry, ok := lf.Get(name) + require.True(t, ok) + assert.Equal(t, want, entry.RequiredBy) + } + }) + } +} + +func TestUpsertEntryAndRemoveEntry(t *testing.T) { + t.Parallel() + root := testRoot(t) + + entry := Entry{Name: "my-skill", Source: "my-skill", Digest: ociDigest(3)} + require.NoError(t, UpsertEntry(root, entry)) + + lf, err := Load(root) + require.NoError(t, err) + require.Len(t, lf.Skills, 1) + assert.Equal(t, "my-skill", lf.Skills[0].Name) + + // Upserting a second entry preserves the first. + other := Entry{Name: "other-skill", Source: "other-skill", Digest: ociDigest(6)} + require.NoError(t, UpsertEntry(root, other)) + lf, err = Load(root) + require.NoError(t, err) + assert.Len(t, lf.Skills, 2) + + require.NoError(t, RemoveEntry(root, "my-skill")) + lf, err = Load(root) + require.NoError(t, err) + require.Len(t, lf.Skills, 1) + assert.Equal(t, "other-skill", lf.Skills[0].Name) + + // Removing a name that isn't present is a no-op, not an error. + require.NoError(t, RemoveEntry(root, "does-not-exist")) +} + +func TestUpdateLeavesLockfileUnchangedOnError(t *testing.T) { + t.Parallel() + root := testRoot(t) + require.NoError(t, UpsertEntry(root, Entry{Name: "my-skill", Source: "my-skill", Digest: ociDigest(3)})) + + wantErr := errors.New("boom") + err := Update(root, func(lf *Lockfile) error { + lf.Upsert(Entry{Name: "should-not-persist", Source: "x", Digest: ociDigest(4)}) + return wantErr + }) + require.ErrorIs(t, err, wantErr) + + lf, err := Load(root) + require.NoError(t, err) + require.Len(t, lf.Skills, 1) + assert.Equal(t, "my-skill", lf.Skills[0].Name) +} + +// TestConcurrentUpsertEntryDoesNotLoseUpdates hammers UpsertEntry from many +// goroutines to confirm the read-modify-write cycle is fully serialized by +// the file lock, not just protected at the Save call. +func TestConcurrentUpsertEntryDoesNotLoseUpdates(t *testing.T) { + t.Parallel() + root := testRoot(t) + + const n = 20 + var wg sync.WaitGroup + errCh := make(chan error, n) + for i := range n { + wg.Add(1) + go func(i int) { + defer wg.Done() + name := skillNameForIndex(i) + errCh <- UpsertEntry(root, Entry{ + Name: name, + Source: name, + Digest: ociDigest(3), + }) + }(i) + } + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent upserts to finish") + } + close(errCh) + for err := range errCh { + require.NoError(t, err) + } + + lf, err := Load(root) + require.NoError(t, err) + require.Len(t, lf.Skills, n, "a concurrent upsert lost an update") +} + +func skillNameForIndex(i int) string { + return "skill-" + string(rune('a'+i/26)) + string(rune('a'+i%26)) +} + +// TestProvenanceRoundTrip covers the Stack-2 schema: a verified entry's +// provenance block and an unsigned entry's exception flag both survive +// Save -> Load intact under schema version 1. +func TestProvenanceRoundTrip(t *testing.T) { + t.Parallel() + root := testRoot(t) + + lf := &Lockfile{Version: CurrentVersion} + signed := Entry{ + Name: "signed-skill", + Source: "ghcr.io/org/signed-skill", + ResolvedReference: "ghcr.io/org/signed-skill:latest", + Digest: ociDigest(1), + Provenance: &Provenance{ + SignerIdentity: "/.github/workflows/release.yml", + CertIssuer: "https://token.actions.githubusercontent.com", + RepositoryURI: "https://github.com/org/signed-skill", + SigstoreURL: "https://rekor.sigstore.dev", + }, + Explicit: true, + } + unsigned := Entry{ + Name: "unsigned-skill", + Source: "unsigned-skill", + Digest: ociDigest(2), + Unsigned: true, + Explicit: true, + } + lf.Upsert(signed) + lf.Upsert(unsigned) + require.NoError(t, lf.Save(root)) + + loaded, err := Load(root) + require.NoError(t, err) + + gotSigned, ok := loaded.Get("signed-skill") + require.True(t, ok) + assert.Equal(t, signed.Provenance, gotSigned.Provenance) + assert.False(t, gotSigned.Unsigned) + assert.Empty(t, gotSigned.Extra, "a typed provenance block must not leak into the inline Extra map") + + gotUnsigned, ok := loaded.Get("unsigned-skill") + require.True(t, ok) + assert.True(t, gotUnsigned.Unsigned) + assert.Nil(t, gotUnsigned.Provenance) +} + +// TestProvenanceGraduatesFromExtraMap: before this schema change, a +// provenance block written by a newer binary round-tripped through the +// Extra inline map. Now that the field is typed, loading such a file must +// parse it into Entry.Provenance — not duplicate it in Extra, which would +// make Save produce a colliding key. +func TestProvenanceGraduatesFromExtraMap(t *testing.T) { + t.Parallel() + root := testRoot(t) + path, err := root.Path() + require.NoError(t, err) + + handWritten := "" + + "version: 1\n" + + "skills:\n" + + " - name: signed-skill\n" + + " source: ghcr.io/org/signed-skill\n" + + " digest: " + ociDigest(1) + "\n" + + " provenance:\n" + + " signerIdentity: dev@example.com\n" + + " certIssuer: https://accounts.example.com\n" + + " explicit: true\n" + require.NoError(t, os.WriteFile(path, []byte(handWritten), 0o644)) + + loaded, err := Load(root) + require.NoError(t, err) + entry, ok := loaded.Get("signed-skill") + require.True(t, ok) + require.NotNil(t, entry.Provenance, "provenance must parse into the typed field") + assert.Equal(t, "dev@example.com", entry.Provenance.SignerIdentity) + assert.NotContains(t, entry.Extra, "provenance", "the typed field must not also appear in Extra") + + // A Save after modification must not produce duplicate keys. + require.NoError(t, UpsertEntry(root, Entry{ + Name: "other-skill", Source: "other-skill", Digest: ociDigest(2), + })) + reloaded, err := Load(root) + require.NoError(t, err) + entry, ok = reloaded.Get("signed-skill") + require.True(t, ok) + assert.Equal(t, "dev@example.com", entry.Provenance.SignerIdentity) +} diff --git a/pkg/skills/lockfile/validation.go b/pkg/skills/lockfile/validation.go new file mode 100644 index 0000000000..a66e3f46ec --- /dev/null +++ b/pkg/skills/lockfile/validation.go @@ -0,0 +1,266 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package lockfile + +import ( + "errors" + "fmt" + "strings" + "unicode" + + nameref "github.com/google/go-containerregistry/pkg/name" + + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/gitresolver" +) + +// ErrUnsupportedVersion indicates the lock file schema version is not +// supported by this build. +var ErrUnsupportedVersion = errors.New("unsupported lock file version") + +const ( + // ContentDigestPrefix is the prefix for contentDigest values in the lock file. + ContentDigestPrefix = "sha256:" + + sha256HexLength = 64 + sha1HexLength = 40 +) + +// validateLockfile checks the schema version, every entry, and +// cross-references requiredBy links against the set of known entries. +func validateLockfile(lf *Lockfile) error { + if lf.Version != CurrentVersion { + return fmt.Errorf("%w: version %d (upgrade thv to read this lock file)", ErrUnsupportedVersion, lf.Version) + } + + names := make(map[string]struct{}, len(lf.Skills)) + for _, entry := range lf.Skills { + if err := validateEntry(entry); err != nil { + return err + } + if _, dup := names[entry.Name]; dup { + return fmt.Errorf("duplicate entry %q", entry.Name) + } + names[entry.Name] = struct{}{} + } + + for _, entry := range lf.Skills { + for _, parent := range entry.RequiredBy { + if parent == entry.Name { + return fmt.Errorf("entry %q: requiredBy references itself", entry.Name) + } + if _, ok := names[parent]; !ok { + return fmt.Errorf("entry %q: requiredBy references unknown parent %q", entry.Name, parent) + } + } + } + + if cycle := findRequiredByCycle(lf.Skills); len(cycle) > 0 { + return fmt.Errorf("requiredBy cycle: %s", strings.Join(cycle, " -> ")) + } + return nil +} + +// findRequiredByCycle detects a cycle in the requiredBy graph and returns +// one such cycle path, or nil. Normal installs can never produce a cycle +// (the install-time Visited set breaks them), but a hand-edited or badly +// merge-resolved lock file can — and a ring of mutually-required, +// non-explicit entries would then be impossible to ever cascade-remove, so +// it is rejected at validation instead of persisting silently. +func findRequiredByCycle(entries []Entry) []string { + requiredBy := make(map[string][]string, len(entries)) + for _, e := range entries { + requiredBy[e.Name] = e.RequiredBy + } + + const ( + unvisited = 0 + inStack = 1 + done = 2 + ) + state := make(map[string]int, len(entries)) + + var visit func(name string, path []string) []string + visit = func(name string, path []string) []string { + state[name] = inStack + path = append(path, name) + for _, parent := range requiredBy[name] { + switch state[parent] { + case inStack: + // Trim the path to start at the cycle entry point. + for i, n := range path { + if n == parent { + return append(path[i:], parent) + } + } + case unvisited: + if cycle := visit(parent, path); cycle != nil { + return cycle + } + } + } + state[name] = done + return nil + } + + for _, e := range entries { + if state[e.Name] == unvisited { + if cycle := visit(e.Name, nil); cycle != nil { + return cycle + } + } + } + return nil +} + +// maxReferenceLength bounds ResolvedReference: longer than any legitimate +// OCI reference or git URL, short enough to stop megabyte-scale garbage +// from a corrupted or hostile lock file reaching the fetch path. +const maxReferenceLength = 512 + +func validateEntry(entry Entry) error { + if err := skills.ValidateSkillName(entry.Name); err != nil { + return fmt.Errorf("entry name: %w", err) + } + if entry.Source == "" { + return fmt.Errorf("entry %q: source is required", entry.Name) + } + if entry.Digest == "" { + return fmt.Errorf("entry %q: digest is required", entry.Name) + } + if err := validateDigest(entry.Digest); err != nil { + return fmt.Errorf("entry %q: digest: %w", entry.Name, err) + } + if entry.ContentDigest != "" { + if err := validateContentDigest(entry.ContentDigest); err != nil { + return fmt.Errorf("entry %q: contentDigest: %w", entry.Name, err) + } + } + if entry.ResolvedReference != "" { + if err := validateResolvedReference(entry.ResolvedReference); err != nil { + return fmt.Errorf("entry %q: resolvedReference: %w", entry.Name, err) + } + } + if entry.Provenance != nil && entry.Unsigned { + return fmt.Errorf("entry %q: provenance and unsigned are mutually exclusive"+ + " — an entry is either a verified signature or a recorded unsigned exception", entry.Name) + } + if entry.Provenance != nil { + if err := validateProvenance(entry.Provenance); err != nil { + return fmt.Errorf("entry %q: provenance: %w", entry.Name, err) + } + } + return nil +} + +// validateProvenance syntactically constrains a provenance block. Like +// resolvedReference, these values are hand-editable and feed the identity +// policy that future verifications are checked against, so they must be +// well-formed graphic strings of bounded length. Validation is purely +// syntactic — whether the identity is trustworthy is the verifier's job. +func validateProvenance(p *Provenance) error { + if p.SignerIdentity == "" { + return errors.New("signerIdentity is required") + } + if p.CertIssuer == "" { + return errors.New("certIssuer is required") + } + fields := map[string]string{ + "signerIdentity": p.SignerIdentity, + "certIssuer": p.CertIssuer, + "repositoryUri": p.RepositoryURI, + "sigstoreUrl": p.SigstoreURL, + } + for name, value := range fields { + if value == "" { + continue + } + if len(value) > maxReferenceLength { + return fmt.Errorf("%s exceeds %d characters", name, maxReferenceLength) + } + if strings.TrimSpace(value) != value { + return fmt.Errorf("%s has leading or trailing whitespace", name) + } + for _, r := range value { + if !unicode.IsGraphic(r) || unicode.IsSpace(r) { + return fmt.Errorf("%s contains non-graphic or whitespace character %q", name, r) + } + } + } + return nil +} + +// validateResolvedReference syntactically constrains the resolvedReference +// field. Sync fetches from this value without re-resolving Source, and the +// lock file is hand-editable, so a value that is not a plausible git:// or +// OCI reference must never reach the fetch path. Validation is purely +// syntactic — no network access, no allow-list policy. +func validateResolvedReference(ref string) error { + if len(ref) > maxReferenceLength { + return fmt.Errorf("exceeds %d characters", maxReferenceLength) + } + if strings.TrimSpace(ref) != ref { + return errors.New("has leading or trailing whitespace") + } + for _, r := range ref { + if !unicode.IsGraphic(r) || unicode.IsSpace(r) { + return fmt.Errorf("contains non-graphic or whitespace character %q", r) + } + } + if gitresolver.IsGitReference(ref) { + if _, err := gitresolver.ParseGitReference(ref); err != nil { + return fmt.Errorf("invalid git reference: %w", err) + } + return nil + } + // StrictValidation requires an explicit tag or digest, matching what the + // install path records (qualifiedOCIRef always includes one) — and, + // unlike weak validation, rejects URL-shaped strings such as + // "http://169.254.169.254/…" that must never reach the fetch path. + if _, err := nameref.ParseReference(ref, nameref.StrictValidation); err != nil { + return fmt.Errorf("not a valid git:// or OCI reference: %w", err) + } + return nil +} + +// validateDigest accepts the two pin formats the lock file records: an OCI +// manifest digest ("sha256:" + 64 hex chars) or a full git commit hash +// (40 hex chars for SHA-1 repositories, 64 for SHA-256 repositories). +// Abbreviated digests are rejected: a truncated pin weakens the guarantee +// the lock file exists to provide. +func validateDigest(d string) error { + if hexPart, ok := strings.CutPrefix(d, ContentDigestPrefix); ok { + if err := validateHex(hexPart, sha256HexLength); err != nil { + return fmt.Errorf("OCI digest: %w", err) + } + return nil + } + if len(d) != sha1HexLength && len(d) != sha256HexLength { + return fmt.Errorf("expected %q + 64 hex chars or a full git commit hash, got %q", ContentDigestPrefix, d) + } + if err := validateHex(d, len(d)); err != nil { + return fmt.Errorf("git commit hash: %w", err) + } + return nil +} + +func validateContentDigest(d string) error { + hexPart, ok := strings.CutPrefix(d, ContentDigestPrefix) + if !ok { + return fmt.Errorf("must start with %q", ContentDigestPrefix) + } + return validateHex(hexPart, sha256HexLength) +} + +func validateHex(s string, wantLen int) error { + if len(s) != wantLen { + return fmt.Errorf("expected %d hex characters, got %d", wantLen, len(s)) + } + for _, c := range s { + if (c < '0' || c > '9') && (c < 'a' || c > 'f') { + return fmt.Errorf("invalid hex character %q", c) + } + } + return nil +} diff --git a/pkg/skills/lockfile/validation_test.go b/pkg/skills/lockfile/validation_test.go new file mode 100644 index 0000000000..a54430d0db --- /dev/null +++ b/pkg/skills/lockfile/validation_test.go @@ -0,0 +1,273 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package lockfile + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var ( + validSHA256Hex = hexDigest(sha256HexLength, 0) + validSHA256Digest = "sha256:" + validSHA256Hex + validGitSHA1 = hexDigest(sha1HexLength, 0) + validContentDigest = ContentDigestPrefix + validSHA256Hex +) + +func TestValidateDigest(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + digest string + wantErr string + }{ + {name: "valid OCI digest", digest: validSHA256Digest}, + {name: "valid git SHA-1 commit", digest: validGitSHA1}, + {name: "valid git SHA-256 commit", digest: validSHA256Hex}, + {name: "empty", digest: "", wantErr: "expected"}, + {name: "abbreviated git hash rejected", digest: "0123456", wantErr: "expected"}, + {name: "oci digest wrong length", digest: "sha256:abc", wantErr: "OCI digest"}, + {name: "oci digest bad hex", digest: "sha256:" + strings.Repeat("z", 64), wantErr: "OCI digest"}, + {name: "git hash bad hex", digest: strings.Repeat("z", 40), wantErr: "git commit hash"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := validateDigest(tt.digest) + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestValidateContentDigest(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + digest string + wantErr string + }{ + {name: "valid", digest: validContentDigest}, + {name: "missing prefix", digest: validSHA256Hex, wantErr: "must start with"}, + {name: "wrong length", digest: ContentDigestPrefix + "abc", wantErr: "expected 64 hex"}, + {name: "bad hex", digest: ContentDigestPrefix + strings.Repeat("z", 64), wantErr: "invalid hex"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := validateContentDigest(tt.digest) + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestValidateResolvedReference(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + ref string + wantErr string + }{ + {name: "valid OCI reference with tag", ref: "ghcr.io/org/skill:1.0.0"}, + {name: "valid OCI reference with digest", ref: "ghcr.io/org/skill@sha256:" + validSHA256Hex}, + {name: "valid git reference", ref: "git://github.com/org/repo@main#skills/my-skill"}, + {name: "too long", ref: "ghcr.io/org/" + strings.Repeat("a", maxReferenceLength), wantErr: "exceeds"}, + {name: "leading whitespace", ref: " ghcr.io/org/skill:1", wantErr: "whitespace"}, + {name: "embedded newline", ref: "ghcr.io/org/\nskill:1", wantErr: "non-graphic"}, + {name: "embedded ANSI escape", ref: "ghcr.io/org/\x1b[31mskill:1", wantErr: "non-graphic"}, + {name: "malformed git reference", ref: "git://", wantErr: "invalid git reference"}, + {name: "not a reference at all", ref: "http://169.254.169.254/latest/meta-data", wantErr: "not a valid"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := validateResolvedReference(tt.ref) + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestValidateLockfile(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + lf Lockfile + wantErr string + }{ + { + name: "valid single entry", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "my-skill", Source: "my-skill", Digest: validSHA256Digest}, + }}, + }, + { + name: "unsupported version", + lf: Lockfile{Version: 99}, + wantErr: "unsupported lock file version", + }, + { + name: "duplicate entry names", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "dup", Source: "a", Digest: validSHA256Digest}, + {Name: "dup", Source: "b", Digest: validSHA256Digest}, + }}, + wantErr: "duplicate entry", + }, + { + name: "requiredBy references unknown parent", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "dep", Source: "dep", Digest: validSHA256Digest, RequiredBy: []string{"ghost"}}, + }}, + wantErr: "unknown parent", + }, + { + name: "requiredBy references itself", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "dep", Source: "dep", Digest: validSHA256Digest, RequiredBy: []string{"dep"}}, + }}, + wantErr: "references itself", + }, + { + name: "missing source", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "my-skill", Digest: validSHA256Digest}, + }}, + wantErr: "source is required", + }, + { + name: "missing digest", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "my-skill", Source: "my-skill"}, + }}, + wantErr: "digest is required", + }, + { + name: "invalid skill name", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "Not_Valid", Source: "x", Digest: validSHA256Digest}, + }}, + wantErr: "entry name", + }, + { + // A mutual requiredBy ring of non-explicit entries would pass the + // per-edge checks yet be impossible to ever cascade-remove. + name: "requiredBy cycle", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "ring-a", Source: "a", Digest: validSHA256Digest, RequiredBy: []string{"ring-b"}}, + {Name: "ring-b", Source: "b", Digest: validSHA256Digest, RequiredBy: []string{"ring-a"}}, + }}, + wantErr: "requiredBy cycle", + }, + { + name: "valid provenance", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "signed", Source: "s", Digest: validSHA256Digest, Provenance: &Provenance{ + SignerIdentity: "dev@example.com", + CertIssuer: "https://accounts.example.com", + }}, + }}, + }, + { + name: "provenance and unsigned are mutually exclusive", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "conflicted", Source: "s", Digest: validSHA256Digest, Unsigned: true, Provenance: &Provenance{ + SignerIdentity: "dev@example.com", + CertIssuer: "https://accounts.example.com", + }}, + }}, + wantErr: "mutually exclusive", + }, + { + name: "provenance missing signer identity", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "signed", Source: "s", Digest: validSHA256Digest, Provenance: &Provenance{ + CertIssuer: "https://accounts.example.com", + }}, + }}, + wantErr: "signerIdentity is required", + }, + { + name: "provenance missing cert issuer", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "signed", Source: "s", Digest: validSHA256Digest, Provenance: &Provenance{ + SignerIdentity: "dev@example.com", + }}, + }}, + wantErr: "certIssuer is required", + }, + { + name: "provenance with control characters rejected", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "signed", Source: "s", Digest: validSHA256Digest, Provenance: &Provenance{ + SignerIdentity: "dev@example.com\x1b[31m", + CertIssuer: "https://accounts.example.com", + }}, + }}, + wantErr: "non-graphic", + }, + { + name: "provenance field too long rejected", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "signed", Source: "s", Digest: validSHA256Digest, Provenance: &Provenance{ + SignerIdentity: strings.Repeat("a", maxReferenceLength+1), + CertIssuer: "https://accounts.example.com", + }}, + }}, + wantErr: "exceeds", + }, + { + name: "unsigned exception alone is valid", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "unsigned", Source: "s", Digest: validSHA256Digest, Unsigned: true}, + }}, + }, + { + name: "requiredBy diamond is not a cycle", + lf: Lockfile{Version: CurrentVersion, Skills: []Entry{ + {Name: "shared-dep", Source: "d", Digest: validSHA256Digest, RequiredBy: []string{"parent-a", "parent-b"}}, + {Name: "parent-a", Source: "a", Digest: validSHA256Digest, RequiredBy: []string{"root"}}, + {Name: "parent-b", Source: "b", Digest: validSHA256Digest, RequiredBy: []string{"root"}}, + {Name: "root", Source: "r", Digest: validSHA256Digest, Explicit: true}, + }}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + lf := tt.lf + err := validateLockfile(&lf) + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} diff --git a/pkg/skills/mocks/mock_lock_service.go b/pkg/skills/mocks/mock_lock_service.go new file mode 100644 index 0000000000..e1c63d4a80 --- /dev/null +++ b/pkg/skills/mocks/mock_lock_service.go @@ -0,0 +1,72 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: lock_service.go +// +// Generated by this command: +// +// mockgen -destination=mocks/mock_lock_service.go -package=mocks -source=lock_service.go SkillLockService +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + reflect "reflect" + + skills "github.com/stacklok/toolhive/pkg/skills" + gomock "go.uber.org/mock/gomock" +) + +// MockSkillLockService is a mock of SkillLockService interface. +type MockSkillLockService struct { + ctrl *gomock.Controller + recorder *MockSkillLockServiceMockRecorder + isgomock struct{} +} + +// MockSkillLockServiceMockRecorder is the mock recorder for MockSkillLockService. +type MockSkillLockServiceMockRecorder struct { + mock *MockSkillLockService +} + +// NewMockSkillLockService creates a new mock instance. +func NewMockSkillLockService(ctrl *gomock.Controller) *MockSkillLockService { + mock := &MockSkillLockService{ctrl: ctrl} + mock.recorder = &MockSkillLockServiceMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSkillLockService) EXPECT() *MockSkillLockServiceMockRecorder { + return m.recorder +} + +// Sync mocks base method. +func (m *MockSkillLockService) Sync(ctx context.Context, opts skills.SyncOptions) (*skills.SyncResult, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Sync", ctx, opts) + ret0, _ := ret[0].(*skills.SyncResult) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Sync indicates an expected call of Sync. +func (mr *MockSkillLockServiceMockRecorder) Sync(ctx, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Sync", reflect.TypeOf((*MockSkillLockService)(nil).Sync), ctx, opts) +} + +// Upgrade mocks base method. +func (m *MockSkillLockService) Upgrade(ctx context.Context, opts skills.UpgradeOptions) (*skills.UpgradeResult, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Upgrade", ctx, opts) + ret0, _ := ret[0].(*skills.UpgradeResult) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Upgrade indicates an expected call of Upgrade. +func (mr *MockSkillLockServiceMockRecorder) Upgrade(ctx, opts any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Upgrade", reflect.TypeOf((*MockSkillLockService)(nil).Upgrade), ctx, opts) +} diff --git a/pkg/skills/options.go b/pkg/skills/options.go index a822c8918f..6975fcbec3 100644 --- a/pkg/skills/options.go +++ b/pkg/skills/options.go @@ -31,18 +31,96 @@ type InstallOptions struct { ProjectRoot string `json:"project_root,omitempty"` // Group is the group name to add the skill to after installation. Group string `json:"group,omitempty"` + // AllowUnsigned permits installing a project-scoped skill whose artifact + // carries no Sigstore signature. Without it, unsigned artifacts are + // rejected; with it, the lock entry records the exception as + // "unsigned: true". Skill content is AI-executed instructions, so this + // is an explicit per-install trust decision, never a default. + AllowUnsigned bool `json:"allow_unsigned,omitempty"` // LayerData is the tar.gz content from an OCI layer. Internal use only — NOT exposed via HTTP API. LayerData []byte `json:"-"` // Reference is the full OCI reference (e.g. ghcr.io/org/skill:v1). Reference string `json:"-"` // Digest is the OCI digest for upgrade detection. Digest string `json:"-"` + // LockSource overrides the value recorded as the lock entry's Source. When + // empty, the entry's Source is Name as given by the caller before any + // internal resolution. Set by Sync/Upgrade, which pass an already-resolved + // Name that must not overwrite the entry's original Source. Internal use + // only — NOT exposed via HTTP API. + LockSource string `json:"-"` + // LockResolvedReference overrides the value recorded as the lock entry's + // ResolvedReference. When empty, the entry's ResolvedReference is + // whatever this install actually resolved to. Set by Sync when + // reinstalling at a pinned reference: without this override, a drift + // repair would overwrite ResolvedReference with the internal pinned + // form (e.g. a digest or commit hash spliced into the reference) + // instead of preserving what Source originally resolved to. Internal + // use only — NOT exposed via HTTP API. + LockResolvedReference string `json:"-"` + // RequiredByParent is set when this install is a transitively materialized + // dependency (toolhive.requires) of another skill, naming that parent. + // Empty means the user explicitly requested this install. Internal use + // only — NOT exposed via HTTP API. + RequiredByParent string `json:"-"` + // Visited tracks skill names already materialized in this dependency + // tree, preventing infinite recursion on a requires cycle. Left nil by + // external callers; Install initializes it on first entry and threads it + // through recursive dependency installs. Internal use only — NOT exposed + // via HTTP API. + Visited map[string]struct{} `json:"-"` + // SyncRestore forces re-extraction to every existing client even when + // Digest matches the currently-installed digest. Set by Sync when + // reinstalling at a pinned reference: the whole point is repairing + // on-disk drift that happened without the pinned digest changing, so the + // normal "same digest means content is already correct" fast path must + // not apply. Internal use only — NOT exposed via HTTP API. + SyncRestore bool `json:"-"` + // AllowSignerChange lets install-time verification re-record the + // observed identity instead of enforcing the lock file's recorded one. + // Internal use only — set by upgrade when its signer-change guard was + // explicitly overridden. + AllowSignerChange bool `json:"-"` + // Unsigned records the trust decision that this install proceeded + // without a verified signature (via AllowUnsigned). Set internally by + // install-time verification; recorded as `unsigned: true` in the lock + // entry. + Unsigned bool `json:"-"` + // Provenance carries the verified signer identity established during + // install-time verification, for recording into the lock entry. Set by + // the verification step, nil when the artifact is unsigned or + // verification did not run. Internal use only — NOT exposed via HTTP API. + Provenance *ProvenanceInfo `json:"-"` + // SigstoreBundle is the serialized Sigstore bundle backing Provenance, + // persisted alongside the install record so sync can re-verify offline. + // Internal use only — NOT exposed via HTTP API. + SigstoreBundle []byte `json:"-"` +} + +// ProvenanceInfo is the verified signer identity of an installed artifact, +// the in-memory mirror of the lock file's provenance block. +type ProvenanceInfo struct { + // SignerIdentity is the certificate subject identity (workflow path for + // GitHub Actions certificates, SAN verbatim otherwise). + SignerIdentity string `json:"-"` + // CertIssuer is the OIDC issuer that authenticated the signer. + CertIssuer string `json:"-"` + // RepositoryURI is the source repository from the certificate + // extensions, when present. + RepositoryURI string `json:"-"` + // SigstoreURL is the Sigstore instance the signature chains to. + SigstoreURL string `json:"-"` } // InstallResult contains the outcome of an Install operation. type InstallResult struct { // Skill is the installed skill. Skill InstalledSkill `json:"skill"` + // PreExisting is the store record as it was before this install, or nil + // when this install created the record. Rollback uses it to restore the + // previous state instead of destructively deleting a record this call + // did not create. Internal use only — NOT exposed via HTTP API. + PreExisting *InstalledSkill `json:"-"` } // UninstallOptions configures the behavior of the Uninstall operation. @@ -53,6 +131,12 @@ type UninstallOptions struct { Scope Scope `json:"scope,omitempty"` // ProjectRoot is the project root path for project-scoped skills. ProjectRoot string `json:"project_root,omitempty"` + // Visited tracks skill names already removed in this cascade-uninstall + // tree, preventing infinite recursion on a requiredBy cycle. Left nil by + // external callers; Uninstall initializes it on first entry and threads + // it through recursive cascade removals. Internal use only — NOT exposed + // via HTTP API. + Visited map[string]struct{} `json:"-"` } // InfoOptions configures the behavior of the Info operation. @@ -133,6 +217,165 @@ type PushOptions struct { Reference string `json:"reference"` } +// SyncOptions configures the behavior of the Sync operation. +type SyncOptions struct { + // ProjectRoot is the project root path whose lock file should be synced. + ProjectRoot string `json:"project_root"` + // Clients lists target clients (e.g., "claude-code"). Empty means every + // skill-supporting client detected on this host. + Clients []string `json:"clients,omitempty"` + // Prune removes project-scoped skills that are installed but not present + // in the lock file. When false, such skills are only reported. + Prune bool `json:"prune,omitempty"` + // Check verifies on-disk content against contentDigest without installing + // or writing anything. + Check bool `json:"check,omitempty"` + // AllowUnsigned permits adopting a skill whose signature state cannot + // be established (no stored bundle), recording an explicit unsigned + // exception — the same trust decision as an unsigned install. + AllowUnsigned bool `json:"allow_unsigned,omitempty"` + // Adopt writes lock entries for existing unmanaged project-scope installs. + Adopt bool `json:"adopt,omitempty"` +} + +// FailureReason is a typed failure reason for sync/upgrade operations, per +// RFC THV-0080's exit-code and automation contract. Only reasons the current +// feature set can actually produce are defined; the RFC's remaining values +// (the Sigstore verification reasons and ref-change-blocked, which surfaces +// as an UpgradeStatus rather than a failure) land together with the code +// that emits them. +type FailureReason string + +// Typed failure reasons for sync/upgrade operations. +const ( + // FailureReasonRegistryUnreachable means the skill's remote source — an + // OCI registry or a git host — could not be reached. + FailureReasonRegistryUnreachable FailureReason = "registry-unreachable" + FailureReasonDigestMissing FailureReason = "digest-missing" + FailureReasonValidationRejected FailureReason = "validation-rejected" + FailureReasonLockWriteFailed FailureReason = "lock-write-failed" + // FailureReasonSignatureInvalid means the artifact carries signature + // material that failed cryptographic verification. + FailureReasonSignatureInvalid FailureReason = "signature-invalid" + // FailureReasonSignerMismatch means the artifact verifies, but against + // an identity other than the one recorded in the lock file. + FailureReasonSignerMismatch FailureReason = "signer-mismatch" + // FailureReasonUnsignedRejected means the artifact is unsigned and the + // operation did not permit unsigned installs. + FailureReasonUnsignedRejected FailureReason = "unsigned-rejected" + FailureReasonUnknown FailureReason = "unknown" +) + +// SyncFailure describes a single skill that failed to sync. +type SyncFailure struct { + // Name is the skill name that failed. + Name string `json:"name"` + // Reason is a typed failure reason for CI and automation. + Reason FailureReason `json:"reason,omitempty"` + // Error is a human-readable description of the failure. + Error string `json:"error"` +} + +// SyncResult contains the outcome of a Sync operation. +type SyncResult struct { + // Installed lists skills that were installed or reinstalled to match the lock file. + Installed []string `json:"installed,omitempty"` + // Drifted lists skills whose on-disk contentDigest differed from the lock + // file. Normally these are reinstalled to match it; when Check is set, + // nothing is written and this field reports the drift only. + Drifted []string `json:"drifted,omitempty"` + // Missing lists lock entries with no corresponding install record at all + // — the fresh-clone state. Normally these are installed at their pinned + // reference; when Check is set, nothing is written and this field + // reports the gap only. + Missing []string `json:"missing,omitempty"` + // AlreadyCurrent lists skills that already matched the lock file. + AlreadyCurrent []string `json:"already_current,omitempty"` + // NeverManaged lists project-scoped skills never recorded as lock-managed. + NeverManaged []string `json:"never_managed,omitempty"` + // RemovedFromLock lists previously managed skills absent from the lock file. + RemovedFromLock []string `json:"removed_from_lock,omitempty"` + // Pruned lists removed-from-lock skills that were uninstalled because Prune was set. + Pruned []string `json:"pruned,omitempty"` + // Failed lists skills that could not be synced, with the reason for each. + // Drift alone is never reported here — see Drifted. + Failed []SyncFailure `json:"failed,omitempty"` +} + +// UpgradeOptions configures the behavior of the Upgrade operation. +type UpgradeOptions struct { + // ProjectRoot is the project root path whose lock file should be upgraded. + ProjectRoot string `json:"project_root"` + // Names restricts the upgrade to specific skill names. Empty means every + // entry in the lock file. + Names []string `json:"names,omitempty"` + // Preview reports what would change without installing (still fetches + // artifacts to compare digests). + Preview bool `json:"preview,omitempty"` + // FailOnChanges exits with an error when any mutable source would upgrade. + FailOnChanges bool `json:"fail_on_changes,omitempty"` + // AllowRefChange permits an upgrade whose candidate lives in a different + // repository. Tag moves within the same repository are always allowed — + // they are how a mutable source advances. + AllowRefChange bool `json:"allow_ref_change,omitempty"` + // AllowSignerChange permits upgrading to an artifact signed by a + // different identity than the one recorded in the lock file; the new + // identity is recorded in its place. + AllowSignerChange bool `json:"allow_signer_change,omitempty"` + // Clients lists target clients (e.g., "claude-code"). Empty means every + // skill-supporting client detected on this host. + Clients []string `json:"clients,omitempty"` +} + +// UpgradeStatus represents the outcome of upgrading a single skill. +type UpgradeStatus string + +const ( + // UpgradeStatusUpgraded indicates the skill was installed at a new digest. + UpgradeStatusUpgraded UpgradeStatus = "upgraded" + // UpgradeStatusUpToDate indicates the resolved source still points at the pinned digest. + UpgradeStatusUpToDate UpgradeStatus = "up-to-date" + // UpgradeStatusNotUpgradable indicates the entry is pinned to an immutable + // reference (an OCI digest or a full git commit hash) and cannot be upgraded. + UpgradeStatusNotUpgradable UpgradeStatus = "not-upgradable" + // UpgradeStatusRefChangeBlocked indicates re-resolution changed resolvedReference. + UpgradeStatusRefChangeBlocked UpgradeStatus = "ref-change-blocked" + // UpgradeStatusSignerChangeBlocked indicates the candidate artifact is + // signed by a different identity (or unsigned) versus the identity the + // lock file records. + UpgradeStatusSignerChangeBlocked UpgradeStatus = "signer-change-blocked" + // UpgradeStatusFailed indicates the upgrade attempt failed. + UpgradeStatusFailed UpgradeStatus = "failed" +) + +// UpgradeOutcome describes the result of attempting to upgrade one skill. +type UpgradeOutcome struct { + // Name is the skill name. + Name string `json:"name"` + // Status is the outcome of the upgrade attempt. + Status UpgradeStatus `json:"status"` + // OldDigest is the digest pinned in the lock file before this operation. + OldDigest string `json:"old_digest,omitempty"` + // NewDigest is the digest the source currently resolves to. Equal to + // OldDigest when Status is UpgradeStatusUpToDate. + NewDigest string `json:"new_digest,omitempty"` + // NewResolvedReference is the new resolvedReference when it changed. + NewResolvedReference string `json:"new_resolved_reference,omitempty"` + // NewSignerIdentity is the candidate's signer identity when it differs + // from the recorded one (empty when the candidate is unsigned). + NewSignerIdentity string `json:"new_signer_identity,omitempty"` + // Reason is a typed failure reason when Status is UpgradeStatusFailed. + Reason FailureReason `json:"reason,omitempty"` + // Error is a human-readable description of the failure, set only when Status is UpgradeStatusFailed. + Error string `json:"error,omitempty"` +} + +// UpgradeResult contains the outcome of an Upgrade operation. +type UpgradeResult struct { + // Outcomes contains one entry per skill considered for upgrade. + Outcomes []UpgradeOutcome `json:"outcomes"` +} + // LocalBuild represents a locally-built OCI skill artifact in the local store. type LocalBuild struct { // Tag is the OCI tag or name used to reference the artifact. diff --git a/pkg/skills/signer/cosign_attach.go b/pkg/skills/signer/cosign_attach.go new file mode 100644 index 0000000000..880dc0c44b --- /dev/null +++ b/pkg/skills/signer/cosign_attach.go @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package signer + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "strings" + + "github.com/google/go-containerregistry/pkg/authn" + "github.com/google/go-containerregistry/pkg/name" + v1 "github.com/google/go-containerregistry/pkg/v1" + "github.com/google/go-containerregistry/pkg/v1/empty" + "github.com/google/go-containerregistry/pkg/v1/mutate" + "github.com/google/go-containerregistry/pkg/v1/remote" + "github.com/google/go-containerregistry/pkg/v1/static" + "github.com/google/go-containerregistry/pkg/v1/types" + "github.com/opencontainers/go-digest" +) + +const ( + mediaTypeCosignSimpleSigningV1JSON = "application/vnd.dev.cosign.simplesigning.v1+json" + annotationCosignSignature = "dev.cosignproject.cosign/signature" +) + +// cosignSimpleSigning is the payload cosign signs: it embeds the artifact's +// manifest digest, binding the signature to the exact artifact content. +type cosignSimpleSigning struct { + Critical cosignCritical `json:"critical"` +} + +type cosignCritical struct { + Identity cosignIdentity `json:"identity"` + Image cosignImage `json:"image"` + Type string `json:"type"` +} + +type cosignIdentity struct { + DockerReference string `json:"docker-reference"` +} + +type cosignImage struct { + DockerManifestDigest string `json:"docker-manifest-digest"` +} + +// SimpleSigningPayload builds the canonical simple-signing payload for the +// artifact at ref pinned to digestStr. This payload — not the manifest +// digest — is what gets signed, per the cosign convention: a verifier +// recovers the payload from the signature manifest's layer, checks the +// signature over it, and reads the bound manifest digest out of it. +// Exported because offline re-verification of a stored key-signed bundle +// must reconstruct exactly these bytes to check the signature's binding. +func SimpleSigningPayload(imageRef, digestStr string) ([]byte, error) { + ref, err := name.ParseReference(imageRef) + if err != nil { + return nil, fmt.Errorf("parsing image reference: %w", err) + } + d, err := parseManifestDigest(digestStr) + if err != nil { + return nil, err + } + payload := cosignSimpleSigning{ + Critical: cosignCritical{ + Identity: cosignIdentity{DockerReference: ref.Context().Name()}, + Image: cosignImage{DockerManifestDigest: d.String()}, + Type: "cosign container image signature", + }, + } + return json.Marshal(payload) +} + +// parseManifestDigest validates and normalizes an artifact manifest digest +// string, defaulting a bare hex value to sha256. +func parseManifestDigest(digestStr string) (digest.Digest, error) { + digestStr = strings.TrimSpace(digestStr) + if digestStr == "" { + return "", fmt.Errorf("digest is required for signing") + } + if !strings.Contains(digestStr, ":") { + digestStr = "sha256:" + digestStr + } + d, err := digest.Parse(digestStr) + if err != nil { + return "", fmt.Errorf("parsing digest: %w", err) + } + return d, nil +} + +// attachCosignSignature writes the cosign signature manifest for the +// artifact: an OCI image at the "sha256-.sig" tag whose single layer is +// the simple-signing payload, carrying the signature in the layer's +// annotations. This is the classic cosign layout, chosen deliberately for +// interop — "cosign verify --key" and any Sigstore-aware registry tooling +// can discover and verify it. +func attachCosignSignature( + ctx context.Context, + keychain authn.Keychain, + imageRef, digestStr string, + payload, signatureBytes []byte, +) error { + ref, err := name.ParseReference(imageRef) + if err != nil { + return fmt.Errorf("parsing image reference: %w", err) + } + d, err := parseManifestDigest(digestStr) + if err != nil { + return err + } + + layer := static.NewLayer(payload, mediaTypeCosignSimpleSigningV1JSON) + img, err := mutate.Append(empty.Image, mutate.Addendum{ + Layer: layer, + Annotations: map[string]string{ + annotationCosignSignature: base64.StdEncoding.EncodeToString(signatureBytes), + }, + MediaType: mediaTypeCosignSimpleSigningV1JSON, + }) + if err != nil { + return fmt.Errorf("building signature manifest: %w", err) + } + img = mutate.MediaType(img, types.OCIManifestSchema1) + + h, err := v1.NewHash(d.String()) + if err != nil { + return fmt.Errorf("parsing digest hash: %w", err) + } + sigTag := ref.Context().Tag(fmt.Sprint(h.Algorithm, "-", h.Hex, ".sig")) + remoteOpts := []remote.Option{remote.WithAuthFromKeychain(keychain), remote.WithContext(ctx)} + if err := remote.Write(sigTag, img, remoteOpts...); err != nil { + return fmt.Errorf("pushing signature manifest: %w", err) + } + return nil +} diff --git a/pkg/skills/signer/key.go b/pkg/skills/signer/key.go new file mode 100644 index 0000000000..1475dd12aa --- /dev/null +++ b/pkg/skills/signer/key.go @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package signer + +import ( + "context" + "crypto" + "crypto/rand" + "fmt" + "os" + "path/filepath" + + protocommon "github.com/sigstore/protobuf-specs/gen/pb-go/common/v1" + "github.com/sigstore/sigstore-go/pkg/sign" + "github.com/sigstore/sigstore/pkg/cryptoutils" +) + +// resolveKeyPath canonicalizes and vets a user-supplied signing key path: +// absolute, cleaned, symlinks resolved, and a regular file. The value comes +// straight from a --key flag, so this is the recognized sanitization +// barrier before the path is opened — no string-inspection shortcuts. +func resolveKeyPath(path string) (string, error) { + if path == "" { + return "", ErrKeyRequired + } + abs, err := filepath.Abs(path) + if err != nil { + return "", fmt.Errorf("resolving signing key path: %w", err) + } + resolved, err := filepath.EvalSymlinks(filepath.Clean(abs)) + if err != nil { + return "", fmt.Errorf("resolving signing key path: %w", err) + } + info, err := os.Stat(resolved) + if err != nil { + return "", fmt.Errorf("reading signing key: %w", err) + } + if !info.Mode().IsRegular() { + return "", fmt.Errorf("signing key %q is not a regular file", resolved) + } + return resolved, nil +} + +// loadKeypair reads a cosign PEM private key from path. Encrypted keys are +// decrypted with the COSIGN_PASSWORD environment variable, matching the +// cosign CLI's behavior. +func loadKeypair(path string) (sign.Keypair, error) { + resolved, err := resolveKeyPath(path) + if err != nil { + return nil, err + } + pemBytes, err := os.ReadFile(resolved) //nolint:gosec // resolved by resolveKeyPath from an explicit --key flag + if err != nil { + return nil, fmt.Errorf("reading signing key: %w", err) + } + priv, err := cryptoutils.UnmarshalPEMToPrivateKey(pemBytes, cosignPassFunc()) + if err != nil { + return nil, fmt.Errorf("decoding signing key: %w", err) + } + signerKey, ok := priv.(crypto.Signer) + if !ok { + return nil, fmt.Errorf("signing key type %T cannot sign", priv) + } + return &fileKeypair{priv: signerKey}, nil +} + +func cosignPassFunc() cryptoutils.PassFunc { + if pw := os.Getenv("COSIGN_PASSWORD"); pw != "" { + return cryptoutils.StaticPasswordFunc([]byte(pw)) + } + return func(_ bool) ([]byte, error) { return nil, nil } +} + +// fileKeypair adapts a file-loaded private key to sigstore-go's signing +// Keypair interface (ECDSA P-256 / SHA-256, the cosign default). +type fileKeypair struct { + priv crypto.Signer +} + +func (*fileKeypair) GetHashAlgorithm() protocommon.HashAlgorithm { + return protocommon.HashAlgorithm_SHA2_256 +} + +func (*fileKeypair) GetSigningAlgorithm() protocommon.PublicKeyDetails { + return protocommon.PublicKeyDetails_PKIX_ECDSA_P256_SHA_256 +} + +func (k *fileKeypair) GetHint() []byte { + pubKeyBytes, err := cryptoutils.MarshalPublicKeyToPEM(k.priv.Public()) + if err != nil { + return nil + } + return pubKeyBytes +} + +func (*fileKeypair) GetKeyAlgorithm() string { + return "ecdsa" +} + +func (k *fileKeypair) GetPublicKey() crypto.PublicKey { + return k.priv.Public() +} + +func (k *fileKeypair) GetPublicKeyPem() (string, error) { + pemBytes, err := cryptoutils.MarshalPublicKeyToPEM(k.priv.Public()) + if err != nil { + return "", err + } + return string(pemBytes), nil +} + +// SignData hashes data with SHA-256 and signs the digest, returning +// (signature, digest) — the order sigstore-go's Keypair contract requires +// (see sign.EphemeralKeypair for the reference implementation). +func (k *fileKeypair) SignData(_ context.Context, data []byte) ([]byte, []byte, error) { + hash := crypto.SHA256.New() + if _, err := hash.Write(data); err != nil { + return nil, nil, err + } + digest := hash.Sum(nil) + sig, err := k.priv.Sign(rand.Reader, digest, crypto.SHA256) + if err != nil { + return nil, nil, err + } + return sig, digest, nil +} diff --git a/pkg/skills/signer/signer.go b/pkg/skills/signer/signer.go new file mode 100644 index 0000000000..245ff69e38 --- /dev/null +++ b/pkg/skills/signer/signer.go @@ -0,0 +1,108 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package signer signs skill OCI artifacts with Sigstore, following the +// cosign key-pair convention: a "simple signing" payload binding the +// artifact's manifest digest is signed with the user's key, attached to the +// registry as a cosign signature manifest (the "sha256-.sig" tag), and +// returned as a serialized Sigstore bundle for durable storage and offline +// re-verification. +package signer + +import ( + "context" + "errors" + "fmt" + + "github.com/google/go-containerregistry/pkg/authn" + verifybundle "github.com/sigstore/sigstore-go/pkg/bundle" + "github.com/sigstore/sigstore-go/pkg/sign" +) + +// ErrKeyRequired indicates no signing key was provided. Keyless (OIDC) +// signing is not implemented yet, so a cosign private key is the only +// supported signing method. +var ErrKeyRequired = errors.New("signing key required: pass --key with a cosign private key") + +// Options configures OCI signing. +type Options struct { + // Key is the path to a cosign PEM-encoded private key file. An + // encrypted key is decrypted with the COSIGN_PASSWORD environment + // variable, matching the cosign CLI. + Key string +} + +// Signer signs skill OCI artifacts and attaches the signature to the +// registry. +type Signer interface { + // SignOCI signs the artifact at ref pinned to the given manifest digest + // ("sha256:..."), attaches the signature as a cosign signature manifest + // next to the artifact, and returns the serialized Sigstore bundle for + // storage. + SignOCI(ctx context.Context, ref, digest string, opts Options) ([]byte, error) +} + +// Default implements Signer with file-based cosign keys. +type Default struct { + keychain authn.Keychain +} + +var _ Signer = (*Default)(nil) + +// NewDefault creates a signer using the given registry auth keychain for +// pushing the signature manifest. A nil keychain falls back to the default +// keychain. +func NewDefault(keychain authn.Keychain) *Default { + if keychain == nil { + keychain = authn.DefaultKeychain + } + return &Default{keychain: keychain} +} + +// SignOCI signs the artifact following the cosign convention: the signature +// is computed over the simple-signing payload (which embeds the manifest +// digest, binding the signature to the artifact), the SAME signature is +// attached to the registry as a cosign signature manifest, and the returned +// bundle carries it with the payload's digest as the signed message. A +// verifier reconstructing the bundle from the registry manifest (or +// re-verifying the stored bundle offline) therefore checks exactly the +// signature that was attached — one signature, two representations. +func (d *Default) SignOCI(ctx context.Context, ref, digestStr string, opts Options) ([]byte, error) { + if opts.Key == "" { + return nil, ErrKeyRequired + } + keypair, err := loadKeypair(opts.Key) + if err != nil { + return nil, err + } + + payload, err := SimpleSigningPayload(ref, digestStr) + if err != nil { + return nil, err + } + + // sign.Bundle invokes the keypair's SignData over the payload and + // records both the signature and the payload digest in the bundle. + pb, err := sign.Bundle(&sign.PlainData{Data: payload}, keypair, sign.BundleOptions{Context: ctx}) + if err != nil { + return nil, fmt.Errorf("building sigstore bundle: %w", err) + } + msgSig := pb.GetMessageSignature() + if msgSig == nil || len(msgSig.GetSignature()) == 0 { + return nil, errors.New("signing produced no message signature") + } + + if err := attachCosignSignature(ctx, d.keychain, ref, digestStr, payload, msgSig.GetSignature()); err != nil { + return nil, fmt.Errorf("attaching signature manifest: %w", err) + } + + bun, err := verifybundle.NewBundle(pb) + if err != nil { + return nil, fmt.Errorf("finalizing sigstore bundle: %w", err) + } + raw, err := bun.MarshalJSON() + if err != nil { + return nil, fmt.Errorf("serializing sigstore bundle: %w", err) + } + return raw, nil +} diff --git a/pkg/skills/signer/signer_test.go b/pkg/skills/signer/signer_test.go new file mode 100644 index 0000000000..7e72332604 --- /dev/null +++ b/pkg/skills/signer/signer_test.go @@ -0,0 +1,336 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package signer + +import ( + "crypto" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "encoding/pem" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/google/go-containerregistry/pkg/crane" + "github.com/google/go-containerregistry/pkg/name" + "github.com/google/go-containerregistry/pkg/registry" + v1 "github.com/google/go-containerregistry/pkg/v1" + "github.com/google/go-containerregistry/pkg/v1/random" + "github.com/google/go-containerregistry/pkg/v1/remote" + protocommon "github.com/sigstore/protobuf-specs/gen/pb-go/common/v1" + "github.com/sigstore/sigstore-go/pkg/bundle" + "github.com/sigstore/sigstore-go/pkg/root" + "github.com/sigstore/sigstore-go/pkg/verify" + "github.com/sigstore/sigstore/pkg/cryptoutils" + "github.com/sigstore/sigstore/pkg/signature" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeTestKey generates an ECDSA P-256 key pair, writes the private key as +// a cosign-style PEM file, and returns its path plus the public key PEM. +func writeTestKey(t *testing.T) (keyPath string, pubPEM []byte) { + t.Helper() + priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + + privPEM, err := cryptoutils.MarshalPrivateKeyToPEM(priv) + require.NoError(t, err) + keyPath = filepath.Join(t.TempDir(), "cosign.key") + require.NoError(t, os.WriteFile(keyPath, privPEM, 0o600)) + + pubPEM, err = cryptoutils.MarshalPublicKeyToPEM(priv.Public()) + require.NoError(t, err) + return keyPath, pubPEM +} + +// pushTestArtifact pushes a random OCI image to the in-process registry and +// returns its reference and manifest digest. +func pushTestArtifact(t *testing.T, registryHost string) (ref string, digestStr string) { + t.Helper() + img, err := random.Image(256, 1) + require.NoError(t, err) + ref = registryHost + "/test/skill:v1" + parsed, err := name.ParseReference(ref) + require.NoError(t, err) + require.NoError(t, remote.Write(parsed, img)) + d, err := img.Digest() + require.NoError(t, err) + return ref, d.String() +} + +// verifyKeyBundle verifies a serialized bundle against the public key and +// the given payload digest — the same semantics toolhive-core's +// VerifyBundleWithKey exposes (used directly here until the core release +// this stack's verifier PR consumes). +func verifyKeyBundle(t *testing.T, raw, pubPEM []byte, payloadDigest []byte) error { + t.Helper() + pub, err := cryptoutils.UnmarshalPEMToPublicKey(pubPEM) + require.NoError(t, err) + sigVerifier, err := signature.LoadVerifier(pub, crypto.SHA256) + require.NoError(t, err) + tm := root.NewTrustedPublicKeyMaterial(func(string) (root.TimeConstrainedVerifier, error) { + return root.NewExpiringKey(sigVerifier, time.Time{}, time.Time{}), nil + }) + sev, err := verify.NewVerifier(tm, verify.WithNoObserverTimestamps()) + require.NoError(t, err) + + parsed := &bundle.Bundle{} + require.NoError(t, parsed.UnmarshalJSON(raw)) + _, err = sev.Verify(parsed, verify.NewPolicy( + verify.WithArtifactDigest("sha256", payloadDigest), + verify.WithKey(), + )) + return err +} + +//nolint:paralleltest // uses t.Setenv indirectly via subtests sharing a registry +func TestSignOCIRoundTrip(t *testing.T) { + reg := httptest.NewServer(registry.New()) + t.Cleanup(reg.Close) + host := strings.TrimPrefix(reg.URL, "http://") + + keyPath, pubPEM := writeTestKey(t) + ref, digestStr := pushTestArtifact(t, host) + + raw, err := NewDefault(nil).SignOCI(t.Context(), ref, digestStr, Options{Key: keyPath}) + require.NoError(t, err) + require.NotEmpty(t, raw) + + // The returned bundle verifies against the signing key over the + // simple-signing payload digest. + payload, err := SimpleSigningPayload(ref, digestStr) + require.NoError(t, err) + payloadDigest := sha256.Sum256(payload) + require.NoError(t, verifyKeyBundle(t, raw, pubPEM, payloadDigest[:]), + "the returned bundle must verify against the signing key") + + // The attached signature manifest reconstructs to the SAME signature: + // this is the layout toolhive-core's bundle retrieval (and cosign) + // discover — one signature, two representations. + h, err := v1.NewHash(digestStr) + require.NoError(t, err) + sigRef := host + "/test/skill:" + h.Algorithm + "-" + h.Hex + ".sig" + manifestBytes, err := crane.Manifest(sigRef) + require.NoError(t, err, "the signature manifest must exist at the cosign .sig tag") + + var manifest struct { + Layers []struct { + MediaType string `json:"mediaType"` + Digest string `json:"digest"` + Annotations map[string]string `json:"annotations"` + } `json:"layers"` + } + require.NoError(t, json.Unmarshal(manifestBytes, &manifest)) + require.Len(t, manifest.Layers, 1) + layer := manifest.Layers[0] + assert.Equal(t, mediaTypeCosignSimpleSigningV1JSON, layer.MediaType) + + // The layer content digest is the payload digest the bundle signed. + wantLayerDigest := "sha256:" + hex.EncodeToString(payloadDigest[:]) + assert.Equal(t, wantLayerDigest, layer.Digest, + "the signature manifest layer must be the exact signed payload") + + // The annotation signature matches the bundle's message signature. + parsed := &bundle.Bundle{} + require.NoError(t, parsed.UnmarshalJSON(raw)) + bundleSig := parsed.Bundle.GetMessageSignature().GetSignature() + annotationSig, err := base64.StdEncoding.DecodeString(layer.Annotations[annotationCosignSignature]) + require.NoError(t, err) + assert.Equal(t, bundleSig, annotationSig, + "the attached annotation and the returned bundle must carry the same signature") +} + +//nolint:paralleltest // shares process env semantics with the round-trip test +func TestSignOCIRejectsWrongKeyVerification(t *testing.T) { + reg := httptest.NewServer(registry.New()) + t.Cleanup(reg.Close) + host := strings.TrimPrefix(reg.URL, "http://") + + keyPath, _ := writeTestKey(t) + _, otherPub := writeTestKey(t) + ref, digestStr := pushTestArtifact(t, host) + + raw, err := NewDefault(nil).SignOCI(t.Context(), ref, digestStr, Options{Key: keyPath}) + require.NoError(t, err) + + payload, err := SimpleSigningPayload(ref, digestStr) + require.NoError(t, err) + payloadDigest := sha256.Sum256(payload) + require.Error(t, verifyKeyBundle(t, raw, otherPub, payloadDigest[:]), + "a different key must not verify the bundle") +} + +func TestSignOCIRequiresKey(t *testing.T) { + t.Parallel() + _, err := NewDefault(nil).SignOCI(t.Context(), "ghcr.io/org/skill:v1", "sha256:abc", Options{}) + require.ErrorIs(t, err, ErrKeyRequired) +} + +func TestSignOCIEncryptedKey(t *testing.T) { + reg := httptest.NewServer(registry.New()) + t.Cleanup(reg.Close) + host := strings.TrimPrefix(reg.URL, "http://") + + // Generate an encrypted cosign-style key: the "ENCRYPTED SIGSTORE + // PRIVATE KEY" PEM type wrapping password-encrypted DER, which is what + // `cosign generate-key-pair` produces. + priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + encDER, err := cryptoutils.MarshalPrivateKeyToEncryptedDER( + priv, cryptoutils.StaticPasswordFunc([]byte("test-password"))) + require.NoError(t, err) + encPEM := pem.EncodeToMemory(&pem.Block{ + Type: string(cryptoutils.EncryptedSigstorePrivateKeyPEMType), + Bytes: encDER, + }) + keyPath := filepath.Join(t.TempDir(), "cosign-enc.key") + require.NoError(t, os.WriteFile(keyPath, encPEM, 0o600)) + + ref, digestStr := pushTestArtifact(t, host) + + t.Setenv("COSIGN_PASSWORD", "test-password") + raw, err := NewDefault(nil).SignOCI(t.Context(), ref, digestStr, Options{Key: keyPath}) + require.NoError(t, err, "an encrypted key must decrypt via COSIGN_PASSWORD") + require.NotEmpty(t, raw) + + t.Setenv("COSIGN_PASSWORD", "wrong-password") + _, err = NewDefault(nil).SignOCI(t.Context(), ref, digestStr, Options{Key: keyPath}) + require.Error(t, err, "a wrong password must fail key decryption") +} + +func TestSignOCIInputErrors(t *testing.T) { + t.Parallel() + keyPath, _ := writeTestKey(t) + + tests := []struct { + name string + key string + ref string + digest string + wantErr string + }{ + { + name: "garbage key file", + key: writeGarbageKey(t), + ref: "example.com/org/skill:v1", + digest: "sha256:" + strings.Repeat("a", 64), + wantErr: "decoding signing key", + }, + { + name: "invalid image reference", + key: keyPath, + ref: "NOT a valid ref!!", + digest: "sha256:" + strings.Repeat("a", 64), + wantErr: "parsing image reference", + }, + { + name: "empty digest", + key: keyPath, + ref: "example.com/org/skill:v1", + digest: "", + wantErr: "digest is required", + }, + { + name: "malformed digest", + key: keyPath, + ref: "example.com/org/skill:v1", + digest: "sha256:tooshort", + wantErr: "parsing digest", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + _, err := NewDefault(nil).SignOCI(t.Context(), tc.ref, tc.digest, Options{Key: tc.key}) + require.ErrorContains(t, err, tc.wantErr) + }) + } +} + +func writeGarbageKey(t *testing.T) string { + t.Helper() + p := filepath.Join(t.TempDir(), "garbage.key") + require.NoError(t, os.WriteFile(p, []byte("not a pem key"), 0o600)) + return p +} + +func TestParseManifestDigestNormalizesBareHex(t *testing.T) { + t.Parallel() + hexDigest := strings.Repeat("a", 64) + d, err := parseManifestDigest(hexDigest) + require.NoError(t, err) + assert.Equal(t, "sha256:"+hexDigest, d.String()) +} + +func TestSimpleSigningPayloadBindsDigestAndRepo(t *testing.T) { + t.Parallel() + digestStr := "sha256:" + strings.Repeat("b", 64) + payload, err := SimpleSigningPayload("example.com/org/skill:v1", digestStr) + require.NoError(t, err) + + var got cosignSimpleSigning + require.NoError(t, json.Unmarshal(payload, &got)) + assert.Equal(t, digestStr, got.Critical.Image.DockerManifestDigest) + // The identity is the repository, not the tag — retagging must not + // invalidate the signature's digest binding. + assert.Equal(t, "example.com/org/skill", got.Critical.Identity.DockerReference) + assert.Equal(t, "cosign container image signature", got.Critical.Type) +} + +func TestFileKeypairMetadata(t *testing.T) { + t.Parallel() + keyPath, pubPEM := writeTestKey(t) + kp, err := loadKeypair(keyPath) + require.NoError(t, err) + + assert.Equal(t, "ecdsa", kp.GetKeyAlgorithm()) + assert.Equal(t, protocommon.HashAlgorithm_SHA2_256, kp.GetHashAlgorithm()) + assert.Equal(t, protocommon.PublicKeyDetails_PKIX_ECDSA_P256_SHA_256, kp.GetSigningAlgorithm()) + gotPEM, err := kp.GetPublicKeyPem() + require.NoError(t, err) + assert.Equal(t, string(pubPEM), gotPEM) + assert.Equal(t, pubPEM, kp.GetHint()) + assert.NotNil(t, kp.GetPublicKey()) +} + +func TestResolveKeyPath(t *testing.T) { + t.Parallel() + + t.Run("regular file resolves", func(t *testing.T) { + t.Parallel() + p := filepath.Join(t.TempDir(), "key.pem") + require.NoError(t, os.WriteFile(p, []byte("x"), 0o600)) + resolved, err := resolveKeyPath(p) + require.NoError(t, err) + assert.True(t, filepath.IsAbs(resolved)) + }) + + t.Run("directory rejected", func(t *testing.T) { + t.Parallel() + _, err := resolveKeyPath(t.TempDir()) + require.Error(t, err) + assert.Contains(t, err.Error(), "not a regular file") + }) + + t.Run("missing file rejected", func(t *testing.T) { + t.Parallel() + _, err := resolveKeyPath(filepath.Join(t.TempDir(), "nope.pem")) + require.Error(t, err) + }) + + t.Run("empty path rejected", func(t *testing.T) { + t.Parallel() + _, err := resolveKeyPath("") + require.ErrorIs(t, err, ErrKeyRequired) + }) +} diff --git a/pkg/skills/skillsvc/content_digest.go b/pkg/skills/skillsvc/content_digest.go new file mode 100644 index 0000000000..839b7552be --- /dev/null +++ b/pkg/skills/skillsvc/content_digest.go @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "fmt" + "os" + + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// skillMDFileName is the well-known skill definition file every installed +// skill directory contains. +const skillMDFileName = "SKILL.md" + +// computeContentDigest hashes the on-disk file set for an installed skill, +// for recording in the lock file at install time. Every client directory is +// written from the same source in the same operation, so the first client's +// copy is representative *at record time*; verification (sync --check) must +// not make that assumption and checks every client directory — see +// entryMatchesInstalled. +func computeContentDigest(pathResolver skills.PathResolver, sk skills.InstalledSkill) (string, error) { + dir, err := installedSkillDir(pathResolver, sk) + if err != nil { + return "", err + } + digest, err := lockfile.ContentDigestFromDir(dir) + if err != nil { + return "", fmt.Errorf("computing content digest: %w", err) + } + return digest, nil +} + +// readSkillMD reads and parses SKILL.md from an installed skill's directory, +// for discovering toolhive.requires dependencies at install time. The +// artifact source (OCI labels or a git checkout) is not normalized to expose +// requires before extraction, but every installed skill has SKILL.md on disk +// once extraction succeeds, so reading it back is the uniform way to recover +// dependency declarations regardless of source type. +func readSkillMD(pathResolver skills.PathResolver, sk skills.InstalledSkill) (*skills.ParseResult, error) { + dir, err := installedSkillDir(pathResolver, sk) + if err != nil { + return nil, err + } + root, err := os.OpenRoot(dir) + if err != nil { + return nil, fmt.Errorf("opening skill directory: %w", err) + } + defer func() { _ = root.Close() }() + + content, err := root.ReadFile(skillMDFileName) + if err != nil { + return nil, fmt.Errorf("reading %s: %w", skillMDFileName, err) + } + parsed, err := skills.ParseSkillMD(content) + if err != nil { + return nil, fmt.Errorf("parsing %s: %w", skillMDFileName, err) + } + return parsed, nil +} + +// installedSkillDir resolves the filesystem path for one of sk's installed +// clients, used as a representative directory for reading back on-disk state. +func installedSkillDir(pathResolver skills.PathResolver, sk skills.InstalledSkill) (string, error) { + if len(sk.Clients) == 0 { + return "", fmt.Errorf("skill %q has no installed clients", sk.Metadata.Name) + } + dir, err := pathResolver.GetSkillPath(sk.Clients[0], sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + if err != nil { + return "", fmt.Errorf("resolving skill path: %w", err) + } + return dir, nil +} diff --git a/pkg/skills/skillsvc/install.go b/pkg/skills/skillsvc/install.go index 028f78530f..f0ac760916 100644 --- a/pkg/skills/skillsvc/install.go +++ b/pkg/skills/skillsvc/install.go @@ -5,15 +5,19 @@ package skillsvc import ( "context" + "errors" "fmt" "log/slog" "net/http" "strings" + nameref "github.com/google/go-containerregistry/pkg/name" + "github.com/stacklok/toolhive-core/httperr" "github.com/stacklok/toolhive/pkg/groups" "github.com/stacklok/toolhive/pkg/skills" "github.com/stacklok/toolhive/pkg/skills/gitresolver" + "github.com/stacklok/toolhive/pkg/skills/lockfile" ) // Install installs a skill. When the Name field contains an OCI reference @@ -22,6 +26,13 @@ import ( // to disk and a full installation record is created. Without LayerData, a // pending record is created. func (s *service) Install(ctx context.Context, opts skills.InstallOptions) (*skills.InstallResult, error) { + // Captured before any internal resolution (version splicing, registry + // lookup, git/OCI dispatch) mutates a *local* copy of opts.Name in the + // functions below. This is the RFC THV-0080 lock entry "source": exactly + // what the caller asked for, preserved verbatim so upgrade can re-resolve + // the same input later. + originalName := opts.Name + scope, projectRoot, err := normalizeProjectRoot(opts.Scope, opts.ProjectRoot) if err != nil { return nil, err @@ -31,63 +42,45 @@ func (s *service) Install(ctx context.Context, opts skills.InstallOptions) (*ski // the same lock key and DB record. opts.ProjectRoot = projectRoot - // Git references (git://host/owner/repo[@ref][#path]) are dispatched first; - // the prefix is unambiguous and cannot collide with OCI references. - if gitresolver.IsGitReference(opts.Name) { - result, err := s.installFromGit(ctx, opts, scope) - if err != nil { - return nil, err - } - return s.installAndRegister(ctx, result, opts.Group, result.Skill.Metadata.Name, scope, opts.ProjectRoot) - } - // When the caller supplies `version` separately and the name is a tag-less // OCI-like reference (contains '/' but no ':' or '@'), splice the version // in as the tag. Without this, parseOCIReference + qualifiedOCIRef would // default the pull to ":latest" and silently drop opts.Version. An // explicit tag in the name still wins (we only splice when none is set). + // Git references are unaffected: the git:// prefix contains "://". if opts.Version != "" && + !gitresolver.IsGitReference(opts.Name) && strings.ContainsRune(opts.Name, '/') && !strings.ContainsAny(opts.Name, ":@") { opts.Name = opts.Name + ":" + opts.Version } - ref, isOCI, err := parseOCIReference(opts.Name) - if err != nil { - return nil, httperr.WithCode( - fmt.Errorf("invalid OCI reference %q: %w", opts.Name, err), - http.StatusBadRequest, - ) - } - if isOCI { - result, ociErr := s.installFromOCI(ctx, opts, scope, ref) - if ociErr == nil { - return s.installAndRegister(ctx, result, opts.Group, opts.Name, scope, opts.ProjectRoot) - } - // OCI pull failed — fall back to registry lookup for names that look - // like a qualified "namespace/name". Names that are unambiguously OCI - // (digest, explicit tag, or multi-segment path) must not trigger a - // registry search. See isUnambiguousOCIRef for the full rule set. - if isUnambiguousOCIRef(opts.Name, ref) { - return nil, ociErr - } - slog.Debug("OCI pull failed, attempting registry fallback", "name", opts.Name, "error", ociErr) - resolved, regErr := s.resolveFromRegistry(opts.Name) - if regErr != nil { - return nil, regErr - } - if resolved != nil { - return s.installFromResolvedRegistry(ctx, opts, scope, resolved) - } - return nil, ociErr - } - - // Plain skill name — validate and proceed with existing flow. - if err := skills.ValidateSkillName(opts.Name); err != nil { - return nil, httperr.WithCode(err, http.StatusBadRequest) - } - - return s.installByName(ctx, opts, scope) + // Route through the shared source-dispatch skeleton (git, direct OCI + // with registry fallback, plain name) — the same sequence upgrade's + // re-resolution uses, so the two cannot drift. + return dispatchSource(ctx, s, opts.Name, sourceOps[*skills.InstallResult]{ + git: func(ctx context.Context, _ string) (*skills.InstallResult, error) { + result, err := s.installFromGit(ctx, &opts, scope) + if err != nil { + return nil, err + } + return s.installAndRegister(ctx, opts, originalName, result, opts.Group, result.Skill.Metadata.Name, scope) + }, + oci: func(ctx context.Context, ref nameref.Reference) (*skills.InstallResult, error) { + result, err := s.installFromOCI(ctx, &opts, scope, ref) + if err != nil { + slog.Debug("OCI pull failed, registry fallback may apply", "name", opts.Name, "error", err) + return nil, err + } + return s.installAndRegister(ctx, opts, originalName, result, opts.Group, opts.Name, scope) + }, + registry: func(ctx context.Context, resolved *registryResolveResult) (*skills.InstallResult, error) { + return s.installFromResolvedRegistry(ctx, opts, originalName, scope, resolved) + }, + plainName: func(ctx context.Context, _ string) (*skills.InstallResult, error) { + return s.installByName(ctx, opts, originalName, scope) + }, + }) } // installByName handles installation for a validated plain skill name. It @@ -95,6 +88,7 @@ func (s *service) Install(ctx context.Context, opts skills.InstallOptions) (*ski func (s *service) installByName( ctx context.Context, opts skills.InstallOptions, + originalName string, scope skills.Scope, ) (*skills.InstallResult, error) { unlock := s.locks.lock(opts.Name, scope, opts.ProjectRoot) @@ -125,16 +119,27 @@ func (s *service) installByName( unlock() locked = false - return s.installFromRegistryLookup(ctx, opts, scope) + return s.installFromRegistryLookup(ctx, opts, originalName, scope) } // resolved: opts hydrated, fall through to installWithExtraction } + // Local-store artifacts and raw layer data carry no registry signature + // to verify — installing them project-scoped is an unsigned trust + // decision that must be explicit. + if shouldVerifyInstall(opts, scope) { + decision, verifyErr := verifyLocalInstall(opts, opts.Name) + if verifyErr != nil { + return nil, verifyErr + } + applyDecisionToOpts(&opts, decision) + } + result, err := s.installWithExtraction(ctx, opts, scope) if err != nil { return nil, err } - return s.installAndRegister(ctx, result, opts.Group, opts.Name, scope, opts.ProjectRoot) + return s.installAndRegister(ctx, opts, originalName, result, opts.Group, opts.Name, scope) } // installFromRegistryLookup resolves a plain skill name via the registry and @@ -142,6 +147,7 @@ func (s *service) installByName( func (s *service) installFromRegistryLookup( ctx context.Context, opts skills.InstallOptions, + originalName string, scope skills.Scope, ) (*skills.InstallResult, error) { resolved, regErr := s.resolveFromRegistry(opts.Name) @@ -149,7 +155,7 @@ func (s *service) installFromRegistryLookup( return nil, regErr } if resolved != nil { - return s.installFromResolvedRegistry(ctx, opts, scope, resolved) + return s.installFromResolvedRegistry(ctx, opts, originalName, scope, resolved) } return nil, httperr.WithCode( @@ -165,6 +171,7 @@ func (s *service) installFromRegistryLookup( func (s *service) installFromResolvedRegistry( ctx context.Context, opts skills.InstallOptions, + originalName string, scope skills.Scope, resolved *registryResolveResult, ) (*skills.InstallResult, error) { @@ -172,22 +179,22 @@ func (s *service) installFromResolvedRegistry( case resolved.OCIRef != nil: slog.Info("resolved skill from registry (OCI)", "name", opts.Name, "oci_reference", resolved.OCIRef.String()) opts.Name = resolved.OCIRef.String() - result, ociErr := s.installFromOCI(ctx, opts, scope, resolved.OCIRef) + result, ociErr := s.installFromOCI(ctx, &opts, scope, resolved.OCIRef) if ociErr != nil { return nil, ociErr } // Use the skill name extracted from the artifact, not opts.Name which // holds the OCI ref string. installFromOCI mutates its own copy of opts // (Go pass-by-value), so the caller never sees the updated name. - return s.installAndRegister(ctx, result, opts.Group, result.Skill.Metadata.Name, scope, opts.ProjectRoot) + return s.installAndRegister(ctx, opts, originalName, result, opts.Group, result.Skill.Metadata.Name, scope) case resolved.GitURL != "": slog.Info("resolved skill from registry (git)", "name", opts.Name, "git_url", resolved.GitURL) opts.Name = resolved.GitURL - result, gitErr := s.installFromGit(ctx, opts, scope) + result, gitErr := s.installFromGit(ctx, &opts, scope) if gitErr != nil { return nil, gitErr } - return s.installAndRegister(ctx, result, opts.Group, result.Skill.Metadata.Name, scope, opts.ProjectRoot) + return s.installAndRegister(ctx, opts, originalName, result, opts.Group, result.Skill.Metadata.Name, scope) } return nil, httperr.WithCode( fmt.Errorf("skill %q resolved from registry but has no installable package", opts.Name), @@ -208,24 +215,110 @@ func (s *service) registerSkillInGroup(ctx context.Context, groupName string, sk return groups.AddSkillToGroup(ctx, s.groupManager, groupName, skillName) } -// installAndRegister registers the just-installed skill in the target group. -// If group registration fails, the DB record is rolled back so that a retry -// starts fresh rather than leaving the system in an inconsistent state (skill -// installed but not in the expected group). +// installAndRegister registers the just-installed skill in the target group +// and, for project-scope installs with the lock file feature enabled (see +// skills.LockFileFeatureEnabled), records it — and any toolhive.requires +// dependencies — in the project's toolhive.lock.yaml. If group registration +// or the lock write fails, the DB record and lock entry are rolled back to +// their pre-install state: restored when this call updated a pre-existing +// record (a --force reinstall must not be destroyed by a transient failure), +// deleted — with a dependency cascade cleaning up any freshly materialized +// orphans — when this call created them. func (s *service) installAndRegister( ctx context.Context, + opts skills.InstallOptions, + originalName string, result *skills.InstallResult, groupName string, skillName string, scope skills.Scope, - projectRoot string, ) (*skills.InstallResult, error) { + lockScoped := scope == skills.ScopeProject && skills.LockFileFeatureEnabled() + + // Snapshot the prior lock entry before anything below can write one, so + // rollback can reinstate it (RequiredBy links from other parents + // included) rather than blindly deleting it. + var prevEntry *lockfile.Entry + if lockScoped { + if root, rootErr := lockfile.OpenRoot(opts.ProjectRoot); rootErr == nil { + if lf, loadErr := lockfile.Load(root); loadErr == nil { + if e, ok := lf.Get(skillName); ok { + prevEntry = &e + } + } + } + } + + rollback := func() { s.rollbackInstall(ctx, opts, result, skillName, scope, lockScoped, prevEntry) } + if err := s.registerSkillInGroup(ctx, groupName, skillName); err != nil { - // Best-effort rollback: remove the DB record so retries start fresh. - // Files on disk are left in place; a fresh install will detect them - // and either overwrite (force) or return a conflict. - _ = s.store.Delete(ctx, skillName, scope, projectRoot) + // Best-effort rollback. Files on disk are left in place; a fresh + // install will detect them and either overwrite (force) or return a + // conflict. + rollback() return nil, fmt.Errorf("registering skill in group: %w", err) } + + if lockScoped { + updated, err := s.recordLockState(ctx, opts, originalName, result.Skill) + if err != nil { + rollback() + // Preserve a specific code already attached deeper in the chain + // — dependency materialization runs inside recordLockState, so a + // dep's 502 (git resolve) or 404 (registry miss) must reach the + // API boundary as itself, not masked to 500. Only a code-less + // failure (e.g. an actual lock write error) defaults to 500. + wrapped := fmt.Errorf("recording skill in project lock file: %w", err) + var coded *httperr.CodedError + if !errors.As(err, &coded) { + wrapped = httperr.WithCode(wrapped, http.StatusInternalServerError) + } + return nil, wrapped + } + result.Skill = updated + } + return result, nil } + +// rollbackInstall undoes installAndRegister's side effects after a failure, +// best-effort. The DB record is restored to its pre-install snapshot when +// one exists (result.PreExisting) and deleted otherwise; the lock entry is +// likewise reinstated from prevEntry or removed. When this call created the +// entry, removal runs the same dependency cascade as uninstall so that +// freshly materialized dependencies — installed, marked managed, and +// required only by the now-rolled-back skill — do not leak as orphans, +// while pre-existing dependencies with other parents (or explicit installs) +// survive with this skill stripped from their RequiredBy. +func (s *service) rollbackInstall( + ctx context.Context, + opts skills.InstallOptions, + result *skills.InstallResult, + skillName string, + scope skills.Scope, + lockScoped bool, + prevEntry *lockfile.Entry, +) { + if result.PreExisting != nil { + _ = s.store.Update(ctx, *result.PreExisting) + } else { + _ = s.store.Delete(ctx, skillName, scope, opts.ProjectRoot) + } + + if !lockScoped { + return + } + if prevEntry != nil { + if root, err := lockfile.OpenRoot(opts.ProjectRoot); err == nil { + _ = lockfile.UpsertEntry(root, *prevEntry) + } + return + } + uninstallOpts := skills.UninstallOptions{Name: skillName, Scope: scope, ProjectRoot: opts.ProjectRoot} + candidates, err := removeLockEntry(uninstallOpts) + if err != nil { + return + } + visited := map[string]struct{}{skillName: {}} + _ = s.cascadeUninstall(ctx, candidates, visited, opts.ProjectRoot, scope) +} diff --git a/pkg/skills/skillsvc/install_extraction.go b/pkg/skills/skillsvc/install_extraction.go index 60a4da681b..4824a5d43e 100644 --- a/pkg/skills/skillsvc/install_extraction.go +++ b/pkg/skills/skillsvc/install_extraction.go @@ -33,12 +33,33 @@ func (s *service) installWithExtraction( return nil, fmt.Errorf("checking existing skill: %w", storeErr) } - if isExtractionNoOp(existing, storeErr, opts, clientTypes) { + result, err := s.dispatchExtraction(ctx, opts, scope, existing, storeErr, clientTypes, clientDirs) + if err == nil && storeErr == nil { + // Preserve the pre-install record so a later rollback (e.g. a failed + // dependency materialization) can restore it rather than delete it. + pre := existing + result.PreExisting = &pre + } + return result, err +} + +// dispatchExtraction routes an extraction-based install to the no-op, +// same-digest, upgrade, or fresh path based on the pre-install store state. +func (s *service) dispatchExtraction( + ctx context.Context, + opts skills.InstallOptions, + scope skills.Scope, + existing skills.InstalledSkill, + storeErr error, + clientTypes []string, + clientDirs map[string]string, +) (*skills.InstallResult, error) { + if !opts.SyncRestore && isExtractionNoOp(existing, storeErr, opts, clientTypes) { return &skills.InstallResult{Skill: existing}, nil } digestMatches := storeErr == nil && existing.Digest == opts.Digest - if digestMatches && storeErr == nil { + if digestMatches && storeErr == nil && !opts.SyncRestore { return s.installExtractionSameDigestNewClients(ctx, opts, scope, existing, clientTypes, clientDirs) } @@ -199,12 +220,13 @@ func buildInstalledSkill( Name: opts.Name, Version: opts.Version, }, - Scope: scope, - ProjectRoot: opts.ProjectRoot, - Reference: opts.Reference, - Digest: opts.Digest, - Status: skills.InstallStatusInstalled, - InstalledAt: time.Now().UTC(), - Clients: clients, + Scope: scope, + ProjectRoot: opts.ProjectRoot, + Reference: opts.Reference, + Digest: opts.Digest, + Status: skills.InstallStatusInstalled, + InstalledAt: time.Now().UTC(), + Clients: clients, + SigstoreBundle: opts.SigstoreBundle, } } diff --git a/pkg/skills/skillsvc/install_git.go b/pkg/skills/skillsvc/install_git.go index 9a9bb672bc..5ea8e247a5 100644 --- a/pkg/skills/skillsvc/install_git.go +++ b/pkg/skills/skillsvc/install_git.go @@ -22,7 +22,7 @@ import ( // same-commit no-op and upgrade detection. func (s *service) installFromGit( ctx context.Context, - opts skills.InstallOptions, + opts *skills.InstallOptions, scope skills.Scope, ) (*skills.InstallResult, error) { if s.gitResolver == nil { @@ -70,6 +70,7 @@ func (s *service) installFromGit( opts.Name = resolved.SkillConfig.Name opts.Reference = gitURL opts.Digest = resolved.CommitHash + if opts.Version == "" && resolved.SkillConfig.Version != "" { opts.Version = resolved.SkillConfig.Version } @@ -77,12 +78,25 @@ func (s *service) installFromGit( unlock := s.locks.lock(opts.Name, scope, opts.ProjectRoot) defer unlock() - clientTypes, clientDirs, err := s.resolveAndValidateClients(opts, opts.Name, scope, opts.ProjectRoot) + // Verify the commit signature before anything is written or recorded. + // This runs under the per-skill lock so concurrent first installs + // cannot both read an absent lock entry and race their TOFU anchors. + if shouldVerifyInstall(*opts, scope) { + decision, verifyErr := s.verifyGitInstall( + ctx, *opts, resolved.SkillConfig.Name, resolved.CommitPayload, resolved.CommitSignature, + ) + if verifyErr != nil { + return nil, verifyErr + } + applyDecisionToOpts(opts, decision) + } + + clientTypes, clientDirs, err := s.resolveAndValidateClients(*opts, opts.Name, scope, opts.ProjectRoot) if err != nil { return nil, err } - return s.applyGitInstall(ctx, opts, scope, clientTypes, clientDirs, resolved.Files) + return s.applyGitInstall(ctx, *opts, scope, clientTypes, clientDirs, resolved.Files) } // applyGitInstall handles the create/upgrade/no-op logic for a git-based skill @@ -102,7 +116,15 @@ func (s *service) applyGitInstall( return nil, fmt.Errorf("checking existing skill: %w", storeErr) } if !isNotFound { - return s.applyGitInstallExisting(ctx, opts, scope, existing, clientTypes, clientDirs, files) + result, err := s.applyGitInstallExisting(ctx, opts, scope, existing, clientTypes, clientDirs, files) + if err == nil { + // Preserve the pre-install record so a later rollback (e.g. a + // failed dependency materialization) can restore it rather than + // delete it. + pre := existing + result.PreExisting = &pre + } + return result, err } return s.applyGitInstallFresh(ctx, opts, scope, clientTypes, clientDirs, files) } @@ -116,7 +138,11 @@ func (s *service) applyGitInstallExisting( clientDirs map[string]string, files []gitresolver.FileEntry, ) (*skills.InstallResult, error) { - if existing.Digest != opts.Digest { + // SyncRestore forces the same full re-extraction path as a digest + // change: sync repairs on-disk drift that happened without the pinned + // digest changing, so the "same digest means content is already + // correct" no-op/skip branches below must not apply. + if existing.Digest != opts.Digest || opts.SyncRestore { allClients, allDirs, err := s.expandToExistingClients( existing.Clients, clientTypes, clientDirs, opts.Name, scope, opts.ProjectRoot) if err != nil { diff --git a/pkg/skills/skillsvc/install_oci.go b/pkg/skills/skillsvc/install_oci.go index f77be955df..b902b18ade 100644 --- a/pkg/skills/skillsvc/install_oci.go +++ b/pkg/skills/skillsvc/install_oci.go @@ -32,7 +32,7 @@ const maxCompressedLayerSize int64 = 50 * 1024 * 1024 // 50 MB // metadata and layer data, then delegates to the standard extraction flow. func (s *service) installFromOCI( ctx context.Context, - opts skills.InstallOptions, + opts *skills.InstallOptions, scope skills.Scope, ref nameref.Reference, ) (*skills.InstallResult, error) { @@ -104,6 +104,7 @@ func (s *service) installFromOCI( opts.LayerData = layerData opts.Reference = ociRef opts.Digest = pulledDigest.String() + if opts.Version == "" && skillConfig.Version != "" { opts.Version = skillConfig.Version } @@ -112,7 +113,20 @@ func (s *service) installFromOCI( unlock := s.locks.lock(opts.Name, scope, opts.ProjectRoot) defer unlock() - return s.installWithExtraction(ctx, opts, scope) + // Verify the artifact signature before anything is extracted or + // recorded; the decision (verified identity or explicit unsigned + // exception) travels on opts into the DB record and lock entry. This + // runs under the per-skill lock so concurrent first installs cannot + // both read an absent lock entry and race their TOFU anchors. + if shouldVerifyInstall(*opts, scope) { + decision, verifyErr := s.verifyOCIInstall(ctx, *opts, skillConfig.Name, ociRef, opts.Digest) + if verifyErr != nil { + return nil, verifyErr + } + applyDecisionToOpts(opts, decision) + } + + return s.installWithExtraction(ctx, *opts, scope) } // resolveFromLocalStore attempts to resolve a skill name against the local diff --git a/pkg/skills/skillsvc/lock.go b/pkg/skills/skillsvc/lock.go new file mode 100644 index 0000000000..7fe341c3ba --- /dev/null +++ b/pkg/skills/skillsvc/lock.go @@ -0,0 +1,210 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "errors" + "fmt" + "slices" + + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// errLockWrite marks failures to write the project lock file, so +// classifySyncFailure can map exactly these — and not every HTTP 500 — to +// the lock-write-failed reason automation keys on. +var errLockWrite = errors.New("lock file write failed") + +// recordLockState updates opts.ProjectRoot's lock file to reflect a +// just-completed project-scope install: an entry for sk, plus recursively +// materialized entries for any toolhive.requires dependencies declared in +// its SKILL.md. It also marks sk as lock-managed in the store. Callers must +// only invoke this for project-scope installs with the lock file feature +// enabled (see skills.LockFileFeatureEnabled) — sk is returned updated so the +// caller can reflect the Managed flag back to its own result. +func (s *service) recordLockState( + ctx context.Context, + opts skills.InstallOptions, + originalName string, + sk skills.InstalledSkill, +) (skills.InstalledSkill, error) { + contentDigest, err := computeContentDigest(s.pathResolver, sk) + if err != nil { + return sk, fmt.Errorf("computing content digest: %w", err) + } + + source := opts.LockSource + if source == "" { + source = originalName + } + resolvedReference := opts.LockResolvedReference + if resolvedReference == "" { + resolvedReference = sk.Reference + } + if err := recordLockEntry(sk.ProjectRoot, lockEntryInput{ + Name: sk.Metadata.Name, + Version: sk.Metadata.Version, + Source: source, + ResolvedReference: resolvedReference, + Digest: sk.Digest, + ContentDigest: contentDigest, + Provenance: provenanceInfoToLock(opts.Provenance), + Unsigned: opts.Unsigned, + RequiredByParent: opts.RequiredByParent, + PreserveExplicit: opts.SyncRestore, + }); err != nil { + return sk, fmt.Errorf("writing lock entry: %w", errors.Join(errLockWrite, err)) + } + + if !sk.Managed { + sk.Managed = true + if err := s.store.Update(ctx, sk); err != nil { + return sk, fmt.Errorf("marking skill as lock-managed: %w", err) + } + } + + // A sync restore reinstalls exactly what the lock file already pins: + // every dependency has its own entry that the sync loop restores + // directly, so re-materializing toolhive.requires here would re-resolve + // each dependency from its mutable source string — silently upgrading + // and re-pinning it, the opposite of a deterministic restore. + if !opts.SyncRestore { + if err := s.materializeDependencies(ctx, opts, source, sk); err != nil { + return sk, fmt.Errorf("materializing dependencies: %w", err) + } + } + return sk, nil +} + +// materializeDependencies installs every toolhive.requires dependency +// declared by sk's SKILL.md, recursively (a dependency may itself declare +// further dependencies). A Visited set threaded through opts prevents +// infinite recursion on a requires cycle; skills.MaxDependencies bounds the +// total number of skills materialized across the whole tree, not just the +// direct dependency list of a single skill. +// +// source is the reference sk was installed from (recordLockState's Source, +// after any LockSource override) — the same kind of string a dependency +// edge names in another skill's toolhive.requires. Visited must be keyed +// consistently by that reference form, not by sk's resolved Metadata.Name: +// a requires edge naming sk by a reference string other than its resolved +// name would otherwise bypass the cycle check and re-materialize sk as its +// own dependency, corrupting its RequiredBy list. +func (s *service) materializeDependencies( + ctx context.Context, + opts skills.InstallOptions, + source string, + sk skills.InstalledSkill, +) error { + parsed, err := readSkillMD(s.pathResolver, sk) + if err != nil { + return err + } + if len(parsed.Requires) == 0 { + return nil + } + + visited := opts.Visited + if visited == nil { + visited = make(map[string]struct{}) + } + visited[source] = struct{}{} + + for _, dep := range parsed.Requires { + if _, seen := visited[dep.Reference]; seen { + continue + } + if len(visited) >= skills.MaxDependencies { + return fmt.Errorf("dependency tree for %q exceeds maximum of %d skills", + sk.Metadata.Name, skills.MaxDependencies) + } + visited[dep.Reference] = struct{}{} + + depOpts := skills.InstallOptions{ + Name: dep.Reference, + Scope: sk.Scope, + ProjectRoot: sk.ProjectRoot, + Clients: sk.Clients, + Group: opts.Group, // deps join the parent's group, not the default one + RequiredByParent: sk.Metadata.Name, + Visited: visited, + } + if _, err := s.Install(ctx, depOpts); err != nil { + return fmt.Errorf("installing dependency %q (required by %q): %w", dep.Reference, sk.Metadata.Name, err) + } + } + return nil +} + +// lockEntryInput carries the fields recordLockEntry needs to upsert a lock +// entry, decoupled from skillsvc's own InstallOptions/InstalledSkill shapes. +type lockEntryInput struct { + Name string + Version string + Source string + ResolvedReference string + Digest string + ContentDigest string + // RequiredByParent names the parent skill when this entry is a + // transitively materialized dependency. Empty means the entry is + // explicit (a direct, user-requested install). + RequiredByParent string + // Provenance is the verified signer identity to record, nil for + // unsigned or unverified entries. + Provenance *lockfile.Provenance + // Unsigned records the explicit unsigned-install exception. + Unsigned bool + // PreserveExplicit keeps the existing entry's Explicit flag verbatim + // instead of deriving it from RequiredByParent. Set by sync restores: a + // restore is not a user install, so it must not promote a non-explicit + // dependency to explicit — that would permanently exempt it from + // cascade removal. + PreserveExplicit bool +} + +// recordLockEntry upserts a single entry into projectRoot's lock file. When +// an entry for the same name already exists, its RequiredBy list is merged +// (not overwritten) so a dependency shared by multiple parents keeps every +// parent, and Explicit is sticky once true. +func recordLockEntry(projectRoot string, in lockEntryInput) error { + root, err := lockfile.OpenRoot(projectRoot) + if err != nil { + return err + } + return lockfile.Update(root, func(lf *lockfile.Lockfile) error { + entry := lockfile.Entry{ + Name: in.Name, + Version: in.Version, + Source: in.Source, + ResolvedReference: in.ResolvedReference, + Digest: in.Digest, + ContentDigest: in.ContentDigest, + Provenance: in.Provenance, + Unsigned: in.Unsigned, + Explicit: in.RequiredByParent == "", + } + existing, exists := lf.Get(in.Name) + if exists { + entry.RequiredBy = existing.RequiredBy + entry.Explicit = entry.Explicit || existing.Explicit + } + if in.PreserveExplicit { + entry.Explicit = exists && existing.Explicit + } + if in.RequiredByParent != "" { + entry.RequiredBy = appendUnique(entry.RequiredBy, in.RequiredByParent) + } + lf.Upsert(entry) + return nil + }) +} + +func appendUnique(list []string, value string) []string { + if slices.Contains(list, value) { + return list + } + return append(list, value) +} diff --git a/pkg/skills/skillsvc/lock_test.go b/pkg/skills/skillsvc/lock_test.go new file mode 100644 index 0000000000..b859e3fd0c --- /dev/null +++ b/pkg/skills/skillsvc/lock_test.go @@ -0,0 +1,602 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "fmt" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/groups" + groupmocks "github.com/stacklok/toolhive/pkg/groups/mocks" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/gitresolver" + gitmocks "github.com/stacklok/toolhive/pkg/skills/gitresolver/mocks" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + skillsmocks "github.com/stacklok/toolhive/pkg/skills/mocks" + verifiermocks "github.com/stacklok/toolhive/pkg/skills/verifier/mocks" + "github.com/stacklok/toolhive/pkg/storage/sqlite" +) + +// newLockTestService builds a service backed by a real (temp-file) SQLite +// store and the real default Installer, so extracted files and lock file +// state can be inspected on disk exactly as they would be in production. +// Only the git resolver and path resolver are test doubles. +func newLockTestService(t *testing.T, gr *gitmocks.MockResolver, extra ...Option) (skills.SkillService, string) { + t.Helper() + t.Setenv(skills.LockFileEnvVar, "true") + + dbPath := filepath.Join(t.TempDir(), "test.db") + db, err := sqlite.Open(t.Context(), dbPath) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + store := sqlite.NewSkillStore(db) + + projectRoot := makeProjectRoot(t) + installBase := filepath.Join(projectRoot, ".claude", "skills") + + ctrl := gomock.NewController(t) + pr := skillsmocks.NewMockPathResolver(ctrl) + pr.EXPECT().GetSkillPath(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes(). + DoAndReturn(func(_, skillName string, _ skills.Scope, _ string) (string, error) { + return filepath.Join(installBase, skillName), nil + }) + // Only reached when a caller omits --clients with no prior DB record to + // fall back on (e.g. sync restoring an entry missing from a fresh + // clone) — the RFC's client-agnostic design expands that to every + // skill-supporting client detected on the host. Deliberately returns + // more than one client so tests can distinguish "preserved the skill's + // existing (narrower) client list" from "fell through to every detected + // client" (see TestUpgrade_PreservesExistingClients). + pr.EXPECT().ListSkillSupportingClients().AnyTimes().Return([]string{"claude-code", "cursor"}) + + // Default verifier: everything verifies as signed by a fixed test + // identity, so tests exercising lock mechanics don't trip install-time + // verification. Tests about verification pass their own WithVerifier + // via extra (later options win). + mv := verifiermocks.NewMockVerifier(ctrl) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(signedResult(), nil) + mv.EXPECT().VerifyOCI(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(signedResult(), nil) + mv.EXPECT().VerifyBundleOffline(gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(nil) + mv.EXPECT().ResultFromBundle(gomock.Any(), gomock.Any()). + AnyTimes().Return(signedResult(), nil) + + opts := append([]Option{WithPathResolver(pr), WithGitResolver(gr), WithVerifier(mv)}, extra...) + svc := New(store, opts...) + return svc, projectRoot +} + +// gitSkillVersion builds SKILL.md content for a git-resolved test fixture +// with a distinct version ("2.0.0"), for tests simulating newer content +// published at the same source (e.g. an upgrade scenario). It never +// declares requires. +func gitSkillVersion(name string) []byte { + fm := fmt.Sprintf("---\nname: %s\ndescription: test skill\nversion: 2.0.0\n---\n# %s\n", name, name) + return []byte(fm) +} + +// gitSkill builds SKILL.md content for a git-resolved test fixture, +// optionally declaring toolhive.requires dependencies. +func gitSkill(name string, requires ...string) []byte { + fm := fmt.Sprintf("---\nname: %s\ndescription: test skill\n", name) + if len(requires) > 0 { + fm += "toolhive.requires:\n" + for _, r := range requires { + fm += " - " + r + "\n" + } + } + fm += "---\n# " + name + "\n" + return []byte(fm) +} + +// gitRef builds the git:// reference and matching *gitresolver.GitReference +// for a fake repo named after the skill it resolves to. ParseGitReference +// picks "http://" instead of "https://" only when TOOLHIVE_DEV=true (see +// gitresolver.isDevMode); this mirrors that so the fixture URL matches +// whatever ParseGitReference actually produces in the environment the test +// runs in, rather than hardcoding one scheme. +func gitRef(name string) (string, string) { + scheme := "https" + if strings.EqualFold(os.Getenv("TOOLHIVE_DEV"), "true") { + scheme = "http" + } + return "git://github.com/test/" + name, scheme + "://github.com/test/" + name +} + +// gitFixture is a single registered skill: its SKILL.md content and the +// parsed name/version installFromGit reads from SkillConfig (it trusts the +// resolver's parse rather than re-parsing Files itself). +type gitFixture struct { + name string + content []byte +} + +// fakeGitFixtures dispatches a MockResolver's Resolve calls by matching the +// reference URL against a fixed set of registered skills, each getting its +// own commit hash so digests differ. Like a real repository, it keeps every +// registered revision: resolving a reference pinned to a full commit hash +// returns that revision's content, while an unpinned reference returns the +// latest — so sync's pinned restores behave the way real git checkouts do. +type fakeGitFixtures struct { + skills map[string]gitFixture // latest content, keyed by URL + history map[string]map[string]gitFixture // URL -> commit hash -> content +} + +// fixtureCommitHash derives the deterministic per-revision commit hash the +// fake resolver reports for a registered fixture. +func fixtureCommitHash(url string, content []byte) string { + return fmt.Sprintf("%040x", len(content)+len(url)) +} + +func newGitResolverMock(t *testing.T) (*gitmocks.MockResolver, *fakeGitFixtures) { + t.Helper() + fx := &fakeGitFixtures{ + skills: make(map[string]gitFixture), + history: make(map[string]map[string]gitFixture), + } + gr := gitmocks.NewMockResolver(gomock.NewController(t)) + gr.EXPECT().Resolve(gomock.Any(), gomock.Any()).AnyTimes(). + DoAndReturn(func(_ context.Context, ref *gitresolver.GitReference) (*gitresolver.ResolveResult, error) { + fixture, ok := fx.skills[ref.URL] + if !ok { + return nil, fmt.Errorf("no fixture registered for %q", ref.URL) + } + if pinned, isPinned := fx.history[ref.URL][ref.Ref]; isPinned { + fixture = pinned + } + return &gitresolver.ResolveResult{ + SkillConfig: &skills.ParseResult{Name: fixture.name}, + Files: []gitresolver.FileEntry{{Path: "SKILL.md", Content: fixture.content, Mode: 0644}}, + CommitHash: fixtureCommitHash(ref.URL, fixture.content), + }, nil + }) + return gr, fx +} + +func (f *fakeGitFixtures) register(name string, content []byte) { + _, url := gitRef(name) + fixture := gitFixture{name: name, content: content} + f.skills[url] = fixture + if f.history[url] == nil { + f.history[url] = make(map[string]gitFixture) + } + f.history[url][fixtureCommitHash(url, content)] = fixture +} + +func mustOpenRoot(t *testing.T, projectRoot string) lockfile.Root { + t.Helper() + root, err := lockfile.OpenRoot(projectRoot) + require.NoError(t, err) + return root +} + +func readLockfile(t *testing.T, projectRoot string) *lockfile.Lockfile { + t.Helper() + lf, err := lockfile.Load(mustOpenRoot(t, projectRoot)) + require.NoError(t, err) + return lf +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_LockFileDisabled_NoLockFileWritten(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + t.Setenv(skills.LockFileEnvVar, "false") // override newLockTestService's default + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + _, err = os.Stat(filepath.Join(projectRoot, lockfile.FileName)) + assert.True(t, os.IsNotExist(err), "lock file must not be written when the feature is disabled") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_RecordsExplicitEntry(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + result, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + assert.True(t, result.Skill.Managed, "project-scope install must be marked lock-managed") + + lf := readLockfile(t, projectRoot) + entry, ok := lf.Get("my-skill") + require.True(t, ok, "expected a lock entry for my-skill") + assert.Equal(t, ref, entry.Source, "source must be exactly what the caller requested") + // For a direct git install, ResolvedReference is the same git:// URL: the + // package preserves it verbatim as the artifact's provenance reference, + // distinct from a registry-resolved install where Source (registry name) + // and ResolvedReference (the concrete git:// URL it resolved to) differ. + assert.Equal(t, ref, entry.ResolvedReference) + assert.NotEmpty(t, entry.Digest) + assert.NotEmpty(t, entry.ContentDigest, "contentDigest must be computed for a project-scope install") + assert.True(t, entry.Explicit) + assert.Empty(t, entry.RequiredBy) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_MaterializesDependency(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("shared-dep") + fx.register("parent-skill", gitSkill("parent-skill", depRef)) + fx.register("shared-dep", gitSkill("shared-dep")) + svc, projectRoot := newLockTestService(t, gr) + + parentRef, _ := gitRef("parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: parentRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + lf := readLockfile(t, projectRoot) + parent, ok := lf.Get("parent-skill") + require.True(t, ok) + assert.True(t, parent.Explicit) + + dep, ok := lf.Get("shared-dep") + require.True(t, ok, "dependency must be materialized and recorded") + assert.False(t, dep.Explicit) + assert.Equal(t, []string{"parent-skill"}, dep.RequiredBy) + assert.NotEmpty(t, dep.ContentDigest) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_SharedDependencyMergesRequiredBy(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("shared-dep") + fx.register("parent-a", gitSkill("parent-a", depRef)) + fx.register("parent-b", gitSkill("parent-b", depRef)) + fx.register("shared-dep", gitSkill("shared-dep")) + svc, projectRoot := newLockTestService(t, gr) + + refA, _ := gitRef("parent-a") + refB, _ := gitRef("parent-b") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: refA, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: refB, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + lf := readLockfile(t, projectRoot) + dep, ok := lf.Get("shared-dep") + require.True(t, ok) + assert.ElementsMatch(t, []string{"parent-a", "parent-b"}, dep.RequiredBy, + "a dependency shared by two parents must keep both, not overwrite on the second install") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestMaterializeDependencies_RejectsTreeExceedingMaxDependencies(t *testing.T) { + gr, _ := newGitResolverMock(t) + svc, projectRoot := newLockTestService(t, gr) + svcImpl := svc.(*service) //nolint:forcetypeassert // white-box test in the same package + + depRef, _ := gitRef("one-more-dep") + sk := skills.InstalledSkill{ + Metadata: skills.SkillMetadata{Name: "root-skill"}, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + } + skillDir := filepath.Join(projectRoot, ".claude", "skills", sk.Metadata.Name) + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), gitSkill("root-skill", depRef), 0o644)) + + // Pre-fill Visited to the cap, as if this were deep in an already-large + // dependency tree; the next never-seen dependency must trip the limit + // rather than silently keep expanding. + visited := make(map[string]struct{}, skills.MaxDependencies) + for i := range skills.MaxDependencies { + visited[fmt.Sprintf("seen-%d", i)] = struct{}{} + } + + err := svcImpl.materializeDependencies(t.Context(), skills.InstallOptions{Visited: visited}, "root-skill", sk) + require.Error(t, err) + assert.Contains(t, err.Error(), "exceeds maximum") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_DependencyCycleTerminates(t *testing.T) { + gr, fx := newGitResolverMock(t) + refA, _ := gitRef("skill-a") + refB, _ := gitRef("skill-b") + fx.register("skill-a", gitSkill("skill-a", refB)) + fx.register("skill-b", gitSkill("skill-b", refA)) // cycle: a -> b -> a + svc, projectRoot := newLockTestService(t, gr) + + done := make(chan error, 1) + go func() { + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: refA, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + done <- err + }() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("timeout: dependency cycle did not terminate") + } + + lf := readLockfile(t, projectRoot) + a, ok := lf.Get("skill-a") + require.True(t, ok) + assert.True(t, a.Explicit) + // The cycle check must trip on skill-b's requires entry for refA — the + // same git:// reference skill-a was installed with, not skill-a's + // resolved bare name — so skill-a is never re-materialized as its own + // dependency. Without that, skill-a would gain itself a spurious + // RequiredBy via skill-b. + assert.Empty(t, a.RequiredBy, "skill-a must not become required-by skill-b through the a<->b cycle") + b, ok := lf.Get("skill-b") + require.True(t, ok) + assert.Equal(t, []string{"skill-a"}, b.RequiredBy) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_LockWriteFailureRollsBackInstall(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + // Replace the lock file location with a directory so writes to it fail. + require.NoError(t, os.MkdirAll(filepath.Join(projectRoot, lockfile.FileName), 0o755)) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.Error(t, err, "install must fail when the lock file cannot be written") + assert.Equal(t, http.StatusInternalServerError, httperr.Code(err), + "a code-less lock-write failure still defaults to 500") + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.Error(t, err, "the DB record must be rolled back so a retry starts fresh") +} + +// TestInstallProjectScope_DependencyFailureRollsBackParentLockEntry covers a +// failure that happens after the parent's own lock entry was already +// written: recordLockEntry succeeds for the parent, then +// materializeDependencies fails on an unresolvable toolhive.requires +// dependency. The parent's lock entry (and DB record) must be rolled back +// too, not just the DB record — otherwise toolhive.lock.yaml claims the +// parent is pinned while Info() reports it doesn't exist. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_DependencyFailureRollsBackParentLockEntry(t *testing.T) { + gr, fx := newGitResolverMock(t) + missingDepRef, _ := gitRef("missing-dep") // never registered — Resolve fails for it + fx.register("parent-skill", gitSkill("parent-skill", missingDepRef)) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.Error(t, err, "install must fail when a declared dependency cannot be resolved") + + lf := readLockfile(t, projectRoot) + _, ok := lf.Get("parent-skill") + assert.False(t, ok, "parent's lock entry must be rolled back when a dependency fails to materialize") + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "parent-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.Error(t, err, "the DB record must be rolled back so a retry starts fresh") + + // Retrying (with Force, since the first attempt's extracted files are + // left on disk per this rollback's documented contract) must succeed + // once the dependency is fixed, not hit a stale, permanently-broken state. + fx.register("missing-dep", gitSkill("missing-dep")) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, Force: true, + }) + require.NoError(t, err, "retry after fixing the dependency must succeed") +} + +// TestInstallProjectScope_SiblingDependencyRolledBackWithParent covers the +// partial-success case: the first declared dependency installs cleanly, then +// a later sibling fails. Rolling back only the parent would leak the +// succeeded sibling as an orphan — installed, managed, and pinned with a +// RequiredBy pointing at a parent that no longer exists. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_SiblingDependencyRolledBackWithParent(t *testing.T) { + gr, fx := newGitResolverMock(t) + okDepRef, _ := gitRef("ok-dep") + missingDepRef, _ := gitRef("missing-dep") // never registered — Resolve fails + fx.register("ok-dep", gitSkill("ok-dep")) + fx.register("parent-skill", gitSkill("parent-skill", okDepRef, missingDepRef)) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.Error(t, err, "install must fail when a later dependency cannot be resolved") + + lf := readLockfile(t, projectRoot) + _, ok := lf.Get("parent-skill") + assert.False(t, ok, "parent's lock entry must be rolled back") + _, ok = lf.Get("ok-dep") + assert.False(t, ok, "the succeeded sibling's lock entry must be rolled back too, not leaked as an orphan") + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "ok-dep", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.Error(t, err, "the succeeded sibling's DB record must be rolled back too") +} + +// TestInstallProjectScope_RollbackRestoresPreExistingState covers the +// destructive-rollback hazard: a --force reinstall of an already-installed, +// lock-managed skill that fails partway (here: on a newly-declared +// dependency) must restore the previously-valid DB record and lock entry — +// including RequiredBy links other parents merged onto it — not delete them. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_RollbackRestoresPreExistingState(t *testing.T) { + gr, fx := newGitResolverMock(t) + sharedRef, _ := gitRef("shared-skill") + fx.register("shared-skill", gitSkill("shared-skill")) + fx.register("other-parent", gitSkill("other-parent", sharedRef)) + svc, projectRoot := newLockTestService(t, gr) + + // shared-skill is installed both explicitly and as other-parent's + // dependency, so its entry carries Explicit=true and RequiredBy. + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: sharedRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + otherRef, _ := gitRef("other-parent") + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: otherRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + before, ok := readLockfile(t, projectRoot).Get("shared-skill") + require.True(t, ok) + require.Equal(t, []string{"other-parent"}, before.RequiredBy) + + // Republish shared-skill with a new, unresolvable dependency and force- + // reinstall it: the install fails after the DB record and lock entry + // were already rewritten. + missingDepRef, _ := gitRef("missing-dep") + fx.register("shared-skill", gitSkill("shared-skill", missingDepRef)) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: sharedRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, Force: true, + }) + require.Error(t, err, "reinstall must fail on the unresolvable dependency") + + after, ok := readLockfile(t, projectRoot).Get("shared-skill") + require.True(t, ok, "a transient failure must not destroy the pre-existing lock entry") + assert.Equal(t, before.Digest, after.Digest, "the previous pin must be restored") + assert.Equal(t, before.RequiredBy, after.RequiredBy, "RequiredBy links from other parents must survive") + assert.True(t, after.Explicit) + + info, err := svc.Info(t.Context(), skills.InfoOptions{Name: "shared-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err, "the pre-existing DB record must survive") + assert.Equal(t, before.Digest, info.InstalledSkill.Digest, "the DB record must be restored to its previous state") +} + +// TestInstallProjectScope_DependencyJoinsParentGroup asserts a transitively +// materialized dependency is registered in the same group the parent was +// installed into, rather than silently falling back to the default group. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_DependencyJoinsParentGroup(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("grouped-dep") + fx.register("grouped-dep", gitSkill("grouped-dep")) + fx.register("grouped-parent", gitSkill("grouped-parent", depRef)) + svc, projectRoot := newLockTestService(t, gr) + + // Track which group each skill lands in via a stub manager. + skillGroups := make(map[string]string) + gm := groupmocks.NewMockManager(gomock.NewController(t)) + gm.EXPECT().Get(gomock.Any(), gomock.Any()).AnyTimes(). + DoAndReturn(func(_ context.Context, name string) (*groups.Group, error) { + return &groups.Group{Name: name}, nil + }) + gm.EXPECT().Update(gomock.Any(), gomock.Any()).AnyTimes(). + DoAndReturn(func(_ context.Context, g *groups.Group) error { + for _, sk := range g.Skills { + skillGroups[sk] = g.Name + } + return nil + }) + svcImpl := svc.(*service) //nolint:forcetypeassert // white-box test in the same package + svcImpl.groupManager = gm + + ref, _ := gitRef("grouped-parent") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, Group: "custom-group", + }) + require.NoError(t, err) + + assert.Equal(t, "custom-group", skillGroups["grouped-parent"]) + assert.Equal(t, "custom-group", skillGroups["grouped-dep"], + "a materialized dependency must join the parent's group, not the default one") +} + +// TestUninstall_LockWriteFailureAbortsBeforeDestruction guards the uninstall +// ordering: the lock entry is removed first, and if that fails the +// uninstall aborts with everything intact. The reverse order left a stale +// lock entry for a gone skill, which the next sync silently reinstalled. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUninstall_LockWriteFailureAbortsBeforeDestruction(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // Make the lock file unwritable: read-only project root blocks the + // temp-file write inside lockfile.Update. + require.NoError(t, os.Chmod(projectRoot, 0o555)) + t.Cleanup(func() { _ = os.Chmod(projectRoot, 0o755) }) + + err = svc.Uninstall(t.Context(), skills.UninstallOptions{ + Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot, + }) + require.Error(t, err, "uninstall must fail when the lock entry cannot be removed") + + require.NoError(t, os.Chmod(projectRoot, 0o755)) + info, err := svc.Info(t.Context(), skills.InfoOptions{Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err, "the skill must remain fully installed — nothing may be destroyed before the lock update") + assert.NotNil(t, info.InstalledSkill) + _, ok := readLockfile(t, projectRoot).Get("my-skill") + assert.True(t, ok, "the lock entry must be untouched") + + skillMD := filepath.Join(projectRoot, ".claude", "skills", "my-skill", "SKILL.md") + _, statErr := os.Stat(skillMD) + require.NoError(t, statErr, "on-disk files must be untouched") +} + +// TestInstallProjectScope_DependencyFailureCodePreserved: a dependency's +// specific failure code must reach the API boundary as itself. Dependency +// materialization runs inside recordLockState, whose failures used to be +// blanket-wrapped as 500 — masking a git resolver's 502 (and turning the +// typed sync failure reason into "unknown" instead of +// "registry-unreachable"). +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallProjectScope_DependencyFailureCodePreserved(t *testing.T) { + gr, fx := newGitResolverMock(t) + missingDepRef, _ := gitRef("missing-dep") // never registered — Resolve fails with 502 + fx.register("parent-skill", gitSkill("parent-skill", missingDepRef)) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.Error(t, err) + assert.Equal(t, http.StatusBadGateway, httperr.Code(err), + "the dependency's 502 must not be masked to 500 by the lock-state wrap") +} diff --git a/pkg/skills/skillsvc/pin.go b/pkg/skills/skillsvc/pin.go new file mode 100644 index 0000000000..73b627fb0f --- /dev/null +++ b/pkg/skills/skillsvc/pin.go @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "fmt" + "strings" + + nameref "github.com/google/go-containerregistry/pkg/name" + + "github.com/stacklok/toolhive/pkg/skills/gitresolver" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// isImmutableSource reports whether a lock entry's source can never produce +// newer content: an OCI digest reference, or a git reference already pinned +// to a full commit hash. Upgrade reports these as not-upgradable rather than +// attempting to re-resolve them. +func isImmutableSource(entry lockfile.Entry) bool { + if gitresolver.IsGitReference(entry.Source) { + ref, err := gitresolver.ParseGitReference(entry.Source) + return err == nil && isFullCommitHash(ref.Ref) + } + ref, err := nameref.ParseReference(entry.Source) + if err != nil { + return false + } + _, isDigest := ref.(nameref.Digest) + return isDigest +} + +// repositoryMoved reports whether two resolved references point at different +// repositories, as opposed to differing only by tag. A tag move within one +// repository is the normal way a mutable source advances; a repository move +// means the artifact now comes from somewhere else, which is the case the +// ref-change guard exists to catch. +// +// Both the registry and the repository path are compared, so ghcr.io -> +// another registry, or a different org or path on the same registry, all +// count as a move. +func repositoryMoved(oldRef, newRef string) bool { + if oldRef == newRef { + return false + } + oldRepo, ok := referenceRepository(oldRef) + if !ok { + return true + } + newRepo, ok := referenceRepository(newRef) + if !ok { + return true + } + return oldRepo != newRepo +} + +// referenceRepository returns the registry and repository portion of an OCI +// reference, dropping any tag or digest. It reports false for git references +// and for anything it cannot parse, so callers fall back to treating the +// reference as moved rather than silently equating two unlike sources. +func referenceRepository(ref string) (string, bool) { + if gitresolver.IsGitReference(ref) { + return "", false + } + parsed, err := nameref.ParseReference(ref) + if err != nil { + return "", false + } + return parsed.Context().String(), true +} + +// isFullCommitHash accepts both hex cases: the sibling git resolver does +// too, and classifying an uppercase-pinned source as mutable would +// needlessly re-clone it on every upgrade despite the pin being immutable. +func isFullCommitHash(ref string) bool { + if len(ref) != 40 { + return false + } + for _, c := range ref { + if (c < '0' || c > '9') && (c < 'a' || c > 'f') && (c < 'A' || c > 'F') { + return false + } + } + return true +} + +// buildPinnedReference returns the exact reference sync must install: entry's +// resolvedReference re-pointed at its pinned digest, never re-resolved from +// source. This is what makes sync a restore operation rather than an upgrade +// — installing this reference always yields entry's exact pinned content. +func buildPinnedReference(entry lockfile.Entry) (string, error) { + if gitresolver.IsGitReference(entry.ResolvedReference) { + return pinGitReference(entry) + } + return pinOCIReference(entry) +} + +func pinOCIReference(entry lockfile.Entry) (string, error) { + ref, err := nameref.ParseReference(entry.ResolvedReference) + if err != nil { + return "", fmt.Errorf("parsing resolvedReference %q: %w", entry.ResolvedReference, err) + } + return ref.Context().String() + "@" + entry.Digest, nil +} + +func pinGitReference(entry lockfile.Entry) (string, error) { + gitRef, err := gitresolver.ParseGitReference(entry.ResolvedReference) + if err != nil { + return "", fmt.Errorf("parsing resolvedReference %q: %w", entry.ResolvedReference, err) + } + hostAndPath := strings.TrimPrefix(strings.TrimPrefix(gitRef.URL, "https://"), "http://") + pinned := "git://" + hostAndPath + "@" + entry.Digest + if gitRef.Path != "" { + pinned += "#" + gitRef.Path + } + return pinned, nil +} diff --git a/pkg/skills/skillsvc/pin_test.go b/pkg/skills/skillsvc/pin_test.go new file mode 100644 index 0000000000..9cd100508b --- /dev/null +++ b/pkg/skills/skillsvc/pin_test.go @@ -0,0 +1,194 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +func TestBuildPinnedReference(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + entry lockfile.Entry + want string + }{ + { + name: "OCI reference pins to digest", + entry: lockfile.Entry{ + ResolvedReference: "ghcr.io/org/code-review:1.0.0", + Digest: "sha256:" + hexDigestForTest(), + }, + want: "ghcr.io/org/code-review@sha256:" + hexDigestForTest(), + }, + { + name: "git reference pins to commit hash, dropping any tag/branch ref", + entry: lockfile.Entry{ + ResolvedReference: "git://github.com/org/skills@main#testing-conventions", + Digest: testCommitHash, + }, + want: "git://github.com/org/skills@" + testCommitHash + "#testing-conventions", + }, + { + name: "git reference without a subdir", + entry: lockfile.Entry{ + ResolvedReference: "git://github.com/org/skills", + Digest: testCommitHash, + }, + want: "git://github.com/org/skills@" + testCommitHash, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := buildPinnedReference(tt.entry) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestBuildPinnedReferenceRejectsUnparsable(t *testing.T) { + t.Parallel() + _, err := buildPinnedReference(lockfile.Entry{ResolvedReference: "not a valid reference!!", Digest: "sha256:abc"}) + require.Error(t, err) +} + +func TestIsImmutableSource(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + entry lockfile.Entry + want bool + }{ + { + name: "OCI digest source is immutable", + entry: lockfile.Entry{Source: "ghcr.io/org/skill@sha256:" + hexDigestForTest()}, + want: true, + }, + { + name: "OCI tag source is mutable", + entry: lockfile.Entry{Source: "ghcr.io/org/skill:1.0.0"}, + want: false, + }, + { + name: "git full commit hash source is immutable", + entry: lockfile.Entry{Source: "git://github.com/org/skill@" + testCommitHash}, + want: true, + }, + { + name: "git branch source is mutable", + entry: lockfile.Entry{Source: "git://github.com/org/skill@main"}, + want: false, + }, + { + // The sibling git resolver accepts uppercase hex; classifying + // an uppercase pin as mutable would needlessly re-clone it on + // every upgrade despite the pin being immutable. + name: "git uppercase full commit hash source is immutable", + entry: lockfile.Entry{Source: "git://github.com/org/skill@ABCDEF0123456789ABCDEF0123456789ABCDEF01"}, + want: true, + }, + { + name: "git source with no ref is mutable", + entry: lockfile.Entry{Source: "git://github.com/org/skill"}, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, isImmutableSource(tt.entry)) + }) + } +} + +// hexDigestForTest returns a fixed, valid 64-char hex string for OCI digest fixtures. +func hexDigestForTest() string { + const s = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789" + return s[:64] +} + +func TestRepositoryMoved(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + oldRef string + newRef string + want bool + }{ + { + name: "identical references have not moved", + oldRef: "ghcr.io/org/skill:v1", + newRef: "ghcr.io/org/skill:v1", + }, + { + name: "a version bump is not a move", + oldRef: "ghcr.io/org/skill:0.1.0", + newRef: "ghcr.io/org/skill:0.2.0", + }, + { + name: "moving to a digest in the same repository is not a move", + oldRef: "ghcr.io/org/skill:0.1.0", + newRef: "ghcr.io/org/skill@" + ociTestDigest(1), + }, + { + name: "an implicit latest tag matches an explicit one", + oldRef: "ghcr.io/org/skill", + newRef: "ghcr.io/org/skill:latest", + }, + { + name: "a different repository path is a move", + oldRef: "ghcr.io/org/skill:v1", + newRef: "ghcr.io/org/other-skill:v1", + want: true, + }, + { + name: "a different org on the same registry is a move", + oldRef: "ghcr.io/org/skill:v1", + newRef: "ghcr.io/attacker/skill:v1", + want: true, + }, + { + name: "a different registry is a move", + oldRef: "ghcr.io/org/skill:v1", + newRef: "elsewhere.io/org/skill:v1", + want: true, + }, + { + name: "git references fall back to exact comparison", + oldRef: "git://github.com/org/repo#skills/a", + newRef: "git://github.com/org/repo#skills/b", + want: true, + }, + { + name: "an OCI reference replaced by a git one is a move", + oldRef: "ghcr.io/org/skill:v1", + newRef: "git://github.com/org/repo", + want: true, + }, + { + name: "unparsable input fails closed", + oldRef: "ghcr.io/org/skill:v1", + newRef: "not a valid reference at all", + want: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.want, repositoryMoved(tc.oldRef, tc.newRef)) + }) + } +} diff --git a/pkg/skills/skillsvc/service.go b/pkg/skills/skillsvc/service.go index ef87558eab..656adf4644 100644 --- a/pkg/skills/skillsvc/service.go +++ b/pkg/skills/skillsvc/service.go @@ -12,6 +12,7 @@ import ( "github.com/stacklok/toolhive/pkg/groups" "github.com/stacklok/toolhive/pkg/skills" "github.com/stacklok/toolhive/pkg/skills/gitresolver" + "github.com/stacklok/toolhive/pkg/skills/verifier" "github.com/stacklok/toolhive/pkg/storage" ) @@ -118,6 +119,16 @@ type service struct { registry ociskills.RegistryClient skillLookup SkillLookup gitResolver gitresolver.Resolver + sigVerifier verifier.Verifier +} + +// WithVerifier sets the signature verifier used for install-time +// verification. Defaults to the Sigstore verifier with the composite +// registry keychain. +func WithVerifier(v verifier.Verifier) Option { + return func(s *service) { + s.sigVerifier = v + } } // New creates a new SkillService backed by the given store. diff --git a/pkg/skills/skillsvc/source_dispatch.go b/pkg/skills/skillsvc/source_dispatch.go new file mode 100644 index 0000000000..27a56fb6e9 --- /dev/null +++ b/pkg/skills/skillsvc/source_dispatch.go @@ -0,0 +1,101 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "fmt" + "net/http" + + nameref "github.com/google/go-containerregistry/pkg/name" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/gitresolver" +) + +// sourceOps are the per-kind leaf operations dispatchSource routes a skill +// source to. +type sourceOps[T any] struct { + // git handles a git:// reference. + git func(ctx context.Context, gitURL string) (T, error) + // oci handles a direct OCI reference. + oci func(ctx context.Context, ref nameref.Reference) (T, error) + // registry handles a source resolved through the registry catalogue + // (either the OCI ambiguity fallback or a plain-name lookup). + registry func(ctx context.Context, resolved *registryResolveResult) (T, error) + // plainName, when non-nil, handles a validated bare skill name instead + // of the default registry lookup — for callers with a richer flow + // (Install checks the local OCI store and takes a per-name lock). + plainName func(ctx context.Context, name string) (T, error) +} + +// dispatchSource is the single source of truth for how a skill source +// string is routed: git prefix first (unambiguous), then OCI-shaped +// references — with a registry-catalogue fallback when a direct OCI +// operation fails for an *ambiguous* "namespace/name" (see +// isUnambiguousOCIRef) — and finally validated plain skill names via the +// registry. Install and upgrade's re-resolution previously each maintained +// their own copy of this sequence, and the two drifted (the upgrade copy +// lost the fallback branch entirely); every consumer must go through this +// function so they cannot diverge again. +// +// Fallback error precedence: when the direct OCI operation fails and the +// registry lookup itself errors, the registry error wins (it is the more +// actionable signal); when the registry simply has no match, the original +// OCI error is returned. +func dispatchSource[T any](ctx context.Context, s *service, source string, ops sourceOps[T]) (T, error) { + var zero T + + if gitresolver.IsGitReference(source) { + return ops.git(ctx, source) + } + + ref, isOCI, err := parseOCIReference(source) + if err != nil { + return zero, httperr.WithCode( + fmt.Errorf("invalid OCI reference %q: %w", source, err), + http.StatusBadRequest, + ) + } + if isOCI { + result, ociErr := ops.oci(ctx, ref) + if ociErr == nil { + return result, nil + } + // Direct OCI failed — an ambiguous "namespace/name" may be a + // registry catalogue name. Names that are unambiguously OCI + // (digest, explicit tag, or multi-segment path) must not trigger a + // registry search. + if isUnambiguousOCIRef(source, ref) { + return zero, ociErr + } + resolved, regErr := s.resolveFromRegistry(source) + if regErr != nil { + return zero, regErr + } + if resolved == nil { + return zero, ociErr + } + return ops.registry(ctx, resolved) + } + + if err := skills.ValidateSkillName(source); err != nil { + return zero, httperr.WithCode(err, http.StatusBadRequest) + } + if ops.plainName != nil { + return ops.plainName(ctx, source) + } + resolved, regErr := s.resolveFromRegistry(source) + if regErr != nil { + return zero, regErr + } + if resolved == nil { + return zero, httperr.WithCode( + fmt.Errorf("skill %q not found in registry", source), + http.StatusNotFound, + ) + } + return ops.registry(ctx, resolved) +} diff --git a/pkg/skills/skillsvc/source_dispatch_test.go b/pkg/skills/skillsvc/source_dispatch_test.go new file mode 100644 index 0000000000..c8f6e3773a --- /dev/null +++ b/pkg/skills/skillsvc/source_dispatch_test.go @@ -0,0 +1,171 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "errors" + "net/http" + "testing" + + nameref "github.com/google/go-containerregistry/pkg/name" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + regtypes "github.com/stacklok/toolhive-core/registry/types" + regmocks "github.com/stacklok/toolhive/pkg/registry/mocks" + storemocks "github.com/stacklok/toolhive/pkg/storage/mocks" +) + +// dispatchProbe records which leaf dispatchSource routed to. +type dispatchProbe struct { + gitURL string + ociRef string + registry *registryResolveResult + plainName string +} + +// probeOps returns ops whose leaves record the routing into probe and +// return the given leaf results. +func probeOps(probe *dispatchProbe, ociErr error) sourceOps[string] { + return sourceOps[string]{ + git: func(_ context.Context, gitURL string) (string, error) { + probe.gitURL = gitURL + return "git", nil + }, + oci: func(_ context.Context, ref nameref.Reference) (string, error) { + probe.ociRef = ref.String() + if ociErr != nil { + return "", ociErr + } + return "oci", nil + }, + registry: func(_ context.Context, resolved *registryResolveResult) (string, error) { + probe.registry = resolved + return "registry", nil + }, + plainName: func(_ context.Context, name string) (string, error) { + probe.plainName = name + return "plain", nil + }, + } +} + +func newDispatchTestService(t *testing.T, lookup *regmocks.MockProvider) *service { + t.Helper() + opts := []Option{} + if lookup != nil { + opts = append(opts, WithSkillLookup(lookup)) + } + svc := New(storemocks.NewMockSkillStore(gomock.NewController(t)), opts...) + return svc.(*service) //nolint:forcetypeassert // white-box test in the same package +} + +func TestDispatchSourceRouting(t *testing.T) { + t.Parallel() + + t.Run("git reference routes to git", func(t *testing.T) { + t.Parallel() + svc := newDispatchTestService(t, nil) + var probe dispatchProbe + got, err := dispatchSource(t.Context(), svc, "git://github.com/org/skill@main", probeOps(&probe, nil)) + require.NoError(t, err) + assert.Equal(t, "git", got) + assert.Equal(t, "git://github.com/org/skill@main", probe.gitURL) + }) + + t.Run("OCI reference routes to oci", func(t *testing.T) { + t.Parallel() + svc := newDispatchTestService(t, nil) + var probe dispatchProbe + got, err := dispatchSource(t.Context(), svc, "ghcr.io/org/skill:v1", probeOps(&probe, nil)) + require.NoError(t, err) + assert.Equal(t, "oci", got) + assert.Equal(t, "ghcr.io/org/skill:v1", probe.ociRef) + }) + + t.Run("plain name routes to plainName leaf when set", func(t *testing.T) { + t.Parallel() + svc := newDispatchTestService(t, nil) + var probe dispatchProbe + got, err := dispatchSource(t.Context(), svc, "my-skill", probeOps(&probe, nil)) + require.NoError(t, err) + assert.Equal(t, "plain", got) + assert.Equal(t, "my-skill", probe.plainName) + }) + + t.Run("plain name without plainName leaf resolves via registry", func(t *testing.T) { + t.Parallel() + lookup := regmocks.NewMockProvider(gomock.NewController(t)) + lookup.EXPECT().SearchSkills("my-skill").Return([]regtypes.Skill{ + {Name: "my-skill", Packages: []regtypes.SkillPackage{ + {RegistryType: "oci", Identifier: "ghcr.io/test/my-skill:v1"}, + }}, + }, nil) + svc := newDispatchTestService(t, lookup) + var probe dispatchProbe + ops := probeOps(&probe, nil) + ops.plainName = nil + got, err := dispatchSource(t.Context(), svc, "my-skill", ops) + require.NoError(t, err) + assert.Equal(t, "registry", got) + require.NotNil(t, probe.registry) + }) + + t.Run("invalid plain name is rejected", func(t *testing.T) { + t.Parallel() + svc := newDispatchTestService(t, nil) + var probe dispatchProbe + _, err := dispatchSource(t.Context(), svc, "Not_A_Name", probeOps(&probe, nil)) + require.Error(t, err) + assert.Equal(t, http.StatusBadRequest, httperr.Code(err)) + }) +} + +// TestDispatchSourceOCIFallback pins the ambiguity-fallback rules whose +// hand-copied divergence caused the upgrade-path majors: a failed direct +// OCI operation falls back to the registry catalogue only for ambiguous +// names, and the error precedence is fixed (registry error wins over the +// OCI error; a registry miss returns the original OCI error). +func TestDispatchSourceOCIFallback(t *testing.T) { + t.Parallel() + ociFailure := errors.New("name unknown") + + t.Run("ambiguous name falls back to registry on OCI failure", func(t *testing.T) { + t.Parallel() + lookup := regmocks.NewMockProvider(gomock.NewController(t)) + lookup.EXPECT().SearchSkills("my-skill").Return([]regtypes.Skill{ + {Namespace: "io.github.test", Name: "my-skill", Packages: []regtypes.SkillPackage{ + {RegistryType: "oci", Identifier: "ghcr.io/test/my-skill:v2"}, + }}, + }, nil) + svc := newDispatchTestService(t, lookup) + var probe dispatchProbe + got, err := dispatchSource(t.Context(), svc, "io.github.test/my-skill", probeOps(&probe, ociFailure)) + require.NoError(t, err) + assert.Equal(t, "registry", got) + require.NotNil(t, probe.registry) + assert.NotNil(t, probe.registry.OCIRef) + }) + + t.Run("unambiguous reference never falls back", func(t *testing.T) { + t.Parallel() + svc := newDispatchTestService(t, nil) // no lookup: a fallback attempt would nil-panic expectations + var probe dispatchProbe + _, err := dispatchSource(t.Context(), svc, "ghcr.io/org/skill:v1", probeOps(&probe, ociFailure)) + require.ErrorIs(t, err, ociFailure) + }) + + t.Run("registry miss returns the original OCI error", func(t *testing.T) { + t.Parallel() + lookup := regmocks.NewMockProvider(gomock.NewController(t)) + lookup.EXPECT().SearchSkills("my-skill").Return(nil, nil) + svc := newDispatchTestService(t, lookup) + var probe dispatchProbe + _, err := dispatchSource(t.Context(), svc, "io.github.test/my-skill", probeOps(&probe, ociFailure)) + require.ErrorIs(t, err, ociFailure) + }) +} diff --git a/pkg/skills/skillsvc/sync.go b/pkg/skills/skillsvc/sync.go new file mode 100644 index 0000000000..4e0b2ce867 --- /dev/null +++ b/pkg/skills/skillsvc/sync.go @@ -0,0 +1,308 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/http" + "strings" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/verifier" + "github.com/stacklok/toolhive/pkg/storage" +) + +// Sync restores a project's installed skills to match its lock file: missing +// or drifted entries are reinstalled at their pinned digest (never +// re-resolved from source — see buildPinnedReference), unmanaged installs are +// reported (or adopted with Adopt), and lock-managed installs no longer in +// the lock file are reported (or removed with Prune). Check performs the +// same reconciliation read-only: nothing is installed, written, or removed. +func (s *service) Sync(ctx context.Context, opts skills.SyncOptions) (*skills.SyncResult, error) { + if !skills.LockFileFeatureEnabled() { + return nil, errExperimentalLockFeature + } + + _, projectRoot, err := normalizeProjectRoot(skills.ScopeProject, opts.ProjectRoot) + if err != nil { + return nil, err + } + opts.ProjectRoot = projectRoot + + root, err := lockfile.OpenRoot(projectRoot) + if err != nil { + return nil, err + } + lf, err := lockfile.Load(root) + if err != nil { + return nil, err + } + + installed, err := s.store.List(ctx, storage.ListFilter{Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + if err != nil { + return nil, fmt.Errorf("listing installed skills: %w", err) + } + installedByName := make(map[string]skills.InstalledSkill, len(installed)) + for _, sk := range installed { + installedByName[sk.Metadata.Name] = sk + } + + result := &skills.SyncResult{} + for _, entry := range lf.Skills { + sk, dbOK := installedByName[entry.Name] + s.syncLockedEntry(ctx, opts, entry, sk, dbOK, result) + } + for _, sk := range installed { + if _, ok := lf.Get(sk.Metadata.Name); ok { + continue // handled by the loop above + } + s.syncUnlockedInstall(ctx, opts, sk, result) + } + + return result, nil +} + +// errExperimentalLockFeature is returned by Sync/Upgrade while the lock file +// feature is behind its rollout gate (skills.LockFileFeatureEnabled). +var errExperimentalLockFeature = httperr.WithCode( + fmt.Errorf("skills lock file support is experimental; set %s=true to use it", skills.LockFileEnvVar), + http.StatusForbidden, +) + +// syncLockedEntry reconciles one lock file entry against installed state, +// appending its outcome to result. Missing (dbOK false) and drifted (digest +// or contentDigest mismatch) entries are reinstalled at the pinned reference +// unless opts.Check is set, in which case nothing is written — both states +// are still reported (Missing/Drifted), never as failures. Recording +// Missing before the Check return is what makes --check a real gate on a +// fresh clone or CI runner, where every entry has a lock entry but no +// install record. +func (s *service) syncLockedEntry( + ctx context.Context, + opts skills.SyncOptions, + entry lockfile.Entry, + sk skills.InstalledSkill, + dbOK bool, + result *skills.SyncResult, +) { + sigOK := true + if dbOK { + if sigErr := s.verifyStoredSignature(entry, sk); sigErr != nil { + // A failed offline re-verification is treated as drift: check + // mode reports it, apply mode reinstalls from the pinned + // reference — where install-time verification enforces the + // locked identity and, on success, heals the stored bundle. + sigOK = false + slog.Warn("stored signature failed offline re-verification", + "skill", entry.Name, "error", sigErr) + } + } + if dbOK && sigOK && entryMatchesInstalled(s.pathResolver, entry, sk) { + result.AlreadyCurrent = append(result.AlreadyCurrent, entry.Name) + return + } + if dbOK { + result.Drifted = append(result.Drifted, entry.Name) + } else { + result.Missing = append(result.Missing, entry.Name) + } + if opts.Check { + return + } + if err := s.reinstallPinned(ctx, opts, entry, sk, dbOK); err != nil { + result.Failed = append(result.Failed, skills.SyncFailure{ + Name: entry.Name, Reason: classifySyncFailure(err), Error: err.Error(), + }) + return + } + result.Installed = append(result.Installed, entry.Name) +} + +// entryMatchesInstalled reports whether the installed skill's pinned digest +// still matches the lock entry and EVERY client directory's on-disk +// contentDigest does too. Checking only one client's copy would leave +// tampering with any other client's materialized files invisible to +// --check — and which directory got checked would depend on install order. +func entryMatchesInstalled(pathResolver skills.PathResolver, entry lockfile.Entry, sk skills.InstalledSkill) bool { + if sk.Digest != entry.Digest { + return false + } + if len(sk.Clients) == 0 { + return false + } + for _, client := range sk.Clients { + dir, err := pathResolver.GetSkillPath(client, sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + if err != nil { + return false + } + contentDigest, err := lockfile.ContentDigestFromDir(dir) + if err != nil || contentDigest != entry.ContentDigest { + return false + } + } + return true +} + +// reinstallPinned reinstalls entry at its pinned reference, preserving its +// recorded Source (never re-resolving) and the clients it was previously +// installed for unless the caller overrides them. +func (s *service) reinstallPinned( + ctx context.Context, opts skills.SyncOptions, entry lockfile.Entry, existing skills.InstalledSkill, dbOK bool, +) error { + pinnedRef, err := buildPinnedReference(entry) + if err != nil { + return fmt.Errorf("pinning %q: %w", entry.Name, err) + } + clients := opts.Clients + if len(clients) == 0 && dbOK { + clients = existing.Clients + } + _, err = s.Install(ctx, skills.InstallOptions{ + Name: pinnedRef, + Scope: skills.ScopeProject, + ProjectRoot: opts.ProjectRoot, + Clients: clients, + Force: true, // sync restores exactly the pinned content over any drifted files + LockSource: entry.Source, + LockResolvedReference: entry.ResolvedReference, // preserve — pinnedRef is a restore form + SyncRestore: true, // reinstall despite unchanged Digest — drift is on disk, not the pin + }) + return err +} + +// syncUnlockedInstall classifies a project-scope install that has no lock +// entry: NeverManaged (optionally adopted) or RemovedFromLock (optionally +// pruned), appending the outcome to result. +func (s *service) syncUnlockedInstall( + ctx context.Context, opts skills.SyncOptions, sk skills.InstalledSkill, result *skills.SyncResult, +) { + if !sk.Managed { + result.NeverManaged = append(result.NeverManaged, sk.Metadata.Name) + if opts.Adopt && !opts.Check { + if err := s.adoptSkill(ctx, opts, sk); err != nil { + result.Failed = append(result.Failed, skills.SyncFailure{ + Name: sk.Metadata.Name, Reason: classifySyncFailure(err), Error: err.Error(), + }) + } + } + return + } + + result.RemovedFromLock = append(result.RemovedFromLock, sk.Metadata.Name) + if opts.Prune && !opts.Check { + if err := s.Uninstall(ctx, skills.UninstallOptions{ + Name: sk.Metadata.Name, Scope: skills.ScopeProject, ProjectRoot: opts.ProjectRoot, + }); err != nil { + result.Failed = append(result.Failed, skills.SyncFailure{ + Name: sk.Metadata.Name, Reason: classifySyncFailure(err), Error: err.Error(), + }) + return + } + result.Pruned = append(result.Pruned, sk.Metadata.Name) + } +} + +// verifyStoredSignature re-verifies the Sigstore bundle stored with an +// installed skill against the identity its lock entry records — entirely +// offline, via the embedded trust root. Entries recorded unsigned or with +// no provenance have nothing to verify. A recorded identity with no stored +// bundle fails closed for OCI installs (the bundle should exist); git +// installs never store a bundle — their signature lives on the commit and +// is re-verified when content is re-resolved. +func (s *service) verifyStoredSignature(entry lockfile.Entry, sk skills.InstalledSkill) error { + if entry.Unsigned || entry.Provenance == nil { + return nil + } + if len(sk.SigstoreBundle) == 0 { + if !strings.Contains(entry.Digest, ":") { + return nil // git install: no stored bundle by design + } + return fmt.Errorf("%w: lock entry records signer %q but no bundle is stored", + verifier.ErrSignatureInvalid, entry.Provenance.SignerIdentity) + } + return s.artifactVerifier().VerifyBundleOffline(sk.SigstoreBundle, entry.Digest, entry.Provenance) +} + +// adoptSkill writes a lock entry for an existing, unmanaged project-scope +// install, pinning its current on-disk state. The install's own Reference is +// used as Source: an adopted install predates (or never went through) lock +// tracking, so the original user-typed request is not recoverable — the +// concrete resolved reference is the closest available fact to pin against. +// +// Trust state is back-filled from the stored Sigstore bundle when one +// exists; otherwise adoption is the same trust decision as an unsigned +// install and requires the explicit AllowUnsigned exception. +func (s *service) adoptSkill(ctx context.Context, opts skills.SyncOptions, sk skills.InstalledSkill) error { + contentDigest, err := computeContentDigest(s.pathResolver, sk) + if err != nil { + return fmt.Errorf("computing content digest: %w", err) + } + var provenance *lockfile.Provenance + unsigned := false + if len(sk.SigstoreBundle) > 0 { + result, verifyErr := s.artifactVerifier().ResultFromBundle(sk.SigstoreBundle, sk.Digest) + if verifyErr != nil { + return fmt.Errorf("verifying stored bundle for adoption: %w", verifyErr) + } + provenance = provenanceInfoToLock(provenanceInfoFromResult(result)) + } + if provenance == nil { + if !opts.AllowUnsigned { + return httperr.WithCode( + fmt.Errorf("%w: adopting %q records it as unsigned; pass --allow-unsigned to accept that", + verifier.ErrUnsigned, sk.Metadata.Name), + http.StatusForbidden, + ) + } + unsigned = true + } + if err := recordLockEntry(sk.ProjectRoot, lockEntryInput{ + Name: sk.Metadata.Name, + Version: sk.Metadata.Version, + Source: sk.Reference, + ResolvedReference: sk.Reference, + Digest: sk.Digest, + ContentDigest: contentDigest, + Provenance: provenance, + Unsigned: unsigned, + }); err != nil { + return fmt.Errorf("writing lock entry: %w", errors.Join(errLockWrite, err)) + } + sk.Managed = true + if err := s.store.Update(ctx, sk); err != nil { + return fmt.Errorf("marking skill as lock-managed: %w", err) + } + return nil +} + +// classifySyncFailure maps an error from the install/uninstall path to an +// RFC THV-0080 typed failure reason using structured signals those paths +// already attach — the errLockWrite sentinel and httperr status codes — +// rather than matching on error message text. Lock-write failures are +// identified by the sentinel specifically: mapping every HTTP 500 to +// lock-write-failed would mislabel unrelated internal errors (e.g. a +// missing resolver) for the automation that keys on this reason. +func classifySyncFailure(err error) skills.FailureReason { + if errors.Is(err, errLockWrite) { + return skills.FailureReasonLockWriteFailed + } + if reason := classifySignatureError(err); reason != "" { + return reason + } + switch httperr.Code(err) { + case http.StatusNotFound: + return skills.FailureReasonDigestMissing + case http.StatusBadGateway, http.StatusGatewayTimeout, http.StatusTooManyRequests: + return skills.FailureReasonRegistryUnreachable + case http.StatusBadRequest, http.StatusUnprocessableEntity, http.StatusConflict: + return skills.FailureReasonValidationRejected + default: + return skills.FailureReasonUnknown + } +} diff --git a/pkg/skills/skillsvc/sync_test.go b/pkg/skills/skillsvc/sync_test.go new file mode 100644 index 0000000000..c7261a1679 --- /dev/null +++ b/pkg/skills/skillsvc/sync_test.go @@ -0,0 +1,420 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "net/http" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + skillsmocks "github.com/stacklok/toolhive/pkg/skills/mocks" + verifiermocks "github.com/stacklok/toolhive/pkg/skills/verifier/mocks" + "github.com/stacklok/toolhive/pkg/storage/sqlite" +) + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_FeatureDisabledReturnsForbidden(t *testing.T) { + gr, _ := newGitResolverMock(t) + svc, projectRoot := newLockTestService(t, gr) + t.Setenv(skills.LockFileEnvVar, "false") + + _, err := svc.(*service).Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.Error(t, err) + assert.Equal(t, http.StatusForbidden, httperr.Code(err)) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_ReportsUpToDateWhenNothingChanged(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.AlreadyCurrent) + assert.Empty(t, result.Installed) + assert.Empty(t, result.Drifted) + assert.Empty(t, result.Failed) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_ReinstallsDriftedContent(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // Tamper with the installed file. + skillMD := filepath.Join(projectRoot, ".claude", "skills", "my-skill", "SKILL.md") + require.NoError(t, os.WriteFile(skillMD, []byte("tampered content"), 0o644)) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.Drifted) + assert.Equal(t, []string{"my-skill"}, result.Installed) + + restored, err := os.ReadFile(skillMD) //nolint:gosec // fixed test path + require.NoError(t, err) + assert.Contains(t, string(restored), "name: my-skill") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_ReinstallPreservesResolvedReference(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + before, ok := readLockfile(t, projectRoot).Get("my-skill") + require.True(t, ok) + require.Equal(t, ref, before.ResolvedReference) + + // Tamper with the installed file so sync has to reinstall at the pinned + // reference — the drift-repair path that must not overwrite + // ResolvedReference with that internal pinned form. + skillMD := filepath.Join(projectRoot, ".claude", "skills", "my-skill", "SKILL.md") + require.NoError(t, os.WriteFile(skillMD, []byte("tampered content"), 0o644)) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + require.Equal(t, []string{"my-skill"}, result.Installed) + + after, ok := readLockfile(t, projectRoot).Get("my-skill") + require.True(t, ok) + assert.Equal(t, before.ResolvedReference, after.ResolvedReference, + "a drift-repair reinstall must preserve ResolvedReference, not overwrite it with the pinned restore form") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_CheckReportsDriftWithoutWriting(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + skillMD := filepath.Join(projectRoot, ".claude", "skills", "my-skill", "SKILL.md") + require.NoError(t, os.WriteFile(skillMD, []byte("tampered content"), 0o644)) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Check: true}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.Drifted) + assert.Empty(t, result.Installed, "check must not install/write anything") + + stillTampered, err := os.ReadFile(skillMD) //nolint:gosec // fixed test path + require.NoError(t, err) + assert.Equal(t, "tampered content", string(stillTampered)) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_MissingInstallIsRestored(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + require.NoError(t, os.RemoveAll(filepath.Join(projectRoot, ".claude", "skills", "my-skill"))) + require.NoError(t, svc.(*service).store.Delete(t.Context(), "my-skill", skills.ScopeProject, projectRoot)) //nolint:forcetypeassert + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.Missing, "an entry with no install record is Missing") + assert.Equal(t, []string{"my-skill"}, result.Installed) + assert.Empty(t, result.Drifted, "a missing (not previously-DB-tracked) install is Missing, not Drifted") + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err) +} + +// TestSync_CheckReportsMissingInstalls guards the fresh-clone CI gate: a lock +// entry with no install record at all (the state of every entry on a fresh +// checkout or CI runner, where SQLite state is per-machine) must be reported +// by --check, not silently skipped. Before the fix, such entries landed in +// no result bucket and the gate exited 0 with nothing installed. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_CheckReportsMissingInstalls(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // Simulate the fresh clone: lock file committed, no local install state. + require.NoError(t, os.RemoveAll(filepath.Join(projectRoot, ".claude", "skills", "my-skill"))) + require.NoError(t, svc.(*service).store.Delete(t.Context(), "my-skill", skills.ScopeProject, projectRoot)) //nolint:forcetypeassert + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Check: true}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.Missing, "--check must report the missing install") + assert.Empty(t, result.Installed, "check must not install anything") + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.Error(t, err, "check must not have installed anything") +} + +// TestSync_RestoreDoesNotReResolveDependencies covers the deterministic- +// restore guarantee: restoring a parent whose SKILL.md declares +// toolhive.requires must NOT re-materialize those dependencies from their +// mutable source strings. Each dependency has its own pinned lock entry the +// sync loop restores directly; re-resolving would silently upgrade and +// re-pin a dependency whose upstream tag moved. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_RestoreDoesNotReResolveDependencies(t *testing.T) { + gr, fx := newGitResolverMock(t) + // The dependency name sorts BEFORE the parent's: the sync loop restores + // it first (correctly, at its pin), and the parent's restore afterwards + // is what would re-resolve and re-pin it if requires recursion ran — + // with nothing after it in the loop to repair the damage. + depRef, _ := gitRef("a-pinned-dep") + fx.register("a-pinned-dep", gitSkill("a-pinned-dep")) + fx.register("z-parent-skill", gitSkill("z-parent-skill", depRef)) + svc, projectRoot := newLockTestService(t, gr) + + parentRef, _ := gitRef("z-parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: parentRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + depBefore, ok := readLockfile(t, projectRoot).Get("a-pinned-dep") + require.True(t, ok) + + // Move the dependency's upstream: same source string, newer content. + fx.register("a-pinned-dep", gitSkillVersion("a-pinned-dep")) + + // Wipe local state so sync must restore both entries from the lock file. + svcImpl := svc.(*service) //nolint:forcetypeassert + require.NoError(t, os.RemoveAll(filepath.Join(projectRoot, ".claude", "skills"))) + require.NoError(t, svcImpl.store.Delete(t.Context(), "z-parent-skill", skills.ScopeProject, projectRoot)) + require.NoError(t, svcImpl.store.Delete(t.Context(), "a-pinned-dep", skills.ScopeProject, projectRoot)) + + result, err := svcImpl.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Empty(t, result.Failed) + + depAfter, ok := readLockfile(t, projectRoot).Get("a-pinned-dep") + require.True(t, ok) + assert.Equal(t, depBefore.Digest, depAfter.Digest, + "restore must reinstall the dependency at its pinned digest, not re-resolve the moved source") + assert.Equal(t, depBefore.ResolvedReference, depAfter.ResolvedReference) +} + +// TestSync_RestorePreservesExplicitFlag: a restore is not a user install, +// so it must not promote a non-explicit dependency's lock entry to +// Explicit: true — that would permanently exempt it from cascade removal +// when its last requiring parent is uninstalled. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_RestorePreservesExplicitFlag(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("dep-skill") + fx.register("dep-skill", gitSkill("dep-skill")) + fx.register("parent-skill", gitSkill("parent-skill", depRef)) + svc, projectRoot := newLockTestService(t, gr) + + parentRef, _ := gitRef("parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: parentRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + before, ok := readLockfile(t, projectRoot).Get("dep-skill") + require.True(t, ok) + require.False(t, before.Explicit, "a materialized dependency starts non-explicit") + + // Tamper with the dependency so sync restores it. + depMD := filepath.Join(projectRoot, ".claude", "skills", "dep-skill", "SKILL.md") + require.NoError(t, os.WriteFile(depMD, []byte("tampered"), 0o644)) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Contains(t, result.Installed, "dep-skill") + + after, ok := readLockfile(t, projectRoot).Get("dep-skill") + require.True(t, ok) + assert.False(t, after.Explicit, "a restore must not promote a dependency to explicit") + assert.Equal(t, before.RequiredBy, after.RequiredBy) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_AdoptsUnmanagedInstall(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("unmanaged-skill", gitSkill("unmanaged-skill")) + svc, projectRoot := newLockTestService(t, gr) + + // Disable the feature for the initial install so it lands unmanaged + // (no lock entry, Managed=false) — simulating a pre-existing install + // from before the lock feature was ever enabled. + t.Setenv(skills.LockFileEnvVar, "false") + ref, _ := gitRef("unmanaged-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + t.Setenv(skills.LockFileEnvVar, "true") + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"unmanaged-skill"}, result.NeverManaged) + + // Adoption of an install with no stored bundle is an unsigned trust + // decision: without the explicit exception it fails... + result, err = syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Adopt: true}) + require.NoError(t, err) + require.Len(t, result.Failed, 1) + assert.Equal(t, skills.FailureReasonUnsignedRejected, result.Failed[0].Reason) + lf := readLockfile(t, projectRoot) + _, ok := lf.Get("unmanaged-skill") + assert.False(t, ok, "adoption without the unsigned exception must not write a lock entry") + + // ...and with it, the entry records the unsigned state. + result, err = syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Adopt: true, AllowUnsigned: true}) + require.NoError(t, err) + assert.Equal(t, []string{"unmanaged-skill"}, result.NeverManaged) + + lf = readLockfile(t, projectRoot) + entry, ok := lf.Get("unmanaged-skill") + require.True(t, ok, "adopt must write a lock entry for the unmanaged install") + assert.NotEmpty(t, entry.ContentDigest) + assert.True(t, entry.Unsigned, "an adoption without verifiable trust must record unsigned") + + sk, err := svc.Info(t.Context(), skills.InfoOptions{Name: "unmanaged-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err) + require.NotNil(t, sk.InstalledSkill) + assert.True(t, sk.InstalledSkill.Managed, "adopt must mark the DB record as managed") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_PrunesRemovedFromLock(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // Simulate the entry being deleted from the lock file by hand (or by a + // teammate's commit) while the DB record and files remain. + require.NoError(t, lockfile.RemoveEntry(mustOpenRoot(t, projectRoot), "my-skill")) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.RemovedFromLock) + + result, err = syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Prune: true}) + require.NoError(t, err) + assert.Equal(t, []string{"my-skill"}, result.Pruned) + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.Error(t, err, "prune must uninstall the skill") +} + +// TestSync_CheckDetectsTamperInAnyClientDir: verification must cover every +// client directory a skill is installed into. Checking only the first +// client's copy would leave tampering with any other client's materialized +// files invisible to --check — and which directory got checked would depend +// on install order. +// +//nolint:paralleltest // uses t.Setenv, incompatible with t.Parallel +func TestSync_CheckDetectsTamperInAnyClientDir(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("multi-skill", gitSkill("multi-skill")) + + t.Setenv(skills.LockFileEnvVar, "true") + dbPath := filepath.Join(t.TempDir(), "test.db") + db, err := sqlite.Open(t.Context(), dbPath) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + store := sqlite.NewSkillStore(db) + + projectRoot := makeProjectRoot(t) + // Unlike newLockTestService's resolver, this one gives every client its + // own directory so the per-client verification actually has two copies + // to distinguish. + ctrl := gomock.NewController(t) + pr := skillsmocks.NewMockPathResolver(ctrl) + pr.EXPECT().GetSkillPath(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes(). + DoAndReturn(func(client, skillName string, _ skills.Scope, _ string) (string, error) { + return filepath.Join(projectRoot, "."+client, "skills", skillName), nil + }) + pr.EXPECT().ListSkillSupportingClients().AnyTimes().Return([]string{"claude-code", "cursor"}) + mv := verifiermocks.NewMockVerifier(ctrl) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(signedResult(), nil) + mv.EXPECT().VerifyBundleOffline(gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(nil) + svc := New(store, WithPathResolver(pr), WithGitResolver(gr), WithVerifier(mv)) + + ref, _ := gitRef("multi-skill") + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, + Clients: []string{"claude-code", "cursor"}, + }) + require.NoError(t, err) + + syncer := svc.(*service) //nolint:forcetypeassert + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Check: true}) + require.NoError(t, err) + require.Equal(t, []string{"multi-skill"}, result.AlreadyCurrent, "untampered install must verify clean") + + // Tamper with the SECOND client's copy only. + cursorMD := filepath.Join(projectRoot, ".cursor", "skills", "multi-skill", "SKILL.md") + require.NoError(t, os.WriteFile(cursorMD, []byte("tampered"), 0o644)) + + result, err = syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Check: true}) + require.NoError(t, err) + assert.Equal(t, []string{"multi-skill"}, result.Drifted, + "tampering with a non-first client's copy must be reported as drift") +} diff --git a/pkg/skills/skillsvc/sync_verify_test.go b/pkg/skills/skillsvc/sync_verify_test.go new file mode 100644 index 0000000000..a7fbe2e699 --- /dev/null +++ b/pkg/skills/skillsvc/sync_verify_test.go @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/verifier" + verifiermocks "github.com/stacklok/toolhive/pkg/skills/verifier/mocks" +) + +// TestSync_StoredSignatureFailureIsDriftThenHeals proves the offline +// re-verification path: a stored bundle that no longer verifies reports as +// drift in check mode, and an apply reinstalls from the pinned reference — +// where install-time verification enforces the locked identity and heals +// the stored state. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestSync_StoredSignatureFailureIsDriftThenHeals(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("sig-drift-skill", gitSkill("sig-drift-skill")) + + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(signedResult(), nil) + // Every offline re-verification of the stored bundle fails. + mv.EXPECT().VerifyBundleOffline(gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(verifier.ErrSignatureInvalid) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + ref, _ := gitRef("sig-drift-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + syncer := svc.(*service) //nolint:forcetypeassert + + result, err := syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot, Check: true}) + require.NoError(t, err) + assert.Equal(t, []string{"sig-drift-skill"}, result.Drifted, + "a failed offline re-verification must report as drift in check mode") + assert.Empty(t, result.AlreadyCurrent) + + result, err = syncer.Sync(t.Context(), skills.SyncOptions{ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"sig-drift-skill"}, result.Installed, + "apply mode must reinstall from the pinned reference, re-verifying the artifact") + assert.Empty(t, result.Failed) +} + +func TestVerifyStoredSignature(t *testing.T) { + t.Parallel() + + provenance := &lockfile.Provenance{ + SignerIdentity: testSignerIdentity, + CertIssuer: testCertIssuer, + } + ociDigest := "sha256:" + strings.Repeat("a", 64) + gitDigest := strings.Repeat("b", 40) + + tests := []struct { + name string + entry lockfile.Entry + bundle []byte + offlineErr error + expectCall bool + wantErr bool + }{ + { + name: "unsigned entry has nothing to verify", + entry: lockfile.Entry{Unsigned: true, Digest: ociDigest}, + }, + { + name: "no provenance has nothing to verify", + entry: lockfile.Entry{Digest: ociDigest}, + }, + { + name: "provenance without stored bundle fails closed for OCI", + entry: lockfile.Entry{Provenance: provenance, Digest: ociDigest}, + wantErr: true, + }, + { + name: "provenance without stored bundle is fine for git", + entry: lockfile.Entry{Provenance: provenance, Digest: gitDigest}, + }, + { + name: "stored bundle delegates to offline verification", + entry: lockfile.Entry{Provenance: provenance, Digest: ociDigest}, + bundle: []byte(`{"bundle":true}`), + expectCall: true, + }, + { + name: "offline verification failure propagates", + entry: lockfile.Entry{Provenance: provenance, Digest: ociDigest}, + bundle: []byte(`{"bundle":true}`), + offlineErr: verifier.ErrSignatureInvalid, + expectCall: true, + wantErr: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + if tc.expectCall { + mv.EXPECT().VerifyBundleOffline(tc.bundle, tc.entry.Digest, tc.entry.Provenance). + Return(tc.offlineErr) + } + svc := &service{sigVerifier: mv} + err := svc.verifyStoredSignature(tc.entry, skills.InstalledSkill{SigstoreBundle: tc.bundle}) + if tc.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + }) + } +} diff --git a/pkg/skills/skillsvc/uninstall.go b/pkg/skills/skillsvc/uninstall.go index 4c74d8ec98..9c52241957 100644 --- a/pkg/skills/skillsvc/uninstall.go +++ b/pkg/skills/skillsvc/uninstall.go @@ -14,9 +14,13 @@ import ( "github.com/stacklok/toolhive-core/httperr" "github.com/stacklok/toolhive/pkg/groups" "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" ) // Uninstall removes an installed skill and cleans up files for all clients. +// For a project-scope, lock-managed skill (see skills.LockFileFeatureEnabled), +// it also removes the skill's lock entry and cascades to any dependency that +// loses its last requiring parent as a result. func (s *service) Uninstall(ctx context.Context, opts skills.UninstallOptions) error { if err := skills.ValidateSkillName(opts.Name); err != nil { return httperr.WithCode(err, http.StatusBadRequest) @@ -29,6 +33,21 @@ func (s *service) Uninstall(ctx context.Context, opts skills.UninstallOptions) e scope = defaultScope(scope) opts.ProjectRoot = projectRoot + return s.uninstallOne(ctx, opts, scope) +} + +// uninstallOne performs a single skill's lock-entry removal, file/DB/group +// cleanup and, when applicable, its dependency cascade. It is called both +// for the top-level Uninstall request and recursively for cascade candidates. +// +// The lock entry is removed FIRST, and its failure aborts the uninstall +// while everything is still intact. The reverse order (files/DB first, lock +// entry best-effort) had a resurrection hazard: a lock-write failure after +// the record and files were gone left a stale entry that the next sync +// silently reinstalled. Failing after the entry is removed leaves the +// opposite, safe inconsistency — an installed-but-unlocked skill that sync +// reports as removed-from-lock and prune can clean up. +func (s *service) uninstallOne(ctx context.Context, opts skills.UninstallOptions, scope skills.Scope) error { unlock := s.locks.lock(opts.Name, scope, opts.ProjectRoot) defer unlock() @@ -38,36 +57,25 @@ func (s *service) Uninstall(ctx context.Context, opts skills.UninstallOptions) e return err } - // Determine the boundary directory for empty-parent cleanup. - stopDir := opts.ProjectRoot - if scope == skills.ScopeUser { - if homeDir, err := os.UserHomeDir(); err == nil { - stopDir = homeDir - } + visited := opts.Visited + if visited == nil { + visited = make(map[string]struct{}) } + visited[opts.Name] = struct{}{} - // Remove files for each client — best-effort: collect errors but don't - // abort on the first failure so we clean up as much as possible. - var cleanupErrs []error - if s.pathResolver != nil { - for _, clientType := range existing.Clients { - skillPath, pathErr := s.pathResolver.GetSkillPath(clientType, opts.Name, scope, opts.ProjectRoot) - if pathErr != nil { - cleanupErrs = append(cleanupErrs, fmt.Errorf("resolving path for client %q: %w", clientType, pathErr)) - continue - } - if rmErr := s.installer.Remove(skillPath); rmErr != nil { - cleanupErrs = append(cleanupErrs, fmt.Errorf("removing files for client %q: %w", clientType, rmErr)) - continue - } - if stopDir != "" { - skills.RemoveEmptyParents(filepath.Dir(skillPath), stopDir) - } + var cascadeCandidates []string + if scope == skills.ScopeProject && existing.Managed && skills.LockFileFeatureEnabled() { + cascadeCandidates, err = removeLockEntry(opts) + if err != nil { + return fmt.Errorf("updating project lock file: %w", err) } } + cleanupErrs := s.removeClientFiles(existing, opts, scope) + if err := s.store.Delete(ctx, opts.Name, scope, opts.ProjectRoot); err != nil { - return err + cleanupErrs = append(cleanupErrs, err) + return errors.Join(cleanupErrs...) } // Remove the skill from all groups — best-effort, same pattern as file cleanup. @@ -77,5 +85,95 @@ func (s *service) Uninstall(ctx context.Context, opts skills.UninstallOptions) e } } + if cascadeErr := s.cascadeUninstall(ctx, cascadeCandidates, visited, opts.ProjectRoot, scope); cascadeErr != nil { + cleanupErrs = append(cleanupErrs, cascadeErr) + } + return errors.Join(cleanupErrs...) } + +// removeClientFiles removes on-disk files for every client existing is +// installed for. It is best-effort: errors are collected rather than +// aborting, so cleanup proceeds as far as possible. +func (s *service) removeClientFiles( + existing skills.InstalledSkill, opts skills.UninstallOptions, scope skills.Scope, +) []error { + if s.pathResolver == nil { + return nil + } + + // Determine the boundary directory for empty-parent cleanup. + stopDir := opts.ProjectRoot + if scope == skills.ScopeUser { + if homeDir, err := os.UserHomeDir(); err == nil { + stopDir = homeDir + } + } + + var errs []error + for _, clientType := range existing.Clients { + skillPath, pathErr := s.pathResolver.GetSkillPath(clientType, opts.Name, scope, opts.ProjectRoot) + if pathErr != nil { + errs = append(errs, fmt.Errorf("resolving path for client %q: %w", clientType, pathErr)) + continue + } + if rmErr := s.installer.Remove(skillPath); rmErr != nil { + errs = append(errs, fmt.Errorf("removing files for client %q: %w", clientType, rmErr)) + continue + } + if stopDir != "" { + skills.RemoveEmptyParents(filepath.Dir(skillPath), stopDir) + } + } + return errs +} + +// removeLockEntry removes opts.Name's lock entry and strips it from every +// other entry's RequiredBy list, returning the names of dependencies that +// consequently lost their last requiring parent and are not explicit — the +// cascade-removal candidates. +func removeLockEntry(opts skills.UninstallOptions) ([]string, error) { + root, err := lockfile.OpenRoot(opts.ProjectRoot) + if err != nil { + return nil, err + } + var cascadeCandidates []string + if err := lockfile.Update(root, func(lf *lockfile.Lockfile) error { + lf.Remove(opts.Name) + cascadeCandidates = lf.RemoveParentFromRequiredBy(opts.Name) + return nil + }); err != nil { + return nil, err + } + return cascadeCandidates, nil +} + +// cascadeUninstall uninstalls each candidate dependency that has not +// already been visited. The visited set prevents infinite recursion on a +// requiredBy cycle in a hand-edited lock. +func (s *service) cascadeUninstall( + ctx context.Context, candidates []string, visited map[string]struct{}, projectRoot string, scope skills.Scope, +) error { + var errs []error + for _, dep := range candidates { + if _, seen := visited[dep]; seen { + continue + } + visited[dep] = struct{}{} + // A lock entry can reference a dependency whose install failed + // partway (or was already removed by hand); skip it rather than + // erroring on a missing DB record. + if _, getErr := s.store.Get(ctx, dep, scope, projectRoot); getErr != nil { + continue + } + if uninstallErr := s.uninstallOne(ctx, skills.UninstallOptions{ + Name: dep, + Scope: scope, + ProjectRoot: projectRoot, + Visited: visited, + }, scope); uninstallErr != nil { + errs = append(errs, fmt.Errorf("cascade-removing dependency %q: %w", dep, uninstallErr)) + } + } + return errors.Join(errs...) +} diff --git a/pkg/skills/skillsvc/uninstall_cascade_test.go b/pkg/skills/skillsvc/uninstall_cascade_test.go new file mode 100644 index 0000000000..3bff0de3be --- /dev/null +++ b/pkg/skills/skillsvc/uninstall_cascade_test.go @@ -0,0 +1,151 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/skills" +) + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUninstall_CascadesToOrphanedDependency(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("shared-dep") + fx.register("parent-skill", gitSkill("parent-skill", depRef)) + fx.register("shared-dep", gitSkill("shared-dep")) + svc, projectRoot := newLockTestService(t, gr) + + parentRef, _ := gitRef("parent-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: parentRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + _, ok := readLockfile(t, projectRoot).Get("shared-dep") + require.True(t, ok, "expected shared-dep to be materialized before uninstalling parent-skill") + + require.NoError(t, svc.Uninstall(t.Context(), skills.UninstallOptions{ + Name: "parent-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot, + })) + + lf := readLockfile(t, projectRoot) + _, ok = lf.Get("parent-skill") + assert.False(t, ok, "parent's own lock entry must be removed") + _, ok = lf.Get("shared-dep") + assert.False(t, ok, "an orphaned, non-explicit dependency must cascade-uninstall") + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "shared-dep", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.Error(t, err, "cascade must remove the dependency's DB record and files too") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUninstall_SharedDependencySurvivesOneParentRemoval(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("shared-dep") + fx.register("parent-a", gitSkill("parent-a", depRef)) + fx.register("parent-b", gitSkill("parent-b", depRef)) + fx.register("shared-dep", gitSkill("shared-dep")) + svc, projectRoot := newLockTestService(t, gr) + + refA, _ := gitRef("parent-a") + refB, _ := gitRef("parent-b") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: refA, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: refB, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + require.NoError(t, svc.Uninstall(t.Context(), skills.UninstallOptions{ + Name: "parent-a", Scope: skills.ScopeProject, ProjectRoot: projectRoot, + })) + + lf := readLockfile(t, projectRoot) + dep, ok := lf.Get("shared-dep") + require.True(t, ok, "dependency must survive while parent-b still requires it") + assert.Equal(t, []string{"parent-b"}, dep.RequiredBy) + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "shared-dep", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err, "the dependency's DB record and files must remain") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUninstall_ExplicitDependencyIsNeverCascaded(t *testing.T) { + gr, fx := newGitResolverMock(t) + depRef, _ := gitRef("also-explicit") + fx.register("parent-skill", gitSkill("parent-skill", depRef)) + fx.register("also-explicit", gitSkill("also-explicit")) + svc, projectRoot := newLockTestService(t, gr) + + // Install the dependency explicitly first, then install a parent that + // also requires it. + explicitRef, _ := gitRef("also-explicit") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: explicitRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + parentRef, _ := gitRef("parent-skill") + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: parentRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + lf := readLockfile(t, projectRoot) + dep, ok := lf.Get("also-explicit") + require.True(t, ok) + assert.True(t, dep.Explicit, "an entry installed explicitly stays explicit even once also required by a parent") + assert.Equal(t, []string{"parent-skill"}, dep.RequiredBy) + + require.NoError(t, svc.Uninstall(t.Context(), skills.UninstallOptions{ + Name: "parent-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot, + })) + + lf = readLockfile(t, projectRoot) + dep, ok = lf.Get("also-explicit") + require.True(t, ok, "an explicit entry must never be cascade-removed") + assert.Empty(t, dep.RequiredBy) + + _, err = svc.Info(t.Context(), skills.InfoOptions{Name: "also-explicit", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUninstall_CascadeCycleTerminates(t *testing.T) { + gr, fx := newGitResolverMock(t) + refA, _ := gitRef("skill-a") + refB, _ := gitRef("skill-b") + fx.register("skill-a", gitSkill("skill-a", refB)) + fx.register("skill-b", gitSkill("skill-b", refA)) + svc, projectRoot := newLockTestService(t, gr) + + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: refA, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + done := make(chan error, 1) + go func() { + done <- svc.Uninstall(t.Context(), skills.UninstallOptions{ + Name: "skill-a", Scope: skills.ScopeProject, ProjectRoot: projectRoot, + }) + }() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("timeout: cascade uninstall did not terminate on a requiredBy cycle") + } + + lf := readLockfile(t, projectRoot) + _, ok := lf.Get("skill-a") + assert.False(t, ok) + _, ok = lf.Get("skill-b") + assert.False(t, ok, "skill-b lost its only (cyclic) parent and must cascade too") +} diff --git a/pkg/skills/skillsvc/upgrade.go b/pkg/skills/skillsvc/upgrade.go new file mode 100644 index 0000000000..1e5ca599d8 --- /dev/null +++ b/pkg/skills/skillsvc/upgrade.go @@ -0,0 +1,344 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + + nameref "github.com/google/go-containerregistry/pkg/name" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/gitresolver" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/verifier" +) + +// var _ ensures *service continues to satisfy the full lock service surface +// now that both Sync (PR4) and Upgrade exist. +var _ skills.SkillLockService = (*service)(nil) + +// Upgrade re-resolves each targeted lock entry's Source and, when the +// resolved digest has changed, installs the newer content and rewrites the +// entry (Source itself is never rewritten — see RFC THV-0080). Entries +// pinned to an immutable reference (an OCI digest or a full git commit hash) +// are reported not-upgradable: there is nothing newer to resolve to. +func (s *service) Upgrade(ctx context.Context, opts skills.UpgradeOptions) (*skills.UpgradeResult, error) { + if !skills.LockFileFeatureEnabled() { + return nil, errExperimentalLockFeature + } + + _, projectRoot, err := normalizeProjectRoot(skills.ScopeProject, opts.ProjectRoot) + if err != nil { + return nil, err + } + opts.ProjectRoot = projectRoot + + root, err := lockfile.OpenRoot(projectRoot) + if err != nil { + return nil, err + } + lf, err := lockfile.Load(root) + if err != nil { + return nil, err + } + + targets, err := selectUpgradeTargets(lf, opts.Names) + if err != nil { + return nil, err + } + + // Resolve every target's latest state first, without installing + // anything. FailOnChanges is a CI freshness gate: it reports the full + // planned outcome set and never runs the apply pass at all — returning + // the outcomes (rather than an error that discards them) lets callers + // see exactly which skills are stale and distinguish "would change" + // from a genuine resolution failure. Exit-code mapping happens in the + // CLI from these outcomes, mirroring how sync --check works. + plans := make([]upgradePlan, len(targets)) + for i, entry := range targets { + plans[i] = s.planUpgrade(ctx, opts, entry) + } + + result := &skills.UpgradeResult{Outcomes: make([]skills.UpgradeOutcome, 0, len(plans))} + for _, p := range plans { + if opts.FailOnChanges { + result.Outcomes = append(result.Outcomes, p.outcome) + continue + } + result.Outcomes = append(result.Outcomes, s.applyUpgrade(ctx, opts, p)) + } + return result, nil +} + +// selectUpgradeTargets returns the lock entries to upgrade: every entry when +// names is empty, or the named subset in the order requested. An unknown +// name is an error — it is almost always a typo, and silently skipping it +// would make a scripted "upgrade these specific skills" call falsely report +// success. +func selectUpgradeTargets(lf *lockfile.Lockfile, names []string) ([]lockfile.Entry, error) { + if len(names) == 0 { + return lf.Skills, nil + } + targets := make([]lockfile.Entry, 0, len(names)) + for _, name := range names { + entry, ok := lf.Get(name) + if !ok { + return nil, httperr.WithCode( + fmt.Errorf("skill %q is not present in the lock file", name), + http.StatusNotFound, + ) + } + targets = append(targets, entry) + } + return targets, nil +} + +// upgradePlan is entry's resolved outcome before any install happens: either +// a terminal status (not-upgradable, up-to-date, ref-change-blocked, or a +// resolution failure) that needs no further action, or the pinned reference +// to install when the upgrade is applied. +type upgradePlan struct { + entry lockfile.Entry + outcome skills.UpgradeOutcome + pinnedRef string // set only when the upgrade needs installing + resolvedRef string // the resolved reference to record as ResolvedReference; set alongside pinnedRef +} + +// planUpgrade resolves entry's current state and determines its outcome, +// without installing anything — this lets Upgrade check --fail-on-changes +// against every target before any of them are applied. When an upgrade is +// warranted, the outcome's digest is pinned into pinnedRef so applyUpgrade +// installs exactly what was resolved here, rather than re-resolving +// entry.Source from scratch (which could pick up a different digest if a +// mutable ref moved between planning and applying). +func (s *service) planUpgrade(ctx context.Context, opts skills.UpgradeOptions, entry lockfile.Entry) upgradePlan { + outcome := skills.UpgradeOutcome{Name: entry.Name, OldDigest: entry.Digest} + + if isImmutableSource(entry) { + outcome.Status = skills.UpgradeStatusNotUpgradable + return upgradePlan{entry: entry, outcome: outcome} + } + + newRef, newDigest, err := s.resolveLatestState(ctx, entry.Source) + if err != nil { + outcome.Status = skills.UpgradeStatusFailed + outcome.Reason = classifySyncFailure(err) + outcome.Error = err.Error() + return upgradePlan{entry: entry, outcome: outcome} + } + outcome.NewDigest = newDigest + + if newDigest == entry.Digest { + outcome.Status = skills.UpgradeStatusUpToDate + return upgradePlan{entry: entry, outcome: outcome} + } + + if newRef != entry.ResolvedReference { + outcome.NewResolvedReference = newRef + // Only a move to a different repository is a supply-chain event. A + // tag moving within the same repository is how a catalog-sourced + // skill advances at all, and blocking it would force automation to + // pass --allow-ref-change on every routine upgrade — which would + // also disable the repository check this guard exists for. + if repositoryMoved(entry.ResolvedReference, newRef) && !opts.AllowRefChange { + outcome.Status = skills.UpgradeStatusRefChangeBlocked + return upgradePlan{entry: entry, outcome: outcome} + } + } + + // Signer-change guard: when the entry records a signer identity, the + // candidate must verify with the same one before the upgrade is even + // planned. Verification happens here (plan-only modes stay + // install-free) with a nil expected identity so a differing signer is + // reported as a change rather than a bare failure; blocked plans carry + // no pinnedRef, exactly like ref changes. + if entry.Provenance != nil && !opts.AllowSignerChange { + if blocked := s.guardSignerChange(ctx, entry, newRef, newDigest, &outcome); blocked { + return upgradePlan{entry: entry, outcome: outcome} + } + } + + pinnedRef, err := buildPinnedReference(lockfile.Entry{ResolvedReference: newRef, Digest: newDigest}) + if err != nil { + outcome.Status = skills.UpgradeStatusFailed + outcome.Reason = skills.FailureReasonUnknown + outcome.Error = fmt.Errorf("pinning resolved reference: %w", err).Error() + return upgradePlan{entry: entry, outcome: outcome} + } + + outcome.Status = skills.UpgradeStatusUpgraded + return upgradePlan{entry: entry, outcome: outcome, pinnedRef: pinnedRef, resolvedRef: newRef} +} + +// guardSignerChange probes the candidate artifact's signer identity and +// fills outcome when the upgrade must not proceed: the candidate is signed +// by a different identity (or unsigned) versus the recorded provenance, or +// its signature cannot be verified at all. Returns true when blocked. +func (s *service) guardSignerChange( + ctx context.Context, + entry lockfile.Entry, + newRef, newDigest string, + outcome *skills.UpgradeOutcome, +) bool { + probe, probeErr := s.probeCandidateSigner(ctx, newRef, newDigest) + switch { + case probeErr != nil && errors.Is(probeErr, verifier.ErrUnsigned): + outcome.Status = skills.UpgradeStatusSignerChangeBlocked + return true + case probeErr != nil: + outcome.Status = skills.UpgradeStatusFailed + outcome.Reason = classifySignatureError(probeErr) + if outcome.Reason == "" { + outcome.Reason = skills.FailureReasonUnknown + } + outcome.Error = probeErr.Error() + return true + case probe.SignerIdentity != entry.Provenance.SignerIdentity || + probe.CertIssuer != entry.Provenance.CertIssuer: + outcome.Status = skills.UpgradeStatusSignerChangeBlocked + outcome.NewSignerIdentity = probe.SignerIdentity + return true + } + return false +} + +// probeCandidateSigner verifies the candidate artifact chain-of-trust-only +// (nil expected identity) and returns the observed identity. Git candidates +// are re-resolved at the pinned commit to obtain the signature material. +func (s *service) probeCandidateSigner(ctx context.Context, newRef, newDigest string) (*verifier.Result, error) { + if strings.Contains(newDigest, ":") { + return s.artifactVerifier().VerifyOCI(ctx, newRef, newDigest, nil) + } + pinned, err := buildPinnedReference(lockfile.Entry{ResolvedReference: newRef, Digest: newDigest}) + if err != nil { + return nil, err + } + gitRef, err := gitresolver.ParseGitReference(pinned) + if err != nil { + return nil, err + } + resolved, err := s.gitResolver.Resolve(ctx, gitRef) + if err != nil { + return nil, err + } + return s.artifactVerifier().VerifyGit(ctx, resolved.CommitPayload, []byte(resolved.CommitSignature), nil) +} + +// applyUpgrade installs plan's pinned content when the plan calls for it. +// Preview mode reports the plan's outcome without installing anything. +func (s *service) applyUpgrade(ctx context.Context, opts skills.UpgradeOptions, plan upgradePlan) skills.UpgradeOutcome { + if plan.pinnedRef == "" || opts.Preview { + return plan.outcome + } + + clients := opts.Clients + if len(clients) == 0 { + if existing, err := s.store.Get(ctx, plan.entry.Name, skills.ScopeProject, opts.ProjectRoot); err == nil { + clients = existing.Clients + } + } + + if _, err := s.Install(ctx, skills.InstallOptions{ + Name: plan.pinnedRef, + Scope: skills.ScopeProject, + ProjectRoot: opts.ProjectRoot, + Clients: clients, + LockSource: plan.entry.Source, + LockResolvedReference: plan.resolvedRef, + AllowSignerChange: opts.AllowSignerChange, + }); err != nil { + outcome := plan.outcome + outcome.Status = skills.UpgradeStatusFailed + outcome.Reason = classifySyncFailure(err) + outcome.Error = err.Error() + return outcome + } + + return plan.outcome +} + +// resolveLatestState re-resolves source (a lock entry's original Source +// value) to its current resolvedReference and digest, using the same +// dispatch order as Install (git, direct OCI with registry fallback, +// registry name), but stopping short of extraction or any DB/lock write. +// For OCI sources this still pulls the artifact into the local store — +// there is no lighter "digest only" primitive in RegistryClient — matching +// the RFC's "preview is not side-effect-free" note; git sources resolve +// without touching disk. +func (s *service) resolveLatestState(ctx context.Context, source string) (resolvedRef, digestStr string, err error) { + // resolvedState carries the two return values through the shared + // source-dispatch skeleton, which routes exactly like Install does + // (git, direct OCI with registry fallback, plain registry name) so the + // two can never drift again. + type resolvedState struct{ ref, digest string } + + state, err := dispatchSource(ctx, s, source, sourceOps[resolvedState]{ + git: func(ctx context.Context, gitURL string) (resolvedState, error) { + r, d, gitErr := s.resolveGitLatest(ctx, gitURL) + return resolvedState{r, d}, gitErr + }, + oci: func(ctx context.Context, ref nameref.Reference) (resolvedState, error) { + r, d, ociErr := s.resolveOCILatest(ctx, ref) + return resolvedState{r, d}, ociErr + }, + registry: func(ctx context.Context, resolved *registryResolveResult) (resolvedState, error) { + r, d, regErr := s.resolveRegistryLatest(ctx, source, resolved) + return resolvedState{r, d}, regErr + }, + }) + return state.ref, state.digest, err +} + +// resolveRegistryLatest resolves the latest state of a registry catalogue +// result, dispatching to the OCI or git resolver it points at. +func (s *service) resolveRegistryLatest( + ctx context.Context, source string, resolved *registryResolveResult, +) (string, string, error) { + switch { + case resolved.OCIRef != nil: + return s.resolveOCILatest(ctx, resolved.OCIRef) + case resolved.GitURL != "": + return s.resolveGitLatest(ctx, resolved.GitURL) + } + return "", "", httperr.WithCode( + fmt.Errorf("skill %q resolved from registry but has no installable package", source), + http.StatusUnprocessableEntity, + ) +} + +func (s *service) resolveGitLatest(ctx context.Context, gitURL string) (string, string, error) { + if s.gitResolver == nil { + return "", "", httperr.WithCode(errors.New("git resolver is not configured"), http.StatusInternalServerError) + } + gitRef, err := gitresolver.ParseGitReference(gitURL) + if err != nil { + return "", "", httperr.WithCode(fmt.Errorf("invalid git reference: %w", err), http.StatusBadRequest) + } + resolved, err := s.gitResolver.Resolve(ctx, gitRef) + if err != nil { + return "", "", httperr.WithCode(fmt.Errorf("resolving git skill: %w", err), http.StatusBadGateway) + } + return gitURL, resolved.CommitHash, nil +} + +func (s *service) resolveOCILatest(ctx context.Context, ref nameref.Reference) (string, string, error) { + if s.registry == nil || s.ociStore == nil { + return "", "", httperr.WithCode(errors.New("OCI registry is not configured"), http.StatusInternalServerError) + } + d, err := s.registry.Pull(ctx, s.ociStore, ref.String()) + if err != nil { + return "", "", httperr.WithCode(fmt.Errorf("pulling %q: %w", ref.String(), err), classifyPullError(err)) + } + // qualifiedOCIRef, not ref.String(): install records the qualified form + // (implicit ":latest" made explicit) in ResolvedReference, and this value + // is compared against it for the ref-change guard. The unqualified form + // would misreport every digest change on a tag-less source as a blocked + // reference change. + return qualifiedOCIRef(ref), d.String(), nil +} diff --git a/pkg/skills/skillsvc/upgrade_oci_test.go b/pkg/skills/skillsvc/upgrade_oci_test.go new file mode 100644 index 0000000000..749122c858 --- /dev/null +++ b/pkg/skills/skillsvc/upgrade_oci_test.go @@ -0,0 +1,233 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "fmt" + "testing" + + godigest "github.com/opencontainers/go-digest" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + ociskills "github.com/stacklok/toolhive-core/oci/skills" + ocimocks "github.com/stacklok/toolhive-core/oci/skills/mocks" + regtypes "github.com/stacklok/toolhive-core/registry/types" + regmocks "github.com/stacklok/toolhive/pkg/registry/mocks" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + storemocks "github.com/stacklok/toolhive/pkg/storage/mocks" +) + +// newOCIUpgradeService builds a service with a mocked OCI registry client, a +// real (temp) OCI store, and an optional registry lookup — the dependencies +// planUpgrade's OCI resolution path touches. The skill store is a mock with +// no expectations: planning never writes. +func newOCIUpgradeService( + t *testing.T, reg ociskills.RegistryClient, lookup *regmocks.MockProvider, +) *service { + t.Helper() + ctrl := gomock.NewController(t) + ociStore, err := ociskills.NewStore(tempDir(t)) + require.NoError(t, err) + + opts := []Option{WithRegistryClient(reg), WithOCIStore(ociStore)} + if lookup != nil { + opts = append(opts, WithSkillLookup(lookup)) + } + svc := New(storemocks.NewMockSkillStore(ctrl), opts...) + return svc.(*service) //nolint:forcetypeassert // white-box test in the same package +} + +func ociTestDigest(seed byte) string { + const alphabet = "0123456789abcdef" + b := make([]byte, 64) + for i := range b { + b[i] = alphabet[(i+int(seed))%len(alphabet)] + } + return "sha256:" + string(b) +} + +// TestPlanUpgrade_OCITagLessSourceUpgrades pins the resolvedReference +// qualification contract: install records qualifiedOCIRef (implicit +// ":latest" made explicit), so upgrade's resolution must produce the same +// form. Before the fix, resolveOCILatest returned ref.String() — the +// unqualified form — so every digest change on a tag-less source was +// misreported as a blocked reference change instead of an upgrade. +func TestPlanUpgrade_OCITagLessSourceUpgrades(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "ghcr.io/org/skill"). + Return(godigest.Digest(ociTestDigest(2)), nil) + svc := newOCIUpgradeService(t, reg, nil) + + entry := lockfile.Entry{ + Name: "my-skill", + Source: "ghcr.io/org/skill", // tag-less: the most common upgrade case + ResolvedReference: "ghcr.io/org/skill:latest", + Digest: ociTestDigest(1), + } + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{}, entry) + + assert.Equal(t, skills.UpgradeStatusUpgraded, plan.outcome.Status, + "a moved digest on a tag-less source is an upgrade, not a blocked ref change") + assert.Equal(t, ociTestDigest(2), plan.outcome.NewDigest) + assert.Equal(t, "ghcr.io/org/skill:latest", plan.resolvedRef, + "upgrade must resolve to the same qualified form install records") +} + +func TestPlanUpgrade_OCITaggedSourceUpToDate(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "ghcr.io/org/skill:v1"). + Return(godigest.Digest(ociTestDigest(1)), nil) + svc := newOCIUpgradeService(t, reg, nil) + + entry := lockfile.Entry{ + Name: "my-skill", + Source: "ghcr.io/org/skill:v1", + ResolvedReference: "ghcr.io/org/skill:v1", + Digest: ociTestDigest(1), + } + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{}, entry) + assert.Equal(t, skills.UpgradeStatusUpToDate, plan.outcome.Status) +} + +// TestPlanUpgrade_OCIRefChangeGuard exercises the guard on a genuine +// reference change (the recorded resolvedReference no longer matches what +// the source resolves to): blocked by default, permitted with +// --allow-ref-change. +func TestPlanUpgrade_OCIRefChangeGuard(t *testing.T) { + t.Parallel() + + entry := lockfile.Entry{ + Name: "my-skill", + Source: "ghcr.io/org/skill", + ResolvedReference: "ghcr.io/org/old-location:latest", // recorded before the repoint + Digest: ociTestDigest(1), + } + + t.Run("blocked without allow-ref-change", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "ghcr.io/org/skill"). + Return(godigest.Digest(ociTestDigest(2)), nil) + svc := newOCIUpgradeService(t, reg, nil) + + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{}, entry) + assert.Equal(t, skills.UpgradeStatusRefChangeBlocked, plan.outcome.Status) + assert.Equal(t, "ghcr.io/org/skill:latest", plan.outcome.NewResolvedReference) + assert.Empty(t, plan.pinnedRef, "a blocked plan must not carry anything to install") + }) + + t.Run("permitted with allow-ref-change", func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "ghcr.io/org/skill"). + Return(godigest.Digest(ociTestDigest(2)), nil) + svc := newOCIUpgradeService(t, reg, nil) + + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{AllowRefChange: true}, entry) + assert.Equal(t, skills.UpgradeStatusUpgraded, plan.outcome.Status) + assert.Equal(t, "ghcr.io/org/skill:latest", plan.outcome.NewResolvedReference) + assert.NotEmpty(t, plan.pinnedRef) + }) +} + +// TestPlanUpgrade_RegistryFallbackSourceUpgrades covers the dropped branch: +// a skill originally installed through Install's OCI-pull -> registry- +// catalogue fallback has an ambiguous "namespace/name" source that fails a +// direct pull. Without mirroring the fallback, such a skill installs +// cleanly but can never be upgraded — its resolution fails exactly as the +// direct pull did at install time. +func TestPlanUpgrade_RegistryFallbackSourceUpgrades(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + // The direct pull of the ambiguous source fails, as it did at install. + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "io.github.test/my-skill"). + Return(godigest.Digest(""), fmt.Errorf("name unknown")) + // The registry catalogue resolves it to a concrete OCI reference. + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "ghcr.io/test/my-skill:v2"). + Return(godigest.Digest(ociTestDigest(3)), nil) + lookup := regmocks.NewMockProvider(ctrl) + lookup.EXPECT().SearchSkills("my-skill").Return([]regtypes.Skill{ + { + Namespace: "io.github.test", + Name: "my-skill", + Packages: []regtypes.SkillPackage{ + {RegistryType: "oci", Identifier: "ghcr.io/test/my-skill:v2"}, + }, + }, + }, nil) + svc := newOCIUpgradeService(t, reg, lookup) + + entry := lockfile.Entry{ + Name: "my-skill", + Source: "io.github.test/my-skill", // registry catalogue name, ambiguous OCI shape + ResolvedReference: "ghcr.io/test/my-skill:v2", + Digest: ociTestDigest(1), + } + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{}, entry) + + assert.Equal(t, skills.UpgradeStatusUpgraded, plan.outcome.Status, + "a registry-fallback-installed skill must be upgradable through the same fallback") + assert.Equal(t, ociTestDigest(3), plan.outcome.NewDigest) +} + +// TestPlanUpgrade_TagMoveWithinRepositoryIsNotBlocked covers the routine +// case a catalog-sourced skill hits on every release: the version moves, so +// the resolved tag moves with it, but the artifact still comes from the same +// repository. Blocking that would make --allow-ref-change mandatory for +// ordinary upgrades and so disable the repository check it exists for (#6213). +func TestPlanUpgrade_TagMoveWithinRepositoryIsNotBlocked(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "ghcr.io/org/skill:v2"). + Return(godigest.Digest(ociTestDigest(2)), nil) + svc := newOCIUpgradeService(t, reg, nil) + + entry := lockfile.Entry{ + Name: "my-skill", + Source: "ghcr.io/org/skill:v2", + ResolvedReference: "ghcr.io/org/skill:v1", + Digest: ociTestDigest(1), + } + + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{}, entry) + assert.Equal(t, skills.UpgradeStatusUpgraded, plan.outcome.Status) + assert.Equal(t, "ghcr.io/org/skill:v2", plan.outcome.NewResolvedReference, + "the move is still reported even though it is permitted") + assert.NotEmpty(t, plan.pinnedRef, "a permitted plan must carry the pinned candidate") +} + +// TestPlanUpgrade_RegistryChangeStaysBlocked keeps the guard meaningful: the +// repository path is identical, only the registry host differs. +func TestPlanUpgrade_RegistryChangeStaysBlocked(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + reg := ocimocks.NewMockRegistryClient(ctrl) + reg.EXPECT().Pull(gomock.Any(), gomock.Any(), "elsewhere.io/org/skill:v1"). + Return(godigest.Digest(ociTestDigest(2)), nil) + svc := newOCIUpgradeService(t, reg, nil) + + entry := lockfile.Entry{ + Name: "my-skill", + Source: "elsewhere.io/org/skill:v1", + ResolvedReference: "ghcr.io/org/skill:v1", + Digest: ociTestDigest(1), + } + + plan := svc.planUpgrade(t.Context(), skills.UpgradeOptions{}, entry) + assert.Equal(t, skills.UpgradeStatusRefChangeBlocked, plan.outcome.Status) + assert.Empty(t, plan.pinnedRef) +} diff --git a/pkg/skills/skillsvc/upgrade_test.go b/pkg/skills/skillsvc/upgrade_test.go new file mode 100644 index 0000000000..90f7146624 --- /dev/null +++ b/pkg/skills/skillsvc/upgrade_test.go @@ -0,0 +1,257 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" +) + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_FeatureDisabledReturnsForbidden(t *testing.T) { + gr, _ := newGitResolverMock(t) + svc, projectRoot := newLockTestService(t, gr) + t.Setenv(skills.LockFileEnvVar, "false") + + _, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.Error(t, err) + assert.Equal(t, http.StatusForbidden, httperr.Code(err)) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_ReportsUpToDateWhenSourceUnchanged(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusUpToDate, result.Outcomes[0].Status) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_InstallsNewerContent(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + before := readLockfile(t, projectRoot) + beforeEntry, _ := before.Get("my-skill") + + // Republish newer content at the same fixture URL — same source, new commit. + fx.register("my-skill", gitSkillVersion("my-skill")) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + outcome := result.Outcomes[0] + assert.Equal(t, skills.UpgradeStatusUpgraded, outcome.Status) + assert.Equal(t, beforeEntry.Digest, outcome.OldDigest) + assert.NotEqual(t, outcome.OldDigest, outcome.NewDigest) + + after := readLockfile(t, projectRoot) + afterEntry, ok := after.Get("my-skill") + require.True(t, ok) + assert.Equal(t, outcome.NewDigest, afterEntry.Digest) + assert.Equal(t, beforeEntry.Source, afterEntry.Source, "Source must never be rewritten by upgrade") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_PreviewDoesNotInstall(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + before := readLockfile(t, projectRoot) + beforeEntry, _ := before.Get("my-skill") + + fx.register("my-skill", gitSkillVersion("my-skill")) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ //nolint:forcetypeassert + ProjectRoot: projectRoot, Preview: true, + }) + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusUpgraded, result.Outcomes[0].Status, "preview still reports what would happen") + + after := readLockfile(t, projectRoot) + afterEntry, ok := after.Get("my-skill") + require.True(t, ok) + assert.Equal(t, beforeEntry.Digest, afterEntry.Digest, "preview must not rewrite the lock file") +} + +// TestUpgrade_PreservesExistingClients guards against Upgrade silently +// expanding a skill to every detected client. Install without an explicit +// InstallOptions.Clients falls back to whatever the caller passed (here, +// a single client); Upgrade must default to that skill's already-installed +// Clients when opts.Clients is empty, matching Sync's reinstallPinned — not +// fall through to Install's own "no clients means every detected client" +// default. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_PreservesExistingClients(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + fx.register("my-skill", gitSkillVersion("my-skill")) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusUpgraded, result.Outcomes[0].Status) + + info, err := svc.Info(t.Context(), skills.InfoOptions{Name: "my-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot}) + require.NoError(t, err) + assert.Equal(t, []string{"claude-code"}, info.InstalledSkill.Clients, + "upgrade must preserve the skill's existing clients, not expand to every detected client") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_NotUpgradableForImmutableSource(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // Hand-edit the lock entry to pin an immutable (full commit hash) source, + // as if the user had originally installed with an explicit @. + lf := readLockfile(t, projectRoot) + entry, ok := lf.Get("my-skill") + require.True(t, ok) + entry.Source = ref + "@" + entry.Digest + lf.Upsert(entry) + require.NoError(t, lf.Save(mustOpenRoot(t, projectRoot))) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusNotUpgradable, result.Outcomes[0].Status) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_UnknownNameReturnsNotFound(t *testing.T) { + gr, _ := newGitResolverMock(t) + svc, projectRoot := newLockTestService(t, gr) + + _, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ //nolint:forcetypeassert + ProjectRoot: projectRoot, Names: []string{"does-not-exist"}, + }) + require.Error(t, err) + assert.Equal(t, http.StatusNotFound, httperr.Code(err)) +} + +// TestUpgrade_FailOnChangesReportsOutcomesWithoutError: the gate returns +// the full outcome set with a 200 rather than an error — callers (and the +// CLI's exit-code mapping) derive the freshness verdict from the outcomes, +// which also distinguishes "would change" from a genuine resolution +// failure and names the stale skills. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_FailOnChangesReportsOutcomesWithoutError(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("my-skill", gitSkill("my-skill")) + svc, projectRoot := newLockTestService(t, gr) + + ref, _ := gitRef("my-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + fx.register("my-skill", gitSkillVersion("my-skill")) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ //nolint:forcetypeassert + ProjectRoot: projectRoot, FailOnChanges: true, + }) + require.NoError(t, err, "fail-on-changes reports outcomes, it does not error") + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusUpgraded, result.Outcomes[0].Status, + "the outcome reports what would change so callers can name the stale skill") +} + +// TestUpgrade_FailOnChangesWithoutPreviewDoesNotMutateAnyEntry covers the +// core fix: --fail-on-changes must behave as a read-only CI gate even +// without --preview. Before the fix, Upgrade installed each target in +// sequence and only checked FailOnChanges after each one, so a multi-entry +// lock file with one changed skill would install that skill for real before +// reporting the conflict — a partial mutation the gate is supposed to +// prevent entirely. +// +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_FailOnChangesWithoutPreviewDoesNotMutateAnyEntry(t *testing.T) { + gr, fx := newGitResolverMock(t) + fx.register("changed-skill", gitSkill("changed-skill")) + fx.register("stable-skill", gitSkill("stable-skill")) + svc, projectRoot := newLockTestService(t, gr) + + changedRef, _ := gitRef("changed-skill") + stableRef, _ := gitRef("stable-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: changedRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: stableRef, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + before := readLockfile(t, projectRoot) + changedBefore, _ := before.Get("changed-skill") + stableBefore, _ := before.Get("stable-skill") + + // Republish newer content for only one of the two skills. + fx.register("changed-skill", gitSkillVersion("changed-skill")) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ //nolint:forcetypeassert + ProjectRoot: projectRoot, FailOnChanges: true, + }) + require.NoError(t, err) + require.Len(t, result.Outcomes, 2) + + after := readLockfile(t, projectRoot) + changedAfter, ok := after.Get("changed-skill") + require.True(t, ok) + assert.Equal(t, changedBefore.Digest, changedAfter.Digest, + "the changed skill must not be installed for real before --fail-on-changes reports the conflict") + stableAfter, ok := after.Get("stable-skill") + require.True(t, ok) + assert.Equal(t, stableBefore.Digest, stableAfter.Digest) +} diff --git a/pkg/skills/skillsvc/upgrade_verify_test.go b/pkg/skills/skillsvc/upgrade_verify_test.go new file mode 100644 index 0000000000..91628f0414 --- /dev/null +++ b/pkg/skills/skillsvc/upgrade_verify_test.go @@ -0,0 +1,143 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/verifier" + verifiermocks "github.com/stacklok/toolhive/pkg/skills/verifier/mocks" +) + +// otherSignerResult is a verification result from a different identity than +// signedResult's. +func otherSignerResult() *verifier.Result { + r := signedResult() + r.SignerIdentity = "/.github/workflows/other.yml" + return r +} + +// signerChangeFixture installs a signed skill, then republishes newer +// content at the same source and returns a service whose verifier reports +// the candidate as signed by candidate() (or unsigned when candidate +// returns nil, err). +func signerChangeFixture( + t *testing.T, + candidate func() (*verifier.Result, error), +) (skills.SkillService, string) { + t.Helper() + gr, fx := newGitResolverMock(t) + fx.register("guarded-skill", gitSkill("guarded-skill")) + + installs := 0 + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes(). + DoAndReturn(func(_ any, _, _ []byte, expected *lockfile.Provenance) (*verifier.Result, error) { + installs++ + if installs == 1 { + return signedResult(), nil // initial install (TOFU) + } + result, err := candidate() + if err != nil { + return nil, err + } + if expected != nil && expected.SignerIdentity != result.SignerIdentity { + return nil, verifier.ErrSignerMismatch + } + return result, nil + }) + mv.EXPECT().VerifyBundleOffline(gomock.Any(), gomock.Any(), gomock.Any()). + AnyTimes().Return(nil) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + ref, _ := gitRef("guarded-skill") + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, Scope: skills.ScopeProject, ProjectRoot: projectRoot, Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // Republish newer content at the same source so an upgrade is planned. + fx.register("guarded-skill", gitSkillVersion("guarded-skill")) + return svc, projectRoot +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_SignerChangeBlocked(t *testing.T) { + svc, projectRoot := signerChangeFixture(t, func() (*verifier.Result, error) { + return otherSignerResult(), nil + }) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusSignerChangeBlocked, result.Outcomes[0].Status) + assert.Equal(t, "/.github/workflows/other.yml", result.Outcomes[0].NewSignerIdentity) + + // Nothing installed, lock unchanged. + lf := readLockfile(t, projectRoot) + entry, ok := lf.Get("guarded-skill") + require.True(t, ok) + require.NotNil(t, entry.Provenance) + assert.Equal(t, testSignerIdentity, entry.Provenance.SignerIdentity) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_UnsignedCandidateBlockedAgainstSignedEntry(t *testing.T) { + svc, projectRoot := signerChangeFixture(t, func() (*verifier.Result, error) { + return nil, verifier.ErrUnsigned + }) + + result, err := svc.(*service).Upgrade(t.Context(), skills.UpgradeOptions{ProjectRoot: projectRoot}) //nolint:forcetypeassert + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusSignerChangeBlocked, result.Outcomes[0].Status) + assert.Empty(t, result.Outcomes[0].NewSignerIdentity, "an unsigned candidate has no identity to report") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_AllowSignerChangeRecordsNewIdentity(t *testing.T) { + svc, projectRoot := signerChangeFixture(t, func() (*verifier.Result, error) { + return otherSignerResult(), nil + }) + + result, err := svc.(*service).Upgrade(t.Context(), //nolint:forcetypeassert + skills.UpgradeOptions{ProjectRoot: projectRoot, AllowSignerChange: true}) + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusUpgraded, result.Outcomes[0].Status) + + lf := readLockfile(t, projectRoot) + entry, ok := lf.Get("guarded-skill") + require.True(t, ok) + require.NotNil(t, entry.Provenance) + assert.Equal(t, "/.github/workflows/other.yml", entry.Provenance.SignerIdentity, + "the explicit override must re-record the new identity") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestUpgrade_SignerChangePreviewParity(t *testing.T) { + svc, projectRoot := signerChangeFixture(t, func() (*verifier.Result, error) { + return otherSignerResult(), nil + }) + + result, err := svc.(*service).Upgrade(t.Context(), //nolint:forcetypeassert + skills.UpgradeOptions{ProjectRoot: projectRoot, Preview: true}) + require.NoError(t, err) + require.Len(t, result.Outcomes, 1) + assert.Equal(t, skills.UpgradeStatusSignerChangeBlocked, result.Outcomes[0].Status, + "preview must report the same signer-change block as apply") + + // Preview installed nothing and rewrote nothing. + lf := readLockfile(t, projectRoot) + entry, ok := lf.Get("guarded-skill") + require.True(t, ok) + assert.Equal(t, testSignerIdentity, entry.Provenance.SignerIdentity) +} diff --git a/pkg/skills/skillsvc/verify.go b/pkg/skills/skillsvc/verify.go new file mode 100644 index 0000000000..c11d0d595b --- /dev/null +++ b/pkg/skills/skillsvc/verify.go @@ -0,0 +1,301 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "context" + "errors" + "fmt" + "net/http" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/container/images" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/verifier" +) + +// artifactVerifier returns the configured signature verifier, defaulting to +// the Sigstore verifier with the composite registry keychain. +func (s *service) artifactVerifier() verifier.Verifier { + if s.sigVerifier != nil { + return s.sigVerifier + } + return verifier.NewDefault(images.NewCompositeKeychain()) +} + +// shouldVerifyInstall reports whether install-time signature verification +// applies: project-scope installs with the lock file feature enabled. The +// lock file is where trust decisions are recorded, so verification is +// scoped to it. +func shouldVerifyInstall(opts skills.InstallOptions, scope skills.Scope) bool { + return scope == skills.ScopeProject && opts.ProjectRoot != "" && skills.LockFileFeatureEnabled() +} + +// provenanceDecision is the outcome of install-time verification: either a +// verified identity (with the bundle backing it) or an explicit unsigned +// exception. +type provenanceDecision struct { + provenance *skills.ProvenanceInfo + unsigned bool + bundle []byte +} + +// applyDecisionToOpts records the verification outcome on the install +// options, from where it flows into the installed-skill record and the lock +// entry. +func applyDecisionToOpts(opts *skills.InstallOptions, decision *provenanceDecision) { + if decision == nil { + return + } + opts.Provenance = decision.provenance + opts.Unsigned = decision.unsigned + opts.SigstoreBundle = decision.bundle +} + +// verifyOCIInstall verifies the signature of the OCI artifact at ref/digest +// before anything is extracted or recorded. The identity expected by the +// lock file (if any) is enforced inside the verifier's Sigstore policy; +// trust on first use records whatever identity verification observes. +func (s *service) verifyOCIInstall( + ctx context.Context, + opts skills.InstallOptions, + skillName, ref, digest string, +) (*provenanceDecision, error) { + expected, expectUnsigned, err := expectedLockTrust(opts.ProjectRoot, skillName) + if err != nil { + return nil, err + } + if opts.AllowSignerChange { + // The signer-change guard was explicitly overridden: verify the + // chain of trust only and re-record whatever identity is observed. + expected, expectUnsigned = nil, false + } + if expectUnsigned { + return unsignedLockedDecision(opts, skillName) + } + + result, verifyErr := s.artifactVerifier().VerifyOCI(ctx, ref, digest, expected) + if verifyErr != nil { + if isAllowedUnsigned(verifyErr, opts, expected) { + return &provenanceDecision{unsigned: true}, nil + } + return nil, classifyInstallVerifyError(verifyErr, skillName, expected) + } + return &provenanceDecision{ + provenance: provenanceInfoFromResult(result), + bundle: result.Bundle, + }, nil +} + +// verifyGitInstall verifies the gitsign signature on the resolved commit +// before anything is written or recorded. +func (s *service) verifyGitInstall( + ctx context.Context, + opts skills.InstallOptions, + skillName string, + payload []byte, + signature string, +) (*provenanceDecision, error) { + expected, expectUnsigned, err := expectedLockTrust(opts.ProjectRoot, skillName) + if err != nil { + return nil, err + } + if opts.AllowSignerChange { + expected, expectUnsigned = nil, false + } + if expectUnsigned { + return unsignedLockedDecision(opts, skillName) + } + + result, verifyErr := s.artifactVerifier().VerifyGit(ctx, payload, []byte(signature), expected) + if verifyErr != nil { + if isAllowedUnsigned(verifyErr, opts, expected) { + return &provenanceDecision{unsigned: true}, nil + } + return nil, classifyInstallVerifyError(verifyErr, skillName, expected) + } + return &provenanceDecision{ + provenance: provenanceInfoFromResult(result), + bundle: result.Bundle, + }, nil +} + +// verifyLocalInstall handles installs sourced from the local OCI store or +// raw layer data: there is no registry signature to verify, so the install +// is an unsigned trust decision. An entry already locked to a signer +// identity refuses a local replacement outright — swapping a verified +// artifact for a local build is exactly the substitution the lock exists to +// catch. +func verifyLocalInstall(opts skills.InstallOptions, skillName string) (*provenanceDecision, error) { + expected, expectUnsigned, err := expectedLockTrust(opts.ProjectRoot, skillName) + if err != nil { + return nil, err + } + if expected != nil { + return nil, httperr.WithCode( + fmt.Errorf("skill %q is locked to signer %q; a local build cannot satisfy it", + skillName, expected.SignerIdentity), + http.StatusForbidden, + ) + } + if expectUnsigned { + return unsignedLockedDecision(opts, skillName) + } + if !opts.AllowUnsigned { + return nil, httperr.WithCode( + fmt.Errorf("local build for %q is unsigned; set allow_unsigned (--allow-unsigned) to record an exception", + skillName), + http.StatusForbidden, + ) + } + return &provenanceDecision{unsigned: true}, nil +} + +// unsignedLockedDecision handles installs of entries the lock file already +// marks unsigned. A lock-driven operation honors the recorded decision (the +// lock file IS the policy it restores); a fresh user-driven install must +// repeat the explicit exception. +// +// SECURITY: this early return means an entry marked unsigned is installed +// without consulting the verifier at all — by design, but it makes a lock +// diff that converts `provenance:` to `unsigned: true` a trust DOWNGRADE +// that sync will honor silently. That conversion is exactly what lock file +// review must catch; it cannot happen without a lock file edit. +func unsignedLockedDecision(opts skills.InstallOptions, skillName string) (*provenanceDecision, error) { + if opts.AllowUnsigned || lockDrivenInstall(opts) { + return &provenanceDecision{unsigned: true}, nil + } + return nil, httperr.WithCode( + fmt.Errorf("skill %q is locked as unsigned; set allow_unsigned (--allow-unsigned) to reinstall it", skillName), + http.StatusForbidden, + ) +} + +// expectedLockTrust reads the trust state recorded in projectRoot's lock +// file for skillName: the expected signer identity (nil on first use — the +// TOFU case), or that the entry was recorded unsigned. +func expectedLockTrust(projectRoot, skillName string) (*lockfile.Provenance, bool, error) { + if projectRoot == "" { + return nil, false, nil + } + root, err := lockfile.OpenRoot(projectRoot) + if err != nil { + return nil, false, err + } + lf, err := lockfile.Load(root) + if err != nil { + return nil, false, err + } + entry, ok := lf.Get(skillName) + if !ok { + return nil, false, nil + } + if entry.Unsigned { + return nil, true, nil + } + return entry.Provenance, false, nil +} + +// isAllowedUnsigned reports whether a verification failure is the unsigned +// case AND the caller may proceed: either the explicit --allow-unsigned +// exception, or a sync restore of an entry with no recorded trust state +// (entries created before verification existed) — a restore materializes +// what install once accepted, and the outcome is recorded as unsigned so +// the trust state stops being ambiguous. An entry locked to a signer +// identity is never replaceable by an unsigned artifact. +func isAllowedUnsigned(verifyErr error, opts skills.InstallOptions, expected *lockfile.Provenance) bool { + if !errors.Is(verifyErr, verifier.ErrUnsigned) || expected != nil { + return false + } + return opts.AllowUnsigned || lockDrivenInstall(opts) +} + +// lockDrivenInstall reports whether this install materializes an existing +// lock entry rather than making a new trust decision: sync restores and +// upgrade re-pins (both set internal-only options). Such operations honor +// the trust state the entry already records. +func lockDrivenInstall(opts skills.InstallOptions) bool { + return opts.SyncRestore || opts.LockResolvedReference != "" +} + +// classifyInstallVerifyError maps a verifier failure to the HTTP-coded +// error surfaced by the install API — always a 403; the allowed-unsigned +// path is handled before this is called. +func classifyInstallVerifyError( + verifyErr error, + skillName string, + expected *lockfile.Provenance, +) error { + switch { + case errors.Is(verifyErr, verifier.ErrUnsigned): + if expected != nil { + return httperr.WithCode( + fmt.Errorf("skill %q is locked to signer %q but the artifact is unsigned", + skillName, expected.SignerIdentity), + http.StatusForbidden, + ) + } + return httperr.WithCode( + fmt.Errorf("unsigned skill %q rejected; set allow_unsigned (--allow-unsigned) to record an exception", + skillName), + http.StatusForbidden, + ) + case errors.Is(verifyErr, verifier.ErrSignerMismatch): + return httperr.WithCode( + fmt.Errorf("signer identity mismatch for %q: %w"+ + " (if the signer change is intended, remove the skill's lock entry"+ + " and reinstall, or upgrade with allow_signer_change)", skillName, verifyErr), + http.StatusForbidden, + ) + default: + return httperr.WithCode( + fmt.Errorf("signature verification failed for %q: %w", skillName, verifyErr), + http.StatusForbidden, + ) + } +} + +// classifySignatureError maps verifier sentinels to typed failure reasons +// for sync/upgrade results. Returns "" when err is not a signature failure. +func classifySignatureError(err error) skills.FailureReason { + switch { + case errors.Is(err, verifier.ErrSignerMismatch): + return skills.FailureReasonSignerMismatch + case errors.Is(err, verifier.ErrUnsigned): + return skills.FailureReasonUnsignedRejected + case errors.Is(err, verifier.ErrSignatureInvalid): + return skills.FailureReasonSignatureInvalid + default: + return "" + } +} + +// provenanceInfoToLock converts the internal provenance shape to the lock +// file's. +func provenanceInfoToLock(p *skills.ProvenanceInfo) *lockfile.Provenance { + if p == nil { + return nil + } + return &lockfile.Provenance{ + SignerIdentity: p.SignerIdentity, + CertIssuer: p.CertIssuer, + RepositoryURI: p.RepositoryURI, + SigstoreURL: p.SigstoreURL, + } +} + +// provenanceInfoFromResult converts a verification result into the internal +// provenance shape recorded on install options. +func provenanceInfoFromResult(r *verifier.Result) *skills.ProvenanceInfo { + if r == nil || !r.Signed || r.SignerIdentity == "" { + return nil + } + return &skills.ProvenanceInfo{ + SignerIdentity: r.SignerIdentity, + CertIssuer: r.CertIssuer, + RepositoryURI: r.RepositoryURI, + SigstoreURL: r.SigstoreURL, + } +} diff --git a/pkg/skills/skillsvc/verify_test.go b/pkg/skills/skillsvc/verify_test.go new file mode 100644 index 0000000000..8c6e02bce4 --- /dev/null +++ b/pkg/skills/skillsvc/verify_test.go @@ -0,0 +1,321 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package skillsvc + +import ( + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/httperr" + "github.com/stacklok/toolhive/pkg/skills" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/verifier" + verifiermocks "github.com/stacklok/toolhive/pkg/skills/verifier/mocks" +) + +const ( + testSignerIdentity = "/.github/workflows/release.yml" + testCertIssuer = "https://token.actions.githubusercontent.com" +) + +func signedResult() *verifier.Result { + return &verifier.Result{ + Signed: true, + SignerIdentity: testSignerIdentity, + CertIssuer: testCertIssuer, + RepositoryURI: "https://github.com/org/repo", + SigstoreURL: "https://rekor.sigstore.dev", + Bundle: []byte(`{"bundle":true}`), + } +} + +// loadLockEntry reads the lock entry for name from projectRoot. +func loadLockEntry(t *testing.T, projectRoot, name string) (lockfile.Entry, bool) { + t.Helper() + root, err := lockfile.OpenRoot(projectRoot) + require.NoError(t, err) + lf, err := lockfile.Load(root) + require.NoError(t, err) + return lf.Get(name) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallVerification_TOFURecordsProvenance(t *testing.T) { + gr, fixtures := newGitResolverMock(t) + ref, _ := gitRef("signed-skill") + fixtures.register("signed-skill", gitSkill("signed-skill")) + + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + // First install: no lock entry yet — trust on first use, nil expected. + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Nil()). + Return(signedResult(), nil) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + result, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + entry, ok := loadLockEntry(t, projectRoot, "signed-skill") + require.True(t, ok) + require.NotNil(t, entry.Provenance, "TOFU must record the observed identity") + assert.Equal(t, testSignerIdentity, entry.Provenance.SignerIdentity) + assert.Equal(t, testCertIssuer, entry.Provenance.CertIssuer) + assert.False(t, entry.Unsigned) + assert.Equal(t, []byte(`{"bundle":true}`), result.Skill.SigstoreBundle, + "the bundle must be persisted with the install record") + + // Second install: the recorded identity must flow into the verifier as + // the expected identity. + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ any, _, _ []byte, expected *lockfile.Provenance) (*verifier.Result, error) { + require.NotNil(t, expected, "the second install must enforce the recorded identity") + assert.Equal(t, testSignerIdentity, expected.SignerIdentity) + return signedResult(), nil + }) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + Force: true, + }) + require.NoError(t, err) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallVerification_UnsignedRejectedWithoutFlag(t *testing.T) { + gr, fixtures := newGitResolverMock(t) + ref, _ := gitRef("unsigned-skill") + fixtures.register("unsigned-skill", gitSkill("unsigned-skill")) + + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Nil()). + Return(nil, verifier.ErrUnsigned) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + }) + require.Error(t, err) + assert.Equal(t, http.StatusForbidden, httperr.Code(err)) + + _, ok := loadLockEntry(t, projectRoot, "unsigned-skill") + assert.False(t, ok, "a rejected install must not write a lock entry") + _, err = svc.Info(t.Context(), skills.InfoOptions{ + Name: "unsigned-skill", Scope: skills.ScopeProject, ProjectRoot: projectRoot, + }) + require.Error(t, err, "a rejected install must not create a DB record") +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallVerification_UnsignedAcceptedWithFlag(t *testing.T) { + gr, fixtures := newGitResolverMock(t) + ref, _ := gitRef("unsigned-ok") + fixtures.register("unsigned-ok", gitSkill("unsigned-ok")) + + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Nil()). + Return(nil, verifier.ErrUnsigned) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + AllowUnsigned: true, + }) + require.NoError(t, err) + + entry, ok := loadLockEntry(t, projectRoot, "unsigned-ok") + require.True(t, ok) + assert.True(t, entry.Unsigned, "the unsigned exception must be recorded") + assert.Nil(t, entry.Provenance) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallVerification_SignerMismatchRejectedAndLockIntact(t *testing.T) { + gr, fixtures := newGitResolverMock(t) + ref, _ := gitRef("pinned-skill") + fixtures.register("pinned-skill", gitSkill("pinned-skill")) + + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Nil()). + Return(signedResult(), nil) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) + + // The re-install is signed by someone else: the verifier reports a + // mismatch (the expected identity was bound into its policy). + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, verifier.ErrSignerMismatch) + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + Force: true, + }) + require.Error(t, err) + assert.Equal(t, http.StatusForbidden, httperr.Code(err)) + + // The prior trusted state is untouched. + entry, ok := loadLockEntry(t, projectRoot, "pinned-skill") + require.True(t, ok) + require.NotNil(t, entry.Provenance) + assert.Equal(t, testSignerIdentity, entry.Provenance.SignerIdentity) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallVerification_LockedUnsignedRequiresFlagAgain(t *testing.T) { + gr, fixtures := newGitResolverMock(t) + ref, _ := gitRef("unsigned-locked") + fixtures.register("unsigned-locked", gitSkill("unsigned-locked")) + + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + mv.EXPECT().VerifyGit(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Nil()). + Return(nil, verifier.ErrUnsigned) + + svc, projectRoot := newLockTestService(t, gr, WithVerifier(mv)) + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + AllowUnsigned: true, + }) + require.NoError(t, err) + + // Reinstall without the flag: the locked unsigned exception does not + // silently renew — the verifier is not even consulted. + _, err = svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeProject, + ProjectRoot: projectRoot, + Clients: []string{"claude-code"}, + Force: true, + }) + require.Error(t, err) + assert.Equal(t, http.StatusForbidden, httperr.Code(err)) +} + +//nolint:paralleltest // uses t.Setenv via newLockTestService, incompatible with t.Parallel +func TestInstallVerification_UserScopeSkipsVerification(t *testing.T) { + gr, fixtures := newGitResolverMock(t) + ref, _ := gitRef("user-skill") + fixtures.register("user-skill", gitSkill("user-skill")) + + // The mock has no expectations: any verifier call fails the test. + mv := verifiermocks.NewMockVerifier(gomock.NewController(t)) + + svc, _ := newLockTestService(t, gr, WithVerifier(mv)) + _, err := svc.Install(t.Context(), skills.InstallOptions{ + Name: ref, + Scope: skills.ScopeUser, + Clients: []string{"claude-code"}, + }) + require.NoError(t, err) +} + +func TestVerifyLocalInstall(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts skills.InstallOptions + entry *lockfile.Entry + wantErr bool + unsigned bool + }{ + { + name: "no flag rejected", + opts: skills.InstallOptions{}, + wantErr: true, + }, + { + name: "flag records unsigned", + opts: skills.InstallOptions{AllowUnsigned: true}, + unsigned: true, + }, + { + name: "locked identity refuses local replacement even with flag", + opts: skills.InstallOptions{AllowUnsigned: true}, + entry: &lockfile.Entry{ + Name: "local-skill", + Source: "example.com/org/local-skill", + ResolvedReference: "example.com/org/local-skill:v1", + Digest: "sha256:" + strings.Repeat("a", 64), + Provenance: &lockfile.Provenance{ + SignerIdentity: testSignerIdentity, + CertIssuer: testCertIssuer, + }, + }, + wantErr: true, + }, + { + name: "locked unsigned honored with flag", + opts: skills.InstallOptions{AllowUnsigned: true}, + entry: &lockfile.Entry{ + Name: "local-skill", + Source: "example.com/org/local-skill", + ResolvedReference: "example.com/org/local-skill:v1", + Digest: "sha256:" + strings.Repeat("a", 64), + Unsigned: true, + }, + unsigned: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + projectRoot := makeProjectRoot(t) + if tc.entry != nil { + root, err := lockfile.OpenRoot(projectRoot) + require.NoError(t, err) + require.NoError(t, lockfile.Update(root, func(lf *lockfile.Lockfile) error { + lf.Upsert(*tc.entry) + return nil + })) + } + opts := tc.opts + opts.ProjectRoot = projectRoot + + decision, err := verifyLocalInstall(opts, "local-skill") + if tc.wantErr { + require.Error(t, err) + assert.Equal(t, http.StatusForbidden, httperr.Code(err)) + return + } + require.NoError(t, err) + assert.Equal(t, tc.unsigned, decision.unsigned) + }) + } +} + +func TestClassifySignatureError(t *testing.T) { + t.Parallel() + assert.Equal(t, skills.FailureReasonSignerMismatch, classifySignatureError(verifier.ErrSignerMismatch)) + assert.Equal(t, skills.FailureReasonUnsignedRejected, classifySignatureError(verifier.ErrUnsigned)) + assert.Equal(t, skills.FailureReasonSignatureInvalid, classifySignatureError(verifier.ErrSignatureInvalid)) + assert.Equal(t, skills.FailureReason(""), classifySignatureError(assert.AnError)) +} diff --git a/pkg/skills/types.go b/pkg/skills/types.go index 224e7b8d00..50b4a27dd8 100644 --- a/pkg/skills/types.go +++ b/pkg/skills/types.go @@ -167,6 +167,14 @@ type InstalledSkill struct { Clients []string `json:"clients,omitempty"` // Dependencies is the list of external skill dependencies. Dependencies []Dependency `json:"dependencies,omitempty"` + // Managed indicates this install is tracked in the project's + // toolhive.lock.yaml. Only ever true for project-scoped installs. No + // omitempty: false is an observable state (unmanaged), not an absence. + Managed bool `json:"managed"` + // SigstoreBundle is the serialized Sigstore bundle captured at install + // time, used for offline re-verification during sync. Nil for unsigned + // installs. Never serialized to API responses. + SigstoreBundle []byte `json:"-"` } // SkillIndexEntry represents a single skill entry in a remote skill index. diff --git a/pkg/skills/verifier/errors.go b/pkg/skills/verifier/errors.go new file mode 100644 index 0000000000..854991df31 --- /dev/null +++ b/pkg/skills/verifier/errors.go @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import "errors" + +var ( + // ErrUnsigned indicates the artifact carries no Sigstore signature + // material in any supported layout. + ErrUnsigned = errors.New("artifact is not signed") + // ErrSignatureInvalid indicates signature material was found but failed + // cryptographic verification (or a stored bundle is malformed). + ErrSignatureInvalid = errors.New("signature verification failed") + // ErrSignerMismatch indicates the signature verifies, but against an + // identity other than the expected one. + ErrSignerMismatch = errors.New("signer identity mismatch") +) diff --git a/pkg/skills/verifier/git.go b/pkg/skills/verifier/git.go new file mode 100644 index 0000000000..cb9c495071 --- /dev/null +++ b/pkg/skills/verifier/git.go @@ -0,0 +1,202 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import ( + "context" + "crypto/x509" + "encoding/pem" + "errors" + "fmt" + "sync" + "time" + + cms "github.com/github/smimesign/ietf-cms" + "github.com/sigstore/sigstore-go/pkg/fulcio/certificate" + "github.com/sigstore/sigstore-go/pkg/root" + "github.com/sigstore/sigstore-go/pkg/verify" + + coreverifier "github.com/stacklok/toolhive-core/container/verifier" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// VerifyGit cryptographically verifies a gitsign commit signature: the CMS +// signature is checked over the commit payload and the signing certificate +// chain is verified against the Fulcio roots in toolhive-core's embedded +// trusted material — no network. A non-nil expected identity must match the +// certificate identity (post-hoc comparison: git signatures have no +// Sigstore bundle to bind a policy into); nil expected is trust on first +// use. +// +// The embedded Rekor transparency-log proof is NOT yet validated — signing +// time is checked against the certificate's own validity window, matching +// gitsign's certificate verifier. Rekor proof validation is a tracked +// follow-up (it requires reconstructing the proof from CMS unsigned +// attributes, which gitsign only exposes internally). +func (*Default) VerifyGit( + ctx context.Context, + payload, signature []byte, + expected *lockfile.Provenance, +) (*Result, error) { + if len(signature) == 0 { + return nil, fmt.Errorf("%w: commit is not signed", ErrUnsigned) + } + if len(payload) == 0 { + return nil, fmt.Errorf("%w: no commit payload to verify", ErrSignatureInvalid) + } + + roots, intermediates, err := fulcioPools() + if err != nil { + return nil, err + } + cert, err := verifyGitSignature(ctx, payload, signature, roots, intermediates) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrSignatureInvalid, err) + } + + identity, err := gitIdentityFromCertificate(cert) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrSignatureInvalid, err) + } + if expected != nil && !gitIdentityMatches(identity, expected) { + return nil, fmt.Errorf("%w: commit is signed by a different identity than %q", + ErrSignerMismatch, expected.SignerIdentity) + } + result := resultFromCore(identity, nil) + // No transparency-log proof is validated yet (see the doc comment), so + // git verification carries a documented assurance gap — marked so the + // lock file shows it rather than implying full verification. + result.Provisional = true + return result, nil +} + +// verifyGitSignature verifies the CMS (PKCS#7) signature gitsign attaches +// to commits: the armored signature is parsed, the message digest and +// signature are checked over the commit payload, and the signing +// certificate chain is verified against the given pools. It returns the +// certificate the signature was authenticated against. +// +// This mirrors gitsign's own CertVerifier (which wraps a fork of the same +// ietf-cms package) without importing the gitsign module, whose package +// graph compiles cosign and the cloud KMS SDKs into the binary. The +// verification time is anchored inside the signing certificate's own +// validity window — Fulcio certificates live for minutes, so "valid at +// time.Now" would reject every real signature; proving the actual signing +// time is the transparency log's job (tracked follow-up). +func verifyGitSignature( + _ context.Context, + payload, signature []byte, + roots, intermediates *x509.CertPool, +) (*x509.Certificate, error) { + // ietf-cms v0.2.0 mutates package-level state during BER decoding + // (protocol.encodeIndent), so concurrent ParseSignedData calls race. + // Verification is not hot-path — serialize instead of forking the + // library like gitsign did. + cmsMu.Lock() + defer cmsMu.Unlock() + + der := signature + if blk, _ := pem.Decode(signature); blk != nil { + der = blk.Bytes + } + sd, err := cms.ParseSignedData(der) + if err != nil { + return nil, fmt.Errorf("parsing CMS signature: %w", err) + } + certs, err := sd.GetCertificates() + if err != nil { + return nil, fmt.Errorf("reading signature certificates: %w", err) + } + leaf := leafCertificate(certs) + if leaf == nil { + return nil, errors.New("no signing certificate found in signature") + } + + chains, err := sd.VerifyDetached(payload, x509.VerifyOptions{ + Roots: roots, + Intermediates: intermediates, + KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning}, + CurrentTime: leaf.NotBefore.Add(time.Minute), + }) + if err != nil { + return nil, fmt.Errorf("verifying signature: %w", err) + } + // Return the leaf of the first verified chain — the certificate the + // CMS verifier actually authenticated the signature against. + if len(chains) == 0 || len(chains[0]) == 0 || len(chains[0][0]) == 0 { + return nil, errors.New("no verified certificate chains returned") + } + return chains[0][0][0], nil +} + +// cmsMu serializes ietf-cms parsing/verification; see verifyGitSignature. +var cmsMu sync.Mutex + +// leafCertificate picks the end-entity certificate from a CMS certificate +// bag (the one that is not a CA), falling back to the first entry. +func leafCertificate(certs []*x509.Certificate) *x509.Certificate { + for _, c := range certs { + if !c.IsCA { + return c + } + } + if len(certs) > 0 { + return certs[0] + } + return nil +} + +// fulcioPools builds root and intermediate certificate pools from the +// Fulcio certificate authorities in the embedded trusted material. +func fulcioPools() (roots, intermediates *x509.CertPool, err error) { + tm, err := coreverifier.OfflineTrustedMaterial() + if err != nil { + return nil, nil, fmt.Errorf("loading trusted material: %w", err) + } + roots = x509.NewCertPool() + intermediates = x509.NewCertPool() + for _, ca := range tm.FulcioCertificateAuthorities() { + fulcioCA, ok := ca.(*root.FulcioCertificateAuthority) + if !ok || fulcioCA.Root == nil { + continue + } + roots.AddCert(fulcioCA.Root) + for _, intermediate := range fulcioCA.Intermediates { + intermediates.AddCert(intermediate) + } + } + return roots, intermediates, nil +} + +// gitIdentityFromCertificate extracts the signer identity from a gitsign +// certificate, reusing core's identity normalization (GitHub Actions +// workflow paths, issuer extraction) by summarizing the certificate the +// same way sigstore-go does for bundle verification results. +func gitIdentityFromCertificate(cert *x509.Certificate) (coreverifier.Identity, error) { + summary, err := certificate.SummarizeCertificate(cert) + if err != nil { + return coreverifier.Identity{}, fmt.Errorf("summarizing certificate: %w", err) + } + return coreverifier.IdentityFromResult(&verify.VerificationResult{ + Signature: &verify.SignatureVerificationResult{Certificate: &summary}, + }) +} + +// gitIdentityMatches compares the observed certificate identity against the +// lock file's expected provenance. Unlike the OCI path there is no Sigstore +// policy to bind the identity into, so the comparison is explicit; empty +// expected fields are not wildcards — SignerIdentity and CertIssuer are +// always recorded together at trust-on-first-use. +func gitIdentityMatches(observed coreverifier.Identity, expected *lockfile.Provenance) bool { + if observed.SignerIdentity != expected.SignerIdentity { + return false + } + if observed.CertIssuer != expected.CertIssuer { + return false + } + if expected.RepositoryURI != "" && observed.SourceRepositoryURI != expected.RepositoryURI { + return false + } + return true +} diff --git a/pkg/skills/verifier/git_test.go b/pkg/skills/verifier/git_test.go new file mode 100644 index 0000000000..b230ecb097 --- /dev/null +++ b/pkg/skills/verifier/git_test.go @@ -0,0 +1,275 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "math/big" + "testing" + "time" + + cms "github.com/github/smimesign/ietf-cms" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + coreverifier "github.com/stacklok/toolhive-core/container/verifier" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// fulcioIssuerOID is the Fulcio extension carrying the OIDC issuer +// (1.3.6.1.4.1.57264.1.1, the raw-string form gitsign certificates use). +var fulcioIssuerOID = []int{1, 3, 6, 1, 4, 1, 57264, 1, 1} + +const testGitIssuer = "https://accounts.example.com" + +// testCA is a synthetic Fulcio-style CA for signing test certificates. +type testCA struct { + cert *x509.Certificate + key *ecdsa.PrivateKey +} + +func newTestCA(t *testing.T) *testCA { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "test-fulcio-root"}, + NotBefore: time.Now().Add(-48 * time.Hour), + NotAfter: time.Now().Add(time.Hour), + IsCA: true, + KeyUsage: x509.KeyUsageCertSign, + BasicConstraintsValid: true, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, key.Public(), key) + require.NoError(t, err) + cert, err := x509.ParseCertificate(der) + require.NoError(t, err) + return &testCA{cert: cert, key: key} +} + +// issueSigningCert issues a short-lived Fulcio-style code-signing +// certificate with the given email SAN and OIDC issuer extension. +func (ca *testCA) issueSigningCert(t *testing.T, email string) (*x509.Certificate, *ecdsa.PrivateKey) { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(2), + EmailAddresses: []string{email}, + // Fulcio certificates live for minutes. + NotBefore: time.Now().Add(-time.Minute), + NotAfter: time.Now().Add(9 * time.Minute), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning}, + ExtraExtensions: []pkix.Extension{ + {Id: fulcioIssuerOID, Value: []byte(testGitIssuer)}, + }, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.cert, key.Public(), ca.key) + require.NoError(t, err) + cert, err := x509.ParseCertificate(der) + require.NoError(t, err) + return cert, key +} + +// issueExpiredSigningCert issues a signing certificate whose validity +// window lies entirely in the past. +func (ca *testCA) issueExpiredSigningCert(t *testing.T, email string) (*x509.Certificate, *ecdsa.PrivateKey) { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(3), + EmailAddresses: []string{email}, + NotBefore: time.Now().Add(-24 * time.Hour), + NotAfter: time.Now().Add(-23 * time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageCodeSigning}, + ExtraExtensions: []pkix.Extension{ + {Id: fulcioIssuerOID, Value: []byte(testGitIssuer)}, + }, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.cert, key.Public(), ca.key) + require.NoError(t, err) + cert, err := x509.ParseCertificate(der) + require.NoError(t, err) + return cert, key +} + +func (ca *testCA) pool() *x509.CertPool { + pool := x509.NewCertPool() + pool.AddCert(ca.cert) + return pool +} + +// signCommitPayload produces the armored CMS detached signature gitsign +// attaches to commits. +func signCommitPayload(t *testing.T, payload []byte, cert *x509.Certificate, key *ecdsa.PrivateKey) []byte { + t.Helper() + der, err := cms.SignDetached(payload, []*x509.Certificate{cert}, key) + require.NoError(t, err) + return pem.EncodeToMemory(&pem.Block{Type: "SIGNED MESSAGE", Bytes: der}) +} + +const testCommitPayload = "tree 0123456789abcdef0123456789abcdef01234567\n" + + "author Test 1700000000 +0000\n" + + "committer Test 1700000000 +0000\n\n" + + "Signed test commit\n" + +func TestVerifyGitSignatureRoundTrip(t *testing.T) { + t.Parallel() + ca := newTestCA(t) + cert, key := ca.issueSigningCert(t, "dev@example.com") + payload := []byte(testCommitPayload) + sig := signCommitPayload(t, payload, cert, key) + + got, err := verifyGitSignature(t.Context(), payload, sig, ca.pool(), x509.NewCertPool()) + require.NoError(t, err, "a signature chaining to the trusted root must verify") + assert.Equal(t, cert.SerialNumber, got.SerialNumber, + "the authenticated certificate must be the signing leaf") + + // The identity extracted from the certificate carries the SAN and the + // Fulcio issuer extension. + identity, err := gitIdentityFromCertificate(got) + require.NoError(t, err) + assert.Equal(t, "dev@example.com", identity.SignerIdentity) + assert.Equal(t, testGitIssuer, identity.CertIssuer) +} + +func TestVerifyGitSignatureRejects(t *testing.T) { + t.Parallel() + ca := newTestCA(t) + cert, key := ca.issueSigningCert(t, "dev@example.com") + payload := []byte(testCommitPayload) + sig := signCommitPayload(t, payload, cert, key) + + t.Run("tampered payload", func(t *testing.T) { + t.Parallel() + tampered := []byte(testCommitPayload + "tampered\n") + _, err := verifyGitSignature(t.Context(), tampered, sig, ca.pool(), x509.NewCertPool()) + require.Error(t, err, "a signature must not verify over modified commit content") + }) + + t.Run("untrusted root", func(t *testing.T) { + t.Parallel() + otherCA := newTestCA(t) + _, err := verifyGitSignature(t.Context(), payload, sig, otherCA.pool(), x509.NewCertPool()) + require.Error(t, err, "a chain to an untrusted root must not verify") + }) + + t.Run("garbage signature", func(t *testing.T) { + t.Parallel() + _, err := verifyGitSignature(t.Context(), payload, []byte("not a signature"), ca.pool(), x509.NewCertPool()) + require.Error(t, err) + }) +} + +// TestVerifyGitSignatureAcceptsExpiredCertificate pins the current time +// anchoring: verification time is the certificate's own validity window, so +// a certificate that expired long ago still verifies. This is the +// documented assurance gap of the git path (no transparency-log proof of +// signing time — provenance is marked Provisional for it). +// +// TODO(rekor): flip this expectation to require.Error once Rekor +// inclusion-proof validation lands and bounds the replay window. +func TestVerifyGitSignatureAcceptsExpiredCertificate(t *testing.T) { + t.Parallel() + ca := newTestCA(t) + cert, key := ca.issueExpiredSigningCert(t, "dev@example.com") + payload := []byte(testCommitPayload) + sig := signCommitPayload(t, payload, cert, key) + + _, err := verifyGitSignature(t.Context(), payload, sig, ca.pool(), x509.NewCertPool()) + require.NoError(t, err, + "expired certificates currently verify (validity-window anchoring); see TODO(rekor)") +} + +func TestVerifyGitInputGuards(t *testing.T) { + t.Parallel() + d := NewDefault(nil) + + _, err := d.VerifyGit(t.Context(), []byte("payload"), nil, nil) + require.ErrorIs(t, err, ErrUnsigned, "an empty signature is the unsigned case") + + _, err = d.VerifyGit(t.Context(), nil, []byte("sig"), nil) + require.ErrorIs(t, err, ErrSignatureInvalid, "a signature without payload cannot verify") + + // A syntactically broken signature fails as invalid (the embedded + // Fulcio pools load fine; parsing fails first). + _, err = d.VerifyGit(t.Context(), []byte("payload"), []byte("garbage"), nil) + require.ErrorIs(t, err, ErrSignatureInvalid) +} + +func TestGitIdentityMatches(t *testing.T) { + t.Parallel() + base := coreverifier.Identity{SignerIdentity: "dev@example.com", CertIssuer: testGitIssuer} + + tests := []struct { + name string + expected *lockfile.Provenance + want bool + }{ + { + name: "exact match", + expected: &lockfile.Provenance{SignerIdentity: "dev@example.com", CertIssuer: testGitIssuer}, + want: true, + }, + { + name: "identity mismatch", + expected: &lockfile.Provenance{SignerIdentity: "other@example.com", CertIssuer: testGitIssuer}, + want: false, + }, + { + name: "issuer mismatch", + expected: &lockfile.Provenance{SignerIdentity: "dev@example.com", CertIssuer: "https://other"}, + want: false, + }, + { + name: "repository mismatch", + expected: &lockfile.Provenance{ + SignerIdentity: "dev@example.com", + CertIssuer: testGitIssuer, + RepositoryURI: "https://github.com/org/repo", + }, + want: false, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.want, gitIdentityMatches(base, tc.expected)) + }) + } +} + +func TestVerifyGitEnforcesExpectedIdentity(t *testing.T) { + t.Parallel() + ca := newTestCA(t) + cert, key := ca.issueSigningCert(t, "dev@example.com") + payload := []byte(testCommitPayload) + sig := signCommitPayload(t, payload, cert, key) + + // Exercise the identity path through the internal helpers with the + // test CA pools (VerifyGit itself pins the embedded Fulcio roots, + // which a synthetic CA cannot chain to). + got, err := verifyGitSignature(t.Context(), payload, sig, ca.pool(), x509.NewCertPool()) + require.NoError(t, err) + identity, err := gitIdentityFromCertificate(got) + require.NoError(t, err) + + assert.True(t, gitIdentityMatches(identity, &lockfile.Provenance{ + SignerIdentity: "dev@example.com", + CertIssuer: testGitIssuer, + })) + assert.False(t, gitIdentityMatches(identity, &lockfile.Provenance{ + SignerIdentity: "attacker@example.com", + CertIssuer: testGitIssuer, + })) +} diff --git a/pkg/skills/verifier/mocks/mock_verifier.go b/pkg/skills/verifier/mocks/mock_verifier.go new file mode 100644 index 0000000000..ee8287fda0 --- /dev/null +++ b/pkg/skills/verifier/mocks/mock_verifier.go @@ -0,0 +1,131 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: verifier.go +// +// Generated by this command: +// +// mockgen -destination=mocks/mock_verifier.go -package=mocks -source=verifier.go Verifier +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + reflect "reflect" + + lockfile "github.com/stacklok/toolhive/pkg/skills/lockfile" + verifier "github.com/stacklok/toolhive/pkg/skills/verifier" + gomock "go.uber.org/mock/gomock" +) + +// MockVerifier is a mock of Verifier interface. +type MockVerifier struct { + ctrl *gomock.Controller + recorder *MockVerifierMockRecorder + isgomock struct{} +} + +// MockVerifierMockRecorder is the mock recorder for MockVerifier. +type MockVerifierMockRecorder struct { + mock *MockVerifier +} + +// NewMockVerifier creates a new mock instance. +func NewMockVerifier(ctrl *gomock.Controller) *MockVerifier { + mock := &MockVerifier{ctrl: ctrl} + mock.recorder = &MockVerifierMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockVerifier) EXPECT() *MockVerifierMockRecorder { + return m.recorder +} + +// ResultFromBundle mocks base method. +func (m *MockVerifier) ResultFromBundle(bundle []byte, digest string) (*verifier.Result, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ResultFromBundle", bundle, digest) + ret0, _ := ret[0].(*verifier.Result) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ResultFromBundle indicates an expected call of ResultFromBundle. +func (mr *MockVerifierMockRecorder) ResultFromBundle(bundle, digest any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResultFromBundle", reflect.TypeOf((*MockVerifier)(nil).ResultFromBundle), bundle, digest) +} + +// VerifyBundleOffline mocks base method. +func (m *MockVerifier) VerifyBundleOffline(bundle []byte, digest string, expected *lockfile.Provenance) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "VerifyBundleOffline", bundle, digest, expected) + ret0, _ := ret[0].(error) + return ret0 +} + +// VerifyBundleOffline indicates an expected call of VerifyBundleOffline. +func (mr *MockVerifierMockRecorder) VerifyBundleOffline(bundle, digest, expected any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "VerifyBundleOffline", reflect.TypeOf((*MockVerifier)(nil).VerifyBundleOffline), bundle, digest, expected) +} + +// VerifyBundleOfflineWithKey mocks base method. +func (m *MockVerifier) VerifyBundleOfflineWithKey(bundle []byte, imageRef, digest string, pubKeyPEM []byte) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "VerifyBundleOfflineWithKey", bundle, imageRef, digest, pubKeyPEM) + ret0, _ := ret[0].(error) + return ret0 +} + +// VerifyBundleOfflineWithKey indicates an expected call of VerifyBundleOfflineWithKey. +func (mr *MockVerifierMockRecorder) VerifyBundleOfflineWithKey(bundle, imageRef, digest, pubKeyPEM any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "VerifyBundleOfflineWithKey", reflect.TypeOf((*MockVerifier)(nil).VerifyBundleOfflineWithKey), bundle, imageRef, digest, pubKeyPEM) +} + +// VerifyGit mocks base method. +func (m *MockVerifier) VerifyGit(ctx context.Context, payload, signature []byte, expected *lockfile.Provenance) (*verifier.Result, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "VerifyGit", ctx, payload, signature, expected) + ret0, _ := ret[0].(*verifier.Result) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// VerifyGit indicates an expected call of VerifyGit. +func (mr *MockVerifierMockRecorder) VerifyGit(ctx, payload, signature, expected any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "VerifyGit", reflect.TypeOf((*MockVerifier)(nil).VerifyGit), ctx, payload, signature, expected) +} + +// VerifyOCI mocks base method. +func (m *MockVerifier) VerifyOCI(ctx context.Context, imageRef, digest string, expected *lockfile.Provenance) (*verifier.Result, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "VerifyOCI", ctx, imageRef, digest, expected) + ret0, _ := ret[0].(*verifier.Result) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// VerifyOCI indicates an expected call of VerifyOCI. +func (mr *MockVerifierMockRecorder) VerifyOCI(ctx, imageRef, digest, expected any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "VerifyOCI", reflect.TypeOf((*MockVerifier)(nil).VerifyOCI), ctx, imageRef, digest, expected) +} + +// VerifyOCIWithKey mocks base method. +func (m *MockVerifier) VerifyOCIWithKey(ctx context.Context, imageRef, digest string, pubKeyPEM []byte) (*verifier.Result, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "VerifyOCIWithKey", ctx, imageRef, digest, pubKeyPEM) + ret0, _ := ret[0].(*verifier.Result) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// VerifyOCIWithKey indicates an expected call of VerifyOCIWithKey. +func (mr *MockVerifierMockRecorder) VerifyOCIWithKey(ctx, imageRef, digest, pubKeyPEM any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "VerifyOCIWithKey", reflect.TypeOf((*MockVerifier)(nil).VerifyOCIWithKey), ctx, imageRef, digest, pubKeyPEM) +} diff --git a/pkg/skills/verifier/oci.go b/pkg/skills/verifier/oci.go new file mode 100644 index 0000000000..d2291c5b3a --- /dev/null +++ b/pkg/skills/verifier/oci.go @@ -0,0 +1,183 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/sigstore/sigstore-go/pkg/root" + "github.com/sigstore/sigstore-go/pkg/verify" + + coreverifier "github.com/stacklok/toolhive-core/container/verifier" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// VerifyOCI discovers and verifies the Sigstore signature for an OCI +// artifact via the keyless (Fulcio) flow. See the interface documentation +// for the expected/TOFU semantics. +func (d *Default) VerifyOCI( + ctx context.Context, + imageRef, digest string, + expected *lockfile.Provenance, +) (*Result, error) { + bundles, err := d.retrieveBundles(ctx, imageRef, digest) + if err != nil { + return nil, err + } + + tm, err := coreverifier.OfflineTrustedMaterial() + if err != nil { + return nil, fmt.Errorf("loading trusted material: %w", err) + } + opts, err := coreverifier.DefaultVerifierOptions() + if err != nil { + return nil, fmt.Errorf("loading verifier options: %w", err) + } + + result, lastErr := verifyKeylessBundles(bundles, tm, opts, expected) + if result != nil { + return result, nil + } + return nil, classifyVerifyFailure(bundles, tm, opts, expected, lastErr) +} + +// VerifyOCIWithKey discovers and verifies the Sigstore signature for an OCI +// artifact against a PEM public key (the cosign key-pair flow). +func (d *Default) VerifyOCIWithKey( + ctx context.Context, + imageRef, digest string, + pubKeyPEM []byte, +) (*Result, error) { + bundles, err := d.retrieveBundles(ctx, imageRef, digest) + if err != nil { + return nil, err + } + + var lastErr error + for _, b := range bundles { + if _, verifyErr := coreverifier.VerifyBundleWithKey(b, pubKeyPEM); verifyErr != nil { + lastErr = verifyErr + continue + } + return resultFromKey(b.Raw), nil + } + return nil, wrapInvalid(lastErr) +} + +// verifyKeylessBundles verifies bundles until one passes the keyless policy, +// returning its result, or nil with the last verification error. +func verifyKeylessBundles( + bundles []coreverifier.Bundle, + tm root.TrustedMaterial, + opts []verify.VerifierOption, + expected *lockfile.Provenance, +) (*Result, error) { + var lastErr error + for _, b := range bundles { + vr, verifyErr := coreverifier.VerifyBundle(b, tm, expectedIdentity(expected), opts...) + if verifyErr != nil { + lastErr = verifyErr + continue + } + identity, idErr := coreverifier.IdentityFromResult(vr) + if idErr != nil { + lastErr = idErr + continue + } + return resultFromCore(identity, b.Raw), nil + } + return nil, lastErr +} + +// retrieveBundles fetches the signature bundles for the artifact pinned to +// the digest, mapping core's unsigned signal to ErrUnsigned. The digest is +// required — verification without a pinned digest would leave tag +// resolution to fetch time — and when imageRef already embeds one, the two +// must agree: verifying the ref's digest while the caller believes the +// parameter's was verified would hide lock corruption. +func (d *Default) retrieveBundles(ctx context.Context, imageRef, digest string) ([]coreverifier.Bundle, error) { + if digest == "" { + return nil, errors.New("artifact digest is required for verification") + } + ref := imageRef + if embedded, ok := splitEmbeddedDigest(imageRef); ok { + if embedded != digest { + return nil, fmt.Errorf("reference %q embeds digest %s but %s was requested — refusing to verify ambiguous input", + imageRef, embedded, digest) + } + } else { + ref = imageRef + "@" + digest + } + bundles, err := coreverifier.RetrieveBundles(ctx, ref, d.keychain) + if errors.Is(err, coreverifier.ErrNoBundles) { + return nil, fmt.Errorf("%w: no signature material found for %s", ErrUnsigned, ref) + } + if err != nil { + return nil, err + } + return bundles, nil +} + +// splitEmbeddedDigest returns the digest embedded in an OCI reference +// ("repo@sha256:..."), if any. +func splitEmbeddedDigest(imageRef string) (string, bool) { + _, embedded, ok := strings.Cut(imageRef, "@") + return embedded, ok +} + +// classifyVerifyFailure distinguishes a signer mismatch from an invalid +// signature. The expected identity is enforced inside the Sigstore policy, +// so a mismatch surfaces as a verification failure; re-verifying without +// the identity constraint tells the two apart: if the chain of trust holds +// without the constraint, the failure was the identity — and the identity +// that DID verify is reported, so an operator can tell a legitimate +// publisher rotation from an artifact substitution. +func classifyVerifyFailure( + bundles []coreverifier.Bundle, + tm root.TrustedMaterial, + opts []verify.VerifierOption, + expected *lockfile.Provenance, + lastErr error, +) error { + if expected != nil { + for _, b := range bundles { + vr, err := coreverifier.VerifyBundle(b, tm, nil, opts...) + if err != nil { + continue + } + return signerMismatchError(vr, expected) + } + } + return wrapInvalid(lastErr) +} + +// signerMismatchError builds the ErrSignerMismatch error, naming both the +// expected identity tuple and the identity the artifact actually verifies +// with (when extractable). +func signerMismatchError(vr *verify.VerificationResult, expected *lockfile.Provenance) error { + observed, idErr := coreverifier.IdentityFromResult(vr) + if idErr != nil { + return fmt.Errorf("%w: artifact is signed by a different identity than %q (issuer %q)", + ErrSignerMismatch, expected.SignerIdentity, expected.CertIssuer) + } + return fmt.Errorf("%w: locked to %q (issuer %q), but the artifact verifies as %q (issuer %q)", + ErrSignerMismatch, + expected.SignerIdentity, expected.CertIssuer, + observed.SignerIdentity, observed.CertIssuer) +} + +// wrapInvalid wraps a verification cause in ErrSignatureInvalid without +// stuttering: core's own verification-failed sentinel prefix is trimmed +// from the display text (the classification value it carried is replaced by +// our sentinel; this is message cosmetics, not error matching). +func wrapInvalid(cause error) error { + if cause == nil { + return ErrSignatureInvalid + } + text := strings.TrimPrefix(cause.Error(), coreverifier.ErrVerificationFailed.Error()+": ") + return fmt.Errorf("%w: %s", ErrSignatureInvalid, text) +} diff --git a/pkg/skills/verifier/offline.go b/pkg/skills/verifier/offline.go new file mode 100644 index 0000000000..d7dc37f854 --- /dev/null +++ b/pkg/skills/verifier/offline.go @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + + coreverifier "github.com/stacklok/toolhive-core/container/verifier" + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/signer" +) + +// VerifyBundleOffline re-verifies a stored bundle against the artifact +// digest without network access. A non-nil expected identity is enforced +// inside the Sigstore policy; a mismatch is reported as ErrSignerMismatch, +// any other verification failure as ErrSignatureInvalid. +func (*Default) VerifyBundleOffline(bundleBytes []byte, digest string, expected *lockfile.Provenance) error { + if len(bundleBytes) == 0 { + // Classified as an invalid signature because a recorded identity + // with nothing backing it cannot be verified; the message leads + // with the actionable fact — the stored bundle is missing, so + // reinstalling (or re-adopting) is the fix. + return fmt.Errorf("%w: no stored bundle to verify — reinstall to restore it", ErrSignatureInvalid) + } + _, err := coreverifier.VerifyBundleOffline(bundleBytes, digest, expectedIdentity(expected)) + if err == nil { + return nil + } + if expected != nil && !errors.Is(err, coreverifier.ErrVerificationFailed) { + // Malformed input never reaches verification; don't reclassify. + return fmt.Errorf("%w: %s", ErrSignatureInvalid, err.Error()) + } + if expected != nil { + // The identity is bound into the policy, so a mismatch surfaces as + // a verification failure; re-verifying without the constraint tells + // mismatch apart from a broken signature — and yields the identity + // that DID verify, which the error reports. + if vr, tofuErr := coreverifier.VerifyBundleOffline(bundleBytes, digest, nil); tofuErr == nil { + return signerMismatchError(vr, expected) + } + } + return wrapInvalid(err) +} + +// VerifyBundleOfflineWithKey re-verifies a stored key-signed bundle against +// the signer's PEM public key — the offline counterpart of +// VerifyOCIWithKey. Key-signed bundles sign the cosign simple-signing +// payload (which embeds the artifact digest), so the payload is +// reconstructed from imageRef and digest and the signature checked over it +// — that reconstruction IS the digest-binding check. +func (*Default) VerifyBundleOfflineWithKey(bundleBytes []byte, imageRef, digest string, pubKeyPEM []byte) error { + if len(bundleBytes) == 0 { + return fmt.Errorf("%w: no stored bundle to verify — reinstall to restore it", ErrSignatureInvalid) + } + payload, err := signer.SimpleSigningPayload(imageRef, digest) + if err != nil { + return fmt.Errorf("%w: %s", ErrSignatureInvalid, err.Error()) + } + payloadDigest := sha256.Sum256(payload) + digestArg := coreverifier.DigestAlgorithmSHA256 + ":" + hex.EncodeToString(payloadDigest[:]) + if _, err := coreverifier.VerifyBundleOfflineWithKey(bundleBytes, digestArg, pubKeyPEM); err != nil { + return wrapInvalid(err) + } + return nil +} + +// ResultFromBundle verifies a stored bundle offline (chain of trust only) +// and returns the observed identity, for back-filling provenance of +// adopted skills. +func (*Default) ResultFromBundle(bundleBytes []byte, digest string) (*Result, error) { + if len(bundleBytes) == 0 { + return nil, fmt.Errorf("%w: no stored bundle to verify", ErrSignatureInvalid) + } + vr, err := coreverifier.VerifyBundleOffline(bundleBytes, digest, nil) + if err != nil { + return nil, wrapInvalid(err) + } + identity, err := coreverifier.IdentityFromResult(vr) + if err != nil { + return nil, wrapInvalid(err) + } + return resultFromCore(identity, bundleBytes), nil +} diff --git a/pkg/skills/verifier/types.go b/pkg/skills/verifier/types.go new file mode 100644 index 0000000000..7f367d93f2 --- /dev/null +++ b/pkg/skills/verifier/types.go @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import ( + coreverifier "github.com/stacklok/toolhive-core/container/verifier" + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +// sigstorePublicGoodRekorURL identifies the transparency log instance the +// embedded trust root belongs to, recorded in lock provenance for humans +// auditing the file. +const sigstorePublicGoodRekorURL = "https://rekor.sigstore.dev" + +// Result contains the outcome of verifying a signed artifact. +type Result struct { + // Signed is true when a signature was found and verified. + Signed bool + // SignerIdentity is the certificate's subject identity (workflow path + // for GitHub-Actions-issued certificates, SAN otherwise). Empty for + // key-signed artifacts, which carry no certificate. + SignerIdentity string + // CertIssuer is the OIDC issuer that authenticated the signer. Empty + // for key-signed artifacts. + CertIssuer string + // RepositoryURI is the source repository from the certificate, if any. + RepositoryURI string + // SigstoreURL is the transparency log instance used for verification. + SigstoreURL string + // Provisional marks a verification with a documented assurance gap + // (git signatures until Rekor proof validation lands). + Provisional bool + // Bundle is the serialized Sigstore bundle for offline re-verification. + Bundle []byte +} + +// ToLockProvenance converts a verification result to a lock file provenance +// block. Key-signed results have no certificate identity and yield nil — +// the lock file records provenance only for identity-bearing signatures. +func (r *Result) ToLockProvenance() *lockfile.Provenance { + if r == nil || !r.Signed || r.SignerIdentity == "" { + return nil + } + return &lockfile.Provenance{ + SignerIdentity: r.SignerIdentity, + CertIssuer: r.CertIssuer, + RepositoryURI: r.RepositoryURI, + SigstoreURL: r.SigstoreURL, + Provisional: r.Provisional, + } +} + +// expectedIdentity converts a lock provenance block into the core Identity +// that gets bound into the Sigstore verification policy. A nil provenance +// (trust on first use) yields nil, which core treats as chain-of-trust-only +// verification. +func expectedIdentity(p *lockfile.Provenance) *coreverifier.Identity { + if p == nil { + return nil + } + return &coreverifier.Identity{ + SignerIdentity: p.SignerIdentity, + CertIssuer: p.CertIssuer, + SourceRepositoryURI: p.RepositoryURI, + } +} + +// resultFromKey builds a Result for a key-signed artifact. It carries no +// certificate identity and no SigstoreURL — the key flow writes no +// transparency-log entry, and recording an instance the signature never +// touched would fabricate provenance. +func resultFromKey(raw []byte) *Result { + return &Result{ + Signed: true, + Bundle: raw, + } +} + +// resultFromCore builds a Result from a core verification outcome. +func resultFromCore(identity coreverifier.Identity, raw []byte) *Result { + return &Result{ + Signed: true, + SignerIdentity: identity.SignerIdentity, + CertIssuer: identity.CertIssuer, + RepositoryURI: identity.SourceRepositoryURI, + SigstoreURL: sigstorePublicGoodRekorURL, + Bundle: raw, + } +} diff --git a/pkg/skills/verifier/verifier.go b/pkg/skills/verifier/verifier.go new file mode 100644 index 0000000000..681fca9c2e --- /dev/null +++ b/pkg/skills/verifier/verifier.go @@ -0,0 +1,82 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package verifier verifies Sigstore signatures on skill artifacts. +// +// It is a thin wrapper over toolhive-core's container/verifier exports: all +// cryptographic verification — including binding an expected identity into +// the Sigstore policy — happens in core. This package adds the +// skills-domain vocabulary: lock file provenance conversion, the +// unsigned/invalid/mismatch error taxonomy, and the trust-on-first-use flow +// (nil expected identity verifies the chain of trust only; the caller +// records the observed identity). +// +// Verification uses the trusted root embedded in toolhive-core — hermetic, +// no TUF fetch — so results are reproducible offline at the cost of +// snapshot freshness (see core's OfflineTrustedMaterial). +package verifier + +import ( + "context" + + "github.com/google/go-containerregistry/pkg/authn" + + "github.com/stacklok/toolhive/pkg/skills/lockfile" +) + +//go:generate mockgen -destination=mocks/mock_verifier.go -package=mocks -source=verifier.go Verifier + +// Verifier verifies Sigstore signatures for skill artifacts. +type Verifier interface { + // VerifyOCI discovers the Sigstore signature material attached to the + // OCI artifact and verifies it (keyless/Fulcio flow). A non-nil + // expected identity is enforced inside the Sigstore verification + // policy; nil expected is the trust-on-first-use case and verifies the + // chain of trust only. Returns ErrUnsigned when the artifact carries + // no signature material. + VerifyOCI(ctx context.Context, imageRef, digest string, expected *lockfile.Provenance) (*Result, error) + + // VerifyOCIWithKey discovers the signature material and verifies it + // against the given PEM public key (the cosign key-pair flow). + // Key-signed bundles carry no certificate identity: trust is the key. + VerifyOCIWithKey(ctx context.Context, imageRef, digest string, pubKeyPEM []byte) (*Result, error) + + // VerifyGit cryptographically verifies a gitsign commit signature over + // the commit payload against the embedded Fulcio roots. A non-nil + // expected identity must match the certificate identity; nil expected + // is the trust-on-first-use case. Returns ErrUnsigned for an empty + // signature. + VerifyGit(ctx context.Context, payload, signature []byte, expected *lockfile.Provenance) (*Result, error) + + // VerifyBundleOffline re-verifies a stored bundle against the artifact + // digest ("sha256:") without network access, enforcing expected + // like VerifyOCI. + VerifyBundleOffline(bundle []byte, digest string, expected *lockfile.Provenance) error + + // VerifyBundleOfflineWithKey re-verifies a stored key-signed bundle + // against the signer's PEM public key without network access — the + // offline counterpart of VerifyOCIWithKey. imageRef and digest + // reconstruct the signed payload, binding the check to the artifact. + VerifyBundleOfflineWithKey(bundle []byte, imageRef, digest string, pubKeyPEM []byte) error + + // ResultFromBundle verifies a stored bundle offline (chain of trust + // only) and returns the observed identity — used to back-fill + // provenance for adopted skills. + ResultFromBundle(bundle []byte, digest string) (*Result, error) +} + +// Default implements Verifier on toolhive-core's Sigstore exports. +type Default struct { + keychain authn.Keychain +} + +var _ Verifier = (*Default)(nil) + +// NewDefault creates a verifier using the given registry auth keychain for +// bundle retrieval. A nil keychain falls back to the default keychain. +func NewDefault(keychain authn.Keychain) *Default { + if keychain == nil { + keychain = authn.DefaultKeychain + } + return &Default{keychain: keychain} +} diff --git a/pkg/skills/verifier/verifier_test.go b/pkg/skills/verifier/verifier_test.go new file mode 100644 index 0000000000..e7e867e612 --- /dev/null +++ b/pkg/skills/verifier/verifier_test.go @@ -0,0 +1,278 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package verifier + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/google/go-containerregistry/pkg/name" + "github.com/google/go-containerregistry/pkg/registry" + "github.com/google/go-containerregistry/pkg/v1/random" + "github.com/google/go-containerregistry/pkg/v1/remote" + "github.com/sigstore/sigstore/pkg/cryptoutils" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/pkg/skills/signer" +) + +// startTestRegistry runs an in-process OCI registry and returns its host. +func startTestRegistry(t *testing.T) string { + t.Helper() + reg := httptest.NewServer(registry.New()) + t.Cleanup(reg.Close) + return strings.TrimPrefix(reg.URL, "http://") +} + +// pushTestArtifact pushes a random OCI image and returns its ref and digest. +func pushTestArtifact(t *testing.T, host string) (ref string, digest string) { + t.Helper() + img, err := random.Image(256, 1) + require.NoError(t, err) + ref = host + "/test/skill:v1" + parsed, err := name.ParseReference(ref) + require.NoError(t, err) + require.NoError(t, remote.Write(parsed, img)) + d, err := img.Digest() + require.NoError(t, err) + return ref, d.String() +} + +// signArtifact signs the artifact with a fresh cosign key via the signer +// package and returns the public key PEM and the bundle it produced — +// exactly the flow `thv skill push --key` performs. +func signArtifact(t *testing.T, ref, digest string) (pubPEM, bundle []byte) { + t.Helper() + priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + privPEM, err := cryptoutils.MarshalPrivateKeyToPEM(priv) + require.NoError(t, err) + keyPath := filepath.Join(t.TempDir(), "cosign.key") + require.NoError(t, os.WriteFile(keyPath, privPEM, 0o600)) + pubPEM, err = cryptoutils.MarshalPublicKeyToPEM(priv.Public()) + require.NoError(t, err) + + raw, err := signer.NewDefault(nil).SignOCI(t.Context(), ref, digest, signer.Options{Key: keyPath}) + require.NoError(t, err) + return pubPEM, raw +} + +func TestVerifyOCIWithKeyRoundTrip(t *testing.T) { + t.Parallel() + host := startTestRegistry(t) + ref, digest := pushTestArtifact(t, host) + pubPEM, _ := signArtifact(t, ref, digest) + + result, err := NewDefault(nil).VerifyOCIWithKey(t.Context(), ref, digest, pubPEM) + require.NoError(t, err, "the artifact signed by the signer package must verify with its key") + assert.True(t, result.Signed) + assert.NotEmpty(t, result.Bundle, "the bundle must be returned for durable storage") + assert.Empty(t, result.SignerIdentity, "key-signed artifacts carry no certificate identity") + assert.Empty(t, result.SigstoreURL, + "the key flow writes no transparency-log entry; recording one would fabricate provenance") + assert.Nil(t, result.ToLockProvenance(), "key-signed results must not fabricate lock provenance") + + // The stored bundle re-verifies offline with the key, and rejects a + // different key. + require.NoError(t, NewDefault(nil).VerifyBundleOfflineWithKey(result.Bundle, ref, digest, pubPEM)) + otherPriv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + otherPub, err := cryptoutils.MarshalPublicKeyToPEM(otherPriv.Public()) + require.NoError(t, err) + require.ErrorIs(t, NewDefault(nil).VerifyBundleOfflineWithKey(result.Bundle, ref, digest, otherPub), + ErrSignatureInvalid) + require.ErrorIs(t, NewDefault(nil).VerifyBundleOfflineWithKey( + result.Bundle, ref, "sha256:"+strings.Repeat("e", 64), pubPEM), ErrSignatureInvalid, + "a different artifact digest reconstructs a different payload and must not verify") +} + +func TestVerifyOCIDigestGuards(t *testing.T) { + t.Parallel() + host := startTestRegistry(t) + ref, digest := pushTestArtifact(t, host) + d := NewDefault(nil) + + _, err := d.VerifyOCI(t.Context(), ref, "", nil) + require.ErrorContains(t, err, "digest is required", + "an empty digest would leave tag resolution to fetch time") + + otherDigest := "sha256:" + strings.Repeat("f", 64) + _, err = d.VerifyOCI(t.Context(), ref+"@"+digest, otherDigest, nil) + require.ErrorContains(t, err, "refusing to verify ambiguous input", + "a ref-embedded digest disagreeing with the parameter is lock corruption") + + // Agreement between the embedded digest and the parameter is fine. + _, err = d.VerifyOCI(t.Context(), ref+"@"+digest, digest, nil) + require.ErrorIs(t, err, ErrUnsigned, "consistent inputs proceed to retrieval") +} + +func TestVerifyOCIWithKeyRejectsWrongKey(t *testing.T) { + t.Parallel() + host := startTestRegistry(t) + ref, digest := pushTestArtifact(t, host) + signArtifact(t, ref, digest) + + otherPriv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + otherPub, err := cryptoutils.MarshalPublicKeyToPEM(otherPriv.Public()) + require.NoError(t, err) + + _, err = NewDefault(nil).VerifyOCIWithKey(t.Context(), ref, digest, otherPub) + require.ErrorIs(t, err, ErrSignatureInvalid) +} + +func TestVerifyOCIUnsignedArtifact(t *testing.T) { + t.Parallel() + host := startTestRegistry(t) + ref, digest := pushTestArtifact(t, host) + + _, err := NewDefault(nil).VerifyOCI(t.Context(), ref, digest, nil) + require.ErrorIs(t, err, ErrUnsigned) + + _, err = NewDefault(nil).VerifyOCIWithKey(t.Context(), ref, digest, nil) + require.ErrorIs(t, err, ErrUnsigned) +} + +func TestVerifyOCIKeylessRejectsKeySignedArtifact(t *testing.T) { + t.Parallel() + host := startTestRegistry(t) + ref, digest := pushTestArtifact(t, host) + signArtifact(t, ref, digest) + + // The keyless flow requires a certificate chain to Fulcio; a key-signed + // bundle has none, so this must fail as an invalid signature — not as + // unsigned, and never as a panic. + _, err := NewDefault(nil).VerifyOCI(t.Context(), ref, digest, nil) + require.ErrorIs(t, err, ErrSignatureInvalid) +} + +func TestVerifyBundleOffline(t *testing.T) { + t.Parallel() + host := startTestRegistry(t) + ref, digest := pushTestArtifact(t, host) + _, bundle := signArtifact(t, ref, digest) + d := NewDefault(nil) + + t.Run("empty bundle rejected", func(t *testing.T) { + t.Parallel() + err := d.VerifyBundleOffline(nil, digest, nil) + require.ErrorIs(t, err, ErrSignatureInvalid) + }) + + t.Run("malformed bundle rejected", func(t *testing.T) { + t.Parallel() + err := d.VerifyBundleOffline([]byte("not a bundle"), digest, nil) + require.ErrorIs(t, err, ErrSignatureInvalid) + }) + + t.Run("key-signed bundle fails keyless offline verification", func(t *testing.T) { + t.Parallel() + // The stored bundle parses and reaches verification, but carries no + // certificate — the offline keyless path must reject it rather than + // trusting it. + err := d.VerifyBundleOffline(bundle, digest, nil) + require.ErrorIs(t, err, ErrSignatureInvalid) + }) + + t.Run("expected identity against unverifiable bundle stays invalid", func(t *testing.T) { + t.Parallel() + // Both the pinned and the TOFU re-verify fail, so this must NOT be + // misclassified as a signer mismatch. + err := d.VerifyBundleOffline(bundle, digest, &lockfile.Provenance{ + SignerIdentity: "/.github/workflows/release.yml", + CertIssuer: "https://token.actions.githubusercontent.com", + }) + require.ErrorIs(t, err, ErrSignatureInvalid) + require.NotErrorIs(t, err, ErrSignerMismatch) + }) +} + +func TestResultFromBundleMalformed(t *testing.T) { + t.Parallel() + d := NewDefault(nil) + + _, err := d.ResultFromBundle(nil, "sha256:abc") + require.ErrorIs(t, err, ErrSignatureInvalid) + + _, err = d.ResultFromBundle([]byte("junk"), "sha256:abc") + require.ErrorIs(t, err, ErrSignatureInvalid) +} + +func TestToLockProvenance(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + result *Result + want *lockfile.Provenance + }{ + {name: "nil result", result: nil, want: nil}, + {name: "unsigned result", result: &Result{Signed: false}, want: nil}, + { + name: "key-signed result has no identity to record", + result: &Result{Signed: true, SigstoreURL: sigstorePublicGoodRekorURL}, + want: nil, + }, + { + name: "provisional result marks the lock provenance", + result: &Result{ + Signed: true, + SignerIdentity: "dev@example.com", + CertIssuer: "https://accounts.example.com", + Provisional: true, + }, + want: &lockfile.Provenance{ + SignerIdentity: "dev@example.com", + CertIssuer: "https://accounts.example.com", + Provisional: true, + }, + }, + { + name: "identity-bearing result maps all fields", + result: &Result{ + Signed: true, + SignerIdentity: "/.github/workflows/release.yml", + CertIssuer: "https://token.actions.githubusercontent.com", + RepositoryURI: "https://github.com/org/repo", + SigstoreURL: sigstorePublicGoodRekorURL, + }, + want: &lockfile.Provenance{ + SignerIdentity: "/.github/workflows/release.yml", + CertIssuer: "https://token.actions.githubusercontent.com", + RepositoryURI: "https://github.com/org/repo", + SigstoreURL: sigstorePublicGoodRekorURL, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.want, tc.result.ToLockProvenance()) + }) + } +} + +func TestExpectedIdentityConversion(t *testing.T) { + t.Parallel() + + assert.Nil(t, expectedIdentity(nil), "TOFU first use must yield a nil core identity") + + got := expectedIdentity(&lockfile.Provenance{ + SignerIdentity: "/.github/workflows/release.yml", + CertIssuer: "https://token.actions.githubusercontent.com", + RepositoryURI: "https://github.com/org/repo", + }) + require.NotNil(t, got) + assert.Equal(t, "/.github/workflows/release.yml", got.SignerIdentity) + assert.Equal(t, "https://token.actions.githubusercontent.com", got.CertIssuer) + assert.Equal(t, "https://github.com/org/repo", got.SourceRepositoryURI) +} diff --git a/pkg/storage/sqlite/migrations/003_add_managed_flag.sql b/pkg/storage/sqlite/migrations/003_add_managed_flag.sql new file mode 100644 index 0000000000..e9dd5dd55f --- /dev/null +++ b/pkg/storage/sqlite/migrations/003_add_managed_flag.sql @@ -0,0 +1,7 @@ +-- +goose Up + +ALTER TABLE installed_skills ADD COLUMN managed INTEGER NOT NULL DEFAULT 0; + +-- +goose Down + +ALTER TABLE installed_skills DROP COLUMN managed; diff --git a/pkg/storage/sqlite/migrations/004_add_skill_sigstore_bundle.sql b/pkg/storage/sqlite/migrations/004_add_skill_sigstore_bundle.sql new file mode 100644 index 0000000000..7c73e8c72d --- /dev/null +++ b/pkg/storage/sqlite/migrations/004_add_skill_sigstore_bundle.sql @@ -0,0 +1,7 @@ +-- +goose Up + +ALTER TABLE installed_skills ADD COLUMN sigstore_bundle BLOB DEFAULT NULL; + +-- +goose Down + +ALTER TABLE installed_skills DROP COLUMN sigstore_bundle; diff --git a/pkg/storage/sqlite/migrations_test.go b/pkg/storage/sqlite/migrations_test.go index 52e1012905..370db0bf16 100644 --- a/pkg/storage/sqlite/migrations_test.go +++ b/pkg/storage/sqlite/migrations_test.go @@ -130,6 +130,95 @@ func newGooseProvider(t *testing.T, db *sql.DB) *goose.Provider { return provider } +// TestMigrations_ManagedFlagAppliesOverPriorState verifies migration 003 adds +// installed_skills.managed with a default of 0 (false) so rows created by an +// earlier version of the schema (001/002) are not implicitly "managed" once +// the column appears. +func TestMigrations_ManagedFlagAppliesOverPriorState(t *testing.T) { + t.Parallel() + ctx := t.Context() + dbPath := filepath.Join(t.TempDir(), "test.db") + + db, err := Open(ctx, dbPath) + require.NoError(t, err) + defer db.Close() + + provider := newGooseProvider(t, db.DB()) + + // Roll back to just after migration 002, before "managed" existed. + _, err = provider.DownTo(ctx, 2) + require.NoError(t, err) + + _, err = db.DB().Exec(`INSERT INTO entries (entry_type, name) VALUES ('skill', 'pre-migration-skill')`) + require.NoError(t, err) + _, err = db.DB().Exec(`INSERT INTO installed_skills (entry_id, scope) VALUES (1, 'user')`) + require.NoError(t, err) + + // Re-apply Up: migration 003 adds the column to the pre-existing row. + _, err = provider.Up(ctx) + require.NoError(t, err) + + var managed int + err = db.DB().QueryRow(`SELECT managed FROM installed_skills WHERE entry_id = 1`).Scan(&managed) + require.NoError(t, err) + assert.Equal(t, 0, managed, "a row created before migration 003 must default to unmanaged") + + // Down for 003 removes the column without disturbing 001/002 tables. + _, err = provider.DownTo(ctx, 2) + require.NoError(t, err) + var count int + err = db.DB().QueryRow( + `SELECT COUNT(*) FROM pragma_table_info('installed_skills') WHERE name = 'managed'`, + ).Scan(&count) + require.NoError(t, err) + assert.Equal(t, 0, count, "managed column should be dropped by 003 Down") + assert.True(t, tableExists(t, db.DB(), "installed_skills"), "installed_skills table itself must remain") +} + +// TestMigrations_SigstoreBundleAppliesOverPriorState verifies migration 004 +// adds installed_skills.sigstore_bundle defaulting to NULL, so rows created by +// an earlier schema read back as unsigned (nil bundle) once the column appears. +func TestMigrations_SigstoreBundleAppliesOverPriorState(t *testing.T) { + t.Parallel() + ctx := t.Context() + dbPath := filepath.Join(t.TempDir(), "test.db") + + db, err := Open(ctx, dbPath) + require.NoError(t, err) + defer db.Close() + + provider := newGooseProvider(t, db.DB()) + + // Roll back to just after migration 003, before "sigstore_bundle" existed. + _, err = provider.DownTo(ctx, 3) + require.NoError(t, err) + + _, err = db.DB().Exec(`INSERT INTO entries (entry_type, name) VALUES ('skill', 'pre-migration-skill')`) + require.NoError(t, err) + _, err = db.DB().Exec(`INSERT INTO installed_skills (entry_id, scope) VALUES (1, 'user')`) + require.NoError(t, err) + + // Re-apply Up: migration 004 adds the column to the pre-existing row. + _, err = provider.Up(ctx) + require.NoError(t, err) + + var bundle []byte + err = db.DB().QueryRow(`SELECT sigstore_bundle FROM installed_skills WHERE entry_id = 1`).Scan(&bundle) + require.NoError(t, err) + assert.Nil(t, bundle, "a row created before migration 004 must default to a NULL bundle") + + // Down for 004 removes the column without disturbing earlier tables. + _, err = provider.DownTo(ctx, 3) + require.NoError(t, err) + var count int + err = db.DB().QueryRow( + `SELECT COUNT(*) FROM pragma_table_info('installed_skills') WHERE name = 'sigstore_bundle'`, + ).Scan(&count) + require.NoError(t, err) + assert.Equal(t, 0, count, "sigstore_bundle column should be dropped by 004 Down") + assert.True(t, tableExists(t, db.DB(), "installed_skills"), "installed_skills table itself must remain") +} + // TestMigrations_DownDropsPluginTables verifies that goose's Down for migration // 002 drops the plugin tables (installed_plugins, plugin_dependencies) while // leaving migration 001's tables (entries, installed_skills, skill_dependencies, diff --git a/pkg/storage/sqlite/skill_store.go b/pkg/storage/sqlite/skill_store.go index 87d197d849..29ffca5539 100644 --- a/pkg/storage/sqlite/skill_store.go +++ b/pkg/storage/sqlite/skill_store.go @@ -39,10 +39,20 @@ var _ storage.SkillStore = (*SkillStore)(nil) // skillColumns is the SELECT column list shared by Get and List queries. const skillColumns = `is_.id, e.name, is_.scope, is_.project_root, is_.reference, is_.tag, is_.digest, is_.version, is_.description, is_.author, json(is_.tags), - json(is_.client_apps), is_.status, is_.installed_at` + json(is_.client_apps), is_.status, is_.installed_at, is_.managed, + is_.sigstore_bundle` + +// errManagedRequiresProjectScope is returned by Create/Update when a +// user-scoped skill has Managed set. Managed pins a skill in a project's +// lock file, which only exists for project-scoped installs. +var errManagedRequiresProjectScope = errors.New("managed skill must be project-scoped") // Create stores a new installed skill. func (s *SkillStore) Create(ctx context.Context, skill skills.InstalledSkill) error { + if skill.Managed && skill.Scope != skills.ScopeProject { + return errManagedRequiresProjectScope + } + tx, err := s.db.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("beginning transaction: %w", err) @@ -85,8 +95,9 @@ func (s *SkillStore) Create(ctx context.Context, skill skills.InstalledSkill) er res, err := tx.ExecContext(ctx, ` INSERT INTO installed_skills ( entry_id, scope, project_root, reference, tag, digest, - version, description, author, tags, client_apps, status - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, jsonb(?), jsonb(?), ?)`, + version, description, author, tags, client_apps, status, managed, + sigstore_bundle + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, jsonb(?), jsonb(?), ?, ?, ?)`, entryID, string(skill.Scope), skill.ProjectRoot, @@ -99,6 +110,8 @@ func (s *SkillStore) Create(ctx context.Context, skill skills.InstalledSkill) er tagsJSON, clientsJSON, string(skill.Status), + boolToInt(skill.Managed), + skill.SigstoreBundle, ) if err != nil { if isUniqueViolation(err) { @@ -221,6 +234,10 @@ func (s *SkillStore) List(ctx context.Context, filter storage.ListFilter) ([]ski // Update modifies an existing installed skill. func (s *SkillStore) Update(ctx context.Context, skill skills.InstalledSkill) error { + if skill.Managed && skill.Scope != skills.ScopeProject { + return errManagedRequiresProjectScope + } + tx, err := s.db.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("beginning transaction: %w", err) @@ -266,7 +283,8 @@ func (s *SkillStore) Update(ctx context.Context, skill skills.InstalledSkill) er if _, err := tx.ExecContext(ctx, ` UPDATE installed_skills SET reference = ?, tag = ?, digest = ?, version = ?, description = ?, - author = ?, tags = jsonb(?), client_apps = jsonb(?), status = ? + author = ?, tags = jsonb(?), client_apps = jsonb(?), status = ?, managed = ?, + sigstore_bundle = ? WHERE id = ?`, skill.Reference, skill.Tag, @@ -277,6 +295,8 @@ func (s *SkillStore) Update(ctx context.Context, skill skills.InstalledSkill) er tagsJSON, clientsJSON, string(skill.Status), + boolToInt(skill.Managed), + skill.SigstoreBundle, installedSkillID, ); err != nil { return fmt.Errorf("updating installed skill: %w", err) @@ -370,12 +390,14 @@ func scanSkillFields(sc scanner) (skills.InstalledSkill, int64, error) { clientsBlob []byte status string installedAtStr string + managed int + sigstoreBundle []byte ) err := sc.Scan( &installedSkillID, &name, &scope, &projectRoot, &reference, &tag, &digest, &version, &description, &author, &tagsBlob, - &clientsBlob, &status, &installedAtStr, + &clientsBlob, &status, &installedAtStr, &managed, &sigstoreBundle, ) if err != nil { if errors.Is(err, sql.ErrNoRows) { @@ -404,14 +426,16 @@ func scanSkillFields(sc scanner) (skills.InstalledSkill, int64, error) { Author: author, Tags: tags, }, - Scope: skills.Scope(scope), - ProjectRoot: projectRoot, - Reference: reference, - Tag: tag, - Digest: digest, - Status: skills.InstallStatus(status), - InstalledAt: installedAt, - Clients: clients, + Scope: skills.Scope(scope), + ProjectRoot: projectRoot, + Reference: reference, + Tag: tag, + Digest: digest, + Status: skills.InstallStatus(status), + InstalledAt: installedAt, + Clients: clients, + Managed: managed != 0, + SigstoreBundle: sigstoreBundle, } return sk, installedSkillID, nil @@ -495,3 +519,11 @@ func isUniqueViolation(err error) bool { // rollback rolls back tx, ignoring errors (tx may already be committed). func rollback(tx *sql.Tx) { _ = tx.Rollback() } + +// boolToInt converts a bool to SQLite's 0/1 INTEGER representation. +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} diff --git a/pkg/storage/sqlite/skill_store_test.go b/pkg/storage/sqlite/skill_store_test.go index 3d6efe578a..7efd267e68 100644 --- a/pkg/storage/sqlite/skill_store_test.go +++ b/pkg/storage/sqlite/skill_store_test.go @@ -72,6 +72,82 @@ func TestSkillStore_Create(t *testing.T) { // InstalledAt is set by the DB DEFAULT, so just assert it is not zero. assert.False(t, got.InstalledAt.IsZero(), "InstalledAt should not be zero") + assert.False(t, got.Managed, "Managed should default to false") + assert.Nil(t, got.SigstoreBundle, "SigstoreBundle should default to nil (unsigned)") +} + +func TestSkillStore_SigstoreBundleRoundTrip(t *testing.T) { + t.Parallel() + store := newTestStore(t) + + bundle := []byte(`{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json"}`) + sk := testSkill("bundle-test") + sk.SigstoreBundle = bundle + require.NoError(t, store.Create(t.Context(), sk)) + + got, err := store.Get(t.Context(), sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + require.NoError(t, err) + assert.Equal(t, bundle, got.SigstoreBundle) + + // Update replaces the stored bundle (e.g. re-install of a re-signed + // artifact), and can clear it back to nil for an unsigned replacement. + newBundle := []byte(`{"replaced":true}`) + got.SigstoreBundle = newBundle + require.NoError(t, store.Update(t.Context(), got)) + got, err = store.Get(t.Context(), sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + require.NoError(t, err) + assert.Equal(t, newBundle, got.SigstoreBundle) + + got.SigstoreBundle = nil + require.NoError(t, store.Update(t.Context(), got)) + got, err = store.Get(t.Context(), sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + require.NoError(t, err) + assert.Nil(t, got.SigstoreBundle) +} + +func TestSkillStore_ManagedFlagRoundTrip(t *testing.T) { + t.Parallel() + store := newTestStore(t) + + sk := testSkill("managed-test") + sk.Scope = skills.ScopeProject + sk.ProjectRoot = "/tmp/project" + sk.Managed = true + require.NoError(t, store.Create(t.Context(), sk)) + + got, err := store.Get(t.Context(), sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + require.NoError(t, err) + assert.True(t, got.Managed) + + // Update can flip managed back to false (e.g. sync --prune leaving the + // record but marking it unmanaged). + got.Managed = false + require.NoError(t, store.Update(t.Context(), got)) + + got, err = store.Get(t.Context(), sk.Metadata.Name, sk.Scope, sk.ProjectRoot) + require.NoError(t, err) + assert.False(t, got.Managed) +} + +// TestSkillStore_ManagedRequiresProjectScope guards the invariant that +// Managed only ever applies to project-scoped installs (it pins a skill in +// a project's lock file, which doesn't exist for user-scoped installs). +func TestSkillStore_ManagedRequiresProjectScope(t *testing.T) { + t.Parallel() + store := newTestStore(t) + + sk := testSkill("managed-user-scope") + sk.Managed = true // sk.Scope is ScopeUser from testSkill + err := store.Create(t.Context(), sk) + require.ErrorIs(t, err, errManagedRequiresProjectScope) + + // Also rejected on Update: create a valid user-scoped record, then try + // to flip Managed on it. + sk.Managed = false + require.NoError(t, store.Create(t.Context(), sk)) + sk.Managed = true + err = store.Update(t.Context(), sk) + require.ErrorIs(t, err, errManagedRequiresProjectScope) } func TestSkillStore_CreateDuplicate(t *testing.T) { diff --git a/pkg/telemetry/config.go b/pkg/telemetry/config.go index 33eddc208b..894ac37529 100644 --- a/pkg/telemetry/config.go +++ b/pkg/telemetry/config.go @@ -17,7 +17,7 @@ import ( sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/trace" - "github.com/stacklok/toolhive/pkg/telemetry/providers" + "github.com/stacklok/toolhive-core/telemetry/providers" "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/pkg/versions" ) diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go index cc17e5f4db..e155212556 100644 --- a/pkg/telemetry/middleware.go +++ b/pkg/telemetry/middleware.go @@ -15,7 +15,6 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -23,6 +22,8 @@ import ( "go.opentelemetry.io/otel/propagation" "go.opentelemetry.io/otel/trace" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/telemetry/metrics" mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/types" ) @@ -41,7 +42,7 @@ const ( // MCPHistogramBuckets are the bucket boundaries defined by the MCP OTEL semantic conventions // for MCP server histograms (e.g. mcp.server.operation.duration). // See https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/mcp.md -var MCPHistogramBuckets = []float64{0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 30, 60, 120, 300} +var MCPHistogramBuckets = metrics.BucketsMCPSemconv() // HTTPMiddleware provides OpenTelemetry instrumentation for HTTP requests. type HTTPMiddleware struct { @@ -396,7 +397,16 @@ func (*HTTPMiddleware) addNetworkAttributes(span trace.Span, r *http.Request, ba } // addMethodSpecificAttributes adds attributes specific to certain MCP methods. +// Despite the name, the mcp.client.name block below runs for every method: under +// the Modern (stateless) MCP revision there is no initialize request, so per-request +// _meta.clientInfo is the only source of client attribution. It is a no-op for +// Legacy requests, which carry no _meta.clientInfo and still get mcp.client.name +// from the "initialize" case below. func (m *HTTPMiddleware) addMethodSpecificAttributes(span trace.Span, parsedMCP *mcpparser.ParsedMCPRequest) { + if name, ok := parsedMCP.ClientInfo["name"].(string); ok && name != "" { + span.SetAttributes(attribute.String("mcp.client.name", name)) + } + switch parsedMCP.Method { case string(mcp.MethodToolsCall): // New gen_ai namespace attributes (always emitted) diff --git a/pkg/telemetry/middleware_test.go b/pkg/telemetry/middleware_test.go index 32797cc25e..a582e5802c 100644 --- a/pkg/telemetry/middleware_test.go +++ b/pkg/telemetry/middleware_test.go @@ -1719,6 +1719,68 @@ func TestHTTPMiddleware_LegacyAttributes_Disabled(t *testing.T) { assert.NotContains(t, span.attributes, "mcp.tool.arguments") }, }, + { + name: "addMethodSpecificAttributes - Modern clientInfo sets mcp.client.name on non-initialize span", + testFunc: func(t *testing.T, middleware *HTTPMiddleware, span *mockSpan) { + t.Helper() + parsedMCP := &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "github_search", + ClientInfo: map[string]interface{}{"name": "acme-client", "version": "1.0"}, + } + + middleware.addMethodSpecificAttributes(span, parsedMCP) + + assert.Equal(t, "acme-client", span.attributes["mcp.client.name"]) + }, + }, + { + name: "addMethodSpecificAttributes - nil ClientInfo on non-initialize method sets nothing (Legacy no-op)", + testFunc: func(t *testing.T, middleware *HTTPMiddleware, span *mockSpan) { + t.Helper() + parsedMCP := &mcpparser.ParsedMCPRequest{ + Method: "tools/list", + } + + middleware.addMethodSpecificAttributes(span, parsedMCP) + + assert.NotContains(t, span.attributes, "mcp.client.name") + }, + }, + { + name: "addMethodSpecificAttributes - Legacy initialize still sets mcp.client.name from ResourceID", + testFunc: func(t *testing.T, middleware *HTTPMiddleware, span *mockSpan) { + t.Helper() + parsedMCP := &mcpparser.ParsedMCPRequest{ + Method: "initialize", + ResourceID: "legacy-client", + } + + middleware.addMethodSpecificAttributes(span, parsedMCP) + + assert.Equal(t, "legacy-client", span.attributes["mcp.client.name"]) + }, + }, + { + name: "addMethodSpecificAttributes - ClientInfo present but name missing or non-string sets nothing, no panic", + testFunc: func(t *testing.T, middleware *HTTPMiddleware, span *mockSpan) { + t.Helper() + parsedMCP := &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + ClientInfo: map[string]interface{}{"version": "1.0"}, + } + assert.NotPanics(t, func() { + middleware.addMethodSpecificAttributes(span, parsedMCP) + }) + assert.NotContains(t, span.attributes, "mcp.client.name") + + parsedMCP.ClientInfo = map[string]interface{}{"name": 42} + assert.NotPanics(t, func() { + middleware.addMethodSpecificAttributes(span, parsedMCP) + }) + assert.NotContains(t, span.attributes, "mcp.client.name") + }, + }, { name: "finalizeSpan - new response names, no legacy", testFunc: func(t *testing.T, middleware *HTTPMiddleware, span *mockSpan) { diff --git a/pkg/telemetry/propagation.go b/pkg/telemetry/propagation.go index dbc836e1a6..fd93edd2b9 100644 --- a/pkg/telemetry/propagation.go +++ b/pkg/telemetry/propagation.go @@ -5,6 +5,7 @@ package telemetry import ( "context" + "maps" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/propagation" @@ -64,7 +65,8 @@ func (c *MetaCarrier) Meta() map[string]interface{} { } // InjectMetaTraceContext injects the current trace context from ctx directly into -// the given meta map using W3C Trace Context format (traceparent, tracestate). +// the given meta map using W3C Trace Context format (traceparent, tracestate) +// plus W3C Baggage, subject to the globally configured OTEL propagator. // // This function operates directly on the meta map contents. Use this when you // already have the _meta map and want to inject trace context fields into it. @@ -75,3 +77,30 @@ func InjectMetaTraceContext(ctx context.Context, meta map[string]interface{}) { carrier := NewMetaCarrier(meta) otel.GetTextMapPropagator().Inject(ctx, carrier) } + +// MetaWithTraceContext returns a COPY of meta enriched with the current W3C +// trace context (traceparent/tracestate/baggage) from ctx, per SEP-414. +// The caller's map is never mutated (copy-before-mutate). Returns nil when the +// result is empty (no caller meta and no active trace context) so callers can +// omit _meta entirely for 2025-11-25 peers that don't expect it. +// +// Overwrite semantics: injecting the CURRENT span's context overwrites any +// inbound traceparent already present in the cloned caller map. This is the +// correct behavior for an intermediary — the outbound request must be a child +// of the span active here, not of whatever the caller inherited upstream. +// +// SECURITY: the injected trace context AND baggage cross the trust boundary into +// backend requests, and backends may be untrusted. Callers must never place +// secrets or PII in baggage, since it is serialized verbatim into the outbound +// _meta. Baggage is intentionally NOT stripped: SEP-414 explicitly reserves it. +func MetaWithTraceContext(ctx context.Context, meta map[string]interface{}) map[string]interface{} { + enriched := maps.Clone(meta) // nil-safe: returns nil for nil input + if enriched == nil { + enriched = make(map[string]interface{}) + } + InjectMetaTraceContext(ctx, enriched) + if len(enriched) == 0 { + return nil + } + return enriched +} diff --git a/pkg/telemetry/propagation_test.go b/pkg/telemetry/propagation_test.go index 3ade00394d..7927d142be 100644 --- a/pkg/telemetry/propagation_test.go +++ b/pkg/telemetry/propagation_test.go @@ -7,7 +7,10 @@ import ( "context" "testing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/baggage" "go.opentelemetry.io/otel/propagation" sdktrace "go.opentelemetry.io/otel/sdk/trace" ) @@ -152,3 +155,70 @@ func TestInjectMetaTraceContext_NilMeta(t *testing.T) { // Should not panic InjectMetaTraceContext(context.Background(), nil) } + +//nolint:paralleltest // Mutates global OTEL propagator +func TestMetaWithTraceContext(t *testing.T) { + oldPropagator := otel.GetTextMapPropagator() + otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})) + defer otel.SetTextMapPropagator(oldPropagator) + + tp := sdktrace.NewTracerProvider() + defer func() { _ = tp.Shutdown(context.Background()) }() + tracer := tp.Tracer("test") + + spanCtx, span := tracer.Start(context.Background(), "test-span") + defer span.End() + + t.Run("active span with nil meta returns traceparent", func(t *testing.T) { + got := MetaWithTraceContext(spanCtx, nil) + require.NotNil(t, got) + tp, ok := got["traceparent"].(string) + require.True(t, ok, "expected traceparent to be present") + assert.NotEmpty(t, tp) + }) + + t.Run("active span with caller meta preserves caller fields without mutating input", func(t *testing.T) { + input := map[string]any{"progressToken": "tok-123"} + got := MetaWithTraceContext(spanCtx, input) + + require.NotNil(t, got) + assert.Equal(t, "tok-123", got["progressToken"]) + tp, ok := got["traceparent"].(string) + require.True(t, ok, "expected traceparent to be present") + assert.NotEmpty(t, tp) + + // Copy-before-mutate: the caller's original map must be untouched. + assert.Equal(t, map[string]any{"progressToken": "tok-123"}, input) + _, present := input["traceparent"] + assert.False(t, present, "input map must not be mutated with traceparent") + }) + + t.Run("no active span with nil meta returns nil", func(t *testing.T) { + got := MetaWithTraceContext(context.Background(), nil) + assert.Nil(t, got) + }) + + t.Run("no active span with caller meta returns unmutated clone", func(t *testing.T) { + input := map[string]any{"progressToken": "tok-456"} + got := MetaWithTraceContext(context.Background(), input) + + require.NotNil(t, got) + assert.Equal(t, input, got) + _, present := got["traceparent"] + assert.False(t, present, "no traceparent expected without an active trace context") + }) + + t.Run("baggage present in ctx is propagated", func(t *testing.T) { + member, err := baggage.NewMember("user.id", "42") + require.NoError(t, err) + bag, err := baggage.New(member) + require.NoError(t, err) + ctx := baggage.ContextWithBaggage(spanCtx, bag) + + got := MetaWithTraceContext(ctx, nil) + require.NotNil(t, got) + bgKey, ok := got["baggage"].(string) + require.True(t, ok, "expected baggage key to be present") + assert.Contains(t, bgKey, "user.id=42") + }) +} diff --git a/pkg/telemetry/providers/otlp/config.go b/pkg/telemetry/providers/otlp/config.go deleted file mode 100644 index 9564c12f23..0000000000 --- a/pkg/telemetry/providers/otlp/config.go +++ /dev/null @@ -1,14 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -// Package otlp provides OpenTelemetry Protocol (OTLP) provider implementations -package otlp - -// Config holds OTLP-specific configuration -type Config struct { - Endpoint string - Headers map[string]string - Insecure bool - SamplingRate float64 - CACertPath string -} diff --git a/pkg/telemetry/providers/otlp/endpoint.go b/pkg/telemetry/providers/otlp/endpoint.go deleted file mode 100644 index 3ceabd613e..0000000000 --- a/pkg/telemetry/providers/otlp/endpoint.go +++ /dev/null @@ -1,37 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import "strings" - -// Default URL path suffixes for the OTLP/HTTP protocol, as defined in the -// OpenTelemetry specification: -// https://opentelemetry.io/docs/specs/otlp/#otlphttp-request -// -// The Go OTLP SDK normally appends these automatically. However, when the user -// provides a custom base path (e.g. "/api/public/otel" for Langfuse), we must -// call WithURLPath which replaces the entire path. In that case we concatenate -// the base path with the appropriate suffix ourselves (e.g. -// "/api/public/otel" + "/v1/traces"). -const ( - otlpTracesPath = "/v1/traces" - otlpMetricsPath = "/v1/metrics" -) - -// splitEndpointPath separates an OTLP endpoint string into its host:port and -// path components. If no path is present, basePath is empty. -// -// The function defensively strips http:// and https:// prefixes so it works -// correctly even when the scheme has not been removed upstream (e.g. the CLI -// path, which does not call NormalizeTelemetryConfig). -func splitEndpointPath(endpoint string) (hostPort, basePath string) { - endpoint = strings.TrimPrefix(endpoint, "https://") - endpoint = strings.TrimPrefix(endpoint, "http://") - - idx := strings.Index(endpoint, "/") - if idx < 0 { - return endpoint, "" - } - return endpoint[:idx], strings.TrimRight(endpoint[idx:], "/") -} diff --git a/pkg/telemetry/providers/otlp/endpoint_test.go b/pkg/telemetry/providers/otlp/endpoint_test.go deleted file mode 100644 index 2f5a3f5e97..0000000000 --- a/pkg/telemetry/providers/otlp/endpoint_test.go +++ /dev/null @@ -1,92 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import ( - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestSplitEndpointPath(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - endpoint string - wantHostPort string - wantBasePath string - }{ - { - name: "host and port only", - endpoint: "localhost:4318", - wantHostPort: "localhost:4318", - wantBasePath: "", - }, - { - name: "hostname without port", - endpoint: "otel-collector.local", - wantHostPort: "otel-collector.local", - wantBasePath: "", - }, - { - name: "Langfuse endpoint with path", - endpoint: "cloud.langfuse.com/api/public/otel", - wantHostPort: "cloud.langfuse.com", - wantBasePath: "/api/public/otel", - }, - { - name: "LangSmith endpoint with port and path", - endpoint: "smith.langchain.com:443/api/v1/otel", - wantHostPort: "smith.langchain.com:443", - wantBasePath: "/api/v1/otel", - }, - { - name: "trailing slash stripped", - endpoint: "cloud.langfuse.com/api/public/otel/", - wantHostPort: "cloud.langfuse.com", - wantBasePath: "/api/public/otel", - }, - { - name: "host:port with trailing slash only", - endpoint: "localhost:4318/", - wantHostPort: "localhost:4318", - wantBasePath: "", - }, - { - name: "https scheme stripped before splitting", - endpoint: "https://cloud.langfuse.com/api/public/otel", - wantHostPort: "cloud.langfuse.com", - wantBasePath: "/api/public/otel", - }, - { - name: "http scheme stripped before splitting", - endpoint: "http://localhost:4318", - wantHostPort: "localhost:4318", - wantBasePath: "", - }, - { - name: "https scheme with host only", - endpoint: "https://api.honeycomb.io", - wantHostPort: "api.honeycomb.io", - wantBasePath: "", - }, - { - name: "empty string", - endpoint: "", - wantHostPort: "", - wantBasePath: "", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - hostPort, basePath := splitEndpointPath(tt.endpoint) - assert.Equal(t, tt.wantHostPort, hostPort) - assert.Equal(t, tt.wantBasePath, basePath) - }) - } -} diff --git a/pkg/telemetry/providers/otlp/logging.go b/pkg/telemetry/providers/otlp/logging.go deleted file mode 100644 index 667570195b..0000000000 --- a/pkg/telemetry/providers/otlp/logging.go +++ /dev/null @@ -1,10 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -// This file is a placeholder for future OTLP logging support -// OpenTelemetry Logs are still in development and will be added when stable - -// TODO: Implement OTLP logging when the specification is stable -// Reference: https://opentelemetry.io/docs/specs/otel/logs/ diff --git a/pkg/telemetry/providers/otlp/metrics.go b/pkg/telemetry/providers/otlp/metrics.go deleted file mode 100644 index 906d2a66da..0000000000 --- a/pkg/telemetry/providers/otlp/metrics.go +++ /dev/null @@ -1,56 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -// Package otlp provides OpenTelemetry Protocol (OTLP) provider implementations -package otlp - -import ( - "context" - "fmt" - - "go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" -) - -// NewMetricReader creates an OTLP metric reader for use in a unified meter provider -func NewMetricReader(ctx context.Context, config Config) (sdkmetric.Reader, error) { - if config.Endpoint == "" { - return nil, fmt.Errorf("OTLP endpoint is required") - } - - exporter, err := createMetricExporter(ctx, config) - if err != nil { - return nil, fmt.Errorf("failed to create OTLP metric exporter: %w", err) - } - - return sdkmetric.NewPeriodicReader(exporter), nil -} - -func createMetricExporter(ctx context.Context, config Config) (sdkmetric.Exporter, error) { - host, basePath := splitEndpointPath(config.Endpoint) - opts := []otlpmetrichttp.Option{ - otlpmetrichttp.WithEndpoint(host), - } - - if basePath != "" { - opts = append(opts, otlpmetrichttp.WithURLPath(basePath+otlpMetricsPath)) - } - - if len(config.Headers) > 0 { - opts = append(opts, otlpmetrichttp.WithHeaders(config.Headers)) - } - - if config.Insecure { - opts = append(opts, otlpmetrichttp.WithInsecure()) - } - - if config.CACertPath != "" { - tlsCfg, err := newTLSConfigFromCA(config.CACertPath) - if err != nil { - return nil, fmt.Errorf("failed to configure TLS for metric exporter: %w", err) - } - opts = append(opts, otlpmetrichttp.WithTLSClientConfig(tlsCfg)) - } - - return otlpmetrichttp.New(ctx, opts...) -} diff --git a/pkg/telemetry/providers/otlp/metrics_test.go b/pkg/telemetry/providers/otlp/metrics_test.go deleted file mode 100644 index 6d52f71e04..0000000000 --- a/pkg/telemetry/providers/otlp/metrics_test.go +++ /dev/null @@ -1,156 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import ( - "context" - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestCreateMetricExporter(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - ctx func() context.Context - wantErr bool - errMsg string - }{ - { - name: "valid config", - config: Config{ - Endpoint: "localhost:4318", - Headers: map[string]string{"x-api-key": "secret"}, - Insecure: true, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "config without headers", - config: Config{ - Endpoint: "localhost:4318", - Insecure: false, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "endpoint with custom path", - config: Config{ - Endpoint: "cloud.langfuse.com/api/public/otel", - Headers: map[string]string{"Authorization": "Basic abc123"}, - Insecure: false, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "error creating metrics exporter due to invalid CA cert", - config: Config{ - Endpoint: "localhost:4318", - Insecure: false, - CACertPath: "/nonexistent/ca.crt", - }, - ctx: func() context.Context { return context.Background() }, - wantErr: true, - errMsg: "failed to configure TLS for metric exporter", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := tt.ctx() - exporter, err := createMetricExporter(ctx, tt.config) - - if tt.wantErr { - assert.Error(t, err) - assert.Nil(t, exporter) - if tt.errMsg != "" { - assert.Contains(t, err.Error(), tt.errMsg) - } - } else { - assert.NoError(t, err) - assert.NotNil(t, exporter) - // Clean up - _ = exporter.Shutdown(ctx) - } - }) - } -} - -func TestNewMetricReader(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - wantErr bool - errMsg string - }{ - { - name: "valid config", - config: Config{ - Endpoint: "localhost:4318", - Headers: map[string]string{"Authorization": "Bearer token"}, - Insecure: true, - }, - wantErr: false, - }, - { - name: "missing endpoint", - config: Config{ - Headers: map[string]string{"Authorization": "Bearer token"}, - }, - wantErr: true, - errMsg: "OTLP endpoint is required", - }, - { - name: "config with custom headers", - config: Config{ - Endpoint: "otel-collector.local:4318", - Headers: map[string]string{ - "x-api-key": "secret", - "x-env": "production", - }, - Insecure: false, - }, - wantErr: false, - }, - { - name: "endpoint with custom path", - config: Config{ - Endpoint: "cloud.langfuse.com/api/public/otel", - Headers: map[string]string{"Authorization": "Basic abc123"}, - Insecure: false, - }, - wantErr: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := context.Background() - reader, err := NewMetricReader(ctx, tt.config) - - if tt.wantErr { - assert.Error(t, err) - assert.Nil(t, reader) - if tt.errMsg != "" { - assert.Contains(t, err.Error(), tt.errMsg) - } - } else { - assert.NoError(t, err) - assert.NotNil(t, reader) - } - }) - } -} diff --git a/pkg/telemetry/providers/otlp/tls.go b/pkg/telemetry/providers/otlp/tls.go deleted file mode 100644 index b038c54981..0000000000 --- a/pkg/telemetry/providers/otlp/tls.go +++ /dev/null @@ -1,37 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import ( - "crypto/tls" - "crypto/x509" - "fmt" - "log/slog" - "os" -) - -// newTLSConfigFromCA creates a tls.Config that trusts certificates from the given -// CA bundle file path. The returned config appends the custom CAs to the system -// pool so both the custom CA and standard public CAs are trusted. -func newTLSConfigFromCA(caCertPath string) (*tls.Config, error) { - caCert, err := os.ReadFile(caCertPath) // #nosec G304 - path comes from operator-controlled mount - if err != nil { - return nil, fmt.Errorf("failed to read CA certificate bundle %q: %w", caCertPath, err) - } - - caCertPool, err := x509.SystemCertPool() - if err != nil { - slog.Warn("System CA pool unavailable, using custom CA only", "error", err) - caCertPool = x509.NewCertPool() - } - - if !caCertPool.AppendCertsFromPEM(caCert) { - return nil, fmt.Errorf("failed to parse CA certificate bundle %q: no valid PEM certificates found", caCertPath) - } - - return &tls.Config{ - RootCAs: caCertPool, - MinVersion: tls.VersionTLS12, - }, nil -} diff --git a/pkg/telemetry/providers/otlp/tls_test.go b/pkg/telemetry/providers/otlp/tls_test.go deleted file mode 100644 index a658b0dad8..0000000000 --- a/pkg/telemetry/providers/otlp/tls_test.go +++ /dev/null @@ -1,117 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import ( - "crypto/ecdsa" - "crypto/elliptic" - "crypto/rand" - "crypto/tls" - "crypto/x509" - "crypto/x509/pkix" - "encoding/pem" - "math/big" - "os" - "path/filepath" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// generateSelfSignedCACert creates a PEM-encoded self-signed CA certificate -// for use in tests. -func generateSelfSignedCACert(t *testing.T) []byte { - t.Helper() - - key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) - require.NoError(t, err) - - template := &x509.Certificate{ - SerialNumber: big.NewInt(1), - Subject: pkix.Name{ - Organization: []string{"ToolHive Test CA"}, - }, - NotBefore: time.Now().Add(-time.Hour), - NotAfter: time.Now().Add(time.Hour), - KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign, - BasicConstraintsValid: true, - IsCA: true, - } - - certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key) - require.NoError(t, err) - - return pem.EncodeToMemory(&pem.Block{ - Type: "CERTIFICATE", - Bytes: certDER, - }) -} - -func TestNewTLSConfigFromCA(t *testing.T) { - t.Parallel() - - // Create a valid CA cert file once for the subtests that need it. - caCertPEM := generateSelfSignedCACert(t) - validCertPath := filepath.Join(t.TempDir(), "ca.crt") - require.NoError(t, os.WriteFile(validCertPath, caCertPEM, 0o600)) - - // Create a file with invalid PEM content. - invalidPEMPath := filepath.Join(t.TempDir(), "bad.crt") - require.NoError(t, os.WriteFile(invalidPEMPath, []byte("not a cert"), 0o600)) - - tests := []struct { - name string - caCertPath string - wantErr bool - errMsg string - validate func(t *testing.T, cfg *tls.Config) - }{ - { - name: "valid CA cert file", - caCertPath: validCertPath, - wantErr: false, - validate: func(t *testing.T, cfg *tls.Config) { - t.Helper() - assert.Equal(t, uint16(tls.VersionTLS12), cfg.MinVersion) - assert.NotNil(t, cfg.RootCAs) - }, - }, - { - name: "non-existent file", - caCertPath: filepath.Join(t.TempDir(), "does-not-exist.crt"), - wantErr: true, - errMsg: "failed to read CA certificate bundle", - }, - { - name: "invalid PEM content", - caCertPath: invalidPEMPath, - wantErr: true, - errMsg: "no valid PEM certificates found", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - cfg, err := newTLSConfigFromCA(tt.caCertPath) - - if tt.wantErr { - assert.Error(t, err) - assert.Nil(t, cfg) - if tt.errMsg != "" { - assert.Contains(t, err.Error(), tt.errMsg) - } - } else { - require.NoError(t, err) - require.NotNil(t, cfg) - if tt.validate != nil { - tt.validate(t, cfg) - } - } - }) - } -} diff --git a/pkg/telemetry/providers/otlp/tracing.go b/pkg/telemetry/providers/otlp/tracing.go deleted file mode 100644 index a76f518582..0000000000 --- a/pkg/telemetry/providers/otlp/tracing.go +++ /dev/null @@ -1,90 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import ( - "context" - "fmt" - - "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" - "go.opentelemetry.io/otel/sdk/resource" - sdktrace "go.opentelemetry.io/otel/sdk/trace" - "go.opentelemetry.io/otel/trace" - tracenoop "go.opentelemetry.io/otel/trace/noop" -) - -func createTraceExporter(ctx context.Context, config Config) (sdktrace.SpanExporter, error) { - host, basePath := splitEndpointPath(config.Endpoint) - opts := []otlptracehttp.Option{ - otlptracehttp.WithEndpoint(host), - } - - if basePath != "" { - opts = append(opts, otlptracehttp.WithURLPath(basePath+otlpTracesPath)) - } - - if len(config.Headers) > 0 { - opts = append(opts, otlptracehttp.WithHeaders(config.Headers)) - } - - if config.Insecure { - opts = append(opts, otlptracehttp.WithInsecure()) - } - - if config.CACertPath != "" { - tlsCfg, err := newTLSConfigFromCA(config.CACertPath) - if err != nil { - return nil, fmt.Errorf("failed to configure TLS for trace exporter: %w", err) - } - opts = append(opts, otlptracehttp.WithTLSClientConfig(tlsCfg)) - } - - exporter, err := otlptracehttp.New(ctx, opts...) - if err != nil { - return nil, fmt.Errorf("failed to create trace exporter: %w", err) - } - return exporter, nil -} - -// NewTracerProviderWithShutdown creates an OTLP tracer provider with a shutdown function. -// Additional span processors (e.g. a Sentry bridge) can be registered via extraProcessors. -// When endpoint is empty but extra processors are provided, a real SDK provider is created -// without an OTLP exporter so the processors still receive spans. -func NewTracerProviderWithShutdown( - ctx context.Context, - config Config, - res *resource.Resource, - extraProcessors ...sdktrace.SpanProcessor, -) (trace.TracerProvider, func(context.Context) error, error) { - // True no-op only when there is nothing at all to do. - if config.Endpoint == "" && len(extraProcessors) == 0 { - return tracenoop.NewTracerProvider(), nil, nil - } - - opts := []sdktrace.TracerProviderOption{ - sdktrace.WithResource(res), - // ParentBased ensures that when an incoming W3C traceparent header marks - // the parent as sampled (e.g. from ToolHive Studio), the child span is - // always sampled regardless of the local ratio. Without ParentBased, a - // bare TraceIDRatioBased sampler could drop a span even when the remote - // parent was sampled, breaking end-to-end distributed trace correlation. - sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(config.SamplingRate))), - } - - // Only wire an OTLP exporter when an endpoint is actually configured. - if config.Endpoint != "" { - exporter, err := createTraceExporter(ctx, config) - if err != nil { - return nil, nil, fmt.Errorf("failed to create trace provider: %w", err) - } - opts = append(opts, sdktrace.WithBatcher(exporter)) - } - - for _, p := range extraProcessors { - opts = append(opts, sdktrace.WithSpanProcessor(p)) - } - - provider := sdktrace.NewTracerProvider(opts...) - return provider, provider.Shutdown, nil -} diff --git a/pkg/telemetry/providers/otlp/tracing_test.go b/pkg/telemetry/providers/otlp/tracing_test.go deleted file mode 100644 index 87ee5b3aad..0000000000 --- a/pkg/telemetry/providers/otlp/tracing_test.go +++ /dev/null @@ -1,220 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package otlp - -import ( - "context" - "fmt" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "go.opentelemetry.io/otel/sdk/resource" - semconv "go.opentelemetry.io/otel/semconv/v1.21.0" -) - -func TestCreateTraceExporter(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - ctx func() context.Context - wantErr bool - errMsg string - }{ - { - name: "valid config", - config: Config{ - Endpoint: "localhost:4318", - Headers: map[string]string{"Authorization": "Bearer token"}, - Insecure: true, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "config with headers", - config: Config{ - Endpoint: "localhost:4318", - Headers: map[string]string{ - "x-api-key": "secret", - "x-env": "test", - }, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "secure config", - config: Config{ - Endpoint: "secure.example.com:4318", - Insecure: false, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "endpoint with custom path", - config: Config{ - Endpoint: "cloud.langfuse.com/api/public/otel", - Headers: map[string]string{"Authorization": "Basic abc123"}, - Insecure: false, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - }, - { - name: "error creating sdk-span-exporter due to error (cancelled context)", - config: Config{ - Endpoint: "secure.example.com:4318", - Insecure: true, - }, - ctx: func() context.Context { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - return ctx - }, - wantErr: true, - errMsg: "context canceled", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := tt.ctx() - exporter, err := createTraceExporter(ctx, tt.config) - - if tt.wantErr { - assert.Error(t, err) - assert.Nil(t, exporter) - if tt.errMsg != "" { - assert.Contains(t, err.Error(), tt.errMsg) - } - } else { - assert.NoError(t, err) - assert.NotNil(t, exporter) - // Clean up - _ = exporter.Shutdown(ctx) - } - }) - } -} - -func TestNewTracerProviderWithShutdown(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - ctx func() context.Context - wantErr bool - errMsg string - expectNoOp bool - expectShutdown bool - }{ - { - name: "valid config with endpoint returns SDK provider with shutdown", - config: Config{ - Endpoint: "localhost:4318", - SamplingRate: 0.5, - Headers: map[string]string{"Authorization": "Bearer token"}, - Insecure: true, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - expectNoOp: false, - expectShutdown: true, - }, - { - name: "no endpoint returns noop provider with nil shutdown", - config: Config{ - SamplingRate: 0.1, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - expectNoOp: true, - expectShutdown: false, - }, - { - name: "config with custom sampling returns SDK provider with shutdown", - config: Config{ - Endpoint: "localhost:4318", - SamplingRate: 1.0, // Always sample - Insecure: true, - }, - ctx: func() context.Context { return context.Background() }, - wantErr: false, - expectNoOp: false, - expectShutdown: true, - }, - { - name: "error creating trace exporter propagates error", - config: Config{ - Endpoint: "localhost:4318", - SamplingRate: 1.0, - Insecure: true, - }, - ctx: func() context.Context { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - return ctx - }, - wantErr: true, - errMsg: "failed to create trace exporter", - expectNoOp: false, - expectShutdown: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := tt.ctx() - res, err := resource.New(ctx, - resource.WithAttributes( - semconv.ServiceName("test-service"), - semconv.ServiceVersion("1.0.0"), - ), - ) - require.NoError(t, err) - - provider, shutdown, err := NewTracerProviderWithShutdown(ctx, tt.config, res) - - if tt.wantErr { - assert.Error(t, err) - assert.Nil(t, provider) - assert.Nil(t, shutdown) - if tt.errMsg != "" { - assert.Contains(t, err.Error(), tt.errMsg) - } - } else { - assert.NoError(t, err) - assert.NotNil(t, provider) - - // Check provider type - providerType := fmt.Sprintf("%T", provider) - if tt.expectNoOp { - assert.Contains(t, providerType, "noop") - } else { - assert.NotContains(t, providerType, "noop") - } - - // Check shutdown function - if tt.expectShutdown { - assert.NotNil(t, shutdown) - // Test that shutdown function works - shutdownCtx := context.Background() - err := shutdown(shutdownCtx) - assert.NoError(t, err) - } else { - assert.Nil(t, shutdown) - } - } - }) - } -} diff --git a/pkg/telemetry/providers/prometheus/prometheus.go b/pkg/telemetry/providers/prometheus/prometheus.go deleted file mode 100644 index 753e6214e8..0000000000 --- a/pkg/telemetry/providers/prometheus/prometheus.go +++ /dev/null @@ -1,54 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -// Package prometheus provides Prometheus metric exporter implementation -package prometheus - -import ( - "fmt" - "net/http" - - promclient "github.com/prometheus/client_golang/prometheus" - "github.com/prometheus/client_golang/prometheus/collectors" - "github.com/prometheus/client_golang/prometheus/promhttp" - "go.opentelemetry.io/otel/exporters/prometheus" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" -) - -// Config holds Prometheus-specific configuration -type Config struct { - // EnableMetricsPath controls whether to expose Prometheus-style /metrics endpoint - EnableMetricsPath bool - // IncludeRuntimeMetrics adds Go runtime metrics to the registry - IncludeRuntimeMetrics bool -} - -// NewReader creates a Prometheus metric reader and HTTP handler for use in a unified meter provider -func NewReader(config Config) (sdkmetric.Reader, http.Handler, error) { - if !config.EnableMetricsPath { - return nil, nil, fmt.Errorf("prometheus provider requires EnableMetricsPath to be true") - } - - // Create a dedicated registry - registry := promclient.NewRegistry() - - // Add runtime metrics if requested - if config.IncludeRuntimeMetrics { - registry.MustRegister(collectors.NewGoCollector()) - registry.MustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{})) - } - - // Create the Prometheus exporter (which is also a Reader) - exporter, err := prometheus.New(prometheus.WithRegisterer(registry)) - if err != nil { - return nil, nil, fmt.Errorf("failed to create prometheus exporter: %w", err) - } - - // Create HTTP handler - handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{ - ErrorHandling: promhttp.ContinueOnError, - ErrorLog: nil, - }) - - return exporter, handler, nil -} diff --git a/pkg/telemetry/providers/prometheus/prometheus_test.go b/pkg/telemetry/providers/prometheus/prometheus_test.go deleted file mode 100644 index e154f01f67..0000000000 --- a/pkg/telemetry/providers/prometheus/prometheus_test.go +++ /dev/null @@ -1,145 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package prometheus - -import ( - "context" - "net/http" - "net/http/httptest" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" - "go.opentelemetry.io/otel/sdk/resource" - semconv "go.opentelemetry.io/otel/semconv/v1.21.0" -) - -func TestNewReader(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - wantErr bool - errMsg string - checkHandler bool - checkRuntimeMetrics bool - }{ - { - name: "valid config with runtime metrics", - config: Config{ - EnableMetricsPath: true, - IncludeRuntimeMetrics: true, - }, - wantErr: false, - checkHandler: true, - checkRuntimeMetrics: true, - }, - { - name: "valid config without runtime metrics", - config: Config{ - EnableMetricsPath: true, - IncludeRuntimeMetrics: false, - }, - wantErr: false, - checkHandler: true, - checkRuntimeMetrics: false, - }, - { - name: "metrics path not enabled", - config: Config{ - EnableMetricsPath: false, - IncludeRuntimeMetrics: true, - }, - wantErr: true, - errMsg: "requires EnableMetricsPath", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - reader, handler, err := NewReader(tt.config) - - if tt.wantErr { - assert.Error(t, err) - assert.Nil(t, reader) - assert.Nil(t, handler) - if tt.errMsg != "" { - assert.Contains(t, err.Error(), tt.errMsg) - } - } else { - assert.NoError(t, err) - assert.NotNil(t, reader) - - if tt.checkHandler { - assert.NotNil(t, handler) - assert.Implements(t, (*http.Handler)(nil), handler) - - // Test that the handler works - req := httptest.NewRequest(http.MethodGet, "/metrics", nil) - rec := httptest.NewRecorder() - handler.ServeHTTP(rec, req) - assert.Equal(t, http.StatusOK, rec.Code) - - // Check for runtime metrics if expected - if tt.checkRuntimeMetrics { - assert.Contains(t, rec.Body.String(), "go_") - assert.Contains(t, rec.Body.String(), "process_") - } - } - } - }) - } -} - -func TestNewReader_Integration(t *testing.T) { - t.Parallel() - - // This test ensures NewReader works correctly with a meter provider - config := Config{ - EnableMetricsPath: true, - IncludeRuntimeMetrics: false, - } - - reader, handler, err := NewReader(config) - require.NoError(t, err) - require.NotNil(t, reader) - require.NotNil(t, handler) - - // Create a meter provider with the reader - ctx := context.Background() - res, err := resource.New(ctx, - resource.WithAttributes( - semconv.ServiceName("test-service"), - semconv.ServiceVersion("1.0.0"), - ), - ) - require.NoError(t, err) - - meterProvider := sdkmetric.NewMeterProvider( - sdkmetric.WithResource(res), - sdkmetric.WithReader(reader), - ) - defer meterProvider.Shutdown(ctx) - - // Create a metric - meter := meterProvider.Meter("test") - counter, err := meter.Int64Counter("test_reader_counter") - require.NoError(t, err) - - // Add some values - counter.Add(ctx, 5) - counter.Add(ctx, 10) - - // Check that metrics appear in the handler output - req := httptest.NewRequest(http.MethodGet, "/metrics", nil) - rec := httptest.NewRecorder() - handler.ServeHTTP(rec, req) - - assert.Equal(t, http.StatusOK, rec.Code) - assert.Contains(t, rec.Body.String(), "test_reader_counter") -} diff --git a/pkg/telemetry/providers/providers.go b/pkg/telemetry/providers/providers.go deleted file mode 100644 index 924d68289f..0000000000 --- a/pkg/telemetry/providers/providers.go +++ /dev/null @@ -1,343 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -// Package providers contains telemetry provider implementations and builder logic -package providers - -import ( - "context" - "fmt" - "log/slog" - "net/http" - "time" - - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/metric" - "go.opentelemetry.io/otel/metric/noop" - "go.opentelemetry.io/otel/sdk/resource" - sdktrace "go.opentelemetry.io/otel/sdk/trace" - semconv "go.opentelemetry.io/otel/semconv/v1.21.0" - "go.opentelemetry.io/otel/trace" - tracenoop "go.opentelemetry.io/otel/trace/noop" -) - -// Config holds the telemetry configuration for all providers. -// It contains service information, OTLP settings, and Prometheus configuration. -type Config struct { - // Service information - ServiceName string // ServiceName identifies the service for telemetry data - ServiceVersion string // ServiceVersion identifies the service version for telemetry data - - // OTLP configuration - OTLPEndpoint string // OTLPEndpoint is the OTLP collector endpoint (e.g., "localhost:4318") - Headers map[string]string // Headers are additional headers to send with OTLP requests - Insecure bool // Insecure enables insecure transport (no TLS) for OTLP - TracingEnabled bool // TracingEnabled controls whether tracing is enabled for OTLP - MetricsEnabled bool // MetricsEnabled controls whether metrics are enabled for OTLP - SamplingRate float64 // SamplingRate controls trace sampling (0.0 to 1.0) - - // Prometheus configuration - EnablePrometheusMetricsPath bool // EnablePrometheusMetricsPath enables Prometheus /metrics endpoint - - // TLS configuration - CACertPath string // CACertPath is the file path to a custom CA certificate bundle for the OTLP endpoint - - // Custom attributes - // CustomAttributes are additional resource attributes to include (as map for JSON serialization) - CustomAttributes map[string]string - - // ExtraSpanProcessors contains additional OTEL span processors to register alongside the - // default OTLP exporter (e.g. a Sentry bridge processor for dual export). - ExtraSpanProcessors []sdktrace.SpanProcessor -} - -// ProviderOption is an option type used to configure the telemetry providers -type ProviderOption func(*Config) error - -// WithServiceName sets the service name -func WithServiceName(serviceName string) ProviderOption { - return func(config *Config) error { - if serviceName == "" { - return fmt.Errorf("service name cannot be empty") - } - config.ServiceName = serviceName - return nil - } -} - -// WithServiceVersion sets the service version -func WithServiceVersion(serviceVersion string) ProviderOption { - return func(config *Config) error { - if serviceVersion == "" { - return fmt.Errorf("service version cannot be empty") - } - config.ServiceVersion = serviceVersion - return nil - } -} - -// WithOTLPEndpoint sets the OTLP endpoint -func WithOTLPEndpoint(endpoint string) ProviderOption { - return func(config *Config) error { - config.OTLPEndpoint = endpoint - return nil - } -} - -// WithHeaders sets the headers -func WithHeaders(headers map[string]string) ProviderOption { - return func(config *Config) error { - config.Headers = headers - return nil - } -} - -// WithInsecure sets the insecure flag -func WithInsecure(insecure bool) ProviderOption { - return func(config *Config) error { - config.Insecure = insecure - return nil - } -} - -// WithCACertPath sets the CA certificate path for the OTLP endpoint -func WithCACertPath(path string) ProviderOption { - return func(config *Config) error { - config.CACertPath = path - return nil - } -} - -// WithTracingEnabled sets the tracing enabled flag -func WithTracingEnabled(tracingEnabled bool) ProviderOption { - return func(config *Config) error { - config.TracingEnabled = tracingEnabled - return nil - } -} - -// WithMetricsEnabled sets the metrics enabled flag -func WithMetricsEnabled(metricsEnabled bool) ProviderOption { - return func(config *Config) error { - config.MetricsEnabled = metricsEnabled - return nil - } -} - -// WithSamplingRate sets the sampling rate -func WithSamplingRate(samplingRate float64) ProviderOption { - return func(config *Config) error { - config.SamplingRate = samplingRate - return nil - } -} - -// WithEnablePrometheusMetricsPath sets the enable prometheus metrics path flag -func WithEnablePrometheusMetricsPath(enablePrometheusMetricsPath bool) ProviderOption { - return func(config *Config) error { - config.EnablePrometheusMetricsPath = enablePrometheusMetricsPath - return nil - } -} - -// WithCustomAttributes sets the custom resource attributes -func WithCustomAttributes(attributes map[string]string) ProviderOption { - return func(config *Config) error { - config.CustomAttributes = attributes - return nil - } -} - -// WithExtraSpanProcessors appends additional OTEL span processors (e.g. a Sentry bridge). -func WithExtraSpanProcessors(processors ...sdktrace.SpanProcessor) ProviderOption { - return func(config *Config) error { - config.ExtraSpanProcessors = append(config.ExtraSpanProcessors, processors...) - return nil - } -} - -// CompositeProvider combines telemetry providers into a single interface. -// It manages tracer providers, meter providers, Prometheus handlers, and cleanup. -type CompositeProvider struct { - tracerProvider trace.TracerProvider // tracerProvider provides distributed tracing - meterProvider metric.MeterProvider // meterProvider provides metrics collection - prometheusHandler http.Handler // prometheusHandler serves Prometheus metrics - shutdownFuncs []func(context.Context) error // shutdownFuncs clean up resources on shutdown -} - -// NewCompositeProvider creates the appropriate providers based on provided options -func NewCompositeProvider( - ctx context.Context, - options ...ProviderOption, -) (*CompositeProvider, error) { - config := Config{} - for _, option := range options { - if err := option(&config); err != nil { - return nil, err - } - } - - // Create resource for all providers - // Start with base attributes - baseAttrs := []attribute.KeyValue{ - semconv.ServiceName(config.ServiceName), - semconv.ServiceVersion(config.ServiceVersion), - } - - // Add custom attributes from CLI flags - if len(config.CustomAttributes) > 0 { - slog.Debug("adding custom attributes to OTEL resource", - "count", len(config.CustomAttributes)) - for key, value := range config.CustomAttributes { - //nolint:gosec // G706: custom attribute key-value from config - slog.Debug("adding custom attribute to resource", - "key", key, "value", value) - baseAttrs = append(baseAttrs, attribute.String(key, value)) - } - } - - // Create resource with base attributes and support for OTEL_RESOURCE_ATTRIBUTES env var - res, err := resource.New(ctx, - resource.WithAttributes(baseAttrs...), - resource.WithFromEnv(), // This reads OTEL_RESOURCE_ATTRIBUTES automatically - resource.WithHost(), // Add host information - ) - if err != nil { - return nil, fmt.Errorf("failed to create resource with service name '%s' and version '%s': %w", - config.ServiceName, config.ServiceVersion, err) - } - - // Use strategy selector to determine provider strategies - selector := &StrategySelector{config: config} - - // Early return for no-op case - if selector.IsFullyNoOp() { - slog.Debug("no telemetry configured, using no-op providers") - return createNoOpProvider(), nil - } - - // Build composite provider using strategies - return buildProviders(ctx, config, selector, res) -} - -func createNoOpProvider() *CompositeProvider { - return &CompositeProvider{ - tracerProvider: tracenoop.NewTracerProvider(), - meterProvider: noop.NewMeterProvider(), - prometheusHandler: nil, - shutdownFuncs: []func(context.Context) error{}, - } -} - -// buildProviders creates a composite provider using the selected strategies -func buildProviders( - ctx context.Context, - config Config, - selector *StrategySelector, - res *resource.Resource, -) (*CompositeProvider, error) { - composite := &CompositeProvider{ - shutdownFuncs: []func(context.Context) error{}, - } - - if err := createMetricsProvider(ctx, config, composite, selector, res); err != nil { - return nil, err - } - - if err := createTracingProvider(ctx, config, composite, selector, res); err != nil { - return nil, err - } - - slog.Debug("telemetry providers created successfully") - return composite, nil -} - -// createMetricsProvider creates the metrics provider for the composite provider -func createMetricsProvider( - ctx context.Context, - config Config, - composite *CompositeProvider, - selector *StrategySelector, - res *resource.Resource, -) error { - // Create meter provider using selected strategy - meterStrategy := selector.SelectMeterStrategy() - meterResult, err := meterStrategy.CreateMeterProvider(ctx, config, res) - if err != nil { - return fmt.Errorf( - "failed to create meter provider with config (endpoint: %s, metrics enabled: %t, prometheus enabled: %t): %w", - config.OTLPEndpoint, - config.MetricsEnabled, - config.EnablePrometheusMetricsPath, - err) - } - - composite.meterProvider = meterResult.MeterProvider - composite.prometheusHandler = meterResult.PrometheusHandler - - if meterResult.ShutdownFunc != nil { - composite.shutdownFuncs = append(composite.shutdownFuncs, meterResult.ShutdownFunc) - } - - return nil -} - -// createTracingProvider creates the tracing provider for the composite provider -func createTracingProvider( - ctx context.Context, - config Config, - composite *CompositeProvider, - selector *StrategySelector, - res *resource.Resource, -) error { - // Create tracer provider using selected strategy - tracerStrategy := selector.SelectTracerStrategy() - tracerProvider, tracerShutdown, err := tracerStrategy.CreateTracerProvider(ctx, config, res) - if err != nil { - return fmt.Errorf("failed to create tracer provider with config (endpoint: %s, tracing enabled: %t): %w", - config.OTLPEndpoint, - config.TracingEnabled, - err) - } - - composite.tracerProvider = tracerProvider - - if tracerShutdown != nil { - composite.shutdownFuncs = append(composite.shutdownFuncs, tracerShutdown) - } - - return nil -} - -// TracerProvider returns the tracer provider -func (p *CompositeProvider) TracerProvider() trace.TracerProvider { - return p.tracerProvider -} - -// MeterProvider returns the primary meter provider -func (p *CompositeProvider) MeterProvider() metric.MeterProvider { - return p.meterProvider -} - -// PrometheusHandler returns the Prometheus metrics handler if configured -func (p *CompositeProvider) PrometheusHandler() http.Handler { - return p.prometheusHandler -} - -// Shutdown gracefully shuts down all providers -func (p *CompositeProvider) Shutdown(ctx context.Context) error { - shutdownCtx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - - var errs []error - for i, shutdown := range p.shutdownFuncs { - if err := shutdown(shutdownCtx); err != nil { - errs = append(errs, fmt.Errorf("provider %d shutdown failed: %w", i, err)) - } - } - - if len(errs) > 0 { - return fmt.Errorf("shutdown failed with %d errors: %v", len(errs), errs) - } - return nil -} diff --git a/pkg/telemetry/providers/providers_strategy.go b/pkg/telemetry/providers/providers_strategy.go deleted file mode 100644 index bb30746425..0000000000 --- a/pkg/telemetry/providers/providers_strategy.go +++ /dev/null @@ -1,256 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package providers - -import ( - "context" - "fmt" - "log/slog" - "net/http" - - "go.opentelemetry.io/otel/metric" - "go.opentelemetry.io/otel/metric/noop" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" - "go.opentelemetry.io/otel/sdk/resource" - "go.opentelemetry.io/otel/trace" - tracenoop "go.opentelemetry.io/otel/trace/noop" - - "github.com/stacklok/toolhive/pkg/telemetry/providers/otlp" - "github.com/stacklok/toolhive/pkg/telemetry/providers/prometheus" -) - -// TracerStrategy defines the interface for creating tracer providers. -// Implementations create trace providers based on configuration and resource information. -type TracerStrategy interface { - // CreateTracerProvider creates a tracer provider with optional shutdown function - CreateTracerProvider(ctx context.Context, config Config, res *resource.Resource) ( - trace.TracerProvider, func(context.Context) error, error) -} - -// NoOpTracerStrategy creates a no-op tracer provider that discards all trace data. -// It's used when tracing is disabled or no OTLP endpoint is configured. -type NoOpTracerStrategy struct{} - -// CreateTracerProvider creates a no-op tracer provider -func (*NoOpTracerStrategy) CreateTracerProvider( - _ context.Context, - _ Config, - _ *resource.Resource, -) (trace.TracerProvider, func(context.Context) error, error) { - slog.Debug("creating no-op tracer provider") - return tracenoop.NewTracerProvider(), nil, nil -} - -// OTLPTracerStrategy creates an OTLP tracer provider that sends traces to an OTLP collector. -// It supports sampling configuration, custom headers, and secure/insecure transport. -type OTLPTracerStrategy struct{} - -// CreateTracerProvider creates an OTLP tracer provider with the configured endpoint and sampling rate -func (*OTLPTracerStrategy) CreateTracerProvider( - ctx context.Context, - config Config, - res *resource.Resource, -) (trace.TracerProvider, func(context.Context) error, error) { - //nolint:gosec // G706: OTLP endpoint from config - slog.Debug("creating OTLP tracer provider", - "endpoint", config.OTLPEndpoint, - "sampling_rate", config.SamplingRate, - "extra_processors", len(config.ExtraSpanProcessors)) - - otlpConfig := otlp.Config{ - Endpoint: config.OTLPEndpoint, - Headers: config.Headers, - Insecure: config.Insecure, - SamplingRate: config.SamplingRate, - CACertPath: config.CACertPath, - } - - provider, shutdown, err := otlp.NewTracerProviderWithShutdown(ctx, otlpConfig, res, config.ExtraSpanProcessors...) - if err != nil { - return nil, nil, fmt.Errorf("failed to create OTLP tracer provider for endpoint %s: %w", config.OTLPEndpoint, err) - } - return provider, shutdown, nil -} - -// MeterResult contains the result of creating a meter provider -type MeterResult struct { - MeterProvider metric.MeterProvider - PrometheusHandler http.Handler - ShutdownFunc func(context.Context) error -} - -// MeterStrategy defines the interface for creating meter providers -type MeterStrategy interface { - CreateMeterProvider(ctx context.Context, config Config, res *resource.Resource) (*MeterResult, error) -} - -// NoOpMeterStrategy creates a no-op meter provider that discards all metric data. -// It's used when both OTLP and Prometheus metrics are disabled. -type NoOpMeterStrategy struct{} - -// CreateMeterProvider creates a no-op meter provider -func (*NoOpMeterStrategy) CreateMeterProvider( - _ context.Context, - _ Config, - _ *resource.Resource, -) (*MeterResult, error) { - slog.Debug("creating no-op meter provider") - return &MeterResult{ - MeterProvider: noop.NewMeterProvider(), - PrometheusHandler: nil, - ShutdownFunc: nil, - }, nil -} - -// UnifiedMeterStrategy creates a meter provider with multiple readers (OTLP and/or Prometheus). -// It can combine OTLP metrics export and Prometheus scraping in a single provider. -type UnifiedMeterStrategy struct { - EnableOTLP bool // EnableOTLP controls whether to add an OTLP metrics reader - EnablePrometheus bool // EnablePrometheus controls whether to add a Prometheus reader -} - -// CreateMeterProvider creates a unified meter provider with OTLP and/or Prometheus readers -func (s *UnifiedMeterStrategy) CreateMeterProvider( - ctx context.Context, - config Config, - res *resource.Resource, -) (*MeterResult, error) { - var readers []sdkmetric.Reader - var prometheusHandler http.Handler - - // Add OTLP reader if enabled - if s.EnableOTLP { - //nolint:gosec // G706: OTLP endpoint from config - slog.Debug("adding OTLP metrics reader", - "endpoint", config.OTLPEndpoint) - - otlpConfig := otlp.Config{ - Endpoint: config.OTLPEndpoint, - Headers: config.Headers, - Insecure: config.Insecure, - SamplingRate: config.SamplingRate, - CACertPath: config.CACertPath, - } - - reader, err := otlp.NewMetricReader(ctx, otlpConfig) - if err != nil { - return nil, fmt.Errorf("failed to create OTLP metric reader for endpoint %s: %w", config.OTLPEndpoint, err) - } - readers = append(readers, reader) - } - - // Add Prometheus reader if enabled - if s.EnablePrometheus { - slog.Debug("adding Prometheus metrics reader") - promConfig := prometheus.Config{ - EnableMetricsPath: true, - IncludeRuntimeMetrics: true, - } - reader, handler, err := prometheus.NewReader(promConfig) - if err != nil { - return nil, fmt.Errorf("failed to create Prometheus metric reader: %w", err) - } - readers = append(readers, reader) - prometheusHandler = handler - } - - // Create meter provider with all readers - if len(readers) == 0 { - return &MeterResult{ - MeterProvider: noop.NewMeterProvider(), - PrometheusHandler: nil, - ShutdownFunc: nil, - }, nil - } - - opts := []sdkmetric.Option{sdkmetric.WithResource(res)} - for _, reader := range readers { - opts = append(opts, sdkmetric.WithReader(reader)) - } - - provider := sdkmetric.NewMeterProvider(opts...) - return &MeterResult{ - MeterProvider: provider, - PrometheusHandler: prometheusHandler, - ShutdownFunc: provider.Shutdown, - }, nil -} - -// StrategySelector determines which strategies to use based on configuration. -// It analyzes the configuration to select appropriate tracer and meter strategies. -type StrategySelector struct { - config Config // config holds the telemetry configuration to analyze -} - -// NewStrategySelector creates a new strategy selector with the given configuration. -// The selector will analyze the config to determine appropriate strategies. -func NewStrategySelector(config Config) *StrategySelector { - return &StrategySelector{config: config} -} - -// SelectTracerStrategy determines the appropriate tracer strategy based on configuration. -func (s *StrategySelector) SelectTracerStrategy() TracerStrategy { - hasEndpoint := s.config.OTLPEndpoint != "" - tracingEnabled := s.config.TracingEnabled - - if hasEndpoint && tracingEnabled { - return &OTLPTracerStrategy{} - } - - // Log informational message when endpoint is configured but tracing is disabled - if hasEndpoint && !tracingEnabled { - slog.Debug("otlp endpoint configured but tracing is disabled") - } - - // Extra processors (e.g. Sentry bridge) need a real SDK tracer provider so spans - // reach the processors. OTLPTracerStrategy handles the no-endpoint case when - // extraProcessors are present — it skips the OTLP exporter but still registers them. - if s.hasExtraProcessors() { - return &OTLPTracerStrategy{} - } - - return &NoOpTracerStrategy{} -} - -// SelectMeterStrategy determines the appropriate meter strategy based on configuration. -func (s *StrategySelector) SelectMeterStrategy() MeterStrategy { - wantsOTLPMetrics := s.hasOTLPMetrics() - wantsPrometheus := s.config.EnablePrometheusMetricsPath - - // Return no-op if no metrics are enabled - if !wantsOTLPMetrics && !wantsPrometheus { - return &NoOpMeterStrategy{} - } - - // Return unified strategy with appropriate readers enabled - return &UnifiedMeterStrategy{ - EnableOTLP: wantsOTLPMetrics, - EnablePrometheus: wantsPrometheus, - } -} - -// IsFullyNoOp returns true if both tracer and meter would be no-op. -func (s *StrategySelector) IsFullyNoOp() bool { - return !s.hasOTLPMetrics() && !s.hasOTLPTracing() && !s.hasPrometheus() && !s.hasExtraProcessors() -} - -// hasOTLPMetrics returns true if OTLP metrics are wanted. -func (s *StrategySelector) hasOTLPMetrics() bool { - return s.config.OTLPEndpoint != "" && s.config.MetricsEnabled -} - -// hasOTLPTracing returns true if OTLP tracing is wanted. -func (s *StrategySelector) hasOTLPTracing() bool { - return s.config.OTLPEndpoint != "" && s.config.TracingEnabled -} - -// hasPrometheus returns true if Prometheus metrics are wanted. -func (s *StrategySelector) hasPrometheus() bool { - return s.config.EnablePrometheusMetricsPath -} - -// hasExtraProcessors returns true if extra span processors (e.g. a Sentry bridge) are registered. -func (s *StrategySelector) hasExtraProcessors() bool { - return len(s.config.ExtraSpanProcessors) > 0 -} diff --git a/pkg/telemetry/providers/providers_strategy_test.go b/pkg/telemetry/providers/providers_strategy_test.go deleted file mode 100644 index 2b2d63562a..0000000000 --- a/pkg/telemetry/providers/providers_strategy_test.go +++ /dev/null @@ -1,600 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package providers - -import ( - "context" - "fmt" - "strings" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "go.opentelemetry.io/otel/metric/noop" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" - "go.opentelemetry.io/otel/sdk/resource" - sdktrace "go.opentelemetry.io/otel/sdk/trace" - semconv "go.opentelemetry.io/otel/semconv/v1.21.0" -) - -// noopSpanProcessor is a minimal SpanProcessor for testing. -type noopSpanProcessor struct{} - -func (noopSpanProcessor) OnStart(_ context.Context, _ sdktrace.ReadWriteSpan) {} -func (noopSpanProcessor) OnEnd(_ sdktrace.ReadOnlySpan) {} -func (noopSpanProcessor) Shutdown(_ context.Context) error { return nil } -func (noopSpanProcessor) ForceFlush(_ context.Context) error { return nil } - -func TestStrategySelector_SelectTracerStrategy(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - expectedType string - }{ - { - name: "OTLP tracer when endpoint and tracing enabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - TracingEnabled: true, - }, - expectedType: "*providers.OTLPTracerStrategy", - }, - { - name: "NoOp tracer when endpoint but tracing disabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - TracingEnabled: false, - }, - expectedType: "*providers.NoOpTracerStrategy", - }, - { - name: "NoOp tracer when no endpoint", - config: Config{ - TracingEnabled: true, - }, - expectedType: "*providers.NoOpTracerStrategy", - }, - { - name: "OTLP tracer when extra processors present without endpoint", - config: Config{ - ExtraSpanProcessors: []sdktrace.SpanProcessor{noopSpanProcessor{}}, - }, - expectedType: "*providers.OTLPTracerStrategy", - }, - { - name: "OTLP tracer when extra processors present with endpoint", - config: Config{ - OTLPEndpoint: "localhost:4318", - TracingEnabled: true, - ExtraSpanProcessors: []sdktrace.SpanProcessor{noopSpanProcessor{}}, - }, - expectedType: "*providers.OTLPTracerStrategy", - }, - } - - for _, tt := range tests { - tt := tt // capture range variable - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - selector := NewStrategySelector(tt.config) - strategy := selector.SelectTracerStrategy() - - assert.NotNil(t, strategy) - assert.Equal(t, tt.expectedType, getTypeName(strategy)) - }) - } -} - -func TestNoOpTracerStrategy_CreateTracerProvider(t *testing.T) { - t.Parallel() - - ctx := context.Background() - res := createTestResource(t) - config := Config{} - - strategy := &NoOpTracerStrategy{} - provider, shutdown, err := strategy.CreateTracerProvider(ctx, config, res) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.Nil(t, shutdown, "Expected no shutdown function for no-op tracer") - - // Verify it's actually a no-op provider - typeName := getTypeName(provider) - assert.Contains(t, typeName, "noop", "Expected no-op tracer provider, got %s", typeName) -} - -func TestOTLPTracerStrategy_CreateTracerProvider(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - expectErr bool - }{ - { - name: "Valid OTLP config", - config: Config{ - OTLPEndpoint: "localhost:4318", - Insecure: true, - SamplingRate: 0.1, - }, - expectErr: false, - }, - { - name: "Valid OTLP config with headers", - config: Config{ - OTLPEndpoint: "localhost:4318", - Insecure: true, - SamplingRate: 1.0, - Headers: map[string]string{"Authorization": "Bearer token"}, - }, - expectErr: false, - }, - { - name: "Valid secure OTLP config", - config: Config{ - OTLPEndpoint: "https://api.example.com:4318", - Insecure: false, - SamplingRate: 0.5, - }, - expectErr: false, - }, - } - - for _, tt := range tests { - tt := tt // capture range variable - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := context.Background() - res := createTestResource(t) - strategy := &OTLPTracerStrategy{} - - provider, shutdown, err := strategy.CreateTracerProvider(ctx, tt.config, res) - - if tt.expectErr { - assert.Error(t, err) - assert.Nil(t, provider) - assert.Nil(t, shutdown) - } else { - require.NoError(t, err) - require.NotNil(t, provider) - require.NotNil(t, shutdown, "Expected shutdown function for OTLP tracer") - - // Verify it's not a no-op provider - typeName := getTypeName(provider) - assert.NotContains(t, typeName, "noop", "Expected non-noop tracer provider, got %s", typeName) - - // Clean up - if shutdown != nil { - err := shutdown(ctx) - assert.NoError(t, err, "Shutdown should not error") - } - } - }) - } -} - -func TestOTLPTracerStrategy_ExtraProcessorsWithoutEndpoint(t *testing.T) { - t.Parallel() - - ctx := context.Background() - res := createTestResource(t) - config := Config{ - SamplingRate: 1.0, - ExtraSpanProcessors: []sdktrace.SpanProcessor{noopSpanProcessor{}}, - } - - strategy := &OTLPTracerStrategy{} - provider, shutdown, err := strategy.CreateTracerProvider(ctx, config, res) - - require.NoError(t, err) - require.NotNil(t, provider) - require.NotNil(t, shutdown, "Expected shutdown function even without OTLP endpoint when processors are registered") - - // Must be a real SDK provider, not a no-op - typeName := getTypeName(provider) - assert.NotContains(t, typeName, "noop", "Expected real SDK tracer provider, got %s", typeName) - - if shutdown != nil { - assert.NoError(t, shutdown(ctx)) - } -} - -func TestStrategySelector_SelectMeterStrategy(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - expectedType string - }{ - { - name: "Unified meter when OTLP metrics enabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - MetricsEnabled: true, - }, - expectedType: "*providers.UnifiedMeterStrategy", - }, - { - name: "Unified meter when Prometheus enabled", - config: Config{ - EnablePrometheusMetricsPath: true, - }, - expectedType: "*providers.UnifiedMeterStrategy", - }, - { - name: "Unified meter when both OTLP and Prometheus", - config: Config{ - OTLPEndpoint: "localhost:4318", - MetricsEnabled: true, - EnablePrometheusMetricsPath: true, - }, - expectedType: "*providers.UnifiedMeterStrategy", - }, - { - name: "NoOp meter when nothing enabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - MetricsEnabled: false, - }, - expectedType: "*providers.NoOpMeterStrategy", - }, - { - name: "NoOp meter when empty config", - config: Config{}, - expectedType: "*providers.NoOpMeterStrategy", - }, - } - - for _, tt := range tests { - tt := tt // capture range variable - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - selector := NewStrategySelector(tt.config) - strategy := selector.SelectMeterStrategy() - - assert.NotNil(t, strategy) - assert.Equal(t, tt.expectedType, getTypeName(strategy)) - }) - } -} - -func TestStrategySelector_IsFullyNoOp(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config Config - expected bool - }{ - { - name: "fully no-op when nothing configured", - config: Config{}, - expected: true, - }, - { - name: "not no-op when OTLP tracing enabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - TracingEnabled: true, - }, - expected: false, - }, - { - name: "not no-op when OTLP metrics enabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - MetricsEnabled: true, - }, - expected: false, - }, - { - name: "not no-op when Prometheus enabled", - config: Config{ - EnablePrometheusMetricsPath: true, - }, - expected: false, - }, - { - name: "fully no-op when endpoint but nothing enabled", - config: Config{ - OTLPEndpoint: "localhost:4318", - TracingEnabled: false, - MetricsEnabled: false, - }, - expected: true, - }, - { - name: "not no-op when extra processors present without endpoint", - config: Config{ - ExtraSpanProcessors: []sdktrace.SpanProcessor{noopSpanProcessor{}}, - }, - expected: false, - }, - } - - for _, tt := range tests { - tt := tt // capture range variable - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - selector := NewStrategySelector(tt.config) - result := selector.IsFullyNoOp() - - assert.Equal(t, tt.expected, result) - }) - } -} - -func TestNoOpMeterStrategy_CreateMeterProvider(t *testing.T) { - t.Parallel() - - ctx := context.Background() - res := createTestResource(t) - config := Config{} - - strategy := &NoOpMeterStrategy{} - provider, err := strategy.CreateMeterProvider(ctx, config, res) - - require.NoError(t, err) - require.NotNil(t, provider) - - // Verify it's actually a no-op provider - assert.Nil(t, provider.PrometheusHandler) - assert.Nil(t, provider.ShutdownFunc) - typeName := getTypeName(provider.MeterProvider) - assert.Contains(t, typeName, "noop", "Expected no-op meter provider, got %s", typeName) -} - -func TestUnifiedMeterStrategy_Configurations(t *testing.T) { - t.Parallel() - - ctx := context.Background() - res := createTestResource(t) - - tests := []struct { - name string - strategy *UnifiedMeterStrategy - config Config - expectPrometheus bool - expectOTLP bool - expectNoOp bool - }{ - { - name: "OTLP only", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: true, - EnablePrometheus: false, - }, - config: Config{ - OTLPEndpoint: "localhost:4318", - Insecure: true, - }, - expectPrometheus: false, - expectOTLP: true, - expectNoOp: false, - }, - { - name: "Prometheus only", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: false, - EnablePrometheus: true, - }, - config: Config{}, - expectPrometheus: true, - expectOTLP: false, - expectNoOp: false, - }, - { - name: "Both OTLP and Prometheus", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: true, - EnablePrometheus: true, - }, - config: Config{ - OTLPEndpoint: "localhost:4318", - Insecure: true, - }, - expectPrometheus: true, - expectOTLP: true, - expectNoOp: false, - }, - { - name: "Neither enabled", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: false, - EnablePrometheus: false, - }, - config: Config{}, - expectPrometheus: false, - expectOTLP: false, - expectNoOp: true, - }, - } - - for _, tt := range tests { - tt := tt // capture range variable - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - result, err := tt.strategy.CreateMeterProvider(ctx, tt.config, res) - - require.NoError(t, err) - require.NotNil(t, result) - require.NotNil(t, result.MeterProvider) - - if tt.expectPrometheus { - assert.NotNil(t, result.PrometheusHandler, "Expected Prometheus handler") - } else { - assert.Nil(t, result.PrometheusHandler, "Expected no Prometheus handler") - } - - // Note: OTLP handler is not exposed in MeterResult, only Prometheus handler - // OTLP functionality is verified through the meter provider type check below - - if tt.expectNoOp { - assert.Contains(t, getTypeName(result.MeterProvider), "noop") - assert.Nil(t, result.ShutdownFunc) - // Verify it's actually a noop provider - need to import noop package - // The noop.MeterProvider is actually noop.meterProvider (unexported) - // so we can't do a direct type assertion. Check the type name instead. - typeName := getTypeName(result.MeterProvider) - assert.Contains(t, typeName, "noop", "Expected noop meter provider, got %s", typeName) - } else { - assert.NotContains(t, getTypeName(result.MeterProvider), "noop") - // Verify it's actually an SDK provider (not noop) - _, isSDKProvider := result.MeterProvider.(*sdkmetric.MeterProvider) - assert.True(t, isSDKProvider, "Expected SDK meter provider, got %T", result.MeterProvider) - // Shutdown may or may not be nil depending on implementation - } - }) - } -} - -// TestUnifiedMeterStrategyConfiguration tests the unified meter strategy configuration -func TestUnifiedMeterStrategyConfiguration(t *testing.T) { - t.Parallel() - ctx := context.Background() - - tests := []struct { - name string - strategy *UnifiedMeterStrategy - config Config - expectError bool - description string - }{ - { - name: "OTLP only", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: true, - EnablePrometheus: false, - }, - config: Config{ - ServiceName: "test", - OTLPEndpoint: "localhost:4318", - Insecure: true, - }, - expectError: false, - description: "Should create meter provider with OTLP reader only", - }, - { - name: "Prometheus only", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: false, - EnablePrometheus: true, - }, - config: Config{ - ServiceName: "test", - }, - expectError: false, - description: "Should create meter provider with Prometheus reader only", - }, - { - name: "Both OTLP and Prometheus", - strategy: &UnifiedMeterStrategy{ - EnableOTLP: true, - EnablePrometheus: true, - }, - config: Config{ - ServiceName: "test", - OTLPEndpoint: "localhost:4318", - Insecure: true, - }, - expectError: false, - description: "Should create meter provider with both readers", - }, - } - - for _, tt := range tests { - tt := tt // capture range variable - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - // Create resource for testing - res := createTestResource(t) - - // Test meter provider creation - result, err := tt.strategy.CreateMeterProvider(ctx, tt.config, res) - - if tt.expectError { - require.Error(t, err) - return - } - - require.NoError(t, err, tt.description) - require.NotNil(t, result) - require.NotNil(t, result.MeterProvider) - - // Validate provider type - if tt.strategy.EnableOTLP || tt.strategy.EnablePrometheus { - // Should be SDK meter provider when any reader is enabled - assert.IsType(t, &sdkmetric.MeterProvider{}, result.MeterProvider, - "Should create SDK meter provider when readers are configured") - } else { - // Should be no-op when nothing is enabled - assert.IsType(t, noop.MeterProvider{}, result.MeterProvider, - "Should create no-op meter provider when no readers are configured") - } - - // Validate Prometheus handler - if tt.strategy.EnablePrometheus { - assert.NotNil(t, result.PrometheusHandler, - "Should have Prometheus handler when Prometheus is enabled") - } else { - assert.Nil(t, result.PrometheusHandler, - "Should not have Prometheus handler when Prometheus is disabled") - } - - // Test shutdown if available - ignore connection errors during test shutdown - if result.ShutdownFunc != nil { - err := result.ShutdownFunc(ctx) - if err != nil && !isConnectionError(err) { - t.Errorf("Shutdown error (non-connection): %v", err) - } - } - }) - } -} - -// isConnectionError checks if the error is a connection-related error that can be ignored in tests -func isConnectionError(err error) bool { - errorStr := err.Error() - return strings.Contains(errorStr, "connection refused") || - strings.Contains(errorStr, "dial tcp") || - strings.Contains(errorStr, "no such host") -} - -// Helper functions - -func getTypeName(v interface{}) string { - if v == nil { - return "nil" - } - return fmt.Sprintf("%T", v) -} - -func createTestResource(t *testing.T) *resource.Resource { - t.Helper() - return createTestResourceWithName(t, "test-service", "1.0.0") -} - -func createTestResourceWithName(t *testing.T, serviceName, serviceVersion string) *resource.Resource { - t.Helper() - res, err := resource.New( - context.Background(), - resource.WithAttributes( - semconv.ServiceName(serviceName), - semconv.ServiceVersion(serviceVersion), - ), - ) - require.NoError(t, err) - return res -} diff --git a/pkg/telemetry/providers/providers_test.go b/pkg/telemetry/providers/providers_test.go deleted file mode 100644 index bec1c562fa..0000000000 --- a/pkg/telemetry/providers/providers_test.go +++ /dev/null @@ -1,600 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package providers - -import ( - "context" - "errors" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "go.opentelemetry.io/otel/metric/noop" - "go.opentelemetry.io/otel/sdk/resource" - "go.opentelemetry.io/otel/trace" - tracenoop "go.opentelemetry.io/otel/trace/noop" -) - -// func TestAssembler_CreateResource(t *testing.T) { -// t.Parallel() - -// ctx := context.Background() -// options := []ProviderOptions{ -// WithServiceName("test-service"), -// WithServiceVersion("1.0.0"), -// WithOTLPEndpoint("localhost:4318"), -// WithInsecure(true), -// WithSamplingRate(0.1), -// WithTracingEnabled(false), -// WithMetricsEnabled(false), -// WithEnablePrometheusMetricsPath(false), -// } - -// options = append(options, WithServiceVersion("1.2.3")) - -// provider, err := NewCompositeProvider(ctx, options...) -// require.NoError(t, err) -// require.NotNil(t, provider) - -// // Check attributes -// attrs := assembler.resource.Attributes() -// hasName := false -// hasVersion := false -// for _, attr := range attrs { -// if attr.Key == "service.name" && attr.Value.AsString() == "test-service" { -// hasName = true -// } -// if attr.Key == "service.version" && attr.Value.AsString() == "1.2.3" { -// hasVersion = true -// } -// } -// assert.True(t, hasName, "service.name attribute should be present") -// assert.True(t, hasVersion, "service.version attribute should be present") -// } - -func TestAssembler_CreateNoOpProvider(t *testing.T) { - t.Parallel() - - provider := createNoOpProvider() - - require.NotNil(t, provider) - assert.NotNil(t, provider.tracerProvider) - assert.NotNil(t, provider.meterProvider) - assert.Nil(t, provider.prometheusHandler) - assert.Empty(t, provider.shutdownFuncs) - - // Verify the providers are actually no-op implementations - // Check if it's a no-op tracer provider interface - assert.IsType(t, tracenoop.NewTracerProvider(), provider.tracerProvider, "tracer provider should be no-op") - - // Check if it's a no-op meter provider interface - assert.IsType(t, noop.NewMeterProvider(), provider.meterProvider, "meter provider should be no-op") -} - -func TestAssembler_Assemble_NoOpCase(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithTracingEnabled(false), - WithMetricsEnabled(false), - WithEnablePrometheusMetricsPath(false), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - - assert.IsType(t, tracenoop.NewTracerProvider(), provider.TracerProvider(), "tracer provider should be no-op") - assert.IsType(t, noop.NewMeterProvider(), provider.MeterProvider(), "meter provider should be no-op") - assert.Nil(t, provider.PrometheusHandler()) -} - -func TestAssembler_Assemble_WithOTLPTracing(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithTracingEnabled(true), - WithMetricsEnabled(false), - WithSamplingRate(0.5), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.TracerProvider()) - assert.IsType(t, noop.NewMeterProvider(), provider.MeterProvider(), "meter provider should be no-op when metrics disabled") - assert.Len(t, provider.shutdownFuncs, 1) // Should have one shutdown function for tracing -} - -func TestAssembler_Assemble_WithPrometheus(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithTracingEnabled(false), - WithMetricsEnabled(false), - WithEnablePrometheusMetricsPath(true), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.MeterProvider()) - assert.NotNil(t, provider.PrometheusHandler()) - assert.NotEmpty(t, provider.shutdownFuncs) // Should have shutdown function for Prometheus -} - -func TestAssembler_Assemble_WithOTLPMetrics(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithTracingEnabled(false), - WithMetricsEnabled(true), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.MeterProvider()) - assert.NotEmpty(t, provider.shutdownFuncs) // Should have shutdown function for OTLP metrics -} - -func TestAssembler_Assemble_WithEverything(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithTracingEnabled(true), - WithMetricsEnabled(true), - WithEnablePrometheusMetricsPath(true), - WithSamplingRate(1.0), - WithHeaders(map[string]string{"test": "header"}), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.TracerProvider()) - assert.NotNil(t, provider.MeterProvider()) - assert.NotNil(t, provider.PrometheusHandler()) - assert.NotEmpty(t, provider.shutdownFuncs) // Should have multiple shutdown functions -} - -func TestCompositeProvider_Accessors(t *testing.T) { - t.Parallel() - - // Create a composite provider manually - tp := tracenoop.NewTracerProvider() - mp := noop.NewMeterProvider() - - provider := &CompositeProvider{ - tracerProvider: tp, - meterProvider: mp, - prometheusHandler: nil, - } - - assert.Equal(t, tp, provider.TracerProvider()) - assert.Equal(t, mp, provider.MeterProvider()) - assert.Nil(t, provider.PrometheusHandler()) -} - -func TestCompositeProvider_Shutdown(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - shutdownFuncs []func(context.Context) error - expectError bool - errorCount int - }{ - { - name: "no shutdown functions", - shutdownFuncs: []func(context.Context) error{}, - expectError: false, - }, - { - name: "single successful shutdown", - shutdownFuncs: []func(context.Context) error{ - func(_ context.Context) error { return nil }, - }, - expectError: false, - }, - { - name: "multiple successful shutdowns", - shutdownFuncs: []func(context.Context) error{ - func(_ context.Context) error { return nil }, - func(_ context.Context) error { return nil }, - func(_ context.Context) error { return nil }, - }, - expectError: false, - }, - { - name: "single failed shutdown", - shutdownFuncs: []func(context.Context) error{ - func(_ context.Context) error { return errors.New("shutdown failed") }, - }, - expectError: true, - errorCount: 1, - }, - { - name: "mixed success and failure", - shutdownFuncs: []func(context.Context) error{ - func(_ context.Context) error { return nil }, - func(_ context.Context) error { return errors.New("provider 1 failed") }, - func(_ context.Context) error { return nil }, - func(_ context.Context) error { return errors.New("provider 3 failed") }, - }, - expectError: true, - errorCount: 2, - }, - { - name: "all failures", - shutdownFuncs: []func(context.Context) error{ - func(_ context.Context) error { return errors.New("error 1") }, - func(_ context.Context) error { return errors.New("error 2") }, - func(_ context.Context) error { return errors.New("error 3") }, - }, - expectError: true, - errorCount: 3, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - provider := &CompositeProvider{ - shutdownFuncs: tt.shutdownFuncs, - } - - ctx := context.Background() - err := provider.Shutdown(ctx) - - if tt.expectError { - require.Error(t, err) - assert.Contains(t, err.Error(), "shutdown failed") - if tt.errorCount > 0 { - assert.Contains(t, err.Error(), "errors:") - } - } else { - require.NoError(t, err) - } - }) - } -} - -// TestAssembler_Shutdown_WithErrors tests shutdown with failing shutdown functions -func TestAssembler_Shutdown_WithErrors(t *testing.T) { - t.Parallel() - - // Create a composite provider with a shutdown function that times out - timeoutShutdown := func(ctx context.Context) error { - select { - case <-time.After(10 * time.Second): - return errors.New("shutdown timed out") - case <-ctx.Done(): - return ctx.Err() - } - } - - errorShutdown := func(_ context.Context) error { - return errors.New("shutdown error") - } - - successShutdown := func(_ context.Context) error { - return nil - } - - provider := &CompositeProvider{ - tracerProvider: tracenoop.NewTracerProvider(), - meterProvider: noop.NewMeterProvider(), - shutdownFuncs: []func(context.Context) error{ - successShutdown, - errorShutdown, - timeoutShutdown, - }, - } - - ctx := context.Background() - err := provider.Shutdown(ctx) - - require.Error(t, err) - assert.Contains(t, err.Error(), "shutdown failed") - assert.Contains(t, err.Error(), "errors:") -} - -func TestCompositeProvider_ShutdownTimeout(t *testing.T) { - t.Parallel() - - slowShutdown := func(ctx context.Context) error { - select { - case <-time.After(10 * time.Second): - return errors.New("shutdown completed too slowly") - case <-ctx.Done(): - return ctx.Err() - } - } - - provider := &CompositeProvider{ - shutdownFuncs: []func(context.Context) error{ - slowShutdown, - }, - } - - ctx := context.Background() - start := time.Now() - err := provider.Shutdown(ctx) - elapsed := time.Since(start) - - // Should timeout within the 5 second limit set in Shutdown - assert.Less(t, elapsed, 6*time.Second) - assert.Error(t, err) -} - -func TestCompositeProvider_MultipleShutdown(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithSamplingRate(0.1), - WithTracingEnabled(false), - WithMetricsEnabled(false), - WithEnablePrometheusMetricsPath(true), - } - - provider, err := NewCompositeProvider(ctx, options...) - require.NoError(t, err) - require.NotNil(t, provider) - - // Multiple shutdowns should not panic - _ = provider.Shutdown(ctx) - _ = provider.Shutdown(ctx) -} - -func TestAssembler_Assemble_WithHeaders(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithHeaders(map[string]string{ - "Authorization": "Bearer token", - "X-Custom": "value", - }), - WithInsecure(true), - WithTracingEnabled(true), - WithMetricsEnabled(true), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.TracerProvider()) - assert.NotNil(t, provider.MeterProvider()) -} - -func TestAssembler_Assemble_DifferentSamplingRates(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - samplingRate float64 - }{ - {"zero sampling", 0.0}, - {"quarter sampling", 0.25}, - {"half sampling", 0.5}, - {"three quarter sampling", 0.75}, - {"full sampling", 1.0}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithTracingEnabled(true), - WithSamplingRate(tt.samplingRate), - } - - provider, err := NewCompositeProvider(ctx, options...) - - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.TracerProvider()) - }) - } -} - -func TestAssembler_Assemble_EdgeCases(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - options []ProviderOption - }{ - { - name: "empty service name and version", - options: []ProviderOption{}, - }, - { - name: "only service name", - options: []ProviderOption{ - WithServiceName("my-service"), - }, - }, - { - name: "only service version", - options: []ProviderOption{ - WithServiceVersion("v1.0.0"), - }, - }, - { - name: "very long service name", - options: []ProviderOption{ - WithServiceName("this-is-a-very-long-service-name-that-might-cause-issues-in-some-systems-but-should-still-work-correctly-when-creating-resources"), - WithServiceVersion("1.0.0"), - }, - }, - { - name: "special characters in service name", - options: []ProviderOption{ - WithServiceName("service-name_with.special@chars"), - WithServiceVersion("1.0.0-beta+assemble.123"), - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - ctx := context.Background() - provider, err := NewCompositeProvider(ctx, tt.options...) - - // All edge cases should still succeed - require.NoError(t, err) - require.NotNil(t, provider) - assert.NotNil(t, provider.TracerProvider()) - assert.NotNil(t, provider.MeterProvider()) - }) - } -} - -// func TestAssembler_AssembleProviders_DirectCall(t *testing.T) { -// t.Parallel() - -// ctx := context.Background() -// options := []ProviderOptions{ -// WithServiceName("test-service"), -// WithServiceVersion("1.0.0"), -// WithOTLPEndpoint("localhost:4318"), -// WithTracingEnabled(true), -// WithMetricsEnabled(true), -// } - -// assembler, err := NewCompositeProvider(ctx, options...) -// require.NoError(t, err) - -// // Create selector -// selector := NewStrategySelector(assembler.config) - -// // Call assembleProviders directly -// composite, err := assembleProviders(ctx, selector, assembler.resource) - -// require.NoError(t, err) -// require.NotNil(t, composite) -// assert.NotNil(t, composite.TracerProvider()) -// assert.NotNil(t, composite.MeterProvider()) -// assert.NotEmpty(t, composite.shutdownFuncs) -// } - -// TestMockErrorStrategy tests error handling using a custom strategy that always fails -type TestMockErrorMeterStrategy struct { - errMsg string -} - -func (s *TestMockErrorMeterStrategy) CreateMeterProvider(_ context.Context, _ Config, _ *resource.Resource) (*MeterResult, error) { - return nil, errors.New(s.errMsg) -} - -type TestMockErrorTracerStrategy struct { - errMsg string -} - -func (s *TestMockErrorTracerStrategy) CreateTracerProvider(_ context.Context, _ Config, _ *resource.Resource) ( - trace.TracerProvider, func(context.Context) error, error, -) { - return nil, nil, errors.New(s.errMsg) -} - -// TestErrorStrategies verifies that strategy errors are properly handled -func TestErrorStrategies(t *testing.T) { - t.Parallel() - - ctx := context.Background() - config := Config{ - ServiceName: "test-service", - ServiceVersion: "1.0.0", - } - - res := resource.Default() - - t.Run("meter strategy error", func(t *testing.T) { - t.Parallel() - strategy := &TestMockErrorMeterStrategy{errMsg: "meter creation failed"} - _, err := strategy.CreateMeterProvider(ctx, config, res) - require.Error(t, err) - assert.Contains(t, err.Error(), "meter creation failed") - }) - - t.Run("tracer strategy error", func(t *testing.T) { - t.Parallel() - strategy := &TestMockErrorTracerStrategy{errMsg: "tracer creation failed"} - _, _, err := strategy.CreateTracerProvider(ctx, config, res) - require.Error(t, err) - assert.Contains(t, err.Error(), "tracer creation failed") - }) -} - -// // TestAssembler_ResourceCreationError tests handling of resource creation errors -// func TestAssembler_ResourceCreationError(t *testing.T) { -// t.Parallel() - -// // Create a assembler with invalid configuration that might cause resource creation issues -// options := []ProviderOptions{ -// ServiceName: string([]byte{0xFF, 0xFE, 0xFD}), // Invalid UTF-8 characters -// ServiceVersion: string([]byte{0xFF, 0xFE, 0xFD}), -// } - -// ctx := context.Background() -// assembler, err := NewCompositeProvider(ctx, options...) -// require.NoError(t, err) - -// // Even with invalid characters, resource creation typically succeeds -// // as OpenTelemetry handles them gracefully -// err := assembler.createResource(ctx) -// // This won't actually error, but we're testing the path -// assert.NoError(t, err) -// if assembler.resource != nil { -// attrs := assembler.resource.Attributes() -// assert.NotNil(t, attrs) -// } -// } diff --git a/pkg/telemetry/providers/unified_test.go b/pkg/telemetry/providers/unified_test.go deleted file mode 100644 index 0d577d1650..0000000000 --- a/pkg/telemetry/providers/unified_test.go +++ /dev/null @@ -1,135 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package providers - -import ( - "context" - "net/http/httptest" - "strings" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestUnifiedMeterProvider_BothProviders(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithSamplingRate(0.1), - WithTracingEnabled(true), - WithMetricsEnabled(true), - WithEnablePrometheusMetricsPath(true), - } - - provider, err := NewCompositeProvider(ctx, options...) - require.NoError(t, err) - require.NotNil(t, provider) - defer provider.Shutdown(ctx) - - // Create a test metric using the meter provider - meter := provider.MeterProvider().Meter("test") - counter, err := meter.Int64Counter("test_unified_counter") - require.NoError(t, err) - - // Record some values - counter.Add(ctx, 1) - counter.Add(ctx, 2) - counter.Add(ctx, 3) - - // Check Prometheus endpoint - require.NotNil(t, provider.PrometheusHandler()) - req := httptest.NewRequest("GET", "/metrics", nil) - rec := httptest.NewRecorder() - provider.PrometheusHandler().ServeHTTP(rec, req) - - assert.Equal(t, 200, rec.Code) - body := rec.Body.String() - - // Verify the metric appears in Prometheus output - assert.True(t, strings.Contains(body, "test_unified_counter"), - "Prometheus should contain our test metric") - - // The total should be 6 (1+2+3) - assert.True(t, strings.Contains(body, "test_unified_counter_total"), - "Prometheus should show the counter with _total suffix") -} - -func TestUnifiedMeterProvider_PrometheusOnly(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithSamplingRate(0.1), - WithTracingEnabled(false), - WithMetricsEnabled(false), - WithEnablePrometheusMetricsPath(true), - } - - provider, err := NewCompositeProvider(ctx, options...) - require.NoError(t, err) - require.NotNil(t, provider) - defer provider.Shutdown(ctx) - - // Create a test metric - meter := provider.MeterProvider().Meter("test") - histogram, err := meter.Float64Histogram("test_histogram") - require.NoError(t, err) - - // Record some values - histogram.Record(ctx, 10.5) - histogram.Record(ctx, 20.3) - - // Check Prometheus endpoint - require.NotNil(t, provider.PrometheusHandler()) - req := httptest.NewRequest("GET", "/metrics", nil) - rec := httptest.NewRecorder() - provider.PrometheusHandler().ServeHTTP(rec, req) - - assert.Equal(t, 200, rec.Code) - assert.Contains(t, rec.Body.String(), "test_histogram") -} - -func TestUnifiedMeterProvider_OTLPOnly(t *testing.T) { - t.Parallel() - - ctx := context.Background() - options := []ProviderOption{ - WithServiceName("test-service"), - WithServiceVersion("1.0.0"), - WithOTLPEndpoint("localhost:4318"), - WithInsecure(true), - WithSamplingRate(0.1), - WithTracingEnabled(true), - WithMetricsEnabled(true), - WithEnablePrometheusMetricsPath(false), - } - - provider, err := NewCompositeProvider(ctx, options...) - require.NoError(t, err) - require.NotNil(t, provider) - defer provider.Shutdown(ctx) - - // Should have meter provider but no Prometheus handler - assert.NotNil(t, provider.MeterProvider()) - assert.Nil(t, provider.PrometheusHandler()) - - // Should still be able to create metrics (they go to OTLP) - meter := provider.MeterProvider().Meter("test") - gauge, err := meter.Int64UpDownCounter("test_gauge") - require.NoError(t, err) - - // Record values - gauge.Add(ctx, 100) - gauge.Add(ctx, -50) -} diff --git a/pkg/transport/bridge.go b/pkg/transport/bridge.go index 796cd563c8..6b439c72a7 100644 --- a/pkg/transport/bridge.go +++ b/pkg/transport/bridge.go @@ -11,11 +11,10 @@ import ( "strings" "sync" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/pkg/versions" ) @@ -105,6 +104,19 @@ func (b *StdioBridge) run(ctx context.Context) { params = map[string]any{} } + // On a *_list_changed notification, re-fetch the upstream capability set + // before forwarding the notification. forwardAll is additive only: + // AddTool/AddResource/AddPrompt upsert by name/URI, so re-running it adds + // or updates capabilities but does NOT prune ones the upstream removed — + // a stale capability stays advertised locally after an upstream removal. + // (SetTools exists if pruning is needed later.) + switch n.Method { + case "notifications/tools/list_changed", + "notifications/resources/list_changed", + "notifications/prompts/list_changed": + b.forwardAll(context.Background()) + } + b.srv.SendNotificationToAllClients(n.Method, params) }) @@ -206,33 +218,27 @@ func (b *StdioBridge) forwardAll(ctx context.Context) { } } - // Resources -> return []mcp.ResourceContents + // Resources -> full result passthrough so upstream _meta reaches the client slog.Debug("forwarding resources from upstream to local stdio server") if lr, err := b.up.ListResources(ctx, mcp.ListResourcesRequest{}); err == nil { for _, res := range lr.Resources { resCopy := res - b.srv.AddResource(resCopy, func(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { - out, err := b.up.ReadResource(ctx, req) - if err != nil { - return nil, err - } - return out.Contents, nil - }) + b.srv.AddResourceWithResult(resCopy, + func(ctx context.Context, req mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return b.up.ReadResource(ctx, req) + }) } } - // Resource templates -> same return type as resources + // Resource templates -> same result-returning shape as resources slog.Debug("forwarding resource templates from upstream to local stdio server") if lt, err := b.up.ListResourceTemplates(ctx, mcp.ListResourceTemplatesRequest{}); err == nil { for _, tpl := range lt.ResourceTemplates { tplCopy := tpl - b.srv.AddResourceTemplate(tplCopy, func(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { - out, err := b.up.ReadResource(ctx, req) - if err != nil { - return nil, err - } - return out.Contents, nil - }) + b.srv.AddResourceTemplateWithResult(tplCopy, + func(ctx context.Context, req mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return b.up.ReadResource(ctx, req) + }) } } diff --git a/pkg/transport/bridge_test.go b/pkg/transport/bridge_test.go new file mode 100644 index 0000000000..c76d0f1292 --- /dev/null +++ b/pkg/transport/bridge_test.go @@ -0,0 +1,544 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transport + +import ( + "context" + "encoding/json" + "net/http/httptest" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/transport/types" +) + +// toolNamesOnServer returns the names of the tools currently registered on the +// given local MCP server, by issuing a synthetic tools/list request through the +// shim's HandleMessage dispatch path. The bridge's forwardAll registers upstream +// tools here, so this is the assertion seam for "what does the bridge advertise +// downstream". Call only after the bridge has fully stopped (run() returned), +// since StdioBridge exposes no synchronization for its srv field while running. +func toolNamesOnServer(t *testing.T, srv *server.MCPServer) []string { + t.Helper() + resp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","method":"tools/list","id":1}`)) + jr, ok := resp.(mcp.JSONRPCResponse) + require.True(t, ok, "tools/list should return a JSONRPCResponse, got %T", resp) + buf, err := json.Marshal(jr.Result) + require.NoError(t, err) + var lt mcp.ListToolsResult + require.NoError(t, json.Unmarshal(buf, <)) + names := make([]string, 0, len(lt.Tools)) + for _, tl := range lt.Tools { + names = append(names, tl.Name) + } + return names +} + +// resourceNamesOnServer returns the URIs of the resources currently registered +// on the given local MCP server, by issuing a synthetic resources/list request +// through the shim's HandleMessage dispatch path. Mirrors toolNamesOnServer; +// call only after the bridge has fully stopped (run() returned). +func resourceNamesOnServer(t *testing.T, srv *server.MCPServer) []string { + t.Helper() + resp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","method":"resources/list","id":1}`)) + jr, ok := resp.(mcp.JSONRPCResponse) + require.True(t, ok, "resources/list should return a JSONRPCResponse, got %T", resp) + buf, err := json.Marshal(jr.Result) + require.NoError(t, err) + var lr mcp.ListResourcesResult + require.NoError(t, json.Unmarshal(buf, &lr)) + uris := make([]string, 0, len(lr.Resources)) + for _, res := range lr.Resources { + uris = append(uris, res.URI) + } + return uris +} + +// promptNamesOnServer returns the names of the prompts currently registered on +// the given local MCP server, by issuing a synthetic prompts/list request +// through the shim's HandleMessage dispatch path. Mirrors toolNamesOnServer; +// call only after the bridge has fully stopped (run() returned). +func promptNamesOnServer(t *testing.T, srv *server.MCPServer) []string { + t.Helper() + resp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","method":"prompts/list","id":1}`)) + jr, ok := resp.(mcp.JSONRPCResponse) + require.True(t, ok, "prompts/list should return a JSONRPCResponse, got %T", resp) + buf, err := json.Marshal(jr.Result) + require.NoError(t, err) + var lp mcp.ListPromptsResult + require.NoError(t, json.Unmarshal(buf, &lp)) + names := make([]string, 0, len(lp.Prompts)) + for _, p := range lp.Prompts { + names = append(names, p.Name) + } + return names +} + +func containsTool(names []string, want string) bool { + for _, n := range names { + if n == want { + return true + } + } + return false +} + +// noopToolHandler is a stand-in tool handler; the bridge re-fetch tests never +// invoke tools, they only assert the advertised set. +func noopToolHandler(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return nil, nil +} + +// listToolsCounter is a server hook that atomically counts tools/list requests +// received by the backend. The bridge's forwardAll issues a tools/list upstream +// on startup and again whenever its OnNotification handler re-syncs on a +// *_list_changed notification, so the counter is the race-free readiness and +// re-fetch signal observed entirely on the backend side. +type listToolsCounter struct { + count atomic.Int64 +} + +func (c *listToolsCounter) hook() server.OnBeforeListToolsFunc { + return func(context.Context, any, *mcp.ListToolsRequest) { c.count.Add(1) } +} + +// sessionHolder captures the single client session that connects to the backend, +// guarded by a mutex because the hook fires on the backend's session goroutine. +type sessionHolder struct { + mu sync.Mutex + _session server.ClientSession +} + +func (h *sessionHolder) set(s server.ClientSession) { + h.mu.Lock() + defer h.mu.Unlock() + h._session = s +} + +func (h *sessionHolder) get() server.ClientSession { + h.mu.Lock() + defer h.mu.Unlock() + return h._session +} + +// TestBridge_ToolsListChanged_TriggersReSync verifies the bridge's +// notifications/tools/list_changed re-fetch fix: when the upstream backend's +// tool set changes after the bridge has connected, the upstream emits a +// tools/list_changed notification, the bridge's OnNotification handler re-runs +// forwardAll, and the newly added tool appears on the bridge's local stdio +// server. +// +// The backend is a real mcpcompat streamable-HTTP server. A tool-set change +// visible to an already-connected client is driven through the per-session +// overlay (SessionWithTools.SetSessionTools): the shim syncs the overlay onto the +// live go-sdk server, which emits notifications/tools/list_changed over the +// standalone SSE stream to the bridge's upstream client (the bridge connects +// with WithContinuousListening, so the SSE stream is established). The bridge's +// OnNotification closure is the fix under test. +// +// ServeStdio blocks reading os.Stdin; this test swaps os.Stdin for a pipe so +// ServeStdio returns cleanly on stdin EOF at teardown. The bridge exposes no +// synchronization for its srv field, so readiness/re-fetch is observed via the +// backend's tools/list counter (race-free) and bridge.srv is read only after +// Shutdown returns (b.wg.Wait() provides the happens-before edge). +// +//nolint:paralleltest // Swaps process-global os.Stdin; cannot run in parallel. +func TestBridge_ToolsListChanged_TriggersReSync(t *testing.T) { + // This test swaps process-global os.Stdin, so it cannot run in parallel + // with any other test touching stdio. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // --- upstream backend with an initial tool "alpha" --- + counter := &listToolsCounter{} + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + hooks.AddBeforeListTools(counter.hook()) + + backend := server.NewMCPServer( + "backend", "1.0", + server.WithToolCapabilities(true), + server.WithHooks(hooks), + ) + backend.AddTool(mcp.NewTool("alpha"), noopToolHandler) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // --- bridge pointing at the backend (streamable-http) --- + bridge, err := NewStdioBridge("test", ts.URL, types.TransportTypeStreamableHTTP) + require.NoError(t, err) + + // Swap os.Stdin for a pipe so ServeStdio (which hardcodes os.Stdin) does not + // touch the real process stdin and can be unblocked at teardown by closing + // the write end. Leave os.Stdout real; the assertions don't read it. + origIn := os.Stdin + pipeR, pipeW, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { + os.Stdin = origIn + _ = pipeR.Close() + _ = pipeW.Close() + }) + os.Stdin = pipeR + + bridge.Start(ctx) + t.Cleanup(func() { + // Close the stdin write end so the go-sdk stdio reader hits EOF and + // ServeStdio returns, letting run() exit and the WaitGroup drain. Then + // Shutdown closes the upstream client and waits for run() to finish. + _ = pipeW.Close() + bridge.Shutdown() + }) + + // Step 1: wait for the bridge to connect and run its initial forwardAll, + // which issues the first upstream tools/list (counter >= 1) after the + // session is registered. Observed entirely on the backend side: race-free. + require.Eventually(t, func() bool { + return holder.get() != nil && counter.count.Load() >= 1 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not complete initial forwardAll (session=%v, listCalls=%d)", + holder.get() != nil, counter.count.Load()) + + // Step 2: mutate the upstream tool set — add "beta" via the per-session + // overlay. SetSessionTools syncs onto the live go-sdk server bound to this + // session, which emits notifications/tools/list_changed over the SSE stream + // to the bridge's upstream client. The bridge's OnNotification handler then + // re-runs forwardAll, issuing a SECOND upstream tools/list (counter >= 2). + backendSession := holder.get() + swt, ok := backendSession.(server.SessionWithTools) + require.True(t, ok, "backend session must implement SessionWithTools") + swt.SetSessionTools(map[string]server.ServerTool{ + "alpha": {Tool: mcp.NewTool("alpha"), Handler: noopToolHandler}, + "beta": {Tool: mcp.NewTool("beta"), Handler: noopToolHandler}, + }) + + // Step 3: wait for the re-fetch — the second tools/list proves the + // OnNotification handler fired and re-ran forwardAll (the fix). + require.Eventually(t, func() bool { + return counter.count.Load() >= 2 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not re-run forwardAll after tools/list_changed (listCalls=%d)", + counter.count.Load()) + + // Step 4: stop the bridge, then read bridge.srv. Shutdown calls + // b.wg.Wait(), which synchronizes with run()'s defer b.wg.Done() (which is + // happens-after every field write run() made), so bridge.srv is safe to read + // here without a data race. + _ = pipeW.Close() // unblock ServeStdio (stdin EOF) + bridge.Shutdown() + + names := toolNamesOnServer(t, bridge.srv) + assert.True(t, containsTool(names, "alpha"), + "alpha must be present after the re-fetch (forwardAll is additive), got %v", names) + assert.True(t, containsTool(names, "beta"), + "beta must be present after the re-sync, got %v", names) +} + +// noopResourceHandler is a stand-in resource handler; the bridge re-fetch +// tests never read resources, they only assert the advertised set. +func noopResourceHandler(_ context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + return nil, nil +} + +// noopPromptHandler is a stand-in prompt handler; the bridge re-fetch tests +// never fetch prompts, they only assert the advertised set. +func noopPromptHandler(_ context.Context, _ mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return nil, nil +} + +// TestBridge_ResourcesListChanged_TriggersReSync verifies the bridge's +// notifications/resources/list_changed re-fetch fix: when the upstream +// backend's resource set changes after the bridge has connected, the upstream +// emits a resources/list_changed notification, the bridge's OnNotification +// handler re-runs forwardAll, and the newly added resource appears on the +// bridge's local stdio server. +// +// Mirrors TestBridge_ToolsListChanged_TriggersReSync: the mutation is driven +// through the per-session overlay (SessionWithResources.SetSessionResources). +// forwardAll re-fetches tools, resources, resource templates, and prompts on +// ANY *_list_changed notification, so the existing tools/list counter +// (listToolsCounter, wired via hooks.AddBeforeListTools) remains the +// race-free readiness/re-fetch signal even though this test mutates +// resources, not tools. +// +//nolint:paralleltest // Swaps process-global os.Stdin; cannot run in parallel. +func TestBridge_ResourcesListChanged_TriggersReSync(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // --- upstream backend with an initial resource "res://alpha" --- + counter := &listToolsCounter{} + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + hooks.AddBeforeListTools(counter.hook()) + + backend := server.NewMCPServer( + "backend", "1.0", + // WithToolCapabilities is declared (though no tool is registered) so the + // backend unambiguously serves tools/list and fires the BeforeListTools + // hook — the counter is the readiness signal for the resource-triggered + // re-sync, since forwardAll re-fetches tools on ANY *_list_changed. + server.WithToolCapabilities(true), + server.WithResourceCapabilities(true, true), + server.WithHooks(hooks), + ) + backend.AddResource(mcp.Resource{URI: "res://alpha", Name: "alpha"}, noopResourceHandler) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // --- bridge pointing at the backend (streamable-http) --- + bridge, err := NewStdioBridge("test", ts.URL, types.TransportTypeStreamableHTTP) + require.NoError(t, err) + + // Swap os.Stdin ONLY for a pipe so ServeStdio does not touch the real + // process stdin and can be unblocked at teardown by closing the write end. + origIn := os.Stdin + pipeR, pipeW, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { + os.Stdin = origIn + _ = pipeR.Close() + _ = pipeW.Close() + }) + os.Stdin = pipeR + + bridge.Start(ctx) + t.Cleanup(func() { + _ = pipeW.Close() + bridge.Shutdown() + }) + + // Step 1: wait for the bridge to connect and run its initial forwardAll + // (counter >= 1 from the tools/list call forwardAll always issues). + require.Eventually(t, func() bool { + return holder.get() != nil && counter.count.Load() >= 1 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not complete initial forwardAll (session=%v, listCalls=%d)", + holder.get() != nil, counter.count.Load()) + + // Step 2: mutate the upstream resource set — add "res://beta" via the + // per-session overlay. SetSessionResources syncs onto the live go-sdk + // server bound to this session, which emits + // notifications/resources/list_changed to the bridge's upstream client. + backendSession := holder.get() + swr, ok := backendSession.(server.SessionWithResources) + require.True(t, ok, "backend session must implement SessionWithResources") + swr.SetSessionResources(map[string]server.ServerResource{ + "res://alpha": {Resource: mcp.Resource{URI: "res://alpha", Name: "alpha"}, Handler: noopResourceHandler}, + "res://beta": {Resource: mcp.Resource{URI: "res://beta", Name: "beta"}, Handler: noopResourceHandler}, + }) + + // Step 3: wait for the re-fetch. forwardAll re-fetches tools on every + // *_list_changed notification (not just tools/list_changed), so the + // second tools/list call is the race-free signal that the resync ran. + require.Eventually(t, func() bool { + return counter.count.Load() >= 2 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not re-run forwardAll after resources/list_changed (listCalls=%d)", + counter.count.Load()) + + // Step 4: stop the bridge, then read bridge.srv. The re-sync forwardAll runs + // on the upstream client's OnNotification goroutine, not run(); Shutdown's + // b.up.Close() joins that goroutine before returning, which is the + // happens-before edge that makes reading bridge.srv here race-free. + _ = pipeW.Close() + bridge.Shutdown() + + uris := resourceNamesOnServer(t, bridge.srv) + assert.Contains(t, uris, "res://alpha", "res://alpha must be present after the re-fetch (forwardAll is additive)") + assert.Contains(t, uris, "res://beta", "res://beta must be present after the re-sync") +} + +// TestBridge_PromptsListChanged_TriggersReSync verifies the bridge's +// notifications/prompts/list_changed re-fetch fix: when the upstream +// backend's prompt set changes after the bridge has connected, the upstream +// emits a prompts/list_changed notification, the bridge's OnNotification +// handler re-runs forwardAll, and the newly added prompt appears on the +// bridge's local stdio server. +// +// Mirrors TestBridge_ToolsListChanged_TriggersReSync/ +// TestBridge_ResourcesListChanged_TriggersReSync: the mutation is driven +// through the per-session overlay (SessionWithPrompts.SetSessionPrompts), and +// readiness/re-fetch is observed via the tools/list counter for the same +// race-free reason (forwardAll always re-fetches tools too). +// +//nolint:paralleltest // Swaps process-global os.Stdin; cannot run in parallel. +func TestBridge_PromptsListChanged_TriggersReSync(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // --- upstream backend with an initial prompt "alpha" --- + counter := &listToolsCounter{} + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + hooks.AddBeforeListTools(counter.hook()) + + backend := server.NewMCPServer( + "backend", "1.0", + // WithToolCapabilities is declared (though no tool is registered) so the + // backend unambiguously serves tools/list and fires the BeforeListTools + // hook — the counter is the readiness signal for the prompt-triggered + // re-sync, since forwardAll re-fetches tools on ANY *_list_changed. + server.WithToolCapabilities(true), + server.WithPromptCapabilities(true), + server.WithHooks(hooks), + ) + backend.AddPrompt(mcp.NewPrompt("alpha"), noopPromptHandler) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // --- bridge pointing at the backend (streamable-http) --- + bridge, err := NewStdioBridge("test", ts.URL, types.TransportTypeStreamableHTTP) + require.NoError(t, err) + + // Swap os.Stdin ONLY for a pipe so ServeStdio does not touch the real + // process stdin and can be unblocked at teardown by closing the write end. + origIn := os.Stdin + pipeR, pipeW, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { + os.Stdin = origIn + _ = pipeR.Close() + _ = pipeW.Close() + }) + os.Stdin = pipeR + + bridge.Start(ctx) + t.Cleanup(func() { + _ = pipeW.Close() + bridge.Shutdown() + }) + + // Step 1: wait for the bridge to connect and run its initial forwardAll. + require.Eventually(t, func() bool { + return holder.get() != nil && counter.count.Load() >= 1 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not complete initial forwardAll (session=%v, listCalls=%d)", + holder.get() != nil, counter.count.Load()) + + // Step 2: mutate the upstream prompt set — add "beta" via the per-session + // overlay. SetSessionPrompts syncs onto the live go-sdk server bound to + // this session, which emits notifications/prompts/list_changed to the + // bridge's upstream client. + backendSession := holder.get() + swp, ok := backendSession.(server.SessionWithPrompts) + require.True(t, ok, "backend session must implement SessionWithPrompts") + swp.SetSessionPrompts(map[string]server.ServerPrompt{ + "alpha": {Prompt: mcp.NewPrompt("alpha"), Handler: noopPromptHandler}, + "beta": {Prompt: mcp.NewPrompt("beta"), Handler: noopPromptHandler}, + }) + + // Step 3: wait for the re-fetch, observed via the tools/list counter + // (forwardAll re-fetches tools on every *_list_changed notification). + require.Eventually(t, func() bool { + return counter.count.Load() >= 2 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not re-run forwardAll after prompts/list_changed (listCalls=%d)", + counter.count.Load()) + + // Step 4: stop the bridge, then read bridge.srv. The re-sync forwardAll runs + // on the upstream client's OnNotification goroutine, not run(); Shutdown's + // b.up.Close() joins that goroutine before returning, which is the + // happens-before edge that makes reading bridge.srv here race-free. + _ = pipeW.Close() + bridge.Shutdown() + + names := promptNamesOnServer(t, bridge.srv) + assert.Contains(t, names, "alpha", "alpha must be present after the re-fetch (forwardAll is additive)") + assert.Contains(t, names, "beta", "beta must be present after the re-sync") +} + +// TestBridge_ProgressAndLoggingNotifications_ForwardedByShim guards the premise +// the bridge's notification forwarding depends on: the mcpcompat client the +// bridge uses must deliver upstream notifications/progress and +// notifications/message to registered OnNotification handlers. The bridge's +// OnNotification handler then relays every method downstream unconditionally +// (SendNotificationToAllClients, in run()), so client delivery is the load- +// bearing link. +// +// An earlier shim iteration did NOT install the go-sdk +// ProgressNotificationHandler/LoggingMessageHandler and dropped these +// notifications; they were wired in toolhive-core (issue #156). This test fails +// if that regresses — a client behind the bridge would silently stop seeing +// progress/logging on long-running tool calls. +func TestBridge_ProgressAndLoggingNotifications_ForwardedByShim(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Live mcpcompat backend; capture the connected session so we can push + // server->client notifications to it once the client is listening. + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + backend := server.NewMCPServer("backend", "1.0", server.WithToolCapabilities(true), server.WithHooks(hooks)) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // The bridge connects its upstream client with the standalone SSE stream on + // (WithContinuousListening); mirror that here so server->client notifications + // have a stream to arrive on. + c, err := client.NewStreamableHttpClient(ts.URL, transport.WithContinuousListening()) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + var mu sync.Mutex + seen := map[string]int{} + c.OnNotification(func(n mcp.JSONRPCNotification) { + mu.Lock() + seen[n.Method]++ + mu.Unlock() + }) + + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "test", Version: "1.0"}, + }}) + require.NoError(t, err) + + // Logging is level-gated per the MCP spec (and per mcp-go, whose SetLevel the + // shim aliases): the client must subscribe before the server delivers + // notifications/message. Progress is not gated. Subscribe so both paths are + // exercised. + require.NoError(t, c.SetLoggingLevel(ctx, "info")) + + // Wait for the session to register (and the standalone SSE stream to attach) + // before pushing notifications, so they are not emitted into the void. + require.Eventually(t, func() bool { return holder.get() != nil }, 5*time.Second, 50*time.Millisecond, + "client did not connect to the backend") + + backend.SendNotificationToAllClients("notifications/progress", + map[string]any{"progressToken": "tok", "progress": 0.5}) + backend.SendNotificationToAllClients("notifications/message", + map[string]any{"level": "info", "data": "hello"}) + + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return seen["notifications/progress"] >= 1 && seen["notifications/message"] >= 1 + }, 5*time.Second, 50*time.Millisecond, + "shim client must forward progress + logging notifications to OnNotification (regression: they were dropped); got %v", seen) +} diff --git a/pkg/transport/factory.go b/pkg/transport/factory.go index eae3218987..f54fb13a1f 100644 --- a/pkg/transport/factory.go +++ b/pkg/transport/factory.go @@ -52,6 +52,7 @@ func (*Factory) Create(config types.Config, opts ...Option) (types.Transport, er config.PrometheusHandler, config.Middlewares..., ) stdio.SetProxyMode(config.ProxyMode) + stdio.SetStrictProtocolValidation(config.StrictProtocolValidation) if config.SessionStorage != nil { stdio.SetSessionStorage(config.SessionStorage) } diff --git a/pkg/transport/middleware/origin/origin.go b/pkg/transport/middleware/origin/origin.go index 302264751d..556a988f66 100644 --- a/pkg/transport/middleware/origin/origin.go +++ b/pkg/transport/middleware/origin/origin.go @@ -36,7 +36,7 @@ const ( // forbiddenBodyFallback is returned if JSON marshalling of the error body // fails (should never happen with simple map types). - forbiddenBodyFallback = `{"jsonrpc":"2.0","error":{"code":-32600,"message":"Origin not allowed"},"id":null}` + forbiddenBodyFallback = `{"jsonrpc":"2.0","error":{"code":-32600,"message":"Origin not allowed"}}` ) // MiddlewareParams holds the parameters for the origin middleware factory. @@ -256,7 +256,12 @@ func isLoopbackHost(host string) bool { return false } -// writeForbidden emits a 403 response with a JSON-RPC error body (id: null). +// writeForbidden emits a 403 response with a JSON-RPC error body. The "id" +// key is omitted: this rejection happens on the Origin header alone, before +// any JSON-RPC body is parsed, so there is no incoming id to echo. MCP +// narrows base JSON-RPC 2.0 here -- schema/2025-11-25 types the error +// response id as optional, string|number (no null), so an absent id must be +// encoded by omitting the key, never as "id":null. func writeForbidden(w http.ResponseWriter) { body, err := json.Marshal(map[string]any{ "jsonrpc": "2.0", @@ -264,7 +269,6 @@ func writeForbidden(w http.ResponseWriter) { "code": jsonRPCCodeInvalidRequest, "message": "Origin not allowed", }, - "id": nil, }) if err != nil { // Marshal of a static map should never fail; fall back to a literal. diff --git a/pkg/transport/middleware/origin/origin_test.go b/pkg/transport/middleware/origin/origin_test.go index 805ddb1ae7..e3b8683500 100644 --- a/pkg/transport/middleware/origin/origin_test.go +++ b/pkg/transport/middleware/origin/origin_test.go @@ -5,7 +5,6 @@ package origin import ( "encoding/json" - "io" "net/http" "net/http/httptest" "testing" @@ -175,27 +174,36 @@ func TestOriginMiddleware_MultipleOriginHeadersRejected(t *testing.T) { // assertForbiddenJSONRPC validates that rec carries a 403 with a canonical // JSON-RPC error body and that the inner handler was never invoked. +// +// The JSONEq below already catches an extra "id":null (it wouldn't match the +// expected body), so the explicit key-absence check that follows is +// belt-and-suspenders -- kept because it names the failure directly, whereas +// a JSONEq diff would just show an unexpected field. func assertForbiddenJSONRPC(t *testing.T, rec *httptest.ResponseRecorder, nextCalled bool) { t.Helper() assert.False(t, nextCalled, "next handler must NOT be invoked") assert.Equal(t, http.StatusForbidden, rec.Code) assert.Equal(t, "application/json", rec.Header().Get("Content-Type")) + assert.JSONEq(t, `{"jsonrpc":"2.0","error":{"code":-32600,"message":"Origin not allowed"}}`, rec.Body.String()) - body, err := io.ReadAll(rec.Body) - require.NoError(t, err) - var parsed struct { - JSONRPC string `json:"jsonrpc"` - Error struct { - Code int64 `json:"code"` - Message string `json:"message"` - } `json:"error"` - ID any `json:"id"` - } - require.NoError(t, json.Unmarshal(body, &parsed)) - assert.Equal(t, "2.0", parsed.JSONRPC) - assert.Equal(t, jsonRPCCodeInvalidRequest, parsed.Error.Code) - assert.Equal(t, "Origin not allowed", parsed.Error.Message) - assert.Nil(t, parsed.ID) + var parsed map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &parsed)) + _, hasID := parsed["id"] + assert.False(t, hasID, "forbidden body must omit \"id\" entirely, never send it as null") +} + +// TestForbiddenBodyFallback verifies the hand-written fallback literal used +// when json.Marshal fails (unreachable in practice with static map types, so +// this is the only path that exercises it) has the same shape as the +// marshaled body: no "id" key, never "null". +func TestForbiddenBodyFallback(t *testing.T) { + t.Parallel() + assert.JSONEq(t, `{"jsonrpc":"2.0","error":{"code":-32600,"message":"Origin not allowed"}}`, forbiddenBodyFallback) + + var parsed map[string]any + require.NoError(t, json.Unmarshal([]byte(forbiddenBodyFallback), &parsed)) + _, hasID := parsed["id"] + assert.False(t, hasID, "fallback body must omit \"id\" entirely, never send it as null") } func TestNewHandler_RejectsDisallowedOrigin(t *testing.T) { diff --git a/pkg/transport/proxy/streamable/dispatcher.go b/pkg/transport/proxy/streamable/dispatcher.go index 4e75c26a61..82e3d7b54a 100644 --- a/pkg/transport/proxy/streamable/dispatcher.go +++ b/pkg/transport/proxy/streamable/dispatcher.go @@ -6,63 +6,222 @@ package streamable import ( "fmt" "log/slog" - "time" "golang.org/x/exp/jsonrpc2" ) -// dispatchResponses routes container responses to the appropriate waiter by request ID. -// Notifications are ignored for streamable HTTP. +// JSON-RPC method names used for server<->client routing decisions across +// this package (dispatcher.go and streamable_proxy.go). Collected here as +// the single source of truth so the two files can never drift on a wire +// string; none of these are exported by any SDK the proxy depends on (see +// #5744's review). +const ( + methodToolsListChanged = "notifications/tools/list_changed" + methodResourcesListChanged = "notifications/resources/list_changed" + methodPromptsListChanged = "notifications/prompts/list_changed" + methodProgress = "notifications/progress" + methodResourcesUpdated = "notifications/resources/updated" + methodLoggingMessage = "notifications/message" + methodResourcesSubscribe = "resources/subscribe" + methodResourcesUnsubscribe = "resources/unsubscribe" +) + +// listChangedNotificationMethods is the set of server->client notification +// methods that describe a server-wide capability change (not tied to any +// particular request or subscription), so delivering one copy to every +// connected session's standalone stream is correct and safe. +var listChangedNotificationMethods = map[string]bool{ + methodToolsListChanged: true, + methodResourcesListChanged: true, + methodPromptsListChanged: true, +} + +// dispatchResponses routes container messages arriving on responseCh to the +// appropriate destination based on their concrete JSON-RPC type: +// - *jsonrpc2.Response routes to the waiter correlated by its composite ID +// (see routeResponseToWaiter). +// - *jsonrpc2.Request with no valid ID is a server->client notification, +// routed by method (see routeNotification). +// - *jsonrpc2.Request with a valid ID is a server->client REQUEST (e.g. +// sampling/createMessage, elicitation/create). It is rejected back to the +// backend rather than forwarded to any client (see +// rejectServerRequestToBackend). func (p *HTTPProxy) dispatchResponses() { for { select { case <-p.shutdownCh: return - case resp := <-p.responseCh: - // Ignore notifications - if isNotification(resp) { - continue - } - r, ok := resp.(*jsonrpc2.Response) - if !ok || !r.ID.IsValid() { - slog.Warn("received invalid message that is not a valid response", - "type", fmt.Sprintf("%T", resp)) - continue - } - - rawID := r.ID.Raw() - // Composite-only routing: responses must carry composite ID (sessID|idKey) - if sID, ok := rawID.(string); ok { - if chVal, ok := p.waiters.Load(sID); ok { - if ch, ok := chVal.(chan jsonrpc2.Message); ok { - select { - case ch <- resp: - default: - slog.Warn("waiter channel full; dropping response", - "composite_key", sID) - } - continue - } + case msg := <-p.responseCh: + switch m := msg.(type) { + case *jsonrpc2.Response: + p.routeResponseToWaiter(m) + case *jsonrpc2.Request: + if !m.ID.IsValid() { + p.routeNotification(m) + continue } - slog.Warn("no waiter found for composite key; dropping", "composite_key", sID) - continue + p.rejectServerRequestToBackend(m) + default: + slog.Warn("received invalid message that is not a valid response", + "type", fmt.Sprintf("%T", msg)) } - slog.Warn("non-string response id (expected composite string); dropping", - "raw_id", fmt.Sprintf("%v", rawID)) } } } -// waitForResponse waits for a response on the given channel with timeout. -func (p *HTTPProxy) waitForResponse(ch <-chan jsonrpc2.Message, timeout time.Duration) jsonrpc2.Message { - timer := time.NewTimer(timeout) - defer timer.Stop() +// routeNotification delivers a server->client notification (m.ID invalid) to +// its correct destination based on m.Method, per the shared-backend routing +// design (see #5744 and docs/arch/03-transport-architecture.md): +// +// - */list_changed (listChangedNotificationMethods): GLOBAL, no +// session-specific payload -- broadcast to every connected session's +// standalone GET stream. +// - notifications/progress: request-scoped -- delivered ONLY to the +// originating request's POST-SSE stream, via routeProgress's +// progress-token correlation. Never broadcast or fanned out. +// - notifications/resources/updated: subscription-scoped -- delivered ONLY +// to sessions currently subscribed to the notification's uri, via +// routeResourceUpdated. +// - notifications/message (logging): SECURE-DROP. The shared backend is one +// process serving every session; a log message it emits carries no +// session attribution, so there is no way to deliver it to the session +// that caused it without either guessing or leaking it to every other +// connected session. Dropping it is the only safe choice until per-session +// backends (see the vMCP architecture) are in the request path. +// - anything else: dropped, logged at Debug. +func (p *HTTPProxy) routeNotification(m *jsonrpc2.Request) { + switch { + case listChangedNotificationMethods[m.Method]: + p.serverStreams.broadcast(m) + case m.Method == methodProgress: + p.routeProgress(m) + case m.Method == methodResourcesUpdated: + p.routeResourceUpdated(m) + case m.Method == methodLoggingMessage: + // SECURE-DROP: shared-backend log content is unattributable to any + // one session and forwarding it (to all sessions, or a guessed one) + // would leak cross-session information. Deferred to per-session + // backend work (#5744). + slog.Debug("dropping logging notification; shared backend cannot attribute it to a session") + default: + //nolint:gosec // G706: method is from parsed JSON-RPC request + slog.Debug("dropping unrecognized server->client notification", "method", m.Method) + } +} + +// routeProgress delivers a notifications/progress message to the single +// in-flight POST-SSE request that asked for it, correlated by the raw +// (proxy-minted) progressToken the backend echoes back (see +// rewriteMetaProgressToken and handleSingleRequestSSE's setupProgressRouting). +// If the token is missing or not currently registered (the request already +// completed, or the token was never ours), the notification is dropped: there +// is no session to deliver it to without guessing, and guessing risks +// cross-session leakage. +func (p *HTTPProxy) routeProgress(m *jsonrpc2.Request) { + ptGlobal, ok := extractStringParam(m.Params, "progressToken") + if !ok || ptGlobal == "" { + slog.Debug("dropping progress notification with no progressToken") + return + } + + route, found := p.routing.lookupProgressToken(ptGlobal) + if !found { + slog.Debug("dropping progress notification for unknown or expired progress token") + return + } + + restored, err := rewriteRequestParam(m, "progressToken", route.originalToken) + if err != nil { + slog.Error("failed to restore client progressToken on progress notification", "error", err) + return + } + + select { + case route.deliver <- restored: + default: + slog.Warn("progress delivery channel full; dropping progress notification") + } +} + +// routeResourceUpdated delivers a notifications/resources/updated message to +// every session currently subscribed to its uri (see +// sessionRouter.subscribersOf), and only those sessions. A notification for a +// uri with no recorded subscriber is dropped. +func (p *HTTPProxy) routeResourceUpdated(m *jsonrpc2.Request) { + uri, ok := extractStringParam(m.Params, "uri") + if !ok || uri == "" { + slog.Debug("dropping resources/updated notification with no uri") + return + } + + subscribers := p.routing.subscribersOf(uri) + if len(subscribers) == 0 { + //nolint:gosec // G706: uri is from parsed JSON-RPC notification + slog.Debug("dropping resources/updated notification with no subscribers", "uri", uri) + return + } + + p.serverStreams.deliverToMany(subscribers, m) +} + +// rejectServerRequestToBackend responds to a server-initiated request (a +// *jsonrpc2.Request with a valid ID, e.g. sampling/createMessage or +// elicitation/create) with a JSON-RPC error, written back to the BACKEND (not +// to any client) via SendMessageToDestination, so the backend's own blocking +// call unblocks with an error instead of hanging until its own timeout. +// +// The shared-backend streamable proxy has no way to route a client's reply +// back to the correct originating backend call when multiple sessions share +// one backend process, so it cannot forward this to any client without either +// guessing (misdelivery) or broadcasting (cross-session leakage). Per-session +// backends (see the vMCP architecture, docs/arch/10-virtual-mcp-architecture.md) +// are the correct place to support server->client requests; see #5744. +func (p *HTTPProxy) rejectServerRequestToBackend(m *jsonrpc2.Request) { + resp := &jsonrpc2.Response{ + ID: m.ID, + Error: jsonrpc2.NewError(-32601, fmt.Sprintf( + "server-initiated %s is not supported by the shared streamable proxy; requires a per-session backend deployment", + m.Method, + )), + } + if err := p.SendMessageToDestination(resp); err != nil { + //nolint:gosec // G706: method is from parsed JSON-RPC request + slog.Error("failed to send server-request-rejected error back to backend", "method", m.Method, "error", err) + } +} + +// routeResponseToWaiter delivers resp to the waiter channel correlated by its +// composite request ID (compositeKey(sessID, idKey)), if one is registered. +func (p *HTTPProxy) routeResponseToWaiter(resp *jsonrpc2.Response) { + if !resp.ID.IsValid() { + slog.Warn("received invalid message that is not a valid response", + "type", fmt.Sprintf("%T", resp)) + return + } + + rawID := resp.ID.Raw() + // Composite-only routing: responses must carry composite ID (sessID|idKey) + sID, ok := rawID.(string) + if !ok { + slog.Warn("non-string response id (expected composite string); dropping", + "raw_id", fmt.Sprintf("%v", rawID)) + return + } + + chVal, ok := p.waiters.Load(sID) + if !ok { + slog.Warn("no waiter found for composite key; dropping", "composite_key", sID) + return + } + ch, ok := chVal.(chan jsonrpc2.Message) + if !ok { + slog.Warn("no waiter found for composite key; dropping", "composite_key", sID) + return + } + select { - case msg := <-ch: - return msg - case <-timer.C: - return nil - case <-p.shutdownCh: - return nil + case ch <- resp: + default: + slog.Warn("waiter channel full; dropping response", "composite_key", sID) } } diff --git a/pkg/transport/proxy/streamable/dispatcher_routing.go b/pkg/transport/proxy/streamable/dispatcher_routing.go new file mode 100644 index 0000000000..9527b99f04 --- /dev/null +++ b/pkg/transport/proxy/streamable/dispatcher_routing.go @@ -0,0 +1,296 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "sync" + + "golang.org/x/exp/jsonrpc2" +) + +// progressRoute is the per-request delivery route for notifications/progress +// messages that echo back a proxy-minted ptGlobal token (see +// rewriteMetaProgressToken and handleSingleRequestSSE's setupProgressRouting). +type progressRoute struct { + // deliver is the in-flight POST-SSE handler's per-request delivery + // channel. It is buffered (see progressDeliverBufferSize) and owned by + // that handler; routeProgress only ever sends to it, never closes it. + deliver chan jsonrpc2.Message + // originalToken is the client's own progressToken value (string or + // number), restored into the notification before delivery so the client + // sees the token it originally asked for, never ptGlobal. + originalToken any +} + +// sessionRouter owns three INDEPENDENT registries used to route server->client +// messages that arrive on the shared backend's response channel to the +// correct session(s) without leaking to any other session: +// +// - progress tokens: request-scoped, correlate a notifications/progress +// reply to the single in-flight POST-SSE request that asked for it. +// - subscriptions: session-scoped, correlate a resources/updated +// notification to the set of sessions currently subscribed to its uri, +// and ref-count resources/subscribe|unsubscribe so the shared backend only +// sees one upstream (un)subscribe per uri. +// - log levels: session-scoped, track each session's requested logging +// verbosity so the shared backend's single process-wide level can be +// reconciled to the maximum (most verbose) requested by any session. +// +// Each registry has its own mutex (see the "one synchronization primitive per +// data structure" style rule): the three are never locked together, and no +// method here locks more than one of ptMu/subMu/logMu at a time. +type sessionRouter struct { + ptMu sync.Mutex + ptToks map[string]progressRoute // key: ptGlobal (proxy-minted, unique per request) + + subMu sync.Mutex + subs map[string]map[string]struct{} // uri -> set(sessID) + + logMu sync.Mutex + logLvl map[string]int // sessID -> requested verbosity rank (higher = more verbose) +} + +// newSessionRouter creates an empty sessionRouter. +func newSessionRouter() *sessionRouter { + return &sessionRouter{ + ptToks: make(map[string]progressRoute), + subs: make(map[string]map[string]struct{}), + logLvl: make(map[string]int), + } +} + +// recordProgressToken registers route under ptGlobal, so a later +// notifications/progress echoing ptGlobal is delivered via route.deliver (see +// lookupProgressToken). The caller (handleSingleRequestSSE) MUST call +// dropProgressToken(ptGlobal) once the request completes -- progress tokens +// are request-scoped, not session-scoped, so they are never reaped by reap. +func (r *sessionRouter) recordProgressToken(ptGlobal string, route progressRoute) { + r.ptMu.Lock() + defer r.ptMu.Unlock() + r.ptToks[ptGlobal] = route +} + +// lookupProgressToken returns the route registered for ptGlobal, if any. +func (r *sessionRouter) lookupProgressToken(ptGlobal string) (progressRoute, bool) { + r.ptMu.Lock() + defer r.ptMu.Unlock() + route, ok := r.ptToks[ptGlobal] + return route, ok +} + +// dropProgressToken removes the route registered for ptGlobal, if any. It is +// always safe to call, including for a ptGlobal that was never recorded. +func (r *sessionRouter) dropProgressToken(ptGlobal string) { + r.ptMu.Lock() + defer r.ptMu.Unlock() + delete(r.ptToks, ptGlobal) +} + +// addSubscription records sess as a subscriber of uri and reports whether +// sess is the FIRST subscriber currently recorded for uri. +// +// The forward-vs-dedup decision is NOT driven by this return value: to record +// a subscription only after the upstream subscribe succeeds (see +// interceptSubscribe), the caller first peeks subscribersOf(uri) to decide +// whether an upstream call is needed, forwards, and only then calls +// addSubscription on success. Once any session is subscribed to uri the shared +// backend is already sending updates for it, so a later subscribe for the same +// uri is served locally without another upstream call. firstForURI is retained +// for tests/observability. +func (r *sessionRouter) addSubscription(uri, sess string) (firstForURI bool) { + r.subMu.Lock() + defer r.subMu.Unlock() + + set, ok := r.subs[uri] + if !ok { + set = make(map[string]struct{}) + r.subs[uri] = set + } + firstForURI = len(set) == 0 + set[sess] = struct{}{} + return firstForURI +} + +// removeSubscription removes sess as a subscriber of uri and reports whether +// sess was the LAST subscriber recorded for uri. The caller uses +// lastForURI to decide whether to forward resources/unsubscribe upstream: the +// backend should only stop sending updates for uri once no session wants them +// anymore. +func (r *sessionRouter) removeSubscription(uri, sess string) (lastForURI bool) { + r.subMu.Lock() + defer r.subMu.Unlock() + + set, ok := r.subs[uri] + if !ok { + return false + } + delete(set, sess) + lastForURI = len(set) == 0 + if lastForURI { + delete(r.subs, uri) + } + return lastForURI +} + +// subscribersOf returns the session IDs currently subscribed to uri (nil if +// none). The caller (dispatcher.go's routeResourceUpdated) uses this as the +// routeKeys for serverStreamRegistry.deliverToMany, so a resources/updated +// notification reaches only sessions that actually subscribed to that uri. +func (r *sessionRouter) subscribersOf(uri string) []string { + r.subMu.Lock() + defer r.subMu.Unlock() + + set, ok := r.subs[uri] + if !ok { + return nil + } + out := make([]string, 0, len(set)) + for sess := range set { + out = append(out, sess) + } + return out +} + +// setLogLevel records sess's requested verbosity rank and reports whether the +// MAXIMUM rank across all currently-tracked sessions changed as a result, +// along with the new maximum. The caller (streamable_proxy.go's +// logging/setLevel handling) uses maxChanged to decide whether to reconcile +// the shared backend's single process-wide log level upstream: forwarding +// each session's own request directly would let the last writer silently +// under-serve an earlier, more-verbose requester, since every session shares +// one backend log level. +func (r *sessionRouter) setLogLevel(sess string, rank int) (maxChanged bool, newMax int) { + r.logMu.Lock() + defer r.logMu.Unlock() + + before := r.maxLogLevelLocked() + r.logLvl[sess] = rank + after := r.maxLogLevelLocked() + return after != before, after +} + +// dropLogLevel removes sess's tracked verbosity rank, if any. It is always +// safe to call, including for a sess that was never recorded. +func (r *sessionRouter) dropLogLevel(sess string) { + r.logMu.Lock() + defer r.logMu.Unlock() + delete(r.logLvl, sess) +} + +// maxLogLevelLocked returns the maximum rank across all tracked sessions, or +// -1 if none are tracked. Callers MUST hold logMu. +func (r *sessionRouter) maxLogLevelLocked() int { + maxRank := -1 + for _, rank := range r.logLvl { + if rank > maxRank { + maxRank = rank + } + } + return maxRank +} + +// purgeSession removes sess from every subscription set and from the +// log-level table. Progress tokens are request-scoped (see +// recordProgressToken's doc comment), not session-scoped, so they are +// untouched here -- an in-flight request's progress token outlives this call +// by design, and is dropped by its own handler's deferred dropProgressToken. +// +// It reports whether removing sess changed the maximum log-level rank, so the +// caller (handleDelete) can decide whether the shared backend's level should +// be reconciled down now that sess is gone. +func (r *sessionRouter) purgeSession(sess string) (maxChanged bool, newMax int) { + r.subMu.Lock() + for uri, set := range r.subs { + if _, ok := set[sess]; ok { + delete(set, sess) + if len(set) == 0 { + delete(r.subs, uri) + } + } + } + r.subMu.Unlock() + + r.logMu.Lock() + defer r.logMu.Unlock() + if _, ok := r.logLvl[sess]; !ok { + return false, r.maxLogLevelLocked() + } + before := r.maxLogLevelLocked() + delete(r.logLvl, sess) + after := r.maxLogLevelLocked() + return after != before, after +} + +// reap removes subs/logLvl entries for every session isActive reports as no +// longer active, bounding the router's memory growth for sessions whose owning +// client disappeared without an explicit DELETE (see purgeSession for that +// path). Progress tokens are request-scoped and are never touched by reap +// (see recordProgressToken's doc comment). +// +// isActive is called OUTSIDE of ptMu/subMu/logMu: it may perform I/O (e.g. a +// Redis-backed session.Manager), so resolving liveness for every candidate +// session happens in a first pass with no lock held, and stale entries are +// then removed in a second pass that only ever does in-memory map work while +// holding a lock. This keeps reap from blocking concurrent +// addSubscription/removeSubscription/setLogLevel calls for the (possibly +// slow) duration of every liveness check. +// +// TOCTOU, the other direction: because isActive is snapshotted in the first +// pass and never re-checked once the second pass's locks are held, a session +// that (a) expires, (b) is re-created with the SAME session ID, and (c) +// re-subscribes or re-sets a log level -- all within the window between this +// tick's first and second pass -- could have its brand-new entry wrongly +// swept away by this same tick, since isActive's stale "false" result for the +// old incarnation is applied to the new one's state too. This is practically +// unreachable: session IDs are server-minted UUIDs (see resolveSessionForRequest), +// never reused by any real client in practice, so the "same ID reappears +// mid-tick" precondition does not occur outside a contrived test. Re-checking +// liveness a second time inside the phase-3 lock would close this gap, but is +// deliberately avoided: isActive's underlying session-store read (e.g. Redis +// GETEX) is I/O, and doing it while holding subMu/logMu would reintroduce the +// exact blocking-under-lock problem the two-pass design exists to avoid. +func (r *sessionRouter) reap(isActive func(sess string) bool) { + seen := make(map[string]struct{}) + + r.subMu.Lock() + for _, set := range r.subs { + for sess := range set { + seen[sess] = struct{}{} + } + } + r.subMu.Unlock() + + r.logMu.Lock() + for sess := range r.logLvl { + seen[sess] = struct{}{} + } + r.logMu.Unlock() + + stale := make(map[string]struct{}, len(seen)) + for sess := range seen { + if !isActive(sess) { + stale[sess] = struct{}{} + } + } + if len(stale) == 0 { + return + } + + r.subMu.Lock() + for uri, set := range r.subs { + for sess := range stale { + delete(set, sess) + } + if len(set) == 0 { + delete(r.subs, uri) + } + } + r.subMu.Unlock() + + r.logMu.Lock() + for sess := range stale { + delete(r.logLvl, sess) + } + r.logMu.Unlock() +} diff --git a/pkg/transport/proxy/streamable/dispatcher_routing_test.go b/pkg/transport/proxy/streamable/dispatcher_routing_test.go new file mode 100644 index 0000000000..1025a7443c --- /dev/null +++ b/pkg/transport/proxy/streamable/dispatcher_routing_test.go @@ -0,0 +1,228 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "strconv" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" +) + +// --- progress tokens --- + +func TestSessionRouter_ProgressTokenRecordLookupDrop(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + + _, ok := r.lookupProgressToken("pt-1") + assert.False(t, ok, "unrecorded token must not be found") + + deliver := make(chan jsonrpc2.Message, 1) + r.recordProgressToken("pt-1", progressRoute{deliver: deliver, originalToken: "client-token"}) + + route, ok := r.lookupProgressToken("pt-1") + require.True(t, ok) + assert.Equal(t, "client-token", route.originalToken) + assert.Equal(t, deliver, route.deliver) + + r.dropProgressToken("pt-1") + _, ok = r.lookupProgressToken("pt-1") + assert.False(t, ok, "dropped token must no longer be found") + + // Dropping an unrecorded token is a no-op, not a panic. + assert.NotPanics(t, func() { r.dropProgressToken("never-recorded") }) +} + +// --- subscriptions --- + +func TestSessionRouter_AddSubscription_ReportsFirstForURI(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + + assert.True(t, r.addSubscription("uri-1", "sess-a"), "first subscriber for a uri must report first=true") + assert.False(t, r.addSubscription("uri-1", "sess-b"), "second subscriber for the same uri must report first=false") + assert.True(t, r.addSubscription("uri-2", "sess-a"), "first subscriber for a DIFFERENT uri must report first=true") + + assert.ElementsMatch(t, []string{"sess-a", "sess-b"}, r.subscribersOf("uri-1")) + assert.ElementsMatch(t, []string{"sess-a"}, r.subscribersOf("uri-2")) +} + +func TestSessionRouter_RemoveSubscription_ReportsLastForURI(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + r.addSubscription("uri-1", "sess-a") + r.addSubscription("uri-1", "sess-b") + + assert.False(t, r.removeSubscription("uri-1", "sess-a"), "removing a non-last subscriber must report last=false") + assert.ElementsMatch(t, []string{"sess-b"}, r.subscribersOf("uri-1")) + + assert.True(t, r.removeSubscription("uri-1", "sess-b"), "removing the final subscriber must report last=true") + assert.Nil(t, r.subscribersOf("uri-1"), "uri entry must be cleaned up once its last subscriber leaves") + + // Removing from a uri with no recorded subscribers at all is a no-op. + assert.False(t, r.removeSubscription("never-subscribed", "sess-a")) +} + +func TestSessionRouter_SubscribersOf_UnknownURI(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + assert.Nil(t, r.subscribersOf("no-such-uri")) +} + +// --- log level --- + +func TestSessionRouter_SetLogLevel_MaxReconciliation(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + + maxChanged, newMax := r.setLogLevel("sess-a", 3) + assert.True(t, maxChanged, "first session's level always changes the max") + assert.Equal(t, 3, newMax) + + maxChanged, newMax = r.setLogLevel("sess-b", 1) + assert.False(t, maxChanged, "a LOWER verbosity from another session must not change the max") + assert.Equal(t, 3, newMax) + + maxChanged, newMax = r.setLogLevel("sess-b", 7) + assert.True(t, maxChanged, "a HIGHER verbosity must raise the max") + assert.Equal(t, 7, newMax) + + // Lowering sess-b (the current max holder) without removing sess-a (rank 3) + // must NOT drop the max below 3. + maxChanged, newMax = r.setLogLevel("sess-b", 0) + assert.True(t, maxChanged) + assert.Equal(t, 3, newMax, "max must fall back to the remaining highest tracked session") + + r.dropLogLevel("sess-a") + r.dropLogLevel("sess-b") + // Dropping already-removed sessions is a no-op. + assert.NotPanics(t, func() { r.dropLogLevel("sess-a") }) +} + +// --- purgeSession --- + +func TestSessionRouter_PurgeSession(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + r.addSubscription("uri-1", "sess-a") + r.addSubscription("uri-1", "sess-b") + r.addSubscription("uri-2", "sess-a") + r.setLogLevel("sess-a", 7) + r.setLogLevel("sess-b", 2) + + maxChanged, newMax := r.purgeSession("sess-a") + assert.True(t, maxChanged, "removing the max-rank session must change the max") + assert.Equal(t, 2, newMax) + + assert.ElementsMatch(t, []string{"sess-b"}, r.subscribersOf("uri-1")) + assert.Nil(t, r.subscribersOf("uri-2"), "uri-2 had only sess-a, so it must be fully removed") + + // Purging a session with no tracked state at all reports no change. + maxChanged, newMax = r.purgeSession("never-tracked") + assert.False(t, maxChanged) + assert.Equal(t, 2, newMax) +} + +// --- reap --- + +func TestSessionRouter_Reap_DropsInactiveSessionState(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + r.addSubscription("uri-1", "active-sess") + r.addSubscription("uri-1", "expired-sess") + r.addSubscription("uri-2", "expired-sess") + r.setLogLevel("active-sess", 3) + r.setLogLevel("expired-sess", 7) + + active := map[string]bool{"active-sess": true, "expired-sess": false} + isActive := func(sess string) bool { return active[sess] } + + r.reap(isActive) + + assert.ElementsMatch(t, []string{"active-sess"}, r.subscribersOf("uri-1")) + assert.Nil(t, r.subscribersOf("uri-2"), "uri-2 had only the expired session, so it must be fully removed") + + _, newMax := r.setLogLevel("active-sess", 3) // re-set to read back the max without mutating anything else + assert.Equal(t, 3, newMax, "expired session's log level must have been reaped") +} + +func TestSessionRouter_Reap_NoStaleSessionsIsNoop(t *testing.T) { + t.Parallel() + + r := newSessionRouter() + r.addSubscription("uri-1", "sess-a") + r.setLogLevel("sess-a", 5) + + r.reap(func(string) bool { return true }) + + assert.ElementsMatch(t, []string{"sess-a"}, r.subscribersOf("uri-1")) +} + +// TestSessionRouter_ConcurrentAccess hammers all three registries +// concurrently across many goroutines and sessions/URIs. Run with -race: the +// invariant under test is that each registry's own mutex fully serializes +// access to it (see the "one synchronization primitive per data structure" +// style rule) -- there is no cross-registry invariant to check here, only the +// absence of races/panics/deadlocks. +func TestSessionRouter_ConcurrentAccess(t *testing.T) { + t.Parallel() + + const ( + numWorkers = 16 + numRounds = 200 + ) + + r := newSessionRouter() + var wg sync.WaitGroup + + for i := 0; i < numWorkers; i++ { + wg.Add(1) + go func(worker int) { + defer wg.Done() + sess := "sess-" + strconv.Itoa(worker%4) + uri := "uri-" + strconv.Itoa(worker%3) + ptGlobal := "pt-" + strconv.Itoa(worker) + + for j := 0; j < numRounds; j++ { + r.addSubscription(uri, sess) + r.subscribersOf(uri) + r.removeSubscription(uri, sess) + + r.setLogLevel(sess, j%8) + r.dropLogLevel(sess) + + r.recordProgressToken(ptGlobal, progressRoute{deliver: make(chan jsonrpc2.Message, 1), originalToken: j}) + r.lookupProgressToken(ptGlobal) + r.dropProgressToken(ptGlobal) + + r.purgeSession(sess) + r.reap(func(string) bool { return true }) + } + }(i) + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(20 * time.Second): + t.Fatal("concurrent sessionRouter stress test timed out") + } +} diff --git a/pkg/transport/proxy/streamable/dispatcher_standalone_sse_integration_test.go b/pkg/transport/proxy/streamable/dispatcher_standalone_sse_integration_test.go new file mode 100644 index 0000000000..7fe41f04bc --- /dev/null +++ b/pkg/transport/proxy/streamable/dispatcher_standalone_sse_integration_test.go @@ -0,0 +1,1386 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" + + "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/transport/types" +) + +// establishSession POSTs an "initialize" request to the proxy and returns the +// Mcp-Session-Id the server assigns. startProxyWithBackend's fake backend +// answers any ID-bearing request (including "initialize") with a generic +// success result, which is sufficient here: session assignment happens in +// resolveSessionForRequest based on the request method alone, independent of +// the backend's response payload. +func establishSession(t *testing.T, port int) string { + t.Helper() + + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + initJSON := `{"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2025-03-26","clientInfo":{"name":"sse-test","version":"0.0.0"},"capabilities":{}}}` + resp, err := http.Post(url, "application/json", bytes.NewReader([]byte(initJSON))) //nolint:noctx // test-only + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + require.Equal(t, http.StatusOK, resp.StatusCode) + + sessID := resp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessID, "server should set Mcp-Session-Id header on initialize response") + return sessID +} + +// openGETStream issues a GET request to the proxy's MCP endpoint carrying +// sessID as Mcp-Session-Id, and returns the live response along with a +// buffered reader over its body. The caller owns closing resp.Body (register +// with t.Cleanup). +func openGETStream(ctx context.Context, t *testing.T, port int, sessID string) (*http.Response, *bufio.Reader) { + t.Helper() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint), nil) + require.NoError(t, err) + req.Header.Set("Mcp-Session-Id", sessID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + + return resp, bufio.NewReader(resp.Body) +} + +// readSSEData reads from r until it finds an SSE "data:" line (skipping +// comment/keep-alive lines), and returns its payload. It fails the test if no +// data line arrives within the timeout, per the "always bound blocking +// receives" testing rule. +func readSSEData(t *testing.T, r *bufio.Reader) string { + t.Helper() + + type result struct { + payload string + err error + } + ch := make(chan result, 1) + go func() { + for { + line, err := r.ReadString('\n') + if err != nil { + ch <- result{err: err} + return + } + trimmed := strings.TrimRight(line, "\r\n") + if data, ok := strings.CutPrefix(trimmed, "data:"); ok { + ch <- result{payload: strings.TrimSpace(data)} + return + } + // Skip blank lines, ":keep-alive" comments, and any other SSE fields. + } + }() + + select { + case res := <-ch: + require.NoError(t, res.err) + return res.payload + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for an SSE data line") + return "" + } +} + +// TestStandaloneSSE_NotificationReachesConnectedClient verifies a GLOBAL +// notification (see listChangedNotificationMethods) pushed via +// ForwardResponseToClients is delivered as a data: frame on a connected GET +// SSE stream. +func TestStandaloneSSE_NotificationReachesConnectedClient(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + + sessID := establishSession(t, port) + resp, reader := openGETStream(ctx, t, port, sessID) + t.Cleanup(func() { _ = resp.Body.Close() }) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "stream was not registered") + + notification, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, notification)) + + payload := readSSEData(t, reader) + assert.Contains(t, payload, "notifications/tools/list_changed") +} + +// TestStandaloneSSE_NonGlobalNotificationDoesNotReachClient verifies a +// notifications/progress message with no registered progress route (see +// dispatcher.go's routeProgress) pushed via ForwardResponseToClients is NOT +// delivered to a connected GET standalone SSE stream: progress belongs on its +// originating request's own SSE stream, never the standalone one, and an +// unrecognized/unregistered token is dropped rather than guessed at (see +// TestPostSSE_ProgressIsolationBetweenSessions for the positive, routed case). +func TestStandaloneSSE_NonGlobalNotificationDoesNotReachClient(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + + sessID := establishSession(t, port) + resp, reader := openGETStream(ctx, t, port, sessID) + t.Cleanup(func() { _ = resp.Body.Close() }) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "stream was not registered") + + nonGlobal, err := jsonrpc2.NewNotification("notifications/progress", map[string]any{"pct": 50}) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, nonGlobal)) + + // Prove the non-global notification did NOT reach the client by racing it + // against a global notification sent right after: if the non-global one + // had been delivered, it would be the first data: frame read here. + global, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, global)) + + payload := readSSEData(t, reader) + assert.Contains(t, payload, "notifications/tools/list_changed") + assert.NotContains(t, payload, "notifications/progress") +} + +// TestStandaloneSSE_FanOutToConcurrentClients verifies a single GLOBAL +// notification (see listChangedNotificationMethods) is delivered to every session with +// a concurrently connected GET SSE stream. It uses two DISTINCT sessions -- +// per MCP 2025-11-25, a single session's stream must never receive a +// notification more than once, so this test uses two different sessions to +// exercise multi-session fan-out without exercising (or accidentally +// validating) same-session duplication. +func TestStandaloneSSE_FanOutToConcurrentClients(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + + sessID1 := establishSession(t, port) + sessID2 := establishSession(t, port) + + resp1, reader1 := openGETStream(ctx, t, port, sessID1) + t.Cleanup(func() { _ = resp1.Body.Close() }) + resp2, reader2 := openGETStream(ctx, t, port, sessID2) + t.Cleanup(func() { _ = resp2.Body.Close() }) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 2 }, + time.Second, 5*time.Millisecond, "both streams were not registered") + + notification, err := jsonrpc2.NewNotification("notifications/resources/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, notification)) + + assert.Contains(t, readSSEData(t, reader1), "notifications/resources/list_changed") + assert.Contains(t, readSSEData(t, reader2), "notifications/resources/list_changed") +} + +// TestStandaloneSSE_ClientDisconnectDrainsRegistry verifies that closing a GET +// SSE client connection deregisters its stream, leaving the registry empty. +func TestStandaloneSSE_ClientDisconnectDrainsRegistry(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + + sessID := establishSession(t, port) + getCtx, getCancel := context.WithCancel(ctx) + resp, _ := openGETStream(getCtx, t, port, sessID) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "stream was not registered") + + getCancel() + _ = resp.Body.Close() + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 0 }, + time.Second, 5*time.Millisecond, "stream was not deregistered after client disconnect") +} + +// TestStandaloneSSE_SecondGetOnSameSessionEvictsFirst verifies that opening a +// second GET SSE stream for the same session evicts the first: the first +// connection's handler returns (its stop channel is closed) and the registry +// retains exactly one stream for the session. +func TestStandaloneSSE_SecondGetOnSameSessionEvictsFirst(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + + sessID := establishSession(t, port) + + firstCtx, firstCancel := context.WithCancel(ctx) + t.Cleanup(firstCancel) + firstResp, _ := openGETStream(firstCtx, t, port, sessID) + t.Cleanup(func() { _ = firstResp.Body.Close() }) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "first stream was not registered") + + secondResp, _ := openGETStream(ctx, t, port, sessID) + t.Cleanup(func() { _ = secondResp.Body.Close() }) + + // Still exactly one stream for the session (the second evicted the first), + // not two. + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "registry must retain exactly one stream per session") + + // The first connection's handler must observe eviction and return, ending + // its response body from the server side (the client observes EOF). + firstClosed := make(chan error, 1) + go func() { + _, err := io.Copy(io.Discard, firstResp.Body) + firstClosed <- err + }() + + select { + case err := <-firstClosed: + assert.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("first GET stream was not closed after eviction") + } +} + +// TestStandaloneSSE_GetSessionDecisionMatrix verifies handleGet's session +// resolution: a missing Mcp-Session-Id header is rejected with 400, an +// unknown one with 404, and a known one opens the SSE stream with 200. +func TestStandaloneSSE_GetSessionDecisionMatrix(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + _, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + + sessID := establishSession(t, port) + + tests := []struct { + name string + sessID string + setHeader bool + wantStatus int + }{ + {name: "missing header", setHeader: false, wantStatus: http.StatusBadRequest}, + {name: "unknown session", sessID: "bogus-session-id", setHeader: true, wantStatus: http.StatusNotFound}, + {name: "known session", sessID: sessID, setHeader: true, wantStatus: http.StatusOK}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + reqCtx, reqCancel := context.WithCancel(ctx) + t.Cleanup(reqCancel) + + req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, + fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint), nil) + require.NoError(t, err) + if tt.setHeader { + req.Header.Set("Mcp-Session-Id", tt.sessID) + } + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + + assert.Equal(t, tt.wantStatus, resp.StatusCode) + if tt.wantStatus == http.StatusOK { + assert.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + } + }) + } +} + +// TestStandaloneSSE_StopWhileStreamOpenReturnsCleanly verifies Stop returns +// without error, and unblocks the open GET handler, even while a client is +// connected to the standalone SSE stream. +func TestStandaloneSSE_StopWhileStreamOpenReturnsCleanly(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy := NewHTTPProxy("127.0.0.1", port, nil, nil) + bgCtx := context.Background() + require.NoError(t, proxy.Start(bgCtx)) + + // Backend responder so establishSession's initialize gets a response. + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + if req, ok := msg.(*jsonrpc2.Request); ok && req.ID.IsValid() { + result := map[string]any{"ok": true} + resp, _ := jsonrpc2.NewResponse(req.ID, result, nil) + _ = proxy.ForwardResponseToClients(bgCtx, resp) + } + case <-proxy.shutdownCh: + return + } + } + }() + + time.Sleep(50 * time.Millisecond) + + sessID := establishSession(t, port) + + // Independent context so the GET request outlives proxy.Stop below. + resp, _ := openGETStream(context.Background(), t, port, sessID) + t.Cleanup(func() { _ = resp.Body.Close() }) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "stream was not registered") + + stopCtx, stopCancel := context.WithTimeout(bgCtx, 5*time.Second) + defer stopCancel() + + done := make(chan error, 1) + go func() { done <- proxy.Stop(stopCtx) }() + + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("Stop did not return while a GET SSE stream was open") + } + + assert.Equal(t, 0, proxy.serverStreams.streamCount(), "streams must be closed by Stop") +} + +// TestStandaloneSSE_ListChangedRefiltersThroughExistingMiddleware is a +// security regression test: it proves the standalone GET SSE stream does not +// bypass tool-filter middleware. With a tool-filter configured to hide +// "filtered_tool": +// 1. a notifications/tools/list_changed notification still reaches the +// client over the GET stream (fan-out isn't blocked by middleware); +// 2. a subsequent tools/list POST response omits the filtered tool (the +// existing filtering behavior is unaffected by this change); and +// 3. a tools/call POST for the filtered tool is still blocked. +// +// All three requests flow through the same applyMiddlewares-wrapped handler +// (see Start's mux.Handle), so there is no separate, unfiltered code path for +// the GET stream. +func TestStandaloneSSE_ListChangedRefiltersThroughExistingMiddleware(t *testing.T) { + t.Parallel() + + filterOpts := []mcp.ToolMiddlewareOption{mcp.WithToolsFilter("allowed_tool")} + listMw, err := mcp.NewListToolsMappingMiddleware(filterOpts...) + require.NoError(t, err) + callMw, err := mcp.NewToolCallMappingMiddleware(filterOpts...) + require.NoError(t, err) + + middlewares := []types.NamedMiddleware{ + {Name: "tool-list-filter", Function: listMw}, + {Name: "tool-call-filter", Function: callMw}, + } + + port := pickFreePort(t) + proxy := NewHTTPProxy("127.0.0.1", port, nil, middlewares) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + require.NoError(t, proxy.Start(ctx)) + t.Cleanup(func() { + stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer stopCancel() + _ = proxy.Stop(stopCtx) + }) + time.Sleep(50 * time.Millisecond) + + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + var result any + if req.Method == "tools/list" { + result = map[string]any{ + "tools": []map[string]any{ + {"name": "allowed_tool", "description": "kept"}, + {"name": "filtered_tool", "description": "hidden"}, + }, + } + } else { + result = map[string]any{"ok": true} + } + resp, _ := jsonrpc2.NewResponse(req.ID, result, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + case <-ctx.Done(): + return + } + } + }() + + sessID := establishSession(t, port) + + // (1) notifications/tools/list_changed reaches the GET SSE client. + getResp, reader := openGETStream(ctx, t, port, sessID) + t.Cleanup(func() { _ = getResp.Body.Close() }) + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "stream was not registered") + + notification, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, notification)) + assert.Contains(t, readSSEData(t, reader), "notifications/tools/list_changed") + + endpoint := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + + // (2) tools/list POST response omits the filtered tool. + listBody := `{"jsonrpc":"2.0","id":"list-1","method":"tools/list"}` + listResp, err := http.Post(endpoint, "application/json", strings.NewReader(listBody)) //nolint:noctx // test-only + require.NoError(t, err) + t.Cleanup(func() { _ = listResp.Body.Close() }) + listRespBody, err := io.ReadAll(listResp.Body) + require.NoError(t, err) + assert.Contains(t, string(listRespBody), "allowed_tool") + assert.NotContains(t, string(listRespBody), "filtered_tool") + + // (3) tools/call POST for the filtered tool is still blocked. Since #5944 the + // tool-call filter returns a JSON-RPC error over HTTP 200 (the client accepts + // JSON) rather than a bare 400, so a filtered tool is indistinguishable from a + // tool that does not exist. + callBody := `{"jsonrpc":"2.0","id":"call-1","method":"tools/call","params":{"name":"filtered_tool"}}` + callResp, err := http.Post(endpoint, "application/json", strings.NewReader(callBody)) //nolint:noctx // test-only + require.NoError(t, err) + t.Cleanup(func() { _ = callResp.Body.Close() }) + assert.Equal(t, http.StatusOK, callResp.StatusCode) + callRespBody, err := io.ReadAll(callResp.Body) + require.NoError(t, err) + assert.Contains(t, string(callRespBody), `"error"`) + assert.Contains(t, string(callRespBody), "tool not found") +} + +// ------------------------- Routing security regressions ------------------------- +// +// The tests below exercise the FULL per-session routing table introduced to +// replace the placeholder global-only filter (progress correlation, resource +// subscriptions, logging reconciliation, and the sampling/elicitation +// rejection path). Each is a security regression: it proves one session's +// (or the backend's) activity is never delivered to a session that has no +// right to see it. + +// postSSE POSTs body to the proxy's MCP endpoint with Accept: +// text/event-stream, for sessID (empty omits the Mcp-Session-Id header +// entirely). It returns once response headers arrive (handleSingleRequestSSE +// flushes headers before waiting for any request-scoped data), so the +// returned reader can be used to read SSE frames as they are written, +// including interim progress frames before the final response. +func postSSE(ctx context.Context, t *testing.T, port int, sessID, body string) (*http.Response, *bufio.Reader) { + t.Helper() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint), strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + if sessID != "" { + req.Header.Set("Mcp-Session-Id", sessID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + return resp, bufio.NewReader(resp.Body) +} + +// postJSON POSTs body to the proxy's MCP endpoint (no Accept header, so the +// proxy answers with a plain JSON response) for sessID, and returns the +// response along with its fully-read body. +func postJSON(t *testing.T, port int, sessID, body string) (*http.Response, []byte) { + t.Helper() + req, err := http.NewRequest(http.MethodPost, //nolint:noctx // test-only + fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint), strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + if sessID != "" { + req.Header.Set("Mcp-Session-Id", sessID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + respBody, err := io.ReadAll(resp.Body) + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + return resp, respBody +} + +// extractMetaProgressToken reads params._meta.progressToken as a string. It +// simulates what a real MCP backend does with an incoming request: read back +// whatever progressToken the caller (here, the proxy, after +// rewriteMetaProgressToken) put in _meta, to later echo it verbatim in a +// notifications/progress reply. It deliberately avoids test assertions (no +// *testing.T) since it runs on backend goroutines, not the test goroutine. +func extractMetaProgressToken(params json.RawMessage) (string, bool) { + var m map[string]any + if err := json.Unmarshal(params, &m); err != nil { + return "", false + } + meta, ok := m["_meta"].(map[string]any) + if !ok { + return "", false + } + token, ok := meta["progressToken"].(string) + return token, ok +} + +// startProxyOnly starts an HTTPProxy on port with no automatic backend +// responder attached, for tests that need a custom backend goroutine (e.g. to +// track subscribe/unsubscribe call counts, or react differently per request). +func startProxyOnly(t *testing.T, port int) (*HTTPProxy, context.Context) { + t.Helper() + proxy := NewHTTPProxy("127.0.0.1", port, nil, nil) + ctx, cancel := context.WithCancel(context.Background()) + require.NoError(t, proxy.Start(ctx)) + t.Cleanup(func() { + cancel() + stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer stopCancel() + _ = proxy.Stop(stopCtx) + }) + time.Sleep(50 * time.Millisecond) + return proxy, ctx +} + +// TestPostSSE_ProgressIsolationBetweenSessions is the progress-isolation +// SECURITY regression: two DIFFERENT sessions each send a POST-SSE request +// carrying the SAME client-visible progressToken ("1"). The backend emits a +// progress notification correlated to session A's (proxy-minted, internal) +// progress token only; the test proves it reaches ONLY A's SSE stream (with +// the token restored to "1"), and that B's stream never receives a progress +// frame at all -- proving isolation holds even when both sessions happen to +// choose the identical client-facing token. +func TestPostSSE_ProgressIsolationBetweenSessions(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + // The backend only emits progress for the request tagged "which":"A", so + // the test controls precisely which session's progress token is used, + // without racing on request arrival order. + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + go func(req *jsonrpc2.Request) { + if req.Method == "tools/call" { + var params map[string]any + _ = json.Unmarshal(req.Params, ¶ms) + if params["which"] == "A" { + if token, ok := extractMetaProgressToken(req.Params); ok { + notif, nErr := jsonrpc2.NewNotification("notifications/progress", + map[string]any{"progressToken": token, "progress": 1}) + require.NoError(t, nErr) + _ = proxy.ForwardResponseToClients(ctx, notif) + } + } + } + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{"ok": true}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + }(req) + case <-ctx.Done(): + return + } + } + }() + + sessA := establishSession(t, port) + sessB := establishSession(t, port) + + bodyA := `{"jsonrpc":"2.0","id":"call-a","method":"tools/call","params":{"which":"A","_meta":{"progressToken":"1"}}}` + bodyB := `{"jsonrpc":"2.0","id":"call-b","method":"tools/call","params":{"which":"B","_meta":{"progressToken":"1"}}}` + + respA, readerA := postSSE(ctx, t, port, sessA, bodyA) + t.Cleanup(func() { _ = respA.Body.Close() }) + respB, readerB := postSSE(ctx, t, port, sessB, bodyB) + t.Cleanup(func() { _ = respB.Body.Close() }) + + // A's stream: progress frame (token restored to "1"), then the final response. + progressA := readSSEData(t, readerA) + assert.Contains(t, progressA, "notifications/progress") + assert.Contains(t, progressA, `"progressToken":"1"`) + finalA := readSSEData(t, readerA) + assert.Contains(t, finalA, `"id":"call-a"`) + + // B's stream: its ONLY frame is the final response -- no progress ever, even + // though B asked for the identical client-visible token "1". + onlyFrameB := readSSEData(t, readerB) + assert.Contains(t, onlyFrameB, `"id":"call-b"`) + assert.NotContains(t, onlyFrameB, "notifications/progress") +} + +// TestPostSSE_ProgressTokenClearedAfterRequestCompletes verifies a +// progress route is dropped once its POST-SSE request completes: a +// notifications/progress echoing that same (now-stale) token sent AFTER +// completion must be dropped rather than misdelivered. +func TestPostSSE_ProgressTokenClearedAfterRequestCompletes(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + capturedToken := make(chan string, 1) + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + if req.Method == "tools/call" { + if token, ok := extractMetaProgressToken(req.Params); ok { + capturedToken <- token + } + } + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{"ok": true}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + case <-ctx.Done(): + return + } + } + }() + + sessID := establishSession(t, port) + body := `{"jsonrpc":"2.0","id":"call-1","method":"tools/call","params":{"_meta":{"progressToken":"1"}}}` + + resp, reader := postSSE(ctx, t, port, sessID, body) + t.Cleanup(func() { _ = resp.Body.Close() }) + + final := readSSEData(t, reader) + assert.Contains(t, final, `"id":"call-1"`) + + var ptGlobal string + select { + case ptGlobal = <-capturedToken: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for backend to observe the rewritten progress token") + } + + // The request has completed (we already read its final frame), so its + // progress route must be gone: this lookup must fail. + require.Eventually(t, func() bool { + _, found := proxy.routing.lookupProgressToken(ptGlobal) + return !found + }, time.Second, 5*time.Millisecond, "progress token must be dropped once its request completes") +} + +// TestPostSSE_ResourcesUpdated_SubscribersOnlyAndUpstreamRefCounted is the +// resources/subscribe|unsubscribe SECURITY + ref-counting regression: +// - A subscribes uri1, B subscribes uri2: a resources/updated{uri1} +// notification reaches ONLY A (delivered on A's standalone GET stream), +// never B. +// - A THIRD session subscribing to uri1 does not cause a second upstream +// resources/subscribe call for uri1 (ref-counted). +// - Upstream resources/unsubscribe for uri1 is sent only once the LAST of +// its subscribers unsubscribes. +func TestPostSSE_ResourcesUpdated_SubscribersOnlyAndUpstreamRefCounted(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + var mu sync.Mutex + subscribeCalls := map[string]int{} + unsubscribeCalls := map[string]int{} + + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + var params map[string]any + _ = json.Unmarshal(req.Params, ¶ms) + uri, _ := params["uri"].(string) + + mu.Lock() + switch req.Method { + case "resources/subscribe": + subscribeCalls[uri]++ + case "resources/unsubscribe": + unsubscribeCalls[uri]++ + } + mu.Unlock() + + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + case <-ctx.Done(): + return + } + } + }() + + sessA := establishSession(t, port) + sessB := establishSession(t, port) + sessC := establishSession(t, port) + sessD := establishSession(t, port) + + getA, readerA := openGETStream(ctx, t, port, sessA) + t.Cleanup(func() { _ = getA.Body.Close() }) + getB, readerB := openGETStream(ctx, t, port, sessB) + t.Cleanup(func() { _ = getB.Body.Close() }) + + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 2 }, + time.Second, 5*time.Millisecond, "both GET streams were not registered") + + subscribe := func(sessID, uri string) { + body := fmt.Sprintf(`{"jsonrpc":"2.0","id":"sub-%s","method":"resources/subscribe","params":{"uri":%q}}`, sessID, uri) + resp, respBody := postJSON(t, port, sessID, body) + require.Equal(t, http.StatusOK, resp.StatusCode, "subscribe must succeed: %s", string(respBody)) + } + unsubscribe := func(sessID, uri string) { + body := fmt.Sprintf(`{"jsonrpc":"2.0","id":"unsub-%s","method":"resources/unsubscribe","params":{"uri":%q}}`, sessID, uri) + resp, respBody := postJSON(t, port, sessID, body) + require.Equal(t, http.StatusOK, resp.StatusCode, "unsubscribe must succeed: %s", string(respBody)) + } + + // A, C, D all subscribe to uri1 (three subscribers); B subscribes to uri2. + subscribe(sessA, "uri1") + subscribe(sessC, "uri1") + subscribe(sessD, "uri1") + subscribe(sessB, "uri2") + + mu.Lock() + assert.Equal(t, 1, subscribeCalls["uri1"], "only the FIRST subscribe for uri1 must reach the backend") + assert.Equal(t, 1, subscribeCalls["uri2"], "uri2's subscribe must independently reach the backend once") + mu.Unlock() + + // updated{uri1} reaches only A (the only GET-stream-connected subscriber of + // uri1); B (subscribed only to uri2) must receive nothing. + notif, err := jsonrpc2.NewNotification("notifications/resources/updated", map[string]any{"uri": "uri1"}) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, notif)) + + assert.Contains(t, readSSEData(t, readerA), "uri1") + + // Prove B got nothing by racing a global sentinel right after. + sentinel, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, sentinel)) + assert.Contains(t, readSSEData(t, readerB), "notifications/tools/list_changed") + + // Unsubscribe A and C (neither is last: D remains) -- no upstream call yet. + unsubscribe(sessA, "uri1") + unsubscribe(sessC, "uri1") + mu.Lock() + assert.Equal(t, 0, unsubscribeCalls["uri1"], "upstream unsubscribe must not fire before the LAST subscriber leaves") + mu.Unlock() + + // D is last: NOW the upstream unsubscribe must fire. + unsubscribe(sessD, "uri1") + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return unsubscribeCalls["uri1"] == 1 + }, time.Second, 5*time.Millisecond, "upstream unsubscribe must fire once the last subscriber leaves") +} + +// TestPostSSE_LoggingDroppedAndUpstreamLevelIsMaxAcrossSessions verifies: +// - a notifications/message the backend sends is NEVER delivered to any +// connected client (shared-backend logs are unattributable to a session); +// - the upstream logging/setLevel level the backend observes is the MAXIMUM +// verbosity requested across all sessions, not merely the last session's +// own request (which would silently under-serve an earlier, more verbose +// requester). +func TestPostSSE_LoggingDroppedAndUpstreamLevelIsMaxAcrossSessions(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + var mu sync.Mutex + var observedLevels []string + + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + if req.Method == "logging/setLevel" { + var params map[string]any + _ = json.Unmarshal(req.Params, ¶ms) + level, _ := params["level"].(string) + mu.Lock() + observedLevels = append(observedLevels, level) + mu.Unlock() + } + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + case <-ctx.Done(): + return + } + } + }() + + sessA := establishSession(t, port) + getA, readerA := openGETStream(ctx, t, port, sessA) + t.Cleanup(func() { _ = getA.Body.Close() }) + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "GET stream was not registered") + + setLevel := func(sessID, level string) { + body := fmt.Sprintf(`{"jsonrpc":"2.0","id":"lvl-%s","method":"logging/setLevel","params":{"level":%q}}`, sessID, level) + resp, respBody := postJSON(t, port, sessID, body) + require.Equal(t, http.StatusOK, resp.StatusCode, "setLevel must succeed: %s", string(respBody)) + } + + sessB := establishSession(t, port) + + // reconcileUpstreamLogLevel runs in a background goroutine (fire-and-forget + // from the client's perspective, see its doc comment), so each step below + // waits for its own reconciliation to be observed before issuing the next + // setLevel -- otherwise the two reconciliations could race each other + // on the shared backend channel and land out of order. + setLevel(sessA, "warning") // rank 4: first request always reconciles upstream + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(observedLevels) == 1 && observedLevels[0] == "warning" + }, 2*time.Second, 5*time.Millisecond, "upstream must reconcile to sessA's warning level") + + setLevel(sessB, "debug") // rank 7 > 4: raises the max, reconciles upstream again + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(observedLevels) == 2 && observedLevels[1] == "debug" + }, 2*time.Second, 5*time.Millisecond, "upstream must reconcile to the new max (debug)") + + // rank 3 < 7 (current max): must NOT reconcile upstream a third time. + // require.Never (rather than a fixed sleep) actively fails on an unexpected + // extra reconcile, including a slow-but-present one that a sleep would miss. + setLevel(sessA, "error") + require.Never(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(observedLevels) > 2 + }, 500*time.Millisecond, 20*time.Millisecond, + "a session lowering below the current max must not reconcile upstream") + + // A logging notification the backend sends is dropped, not delivered. + logMsg, err := jsonrpc2.NewNotification("notifications/message", + map[string]any{"level": "debug", "data": "secret log line"}) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, logMsg)) + + sentinel, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, sentinel)) + + got := readSSEData(t, readerA) + assert.Contains(t, got, "notifications/tools/list_changed") + assert.NotContains(t, got, "secret log line") +} + +// TestPostSSE_SamplingRequestNeverReachesClientAndBackendGetsError verifies a +// server-initiated sampling/createMessage request is NEVER delivered to any +// connected client, and instead unblocks the backend with a JSON-RPC error +// response on its own messageCh (see rejectServerRequestToBackend). +func TestPostSSE_SamplingRequestNeverReachesClientAndBackendGetsError(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + // backendReceivedError captures the rejection response the backend loop + // below observes on messageCh, so the test can assert on it without a + // second, racing reader of the same channel. + backendReceivedError := make(chan *jsonrpc2.Response, 1) + + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + switch m := msg.(type) { + case *jsonrpc2.Request: + if !m.ID.IsValid() { + continue + } + resp, _ := jsonrpc2.NewResponse(m.ID, map[string]any{"ok": true}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + case *jsonrpc2.Response: + backendReceivedError <- m + } + case <-ctx.Done(): + return + } + } + }() + + sessID := establishSession(t, port) + getResp, reader := openGETStream(ctx, t, port, sessID) + t.Cleanup(func() { _ = getResp.Body.Close() }) + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "GET stream was not registered") + + samplingReq, err := jsonrpc2.NewCall(jsonrpc2.StringID("sampling-1"), "sampling/createMessage", + map[string]any{"messages": []any{}}) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, samplingReq)) + + // The rejection error is written back to the backend on messageCh. + select { + case resp := <-backendReceivedError: + assert.Equal(t, "sampling-1", resp.ID.Raw()) + require.Error(t, resp.Error) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for the backend to receive the rejection error") + } + + // Prove the sampling request never reached the client by racing a global + // sentinel right after. + sentinel, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, sentinel)) + + got := readSSEData(t, reader) + assert.Contains(t, got, "notifications/tools/list_changed") + assert.NotContains(t, got, "sampling/createMessage") +} + +// TestPostSSE_FailedFirstSubscribeIsNotRecordedAndSecondSubscriberRetries is +// the FIX 1 regression from #5744's review: interceptSubscribe used to record +// a subscription BEFORE forwarding the FIRST resources/subscribe upstream, so +// a backend rejection (or timeout) still left the uri's subscription table +// entry in place -- every later session's subscribe for the SAME uri was then +// dedup-served a synthesized success and never actually reached the backend, +// permanently starving the uri of a real subscription. +// +// This proves all three parts of the fix: (a) the FIRST client whose backend +// subscribe fails gets the real JSON-RPC error back, not a synthesized +// success; (b) the routing table has NO subscriber recorded for the uri after +// that failure; and (c) a SECOND client subscribing to the SAME uri actually +// forwards upstream again (the backend's subscribe call count increments to +// 2, proving it was not dedup-served) rather than being silently starved of a +// real subscription. +func TestPostSSE_FailedFirstSubscribeIsNotRecordedAndSecondSubscriberRetries(t *testing.T) { + t.Parallel() + + const uri = "err-uri" + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + var mu sync.Mutex + subscribeCalls := 0 + + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + if req.Method != "resources/subscribe" { + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + continue + } + + mu.Lock() + subscribeCalls++ + callNum := subscribeCalls + mu.Unlock() + + // The FIRST upstream resources/subscribe call fails (e.g. the + // resource does not exist); every subsequent call succeeds. + var resp *jsonrpc2.Response + if callNum == 1 { + resp = &jsonrpc2.Response{ID: req.ID, Error: jsonrpc2.NewError(-32002, "resource not found")} + } else { + resp, _ = jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + } + _ = proxy.ForwardResponseToClients(ctx, resp) + case <-ctx.Done(): + return + } + } + }() + + sessA := establishSession(t, port) + sessB := establishSession(t, port) + + // (a) sessA's subscribe gets the real backend error back. + subBody := fmt.Sprintf(`{"jsonrpc":"2.0","id":"sub-a","method":"resources/subscribe","params":{"uri":%q}}`, uri) + respA, bodyA := postJSON(t, port, sessA, subBody) + assert.Equal(t, http.StatusOK, respA.StatusCode, "a JSON-RPC error is still an HTTP 200 envelope") + assert.Contains(t, string(bodyA), "resource not found", "the client must see the real backend error") + assert.Contains(t, string(bodyA), `"error"`) + + // (b) the routing table must have NO subscriber recorded for uri. + assert.Empty(t, proxy.routing.subscribersOf(uri), + "a failed first subscribe must never be recorded in the routing table") + + // (c) sessB's subscribe for the SAME uri must actually reach the backend + // again (call count 2), not be dedup-served a synthesized success. + subBodyB := fmt.Sprintf(`{"jsonrpc":"2.0","id":"sub-b","method":"resources/subscribe","params":{"uri":%q}}`, uri) + respB, bodyB := postJSON(t, port, sessB, subBodyB) + assert.Equal(t, http.StatusOK, respB.StatusCode) + assert.NotContains(t, string(bodyB), `"error"`, "sessB's subscribe should succeed against the backend") + + mu.Lock() + assert.Equal(t, 2, subscribeCalls, "sessB's subscribe must actually forward upstream, not be deduped") + mu.Unlock() + + assert.ElementsMatch(t, []string{sessB}, proxy.routing.subscribersOf(uri), + "only sessB (the successful subscriber) should be recorded") +} + +// TestPostSSE_AuthzDeniedSubscribeNeverRecordedInRouting is the authz +// non-bypass SECURITY regression: a resources/subscribe request blocked by +// authorization middleware must NEVER be recorded in sessionRouter's +// subscription table, so a later resources/updated for that uri reaches +// nobody -- proving the routing table only ever reflects requests that were +// actually admitted upstream, not merely attempted by a client. +func TestPostSSE_AuthzDeniedSubscribeNeverRecordedInRouting(t *testing.T) { + t.Parallel() + + denyAll := types.NamedMiddleware{ + Name: "deny-all", + Function: func(http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + }) + }, + } + + port := pickFreePort(t) + proxy := NewHTTPProxy("127.0.0.1", port, nil, []types.NamedMiddleware{denyAll}) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + require.NoError(t, proxy.Start(ctx)) + t.Cleanup(func() { + stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer stopCancel() + _ = proxy.Stop(stopCtx) + }) + time.Sleep(50 * time.Millisecond) + + body := `{"jsonrpc":"2.0","id":"sub-1","method":"resources/subscribe","params":{"uri":"secret-uri"}}` + req, err := http.NewRequest(http.MethodPost, //nolint:noctx // test-only + fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint), strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", "does-not-matter-authz-denies-first") + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + require.Equal(t, http.StatusForbidden, resp.StatusCode) + + assert.Nil(t, proxy.routing.subscribersOf("secret-uri"), + "an authz-denied subscribe must never be recorded in the routing table") +} + +// TestHandleDelete_PurgesRoutingState verifies handleDelete purges a +// session's subscription and log-level state from routing, alongside +// deleting the session itself, AND (FIX 2 in #5744's review) tears down the +// session's standalone GET SSE stream: its handler goroutine must return +// (ending the response from the server side) and the registry must no longer +// retain an entry for the deleted session -- neither survives waiting for the +// client to eventually close the underlying socket on its own. +func TestHandleDelete_PurgesRoutingState(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + if req, ok := msg.(*jsonrpc2.Request); ok && req.ID.IsValid() { + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + } + case <-ctx.Done(): + return + } + } + }() + + sessID := establishSession(t, port) + postJSON(t, port, sessID, + `{"jsonrpc":"2.0","id":"sub-1","method":"resources/subscribe","params":{"uri":"uri1"}}`) + require.NotEmpty(t, proxy.routing.subscribersOf("uri1"), "subscription must be recorded before delete") + + getResp, _ := openGETStream(ctx, t, port, sessID) + t.Cleanup(func() { _ = getResp.Body.Close() }) + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "GET stream was not registered") + + delReq, err := http.NewRequest(http.MethodDelete, //nolint:noctx // test-only + fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint), nil) + require.NoError(t, err) + delReq.Header.Set("Mcp-Session-Id", sessID) + delResp, err := http.DefaultClient.Do(delReq) + require.NoError(t, err) + t.Cleanup(func() { _ = delResp.Body.Close() }) + require.Equal(t, http.StatusNoContent, delResp.StatusCode) + + assert.Nil(t, proxy.routing.subscribersOf("uri1"), "subscription must be purged by handleDelete") + + // The GET stream's handler goroutine must observe the closed stop channel + // and return, ending the response body from the server side (the client + // observes EOF) rather than staying open to keep receiving broadcasts. + getClosed := make(chan error, 1) + go func() { + _, err := io.Copy(io.Discard, getResp.Body) + getClosed <- err + }() + select { + case err := <-getClosed: + assert.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("GET stream was not torn down by handleDelete") + } + assert.Equal(t, 0, proxy.serverStreams.streamCount(), + "registry must no longer retain the deleted session's stream") +} + +// TestReapServerStreams_ClosesStreamForInactiveSession is the FIX 2 reaper +// regression from #5744's review: reapRoutingState used to leave +// serverStreams entirely untouched, so a session whose owning client +// disappeared without ever sending DELETE (and without the proxy observing +// its GET request's context ending) would keep its standalone GET stream -- +// and the handler goroutine consuming it -- alive indefinitely. This verifies +// reapServerStreams closes (and removes) the stream for a session isActive no +// longer reports as live, while leaving a still-active session's stream +// alone, using a real serverStreamRegistry and a stub isActive rather than +// the production reapRoutingState goroutine's real timer (too slow to wait +// out in a unit test). +func TestReapServerStreams_ClosesStreamForInactiveSession(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("127.0.0.1", 0, nil, nil) + goneStream := proxy.serverStreams.register("gone-session") + proxy.serverStreams.register("still-here-session") + + proxy.reapServerStreams(func(sess string) bool { return sess != "gone-session" }) + + assert.Equal(t, 1, proxy.serverStreams.streamCount(), + "only the inactive session's stream should be closed") + select { + case <-goneStream.stop: + // Expected: the reaper closed it. + case <-time.After(2 * time.Second): + t.Fatal("expired session's stream was not closed by the reaper") + } +} + +// TestSessionRouter_Reap_IntegratesWithProxy verifies the reaper wired into +// HTTPProxy.Start drops routing state for a session isActive reports as +// gone, using a real sessionRouter and a stub isActive rather than the +// production reapRoutingState goroutine's real timer (which runs on +// sessionTTL/2 -- too slow to wait out in a unit test). +func TestSessionRouter_Reap_IntegratesWithProxy(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("127.0.0.1", 0, nil, nil) + proxy.routing.addSubscription("uri1", "gone-session") + proxy.routing.setLogLevel("gone-session", 7) + proxy.routing.addSubscription("uri1", "still-here-session") + + proxy.routing.reap(func(sess string) bool { return sess != "gone-session" }) + + assert.ElementsMatch(t, []string{"still-here-session"}, proxy.routing.subscribersOf("uri1")) +} + +// TestPostSSE_ConcurrentLastUnsubscribeAndFirstSubscribe_SameURIAreOrdered is +// the FIX 1 regression from #5744's review: interceptSessionScopedRequest's +// ref-count decision and the upstream forward it triggers were not atomic +// with each other across sessions, so a concurrent last-unsubscribe (sessA) +// and first-subscribe (sessB) for the SAME uri could reach the backend out of +// order, leaving the backend unsubscribed while the routing table said sessB +// was subscribed -- sessB would then silently stop receiving +// resources/updated. handlePost now serializes both the ref-count decision +// and the upstream forward per-uri via uriLocks (see +// resourceSubscriptionURI), so the backend must observe sessA's unsubscribe +// and sessB's subscribe in an order consistent with the FINAL ref count +// (subscribed, by sessB alone) -- and sessB must actually receive a +// subsequent resources/updated for the uri. +// +// This proves ordering deterministically, not via a timing race: the backend +// gates its response to sessA's unsubscribe until the test explicitly +// releases it (releaseUnsub), and sessB's subscribe request is only issued +// after the test has confirmed -- via a channel, never a sleep -- that the +// backend has already received sessA's unsubscribe, i.e. that sessA's +// handlePost is inside its uriLocks critical section, blocked waiting for +// that gated backend response. With the fix, sessB's handlePost cannot send +// ANYTHING to the backend until sessA's entire critical section (decision, +// forward, backend response, HTTP response write) has completed and released +// uriLocks for the uri: that is a structural guarantee of the mutex, not a +// probabilistic race, so the backend's observed order is deterministic +// regardless of goroutine scheduling. Run with -race: the property under +// test is ordering, and a data race would indicate the lock is not doing its +// job. +func TestPostSSE_ConcurrentLastUnsubscribeAndFirstSubscribe_SameURIAreOrdered(t *testing.T) { + t.Parallel() + + const uri = "shared-uri" + + port := pickFreePort(t) + proxy, ctx := startProxyOnly(t, port) + + var mu sync.Mutex + var order []string + + unsubReceived := make(chan struct{}) + releaseUnsub := make(chan struct{}) + + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok || !req.ID.IsValid() { + continue + } + switch req.Method { + case "resources/unsubscribe": + mu.Lock() + order = append(order, "unsubscribe") + mu.Unlock() + close(unsubReceived) + <-releaseUnsub // hold the response until the test releases it + case "resources/subscribe": + mu.Lock() + order = append(order, "subscribe") + mu.Unlock() + } + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + case <-ctx.Done(): + return + } + } + }() + + sessA := establishSession(t, port) + sessB := establishSession(t, port) + + // sessB's GET stream is opened up front so it can observe the eventual + // resources/updated once it becomes uri's subscriber. + getB, readerB := openGETStream(ctx, t, port, sessB) + t.Cleanup(func() { _ = getB.Body.Close() }) + require.Eventually(t, func() bool { return proxy.serverStreams.streamCount() == 1 }, + time.Second, 5*time.Millisecond, "sessB's GET stream was not registered") + + // Precondition: sessA is the sole (first) subscriber of uri. + subA, subABody := postJSON(t, port, sessA, + fmt.Sprintf(`{"jsonrpc":"2.0","id":"sub-a","method":"resources/subscribe","params":{"uri":%q}}`, uri)) + require.Equal(t, http.StatusOK, subA.StatusCode, "sessA's initial subscribe must succeed: %s", string(subABody)) + require.ElementsMatch(t, []string{sessA}, proxy.routing.subscribersOf(uri)) + + mu.Lock() + order = nil // only the concurrent unsubscribe/subscribe pair below is under test + mu.Unlock() + + // sessA's last-unsubscribe fires concurrently; it will block in the + // backend's gate (above) until releaseUnsub is closed below. + unsubDone := make(chan struct{}) + go func() { + defer close(unsubDone) + resp, body := postJSON(t, port, sessA, + fmt.Sprintf(`{"jsonrpc":"2.0","id":"unsub-a","method":"resources/unsubscribe","params":{"uri":%q}}`, uri)) + assert.Equal(t, http.StatusOK, resp.StatusCode, "sessA's unsubscribe must succeed: %s", string(body)) + }() + + select { + case <-unsubReceived: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for the backend to receive sessA's unsubscribe") + } + + // sessB's first-subscribe is issued only now, while sessA's request is + // provably still holding uriLocks for uri (blocked on the gated backend + // response above). With the fix, sessB's handlePost must block acquiring + // the SAME per-uri lock until sessA's request fully completes. + subDone := make(chan struct{}) + go func() { + defer close(subDone) + resp, body := postJSON(t, port, sessB, + fmt.Sprintf(`{"jsonrpc":"2.0","id":"sub-b","method":"resources/subscribe","params":{"uri":%q}}`, uri)) + assert.Equal(t, http.StatusOK, resp.StatusCode, "sessB's subscribe must succeed: %s", string(body)) + }() + + // THE discriminating check: while sessA's unsubscribe is still gated + // (its handlePost provably still holds uriLocks for uri), sessB's + // subscribe decision must not have taken effect yet. Without FIX 1, sessB's + // decision is not ordered by any per-uri lock and would run as soon as + // sessB's request is scheduled -- independent of sessA's still-in-flight + // upstream forward -- so this would very quickly observe sessB already + // recorded as a subscriber, which is exactly the desynchronization the fix + // prevents. With the fix this is a structural guarantee (sessB's + // handlePost is blocked acquiring the same lock sessA holds), so it holds + // for the ENTIRE waitFor window below, not merely at one sampled instant. + require.Never(t, func() bool { + return len(proxy.routing.subscribersOf(uri)) > 0 + }, 200*time.Millisecond, 5*time.Millisecond, + "sessB's subscribe decision must not take effect while sessA's unsubscribe is still gated on the backend") + + close(releaseUnsub) + + select { + case <-unsubDone: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for sessA's unsubscribe request to complete") + } + select { + case <-subDone: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for sessB's subscribe request to complete") + } + + // The backend must have observed the unsubscribe strictly before the + // subscribe: the order matching the FINAL ref count (subscribed, by + // sessB), never the reverse -- the reverse would leave the backend + // unsubscribed while the routing table said sessB was subscribed, the + // #5744 bug this test guards against. + mu.Lock() + assert.Equal(t, []string{"unsubscribe", "subscribe"}, order, + "backend must observe sessA's unsubscribe strictly before sessB's subscribe") + mu.Unlock() + + assert.ElementsMatch(t, []string{sessB}, proxy.routing.subscribersOf(uri), + "routing table must end with sessB as the sole subscriber") + + // sessB must actually receive a subsequent resources/updated for uri, + // proving the backend and the routing table ended up CONSISTENT with each + // other (both "subscribed"), not merely that the routing table says so. + notif, err := jsonrpc2.NewNotification("notifications/resources/updated", map[string]any{"uri": uri}) + require.NoError(t, err) + require.NoError(t, proxy.ForwardResponseToClients(ctx, notif)) + assert.Contains(t, readSSEData(t, readerB), uri) +} diff --git a/pkg/transport/proxy/streamable/dispatcher_streams.go b/pkg/transport/proxy/streamable/dispatcher_streams.go new file mode 100644 index 0000000000..877b71811d --- /dev/null +++ b/pkg/transport/proxy/streamable/dispatcher_streams.go @@ -0,0 +1,246 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "log/slog" + "sync" + + "golang.org/x/exp/jsonrpc2" +) + +// serverStreamBufferSize is the buffer size for each registered stream's data +// channel. +const serverStreamBufferSize = 100 + +// serverStream is a single MCP session's standalone server->client stream. +// data delivers unsolicited notifications to the stream's consumer (the GET +// handler in streamable_proxy.go) and is NEVER closed -- eviction and +// shutdown are signaled exclusively via stop, so a concurrent dispatch can +// never race a close of data (see registry method docs for the panic this +// avoids). +type serverStream struct { + // data is buffered (serverStreamBufferSize) and never closed. + data chan jsonrpc2.Message + // stop is closed exactly once, either by an evicting register call or by + // closeAll, to signal the consumer to stop reading data and return. + stop chan struct{} +} + +// serverStreamRegistry is a transport-agnostic registry of live server->client +// streams, used to deliver unsolicited container notifications to the +// standalone stream for each MCP session. Per the MCP spec, a server MUST NOT +// deliver the same notification more than once per session: this registry +// enforces that by keying on session ID and allowing at most one stream per +// session. It has no knowledge of HTTP: callers (the GET handler in +// streamable_proxy.go) own translating a registered stream's messages into +// wire frames. +// +// Forward-compatibility note: the 2026-07-28 (Modern) revision's +// "subscriptions/listen" method (see mcp.MCPVersionModern in +// pkg/mcp/revision.go) is a reasonable candidate to fan out through this SAME +// decoupled registry rather than inventing a second mechanism, but it is NOT +// a drop-in reuse. Per the draft spec, a future Modern handler will need real +// adaptation on top of what is here today: "subscriptions/listen" is a +// long-lived POST (not a GET), Modern has no sessions so streams must be +// re-keyed per-subscription-id rather than per session ID, delivery requires +// per-notification-type AND per-URI opt-in filtering (not blanket fan-out), +// an initial notifications/subscriptions/acknowledged must be sent, and +// deliveries must be tagged with io.modelcontextprotocol/subscriptionId. Keep +// this type free of HTTP/GET-specific assumptions so that adaptation, when it +// happens, does not also require a rewrite of the fan-out primitives +// themselves. +// +// All fields are guarded by mu; do not add another synchronization primitive +// over streams (see the "one synchronization primitive per data structure" +// style rule). The only channel operation performed while mu is held is a +// non-blocking select-send in broadcast, dispatchTo, and deliverToMany, which +// cannot deadlock. +type serverStreamRegistry struct { + mu sync.Mutex + streams map[string]*serverStream +} + +// newServerStreamRegistry creates an empty serverStreamRegistry. +func newServerStreamRegistry() *serverStreamRegistry { + return &serverStreamRegistry{streams: make(map[string]*serverStream)} +} + +// register creates and stores a new stream for the given MCP session ID, +// returning it for the caller to consume from (data) and select on (stop). +// If a stream is already registered for sessionID (e.g. a second concurrent +// GET for the same session), the prior stream is evicted: its stop channel is +// closed so its consumer goroutine observes the eviction and returns, and the +// new stream replaces it in the map (last-writer-wins, one stream per +// session). The caller MUST call deregister(sessionID, s) when the stream +// ends (client disconnect, proxy shutdown) to avoid leaking the map entry. +func (r *serverStreamRegistry) register(sessionID string) *serverStream { + r.mu.Lock() + defer r.mu.Unlock() + + if old, ok := r.streams[sessionID]; ok { + close(old.stop) + } + + s := &serverStream{ + data: make(chan jsonrpc2.Message, serverStreamBufferSize), + stop: make(chan struct{}), + } + r.streams[sessionID] = s + return s +} + +// deregister removes the stream for sessionID, but only if s is still the +// currently registered stream for that session (identity check). This makes +// a late deregister from an already-evicted handler (see register) a correct +// no-op instead of deleting the newer stream that replaced it. It never +// closes any channel: data is never closed (see serverStream's doc comment), +// and stop is closed exactly once by whichever of register (eviction) or +// closeAll observes it first. +func (r *serverStreamRegistry) deregister(sessionID string, s *serverStream) { + r.mu.Lock() + defer r.mu.Unlock() + + if r.streams[sessionID] == s { + delete(r.streams, sessionID) + } +} + +// broadcast delivers msg once to EVERY currently registered session's stream. +// It is used only for GLOBAL notifications (see the */list_changed handling +// in dispatcher.go) that describe a server-wide capability change with no +// session-specific payload, so fanning out one copy to every session is +// correct and safe. +// +// Delivery is non-blocking and best-effort per the locked product decision: a +// stream whose data channel is full has a slow or stuck consumer, and +// dropping the message for that stream is preferable to blocking delivery to +// every other session's stream. Because data is never closed, this select can +// never panic on a send to a closed channel, even if a concurrent deregister +// or closeAll is racing this call. +func (r *serverStreamRegistry) broadcast(msg jsonrpc2.Message) { + r.mu.Lock() + defer r.mu.Unlock() + + for sessionID, s := range r.streams { + select { + case s.data <- msg: + default: + slog.Warn("server stream channel full; dropping notification", "session_id", sessionID) + } + } +} + +// dispatchTo delivers msg to the single stream registered for routeKey (a +// session ID), if one is currently connected. Delivery is non-blocking and +// best-effort, matching broadcast: an absent stream (no GET SSE connection for +// that session right now) or a full one both simply drop msg, logged at +// Debug/Warn respectively, rather than blocking or erroring. Callers use this +// for messages that are correlated to exactly one session (e.g. a +// resources/updated notification's subscribers) and MUST NOT ever reach a +// different session's stream. +func (r *serverStreamRegistry) dispatchTo(routeKey string, msg jsonrpc2.Message) { + r.mu.Lock() + defer r.mu.Unlock() + + s, ok := r.streams[routeKey] + if !ok { + slog.Debug("no server stream registered for route key; dropping notification", "route_key", routeKey) + return + } + select { + case s.data <- msg: + default: + slog.Warn("server stream channel full; dropping notification", "route_key", routeKey) + } +} + +// deliverToMany delivers msg to each of routeKeys' streams that is currently +// registered, taking mu once for the whole batch rather than once per key. +// Delivery to each present stream is non-blocking and best-effort, matching +// broadcast/dispatchTo. Callers use this for messages correlated to a known, +// bounded set of sessions (e.g. a resource's subscribers) -- routeKeys not +// present in the registry are silently skipped (no connected GET stream right +// now), not an error. +func (r *serverStreamRegistry) deliverToMany(routeKeys []string, msg jsonrpc2.Message) { + r.mu.Lock() + defer r.mu.Unlock() + + for _, routeKey := range routeKeys { + s, ok := r.streams[routeKey] + if !ok { + slog.Debug("no server stream registered for route key; dropping notification", "route_key", routeKey) + continue + } + select { + case s.data <- msg: + default: + slog.Warn("server stream channel full; dropping notification", "route_key", routeKey) + } + } +} + +// streamCount reports how many sessions currently have a registered stream. +func (r *serverStreamRegistry) streamCount() int { + r.mu.Lock() + defer r.mu.Unlock() + + return len(r.streams) +} + +// closeAll signals every currently registered stream to stop (by closing its +// stop channel) and resets the registry to empty. It is used on proxy Stop so +// that in-flight GET handlers observe a closed per-stream stop signal and +// return, in addition to (not instead of) selecting on shutdownCh. +func (r *serverStreamRegistry) closeAll() { + r.mu.Lock() + defer r.mu.Unlock() + + for _, s := range r.streams { + close(s.stop) + } + r.streams = make(map[string]*serverStream) +} + +// closeStream signals the single stream currently registered for routeKey (a +// session ID) to stop, mirroring closeAll's single-close discipline but +// scoped to one session: it closes stop (never data -- see serverStream's +// doc comment) exactly once and removes routeKey's entry, so the stream's +// consumer (handleGet) observes the closed stop channel and returns instead +// of surviving until its underlying socket happens to close on its own (see +// handleDelete and reapServerStreams, and FIX 2 in #5744's review). Unlike +// deregister, this does not take an identity parameter: callers here (a +// session teardown path) always want to evict WHATEVER stream is currently +// registered for routeKey, regardless of which register call produced it -- +// there is no "stale handler" case to guard against as there is in +// register's eviction path. It reports whether a stream was found and +// closed, so callers can distinguish "nothing to tear down" from "torn down". +func (r *serverStreamRegistry) closeStream(routeKey string) bool { + r.mu.Lock() + defer r.mu.Unlock() + + s, ok := r.streams[routeKey] + if !ok { + return false + } + close(s.stop) + delete(r.streams, routeKey) + return true +} + +// streamKeys returns a snapshot of every session ID that currently has a +// registered stream. Used by reapServerStreams to find candidate streams +// whose owning session may no longer be active, without holding mu while +// resolving each candidate's liveness (see reapServerStreams's doc comment +// for why that matters). +func (r *serverStreamRegistry) streamKeys() []string { + r.mu.Lock() + defer r.mu.Unlock() + + keys := make([]string, 0, len(r.streams)) + for k := range r.streams { + keys = append(keys, k) + } + return keys +} diff --git a/pkg/transport/proxy/streamable/dispatcher_streams_test.go b/pkg/transport/proxy/streamable/dispatcher_streams_test.go new file mode 100644 index 0000000000..7aa60d6683 --- /dev/null +++ b/pkg/transport/proxy/streamable/dispatcher_streams_test.go @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "strconv" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" +) + +func newTestNotification(t *testing.T, method string) jsonrpc2.Message { + t.Helper() + msg, err := jsonrpc2.NewNotification(method, map[string]any{"ok": true}) + require.NoError(t, err) + return msg +} + +// recvWithTimeout receives a single message from ch or fails the test after a +// bounded wait, per the "always add timeouts to blocking barriers" testing rule. +func recvWithTimeout(t *testing.T, ch <-chan jsonrpc2.Message) jsonrpc2.Message { + t.Helper() + select { + case msg := <-ch: + return msg + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for dispatched message") + return nil + } +} + +// assertNoDelivery proves ch does not receive a message within a short bounded +// wait. It is intentionally shorter than recvWithTimeout's timeout since it is +// expected to always run to completion (there is nothing to wait on). +func assertNoDelivery(t *testing.T, ch <-chan jsonrpc2.Message) { + t.Helper() + select { + case msg := <-ch: + t.Fatalf("expected no delivery, but received: %v", msg) + case <-time.After(100 * time.Millisecond): + } +} + +func TestServerStreamRegistry_RegisterDispatchDeregister(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + assert.Equal(t, 0, r.streamCount(), "registry should start empty") + + s := r.register("session-1") + assert.Equal(t, 1, r.streamCount()) + + msg := newTestNotification(t, "notifications/progress") + r.broadcast(msg) + + got := recvWithTimeout(t, s.data) + assert.Equal(t, msg, got) + + r.deregister("session-1", s) + assert.Equal(t, 0, r.streamCount()) + + // A dispatch after deregister must not reach the evicted stream's data + // channel (the channel itself is never closed -- see serverStream's doc). + r.broadcast(newTestNotification(t, "notifications/progress")) + assertNoDelivery(t, s.data) +} + +func TestServerStreamRegistry_PerSessionSingleDelivery(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + s1 := r.register("session-1") + s2 := r.register("session-2") + t.Cleanup(func() { + r.deregister("session-1", s1) + r.deregister("session-2", s2) + }) + + msg := newTestNotification(t, "notifications/resources/updated") + r.broadcast(msg) + + assert.Equal(t, msg, recvWithTimeout(t, s1.data)) + assert.Equal(t, msg, recvWithTimeout(t, s2.data)) + + // Each session must receive exactly one copy -- no duplicate second delivery. + assertNoDelivery(t, s1.data) + assertNoDelivery(t, s2.data) +} + +func TestServerStreamRegistry_OneStreamPerSessionEvicts(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + first := r.register("session-1") + + second := r.register("session-1") + t.Cleanup(func() { r.deregister("session-1", second) }) + + assert.Equal(t, 1, r.streamCount(), "second register must evict, not add, an entry") + + // The first stream's stop channel must be closed to signal eviction. + select { + case _, ok := <-first.stop: + assert.False(t, ok, "first stream's stop channel must be closed") + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for evicted stream's stop channel to close") + } + + msg := newTestNotification(t, "notifications/progress") + r.broadcast(msg) + + // Only the second (current) stream receives the dispatched message. + assert.Equal(t, msg, recvWithTimeout(t, second.data)) + assertNoDelivery(t, first.data) +} + +// TestServerStreamRegistry_FullChannelDropsWithoutBlocking proves dispatch +// never blocks on a full consumer, and that a slow stream does not prevent +// delivery to other, healthy streams. +func TestServerStreamRegistry_FullChannelDropsWithoutBlocking(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + full := r.register("full") + healthy := r.register("healthy") + t.Cleanup(func() { + r.deregister("full", full) + r.deregister("healthy", healthy) + }) + + // Saturate the "full" stream's buffer. + for i := 0; i < serverStreamBufferSize; i++ { + r.broadcast(newTestNotification(t, "notifications/progress")) + } + // Drain "healthy" of the same backlog so it isn't itself full for the + // assertion below. + for i := 0; i < serverStreamBufferSize; i++ { + recvWithTimeout(t, healthy.data) + } + + done := make(chan struct{}) + go func() { + defer close(done) + r.broadcast(newTestNotification(t, "notifications/message")) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("dispatch blocked on a full channel instead of dropping") + } + + // The healthy stream still receives the new message even though the + // other stream's buffer was full. + msg := recvWithTimeout(t, healthy.data) + assert.NotNil(t, msg) +} + +func TestServerStreamRegistry_DeregisterIsIdempotentAndCloseAllClears(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + sa := r.register("a") + r.register("b") + require.Equal(t, 2, r.streamCount()) + + r.closeAll() + assert.Equal(t, 0, r.streamCount()) + + // Deregistering an already-removed stream is a no-op, not a panic. + assert.NotPanics(t, func() { r.deregister("a", sa) }) +} + +// TestServerStreamRegistry_DispatchAfterDeregisterIsNoop proves a dispatch that +// races a deregister neither panics nor blocks: the removed stream simply +// does not receive the message. +func TestServerStreamRegistry_DispatchAfterDeregisterIsNoop(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + s := r.register("gone") + r.deregister("gone", s) + + assert.NotPanics(t, func() { + r.broadcast(newTestNotification(t, "notifications/progress")) + }) + assert.Equal(t, 0, r.streamCount()) +} + +// TestServerStreamRegistry_ConcurrentDispatchAndEvict is the Blocker-1 +// regression test: it stresses concurrent dispatch (which sends on +// serverStream.data) against concurrent register-on-same-session (which +// evicts, closing the evicted stream's stop channel), deregister, and a +// drainer reading data for each stream. Run with -race; the invariant under +// test is that data is never closed, so dispatch's send can never race a +// close and panic. +func TestServerStreamRegistry_ConcurrentDispatchAndEvict(t *testing.T) { + t.Parallel() + + const ( + sessionID = "stress-session" + numWorkers = 8 + numRounds = 500 + waitTimeout = 20 * time.Second + drainTimeout = 5 * time.Second + ) + + r := newServerStreamRegistry() + + var workersWG sync.WaitGroup + var drainersWG sync.WaitGroup + testDone := make(chan struct{}) + + // Dispatcher goroutines: hammer dispatch concurrently. + for i := 0; i < numWorkers; i++ { + workersWG.Add(1) + go func() { + defer workersWG.Done() + for j := 0; j < numRounds; j++ { + r.broadcast(newTestNotification(t, "notifications/progress")) + } + }() + } + + // Register/evict/deregister goroutines: force eviction of the same session + // repeatedly while dispatch is hammering it. Each newly registered stream + // gets its own drainer so dispatch's non-blocking send has somewhere to go + // in the common case, though correctness (no panic, no block) must hold + // regardless of whether anything is draining. + for i := 0; i < numWorkers; i++ { + workersWG.Add(1) + go func() { + defer workersWG.Done() + for j := 0; j < numRounds; j++ { + s := r.register(sessionID) + + drainersWG.Add(1) + go func(s *serverStream) { + defer drainersWG.Done() + for { + select { + case <-s.data: + case <-s.stop: + return + case <-testDone: + return + } + } + }(s) + + r.deregister(sessionID, s) + } + }() + } + + allWorkersDone := make(chan struct{}) + go func() { + workersWG.Wait() + close(allWorkersDone) + }() + + select { + case <-allWorkersDone: + case <-time.After(waitTimeout): + close(testDone) + t.Fatal("concurrent dispatch/evict/deregister stress test timed out") + } + + // Signal any drainer whose stream was never evicted (i.e. it was the last + // one registered) to stop, then bound how long we wait for the drainers to + // exit cleanly. + close(testDone) + allDrainersDone := make(chan struct{}) + go func() { + drainersWG.Wait() + close(allDrainersDone) + }() + select { + case <-allDrainersDone: + case <-time.After(drainTimeout): + t.Fatal("drainer goroutines did not exit after testDone was closed") + } + + // At most one stream (the last registered) should remain. + assert.LessOrEqual(t, r.streamCount(), 1) +} + +func TestServerStreamRegistry_ManySessionsStreamCount(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + streams := make([]*serverStream, 0, 10) + for i := 0; i < 10; i++ { + streams = append(streams, r.register("session-"+strconv.Itoa(i))) + } + t.Cleanup(func() { + for i, s := range streams { + r.deregister("session-"+strconv.Itoa(i), s) + } + }) + + assert.Equal(t, 10, r.streamCount()) +} + +// TestServerStreamRegistry_DispatchToHitsOnlyTarget verifies dispatchTo +// delivers a message to the single stream registered for its route key, and +// no other currently-registered stream receives it -- the property routeProgress +// and routeResourceUpdated depend on to avoid cross-session leakage. +func TestServerStreamRegistry_DispatchToHitsOnlyTarget(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + target := r.register("target-session") + other := r.register("other-session") + t.Cleanup(func() { + r.deregister("target-session", target) + r.deregister("other-session", other) + }) + + msg := newTestNotification(t, "notifications/progress") + r.dispatchTo("target-session", msg) + + assert.Equal(t, msg, recvWithTimeout(t, target.data)) + assertNoDelivery(t, other.data) +} + +// TestServerStreamRegistry_DispatchToUnknownKeyIsNoop verifies dispatchTo for +// a route key with no registered stream neither panics nor blocks. +func TestServerStreamRegistry_DispatchToUnknownKeyIsNoop(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + assert.NotPanics(t, func() { + r.dispatchTo("no-such-session", newTestNotification(t, "notifications/progress")) + }) +} + +// TestServerStreamRegistry_DeliverToMany verifies deliverToMany delivers a +// message to every present stream among routeKeys, and does not deliver it to +// a registered stream whose session is NOT in routeKeys. +func TestServerStreamRegistry_DeliverToMany(t *testing.T) { + t.Parallel() + + r := newServerStreamRegistry() + subA := r.register("subscriber-a") + subB := r.register("subscriber-b") + notSub := r.register("not-a-subscriber") + t.Cleanup(func() { + r.deregister("subscriber-a", subA) + r.deregister("subscriber-b", subB) + r.deregister("not-a-subscriber", notSub) + }) + + msg := newTestNotification(t, "notifications/resources/updated") + r.deliverToMany([]string{"subscriber-a", "subscriber-b", "unknown-session"}, msg) + + assert.Equal(t, msg, recvWithTimeout(t, subA.data)) + assert.Equal(t, msg, recvWithTimeout(t, subB.data)) + assertNoDelivery(t, notSub.data) +} diff --git a/pkg/transport/proxy/streamable/dispatcher_test.go b/pkg/transport/proxy/streamable/dispatcher_test.go new file mode 100644 index 0000000000..a205111774 --- /dev/null +++ b/pkg/transport/proxy/streamable/dispatcher_test.go @@ -0,0 +1,240 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" +) + +// newDispatcherTestProxy creates an HTTPProxy with dispatchResponses running in +// the background, without starting the HTTP server. t.Cleanup stops it. +func newDispatcherTestProxy(t *testing.T) *HTTPProxy { + t.Helper() + proxy := NewHTTPProxy("localhost", 0, nil, nil) + go proxy.dispatchResponses() + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = proxy.Stop(stopCtx) + }) + return proxy +} + +// TestDispatchResponses_ResponseRoutesToWaiter verifies a *jsonrpc2.Response +// carrying a known composite ID is delivered to its registered waiter. +func TestDispatchResponses_ResponseRoutesToWaiter(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + waitCh, cleanup := proxy.createWaiter("sess-1", jsonrpc2.StringID("req-1")) + defer cleanup() + ck := compositeKey("sess-1", idKeyFromID(jsonrpc2.StringID("req-1"))) + + resp, err := jsonrpc2.NewResponse(jsonrpc2.StringID(ck), "result", nil) + require.NoError(t, err) + + proxy.responseCh <- resp + + select { + case got := <-waitCh: + assert.Equal(t, resp, got) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for response to reach waiter") + } +} + +// TestDispatchResponses_UnknownCompositeKeyResponseIsDropped verifies a +// *jsonrpc2.Response whose composite ID has no registered waiter is dropped +// without blocking or panicking. +func TestDispatchResponses_UnknownCompositeKeyResponseIsDropped(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + resp, err := jsonrpc2.NewResponse(jsonrpc2.StringID("no-such-waiter|nil"), "result", nil) + require.NoError(t, err) + + done := make(chan struct{}) + go func() { + defer close(done) + proxy.responseCh <- resp + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("timed out sending response with unknown composite key") + } + + // Give the dispatcher goroutine a moment to process; there is nothing to + // assert beyond "did not panic/block", since there is no waiter to observe. + time.Sleep(50 * time.Millisecond) +} + +// TestDispatchResponses_ListChangedBroadcastsToEverySession verifies every +// */list_changed notification (listChangedNotificationMethods) is broadcast +// to every connected session's standalone GET stream. +func TestDispatchResponses_ListChangedBroadcastsToEverySession(t *testing.T) { + t.Parallel() + + methods := []string{ + "notifications/tools/list_changed", + "notifications/resources/list_changed", + "notifications/prompts/list_changed", + } + + for _, method := range methods { + t.Run(method, func(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + s1 := proxy.serverStreams.register("session-1") + s2 := proxy.serverStreams.register("session-2") + t.Cleanup(func() { + proxy.serverStreams.deregister("session-1", s1) + proxy.serverStreams.deregister("session-2", s2) + }) + + notification, err := jsonrpc2.NewNotification(method, nil) + require.NoError(t, err) + proxy.responseCh <- notification + + assert.Equal(t, notification, recvWithTimeout(t, s1.data)) + assert.Equal(t, notification, recvWithTimeout(t, s2.data)) + }) + } +} + +// TestDispatchResponses_ProgressDeliveredOnlyToRecordedRoute verifies a +// notifications/progress message is delivered ONLY to the progress route +// recorded for its progressToken (with the client's original token restored), +// and that neither a different token nor no recorded route at all delivers +// anything -- proving progress can never leak across sessions/requests. +func TestDispatchResponses_ProgressDeliveredOnlyToRecordedRoute(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + deliverA := make(chan jsonrpc2.Message, 1) + proxy.routing.recordProgressToken("ptGlobal-A", progressRoute{deliver: deliverA, originalToken: "client-token-A"}) + t.Cleanup(func() { proxy.routing.dropProgressToken("ptGlobal-A") }) + + // Unregistered token: must not be delivered anywhere. + progressUnknown, err := jsonrpc2.NewNotification("notifications/progress", + map[string]any{"progressToken": "ptGlobal-unknown", "progress": 1}) + require.NoError(t, err) + proxy.responseCh <- progressUnknown + + // Recorded token: must be delivered on deliverA, with the ORIGINAL client + // token restored (not the internal ptGlobal-A). + progressA, err := jsonrpc2.NewNotification("notifications/progress", + map[string]any{"progressToken": "ptGlobal-A", "progress": 42}) + require.NoError(t, err) + proxy.responseCh <- progressA + + got := recvWithTimeout(t, deliverA) + gotReq, ok := got.(*jsonrpc2.Request) + require.True(t, ok) + token, ok := extractStringParam(gotReq.Params, "progressToken") + require.True(t, ok) + assert.Equal(t, "client-token-A", token, "delivered notification must carry the client's original token, not ptGlobal") + + // Prove the unknown-token notification was never delivered (it would have + // arrived first, before progressA, if it had been). + assertNoDelivery(t, deliverA) +} + +// TestDispatchResponses_ResourcesUpdatedRoutesToSubscribersOnly verifies a +// notifications/resources/updated message reaches only sessions subscribed to +// its uri, delivered on their standalone GET stream -- and NOT a connected +// session that never subscribed. +func TestDispatchResponses_ResourcesUpdatedRoutesToSubscribersOnly(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + subscriber := proxy.serverStreams.register("subscriber-session") + nonSubscriber := proxy.serverStreams.register("non-subscriber-session") + t.Cleanup(func() { + proxy.serverStreams.deregister("subscriber-session", subscriber) + proxy.serverStreams.deregister("non-subscriber-session", nonSubscriber) + }) + + proxy.routing.addSubscription("file:///watched.txt", "subscriber-session") + + notification, err := jsonrpc2.NewNotification("notifications/resources/updated", + map[string]any{"uri": "file:///watched.txt"}) + require.NoError(t, err) + proxy.responseCh <- notification + + assert.Equal(t, notification, recvWithTimeout(t, subscriber.data)) + assertNoDelivery(t, nonSubscriber.data) +} + +// TestDispatchResponses_LoggingMessageIsDropped verifies notifications/message +// (logging) is dropped rather than delivered to any connected session's +// stream: the shared backend cannot attribute log content to a session, so +// forwarding it (to any session) would be a cross-session leak. +func TestDispatchResponses_LoggingMessageIsDropped(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + s := proxy.serverStreams.register("session-1") + t.Cleanup(func() { proxy.serverStreams.deregister("session-1", s) }) + + logMsg, err := jsonrpc2.NewNotification("notifications/message", + map[string]any{"level": "info", "data": "secret log line"}) + require.NoError(t, err) + proxy.responseCh <- logMsg + + // Race the dropped notification against a sentinel GLOBAL notification: if + // the log message had been delivered, it would arrive on s.data first. + sentinel, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + proxy.responseCh <- sentinel + + got := recvWithTimeout(t, s.data) + assert.Equal(t, sentinel, got, "the dropped logging notification must not appear on any stream") +} + +// TestDispatchResponses_ServerRequestRejectedToBackend verifies a +// server->client REQUEST with a valid ID (sampling/elicitation) is never +// delivered to any client stream, and instead gets a JSON-RPC error response +// written back to the BACKEND (via SendMessageToDestination/messageCh) so the +// backend's blocking call unblocks. +func TestDispatchResponses_ServerRequestRejectedToBackend(t *testing.T) { + t.Parallel() + proxy := newDispatcherTestProxy(t) + + s := proxy.serverStreams.register("session-1") + t.Cleanup(func() { proxy.serverStreams.deregister("session-1", s) }) + + req, err := jsonrpc2.NewCall(jsonrpc2.StringID("sampling-1"), "sampling/createMessage", nil) + require.NoError(t, err) + proxy.responseCh <- req + + // The error response is written back to the backend on messageCh. + select { + case msg := <-proxy.GetMessageChannel(): + resp, ok := msg.(*jsonrpc2.Response) + require.True(t, ok, "expected a *jsonrpc2.Response written back to the backend") + assert.Equal(t, "sampling-1", resp.ID.Raw()) + require.Error(t, resp.Error) + assert.Contains(t, resp.Error.Error(), "sampling/createMessage") + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for error response to reach the backend") + } + + // Prove it did NOT reach any client stream by racing it against a GLOBAL + // notification sent right after. + sentinel, err := jsonrpc2.NewNotification("notifications/tools/list_changed", nil) + require.NoError(t, err) + proxy.responseCh <- sentinel + + got := recvWithTimeout(t, s.data) + assert.Equal(t, sentinel, got, "the rejected server->client request must not appear on any client stream") +} diff --git a/pkg/transport/proxy/streamable/initialize_cache.go b/pkg/transport/proxy/streamable/initialize_cache.go new file mode 100644 index 0000000000..654c22e07d --- /dev/null +++ b/pkg/transport/proxy/streamable/initialize_cache.go @@ -0,0 +1,180 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "context" + "encoding/json" + "errors" + "sync/atomic" + + "golang.org/x/exp/jsonrpc2" + "golang.org/x/sync/singleflight" + + sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// errUnexpectedInitializeReply is returned when the backend answers initialize +// with something that is not a *jsonrpc2.Response, so it cannot be re-addressed +// to the calling client. +var errUnexpectedInitializeReply = errors.New("backend returned an unexpected reply to initialize") + +// initializeCache holds the InitializeResult the backend produced for the +// first client handshake, so later handshakes are answered locally instead of +// reaching the backend again. +// +// A stdio MCP server is one process and therefore exactly one MCP session, +// while this proxy fans an arbitrary number of HTTP client sessions onto it. +// Forwarding every client's initialize to that single session means only the +// first can succeed: go-sdk v1.7+ rejects a second handshake on an +// already-initialized session with `duplicate "initialize" received`, which +// takes down every client after the first -- and, behind vMCP, every health +// check and tool call as well (see #5890, #1982). +// +// Answering from a cache restores the behaviour go-sdk v1.6.1 and earlier had, +// where a repeated initialize returned the original result, but places it in +// the component that actually owns the multiplexing rather than asking every +// MCP server to tolerate duplicate handshakes. +// +// The cached value is process-scoped rather than session-scoped: an +// InitializeResult carries the backend's protocol version, capabilities, +// instructions and server info, none of which vary by client. +type initializeCache struct { + // flight collapses concurrent handshakes into a single upstream call whose + // outcome every waiter shares. Serializing them behind a mutex instead + // would queue each waiter's own round trip: with an unresponsive backend + // the Nth client waits N*requestTimeout (60s each) before it may even + // start. + flight singleflight.Group + + // result is the raw InitializeResult of the first successful handshake, + // nil until one completes. + // + // UNSYNCHRONISED BY DESIGN: only the flight callback in interceptInitialize + // may read or write this field. singleflight runs at most one callback per + // key at a time and a later callback starts only after the previous one has + // completed through the group's own mutex, so callback-only access can + // never overlap. Reading it anywhere else -- a status endpoint, a helper + // accessor -- reintroduces a data race, which is why no accessor exists. + // Once assigned it is never mutated, so the same bytes may be handed to + // concurrent responses. + result json.RawMessage + + // initializedForwarded records whether a notifications/initialized has + // already been passed to the backend. This deliberately does NOT use the + // flight: singleflight de-duplicates rather than serializes, so two + // concurrent claims would share one callback's answer and both forward. + // CompareAndSwap gives exactly one winner. + initializedForwarded atomic.Bool +} + +// handshake is the id-free half of an initialize outcome, shared by a flight +// with all of its waiters. Exactly one of the fields is set. Sharing this +// rather than a whole *jsonrpc2.Response means each caller builds its own +// reply, so a waiter cannot receive the forwarding caller's JSON-RPC id. +type handshake struct { + result json.RawMessage + rpcErr error +} + +// initializeFlightKey is the sole singleflight key: the cached handshake is +// process-scoped, so all callers share one flight. +const initializeFlightKey = "initialize" + +// newInitializeCache creates an empty initializeCache. +func newInitializeCache() *initializeCache { + return &initializeCache{} +} + +// interceptInitialize answers a client's initialize request: the first caller's +// request is forwarded to the backend and its result cached, and every later +// caller is served that cached result without an upstream round trip. +// +// Only a successful handshake is cached. A transport failure or a JSON-RPC +// error from the backend leaves the cache empty so the next client retries the +// handshake for real, rather than pinning every future client to one bad +// response. +// +// Concurrent first handshakes are collapsed by singleflight, so exactly one +// reaches the backend and the rest share its outcome. The upstream call runs on +// a context detached from the caller's request (context.WithoutCancel): the +// result belongs to every waiter, so the first caller disconnecting must not +// cancel the handshake for the others. forwardUpstream still applies +// p.requestTimeout to that detached context, so it cannot run unbounded. +// +// Every response is rebuilt with the calling request's own JSON-RPC id. A +// shared flight result carries the id of whichever caller performed the +// forward, and returning it verbatim would hand later waiters an id they never +// sent. +// +// One client-visible consequence: the protocol version negotiated by the first +// client is replayed to every later client. A client that asked for a +// different version receives the first client's instead, and per the MCP +// lifecycle spec should disconnect if it cannot speak it. This is inherent to +// collapsing N client sessions onto one backend session, and matches what +// go-sdk v1.6.1 and earlier did when replaying a cached result. +func (p *HTTPProxy) interceptInitialize( + ctx context.Context, sessID string, req *jsonrpc2.Request, +) jsonrpc2.Message { + c := p.initialize + v, err, _ := c.flight.Do(initializeFlightKey, func() (any, error) { + // A handshake already succeeded: replay it, no backend round trip. + if c.result != nil { + return handshake{result: c.result}, nil + } + + // Detached from the caller's context: this result belongs to every + // waiter, so whichever caller performs the forward must not cancel the + // handshake for the others by disconnecting. forwardUpstream still + // bounds it with p.requestTimeout. + msg, ferr := p.forwardUpstream(context.WithoutCancel(ctx), sessID, req) + if ferr != nil { + return nil, ferr // nothing cached; the next client retries for real + } + + resp, ok := msg.(*jsonrpc2.Response) + if !ok { + return nil, errUnexpectedInitializeReply + } + if resp.Error != nil { + // Shared with this flight's waiters but deliberately not cached, so + // a transient failure cannot pin every future client to one + // response. + return handshake{rpcErr: resp.Error}, nil + } + if len(resp.Result) == 0 { + return nil, errUnexpectedInitializeReply // no result and no error + } + + c.result = resp.Result + return handshake{result: resp.Result}, nil + }) + if err != nil { + return upstreamErrorResponse(req.ID, err) + } + + h := v.(handshake) // the callback returns a handshake or a non-nil error + return &jsonrpc2.Response{ID: req.ID, Result: h.result, Error: h.rpcErr} +} + +// claimInitializedForward reports whether a client's notifications/initialized +// should be passed to the backend, and records that it has been. +// +// The backend completes a single handshake and so expects a single initialized +// notification; the first client's is forwarded and every later one is +// swallowed. Swallowed notifications are still acknowledged with 202, which is +// all the sending client expects -- notifications carry no response. +func (p *HTTPProxy) claimInitializedForward() bool { + return p.initialize.initializedForwarded.CompareAndSwap(false, true) +} + +// isInitializedNotification reports whether msg is a client's +// notifications/initialized (a JSON-RPC notification, so no id). +func isInitializedNotification(msg jsonrpc2.Message) bool { + req, ok := msg.(*jsonrpc2.Request) + if !ok { + return false + } + return req.ID.Raw() == nil && req.Method == string(sdkmcp.MethodNotificationInitialized) +} diff --git a/pkg/transport/proxy/streamable/initialize_cache_test.go b/pkg/transport/proxy/streamable/initialize_cache_test.go new file mode 100644 index 0000000000..e6127ecf0f --- /dev/null +++ b/pkg/transport/proxy/streamable/initialize_cache_test.go @@ -0,0 +1,385 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" +) + +// fakeStdioBackend stands in for the single stdio MCP session behind the +// proxy: it drains the proxy's outbound message channel, records every method +// that reaches it, and answers id-bearing requests via respond. +type fakeStdioBackend struct { + mu sync.Mutex + methods []string +} + +func (b *fakeStdioBackend) record(method string) { + b.mu.Lock() + defer b.mu.Unlock() + b.methods = append(b.methods, method) +} + +// count returns how many times method reached the backend. +func (b *fakeStdioBackend) count(method string) int { + b.mu.Lock() + defer b.mu.Unlock() + n := 0 + for _, m := range b.methods { + if m == method { + n++ + } + } + return n +} + +// startFakeStdioBackend wires a fakeStdioBackend to proxy and returns it. +// respond is consulted only for id-bearing requests; notifications are +// recorded and dropped, as a real backend would handle them. +func startFakeStdioBackend( + ctx context.Context, proxy *HTTPProxy, respond func(*jsonrpc2.Request) jsonrpc2.Message, +) *fakeStdioBackend { + backend := &fakeStdioBackend{} + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok { + continue + } + backend.record(req.Method) + if !req.ID.IsValid() { + continue + } + if reply := respond(req); reply != nil { + _ = proxy.ForwardResponseToClients(ctx, reply) + } + case <-ctx.Done(): + return + } + } + }() + return backend +} + +// testProtocolVersion is the protocol version the fake backend negotiates. +const testProtocolVersion = "2025-11-25" + +// initializeResultFor builds a realistic InitializeResult response for req. +func initializeResultFor(req *jsonrpc2.Request) jsonrpc2.Message { + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{ + "protocolVersion": testProtocolVersion, + "capabilities": map[string]any{"tools": map[string]any{}}, + "serverInfo": map[string]any{"name": "fake-stdio-server", "version": "1.0.0"}, + }, nil) + return resp +} + +// strictInitializeResponder returns a backend responder that reproduces +// go-sdk v1.7+ server-session semantics: the first initialize is answered with +// an InitializeResult and every later one is rejected with the same +// `duplicate "initialize" received` error ServerSession.initialize returns +// (go-sdk mcp/server.go:1978). go-sdk rejects a repeated +// notifications/initialized the same way (mcp/server.go:1424). +// +// Tests use this rather than a permissive stub so they assert the outcome that +// actually matters -- a later client's handshake succeeds -- instead of only +// counting how many messages reached the backend. A duplicate that escapes the +// cache fails the handshake here exactly as it does against a real server. +func strictInitializeResponder() func(*jsonrpc2.Request) jsonrpc2.Message { + var mu sync.Mutex + var handshakeDone bool + return func(req *jsonrpc2.Request) jsonrpc2.Message { + if req.Method != "initialize" { + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + return resp + } + mu.Lock() + defer mu.Unlock() + if handshakeDone { + resp, _ := jsonrpc2.NewResponse(req.ID, nil, + jsonrpc2.NewError(-32603, `duplicate "initialize" received`)) + return resp + } + handshakeDone = true + return initializeResultFor(req) + } +} + +// startInitTestProxy starts a proxy on a free port and returns it with its +// /mcp endpoint URL. +func startInitTestProxy(t *testing.T) (*HTTPProxy, string) { + t.Helper() + port := getFreePort(t) + proxy := NewHTTPProxy("localhost", port, nil, nil) + require.NoError(t, proxy.Start(t.Context())) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = proxy.Stop(stopCtx) + }) + // Give the listener a moment to accept connections. + time.Sleep(100 * time.Millisecond) + return proxy, fmt.Sprintf("http://localhost:%d%s", port, StreamableHTTPEndpoint) +} + +// postInitialize sends an initialize request and returns the decoded JSON-RPC +// response body. +func postInitialize(t *testing.T, url string, id string) map[string]any { + t.Helper() + body := fmt.Sprintf( + `{"jsonrpc":"2.0","id":%q,"method":"initialize","params":{"protocolVersion":%q}}`, + id, testProtocolVersion) + resp, err := http.Post(url, "application/json", bytes.NewReader([]byte(body))) //nolint:gosec // test-local URL + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + var decoded map[string]any + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + return decoded +} + +// TestEveryClientHandshakeSucceedsAgainstStrictBackend verifies the behaviour +// #5890 reports: with a backend that rejects a repeated initialize, every +// client after the first used to fail. Only the first handshake may reach the +// single stdio session; the rest are answered from the cached InitializeResult +// with their own JSON-RPC id echoed back. +func TestEveryClientHandshakeSucceedsAgainstStrictBackend(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + backend := startFakeStdioBackend(t.Context(), proxy, strictInitializeResponder()) + + var results []any + for i := range 3 { + decoded := postInitialize(t, url, fmt.Sprintf("client-%d", i)) + assert.Equal(t, fmt.Sprintf("client-%d", i), decoded["id"], + "each client must get its own JSON-RPC id back, not the first client's") + assert.Nil(t, decoded["error"], + "no client may see the backend's duplicate-initialize rejection") + require.NotNil(t, decoded["result"], "handshake should succeed") + results = append(results, decoded["result"]) + } + + assert.Equal(t, 1, backend.count("initialize"), + "only the first handshake may reach the single stdio session") + assert.Equal(t, results[0], results[1], "cached result must match the original") + assert.Equal(t, results[0], results[2], "cached result must match the original") +} + +// TestConcurrentInitializeReachesBackendOnce verifies that simultaneous +// handshakes are single-flighted: without the lock held across the first +// upstream round trip, every concurrent client would see an empty cache and +// race a second initialize to the backend. +func TestConcurrentInitializeReachesBackendOnce(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + strict := strictInitializeResponder() + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + // Slow the first handshake so the others pile up behind it. + time.Sleep(150 * time.Millisecond) + return strict(req) + }) + + const clients = 5 + var wg sync.WaitGroup + for i := range clients { + wg.Go(func() { + id := fmt.Sprintf("client-%d", i) + decoded := postInitialize(t, url, id) + assert.Nil(t, decoded["error"], + "a handshake that raced past the cache would be rejected by the backend") + assert.NotNil(t, decoded["result"], "every concurrent handshake should succeed") + // Waiters share one singleflight result, which carries the id of + // whichever caller performed the forward. Each response must be + // re-addressed to its own caller. + assert.Equal(t, id, decoded["id"], + "a waiter must never receive another caller's JSON-RPC id") + }) + } + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent handshakes to complete") + } + + assert.Equal(t, 1, backend.count("initialize"), + "concurrent handshakes must collapse into a single upstream initialize") +} + +// TestFailedInitializeIsNotCached verifies that a backend error is not stored, +// so a transient failure cannot pin every future client to one bad response. +func TestFailedInitializeIsNotCached(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + + var attempts int + var mu sync.Mutex + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + mu.Lock() + attempts++ + first := attempts == 1 + mu.Unlock() + if first { + resp, _ := jsonrpc2.NewResponse(req.ID, nil, jsonrpc2.NewError(-32603, "backend not ready")) + return resp + } + return initializeResultFor(req) + }) + + failed := postInitialize(t, url, "client-0") + assert.NotNil(t, failed["error"], "first handshake should surface the backend error") + + succeeded := postInitialize(t, url, "client-1") + assert.NotNil(t, succeeded["result"], "second handshake should be retried against the backend") + + assert.Equal(t, 2, backend.count("initialize"), + "a failed handshake must not be cached; the next client must retry for real") +} + +// TestDuplicateInitializedNotificationIsNotForwarded verifies that only the +// first notifications/initialized reaches the backend. The backend completed a +// single handshake, so a second initialized notification is out of lifecycle +// for it; later ones are acknowledged locally instead. +func TestDuplicateInitializedNotificationIsNotForwarded(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + backend := startFakeStdioBackend(t.Context(), proxy, strictInitializeResponder()) + + for range 3 { + resp, err := http.Post(url, "application/json", //nolint:gosec // test-local URL + bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"notifications/initialized"}`))) + require.NoError(t, err) + require.Equal(t, http.StatusAccepted, resp.StatusCode, + "every client's initialized notification must still be acknowledged") + _ = resp.Body.Close() + } + + // The backend goroutine records asynchronously; give it a moment to drain. + assert.Eventually(t, func() bool { + return backend.count("notifications/initialized") >= 1 + }, 2*time.Second, 20*time.Millisecond, "the first notification should reach the backend") + + assert.Equal(t, 1, backend.count("notifications/initialized"), + "only one initialized notification may reach the single stdio session") +} + +// TestInitializeSurvivesFirstCallerDisconnect verifies that a client giving up +// mid-handshake does not cancel the shared upstream call for everyone else. +// +// The handshake is collapsed by singleflight, and singleflight runs the shared +// function with whichever context the first caller supplied. Passing that +// context straight through would let one client's disconnect fail every waiter, +// so interceptInitialize detaches it with context.WithoutCancel. +func TestInitializeSurvivesFirstCallerDisconnect(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + strict := strictInitializeResponder() + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + // Long enough that the first caller aborts while the handshake is in + // flight, short enough to keep the test quick. + time.Sleep(400 * time.Millisecond) + return strict(req) + }) + + // First caller: gives up well before the backend replies. + var wg sync.WaitGroup + wg.Go(func() { + ctx, cancel := context.WithTimeout(t.Context(), 80*time.Millisecond) + defer cancel() + body := fmt.Sprintf( + `{"jsonrpc":"2.0","id":%q,"method":"initialize","params":{"protocolVersion":%q}}`, + "quitter", testProtocolVersion) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body))) + if !assert.NoError(t, err) { + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + resp, err := http.DefaultClient.Do(req) //nolint:bodyclose // err path returns no body + assert.Error(t, err, "the first caller is expected to abort mid-handshake") + if err == nil { + _ = resp.Body.Close() + } + }) + + // Give the first caller time to enter the flight and then abort. + time.Sleep(150 * time.Millisecond) + + // Second caller arrives while the shared handshake is still running. + decoded := postInitialize(t, url, "survivor") + assert.Nil(t, decoded["error"], + "the shared handshake must not be cancelled by the first caller disconnecting") + require.NotNil(t, decoded["result"], "second caller should still complete its handshake") + assert.Equal(t, "survivor", decoded["id"]) + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for the aborting caller to finish") + } + + assert.Equal(t, 1, backend.count("initialize"), + "the disconnect must not cause a second upstream handshake") +} + +// TestMalformedInitializeReplyBecomesInternalError pins an intentional +// behaviour change made when the flight started sharing an id-free handshake +// rather than a whole response. +// +// A JSON-RPC success carrying neither an error nor a result is malformed for +// initialize -- InitializeResult requires protocolVersion, capabilities and +// serverInfo. It used to be forwarded to the client verbatim; it is now +// reported as an internal error, and nothing is cached, so the next client +// still gets a real handshake attempt. +func TestMalformedInitializeReplyBecomesInternalError(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + var mu sync.Mutex + var calls int + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + mu.Lock() + calls++ + first := calls == 1 + mu.Unlock() + if first { + // Success envelope with no result and no error. + return &jsonrpc2.Response{ID: req.ID} + } + return initializeResultFor(req) + }) + + malformed := postInitialize(t, url, "client-0") + require.NotNil(t, malformed["error"], "a result-less success must not be replayed to the client") + assert.Equal(t, "client-0", malformed["id"]) + + recovered := postInitialize(t, url, "client-1") + assert.Nil(t, recovered["error"], "the malformed reply must not be cached") + require.NotNil(t, recovered["result"], "the next client should get a real handshake") + + assert.Equal(t, 2, backend.count("initialize"), + "a malformed reply must leave the cache empty so the next client retries") +} diff --git a/pkg/transport/proxy/streamable/keyed_mutex.go b/pkg/transport/proxy/streamable/keyed_mutex.go new file mode 100644 index 0000000000..3a963b55e6 --- /dev/null +++ b/pkg/transport/proxy/streamable/keyed_mutex.go @@ -0,0 +1,104 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import "sync" + +// keyedMutex provides per-key mutual exclusion: two callers locking the SAME +// key are fully serialized (the second blocks until the first releases), +// while callers locking DIFFERENT keys proceed concurrently, never blocking +// on each other. HTTPProxy uses it (as uriLocks) to order a +// resources/subscribe|unsubscribe ref-count decision (see +// sessionRouter.addSubscription/removeSubscription) atomically with the +// upstream forward it may trigger, per uri -- see handlePost and FIX 1 in +// #5744's review: without this, a concurrent last-unsubscribe and +// first-subscribe for the SAME uri from different sessions could reach the +// backend out of order, leaving the backend desynchronized from the routing +// table's ref count. +// +// Unlike sessionRouter's own mutexes (subMu et al.), which only ever guard a +// map lookup/mutation and are never held across a blocking call, a key's +// entry here IS deliberately held across a blocking upstream request/response +// round trip -- that is the whole point of this type. Do not reuse it to +// guard plain in-memory state that has no blocking call in its critical +// section; a plain mutex is simpler and sufficient there. +// +// Entries are ref-counted and removed once no caller holds or is waiting for +// them, bounding keyedMutex's memory footprint to the number of CURRENTLY +// in-flight keys rather than the number of keys ever seen (see the +// resource-leak style rule). This differs deliberately from, e.g., +// pluginsvc's never-evicted pluginLock: that type's key cardinality is small +// and bounded by installed plugins, but uris here are arbitrary and +// client-supplied over the life of the proxy, so never evicting would leak +// memory. +type keyedMutex struct { + mu sync.Mutex + locks map[string]*keyedMutexEntry +} + +// keyedMutexEntry is one key's mutex plus a count of callers that currently +// hold or are waiting to acquire it. The count lets lock's returned unlock +// func safely delete the map entry once it reaches zero, without racing a +// concurrent lock(key) call that is (or is about to be) waiting on the same +// entry -- see lock's doc comment for why this two-phase acquire (register +// intent under mu, then block on the entry's own mutex) is race-free. +type keyedMutexEntry struct { + mu sync.Mutex + refCount int +} + +// newKeyedMutex creates an empty keyedMutex. +func newKeyedMutex() *keyedMutex { + return &keyedMutex{locks: make(map[string]*keyedMutexEntry)} +} + +// lock acquires exclusive ownership of key, blocking until any other caller +// currently holding (or already waiting for) the same key releases it. It +// returns an unlock func the caller MUST call exactly once (typically via +// defer, so it runs on every return/error path) to release key and, if no +// other caller is currently holding or waiting for it, remove its entry. +// +// Acquisition is two-phase: first, under k.mu, the entry is looked up (or +// created) and its refCount incremented -- this is what makes a concurrent +// unlock's decrement-and-maybe-delete race-free, since any new lock(key) +// caller either observes the entry before it is deleted (and increments it, +// preventing deletion) or observes it already deleted (and creates a fresh +// one), never a stale, about-to-be-deleted entry. Second, k.mu is released +// and the caller blocks on the entry's own mutex, so holding one key never +// blocks callers locking a different key. +// +// lock is reentrancy-free: calling lock(key) again for the same key from the +// same goroutine before releasing the first would deadlock, exactly like a +// plain sync.Mutex -- callers must never nest calls for the same key. +// +// The returned unlock is wrapped in sync.Once as defense-in-depth: it must +// still be called, but an accidental double-call becomes a no-op instead of +// unlocking a mutex a different caller now holds (releasing someone else's +// critical section) and driving refCount negative (a spurious later eviction). +func (k *keyedMutex) lock(key string) (unlock func()) { + k.mu.Lock() + e, ok := k.locks[key] + if !ok { + e = &keyedMutexEntry{} + k.locks[key] = e + } + e.refCount++ + k.mu.Unlock() + + e.mu.Lock() + + var once sync.Once + return func() { + once.Do(func() { + e.mu.Unlock() + + k.mu.Lock() + e.refCount-- + if e.refCount == 0 { + delete(k.locks, key) + } + k.mu.Unlock() + }) + } +} diff --git a/pkg/transport/proxy/streamable/keyed_mutex_test.go b/pkg/transport/proxy/streamable/keyed_mutex_test.go new file mode 100644 index 0000000000..179da5dc6d --- /dev/null +++ b/pkg/transport/proxy/streamable/keyed_mutex_test.go @@ -0,0 +1,118 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestKeyedMutex_DifferentKeysProceedConcurrently pins the whole point of the +// type: two callers holding DIFFERENT keys must not block each other. A +// regression to a single global lock would still pass every integration test +// (just serialized), so this is asserted directly here. +func TestKeyedMutex_DifferentKeysProceedConcurrently(t *testing.T) { + t.Parallel() + km := newKeyedMutex() + + // Hold key "a" for the duration. + unlockA := km.lock("a") + t.Cleanup(unlockA) + + // Locking a different key must succeed while "a" is still held. + got := make(chan func(), 1) + go func() { got <- km.lock("b") }() + + select { + case unlockB := <-got: + unlockB() + case <-time.After(2 * time.Second): + t.Fatal("lock on a different key blocked behind an unrelated held key") + } +} + +// TestKeyedMutex_SameKeySerializes verifies mutual exclusion for the same key: +// a second lock blocks until the first releases. +func TestKeyedMutex_SameKeySerializes(t *testing.T) { + t.Parallel() + km := newKeyedMutex() + + unlock1 := km.lock("k") + + acquired := make(chan struct{}) + go func() { + unlock2 := km.lock("k") + close(acquired) + unlock2() + }() + + // Must not acquire while the first holder is active. + select { + case <-acquired: + t.Fatal("second lock on the same key acquired before the first released") + case <-time.After(100 * time.Millisecond): + } + + unlock1() + + select { + case <-acquired: + case <-time.After(2 * time.Second): + t.Fatal("second lock on the same key never acquired after release") + } +} + +// TestKeyedMutex_EvictsAtZeroRefCount verifies the documented memory-leak +// guard: an entry is removed once no caller holds or is waiting for it. URIs +// are client-supplied and unbounded, so failing to evict would leak. +func TestKeyedMutex_EvictsAtZeroRefCount(t *testing.T) { + t.Parallel() + km := newKeyedMutex() + + unlock := km.lock("k") + + km.mu.Lock() + _, present := km.locks["k"] + km.mu.Unlock() + require.True(t, present, "entry must exist while the key is held") + + unlock() + + km.mu.Lock() + _, present = km.locks["k"] + remaining := len(km.locks) + km.mu.Unlock() + assert.False(t, present, "entry must be evicted once refCount reaches 0") + assert.Equal(t, 0, remaining, "no entries should remain after all locks are released") +} + +// TestKeyedMutex_DoubleUnlockIsNoop verifies the sync.Once guard: calling the +// returned unlock more than once is a safe no-op — it must not panic, must not +// drive refCount negative, and must not release a lock a later caller holds. +func TestKeyedMutex_DoubleUnlockIsNoop(t *testing.T) { + t.Parallel() + km := newKeyedMutex() + + unlock := km.lock("k") + unlock() + assert.NotPanics(t, unlock, "second unlock must be a no-op, not a panic") + + // A fresh acquisition must see a clean refCount of exactly 1 (a stray + // double-unlock decrementing a live entry would skew this and corrupt + // eviction). + u2 := km.lock("k") + km.mu.Lock() + rc := km.locks["k"].refCount + km.mu.Unlock() + assert.Equal(t, 1, rc, "refCount must reflect exactly one active holder") + + u2() + km.mu.Lock() + remaining := len(km.locks) + km.mu.Unlock() + assert.Equal(t, 0, remaining, "entry must be evicted after the fresh holder releases") +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy.go b/pkg/transport/proxy/streamable/streamable_proxy.go index f0f7c098db..7c9fb11edd 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy.go +++ b/pkg/transport/proxy/streamable/streamable_proxy.go @@ -5,7 +5,6 @@ package streamable import ( - "bytes" "context" "encoding/json" "errors" @@ -22,9 +21,11 @@ import ( "github.com/google/uuid" "golang.org/x/exp/jsonrpc2" + sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/bodylimit" "github.com/stacklok/toolhive/pkg/healthcheck" + "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/transport/types" ) @@ -46,6 +47,10 @@ const ( // does not affect responses, so SSE response streams are unaffected. Matches // the vMCP server default. defaultReadTimeout = 30 * time.Second + + // proxyChannelBufferSize is the buffer size for the messageCh and + // responseCh channels connecting the HTTP proxy to the container runner. + proxyChannelBufferSize = 100 ) // HTTPProxy implements a proxy for streamable HTTP transport. @@ -91,6 +96,44 @@ type HTTPProxy struct { // to restore before replying. Same uniqueness requirement as `waiters`. idRestore sync.Map // map[string]jsonrpc2.ID + // serverStreams delivers server->client notifications to connected + // sessions' SSE streams (the standalone GET stream for global + // notifications, or a specific session's stream for subscription-scoped + // ones). See dispatcher_streams.go. + serverStreams *serverStreamRegistry + + // routing owns the per-session routing state (progress tokens, + // subscriptions, log levels) that dispatcher.go's routeNotification uses + // to deliver server->client messages to the correct session(s) without + // cross-session leakage. See dispatcher_routing.go. + routing *sessionRouter + + // initialize caches the backend's InitializeResult so only the first + // client handshake reaches the single stdio session behind this proxy and + // every later one is answered locally. See initialize_cache.go. + initialize *initializeCache + + // uriLocks serializes each uri's resources/subscribe|unsubscribe ref-count + // decision with the upstream forward it may trigger, so concurrent + // subscribe/unsubscribe calls for the SAME uri from different sessions can + // never reach the backend out of order. See keyed_mutex.go and + // handlePost/interceptSessionScopedRequest. + uriLocks *keyedMutex + + // standaloneSSE controls whether GET /mcp opens a standalone SSE stream for + // server->client notifications (the default) or returns 405. Set via + // WithStandaloneSSE(false) to opt out. + standaloneSSE bool + + // strictProtocolValidation controls whether handlePost rejects a request + // whose MCP-Protocol-Version header names an unknown/unsupported MCP + // revision (see isSupportedMCPVersion) with HTTP 400. Default false: any + // version string is accepted, since this proxy is transport-level and does + // not depend on a specific MCP revision. An absent header is always + // accepted, in either mode, per the streamable HTTP spec's rule to assume + // 2025-03-26 when the header is missing. Set via WithStrictProtocolValidation. + strictProtocolValidation bool + // Health checker healthChecker *healthcheck.HealthChecker @@ -152,6 +195,25 @@ func WithPrefixHandlers(handlers map[string]http.Handler) Option { } } +// WithStandaloneSSE controls whether GET /mcp opens a standalone SSE stream +// for server->client notifications. It defaults to enabled; pass false to opt +// out and restore the prior behavior of returning 405 for GET requests. +func WithStandaloneSSE(enabled bool) Option { + return func(p *HTTPProxy) { + p.standaloneSSE = enabled + } +} + +// WithStrictProtocolValidation enables strict MCP-Protocol-Version checking. +// When enabled, a request whose MCP-Protocol-Version header names an +// unsupported/unknown MCP revision is rejected with HTTP 400. An absent +// header is still accepted (the streamable HTTP spec says to assume +// 2025-03-26). Default (false) preserves the version-agnostic behavior: +// any version string is accepted. +func WithStrictProtocolValidation(enabled bool) Option { + return func(p *HTTPProxy) { p.strictProtocolValidation = enabled } +} + // NewHTTPProxy creates a new HTTPProxy for streamable HTTP transport. func NewHTTPProxy( host string, @@ -171,9 +233,14 @@ func NewHTTPProxy( shutdownCh: make(chan struct{}), prometheusHandler: prometheusHandler, middlewares: middlewares, - messageCh: make(chan jsonrpc2.Message, 100), - responseCh: make(chan jsonrpc2.Message, 100), + messageCh: make(chan jsonrpc2.Message, proxyChannelBufferSize), + responseCh: make(chan jsonrpc2.Message, proxyChannelBufferSize), sessionTTL: session.ResolveSessionTTLFromEnv(), + serverStreams: newServerStreamRegistry(), + routing: newSessionRouter(), + initialize: newInitializeCache(), + uriLocks: newKeyedMutex(), + standaloneSSE: true, } for _, opt := range opts { @@ -233,6 +300,11 @@ func (p *HTTPProxy) Start(_ context.Context) error { // Route container responses to matching waiter channels go p.dispatchResponses() + // Periodically bound routing's subscription/log-level maps, and + // serverStreams' standalone GET streams, to active sessions (see + // reapRoutingState's doc comment for the isActive tradeoff). + go p.reapRoutingState() + go func() { slog.Debug("streamable HTTP proxy started", "port", p.port) //nolint:gosec // G706: logging configured host and port @@ -253,6 +325,11 @@ func (p *HTTPProxy) Stop(ctx context.Context) error { p.stopOnce.Do(func() { close(p.shutdownCh) + // Signal every live GET SSE stream's stop channel so in-flight handlers + // observe a closed per-stream stop signal and unblock/return promptly, + // rather than relying solely on shutdownCh (which they also select on). + p.serverStreams.closeAll() + // Stop session manager cleanup; active sessions expire via TTL if p.sessionManager != nil { if err := p.sessionManager.Stop(); err != nil { @@ -331,9 +408,76 @@ func (p *HTTPProxy) handleStreamableRequest(w http.ResponseWriter, r *http.Reque } } -func (*HTTPProxy) handleGet(w http.ResponseWriter, _ *http.Request) { - // SSE not offered here; explicit 405 is spec-compliant - writeHTTPError(w, http.StatusMethodNotAllowed, "SSE not supported on this endpoint") +// handleGet serves a standalone GET SSE stream for unsolicited server->client +// notifications (progress, resources/updated, */list_changed, logging/message). +// It is enabled by default; WithStandaloneSSE(false) restores the prior 405 +// behavior. The caller MUST provide a known Mcp-Session-Id header: an empty +// header is rejected with 400, and an unknown one with 404 (mirroring +// handleDelete). A session's stream is one-per-session: a second GET for the +// same session evicts the first (see serverStreamRegistry.register). Per the +// locked product decision, a notification dispatched while no GET stream is +// connected for a session is dropped -- there is no replay or pending queue +// for this stream. +func (p *HTTPProxy) handleGet(w http.ResponseWriter, r *http.Request) { + if !p.standaloneSSE { + // SSE not offered here; explicit 405 is spec-compliant + writeHTTPError(w, http.StatusMethodNotAllowed, "SSE not supported on this endpoint") + return + } + + flusher, ok := assertFlushable(w) + if !ok { + return + } + + sessID := r.Header.Get("Mcp-Session-Id") + if sessID == "" { + writeHTTPError(w, http.StatusBadRequest, "Mcp-Session-Id header required for standalone SSE") + return + } + if _, ok := p.sessionManager.Get(sessID); !ok { + session.WriteNotFound(w, nil) + return + } + + setSSEHeaders(w) + + s := p.serverStreams.register(sessID) + defer p.serverStreams.deregister(sessID, s) + + // Send headers immediately so the client knows the stream is open. + flusher.Flush() + + keepAliveTicker := time.NewTicker(sseKeepAliveInterval) + defer keepAliveTicker.Stop() + + for { + select { + case msg := <-s.data: + // No ok-check: data is never closed (see serverStream's doc comment). + data, err := jsonrpc2.EncodeMessage(msg) + if err != nil { + slog.Error("failed to encode server->client notification", "error", err) + continue + } + if err := writeSSEData(w, flusher, data); err != nil { + slog.Debug("failed to write notification to GET SSE stream", "error", err) + return + } + case <-keepAliveTicker.C: + if err := writeSSEKeepAlive(w, flusher); err != nil { + slog.Debug("failed to write keep-alive", "error", err) + return + } + case <-s.stop: + // Evicted by a second GET for this session, or proxy Stop's closeAll. + return + case <-r.Context().Done(): + return + case <-p.shutdownCh: + return + } + } } func (p *HTTPProxy) handleDelete(w http.ResponseWriter, r *http.Request) { @@ -350,15 +494,41 @@ func (p *HTTPProxy) handleDelete(w http.ResponseWriter, r *http.Request) { //nolint:gosec // G706: session ID is from validated request header slog.Debug("failed to delete session", "session_id", sessID, "error", err) } + + // Purge routing state (subscriptions, log level) for the deleted session so + // it cannot outlive the session itself. maxChanged is intentionally not + // reconciled upstream here: if the deleted session held the max-verbosity + // log level, the shared backend simply stays MORE verbose than strictly + // necessary for the remaining sessions until one of them changes the level + // again -- extra log volume, not a correctness or security issue, so it is + // not worth an extra upstream round-trip on every DELETE. + p.routing.purgeSession(sessID) + + // Tear down the session's standalone GET SSE stream (if any), so its + // handler goroutine (and the list_changed/subscription broadcasts it would + // otherwise keep receiving) stops immediately instead of surviving until + // the client eventually closes the underlying socket on its own (see FIX 2 + // in #5744's review). Ordering matches the durable-before-in-memory style + // rule: the durable session delete and routing purge above have already + // happened, so nothing durable or in routing state can reference this + // session by the time its stream is closed. + p.serverStreams.closeStream(sessID) + w.WriteHeader(http.StatusNoContent) } func (p *HTTPProxy) handlePost(w http.ResponseWriter, r *http.Request) { ctx := r.Context() - // Optionally validate MCP-Protocol-Version; accept missing for compatibility + // MCP-Protocol-Version validation is opt-in via strictProtocolValidation + // (WithStrictProtocolValidation). Default (false) is version-agnostic: + // any header value is accepted, since this proxy is transport-level and + // does not depend on a specific MCP revision. When strict mode is + // enabled, a present-but-unrecognized version is rejected with 400; an + // absent header is always accepted in either mode, per the streamable + // HTTP spec's rule to assume 2025-03-26 when the header is missing. protoVer := r.Header.Get("MCP-Protocol-Version") - if protoVer != "" && !isSupportedMCPVersion(protoVer) { + if p.strictProtocolValidation && protoVer != "" && !isSupportedMCPVersion(protoVer) { writeHTTPError(w, http.StatusBadRequest, "Unsupported MCP-Protocol-Version") return } @@ -377,13 +547,13 @@ func (p *HTTPProxy) handlePost(w http.ResponseWriter, r *http.Request) { return } - // Batch vs single message - if isBatch(body) { - sessID, err := p.resolveSessionForBatch(w, r) - if err != nil { - return - } - p.handleBatchRequest(w, body, sessID) + // Reject JSON-RPC batches outright. Batching was removed in MCP revision + // 2025-06-18 and ToolHive serves only 2025-11-25 and 2026-07-28. Rejecting + // at the executor (not only in ParsingMiddleware) makes this independent of + // middleware presence, ordering, and Content-Type, so a batch can never + // reach the backend uninspected by authz/audit/tool-filtering (see #5745). + if mcp.IsBatchRequest(body) { + mcp.WriteBatchUnsupportedError(w) return } @@ -409,6 +579,39 @@ func (p *HTTPProxy) handlePost(w http.ResponseWriter, r *http.Request) { return } + // Intercept resources/subscribe, resources/unsubscribe, and + // logging/setLevel for session-bearing (Legacy) requests, AFTER + // applyMiddlewares/authz has admitted req (Start wraps handleStreamableRequest + // in the middleware chain, and handlePost is only ever reached through it), + // so an authz-denied request never gets recorded in routing. + // + // For resources/subscribe|unsubscribe, uriLocks serializes the ref-count + // decision made inside interceptSessionScopedRequest WITH the upstream + // forward it may perform itself (see interceptSubscribe/ + // interceptUnsubscribe), per uri, so a concurrent last-unsubscribe and + // first-subscribe for the SAME uri from different sessions can never reach + // the backend out of order (see uriLocks's doc comment and FIX 1 in + // #5744's review). The lock is scoped to ONLY this decision+forward+record + // -- acquired immediately before interceptSessionScopedRequest and released + // immediately after it returns, BEFORE writeInterceptedResponse or any + // client-facing streaming -- so a slow or SSE-streamed response never + // head-of-line-blocks a different session's request for the SAME uri (see + // FIX 3 in #5744's review). A different uri, or any non-subscription + // request, is never blocked by this at all: see resourceSubscriptionURI. + interceptedMsg, handled := func() (jsonrpc2.Message, bool) { + uri, ok := resourceSubscriptionURI(req) + if !ok { + return p.interceptSessionScopedRequest(ctx, sessID, req) + } + unlock := p.uriLocks.lock(uri) + defer unlock() + return p.interceptSessionScopedRequest(ctx, sessID, req) + }() + if handled { + writeInterceptedResponse(w, r, interceptedMsg, sessID, setSessionHeader) + return + } + // If client accepts SSE, stream the response on an SSE stream for this request if strings.Contains(r.Header.Get("Accept"), "text/event-stream") { p.handleSingleRequestSSE(ctx, w, sessID, req, setSessionHeader) @@ -419,43 +622,16 @@ func (p *HTTPProxy) handlePost(w http.ResponseWriter, r *http.Request) { p.handleSingleRequest(ctx, w, sessID, req, setSessionHeader) } -// handleBatchRequest processes a batch JSON-RPC request and writes a batch response. -func (p *HTTPProxy) handleBatchRequest(w http.ResponseWriter, body []byte, sessID string) { - rawMessages, ok := decodeBatch(w, body) - if !ok { - return - } - - var responses []json.RawMessage - hadRequest := false - - for _, raw := range rawMessages { - // Detect if this element is a request with an ID - if msg, err := jsonrpc2.DecodeMessage(raw); err == nil { - if req, ok := msg.(*jsonrpc2.Request); ok && req.ID.IsValid() { - hadRequest = true - } - } - resp := p.processSingleMessage(sessID, raw) - if resp != nil { - responses = append(responses, resp) - } - } - - if !hadRequest { - // Per spec: batches containing only notifications/responses -> 202 Accepted, no body - w.WriteHeader(http.StatusAccepted) - return - } - - w.Header().Set("Content-Type", "application/json") - // It's valid to return an empty array if requests produced no responses - if err := json.NewEncoder(w).Encode(responses); err != nil { - slog.Error("failed to encode batch response", "error", err) - } -} - // handleSingleRequest handles a single JSON-RPC request message end-to-end. +// +// If req carries a params._meta.progressToken, it is forwarded upstream +// UNREWRITTEN: a plain JSON POST response has no stream to carry interim +// progress on, so there is nothing to correlate a progress route to, and no +// progress routing is set up (see handleSingleRequestSSE for the SSE path, +// which does). Any notifications/progress the backend later sends for this +// request's original token will simply find no registered route in +// dispatcher.go's routeProgress and be dropped, which is correct: this +// non-streaming client could never have received them anyway. func (p *HTTPProxy) handleSingleRequest( ctx context.Context, w http.ResponseWriter, @@ -488,6 +664,20 @@ func (p *HTTPProxy) handleSingleRequest( } } +// handleSingleRequestSSE handles a single JSON-RPC request whose client asked +// for an SSE response (Accept: text/event-stream). Unlike handleSingleRequest, +// this stream stays open for the lifetime of the request, so it is the +// correct destination for any interim notifications/progress the backend +// sends while it works -- per MCP, request-related messages belong on the +// originating request's stream, not the standalone GET stream (see +// dispatcher.go's routeProgress and docs/arch/03-transport-architecture.md). +// +// If req carries a params._meta.progressToken, setupProgressRouting mints a +// proxy-only token, rewrites the outgoing request to carry it, and registers a +// delivery route for it; the select loop below then interleaves any progress +// notifications delivered on that route with the final response, writing each +// as its own SSE data: frame, and only returns once the final response (or a +// context/shutdown signal) arrives. func (p *HTTPProxy) handleSingleRequestSSE( ctx context.Context, w http.ResponseWriter, @@ -499,132 +689,123 @@ func (p *HTTPProxy) handleSingleRequestSSE( defer cancel() // Prepare SSE response headers - flusher, ok := w.(http.Flusher) - if !ok { - writeHTTPError(w, http.StatusInternalServerError, "Streaming not supported") - return - } - w.Header().Set("Content-Type", "text/event-stream") - w.Header().Set("Cache-Control", "no-cache") - w.Header().Set("Connection", "keep-alive") + var setMcpSessionID func() if setSessionHeader { - w.Header().Set("Mcp-Session-Id", sessID) + setMcpSessionID = func() { w.Header().Set("Mcp-Session-Id", sessID) } } - flusher.Flush() - - msg, err := p.doRequest(ctx, sessID, req) - if err != nil { - // Send a best-effort error event - errMsg := "Internal error" - code := -32603 - if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { - errMsg = "Timeout" - code = -32000 - } - errObj := map[string]any{ - "jsonrpc": "2.0", - "id": req.ID.Raw(), - "error": map[string]any{ - "code": code, - "message": errMsg, - }, - } - if data, mErr := json.Marshal(errObj); mErr == nil { - if _, err := fmt.Fprintf(w, "data: %s\n\n", data); err != nil { - slog.Debug("failed to write error message", "error", err) - return - } - flusher.Flush() - } + flusher, ok := startSSEStream(w, setMcpSessionID) + if !ok { return } - data, err := jsonrpc2.EncodeMessage(msg) - if err != nil { - slog.Error("failed to encode JSON-RPC response", "error", err) - writeHTTPError(w, http.StatusInternalServerError, "Failed to encode response") - return - } - // Write SSE event with the JSON-RPC response and flush - if _, err := fmt.Fprintf(w, "data: %s\n\n", data); err != nil { //nolint:gosec // G705: SSE data from MCP protocol - slog.Debug("failed to write response", "error", err) - return - } - flusher.Flush() -} + outgoingReq, deliverCh, dropProgressRoute := p.setupProgressRouting(req) + defer dropProgressRoute() -// processSingleMessage processes one raw JSON-RPC in a batch and returns encoded response bytes or nil. -func (p *HTTPProxy) processSingleMessage(sessID string, raw json.RawMessage) json.RawMessage { - // Note: batch processing path - msg, err := jsonrpc2.DecodeMessage(raw) + waitCh, ck, cleanup, err := p.sendCorrelatedRequest(sessID, outgoingReq) if err != nil { - //nolint:gosec // G706: logging raw JSON-RPC data from HTTP request body - slog.Warn("skipping invalid message in batch", "raw", string(raw)) - return nil + //nolint:gosec // G706: method is from parsed JSON-RPC request + slog.Error("failed to send request upstream", "method", req.Method, "error", err) + writeSSEErrorEvent(w, flusher, req.ID, err) + return } + defer cleanup() - // Notifications: just forward and continue - if isNotification(msg) { - if err := p.SendMessageToDestination(msg); err != nil { - slog.Error("failed to send notification to destination", "error", err) + for { + select { + case progressMsg := <-deliverCh: + // deliverCh is nil when req carried no progressToken (see + // setupProgressRouting); a nil channel's case is never selected, + // so this arm only fires when progress routing is active. + data, err := jsonrpc2.EncodeMessage(progressMsg) + if err != nil { + slog.Error("failed to encode progress notification", "error", err) + continue + } + if err := writeSSEData(w, flusher, data); err != nil { + slog.Debug("failed to write progress notification to SSE stream", "error", err) + return + } + // Progress does not end the request; keep waiting for more + // progress or the final response. + case msg := <-waitCh: + p.writeSingleRequestSSEFinalResponse(w, flusher, msg, ck) + return + case <-ctx.Done(): + writeSSEErrorEvent(w, flusher, req.ID, ctx.Err()) + return + case <-p.shutdownCh: + // Matches the prior (pre-progress) behavior: proxy shutdown while a + // request is in flight is reported to the client as a best-effort + // Timeout error event, the same as context.Canceled. + writeSSEErrorEvent(w, flusher, req.ID, context.Canceled) + return } - return nil } +} - // Client responses: accept and forward, no HTTP body - if _, ok := msg.(*jsonrpc2.Response); ok { - if err := p.SendMessageToDestination(msg); err != nil { - slog.Error("failed to forward client response to destination", "error", err) +// writeSingleRequestSSEFinalResponse restores msg's original client ID (if it +// is a correlated *jsonrpc2.Response) and writes it as the final SSE data: +// frame for handleSingleRequestSSE's request. Errors encoding/restoring are +// logged; the client simply does not get a final frame in that case, matching +// the prior (pre-progress) behavior's error handling. +func (p *HTTPProxy) writeSingleRequestSSEFinalResponse( + w http.ResponseWriter, flusher http.Flusher, msg jsonrpc2.Message, ck string, +) { + finalMsg := msg + if r, ok := msg.(*jsonrpc2.Response); ok && r.ID.IsValid() { + restored, err := p.restoreResponseID(r, ck) + if err != nil { + slog.Error("failed to restore response id", "error", err) + return } - return nil - } - - req, ok := msg.(*jsonrpc2.Request) - if !ok || !req.ID.IsValid() { - slog.Warn("skipping invalid batch item (not a request with ID/response/notification)", - "type", fmt.Sprintf("%T", msg)) - return nil + finalMsg = restored } - waitCh, cleanup := p.createWaiter(sessID, req.ID) - defer cleanup() - bkey := idKeyFromID(req.ID) - - // Transform outgoing request ID to composite and send - ck := compositeKey(sessID, bkey) - proxiedMsg, err := encodeRequestWithID(req, ck) + data, err := jsonrpc2.EncodeMessage(finalMsg) if err != nil { - slog.Error("failed to encode batch request", "error", err) - return nil + slog.Error("failed to encode JSON-RPC response", "error", err) + return } - if err := p.SendMessageToDestination(proxiedMsg); err != nil { - slog.Error("failed to send message to destination", "error", err) - return nil + if err := writeSSEData(w, flusher, data); err != nil { + slog.Debug("failed to write response", "error", err) } +} - response := p.waitForResponse(waitCh, p.requestTimeout) - if response == nil { - slog.Warn("streamableHTTP: batch timeout waiting for key", "key", bkey) - return nil +// writeSSEErrorEvent writes a best-effort JSON-RPC error as a single SSE +// message event (via writeSSEData), for a request whose response headers +// (200 + text/event-stream) have already been sent -- an HTTP error status can +// no longer be set at this point, so the error must be communicated in-band as +// the response payload. +// +// The "id" key is included only if id.IsValid(); MCP narrows base JSON-RPC +// 2.0 here (schema/2025-11-25 types the error response id as optional, +// string|number, no null), so an invalid/absent id must be encoded by +// omitting the key, never as "id":null. +func writeSSEErrorEvent(w http.ResponseWriter, flusher http.Flusher, id jsonrpc2.ID, err error) { + errMsg := "Internal error" + code := -32603 + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + errMsg = "Timeout" + code = -32000 + } + errObj := map[string]any{ + "jsonrpc": "2.0", + "error": map[string]any{ + "code": code, + "message": errMsg, + }, + } + if id.IsValid() { + errObj["id"] = id.Raw() + } + data, mErr := json.Marshal(errObj) + if mErr != nil { + slog.Error("failed to encode SSE error event", "error", mErr) + return } - - if r, ok := response.(*jsonrpc2.Response); ok && r.ID.IsValid() { - restored, err := p.restoreResponseID(r, ck) - if err != nil { - slog.Error("failed to restore response ID", "error", err) - return nil - } - data, err := jsonrpc2.EncodeMessage(restored) - if err != nil { - slog.Error("failed to encode JSON-RPC response", "error", err) - return nil - } - return data + if err := writeSSEData(w, flusher, data); err != nil { + slog.Debug("failed to write error message", "error", err) } - - slog.Warn("received invalid message that is not a valid response", - "type", fmt.Sprintf("%T", response)) - return nil } func encodeRequestWithID(req *jsonrpc2.Request, newID string) (jsonrpc2.Message, error) { @@ -644,6 +825,99 @@ func encodeRequestWithID(req *jsonrpc2.Request, newID string) (jsonrpc2.Message, return jsonrpc2.DecodeMessage(data2) } +// progressDeliverBufferSize bounds how many not-yet-written progress +// notifications may queue for a single in-flight POST-SSE request before +// further ones are dropped (see routeProgress's non-blocking send). +const progressDeliverBufferSize = 16 + +// rewriteMetaProgressToken returns a request identical to req except that, if +// params._meta.progressToken is present, it is replaced with ptGlobal. This +// lets dispatcher.go's routeProgress correlate the backend's later +// notifications/progress replies (which echo the token verbatim) back to the +// specific POST-SSE request that asked for them, without ever exposing +// ptGlobal -- or any other session's progress -- to more than the one client +// that requested it. +// +// It returns hadToken=false (with rewritten == req, unchanged) if req carried +// no progressToken, telling the caller not to set up progress routing at all. +// originalToken is the client's own progressToken value (string or number), +// to be restored before delivering progress back to the client. +// +// Per the "copy before mutating caller input" style rule, req's params map is +// decoded into fresh copies before any mutation, so req itself (and any +// caller-held reference to it) is never modified. +func rewriteMetaProgressToken( + req *jsonrpc2.Request, ptGlobal string, +) (originalToken any, hadToken bool, rewritten *jsonrpc2.Request, err error) { + if len(req.Params) == 0 { + return nil, false, req, nil + } + var paramsMap map[string]any + if err := json.Unmarshal(req.Params, ¶msMap); err != nil { + return nil, false, nil, err + } + meta, ok := paramsMap["_meta"].(map[string]any) + if !ok { + return nil, false, req, nil + } + original, ok := meta["progressToken"] + if !ok { + return nil, false, req, nil + } + + metaCopy := maps.Clone(meta) + metaCopy["progressToken"] = ptGlobal + paramsCopy := maps.Clone(paramsMap) + paramsCopy["_meta"] = metaCopy + + data, err := json.Marshal(paramsCopy) + if err != nil { + return nil, false, nil, err + } + return original, true, &jsonrpc2.Request{ID: req.ID, Method: req.Method, Params: data}, nil +} + +// setupProgressRouting mints a fresh proxy-only progress-routing token and +// rewrites req's params._meta.progressToken to it (see +// rewriteMetaProgressToken), registering a progressRoute so the backend's +// later notifications/progress messages that echo the token are delivered on +// the returned channel rather than leaking to the standalone GET stream or +// another session. +// +// If req carries no progressToken (or minting/rewriting fails, logged as a +// warning), the returned request is req itself, unmodified, and the returned +// channel is nil: a select on a nil channel never fires, so callers can select +// on it unconditionally without special-casing "no progress requested". +// +// The returned cleanup func drops the progress token and MUST be deferred by +// the caller; it is always safe to call, including when no route was +// registered. +func (p *HTTPProxy) setupProgressRouting(req *jsonrpc2.Request) (*jsonrpc2.Request, <-chan jsonrpc2.Message, func()) { + noop := func() {} + + ptGlobal, err := uuid.NewRandom() + if err != nil { + slog.Warn("failed to mint progress routing token; interim progress will not be delivered for this request", + "error", err) + return req, nil, noop + } + + originalToken, hadToken, rewritten, err := rewriteMetaProgressToken(req, ptGlobal.String()) + if err != nil { + slog.Warn("failed to rewrite progress token; interim progress will not be delivered for this request", + "error", err) + return req, nil, noop + } + if !hadToken { + return req, nil, noop + } + + key := ptGlobal.String() + deliver := make(chan jsonrpc2.Message, progressDeliverBufferSize) + p.routing.recordProgressToken(key, progressRoute{deliver: deliver, originalToken: originalToken}) + return rewritten, deliver, func() { p.routing.dropProgressToken(key) } +} + func (p *HTTPProxy) restoreResponseID(resp *jsonrpc2.Response, ck string) (jsonrpc2.Message, error) { orig, ok := p.idRestore.Load(ck) if !ok { @@ -668,20 +942,42 @@ func (p *HTTPProxy) restoreResponseID(resp *jsonrpc2.Response, ck string) (jsonr return jsonrpc2.DecodeMessage(data2) } -func (p *HTTPProxy) doRequest(ctx context.Context, sessID string, req *jsonrpc2.Request) (jsonrpc2.Message, error) { +// sendCorrelatedRequest mints the composite-key waiter for req, rewrites its +// wire ID to that composite key, and sends it upstream. It returns the waiter +// channel and ck (needed to restore the original ID on the eventual response, +// see restoreResponseID), plus a cleanup func the caller MUST call (typically +// via defer) to release the waiter/idRestore entries once done waiting. +// +// This is split out from doRequest so callers that need to interleave OTHER +// events (e.g. handleSingleRequestSSE's progress notifications) with the wait +// for the final response can select over both, instead of being limited to +// doRequest's single blocking wait. +func (p *HTTPProxy) sendCorrelatedRequest( + sessID string, req *jsonrpc2.Request, +) (waitCh chan jsonrpc2.Message, ck string, cleanup func(), err error) { key := idKeyFromID(req.ID) - ck := compositeKey(sessID, key) + ck = compositeKey(sessID, key) - waitCh, cleanup := p.createWaiter(sessID, req.ID) - defer cleanup() + waitCh, cleanup = p.createWaiter(sessID, req.ID) proxiedMsg, err := encodeRequestWithID(req, ck) if err != nil { - return nil, fmt.Errorf("encode request: %w", err) + cleanup() + return nil, "", nil, fmt.Errorf("encode request: %w", err) } if err := p.SendMessageToDestination(proxiedMsg); err != nil { - return nil, fmt.Errorf("send message: %w", err) + cleanup() + return nil, "", nil, fmt.Errorf("send message: %w", err) + } + return waitCh, ck, cleanup, nil +} + +func (p *HTTPProxy) doRequest(ctx context.Context, sessID string, req *jsonrpc2.Request) (jsonrpc2.Message, error) { + waitCh, ck, cleanup, err := p.sendCorrelatedRequest(sessID, req) + if err != nil { + return nil, err } + defer cleanup() select { case msg := <-waitCh: @@ -716,46 +1012,27 @@ func (p *HTTPProxy) ensureSession(id string) error { return p.sessionManager.AddWithID(id) } -// resolveSessionForBatch resolves the session for batch POSTs. -// Writes appropriate HTTP errors and returns an error when handling should stop. -// -// Sessionless POSTs receive a per-request UUID used solely as an in-process -// routing token. Sessionless routing tokens MUST be unique per request: -// sharing one (e.g. the empty string) across concurrent sessionless requests -// causes waiters/idRestore overwrites in the in-process sync.Maps, which leaks -// one client's response payload to another (with the JSON-RPC id rewritten to -// the receiver's). This is a confidentiality bug, not a performance issue -- -// do not collapse the token. The UUID is not registered with sessionManager, -// so no session object is created in any storage backend. -func (p *HTTPProxy) resolveSessionForBatch(w http.ResponseWriter, r *http.Request) (string, error) { - sessID := r.Header.Get("Mcp-Session-Id") - if sessID == "" { - token, err := uuid.NewRandom() - if err != nil { - writeHTTPError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to generate routing token: %v", err)) - return "", fmt.Errorf("generate routing token: %w", err) - } - return token.String(), nil - } - if _, ok := p.sessionManager.Get(sessID); !ok { - session.WriteNotFound(w, nil) - return "", fmt.Errorf("session not found") - } - return sessID, nil -} - // resolveSessionForRequest resolves session rules for a single JSON-RPC request. -// On initialize, assigns a new session ID if none is provided and returns setSessionHeader=true. -// A provided but unknown session ID returns 404. // -// Sessionless non-initialize requests receive a per-request UUID used solely as -// an in-process routing token (not registered with sessionManager). Sessionless -// routing tokens MUST be unique per request: sharing one (e.g. the empty string) -// across concurrent sessionless requests with the same JSON-RPC id collapses -// them onto the same compositeKey(sessID, idKey) and overwrites entries in the -// waiters / idRestore sync.Maps, leaking one client's response payload to -// another. This is a confidentiality bug, not a performance issue -- do not -// collapse the token. +// It first classifies the request as Modern or Legacy MCP (mcp.ClassifyRevision). +// A classification error is rejected outright with an HTTP 400 JSON-RPC error +// response. A Modern request is always sessionless: it gets a fresh per-request +// routing token and never touches sessionManager, regardless of any +// Mcp-Session-Id header it carries (see the confidentiality note below). A +// Legacy request falls through to the existing session rules: on initialize, +// assigns a new session ID if none is provided and returns setSessionHeader=true; +// a provided but unknown session ID returns 404. +// +// Sessionless requests (both the Modern branch and Legacy's sessionless +// non-initialize branch) receive a per-request UUID used solely as an +// in-process routing token (not registered with sessionManager). Sessionless +// routing tokens MUST be unique per request: sharing one (e.g. the empty string, +// or a stale/foreign Mcp-Session-Id) across concurrent sessionless requests with +// the same JSON-RPC id collapses them onto the same compositeKey(sessID, idKey) +// and overwrites entries in the waiters / idRestore sync.Maps, leaking one +// client's response payload to another. This is a confidentiality bug, not a +// performance issue -- do not collapse the token, and do not let a Modern +// request fall through to the session-lookup path below. // // Writes HTTP errors on failure and returns error to stop handling. func (p *HTTPProxy) resolveSessionForRequest( @@ -763,6 +1040,39 @@ func (p *HTTPProxy) resolveSessionForRequest( r *http.Request, req *jsonrpc2.Request, ) (string, bool, error) { + // Classification/routing here applies only to the single id-bearing request + // path; batches are rejected earlier in handlePost, and the + // notification/client-response path is Legacy-only by construction. + meta := mcp.ExtractMeta(req.Params) + protoHeader := r.Header.Get("MCP-Protocol-Version") + rev, err := mcp.ClassifyRevision(req.Method, meta, protoHeader) + if err != nil { + mcp.WriteClassificationError(w, req.ID.Raw(), err) + return "", false, err + } + + if rev == mcp.RevisionModern { + // Modern is stateless: mint a fresh routing token unconditionally and + // ignore any client-supplied Mcp-Session-Id. Never fall through to the + // session-lookup path below with it -- see the confidentiality note + // in this function's doc comment. + // + // This is safe only because, in this proxy, Mcp-Session-Id has no + // downstream authority: it is purely an in-process response-correlation + // token (the backend is stateless stdio, unlike the transparent proxy, + // which maps it to a backend-assigned SID via session metadata). If this + // proxy ever gains a shared session store or a stateful backend mapping, + // "ignore the client SID on Modern" must stay true -- silently reusing or + // looking up a known client SID here would reopen a session-validation + // bypass. See TestModernNeverReusesClientSessionIDAsRoutingToken. + token, err := uuid.NewRandom() + if err != nil { + writeHTTPError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to generate routing token: %v", err)) + return "", false, fmt.Errorf("generate routing token: %w", err) + } + return token.String(), false, nil + } + var setSessionHeader bool sessID := r.Header.Get("Mcp-Session-Id") @@ -803,20 +1113,427 @@ func (p *HTTPProxy) resolveSessionForRequest( return sessID, false, nil } -func isBatch(body []byte) bool { - t := bytes.TrimSpace(body) - return len(t) > 0 && t[0] == '[' +// resourceSubscriptionURI returns req's uri param and true if req is a +// resources/subscribe or resources/unsubscribe request carrying a non-empty +// uri -- the only requests whose ref-count decision and upstream forward (see +// interceptSubscribe/interceptUnsubscribe, called from +// interceptSessionScopedRequest) handlePost must serialize per-uri via +// uriLocks. Any other request -- including logging/setLevel (whose +// reconciliation is best-effort-ordered, see reconcileUpstreamLogLevel) and a +// subscribe/unsubscribe with no/empty uri (which interceptSessionScopedRequest +// itself no-ops on) -- returns ("", false), so handlePost skips locking +// entirely for it. +func resourceSubscriptionURI(req *jsonrpc2.Request) (string, bool) { + switch req.Method { + case methodResourcesSubscribe, methodResourcesUnsubscribe: + uri, ok := extractStringParam(req.Params, "uri") + if !ok || uri == "" { + return "", false + } + return uri, true + default: + return "", false + } } -func decodeBatch(w http.ResponseWriter, body []byte) ([]json.RawMessage, bool) { - var rawMessages []json.RawMessage - if err := json.Unmarshal(bytes.TrimSpace(body), &rawMessages); err != nil { - //nolint:gosec // G706: logging raw JSON-RPC batch data from HTTP request - slog.Warn("failed to decode batch JSON-RPC", "body", string(body)) - writeHTTPError(w, http.StatusBadRequest, "Invalid batch JSON-RPC") +// interceptSessionScopedRequest applies ref-counted resources/subscribe and +// resources/unsubscribe forwarding, and logging/setLevel max-verbosity +// reconciliation, for a request that belongs to a real, persisted Legacy +// session. It MUST be called only after applyMiddlewares/authz has admitted +// req (handlePost is only ever reached through Start's applyMiddlewares- +// wrapped mux entry), so an authz-denied subscribe is never recorded in +// routing -- see the "authz non-bypass" security test. +// +// For resources/subscribe|unsubscribe, the caller (handlePost) holds uriLocks +// for req's uri around this call, scoped to ONLY this call (see FIX 3 in +// #5744's review): interceptSubscribe/interceptUnsubscribe perform their own +// upstream round-trip (via p.doRequest, bounded by p.requestTimeout so the +// lock is never held indefinitely) synchronously, INSIDE this call, rather +// than reporting "forward it yourself" back to handlePost -- this is what +// lets the ref-count decision, the forward, and recording the subscription be +// atomic with respect to any other session's concurrent (un)subscribe for the +// SAME uri (see resourceSubscriptionURI and FIX 1 in #5744's review). +// +// sessID identifies "a real session" (as opposed to a per-request routing +// token minted for a Modern or Legacy-sessionless request, see +// resolveSessionForRequest) by whether sessionManager currently has it: only +// real sessions are ever stored there, so this single check both selects the +// right requests AND naturally excludes ones with nothing durable to track. +// +// It returns (msg, true) when the caller should write msg to the client and +// nothing further needs to happen: msg may be a synthesized success (a +// non-first subscribe, a non-last unsubscribe, or logging/setLevel), or the +// real upstream response (success OR error) for a first subscribe/last +// unsubscribe that this call forwarded itself. It returns (nil, false) only +// when the method isn't one of the three above, or the request isn't +// session-bearing -- in which case the caller must proceed with the normal +// upstream request/response flow. +func (p *HTTPProxy) interceptSessionScopedRequest( + ctx context.Context, sessID string, req *jsonrpc2.Request, +) (jsonrpc2.Message, bool) { + if _, isSession := p.sessionManager.Get(sessID); !isSession { + return nil, false + } + + switch req.Method { + case string(sdkmcp.MethodInitialize): + // Served from the shared InitializeResult cache after the first + // handshake, so the single stdio session behind this proxy never + // receives a second initialize. See initialize_cache.go. + return p.interceptInitialize(ctx, sessID, req), true + + case methodResourcesSubscribe: + uri, ok := extractStringParam(req.Params, "uri") + if !ok || uri == "" { + return nil, false + } + return p.interceptSubscribe(ctx, sessID, uri, req), true + + case methodResourcesUnsubscribe: + uri, ok := extractStringParam(req.Params, "uri") + if !ok || uri == "" { + return nil, false + } + return p.interceptUnsubscribe(ctx, sessID, uri, req), true + + case string(sdkmcp.MethodSetLogLevel): + // Unlike resources/subscribe|unsubscribe, this decision is NOT ordered + // with its upstream forward via uriLocks (there is no per-uri key to + // order on, and no uriLocks acquisition happens for this method -- see + // resourceSubscriptionURI). Concurrent maxChanged transitions from + // different sessions may therefore reach the backend out of order; see + // reconcileUpstreamLogLevel's doc comment for why that is an accepted, + // harmless tradeoff for now. The client's own synthesized success below + // is likewise a fire-and-forget, optimistic ack -- like subscribe's + // dedup path, it is served without waiting for reconcileUpstreamLogLevel + // to complete, which is intentional and low-stakes: notifications/message + // is unconditionally dropped (see routeNotification's SECURE-DROP), so + // there is no client-visible effect of the level change to get wrong. + level, _ := extractStringParam(req.Params, "level") + if maxChanged, newMax := p.routing.setLogLevel(sessID, logLevelRank(level)); maxChanged { + p.reconcileUpstreamLogLevel(logLevelName(newMax)) + } + return synthesizeSuccessResponse(req.ID), true + + default: return nil, false } - return rawMessages, true +} + +// interceptSubscribe handles a session-bearing resources/subscribe request +// for uri, called from interceptSessionScopedRequest while the caller +// (handlePost) holds uriLocks for uri. If uri already has at least one +// recorded subscriber, the backend has already admitted the FIRST subscribe +// and is sending updates for uri, so sess is simply recorded as an +// additional subscriber and a synthesized success is returned WITHOUT another +// upstream call -- dedup is safe here because the thing being deduped against +// is a call that is known to have succeeded. +// +// Otherwise sess is (candidate) first subscriber: the request is forwarded +// upstream via forwardUpstream, and sess is recorded as a subscriber ONLY if +// that forward succeeds with a non-error response. If the backend errors, +// times out, or the transport call itself fails, sess is NOT recorded and the +// real failure is returned to the client -- recording on failure would leave +// a phantom subscription entry that permanently starves every later +// session's subscribe for the SAME uri (deduping it against an upstream +// subscribe the backend never actually granted), and would violate the +// "write durable(upstream) storage before updating in-memory state" style +// rule. See FIX 1 in #5744's review. +func (p *HTTPProxy) interceptSubscribe(ctx context.Context, sessID, uri string, req *jsonrpc2.Request) jsonrpc2.Message { + if len(p.routing.subscribersOf(uri)) > 0 { + p.routing.addSubscription(uri, sessID) + return synthesizeSuccessResponse(req.ID) + } + + resp, err := p.forwardUpstream(ctx, sessID, req) + if err != nil { + return upstreamErrorResponse(req.ID, err) + } + if isErrorResponse(resp) { + // The backend rejected the first subscribe (e.g. resource not found): + // do not record. A later session's subscribe for the SAME uri must be + // free to try again upstream, not be dedup-served a fake success. + return resp + } + p.routing.addSubscription(uri, sessID) + return resp +} + +// interceptUnsubscribe handles a session-bearing resources/unsubscribe +// request for uri, called from interceptSessionScopedRequest while the +// caller (handlePost) holds uriLocks for uri. sess is removed from uri's +// recorded subscribers immediately, regardless of what happens next; the +// request is forwarded upstream only if sess was the LAST recorded +// subscriber (the backend should keep sending updates for uri as long as +// anyone still wants them). +// +// Unlike subscribe, a failed/erroring upstream unsubscribe is not rolled +// back: removeSubscription has already run, and there is no OTHER session to +// re-attribute the subscription to even if we wanted to undo it. The worst +// case is the backend keeps a subscription nobody is listening for anymore -- +// wasted upstream notifications that are silently dropped by +// routeResourceUpdated finding no subscribers, not a correctness or security +// issue -- so it is not worth the complexity of reversing removeSubscription +// here. +func (p *HTTPProxy) interceptUnsubscribe(ctx context.Context, sessID, uri string, req *jsonrpc2.Request) jsonrpc2.Message { + last := p.routing.removeSubscription(uri, sessID) + if !last { + return synthesizeSuccessResponse(req.ID) + } + + resp, err := p.forwardUpstream(ctx, sessID, req) + if err != nil { + return upstreamErrorResponse(req.ID, err) + } + return resp +} + +// forwardUpstream wraps ctx with p.requestTimeout and forwards req upstream +// via doRequest, for interceptSubscribe/interceptUnsubscribe's synchronous +// forward performed while still holding handlePost's per-uri uriLocks lock +// for req's uri: the timeout bounds how long that lock can ever be held, so a +// slow or unresponsive backend cannot indefinitely block other sessions' +// resources/subscribe|unsubscribe requests for the SAME uri (see FIX 1 in +// #5744's review). +func (p *HTTPProxy) forwardUpstream(ctx context.Context, sessID string, req *jsonrpc2.Request) (jsonrpc2.Message, error) { + ctx, cancel := context.WithTimeout(ctx, p.requestTimeout) + defer cancel() + return p.doRequest(ctx, sessID, req) +} + +// isErrorResponse reports whether msg is a *jsonrpc2.Response carrying a +// JSON-RPC error, used by interceptSubscribe to distinguish a successful +// upstream resources/subscribe from one the backend rejected. +func isErrorResponse(msg jsonrpc2.Message) bool { + resp, ok := msg.(*jsonrpc2.Response) + return ok && resp.Error != nil +} + +// upstreamErrorResponse builds a JSON-RPC error response for id when the +// upstream round-trip inside interceptSubscribe/interceptUnsubscribe itself +// fails (timeout, proxy shutdown, or a transport-level send error, as opposed +// to a JSON-RPC error the backend returned normally), mirroring the +// timeout-vs-internal error code mapping writeSSEErrorEvent uses for the +// normal request path. +func upstreamErrorResponse(id jsonrpc2.ID, err error) jsonrpc2.Message { + var code int64 = -32603 + msg := "Internal error" + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + code = -32000 + msg = "Timeout" + } + return &jsonrpc2.Response{ID: id, Error: jsonrpc2.NewError(code, msg)} +} + +// synthesizeSuccessResponse builds a locally-generated JSON-RPC success +// response for id with an empty result object, used by +// interceptSessionScopedRequest when a request is served without an upstream +// round-trip. +func synthesizeSuccessResponse(id jsonrpc2.ID) jsonrpc2.Message { + resp, err := jsonrpc2.NewResponse(id, map[string]any{}, nil) + if err != nil { + // marshalToRaw(map[string]any{}) cannot fail; NewResponse's only + // error path is result marshaling. + slog.Error("failed to synthesize success response", "error", err) + return &jsonrpc2.Response{ID: id, Error: jsonrpc2.NewError(-32603, "internal error")} + } + return resp +} + +// writeInterceptedResponse writes msg (from interceptSessionScopedRequest) to +// the client, honoring the same Accept-header-driven JSON-vs-SSE choice and +// Mcp-Session-Id header behavior as the normal handleSingleRequest / +// handleSingleRequestSSE paths, but without any upstream correlation (msg is +// already the complete, final response). +func writeInterceptedResponse( + w http.ResponseWriter, r *http.Request, msg jsonrpc2.Message, sessID string, setSessionHeader bool, +) { + if setSessionHeader { + w.Header().Set("Mcp-Session-Id", sessID) + } + + if !strings.Contains(r.Header.Get("Accept"), "text/event-stream") { + if err := writeJSONRPC(w, msg); err != nil { + slog.Error("failed to write JSON-RPC response", "error", err) + } + return + } + + flusher, ok := startSSEStream(w, nil) + if !ok { + return + } + + data, err := jsonrpc2.EncodeMessage(msg) + if err != nil { + slog.Error("failed to encode JSON-RPC response", "error", err) + return + } + if err := writeSSEData(w, flusher, data); err != nil { + slog.Debug("failed to write response", "error", err) + } +} + +// logLevelRanks assigns each MCP logging level a verbosity rank following the +// RFC 5424 syslog severity numbers MCP's logging capability is aligned with: +// higher number = more verbose = lower severity. "debug" (7) is the most +// verbose; "emergency" (0) is the least. This makes "the maximum rank across +// sessions" exactly the level that satisfies every session's request without +// under-serving the most verbose one (see sessionRouter.setLogLevel). +var logLevelRanks = map[string]int{ + string(sdkmcp.LoggingLevelEmergency): 0, + string(sdkmcp.LoggingLevelAlert): 1, + string(sdkmcp.LoggingLevelCritical): 2, + string(sdkmcp.LoggingLevelError): 3, + string(sdkmcp.LoggingLevelWarning): 4, + string(sdkmcp.LoggingLevelNotice): 5, + string(sdkmcp.LoggingLevelInfo): 6, + string(sdkmcp.LoggingLevelDebug): 7, +} + +// logLevelRank returns level's verbosity rank (see logLevelRanks). An unknown +// or empty level is treated as the MOST verbose (debug's rank): if we can't +// recognize what the client asked for, we would rather over-serve (extra log +// volume) than silently under-serve a session that asked for a level we +// failed to parse. +func logLevelRank(level string) int { + if rank, ok := logLevelRanks[level]; ok { + return rank + } + slog.Debug("unrecognized logging level; treating as most verbose", "level", level) + return logLevelRanks[string(sdkmcp.LoggingLevelDebug)] +} + +// logLevelName returns the level name for rank (the inverse of logLevelRank), +// used to build the reconciled upstream logging/setLevel request. +func logLevelName(rank int) string { + for name, r := range logLevelRanks { + if r == rank { + return name + } + } + return string(sdkmcp.LoggingLevelDebug) +} + +// reconcileUpstreamLogLevel sends a logging/setLevel request upstream with +// level, reconciling the shared backend's single process-wide log level to +// the maximum verbosity requested by any session (see +// interceptSessionScopedRequest). It is fire-and-forget from the calling +// client request's perspective: the reconciliation happens in the background +// via the normal composite-key request/response correlation (so the eventual +// response, if any, is routed and cleaned up correctly instead of leaking a +// waiter), but its result is not surfaced to any client -- the client that +// triggered reconciliation already received its own synthesized success (see +// interceptSessionScopedRequest). +// +// Ordering tradeoff: this reconciliation is best-effort-ordered, NOT +// serialized the way resources/subscribe|unsubscribe is via uriLocks (see FIX +// 1 in #5744's review). Each call runs in its own goroutine, so two +// concurrent maxChanged transitions from different sessions' logging/setLevel +// requests can reach the backend in either order, and the backend could +// briefly end up at a lower verbosity than the CURRENT true max until the +// next transition reconciles it again. This is harmless while +// notifications/message is secure-dropped (see routeNotification): the +// backend's log level is not observable by any client either way, only its +// log volume is affected. Proper ordering (e.g. via a keyed lock analogous to +// uriLocks, or a monotonic sequence check) is deferred to the per-session +// backend / log-delivery follow-up (see the vMCP architecture, #5744). +func (p *HTTPProxy) reconcileUpstreamLogLevel(level string) { + id, err := uuid.NewRandom() + if err != nil { + slog.Warn("failed to mint id for upstream log level reconciliation", "error", err) + return + } + req, err := jsonrpc2.NewCall(jsonrpc2.StringID(id.String()), string(sdkmcp.MethodSetLogLevel), map[string]any{"level": level}) + if err != nil { + slog.Warn("failed to build upstream log level reconciliation request", "error", err) + return + } + + go func() { + // internalReconciliationSessionID keys this request's waiter under a + // session ID no real request ever uses: resolveSessionForRequest + // always assigns a non-empty UUID (real session or per-request + // routing token), so "" can never collide with one. + const internalReconciliationSessionID = "" + ctx, cancel := context.WithTimeout(context.Background(), p.requestTimeout) + defer cancel() + if _, err := p.doRequest(ctx, internalReconciliationSessionID, req); err != nil { + slog.Debug("upstream log level reconciliation did not complete", "error", err) + } + }() +} + +// reapRoutingState periodically purges sessionRouter subscription and +// log-level entries, AND serverStreams' standalone GET streams, for sessions +// that are no longer active, bounding the otherwise-unbounded growth/leakage +// of that state for a session whose owning client disappears without ever +// sending DELETE (see handleDelete's purgeSession/closeStream calls for the +// explicit-teardown path). It ticks at the same cadence as the session +// manager's own cleanup routine (sessionTTL/2, see manager.go's +// cleanupRoutine), and exits when shutdownCh is closed. +// +// isActive is p.sessionManager.Get: the only session liveness check the +// session.Manager/Storage interfaces currently expose. Every Storage +// implementation intentionally refreshes its backend's TTL on every read +// (LocalStorage's Load, RedisStorage's GETEX) to keep genuinely active +// sessions alive -- there is no "peek without refreshing" method today (see +// manager.go and storage.go). Using Get here as isActive means: a session +// that still owns routing state (e.g. an open subscription) has its TTL +// refreshed by this reaper's own liveness check, which can delay that +// session's expiry by up to one reap tick (sessionTTL/2) beyond actual client +// inactivity. This is an accepted, documented tradeoff, not a security gap: +// Get on an ALREADY-deleted session correctly returns false without +// refreshing anything, so reap can never retain routing state past a +// session's true deletion -- it can only delay reaping state for a session +// that reap's own check just observed to still be nominally alive. +func (p *HTTPProxy) reapRoutingState() { + ticker := time.NewTicker(p.sessionTTL / 2) + defer ticker.Stop() + + isActive := func(sess string) bool { + _, ok := p.sessionManager.Get(sess) + return ok + } + + for { + select { + case <-ticker.C: + p.routing.reap(isActive) + p.reapServerStreams(isActive) + case <-p.shutdownCh: + return + } + } +} + +// reapServerStreams closes the standalone GET SSE stream (see +// serverStreamRegistry.closeStream) for any session serverStreams currently +// has registered that isActive no longer reports as live -- a session whose +// owning client disappeared (crashed, or otherwise never sent DELETE) without +// the proxy ever observing r.Context().Done() on its GET request. This mirrors +// handleDelete's explicit p.serverStreams.closeStream(sessID) call for the +// case where the client never disconnects and never calls DELETE either (see +// FIX 2 in #5744's review). +// +// Liveness is resolved for every candidate stream key in a first pass with no +// registry lock held, exactly like sessionRouter.reap's two-phase discipline: +// isActive may perform I/O (e.g. a Redis-backed session.Manager), and running +// it while holding serverStreams.mu would block every concurrent +// broadcast/dispatchTo/deliverToMany/register/deregister call for the +// (possibly slow) duration of every liveness check. closeStream itself only +// ever takes the lock for the in-memory work of removing one already-decided +// stale entry, never across isActive. +func (p *HTTPProxy) reapServerStreams(isActive func(sess string) bool) { + for _, key := range p.serverStreams.streamKeys() { + if isActive(key) { + continue + } + if p.serverStreams.closeStream(key) { + //nolint:gosec // G706: session ID is server-minted, not user-controlled free text + slog.Debug("reaped standalone GET stream for inactive session", "session_id", key) + } + } } // decodeJSONRPCMessage decodes a JSON-RPC message from the request body. @@ -831,12 +1548,22 @@ func decodeJSONRPCMessage(w http.ResponseWriter, body []byte) (jsonrpc2.Message, return msg, true } +// handleNotificationOrClientResponse handles notifications and client responses +// as Legacy today; Modern-aware handling of this path is deferred. func (p *HTTPProxy) handleNotificationOrClientResponse(w http.ResponseWriter, sessID string, msg jsonrpc2.Message) bool { - if isNotification(msg) || (func() bool { _, ok := msg.(*jsonrpc2.Response); return ok })() { + if isNotification(msg) || isClientResponse(msg) { // Refresh TTL so a client sending only notifications doesn't get evicted. if sessID != "" { p.sessionManager.Get(sessID) } + // The backend completed a single handshake and expects a single + // notifications/initialized. Later clients' copies are acknowledged + // here without being forwarded, mirroring the initialize interception + // that answered their handshakes from cache. + if isInitializedNotification(msg) && !p.claimInitializedForward() { + w.WriteHeader(http.StatusAccepted) + return true + } if err := p.SendMessageToDestination(msg); err != nil { slog.Error("failed to send message to destination", "error", err) } diff --git a/pkg/transport/proxy/streamable/streamable_proxy_integration_test.go b/pkg/transport/proxy/streamable/streamable_proxy_integration_test.go index 66208d6e81..899391267b 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy_integration_test.go +++ b/pkg/transport/proxy/streamable/streamable_proxy_integration_test.go @@ -8,8 +8,11 @@ import ( "context" "encoding/json" "fmt" + "io" "net" "net/http" + "runtime" + "sync" "testing" "time" @@ -88,23 +91,17 @@ func TestHTTPRequestIgnoresNotifications(t *testing.T) { assert.Equal(t, "req-123", responseData["id"]) assert.Equal(t, "operation complete", responseData["result"]) - // Test batch request + // A batch request is rejected outright (see TestBatchRequestsRejected and + // #5745), not executed — batching was removed in MCP 2025-06-18. batchJSON := `[{"jsonrpc": "2.0", "method": "test.batch", "id": "batch-1"}]` resp2, err := http.Post(proxyURL, "application/json", bytes.NewReader([]byte(batchJSON))) require.NoError(t, err) defer resp2.Body.Close() - assert.Equal(t, http.StatusOK, resp2.StatusCode) - - var batchResponse []map[string]interface{} - err = json.NewDecoder(resp2.Body).Decode(&batchResponse) + assert.Equal(t, http.StatusBadRequest, resp2.StatusCode) + batchBody, err := io.ReadAll(resp2.Body) require.NoError(t, err) - - // Should have one response (notifications filtered out) - require.Len(t, batchResponse, 1) - assert.Equal(t, "2.0", batchResponse[0]["jsonrpc"]) - assert.Equal(t, "batch-1", batchResponse[0]["id"]) - assert.Equal(t, "operation complete", batchResponse[0]["result"]) + assert.Contains(t, string(batchBody), `"code":-32600`) } // TestHTTPProxy_StartMountsAuthDiscoveryEndpoint is an integration test that @@ -149,3 +146,95 @@ func TestHTTPProxy_StartMountsAuthDiscoveryEndpoint(t *testing.T) { require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) assert.Equal(t, "https://example.com", body["resource"]) } + +// syncMapLen counts entries in a sync.Map for test assertions. +func syncMapLen(m *sync.Map) int { + n := 0 + m.Range(func(_, _ any) bool { n++; return true }) + return n +} + +// TestModernLoadDoesNotAccumulateState drives a burst of concurrent Modern +// (2026-07-28) requests through the full per-request lifecycle -- +// resolveSessionForRequest mints a routing token, doRequest registers a +// waiter/idRestore entry keyed on it and deletes both via its deferred +// cleanup once the response arrives -- and proves nothing is left behind. +// This is the regression test for the unbounded-growth concern behind +// toolhive-core#169. +// +// waiters/idRestore/sessionManager are asserted as hard, exact-zero checks: +// createWaiter's cleanup runs synchronously inside doRequest before +// handleSingleRequest ever writes the HTTP response, so by the time every +// client in the burst has read its response, every cleanup has already run -- +// no settle delay needed, and a regression that skips cleanup fails this +// deterministically. +// +// The goroutine count is checked too, but only as a generous, best-effort +// bound: this package's request path spawns no per-request goroutine, so any +// growth here would come from net/http's own idle-connection handling, not +// from code owned by this package -- hence the loose tolerance rather than a +// strict equality. +// +//nolint:paralleltest // reads process-wide runtime.NumGoroutine() as a baseline and starts an HTTP proxy; must run serially +func TestModernLoadDoesNotAccumulateState(t *testing.T) { + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + t.Cleanup(func() { _ = proxy.Stop(ctx) }) + + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + + const requestCount = 1000 + // Cap in-flight requests well under the proxy's 100-slot message-channel + // buffer. Firing all 1000 at once can exceed that buffer and produce + // dropped-response 504s (#5952, a proxy backpressure gap) -- + // unrelated to the accumulation invariant this test checks, so throttle + // around it rather than let it flake the run. + const maxInFlight = 50 + sem := make(chan struct{}, maxInFlight) + baseline := runtime.NumGoroutine() + + var wg sync.WaitGroup + for i := 0; i < requestCount; i++ { + wg.Add(1) + sem <- struct{}{} + go func(id int) { + defer wg.Done() + defer func() { <-sem }() + + req, err := http.NewRequest( + http.MethodPost, url, bytes.NewReader([]byte(modernBody(id, "tools/list")))) + if !assert.NoError(t, err) { + return + } + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) { + return + } + defer func() { + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + }() + assert.Equal(t, http.StatusOK, resp.StatusCode) + }(i) + } + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("timeout waiting for the request burst to complete") + } + + assert.Zero(t, proxy.sessionManager.Count(), "Modern requests must never register a session") + assert.Zero(t, syncMapLen(&proxy.waiters), "waiters must be fully cleaned up once all responses are delivered") + assert.Zero(t, syncMapLen(&proxy.idRestore), "idRestore must be fully cleaned up once all responses are delivered") + + http.DefaultClient.CloseIdleConnections() + assert.Eventually(t, func() bool { + return runtime.NumGoroutine() <= baseline+20 + }, time.Second, 20*time.Millisecond, "goroutine count did not settle back near baseline after the burst") +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go b/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go index f085b12ad8..475ead6d05 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go +++ b/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go @@ -11,11 +11,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) const ( diff --git a/pkg/transport/proxy/streamable/streamable_proxy_modern_test.go b/pkg/transport/proxy/streamable/streamable_proxy_modern_test.go new file mode 100644 index 0000000000..d226792cbd --- /dev/null +++ b/pkg/transport/proxy/streamable/streamable_proxy_modern_test.go @@ -0,0 +1,292 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" + + "github.com/stacklok/toolhive/pkg/mcp" +) + +// modernBodyWithVersion builds a single Modern JSON-RPC request body carrying the +// reserved io.modelcontextprotocol/* _meta keys ClassifyRevision inspects. +func modernBodyWithVersion(id int, method, version string) string { + return fmt.Sprintf( + `{"jsonrpc":"2.0","id":%d,"method":%q,"params":{"_meta":{`+ + `"io.modelcontextprotocol/protocolVersion":%q,`+ + `"io.modelcontextprotocol/clientCapabilities":{}}}}`, + id, method, version) +} + +// modernBody builds a well-formed Modern request for the version this build supports. +func modernBody(id int, method string) string { + return modernBodyWithVersion(id, method, "2026-07-28") +} + +// TestModernConcurrentRequestsAreNotMixed is the Modern analogue of +// TestSessionlessConcurrentRequestsAreNotMixed, and the critical confidentiality +// guard for the stateless routing path: two concurrent Modern POSTs sharing +// JSON-RPC id 1 must each receive their own payload. One request additionally +// carries a foreign Mcp-Session-Id header — it must NOT fall through to the +// shared-key session path (which would collapse both requests onto +// compositeKey(sessID, idKey) and leak one client's response to the other). +// +// t.Setenv requires a non-parallel test; the short proxy timeout makes a +// regression fail fast rather than hanging on the 60s default. +func TestModernConcurrentRequestsAreNotMixed(t *testing.T) { + t.Setenv(proxyRequestTimeoutEnv, "3s") + + port := pickFreePort(t) + proxy := NewHTTPProxy("127.0.0.1", port, nil, nil) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + require.NoError(t, proxy.Start(ctx)) + t.Cleanup(func() { _ = proxy.Stop(ctx) }) + + time.Sleep(50 * time.Millisecond) + + // Barrier backend: buffer both requests before responding, so both proxy + // waiters are registered first — the precondition under which a shared + // routing token would overwrite one waiter. Echo req.Method to prove each + // client got its own payload. + go func() { + buffered := make([]*jsonrpc2.Request, 0, 2) + for len(buffered) < 2 { + select { + case msg := <-proxy.GetMessageChannel(): + if req, ok := msg.(*jsonrpc2.Request); ok && req.ID.IsValid() { + buffered = append(buffered, req) + } + case <-ctx.Done(): + return + } + } + for _, req := range buffered { + result := map[string]any{"echoed_method": req.Method} + resp, _ := jsonrpc2.NewResponse(req.ID, result, nil) + _ = proxy.ForwardResponseToClients(ctx, resp) + } + }() + + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + + type result struct { + method string + status int + body map[string]any + err error + } + + // Both POSTs are Modern and share JSON-RPC id 1. sessionID is a foreign + // Mcp-Session-Id header the Modern path must ignore; empty means no header. + fire := func(method, sessionID string) result { + req, err := http.NewRequest( + http.MethodPost, url, bytes.NewReader([]byte(modernBody(1, method)))) + if err != nil { + return result{method: method, err: err} + } + req.Header.Set("Content-Type", "application/json") + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + client := &http.Client{Timeout: 8 * time.Second} + resp, err := client.Do(req) + if err != nil { + return result{method: method, err: err} + } + defer func() { + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + }() + var decoded map[string]any + _ = json.NewDecoder(resp.Body).Decode(&decoded) + return result{method: method, status: resp.StatusCode, body: decoded} + } + + resCh := make(chan result, 2) + go func() { resCh <- fire("tools/list", "foreign-session-id") }() + go func() { resCh <- fire("resources/list", "") }() + + received := map[string]result{} + for i := 0; i < 2; i++ { + select { + case r := <-resCh: + received[r.method] = r + case <-time.After(15 * time.Second): + t.Fatalf("timeout waiting for concurrent Modern responses; received so far: %v", received) + } + } + + for _, method := range []string{"tools/list", "resources/list"} { + r := received[method] + require.NoError(t, r.err, "client %q HTTP error", method) + require.Equal(t, http.StatusOK, r.status, "client %q HTTP status", method) + require.NotNil(t, r.body, "client %q empty body", method) + res, ok := r.body["result"].(map[string]any) + require.True(t, ok, "client %q missing result: %v", method, r.body) + assert.Equal(t, method, res["echoed_method"], + "client %q received the other client's payload (response cross-talk)", method) + } +} + +// TestModernRequestIgnoresClientSessionID verifies that a Modern POST carrying a +// bogus Mcp-Session-Id is served statelessly: it returns 200 (a Legacy request +// with an unknown session id returns 404, so status discriminates the paths), +// sets no Mcp-Session-Id response header, and never registers the foreign id +// with the session manager. +func TestModernRequestIgnoresClientSessionID(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + t.Cleanup(func() { _ = proxy.Stop(ctx) }) + + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + + const bogusSession = "bogus-session-id" + req, err := http.NewRequest( + http.MethodPost, url, bytes.NewReader([]byte(modernBody(1, "tools/list")))) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", bogusSession) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Empty(t, resp.Header.Get("Mcp-Session-Id"), + "Modern response must not set a session header") + + _, ok := proxy.sessionManager.Get(bogusSession) + assert.False(t, ok, "Modern path must not register the foreign session id") +} + +// TestModernNeverReusesClientSessionIDAsRoutingToken is the mutation-check +// analogue of the transparent proxy's TestGuardUnknownSessionFiresDespiteForgedModernRevision. +// Where that test proves an UNKNOWN session ID still triggers the Legacy +// guard, this one closes the complementary gap: it registers a REAL, known +// session in sessionManager and proves resolveSessionForRequest still mints +// a fresh routing token instead of reusing it. +// +// TestModernRequestIgnoresClientSessionID (above) only exercises a bogus, +// unregistered session ID, so it cannot detect a regression that consults +// sessionManager as a fallback before minting: on a lookup miss such code +// would still mint a fresh token and pass that test. Using a real registered +// session here closes that hole -- this test FAILS if the Modern branch is +// ever re-keyed on Mcp-Session-Id presence (i.e. reuses a known client SID as +// the routing token, or calls sessionManager.Get/AddWithID for it at all). +func TestModernNeverReusesClientSessionIDAsRoutingToken(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("127.0.0.1", 0, nil, nil) + t.Cleanup(func() { _ = proxy.Stop(context.Background()) }) + + knownSessionID := uuid.New().String() + require.NoError(t, proxy.sessionManager.AddWithID(knownSessionID)) + + newModernRequest := func() *jsonrpc2.Request { + req, err := jsonrpc2.NewCall(jsonrpc2.Int64ID(1), "tools/call", json.RawMessage( + `{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28",`+ + `"io.modelcontextprotocol/clientCapabilities":{}}}`)) + require.NoError(t, err) + return req + } + + resolve := func() string { + httpReq := httptest.NewRequest(http.MethodPost, StreamableHTTPEndpoint, http.NoBody) + httpReq.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + httpReq.Header.Set("Mcp-Session-Id", knownSessionID) + + token, setHeader, err := proxy.resolveSessionForRequest(httptest.NewRecorder(), httpReq, newModernRequest()) + require.NoError(t, err) + assert.False(t, setHeader, "Modern branch must never ask the client to adopt a session") + return token + } + + token1 := resolve() + token2 := resolve() + + assert.NotEqual(t, knownSessionID, token1, "Modern branch must not reuse the known client session id as its routing token") + assert.NotEqual(t, knownSessionID, token2, "Modern branch must not reuse the known client session id as its routing token") + assert.NotEqual(t, token1, token2, "each Modern request must mint its own fresh routing token") + + _, err := uuid.Parse(token1) + assert.NoError(t, err, "routing token must be a freshly minted UUID, not a passthrough of client input") + _, err = uuid.Parse(token2) + assert.NoError(t, err, "routing token must be a freshly minted UUID, not a passthrough of client input") +} + +// TestModernClassificationErrorsReturn400 verifies that ClassifyRevision errors +// on the single-request path are rendered as HTTP 400 JSON-RPC error responses +// with the spec-defined error code. +func TestModernClassificationErrorsReturn400(t *testing.T) { + t.Parallel() + + port := pickFreePort(t) + proxy, ctx, cancel := startProxyWithBackend(t, port) + t.Cleanup(cancel) + t.Cleanup(func() { _ = proxy.Stop(ctx) }) + + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + + tests := []struct { + name string + body string + protoHeader string + wantCode int64 + }{ + { + name: "header/body version mismatch", + body: modernBodyWithVersion(1, "tools/list", "2026-07-28"), + protoHeader: "2025-11-25", + wantCode: mcp.CodeHeaderMismatch, + }, + { + name: "unsupported body version", + body: modernBodyWithVersion(1, "tools/list", "1999-01-01"), + wantCode: mcp.CodeUnsupportedProtocolVersion, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(tt.body))) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + if tt.protoHeader != "" { + req.Header.Set("MCP-Protocol-Version", tt.protoHeader) + } + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var decoded struct { + Error struct { + Code int64 `json:"code"` + } `json:"error"` + } + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + assert.Equal(t, tt.wantCode, decoded.Error.Code) + }) + } +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy_protocol_version_test.go b/pkg/transport/proxy/streamable/streamable_proxy_protocol_version_test.go new file mode 100644 index 0000000000..52572901ae --- /dev/null +++ b/pkg/transport/proxy/streamable/streamable_proxy_protocol_version_test.go @@ -0,0 +1,81 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestHandlePost_StrictProtocolValidation drives handlePost directly (no +// server/backend needed: the notification body used below returns 202 +// without any upstream round-trip) to verify the MCP-Protocol-Version gate: +// - strict mode rejects a present-but-unrecognized version with 400 +// - strict mode accepts a recognized version +// - strict mode accepts an absent header (assume 2025-03-26 per spec) +// - permissive mode (default, strict disabled) accepts any version, +// preserving the proxy's original version-agnostic behavior +func TestHandlePost_StrictProtocolValidation(t *testing.T) { + t.Parallel() + + // A notification (no id) is accepted and forwarded with 202 by + // handleNotificationOrClientResponse without waiting on any backend + // response, so these tests never need to start the proxy or a backend. + const notificationBody = `{"jsonrpc":"2.0","method":"progress","params":{"pct":50}}` + + tests := []struct { + name string + strict bool + protocolHeader string + wantStatus int + }{ + { + name: "strict on, unsupported version rejected", + strict: true, + protocolHeader: "1999-01-01", + wantStatus: http.StatusBadRequest, + }, + { + name: "strict on, supported version accepted", + strict: true, + protocolHeader: "2025-11-25", + wantStatus: http.StatusAccepted, + }, + { + name: "strict on, absent header accepted", + strict: true, + protocolHeader: "", + wantStatus: http.StatusAccepted, + }, + { + name: "strict off (default), unsupported version accepted", + strict: false, + protocolHeader: "1999-01-01", + wantStatus: http.StatusAccepted, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("127.0.0.1", 0, nil, nil, WithStrictProtocolValidation(tt.strict)) + + req := httptest.NewRequest(http.MethodPost, StreamableHTTPEndpoint, bytes.NewReader([]byte(notificationBody))) + req.Header.Set("Content-Type", "application/json") + if tt.protocolHeader != "" { + req.Header.Set("MCP-Protocol-Version", tt.protocolHeader) + } + rec := httptest.NewRecorder() + + proxy.handlePost(rec, req) + + require.Equal(t, tt.wantStatus, rec.Code) + }) + } +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go b/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go index 1574462661..74b290b1b8 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go +++ b/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go @@ -11,6 +11,7 @@ import ( "io" "net" "net/http" + "net/http/httptest" "testing" "time" @@ -53,12 +54,15 @@ func startProxyWithBackend(t *testing.T, port int) (*HTTPProxy, context.Context, return proxy, ctx, cancel } -// TestGETReturns405 validates that GET on MCP endpoint returns 405 (server does not offer SSE here). -func TestGETReturns405(t *testing.T) { +// TestGETReturns405WhenStandaloneSSEDisabled validates that GET on the MCP +// endpoint returns 405 when standalone SSE is explicitly disabled via +// WithStandaloneSSE(false). By default GET opens a standalone SSE stream +// instead -- see streamable_proxy_test.go's TestHandleGet_* for that behavior. +func TestGETReturns405WhenStandaloneSSEDisabled(t *testing.T) { t.Parallel() const port = 8101 - proxy := NewHTTPProxy("127.0.0.1", port, nil, nil) + proxy := NewHTTPProxy("127.0.0.1", port, nil, nil, WithStandaloneSSE(false)) ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -173,69 +177,96 @@ func TestPOSTNotificationOnlyAccepted(t *testing.T) { assert.Equal(t, 0, len(body), "202 should have no body") } -// TestBatchOnlyNotificationsAccepted returns 202 and no body per spec when no requests are present. -func TestBatchOnlyNotificationsAccepted(t *testing.T) { +// TestBatchRequestsRejected verifies that the streamable proxy rejects every +// JSON-RPC batch (a top-level array) with an HTTP 400 JSON-RPC "Invalid Request" +// (-32600) error instead of executing it. Batching was removed in MCP 2025-06-18 +// and ToolHive serves only 2025-11-25 and 2026-07-28, so a batch must never +// reach the backend where its nested calls would bypass authz/audit/tool +// filtering (see #5745). A backend is deliberately not started: reaching it +// would itself be the bug under test. +func TestBatchRequestsRejected(t *testing.T) { t.Parallel() const port = 8105 proxy := NewHTTPProxy("127.0.0.1", port, nil, nil) ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + t.Cleanup(cancel) require.NoError(t, proxy.Start(ctx), "proxy start") - defer func() { _ = proxy.Stop(ctx) }() + t.Cleanup(func() { _ = proxy.Stop(ctx) }) time.Sleep(50 * time.Millisecond) - url := "http://127.0.0.1:8105" + StreamableHTTPEndpoint - // Batch of only notifications (no ids) - batch := `[{"jsonrpc":"2.0","method":"progress","params":{"pct":10}},{"jsonrpc":"2.0","method":"progress","params":{"pct":20}}]` - req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(batch))) - req.Header.Set("Content-Type", "application/json") - - resp, err := http.DefaultClient.Do(req) - require.NoError(t, err) - defer resp.Body.Close() - - assert.Equal(t, http.StatusAccepted, resp.StatusCode) - body, _ := io.ReadAll(resp.Body) - assert.Equal(t, 0, len(body)) -} - -// TestBatchMixedNotificationsAndRequest returns an array with one response corresponding to the request id. -func TestBatchMixedNotificationsAndRequest(t *testing.T) { - t.Parallel() - - const port = 8106 - proxy, ctx, cancel := startProxyWithBackend(t, port) - defer cancel() - defer func() { _ = proxy.Stop(ctx) }() - - url := "http://127.0.0.1:8106" + StreamableHTTPEndpoint - - // Batch includes a notification and a request "r1" - batch := `[{"jsonrpc":"2.0","method":"progress","params":{"pct":99}},{"jsonrpc":"2.0","id":"r1","method":"tools/list","params":{}}]` - req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(batch))) - req.Header.Set("Content-Type", "application/json") - - resp, err := http.DefaultClient.Do(req) - require.NoError(t, err) - defer resp.Body.Close() + tests := []struct { + name string + body string + contentType string + sessionID string + }{ + { + name: "only notifications", + body: `[{"jsonrpc":"2.0","method":"progress","params":{"pct":10}},{"jsonrpc":"2.0","method":"progress","params":{"pct":20}}]`, + contentType: "application/json", + }, + { + name: "mixed notification and request", + body: `[{"jsonrpc":"2.0","method":"progress"},{"jsonrpc":"2.0","id":"r1","method":"tools/list","params":{}}]`, + contentType: "application/json", + }, + { + name: "with stale session id", + body: `[{"jsonrpc":"2.0","id":"b1","method":"tools/list"},{"jsonrpc":"2.0","id":"b2","method":"tools/list"}]`, + contentType: "application/json", + sessionID: "expired-session-id", + }, + { + // Regression for #5745: a batch smuggled under a non-JSON content + // type skips ParsingMiddleware/authz but must still be rejected at + // the executor. + name: "non-json content type", + body: `[{"jsonrpc":"2.0","id":"c1","method":"tools/call","params":{"name":"danger"}}]`, + contentType: "text/plain", + }, + { + // Regression for #5745: leading Unicode whitespace must not let a + // batch evade detection (detector and json decoder must agree). + name: "leading vertical tab", + body: "\v[{\"jsonrpc\":\"2.0\",\"id\":\"d1\",\"method\":\"tools/call\"}]", + contentType: "application/json", + }, + } - assert.Equal(t, http.StatusOK, resp.StatusCode) - assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(tt.body))) + require.NoError(t, err) + req.Header.Set("Content-Type", tt.contentType) + if tt.sessionID != "" { + req.Header.Set("Mcp-Session-Id", tt.sessionID) + } - var arr []map[string]any - err = json.NewDecoder(resp.Body).Decode(&arr) - require.NoError(t, err) - require.Len(t, arr, 1) - - // Response should correspond to id "r1" - assert.Equal(t, "2.0", arr[0]["jsonrpc"]) - assert.Equal(t, "r1", arr[0]["id"]) - // And include a result map as sent by backend - _, hasResult := arr[0]["result"] - assert.True(t, hasResult, "batch response should include result") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"code":-32600`) + + // A batch has no single request id to echo. MCP encodes that by + // omitting the "id" key entirely (schema/2025-11-25 types the + // error response id as optional, not nullable), never as null -- + // a substring check for "id":null would miss the key being + // present-but-null, so the key's absence is checked directly. + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + _, hasID := parsed["id"] + assert.False(t, hasID, `"id" key must be omitted, not present as null`) + }) + } } // TestDeleteUnknownSessionReturnsJSONRPCError verifies that DELETE with an @@ -261,23 +292,30 @@ func TestDeleteUnknownSessionReturnsJSONRPCError(t *testing.T) { body, err := io.ReadAll(resp.Body) require.NoError(t, err) assert.Contains(t, string(body), `"code":-32001`) - assert.Contains(t, string(body), `"id":null`) + + // A DELETE carries no JSON-RPC id, and MCP encodes an absent id by + // omitting the "id" key entirely (schema/2025-11-25), never as null. + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + _, ok := parsed["id"] + assert.False(t, ok, `"id" key should be absent`) } -// TestBatchWithStaleSessionReturnsJSONRPCError verifies that a batch POST -// with a stale session ID returns HTTP 404 with a JSON-RPC error body. -func TestBatchWithStaleSessionReturnsJSONRPCError(t *testing.T) { +// TestSingleRequestWithStaleSessionIncludesRequestID verifies that a single +// POST with a stale session ID returns a JSON-RPC error that includes the +// request's original ID. +func TestSingleRequestWithStaleSessionIncludesRequestID(t *testing.T) { t.Parallel() - const port = 8108 + const port = 8109 proxy, ctx, cancel := startProxyWithBackend(t, port) defer cancel() defer func() { _ = proxy.Stop(ctx) }() - url := "http://127.0.0.1:8108" + StreamableHTTPEndpoint + url := "http://127.0.0.1:8109" + StreamableHTTPEndpoint - batch := `[{"jsonrpc":"2.0","id":"b1","method":"tools/list","params":{}},{"jsonrpc":"2.0","id":"b2","method":"tools/list","params":{}}]` - req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(batch))) + reqJSON := `{"jsonrpc":"2.0","id":"test-42","method":"tools/list","params":{}}` + req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(reqJSON))) req.Header.Set("Content-Type", "application/json") req.Header.Set("Mcp-Session-Id", "expired-session-id") @@ -287,40 +325,72 @@ func TestBatchWithStaleSessionReturnsJSONRPCError(t *testing.T) { assert.Equal(t, http.StatusNotFound, resp.StatusCode) assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) + body, err := io.ReadAll(resp.Body) require.NoError(t, err) assert.Contains(t, string(body), `"code":-32001`) + assert.Contains(t, string(body), `"id":"test-42"`) } -// TestSingleRequestWithStaleSessionIncludesRequestID verifies that a single -// POST with a stale session ID returns a JSON-RPC error that includes the -// request's original ID. -func TestSingleRequestWithStaleSessionIncludesRequestID(t *testing.T) { +// TestSSEResponseIncludesEventMessage verifies that a POST with +// Accept: text/event-stream returns an SSE frame that starts with +// "event: message" before the JSON-RPC data line. Regression for #5655: +// data-only frames hang clients that require an explicit event name. +func TestSSEResponseIncludesEventMessage(t *testing.T) { t.Parallel() - const port = 8109 + port := pickFreePort(t) proxy, ctx, cancel := startProxyWithBackend(t, port) - defer cancel() - defer func() { _ = proxy.Stop(ctx) }() + t.Cleanup(cancel) + t.Cleanup(func() { _ = proxy.Stop(ctx) }) - url := "http://127.0.0.1:8109" + StreamableHTTPEndpoint + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, StreamableHTTPEndpoint) + body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"sse-event-test","version":"1.0"}}}` - reqJSON := `{"jsonrpc":"2.0","id":"test-42","method":"tools/list","params":{}}` - req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(reqJSON))) + req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(body))) + require.NoError(t, err) req.Header.Set("Content-Type", "application/json") - req.Header.Set("Mcp-Session-Id", "expired-session-id") + req.Header.Set("Accept", "application/json, text/event-stream") - resp, err := http.DefaultClient.Do(req) + client := &http.Client{Timeout: 10 * time.Second} + resp, err := client.Do(req) require.NoError(t, err) defer resp.Body.Close() - assert.Equal(t, http.StatusNotFound, resp.StatusCode) - assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) + require.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) - body, err := io.ReadAll(resp.Body) + raw, err := io.ReadAll(resp.Body) require.NoError(t, err) - assert.Contains(t, string(body), `"code":-32001`) - assert.Contains(t, string(body), `"id":"test-42"`) + got := string(raw) + + assert.True(t, bytes.HasPrefix(raw, []byte("event: message\n")), + "SSE body must start with event: message; got %q", got) + assert.Contains(t, got, "\ndata: ") + assert.Contains(t, got, `"jsonrpc":"2.0"`) + assert.Contains(t, got, `"id":1`) + // Defend against a regression that writes data: first then event: later. + eventIdx := bytes.Index(raw, []byte("event: message\n")) + dataIdx := bytes.Index(raw, []byte("\ndata: ")) + require.GreaterOrEqual(t, eventIdx, 0) + require.Greater(t, dataIdx, eventIdx, "event: message must appear before data:") +} + +// TestSSEErrorEventIncludesEventMessage verifies the error SSE writer uses the +// same event: message framing as the success path (writeSSEData). +func TestSSEErrorEventIncludesEventMessage(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeSSEErrorEvent(rec, rec, jsonrpc2.StringID("7"), context.DeadlineExceeded) + + raw := rec.Body.Bytes() + got := string(raw) + assert.True(t, bytes.HasPrefix(raw, []byte("event: message\n")), + "error SSE frame must start with event: message; got %q", got) + assert.Contains(t, got, `"id":"7"`) + assert.Contains(t, got, `"code":-32000`) + assert.Contains(t, got, `"Timeout"`) } // pickFreePort returns a TCP port the OS reports as available. There is a small diff --git a/pkg/transport/proxy/streamable/streamable_proxy_test.go b/pkg/transport/proxy/streamable/streamable_proxy_test.go index 3f27d7636e..0a3908c2f3 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy_test.go +++ b/pkg/transport/proxy/streamable/streamable_proxy_test.go @@ -4,12 +4,18 @@ package streamable import ( + "bytes" "context" + "encoding/json" + "errors" "fmt" "net/http" + "net/http/httptest" + "strings" "testing" "time" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" @@ -361,6 +367,7 @@ func TestNewHTTPProxy_PrefixHandlerDoesNotShadowMCP(t *testing.T) { port := getFreePort(t) proxy := NewHTTPProxy("localhost", port, nil, nil, WithPrefixHandlers(map[string]http.Handler{"/oauth/token": sentinel}), + WithStandaloneSSE(false), // keep the 405 GET response so this test terminates promptly ) ctx := t.Context() @@ -432,3 +439,202 @@ func TestNewHTTPProxy_ExactWellKnownBeatsSubtree(t *testing.T) { assert.Equal(t, http.StatusOK, resp2.StatusCode, "/.well-known/ subtree must still reach authInfoHandler") } + +// nonFlushableResponseWriter implements http.ResponseWriter but deliberately +// does NOT implement http.Flusher, to exercise handleGet's capability check. +type nonFlushableResponseWriter struct { + header http.Header + body bytes.Buffer + statusCode int +} + +func newNonFlushableResponseWriter() *nonFlushableResponseWriter { + return &nonFlushableResponseWriter{header: make(http.Header)} +} + +func (w *nonFlushableResponseWriter) Header() http.Header { return w.header } + +func (w *nonFlushableResponseWriter) Write(b []byte) (int, error) { return w.body.Write(b) } + +func (w *nonFlushableResponseWriter) WriteHeader(statusCode int) { w.statusCode = statusCode } + +// TestHandleGet_StandaloneSSEDisabledReturns405 verifies handleGet returns 405 +// when the proxy was constructed with WithStandaloneSSE(false), regardless of +// whether the underlying ResponseWriter supports flushing. +func TestHandleGet_StandaloneSSEDisabledReturns405(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("localhost", 0, nil, nil, WithStandaloneSSE(false)) + req := httptest.NewRequest(http.MethodGet, StreamableHTTPEndpoint, nil) + rec := httptest.NewRecorder() + + proxy.handleGet(rec, req) + + assert.Equal(t, http.StatusMethodNotAllowed, rec.Code) + assert.Equal(t, 0, proxy.serverStreams.streamCount(), "no stream should be registered when SSE is disabled") +} + +// TestHandleGet_NonFlushableWriterReturns500 verifies handleGet fails loudly +// with 500 when the ResponseWriter does not support streaming, rather than +// silently degrading (standaloneSSE defaults to true here). +func TestHandleGet_NonFlushableWriterReturns500(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("localhost", 0, nil, nil) + req := httptest.NewRequest(http.MethodGet, StreamableHTTPEndpoint, nil) + rec := newNonFlushableResponseWriter() + + proxy.handleGet(rec, req) + + assert.Equal(t, http.StatusInternalServerError, rec.statusCode) + assert.Equal(t, 0, proxy.serverStreams.streamCount(), "stream must not remain registered after the capability check fails") +} + +// TestHandleGet_DefaultEnabledRegistersAndDeregistersStream verifies that, +// with standaloneSSE at its default (enabled) and a known Mcp-Session-Id, a +// GET request registers a stream for the lifetime of the connection and +// deregisters it on client disconnect (request context cancellation). +func TestHandleGet_DefaultEnabledRegistersAndDeregistersStream(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("localhost", 0, nil, nil) + sessID := uuid.NewString() + require.NoError(t, proxy.sessionManager.AddWithID(sessID)) + + ctx, cancel := context.WithCancel(context.Background()) + req := httptest.NewRequest(http.MethodGet, StreamableHTTPEndpoint, nil).WithContext(ctx) + req.Header.Set("Mcp-Session-Id", sessID) + rec := httptest.NewRecorder() + + done := make(chan struct{}) + go func() { + defer close(done) + proxy.handleGet(rec, req) + }() + + require.Eventually(t, func() bool { + return proxy.serverStreams.streamCount() == 1 + }, time.Second, 5*time.Millisecond, "stream was not registered") + + cancel() // simulate client disconnect + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("handleGet did not return after client disconnect") + } + + assert.Equal(t, 0, proxy.serverStreams.streamCount(), "stream must be deregistered after disconnect") + assert.Equal(t, "text/event-stream", rec.Header().Get("Content-Type")) +} + +// TestHandleGet_SessionDecisionMatrix verifies handleGet's Mcp-Session-Id +// resolution matrix directly (without a full HTTP round-trip): no header is +// rejected with 400, an unknown session with 404, and a known session opens +// the stream with 200. +func TestHandleGet_SessionDecisionMatrix(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + setHeader bool + useKnownSess bool + wantStatus int + wantStreamReg bool + }{ + {name: "no header", setHeader: false, wantStatus: http.StatusBadRequest}, + {name: "unknown session", setHeader: true, useKnownSess: false, wantStatus: http.StatusNotFound}, + {name: "known session", setHeader: true, useKnownSess: true, wantStatus: http.StatusOK, wantStreamReg: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + proxy := NewHTTPProxy("localhost", 0, nil, nil) + sessID := uuid.NewString() + require.NoError(t, proxy.sessionManager.AddWithID(sessID)) + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + req := httptest.NewRequest(http.MethodGet, StreamableHTTPEndpoint, nil).WithContext(ctx) + if tt.setHeader { + if tt.useKnownSess { + req.Header.Set("Mcp-Session-Id", sessID) + } else { + req.Header.Set("Mcp-Session-Id", uuid.NewString()) + } + } + rec := httptest.NewRecorder() + + done := make(chan struct{}) + go func() { + defer close(done) + proxy.handleGet(rec, req) + }() + + if tt.wantStreamReg { + require.Eventually(t, func() bool { + return proxy.serverStreams.streamCount() == 1 + }, time.Second, 5*time.Millisecond, "stream was not registered") + cancel() + } + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("handleGet did not return") + } + + assert.Equal(t, tt.wantStatus, rec.Code) + }) + } +} + +// TestWriteSSEErrorEvent verifies the "id" key is present and echoed for a +// valid jsonrpc2.ID, and omitted entirely -- never sent as "id":null -- for +// the zero-value ID (jsonrpc2.ID{}.IsValid() == false), in both cases +// preserving the SSE "event: message" + "data: ..." framing writeSSEData produces. +// +// The body is decoded to map[string]any rather than a tagged struct: a +// struct field of type `any` cannot distinguish an absent key from a +// present-but-null one, so it would not detect a regression to "id":null. +func TestWriteSSEErrorEvent(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + id jsonrpc2.ID + wantID bool + }{ + {name: "valid id is present and echoed", id: jsonrpc2.Int64ID(7), wantID: true}, + {name: "zero-value id is omitted, not sent as null", id: jsonrpc2.ID{}, wantID: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + rec := httptest.NewRecorder() + + writeSSEErrorEvent(rec, rec, tt.id, errors.New("boom")) + + assert.True(t, rec.Flushed, "writeSSEErrorEvent must flush the SSE frame to the client") + + body := rec.Body.String() + require.True(t, strings.HasPrefix(body, "event: message\n"), "SSE frame must include event: message, got %q", body) + require.True(t, strings.HasSuffix(body, "\n\n"), "SSE frame must end with a blank line, got %q", body) + lines := strings.Split(strings.TrimSuffix(body, "\n\n"), "\n") + require.Len(t, lines, 2, "expected event + data lines, got %q", body) + require.Equal(t, "event: message", lines[0]) + payload := strings.TrimPrefix(lines[1], "data: ") + + var parsed map[string]any + require.NoError(t, json.Unmarshal([]byte(payload), &parsed)) + _, hasID := parsed["id"] + assert.Equal(t, tt.wantID, hasID) + if tt.wantID { + assert.EqualValues(t, tt.id.Raw(), parsed["id"]) + } + }) + } +} diff --git a/pkg/transport/proxy/streamable/utils.go b/pkg/transport/proxy/streamable/utils.go index 1281cf970c..886816df3e 100644 --- a/pkg/transport/proxy/streamable/utils.go +++ b/pkg/transport/proxy/streamable/utils.go @@ -6,12 +6,22 @@ package streamable import ( "encoding/json" "fmt" + "io" "log/slog" "net/http" + "time" "golang.org/x/exp/jsonrpc2" + + "github.com/stacklok/toolhive/pkg/mcp" ) +// sseKeepAliveInterval is the cadence at which an otherwise-idle SSE stream +// writes a ": keep-alive\n\n" comment (see writeSSEKeepAlive), so +// intermediary proxies/load balancers do not close the connection for +// inactivity. +const sseKeepAliveInterval = 30 * time.Second + // isNotification returns true if the JSON-RPC message is a notification (no ID). func isNotification(msg jsonrpc2.Message) bool { if req, ok := msg.(*jsonrpc2.Request); ok { @@ -20,6 +30,15 @@ func isNotification(msg jsonrpc2.Message) bool { return false } +// isClientResponse returns true if msg is a *jsonrpc2.Response, i.e. a +// client's reply to a server-initiated request (see +// rejectServerRequestToBackend) rather than a client-initiated request or +// notification. +func isClientResponse(msg jsonrpc2.Message) bool { + _, ok := msg.(*jsonrpc2.Response) + return ok +} + // writeHTTPError writes a plain HTTP error with status. func writeHTTPError(w http.ResponseWriter, status int, msg string) { http.Error(w, msg, status) @@ -46,6 +65,87 @@ func writeJSONRPCError(w http.ResponseWriter, id jsonrpc2.ID, msg string) { } } +// writeSSEData writes a single JSON-RPC SSE frame with an explicit +// "event: message" name plus a "data:" line, then flushes it. The event name +// matches MCP reference transports and ToolHive's own ssecommon serializer; +// clients that only dispatch named "message" events (for example @ai-sdk/mcp) +// drop data-only frames. It returns the underlying write error (if any) so +// the caller can decide whether to end the stream; it does not itself log or +// close anything. +func writeSSEData(w io.Writer, flusher http.Flusher, data []byte) error { + //nolint:gosec // G705: SSE data from MCP protocol + if _, err := fmt.Fprintf(w, "event: message\ndata: %s\n\n", data); err != nil { + return err + } + flusher.Flush() + return nil +} + +// assertFlushable type-asserts w as http.Flusher, which every SSE response in +// this package requires so data can be pushed to the client as it is +// written, rather than buffered until the handler returns. On failure it +// writes the standard 500 "Streaming not supported" error to w and returns +// (nil, false); the caller must return immediately without writing anything +// else. On success it returns (w.(http.Flusher), true). +func assertFlushable(w http.ResponseWriter) (http.Flusher, bool) { + flusher, ok := w.(http.Flusher) + if !ok { + writeHTTPError(w, http.StatusInternalServerError, "Streaming not supported") + return nil, false + } + return flusher, true +} + +// setSSEHeaders sets the standard response headers for an SSE stream: +// Content-Type: text/event-stream, Cache-Control: no-cache, and +// Connection: keep-alive. +func setSSEHeaders(w http.ResponseWriter) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") +} + +// startSSEStream asserts that w supports flushing (see assertFlushable), sets +// the standard SSE headers (see setSSEHeaders), and flushes them to the +// client immediately so it knows the stream is open without waiting for the +// first data: frame. +// +// If extra is non-nil, it is called after the SSE headers are set but before +// the flush, letting a caller add further response headers (e.g. +// Mcp-Session-Id) while they are still guaranteed to reach the client: +// net/http locks in headers at the first Flush/WriteHeader call, so any +// header set after this function returns would be silently dropped. +// +// It returns (flusher, true) on success. If w does not implement +// http.Flusher, it writes the "Streaming not supported" 500 itself (see +// assertFlushable) and returns (nil, false); the caller must return without +// writing anything else, and extra is never called. +func startSSEStream(w http.ResponseWriter, extra func()) (http.Flusher, bool) { + flusher, ok := assertFlushable(w) + if !ok { + return nil, false + } + setSSEHeaders(w) + if extra != nil { + extra() + } + flusher.Flush() + return flusher, true +} + +// writeSSEKeepAlive writes a single SSE comment line (": keep-alive\n\n") to +// w and flushes it, so intermediary proxies/load balancers do not close an +// otherwise-idle SSE connection for inactivity. It returns the underlying +// write error (if any), mirroring writeSSEData, so the caller can decide +// whether to end the stream; it does not itself log or close anything. +func writeSSEKeepAlive(w io.Writer, flusher http.Flusher) error { + if _, err := fmt.Fprint(w, ": keep-alive\n\n"); err != nil { //nolint:gosec // G705: fixed literal, not derived from any request + return err + } + flusher.Flush() + return nil +} + // writeJSONRPC writes a jsonrpc2.Message using the library's encoder to ensure proper serialization. func writeJSONRPC(w http.ResponseWriter, msg jsonrpc2.Message) error { w.Header().Set("Content-Type", "application/json") @@ -81,9 +181,65 @@ func compositeKey(sessID, idKey string) string { return sessID + "|" + idKey } -// isSupportedMCPVersion is intentionally permissive: we accept any present version string. -// This avoids being pedantic and breaking on new protocol dates while remaining compliant, -// since this proxy is transport-level and does not depend on specific MCP versions. -func isSupportedMCPVersion(_ string) bool { - return true +// extractStringParam returns params[key] as a string, and whether it was +// present and string-typed. Absent params, params that don't decode as a JSON +// object, or a non-string value at key all yield ("", false) rather than an +// error -- callers treat a missing/malformed field as "nothing to route on". +func extractStringParam(params json.RawMessage, key string) (string, bool) { + if len(params) == 0 { + return "", false + } + var m map[string]any + if err := json.Unmarshal(params, &m); err != nil { + return "", false + } + s, ok := m[key].(string) + return s, ok +} + +// rewriteRequestParam returns a new *jsonrpc2.Request identical to req except +// that its top-level params[key] is set to value (params is created if req +// had none). It never mutates req: per the "copy before mutating caller +// input" style rule, the params map is freshly decoded from req.Params before +// the assignment, so no caller-held reference (including req itself) is +// modified. +func rewriteRequestParam(req *jsonrpc2.Request, key string, value any) (*jsonrpc2.Request, error) { + m := map[string]any{} + if len(req.Params) > 0 { + if err := json.Unmarshal(req.Params, &m); err != nil { + return nil, err + } + } + m[key] = value + data, err := json.Marshal(m) + if err != nil { + return nil, err + } + return &jsonrpc2.Request{ID: req.ID, Method: req.Method, Params: data}, nil +} + +// supportedMCPVersions is the set of MCP protocol revision dates this proxy +// recognizes when strict validation is enabled (WithStrictProtocolValidation). +// The upcoming stateless revision is sourced from mcp.MCPVersionModern rather +// than duplicated as a literal, so the strict gate and the ClassifyRevision +// routing path (which keys off the same constant) cannot drift if that draft +// date changes before it is finalized. +var supportedMCPVersions = map[string]struct{}{ + "2024-11-05": {}, + "2025-03-26": {}, + "2025-06-18": {}, + "2025-11-25": {}, + mcp.MCPVersionModern: {}, +} + +// isSupportedMCPVersion reports whether version is one of the known MCP +// protocol revision dates in supportedMCPVersions. It is only consulted when +// the proxy's strict protocol validation is enabled (see +// WithStrictProtocolValidation); in the default permissive mode, handlePost +// never calls it and any MCP-Protocol-Version header value is accepted, since +// this proxy is transport-level and does not depend on a specific MCP +// revision. +func isSupportedMCPVersion(version string) bool { + _, ok := supportedMCPVersions[version] + return ok } diff --git a/pkg/transport/proxy/streamable/utils_test.go b/pkg/transport/proxy/streamable/utils_test.go new file mode 100644 index 0000000000..70eebb4be2 --- /dev/null +++ b/pkg/transport/proxy/streamable/utils_test.go @@ -0,0 +1,76 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestIsSupportedMCPVersion verifies membership in supportedMCPVersions: every +// known MCP revision date returns true, and unknown/garbage/empty strings +// return false. This function is only consulted by handlePost when strict +// protocol validation is enabled (see WithStrictProtocolValidation). +func TestIsSupportedMCPVersion(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + version string + want bool + }{ + {name: "2024-11-05 is supported", version: "2024-11-05", want: true}, + {name: "2025-03-26 is supported", version: "2025-03-26", want: true}, + {name: "2025-06-18 is supported", version: "2025-06-18", want: true}, + {name: "2025-11-25 is supported", version: "2025-11-25", want: true}, + {name: "2026-07-28 is supported", version: "2026-07-28", want: true}, + {name: "unknown future date is not supported", version: "1999-01-01", want: false}, + {name: "garbage string is not supported", version: "not-a-version", want: false}, + {name: "empty string is not supported", version: "", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, isSupportedMCPVersion(tt.version)) + }) + } +} + +// recordingFlusher is a minimal http.Flusher that records Flush calls. +type recordingFlusher struct { + flushed int +} + +func (f *recordingFlusher) Flush() { f.flushed++ } + +// TestWriteSSEDataIncludesEventMessage verifies writeSSEData emits an explicit +// SSE event name before the data line. Spec-lenient MCP clients (for example +// @ai-sdk/mcp) only dispatch frames where event === "message" and drop +// data-only frames; ToolHive's own SSE serializer and the MCP SDK reference +// transports always include "event: message". +// +// Regression for #5655: every JSON-RPC SSE writer in this package goes through +// writeSSEData (POST success, POST error, GET standalone notifications, and +// progress), so this unit test covers all of those paths. +func TestWriteSSEDataIncludesEventMessage(t *testing.T) { + t.Parallel() + + var buf bytes.Buffer + flusher := &recordingFlusher{} + payload := []byte(`{"jsonrpc":"2.0","id":1,"result":{"ok":true}}`) + + err := writeSSEData(&buf, flusher, payload) + require.NoError(t, err) + assert.Equal(t, 1, flusher.flushed) + + got := buf.String() + assert.Equal(t, "event: message\ndata: "+string(payload)+"\n\n", got, + "SSE frame must include event: message before data (raw bytes: %q)", got) + assert.True(t, bytes.HasPrefix([]byte(got), []byte("event: message\n")), + "frame must start with event: message, not a bare data: line") +} diff --git a/pkg/transport/proxy/transparent/backend_routing_test.go b/pkg/transport/proxy/transparent/backend_routing_test.go index 0e70e7e449..b6a0709b74 100644 --- a/pkg/transport/proxy/transparent/backend_routing_test.go +++ b/pkg/transport/proxy/transparent/backend_routing_test.go @@ -5,6 +5,7 @@ package transparent import ( "context" + "encoding/json" "io" "net/http" "net/http/httptest" @@ -208,6 +209,100 @@ func TestRoundTripReturns404ForUnknownSession(t *testing.T) { assert.Contains(t, string(body), `"code":-32001`) } +// TestRoundTrip404EchoesRequestID pins #5945: the unknown-session 404 must echo +// the incoming JSON-RPC id so a client correlating responses by id can match the +// error to the request that caused it. Before the fix the transparent proxy +// hardcoded a null id here, while the sibling classification-error path on the +// same RoundTrip already echoed it. +// +// Note TestRoundTripReturns404ForUnknownSession above cannot catch this: its body +// carries no "id" key either way, so it can't tell a present-but-omitted id apart +// from this test's cases. +func TestRoundTrip404EchoesRequestID(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + wantIDKey bool + wantID string // only checked when wantIDKey + }{ + { + name: "numeric id is echoed verbatim", + body: `{"jsonrpc":"2.0","id":7,"method":"tools/list"}`, + wantIDKey: true, + wantID: `"id":7`, + }, + { + name: "string id is echoed verbatim", + body: `{"jsonrpc":"2.0","id":"abc-123","method":"tools/list"}`, + wantIDKey: true, + wantID: `"id":"abc-123"`, + }, + { + // A notification has no id by definition. MCP encodes that by + // omitting the "id" key entirely, not by emitting null. + name: "notification omits the id key", + body: `{"jsonrpc":"2.0","method":"notifications/initialized"}`, + wantIDKey: false, + }, + { + // An explicit null id is not a correlatable id either, so it is + // also encoded by omitting the key. + name: "explicit null id omits the id key", + body: `{"jsonrpc":"2.0","id":null,"method":"tools/list"}`, + wantIDKey: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + // Reaching the backend would mean the guard did not fire; the + // assertions below would fail on the 200 that results. + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + + tt2 := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions( + "localhost", 0, backend.URL, + nil, nil, nil, + false, false, "sse", + nil, nil, "", false, + nil, + )) + + req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp", strings.NewReader(tt.body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", uuid.New().String()) // unknown to the store + + resp, err := tt2.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusNotFound, resp.StatusCode) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + _ = resp.Body.Close() + + assert.Contains(t, string(body), `"code":-32001`) + + // The body must remain a single valid JSON-RPC error object -- echoing + // a raw id must not corrupt the envelope. + var decoded map[string]any + require.NoError(t, json.Unmarshal(body, &decoded)) + assert.Equal(t, "2.0", decoded["jsonrpc"]) + + _, ok := decoded["id"] + assert.Equal(t, tt.wantIDKey, ok, `"id" key presence`) + if tt.wantIDKey { + assert.Contains(t, string(body), tt.wantID) + } + }) + } +} + // TestRoundTripAllowsInitializeWithUnknownSession verifies that an initialize // call is forwarded even when its Mcp-Session-Id is not yet in the session store. func TestRoundTripAllowsInitializeWithUnknownSession(t *testing.T) { @@ -239,12 +334,13 @@ func TestRoundTripAllowsInitializeWithUnknownSession(t *testing.T) { _ = resp.Body.Close() } -// TestRoundTripPreservesExplicitHost verifies that routing via URL.Host does not -// overwrite an explicit logical Host header set by the reverse proxy. +// TestRoundTripPreservesExplicitHost verifies that routing via URL.Host (a pod IP +// from backend_url) does not overwrite the configured target host in the Host +// header. Any other Host value is rewritten (TestRegression_TracingTransport_*). func TestRoundTripPreservesExplicitHost(t *testing.T) { t.Parallel() - const logicalHost = "mcp-aws-mcp" + const logicalHost = "mcp-aws-mcp:8080" var receivedHost atomic.Value backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { receivedHost.Store(r.Host) @@ -253,7 +349,7 @@ func TestRoundTripPreservesExplicitHost(t *testing.T) { defer backend.Close() proxy := NewTransparentProxyWithOptions( - "localhost", 0, backend.URL, + "localhost", 0, "http://"+logicalHost, nil, nil, nil, false, false, "sse", nil, nil, "", false, @@ -274,36 +370,79 @@ func TestRoundTripPreservesExplicitHost(t *testing.T) { assert.Equal(t, logicalHost, receivedHost.Load()) } -// TestRoundTripAllowsBatchInitializeWithUnknownSession verifies that a JSON-RPC -// batch payload containing an initialize call is forwarded even when its -// Mcp-Session-Id is not yet in the session store. -func TestRoundTripAllowsBatchInitializeWithUnknownSession(t *testing.T) { +// TestRoundTripRejectsBatch verifies that the transparent proxy rejects JSON-RPC +// batches with an HTTP 400 "Invalid Request" (-32600) instead of forwarding them +// to the backend. Batching was removed in MCP 2025-06-18; forwarding a batch +// would let its nested calls bypass authz/audit/tool-filtering (see #5745). +// Rejection is Content-Type-independent: a batch under a non-JSON content type +// must be rejected too (it skips the content-type-gated ParsingMiddleware). The +// backend fails the test if it is ever reached. +func TestRoundTripRejectsBatch(t *testing.T) { t.Parallel() - backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Mcp-Session-Id", uuid.New().String()) - w.WriteHeader(http.StatusOK) + backend := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + t.Error("backend must not be reached: batch should be rejected at the proxy") })) - defer backend.Close() - - tt := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions( - "localhost", 0, backend.URL, - nil, nil, nil, - false, false, "sse", - nil, nil, "", false, - nil, - )) - - req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp", - strings.NewReader(`[{"method":"initialize"},{"method":"tools/list"}]`)) - require.NoError(t, err) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Mcp-Session-Id", uuid.New().String()) // unknown but batch contains initialize + t.Cleanup(backend.Close) + + tests := []struct { + name string + body string + contentType string + }{ + { + name: "batch containing initialize", + body: `[{"method":"initialize"},{"method":"tools/list"}]`, + contentType: "application/json", + }, + { + name: "batch of tool calls", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"danger"}}]`, + contentType: "application/json", + }, + { + name: "batch under non-json content type", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"danger"}}]`, + contentType: "text/plain", + }, + } - resp, err := tt.RoundTrip(req) - require.NoError(t, err) - require.NotEqual(t, http.StatusBadRequest, resp.StatusCode) - _ = resp.Body.Close() + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + transport := newTracingTransport(http.DefaultTransport, NewTransparentProxyWithOptions( + "localhost", 0, backend.URL, + nil, nil, nil, + false, false, "sse", + nil, nil, "", false, + nil, + )) + + req, err := http.NewRequest(http.MethodPost, backend.URL+"/mcp", strings.NewReader(tt.body)) + require.NoError(t, err) + req.Header.Set("Content-Type", tt.contentType) + req.Header.Set("Mcp-Session-Id", uuid.New().String()) + + resp, err := transport.RoundTrip(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"code":-32600`) + + // A batch has no single request id to echo. MCP encodes that by + // omitting the "id" key entirely (schema/2025-11-25 types the + // error response id as optional, not nullable), never as null -- + // a substring check for "id":null would miss the key being + // present-but-null, so the key's absence is checked directly. + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + _, hasID := parsed["id"] + assert.False(t, hasID, `"id" key must be omitted, not present as null`) + }) + } } // TestRoundTripStoresBackendURLOnInitialize verifies that when an initialize @@ -715,3 +854,125 @@ func TestRoundTripReinitializesOnDialError(t *testing.T) { assert.Equal(t, http.StatusOK, resp.StatusCode, "client should see 200 after transparent re-init on dial error") assert.GreaterOrEqual(t, freshHit.Load(), int32(2), "fresh backend should receive initialize + replay") } + +// TestRoundTripStripsSessionIDFromInitialize verifies that a client-initiated +// initialize reaches the backend with no Mcp-Session-Id header, whether or not +// the session it carries maps to a backend session ID recorded by a prior +// transparent re-initialization. +// +// MCP requires a re-initializing client to "start a new session by sending a +// new InitializeRequest without a session ID attached". Forwarding the client's +// session ID — or the backend SID mapped from it — hands an already-initialized +// backend session a second handshake, which go-sdk v1.7+ rejects with +// `duplicate "initialize" received`, turning a client-side handshake retry into +// a hard failure. +func TestRoundTripStripsSessionIDFromInitialize(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + priorBackend bool // session carries a backend_sid from an earlier re-init + }{ + {name: "session mapped to a backend SID", priorBackend: true}, + {name: "session with no backend SID", priorBackend: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var receivedSIDs []string + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + receivedSIDs = append(receivedSIDs, r.Header.Get("Mcp-Session-Id")) + mu.Unlock() + w.Header().Set("Mcp-Session-Id", uuid.New().String()) + w.WriteHeader(http.StatusOK) + })) + defer backend.Close() + + proxy, addr := startProxy(t, backend.URL) + + clientSessionID := uuid.New().String() + sess := session.NewProxySession(clientSessionID) + if tt.priorBackend { + sess.SetMetadata(sessionMetadataBackendSID, uuid.New().String()) + } + require.NoError(t, proxy.sessionManager.AddSession(sess)) + + ctx := context.Background() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + "http://"+addr+"/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", clientSessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + _ = resp.Body.Close() + + mu.Lock() + defer mu.Unlock() + require.Len(t, receivedSIDs, 1, "backend should receive exactly one request") + assert.Empty(t, receivedSIDs[0], "initialize must reach the backend with no Mcp-Session-Id") + }) + } +} + +// TestRoundTripDoesNotReplayInitializeOnDialError verifies that a dial error on +// an initialize does not trigger transparent re-initialization. +// +// reinitializeAndReplay sends its own initialize and then replays the original +// request; for an initialize that would give the freshly created backend session +// a second handshake — the exact failure the session-ID strip exists to avoid. +// The client is already starting a new session, so the proxy instead unpins the +// session from the unreachable pod and lets the error surface, leaving the +// client's retry to route via the target service. +func TestRoundTripDoesNotReplayInitializeOnDialError(t *testing.T) { + t.Parallel() + + // A server closed immediately after creation: its URL refuses connections. + dead := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})) + deadURL := dead.URL + dead.Close() + + var targetHits atomic.Int32 + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + targetHits.Add(1) + w.Header().Set("Mcp-Session-Id", uuid.New().String()) + w.WriteHeader(http.StatusOK) + })) + defer target.Close() + + proxy, addr := startProxy(t, target.URL) + + clientSessionID := uuid.New().String() + sess := session.NewProxySession(clientSessionID) + sess.SetMetadata(sessionMetadataBackendURL, deadURL) + sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`) + require.NoError(t, proxy.sessionManager.AddSession(sess)) + + ctx := context.Background() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + "http://"+addr+"/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", clientSessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + _ = resp.Body.Close() + + assert.Zero(t, targetHits.Load(), + "the proxy must not send its own initialize to the target service for an initialize dial error") + + updated, ok := proxy.sessionManager.Get(normalizeSessionID(clientSessionID)) + require.True(t, ok, "session should survive the dial error") + backendURL, exists := updated.GetMetadataValue(sessionMetadataBackendURL) + require.True(t, exists, "backend_url should still be set") + assert.Equal(t, target.URL, backendURL, + "session must be unpinned from the unreachable pod so the client's retry routes via the target service") +} diff --git a/pkg/transport/proxy/transparent/host_header_regression_test.go b/pkg/transport/proxy/transparent/host_header_regression_test.go new file mode 100644 index 0000000000..819b663761 --- /dev/null +++ b/pkg/transport/proxy/transparent/host_header_regression_test.go @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestRegression_LocalhostProxy_NonLocalhostHostHeaderRewritten verifies that +// tracingTransport.RoundTrip rewrites a non-localhost Host header to match the +// target URL's host. Without this, an attacker could inject a Host header to +// bypass host-based validation on the upstream server. +func TestRegression_LocalhostProxy_NonLocalhostHostHeaderRewritten(t *testing.T) { + t.Parallel() + + // Capture the Host header the upstream server actually receives. + var receivedHost string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHost = r.Host + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + targetURL, err := url.Parse(upstream.URL) + require.NoError(t, err) + + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, "streamable-http", nil, nil, "", false) + + proxy := &httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.SetURL(targetURL) + pr.SetXForwarded() + }, + FlushInterval: -1, + Transport: newTracingTransport(http.DefaultTransport, p), + ModifyResponse: p.modifyResponse, + } + + // Send a request with a malicious Host header. + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "http://evil.example.com/some/path", nil) + req.Host = "evil.example.com" + + proxy.ServeHTTP(rec, req) + + // The outbound request's Host must be the target URL's host, + // not the attacker-supplied value. + assert.Equal(t, targetURL.Host, receivedHost, + "tracingTransport must rewrite Host header from attacker value to target URL host") + assert.NotEqual(t, "evil.example.com", receivedHost, + "attacker Host header must not reach the upstream server") +} + +// TestRegression_TracingTransport_RewritesHostHeader exercises the Host rewrite +// directly against tracingTransport.RoundTrip. The end-to-end test above routes +// through httputil.ReverseProxy, whose SetURL sets Out.Host="" so net/http would +// derive the sent Host from URL.Host even if the tracingTransport rewrite were +// deleted. Driving RoundTrip directly with a request whose req.Host differs from +// req.URL.Host gates the rewrite itself: if the RoundTrip Host-rewrite is removed, +// the attacker-supplied req.Host reaches the upstream and this test fails. +func TestRegression_TracingTransport_RewritesHostHeader(t *testing.T) { + t.Parallel() + + var receivedHost string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHost = r.Host + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + targetURL, err := url.Parse(upstream.URL) + require.NoError(t, err) + + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, "streamable-http", nil, nil, "", false) + transport := newTracingTransport(http.DefaultTransport, p) + + // The request targets the real upstream (req.URL.Host) but carries an + // attacker-supplied Host header (req.Host) that differs from it. + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, upstream.URL+"/some/path", nil) + require.NoError(t, err) + req.Host = "evil.example.com" + require.NotEqual(t, req.Host, req.URL.Host, "test precondition: req.Host must differ from req.URL.Host") + + resp, err := transport.RoundTrip(req) + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + require.Equal(t, http.StatusOK, resp.StatusCode) + + // RoundTrip must have rewritten req.Host to the target URL host before + // forwarding, so the upstream never observes the attacker value. + assert.Equal(t, targetURL.Host, receivedHost, + "tracingTransport.RoundTrip must rewrite Host header to the target URL host") + assert.NotEqual(t, "evil.example.com", receivedHost, + "attacker Host header must not reach the upstream server") +} diff --git a/pkg/transport/proxy/transparent/method_gate_test.go b/pkg/transport/proxy/transparent/method_gate_test.go index 9bfc2e35c5..4f761e9f51 100644 --- a/pkg/transport/proxy/transparent/method_gate_test.go +++ b/pkg/transport/proxy/transparent/method_gate_test.go @@ -4,69 +4,83 @@ package transparent import ( + "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/mcp" ) -func TestStatelessMethodGate(t *testing.T) { +func TestMethodGate(t *testing.T) { t.Parallel() - inner := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - }) - tests := []struct { name string + stateless bool + protocolHeader string method string - expectedStatus int - expectAllow bool + expectGated bool }{ - { - name: "GET returns 405 with Allow header", - method: http.MethodGet, - expectedStatus: http.StatusMethodNotAllowed, - expectAllow: true, - }, - { - name: "HEAD returns 405 with Allow header", - method: http.MethodHead, - expectedStatus: http.StatusMethodNotAllowed, - expectAllow: true, - }, - { - name: "DELETE returns 405 with Allow header", - method: http.MethodDelete, - expectedStatus: http.StatusMethodNotAllowed, - expectAllow: true, - }, - { - name: "POST is forwarded", - method: http.MethodPost, - expectedStatus: http.StatusOK, - }, - { - name: "PUT is forwarded", - method: http.MethodPut, - expectedStatus: http.StatusOK, - }, + // stateless=true gates GET/HEAD/DELETE regardless of header + {"stateless GET gated", true, "", http.MethodGet, true}, + {"stateless HEAD gated", true, "", http.MethodHead, true}, + {"stateless DELETE gated", true, "", http.MethodDelete, true}, + {"stateless POST passes", true, "", http.MethodPost, false}, + {"stateless with Modern header GET gated", true, mcp.MCPVersionModern, http.MethodGet, true}, + + // stateless=false, Modern header gates GET/HEAD/DELETE + {"Modern header GET gated", false, mcp.MCPVersionModern, http.MethodGet, true}, + {"Modern header HEAD gated", false, mcp.MCPVersionModern, http.MethodHead, true}, + {"Modern header DELETE gated", false, mcp.MCPVersionModern, http.MethodDelete, true}, + {"Modern header POST passes", false, mcp.MCPVersionModern, http.MethodPost, false}, + + // stateless=false, no/other header: legacy behavior, nothing gated + {"legacy GET passes", false, "", http.MethodGet, false}, + {"legacy HEAD passes", false, "", http.MethodHead, false}, + {"legacy DELETE passes", false, "", http.MethodDelete, false}, + {"legacy POST passes", false, "", http.MethodPost, false}, + {"other header GET passes", false, "2025-11-25", http.MethodGet, false}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - handler := statelessMethodGate(inner) + var gotBody string + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + gotBody = string(body) + w.WriteHeader(http.StatusOK) + }) + + p := &TransparentProxy{stateless: tc.stateless} + handler := p.methodGate(inner) rec := httptest.NewRecorder() - req := httptest.NewRequest(tc.method, "/", nil) + + var body string + if tc.method == http.MethodPost { + body = "request-body-payload" + } + req := httptest.NewRequest(tc.method, "/", strings.NewReader(body)) + if tc.protocolHeader != "" { + req.Header.Set("MCP-Protocol-Version", tc.protocolHeader) + } handler.ServeHTTP(rec, req) - assert.Equal(t, tc.expectedStatus, rec.Code) - if tc.expectAllow { + if tc.expectGated { + assert.Equal(t, http.StatusMethodNotAllowed, rec.Code) assert.Equal(t, "POST, OPTIONS", rec.Header().Get("Allow")) + } else { + assert.Equal(t, http.StatusOK, rec.Code) + if tc.method == http.MethodPost { + require.Equal(t, body, gotBody, "POST body must reach next handler intact") + } } }) } diff --git a/pkg/transport/proxy/transparent/revision_classification_test.go b/pkg/transport/proxy/transparent/revision_classification_test.go new file mode 100644 index 0000000000..5b6ddd8e3b --- /dev/null +++ b/pkg/transport/proxy/transparent/revision_classification_test.go @@ -0,0 +1,290 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/mcp" +) + +// roundTripFunc adapts a function to http.RoundTripper, letting tests spy on +// whether the backend transport was invoked without standing up a real server. +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func TestParseRPCRequest(t *testing.T) { + t.Parallel() + + tp := &tracingTransport{p: &TransparentProxy{targetURI: "http://backend"}} + + tests := []struct { + name string + body string + wantMethod string + wantID string // "" means nil/absent + wantSingleRequest bool + wantSawInitialize bool + }{ + { + name: "single request with id", + body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`, + wantMethod: "tools/list", + wantID: "1", + wantSingleRequest: true, + }, + { + name: "single initialize", + body: `{"jsonrpc":"2.0","id":1,"method":"initialize"}`, + wantMethod: "initialize", + wantID: "1", + wantSingleRequest: true, + wantSawInitialize: true, + }, + { + name: "notification has no id", + body: `{"jsonrpc":"2.0","method":"notifications/progress"}`, + wantMethod: "notifications/progress", + wantSingleRequest: false, + }, + { + name: "explicit null id is not a valid request id", + body: `{"jsonrpc":"2.0","id":null,"method":"tools/list"}`, + wantMethod: "tools/list", + wantID: "null", + wantSingleRequest: false, + }, + { + name: "response-shaped body has no method", + body: `{"jsonrpc":"2.0","id":1,"result":{}}`, + wantMethod: "", + wantID: "1", + wantSingleRequest: false, + }, + { + name: "batch with initialize", + body: `[{"jsonrpc":"2.0","id":1,"method":"initialize"},{"jsonrpc":"2.0","id":2,"method":"tools/list"}]`, + wantSingleRequest: false, + wantSawInitialize: true, + }, + { + name: "batch without initialize", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"tools/call"}]`, + wantSingleRequest: false, + }, + { + name: "malformed JSON", + body: `{not valid json`, + wantSingleRequest: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + method, _, id, singleRequest, sawInitialize := tp.parseRPCRequest([]byte(tc.body)) + + assert.Equal(t, tc.wantMethod, method) + assert.Equal(t, tc.wantSingleRequest, singleRequest) + assert.Equal(t, tc.wantSawInitialize, sawInitialize) + if tc.wantID == "" { + assert.Empty(t, string(id)) + } else { + assert.Equal(t, tc.wantID, string(id)) + } + }) + } +} + +// TestRoundTripClassifiesModernRequests drives tracingTransport.RoundTrip +// directly (bypassing httputil.ReverseProxy) with a spy backend RoundTripper, +// so tests can assert both the returned response and whether the backend was +// ever contacted. +func TestRoundTripClassifiesModernRequests(t *testing.T) { + t.Parallel() + + newProxy := func(spy http.RoundTripper) (*tracingTransport, *TransparentProxy) { + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, + "streamable-http", nil, nil, "", false) + return newTracingTransport(spy, p), p + } + + t.Run("malformed Modern single-request is rejected before the backend is called", func(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return httptest.NewRecorder().Result(), nil + }) + tt, _ := newProxy(spy) + + // Header claims Modern but the body carries no _meta at all: a + // HeaderMismatchError (bad/absent body version, non-empty header). + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + assert.False(t, backendCalled.Load(), "backend must not be contacted for a rejected request") + }) + + t.Run("well-formed Modern single-request falls through to the backend", func(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, _ := newProxy(spy) + + // Not "initialize", has a valid id, and _meta carries a matching + // protocolVersion plus clientCapabilities: ClassifyRevision returns + // (RevisionModern, nil), so the request must reach the backend. + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, http.StatusOK, resp.StatusCode, "a well-formed Modern request must not be rejected") + assert.True(t, backendCalled.Load(), "backend must be contacted for a well-formed Modern request") + }) + + t.Run("batch is rejected before classification, backend not contacted", func(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, _ := newProxy(spy) + + // Batches are rejected outright (batching was removed in MCP 2025-06-18) + // before revision classification or forwarding — see #5745. + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`[{"jsonrpc":"2.0","id":1,"method":"tools/call"}]`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "batches must be rejected, not forwarded") + assert.False(t, backendCalled.Load(), "backend must not be contacted for a rejected batch") + }) + + t.Run("large-integer id is preserved verbatim in the 400 body", func(t *testing.T) { + t.Parallel() + + const largeID = "9007199254740993" // 2^53 + 1: loses precision if round-tripped through float64 + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("backend must not be contacted for a rejected request") + return nil, nil + }) + tt, _ := newProxy(spy) + + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":`+largeID+`,"method":"tools/call"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var decoded struct { + ID json.RawMessage `json:"id"` + } + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + assert.Equal(t, largeID, string(decoded.ID), "large integer id must not lose precision") + }) + + t.Run("Modern 200 response flips serverInitialized", func(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newProxy(spy) + require.False(t, p.serverInitialized(), "precondition: latch starts unset") + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + assert.True(t, p.serverInitialized(), "a 200 to a Modern request must flip the readiness latch") + }) + + t.Run("Modern non-200 response does not flip serverInitialized", func(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusInternalServerError, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newProxy(spy) + require.False(t, p.serverInitialized(), "precondition: latch starts unset") + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusInternalServerError, resp.StatusCode) + assert.False(t, p.serverInitialized(), "a non-200 response must not flip the readiness latch") + }) + + t.Run("Legacy 200 response without session header or initialize does not flip serverInitialized", func(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newProxy(spy) + require.False(t, p.serverInitialized(), "precondition: latch starts unset") + + // Legacy request (no MCP-Protocol-Version header, not initialize) whose + // 200 response carries no Mcp-Session-Id: neither existing branch fires, + // and the new Modern branch must not broaden to cover this case either. + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + req.Header.Set("Content-Type", "application/json") + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + assert.False(t, p.serverInitialized(), "a Legacy 200 with no session header must not flip the readiness latch") + }) +} diff --git a/pkg/transport/proxy/transparent/revision_guard_regression_test.go b/pkg/transport/proxy/transparent/revision_guard_regression_test.go new file mode 100644 index 0000000000..8557c5f6e4 --- /dev/null +++ b/pkg/transport/proxy/transparent/revision_guard_regression_test.go @@ -0,0 +1,247 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/transport/session" +) + +// This file is a SECURITY regression suite: it proves that a client-forgeable +// Modern revision signal (MCP-Protocol-Version header and/or params._meta) +// cannot bypass the Legacy session machinery in tracingTransport.RoundTrip. +// That machinery is deliberately keyed on Mcp-Session-Id PRESENCE, never on +// the classified revision, because revision is derived entirely from +// client-controlled input while Mcp-Session-Id is validated against +// server-side session state. See transparent_proxy.go's RoundTrip guard +// comment and revision_classification_test.go for the classification-only +// coverage this file complements. +// +// Scope: this suite covers single-request Modern-signal forgery only. Batch +// payloads are rejected outright before revision classification (batching was +// removed from MCP in 2025-06-18; see TestRoundTripRejectsBatch and #5745), so +// they can neither carry a Modern signal nor reach the unknown-session guard +// and are out of scope here. + +// modernToolsCallBody is a well-formed Modern (2026-07-28) JSON-RPC request: +// not "initialize", a valid non-null id, and params._meta carrying both +// reserved keys required for a nil-error classification. Reused verbatim (and +// with a missing header in one case) across every test below so each test +// starts from a request that ClassifyRevision genuinely accepts as Modern. +const modernToolsCallBody = `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + +// newGuardTransport builds a tracingTransport backed by a spy RoundTripper, +// mirroring the harness in revision_classification_test.go. targetURI is a +// syntactically valid placeholder: backendRecovery.reinitializeAndReplay +// parses it with url.Parse but the spy intercepts every call, so it is never +// actually dialed. +func newGuardTransport(spy http.RoundTripper) (*tracingTransport, *TransparentProxy) { + p := NewTransparentProxy("127.0.0.1", 0, "http://backend", nil, nil, nil, false, false, + "streamable-http", nil, nil, "", false) + return newTracingTransport(spy, p), p +} + +// assertClassifiesModernNil is a precondition check, not the regression +// itself: it confirms the given body/header combination really does +// classify as Modern with a nil error. If this fails, the test below it is +// not exercising a forged-Modern scenario at all. +func assertClassifiesModernNil(t *testing.T, tt *tracingTransport, body []byte, protoHeader string) { + t.Helper() + method, params, _, singleRequest, _ := tt.parseRPCRequest(body) + require.True(t, singleRequest, "precondition: body must parse as a single JSON-RPC request") + meta := mcp.ExtractMeta(params) + rev, err := mcp.ClassifyRevision(method, meta, protoHeader) + require.NoError(t, err, "precondition: body must classify Modern with a nil error") + require.Equal(t, mcp.RevisionModern, rev, "precondition: body must classify as Modern") +} + +// TestGuardUnknownSessionFiresDespiteForgedModernRevision is the core +// mutation-check case: a well-formed Modern request with an unknown/stale +// Mcp-Session-Id must still be rejected with the session-not-found response, +// exactly as it would be for a Legacy request (see +// TestRoundTripReturns404ForUnknownSession in backend_routing_test.go). +// +// This test FAILS if the guard in tracingTransport.RoundTrip is ever +// re-keyed on `revision == mcp.RevisionModern` (or any check derived from +// the classified revision) instead of Mcp-Session-Id presence — that would +// let a forged Modern signal skip session validation entirely. +func TestGuardUnknownSessionFiresDespiteForgedModernRevision(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, _ := newGuardTransport(spy) + + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), mcp.MCPVersionModern) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + req.Header.Set("Mcp-Session-Id", uuid.New().String()) // unknown/stale, never added to the session manager + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, http.StatusNotFound, resp.StatusCode) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"code":-32001`) + assert.False(t, backendCalled.Load(), + "backend must not be contacted: the unknown-session guard must fire regardless of the classified revision") +} + +// TestGuardBackendSIDRewriteStillHappensWithForgedModernRevision verifies +// that a well-formed Modern request against a KNOWN session whose metadata +// carries a backend_sid still has its outbound Mcp-Session-Id rewritten to +// that backend SID — the rewrite (like the guard above) is driven by session +// metadata, not by the classified revision. Mirrors the assertion style of +// backend_routing_test.go's TestRoundTripReinitializesOnBackend404 (which +// checks sessionMetadataBackendSID) but observes the header directly via a +// spy, as in revision_classification_test.go. +func TestGuardBackendSIDRewriteStillHappensWithForgedModernRevision(t *testing.T) { + t.Parallel() + + const backendSID = "backend-assigned-opaque-sid" + var gotSID atomic.Value + spy := roundTripFunc(func(r *http.Request) (*http.Response, error) { + gotSID.Store(r.Header.Get("Mcp-Session-Id")) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newGuardTransport(spy) + + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), mcp.MCPVersionModern) + + clientSID := uuid.New().String() + sess := session.NewProxySession(clientSID) + sess.SetMetadata(sessionMetadataBackendSID, backendSID) + require.NoError(t, p.sessionManager.AddSession(sess)) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + req.Header.Set("Mcp-Session-Id", clientSID) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + assert.Equal(t, backendSID, gotSID.Load(), + "outbound Mcp-Session-Id must be rewritten to backend_sid even though the request forges a Modern revision signal") +} + +// TestGuardReinitRecoveryStillTriggersWithForgedModernRevision verifies that +// the transparent re-initialize-and-replay recovery path (triggered by a 404 +// from a known session with a stored init body) still fires for a +// well-formed Modern request. Mirrors backend_routing_test.go's +// TestRoundTripReinitializesOnBackend404, but the request itself carries a +// forged Modern signal to prove recovery is unaffected by classified +// revision. +func TestGuardReinitRecoveryStillTriggersWithForgedModernRevision(t *testing.T) { + t.Parallel() + + freshSessionID := uuid.New().String() + var initCalls, otherCalls atomic.Int32 + spy := roundTripFunc(func(r *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(r.Body) + if strings.Contains(string(body), `"initialize"`) { + initCalls.Add(1) + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Mcp-Session-Id": []string{freshSessionID}}, + Body: http.NoBody, + }, nil + } + n := otherCalls.Add(1) + if n == 1 { + // First non-initialize forward: simulate the backend pod having + // lost its in-memory session state. + return &http.Response{StatusCode: http.StatusNotFound, Header: make(http.Header), Body: http.NoBody}, nil + } + // Second non-initialize forward is the replay after re-init. + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newGuardTransport(spy) + + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), mcp.MCPVersionModern) + + clientSID := uuid.New().String() + sess := session.NewProxySession(clientSID) + sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`) + require.NoError(t, p.sessionManager.AddSession(sess)) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + req.Header.Set("Mcp-Session-Id", clientSID) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, + "client must see 200 after transparent re-init, even for a forged-Modern request") + assert.Equal(t, int32(1), initCalls.Load(), "exactly one re-initialize call expected") + assert.Equal(t, int32(2), otherCalls.Load(), "original forward + replay expected") + + updated, ok := p.sessionManager.Get(normalizeSessionID(clientSID)) + require.True(t, ok, "session should still exist after re-init") + backendSID, exists := updated.GetMetadataValue(sessionMetadataBackendSID) + require.True(t, exists) + assert.Equal(t, freshSessionID, backendSID, + "backend_sid must be updated by the recovery path regardless of the classified revision") +} + +// TestGuardDeleteCleanupStillWorksWithBodyMetaButNoHeader verifies that DELETE +// session cleanup (keyed on Mcp-Session-Id, req.Method and response status — +// see RoundTrip's DELETE cleanup block) still runs when the DELETE body +// carries Modern _meta but no MCP-Protocol-Version header. Note: a DELETE +// that also carried the MCP-Protocol-Version: 2026-07-28 header would be +// rejected with 405 at the header-only methodGate and would never reach +// RoundTrip at all (see method_gate_test.go's "Modern header DELETE gated" +// case) — this test exercises the one shape of forged-Modern DELETE that +// does reach RoundTrip. Mirrors delete_session_test.go's cleanup assertion. +func TestGuardDeleteCleanupStillWorksWithBodyMetaButNoHeader(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newGuardTransport(spy) + + // No MCP-Protocol-Version header: classification relies solely on the + // reserved _meta keys, per mcp.ClassifyRevision's documented signal rules. + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), "") + + clientSID := uuid.New().String() + sess := session.NewProxySession(clientSID) + require.NoError(t, p.sessionManager.AddSession(sess)) + + req := httptest.NewRequest(http.MethodDelete, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", clientSID) + // Deliberately no MCP-Protocol-Version header. + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + _, ok := p.sessionManager.Get(normalizeSessionID(clientSID)) + assert.False(t, ok, + "DELETE session cleanup must still key off Mcp-Session-Id presence even when the body carries Modern _meta") +} diff --git a/pkg/transport/proxy/transparent/transparent_proxy.go b/pkg/transport/proxy/transparent/transparent_proxy.go index cd5b4b4bd7..3db98bf89f 100644 --- a/pkg/transport/proxy/transparent/transparent_proxy.go +++ b/pkg/transport/proxy/transparent/transparent_proxy.go @@ -33,6 +33,7 @@ import ( "github.com/stacklok/toolhive/pkg/bodylimit" rt "github.com/stacklok/toolhive/pkg/container/runtime" "github.com/stacklok/toolhive/pkg/healthcheck" + "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/proxy/socket" "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/transport/types" @@ -267,7 +268,7 @@ func WithRemoteRawQuery(rawQuery string) Option { } // WithStateless configures the proxy for stateless streamable-HTTP servers. -// In stateless mode, incoming GET and DELETE requests receive 405 Method Not Allowed +// In stateless mode, incoming GET, HEAD, and DELETE requests receive 405 Method Not Allowed // instead of being forwarded, and health checks use POST ping instead of GET. func WithStateless() Option { return func(p *TransparentProxy) { @@ -540,6 +541,14 @@ type tracingTransport struct { recovery *backendRecovery } +// targetHost returns the host of the configured target URI, or "" if it has none. +func (t *tracingTransport) targetHost() string { + if u, err := url.Parse(t.p.targetURI); err == nil { + return u.Host + } + return "" +} + func newTracingTransport(base http.RoundTripper, p *TransparentProxy) *tracingTransport { return &tracingTransport{ p: p, @@ -565,11 +574,11 @@ func (p *TransparentProxy) serverInitialized() bool { // nolint:gocyclo // This function handles multiple request types and is complex by design func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) { - // Keep any Host value set by ReverseProxy.SetURL. The request URL host can - // be changed later to a routing-only backend_url (for example a Kubernetes - // ClusterIP/pod IP); strict backend host validation must still see the - // logical target host, not that routing address. - if req.Host == "" { + // Rewrite the Host header to the URL host, so an injected Host never reaches the + // backend, except for the configured target host: backend_url can change the URL + // host to a routing-only pod IP, and strict backend host validation must still + // see the logical target host. + if req.Host != req.URL.Host && req.Host != t.targetHost() { req.Host = req.URL.Host } @@ -588,17 +597,50 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) return plainResponse(req, http.StatusRequestEntityTooLarge, "Request Entity Too Large"), nil } + // Reject JSON-RPC batches before forwarding to the backend. This runs + // independently of Content-Type (unlike ParsingMiddleware), so a batch + // smuggled under a non-JSON content type cannot reach the backend where its + // nested calls would bypass authz/audit/tool-filtering. Batching was removed + // in MCP revision 2025-06-18; ToolHive serves only 2025-11-25 and 2026-07-28 + // (see #5745). + if mcp.IsBatchRequest(reqBody) { + return mcp.BatchUnsupportedResponse(req), nil + } + // thv proxy does not provide the transport type, so we need to detect it from the request path := req.URL.Path isMCP := strings.HasPrefix(path, "/mcp") isJSON := strings.Contains(req.Header.Get("Content-Type"), "application/json") sawInitialize := false + revision := mcp.RevisionLegacy + // requestID carries the incoming JSON-RPC id so an error response built below + // can echo it, letting a client correlate the error with its request. It is set + // only when parseRPCRequest reports a single request, which by construction + // means the id is present and non-null; it stays nil for a notification, a + // batch, or a bodiless GET/DELETE. session.NotFoundResponse (via NotFoundBody) + // omits the "id" key entirely for a nil requestID rather than emitting a null + // id — MCP narrows base JSON-RPC to make that omission the correct encoding + // of "no id" (see session.HasJSONRPCID). + var requestID any if len(reqBody) > 0 && ((isMCP && isJSON) || t.p.transportType == types.TransportTypeStreamableHTTP.String()) { - sawInitialize = t.detectInitialize(reqBody) + method, params, id, singleRequest, isInit := t.parseRPCRequest(reqBody) + sawInitialize = isInit + if singleRequest { + requestID = id + meta := mcp.ExtractMeta(params) + rev, cerr := mcp.ClassifyRevision(method, meta, req.Header.Get("MCP-Protocol-Version")) + if cerr != nil { + // Malformed Modern request: reject before the backend is ever contacted. + return mcp.ClassificationErrorResponse(req, id, cerr), nil + } + revision = rev + } } + //nolint:gosec // G706: logging target URI from config + slog.Debug("classified request revision", "modern", revision == mcp.RevisionModern, "target", t.p.targetURI) // Guard: reject non-initialize requests with unknown session IDs. // When multiple proxyrunner replicas share a Redis session store, @@ -611,20 +653,35 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) slog.Error("session store lookup failed", "error", err) return plainResponse(req, http.StatusServiceUnavailable, "session store unavailable"), nil } - return session.NotFoundResponse(req), nil + return session.NotFoundResponse(req, requestID), nil } } - // Capture the client-facing session ID before the backend SID rewrite below. - // Recovery and session cleanup paths must look up sessions by the client SID - // (the store key), not the backend SID that is written into the header. + // Capture the client-facing session ID before the header is rewritten or + // stripped below. Recovery and session cleanup paths must look up sessions by + // the client SID (the store key), not the backend SID that is written into the + // header. clientSID := req.Header.Get("Mcp-Session-Id") - // Rewrite the outbound Mcp-Session-Id to the backend's assigned session ID when - // the proxy transparently re-initialized the backend session. This is done here - // (after the guard check above) so the guard always sees the original client - // session ID and can look it up correctly in the session store. - if clientSID != "" { + switch { + case sawInitialize: + // An initialize request must never carry a session ID. MCP requires a + // client that is re-initializing to "start a new session by sending a new + // InitializeRequest without a session ID attached", and a backend session + // that has already completed its handshake rejects a second initialize on + // it -- go-sdk v1.7+ answers `duplicate "initialize" received`. Forwarding + // the client's stale SID -- or, worse, the backend SID mapped from it -- + // turns a client-side handshake retry into a hard failure. Strip it so the + // backend mints a fresh session instead; the response's Mcp-Session-Id is + // then adopted as a new session below. reinitializeAndReplay already does + // the same on the recovery path. + req.Header.Del("Mcp-Session-Id") + case clientSID != "": + // Rewrite the outbound Mcp-Session-Id to the backend's assigned session ID + // when the proxy transparently re-initialized the backend session. This is + // done here (after the guard check above) so the guard always sees the + // original client session ID and can look it up correctly in the session + // store. if sess, ok := t.p.sessionManager.Get(normalizeSessionID(clientSID)); ok { if backendSID, exists := sess.GetMetadataValue(sessionMetadataBackendSID); exists && backendSID != "" { req.Header.Set("Mcp-Session-Id", backendSID) @@ -654,10 +711,22 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) } // Dial error against a stored pod IP means the pod has been replaced. // Attempt transparent re-initialization so the client sees no error. + // + // An initialize request is excluded, as it already is on the 404 path + // below: reinitializeAndReplay sends its own initialize and then replays + // the original request, which for an initialize would hand the freshly + // created backend session a second handshake -- the very failure this + // function strips the session ID to avoid. The client is already starting + // a new session, so instead unpin it from the dead pod and let the error + // surface; the client's retry is then routed via the target service. if isDialError(err) { - req.Header.Set("Mcp-Session-Id", clientSID) - if reInitResp, reInitErr := t.recovery.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil { - return reInitResp, reInitErr + if sawInitialize { + t.recovery.unpinSession(clientSID) + } else { + req.Header.Set("Mcp-Session-Id", clientSID) + if reInitResp, reInitErr := t.recovery.reinitializeAndReplay(req, reqBody); reInitResp != nil || reInitErr != nil { + return reInitResp, reInitErr + } } } slog.Error("failed to forward request", "error", err) @@ -743,6 +812,14 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) t.p.setServerInitialized() return resp, nil } + + // Modern (2026-07-28, stateless) has no initialize handshake, so use the + // first successful (HTTP 200) Modern request as the readiness signal: it + // lets the background health monitor begin probing this backend, which + // otherwise never starts for a pure-Modern server. See #5831. + if revision == mcp.RevisionModern && !t.p.serverInitialized() { + t.p.setServerInitialized() + } } return resp, nil @@ -756,14 +833,17 @@ func isBackendSessionLostResponse(resp *http.Response) bool { return false } - body, err := io.ReadAll(resp.Body) + // Read only the start of the body (the error text is short); the client + // still gets the whole body. + body, err := io.ReadAll(io.LimitReader(resp.Body, 4<<10)) + resp.Body = struct { + io.Reader + io.Closer + }{io.MultiReader(bytes.NewReader(body), resp.Body), resp.Body} if err != nil { slog.Debug("failed to inspect bad request response body for backend session loss", "error", err) - resp.Body = io.NopCloser(bytes.NewReader(nil)) return false } - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(body)) lower := strings.ToLower(string(body)) // TypeScript SDK servers answer a sessionless request with "no valid session @@ -811,36 +891,56 @@ func readRequestBody(req *http.Request) ([]byte, error) { return reqBody, nil } -func (t *tracingTransport) detectInitialize(body []byte) bool { - type rpcMethod struct { - Method string `json:"method"` +// parseRPCRequest parses a POST body as a JSON-RPC request in a single pass. +// +// If the body is a single JSON-RPC object, method/params/id are populated +// from it and singleRequest reports whether it is a real request (has both a +// method and a valid, non-null id) as opposed to a notification (no id) or a +// response-shaped body (no method) — either of which is not eligible for +// mcp.ClassifyRevision. sawInitialize reports whether method == "initialize". +// +// If the body is not a single JSON object (e.g. a JSON-RPC batch), method, +// params and id are zero and singleRequest is false: batches are never +// classified. sawInitialize is still computed by scanning the batch for any +// member whose method is "initialize", preserving prior behavior. +func (t *tracingTransport) parseRPCRequest( + body []byte, +) (method string, params, id json.RawMessage, singleRequest, sawInitialize bool) { + type rpcRequest struct { + Method string `json:"method"` + Params json.RawMessage `json:"params"` + ID json.RawMessage `json:"id"` } - // Single JSON-RPC object. - var single rpcMethod + var single rpcRequest if err := json.Unmarshal(body, &single); err == nil { - if single.Method == "initialize" { + sawInitialize = single.Method == "initialize" + if sawInitialize { //nolint:gosec // G706: logging target URI from config slog.Debug("detected initialize method call", "target", t.p.targetURI) - return true } - return false + singleRequest = single.Method != "" && len(single.ID) > 0 && string(single.ID) != "null" + return single.Method, single.Params, single.ID, singleRequest, sawInitialize } - // JSON-RPC batch: array of objects. Return true if any member is initialize. + // JSON-RPC batch: array of objects. Only sawInitialize is computed; a + // batch is never eligible for revision classification. + type rpcMethod struct { + Method string `json:"method"` + } var batch []rpcMethod if err := json.Unmarshal(body, &batch); err != nil { slog.Debug("failed to parse JSON-RPC body", "error", err) - return false + return "", nil, nil, false, false } for _, rpc := range batch { if rpc.Method == "initialize" { //nolint:gosec // G706: logging target URI from config slog.Debug("detected initialize method call in batch", "target", t.p.targetURI) - return true + return "", nil, nil, false, true } } - return false + return "", nil, nil, false, false } // podBackendURL constructs a backend URL that targets the specific pod IP captured @@ -962,6 +1062,29 @@ func isRedirectStatus(code int) bool { } } +// unpinSession clears the backend_url pin on the session identified by +// clientSID so the next request carrying it is routed to the target service +// rather than to a pod address that is no longer reachable. It is a no-op for +// an empty or unknown session ID. +// +// This is the same fallback reinitializeAndReplay applies when it has no stored +// initialize body to replay, reused for the initialize case where issuing the +// new handshake is the client's job, not the proxy's. +func (r *backendRecovery) unpinSession(clientSID string) { + if clientSID == "" { + return + } + sess, ok := r.sessions.Get(normalizeSessionID(clientSID)) + if !ok { + return + } + sess.SetMetadata(sessionMetadataBackendURL, r.targetURI) + if err := r.sessions.UpsertSession(sess); err != nil { + slog.Debug("failed to unpin session after dial error on initialize", + "session_id", clientSID, "error", err) + } +} + // reinitializeAndReplay is called when the proxy detects that the backend pod // that owned a session is no longer reachable (dial error) or has lost its // in-memory session state (backend returned 404). It transparently: @@ -1262,10 +1385,7 @@ func (p *TransparentProxy) Start(ctx context.Context) error { // 5. Catch-all proxy handler (least specific - ServeMux routing handles precedence) // Note: No manual path checking needed - ServeMux longest-match routing ensures // more specific paths registered above take precedence over this catch-all. - // In stateless mode, wrap with a method gate that rejects GET/DELETE with 405. - if p.stateless { - finalHandler = statelessMethodGate(finalHandler) - } + finalHandler = p.methodGate(finalHandler) mux.Handle("/", finalHandler) // Use ListenConfig with SO_REUSEADDR to allow port reuse after unclean shutdown @@ -1546,15 +1666,27 @@ func (*TransparentProxy) ForwardResponseToClients(_ context.Context, _ jsonrpc2. return fmt.Errorf("ForwardResponseToClients not implemented for TransparentProxy") } -// statelessMethodGate wraps a handler to reject GET, HEAD, and DELETE requests with 405. -// Used in stateless mode where the server only supports POST. -// HEAD is blocked alongside GET because HEAD is semantically a GET without a response body; -// a server that cannot handle GET will not handle HEAD either. -func statelessMethodGate(next http.Handler) http.Handler { +// methodGate wraps a handler to reject GET, HEAD, and DELETE requests with 405 +// when the proxy is running stateless (--stateless) or the request is tagged +// with the Modern (stateless-only) MCP-Protocol-Version: neither has any +// protocol use for GET/HEAD/DELETE, since Modern is unary request/response +// over POST only. +// +// This check runs outermost, before auth and any other middleware: a method +// rejection needs no identity and reveals nothing auth-gated. It is +// deliberately header-only and never reads or otherwise touches r.Body, so it +// must not call mcp.ClassifyRevision (which requires a body) or consume the +// body in any way — doing so would break the body for downstream handlers. +// +// HEAD is blocked alongside GET because HEAD is semantically a GET without a +// response body; a server that cannot handle GET will not handle HEAD either. +func (p *TransparentProxy) methodGate(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodDelete { + isGateableMethod := r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodDelete + isModern := r.Header.Get("MCP-Protocol-Version") == mcp.MCPVersionModern + if isGateableMethod && (p.stateless || isModern) { w.Header().Set("Allow", "POST, OPTIONS") - http.Error(w, "method not allowed: server is stateless (POST only)", http.StatusMethodNotAllowed) + http.Error(w, "method not allowed: server is stateless / stateless MCP revision (POST only)", http.StatusMethodNotAllowed) return } next.ServeHTTP(w, r) diff --git a/pkg/transport/session/jsonrpc_errors.go b/pkg/transport/session/jsonrpc_errors.go index 5083642f96..441396336d 100644 --- a/pkg/transport/session/jsonrpc_errors.go +++ b/pkg/transport/session/jsonrpc_errors.go @@ -24,8 +24,20 @@ const ( // NotFoundBody returns the JSON-encoded body for a session-not-found // JSON-RPC error response. The requestID is the "id" from the incoming -// JSON-RPC request; pass nil when the request ID is not available (e.g., -// DELETE requests, batch pre-parse, transparent proxy). +// JSON-RPC request, echoed so a client correlating by id can match this error +// to the request that caused it. +// +// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error +// response as `id?: RequestId` where RequestId = string | number, so an +// absent id is encoded by omitting the "id" key entirely, never as null. The +// reference TypeScript SDK enforces this with a .strict() schema and a +// throwing parse(), so an "id":null response crashes a conformant client's +// transport. See HasJSONRPCID. +// +// Pass nil only when the request genuinely carries no id: a bodiless GET +// (standalone SSE) or DELETE, a notification, or a batch. Do NOT pass nil +// merely because threading the id to the call site is inconvenient — that +// was the asymmetry fixed in #5945. func NotFoundBody(requestID any) []byte { resp := map[string]any{ "jsonrpc": "2.0", @@ -33,17 +45,45 @@ func NotFoundBody(requestID any) []byte { "code": CodeSessionNotFound, "message": MessageSessionNotFound, }, - "id": requestID, + } + if HasJSONRPCID(requestID) { + resp["id"] = requestID } data, err := json.Marshal(resp) if err != nil { // This should never happen with simple map types, but return a // hand-crafted fallback to guarantee a valid JSON-RPC error. - return []byte(`{"jsonrpc":"2.0","error":{"code":-32001,"message":"Session not found"},"id":null}`) + return []byte(`{"jsonrpc":"2.0","error":{"code":-32001,"message":"Session not found"}}`) } return data } +// HasJSONRPCID reports whether requestID is a usable JSON-RPC id, i.e. whether +// callers should emit an "id" key at all. +// +// MCP narrows base JSON-RPC here: schema/2025-11-25 types the error response as +// `id?: RequestId` where RequestId = string | number, so an absent id is encoded +// by omitting the key, never as null. The reference TypeScript SDK enforces this +// with a .strict() schema and a throwing parse(), so "id":null crashes a +// conformant client's transport. +// +// The transparent proxy threads the incoming id through as raw bytes, so a +// json.RawMessage holding literal null -- or nothing -- is an absent id even +// though the interface value is non-nil. +// +// Do not pass a typed jsonrpc2.ID here: its zero value is a non-nil interface +// holding an empty struct, so this predicate would (wrongly) report it as +// present and callers would emit "id":{}. Use id.IsValid() instead. +func HasJSONRPCID(requestID any) bool { + if requestID == nil { + return false + } + if raw, ok := requestID.(json.RawMessage); ok { + return len(raw) > 0 && string(raw) != "null" + } + return true +} + // WriteNotFound writes an HTTP 404 response with a JSON-RPC error body // for session-not-found. Use this with http.ResponseWriter in the streamable // HTTP and SSE proxies. @@ -57,8 +97,12 @@ func WriteNotFound(w http.ResponseWriter, requestID any) { // NotFoundResponse constructs an *http.Response with HTTP 404 and a // JSON-RPC error body. Use this in httputil.ReverseProxy.ModifyResponse // (transparent proxy) where no http.ResponseWriter is available. -func NotFoundResponse(req *http.Request) *http.Response { - body := NotFoundBody(nil) +// +// requestID follows NotFoundBody: pass the incoming request's JSON-RPC id so the +// error echoes it, or nil when the request has none. Callers that reject a +// request before parsing a body (GET/DELETE) pass nil. +func NotFoundResponse(req *http.Request, requestID any) *http.Response { + body := NotFoundBody(requestID) hdr := make(http.Header) hdr.Set("Content-Type", "application/json") return &http.Response{ diff --git a/pkg/transport/session/jsonrpc_errors_test.go b/pkg/transport/session/jsonrpc_errors_test.go index 9488bc80f4..594294a468 100644 --- a/pkg/transport/session/jsonrpc_errors_test.go +++ b/pkg/transport/session/jsonrpc_errors_test.go @@ -20,26 +20,46 @@ func TestNotFoundBody(t *testing.T) { tests := []struct { name string requestID any - expectedID any // expected value after JSON round-trip + wantIDKey bool + expectedID any // expected value after JSON round-trip, only checked when wantIDKey }{ { - name: "nil request ID", - requestID: nil, - expectedID: nil, + name: "nil request ID omits id key", + requestID: nil, + wantIDKey: false, + }, + { + name: "raw null request ID omits id key", + requestID: json.RawMessage(`null`), + wantIDKey: false, + }, + { + name: "nil RawMessage omits id key", + requestID: json.RawMessage(nil), + wantIDKey: false, }, { name: "integer request ID", requestID: 42, + wantIDKey: true, expectedID: float64(42), // JSON numbers decode as float64 }, + { + name: "zero request ID is still present", + requestID: 0, + wantIDKey: true, + expectedID: float64(0), + }, { name: "string request ID", requestID: "abc-123", + wantIDKey: true, expectedID: "abc-123", }, { name: "float64 request ID", requestID: float64(7), + wantIDKey: true, expectedID: float64(7), }, } @@ -56,7 +76,12 @@ func TestNotFoundBody(t *testing.T) { // Check JSON-RPC fields assert.Equal(t, "2.0", parsed["jsonrpc"]) - assert.Equal(t, tt.expectedID, parsed["id"]) + + id, ok := parsed["id"] + assert.Equal(t, tt.wantIDKey, ok, `"id" key presence`) + if tt.wantIDKey { + assert.Equal(t, tt.expectedID, id) + } errObj, ok := parsed["error"].(map[string]any) require.True(t, ok, "error field should be an object") @@ -69,6 +94,54 @@ func TestNotFoundBody(t *testing.T) { } } +func TestNotFoundBodyMarshalFallback(t *testing.T) { + t.Parallel() + + // A channel is not JSON-marshalable, forcing json.Marshal to fail so the + // hand-crafted fallback literal is exercised. Since requestID is non-nil + // and not a json.RawMessage, HasJSONRPCID reports it as present, so the + // fallback must still omit "id" per the marshal-failure branch, not echo it. + body := NotFoundBody(make(chan int)) + + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + assert.Equal(t, "2.0", parsed["jsonrpc"]) + _, ok := parsed["id"] + assert.False(t, ok, `fallback body must omit "id" key`) + + errObj, ok := parsed["error"].(map[string]any) + require.True(t, ok, "error field should be an object") + assert.Equal(t, float64(CodeSessionNotFound), errObj["code"]) + assert.Equal(t, MessageSessionNotFound, errObj["message"]) +} + +func TestHasJSONRPCID(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + requestID any + want bool + }{ + {name: "nil", requestID: nil, want: false}, + {name: "raw null", requestID: json.RawMessage(`null`), want: false}, + {name: "raw nil slice", requestID: json.RawMessage(nil), want: false}, + {name: "raw empty slice", requestID: json.RawMessage(""), want: false}, + {name: "raw numeric", requestID: json.RawMessage(`42`), want: true}, + {name: "raw string", requestID: json.RawMessage(`"abc"`), want: true}, + {name: "string", requestID: "abc-123", want: true}, + {name: "integer zero", requestID: 0, want: true}, + {name: "integer", requestID: 42, want: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, HasJSONRPCID(tt.requestID)) + }) + } +} + func TestWriteNotFound(t *testing.T) { t.Parallel() @@ -84,15 +157,87 @@ func TestWriteNotFound(t *testing.T) { func TestNotFoundResponse(t *testing.T) { t.Parallel() - req := httptest.NewRequest(http.MethodPost, "/mcp", nil) - resp := NotFoundResponse(req) + tests := []struct { + name string + requestID any + wantIDKey bool + wantID string // only checked when wantIDKey + }{ + { + // A request with no id of its own: bodiless GET/DELETE, or a + // notification. MCP narrows base JSON-RPC to omit the key here + // rather than emit "id":null (schema/2025-11-25, HasJSONRPCID). + name: "nil id omits the id key", + requestID: nil, + wantIDKey: false, + }, + { + // The transparent proxy threads the incoming id through as raw + // bytes; literal null must be treated the same as a missing id. + name: "raw null id omits the id key", + requestID: json.RawMessage(`null`), + wantIDKey: false, + }, + { + name: "nil RawMessage omits the id key", + requestID: json.RawMessage(nil), + wantIDKey: false, + }, + { + name: "string id is echoed", + requestID: "req-1", + wantIDKey: true, + wantID: `"id":"req-1"`, + }, + { + name: "zero id is echoed, not treated as absent", + requestID: 0, + wantIDKey: true, + wantID: `"id":0`, + }, + { + // The shape the transparent proxy passes: the raw bytes of the + // incoming "id", forwarded verbatim so a numeric id stays numeric + // rather than being coerced to a float or a string (#5945). + name: "raw numeric id is echoed verbatim", + requestID: json.RawMessage(`42`), + wantIDKey: true, + wantID: `"id":42`, + }, + { + name: "raw string id is echoed verbatim", + requestID: json.RawMessage(`"abc"`), + wantIDKey: true, + wantID: `"id":"abc"`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + resp := NotFoundResponse(req, tt.requestID) - assert.Equal(t, http.StatusNotFound, resp.StatusCode) - assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) - assert.Equal(t, req, resp.Request) + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) + assert.Equal(t, req, resp.Request) - body, err := io.ReadAll(resp.Body) - require.NoError(t, err) - assert.Contains(t, string(body), `"code":-32001`) - assert.Contains(t, string(body), `"id":null`) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"code":-32001`) + + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + _, ok := parsed["id"] + assert.Equal(t, tt.wantIDKey, ok, `"id" key presence`) + if tt.wantIDKey { + assert.Contains(t, string(body), tt.wantID) + } + + // ContentLength must match the body actually written, or a client + // reading exactly that many bytes truncates or blocks. + assert.Equal(t, int64(len(body)), resp.ContentLength) + }) + } } diff --git a/pkg/transport/stdio.go b/pkg/transport/stdio.go index 836d3037e9..4a8c9b262c 100644 --- a/pkg/transport/stdio.go +++ b/pkg/transport/stdio.go @@ -68,6 +68,11 @@ type StdioTransport struct { authInfoHandler http.Handler prefixHandlers map[string]http.Handler + // strictProtocolValidation controls whether the streamable HTTP proxy + // rejects an unknown/unsupported MCP-Protocol-Version header with 400. + // See streamable.WithStrictProtocolValidation. + strictProtocolValidation bool + // Mutex for protecting shared state mutex sync.Mutex @@ -139,6 +144,14 @@ func (t *StdioTransport) SetProxyMode(mode types.ProxyMode) { t.proxyMode = mode } +// SetStrictProtocolValidation configures whether the streamable HTTP proxy +// rejects requests carrying an unknown/unsupported MCP-Protocol-Version +// header with HTTP 400. Default false preserves the version-agnostic +// behavior of accepting any header value. +func (t *StdioTransport) SetStrictProtocolValidation(strict bool) { + t.strictProtocolValidation = strict +} + // SetSessionStorage configures a custom session storage backend. // When set, the underlying proxy will use this storage instead of the default // in-memory store, enabling session sharing across replicas (e.g. Redis-backed). @@ -264,6 +277,7 @@ func (t *StdioTransport) streamableProxyOptions() []streamable.Option { return append(opts, streamable.WithAuthInfoHandler(t.authInfoHandler), streamable.WithPrefixHandlers(t.prefixHandlers), + streamable.WithStrictProtocolValidation(t.strictProtocolValidation), ) } diff --git a/pkg/transport/types/transport.go b/pkg/transport/types/transport.go index 1096827531..f94254384a 100644 --- a/pkg/transport/types/transport.go +++ b/pkg/transport/types/transport.go @@ -249,6 +249,12 @@ type Config struct { // TrustProxyHeaders indicates whether to trust X-Forwarded-* headers from reverse proxies TrustProxyHeaders bool + // StrictProtocolValidation enables strict MCP-Protocol-Version validation + // on the streamable HTTP proxy: a present-but-unsupported header value is + // rejected with HTTP 400. Default false accepts any version string. Only + // consulted for the stdio/streamable-HTTP proxy path (see factory.go). + StrictProtocolValidation bool + // ProxyMode is the proxy mode for stdio transport ("sse" or "streamable-http") ProxyMode ProxyMode diff --git a/pkg/tui/inspector.go b/pkg/tui/inspector.go index bb431866f5..2b4bc9ce03 100644 --- a/pkg/tui/inspector.go +++ b/pkg/tui/inspector.go @@ -14,9 +14,9 @@ import ( "github.com/charmbracelet/bubbles/textinput" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" ) diff --git a/pkg/tui/inspector_test.go b/pkg/tui/inspector_test.go index 2bff458729..a249a68406 100644 --- a/pkg/tui/inspector_test.go +++ b/pkg/tui/inspector_test.go @@ -7,10 +7,10 @@ import ( "testing" "github.com/charmbracelet/bubbles/textinput" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" ) diff --git a/pkg/tui/model.go b/pkg/tui/model.go index d1b1f2258b..8318c9de55 100644 --- a/pkg/tui/model.go +++ b/pkg/tui/model.go @@ -10,9 +10,9 @@ import ( "github.com/charmbracelet/bubbles/textinput" "github.com/charmbracelet/bubbles/viewport" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" regtypes "github.com/stacklok/toolhive-core/registry/types" "github.com/stacklok/toolhive/pkg/core" "github.com/stacklok/toolhive/pkg/registry" diff --git a/pkg/tui/tools.go b/pkg/tui/tools.go index e0e3c2fa06..e28ba07fd2 100644 --- a/pkg/tui/tools.go +++ b/pkg/tui/tools.go @@ -8,9 +8,9 @@ import ( "errors" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" thclient "github.com/stacklok/toolhive/pkg/mcp/client" ) diff --git a/pkg/tui/update.go b/pkg/tui/update.go index efbdd3d186..c4614ebdb4 100644 --- a/pkg/tui/update.go +++ b/pkg/tui/update.go @@ -9,9 +9,9 @@ import ( "time" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" "github.com/stacklok/toolhive/pkg/runner" ) diff --git a/pkg/tui/update_inspector.go b/pkg/tui/update_inspector.go index eb41e4298d..79997b87b5 100644 --- a/pkg/tui/update_inspector.go +++ b/pkg/tui/update_inspector.go @@ -9,7 +9,8 @@ import ( "github.com/atotto/clipboard" "github.com/charmbracelet/bubbles/key" tea "github.com/charmbracelet/bubbletea" - "github.com/mark3labs/mcp-go/mcp" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // handleInspectorKey handles key input when the inspector panel is active. diff --git a/pkg/tui/update_navigation.go b/pkg/tui/update_navigation.go index c8f4a6993e..4d07fa63ad 100644 --- a/pkg/tui/update_navigation.go +++ b/pkg/tui/update_navigation.go @@ -9,8 +9,8 @@ import ( "github.com/atotto/clipboard" "github.com/charmbracelet/bubbles/key" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" "github.com/stacklok/toolhive/pkg/core" "github.com/stacklok/toolhive/pkg/runner" types "github.com/stacklok/toolhive/pkg/transport/types" diff --git a/pkg/tui/view_inspector.go b/pkg/tui/view_inspector.go index 5ddec4ab30..e883eda658 100644 --- a/pkg/tui/view_inspector.go +++ b/pkg/tui/view_inspector.go @@ -9,8 +9,8 @@ import ( "strings" "github.com/charmbracelet/lipgloss" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/cmd/thv/app/ui" ) diff --git a/pkg/vmcp/aggregator/advertised_backendid_test.go b/pkg/vmcp/aggregator/advertised_backendid_test.go index 3e0678ea6b..14b9d23a4f 100644 --- a/pkg/vmcp/aggregator/advertised_backendid_test.go +++ b/pkg/vmcp/aggregator/advertised_backendid_test.go @@ -265,6 +265,53 @@ func TestDefaultAggregator_AdvertisingFilterPreservesBackendID(t *testing.T) { }, wantAdvertised: []string{"be_tool_a"}, }, + { + // Backward compatibility: a config predating DefaultToolVisibility leaves it + // "" and must keep advertising everything. + name: "unset DefaultToolVisibility advertises an unlisted backend", + aggCfg: &config.AggregationConfig{}, + wantAdvertised: []string{"be_tool_a", "be_tool_b"}, + }, + { + name: "explicit DefaultToolVisibility allow advertises an unlisted backend", + aggCfg: &config.AggregationConfig{DefaultToolVisibility: config.DefaultToolVisibilityAllow}, + wantAdvertised: []string{"be_tool_a", "be_tool_b"}, + }, + { + // The #6073 case: a backend nobody listed contributes nothing. + name: "DefaultToolVisibility deny hides an unlisted backend's tools", + aggCfg: &config.AggregationConfig{DefaultToolVisibility: config.DefaultToolVisibilityDeny}, + wantAdvertised: nil, + }, + { + // A listed backend is opted in by its entry: deny governs only backends + // absent from Tools, so Filter still decides which of its tools show. + name: "DefaultToolVisibility deny leaves a listed backend's Filter in charge", + aggCfg: &config.AggregationConfig{ + DefaultToolVisibility: config.DefaultToolVisibilityDeny, + Tools: []*config.WorkloadToolConfig{{Workload: backendID, Filter: []string{"tool_a"}}}, + }, + wantAdvertised: []string{"be_tool_a"}, + }, + { + // An entry with neither ExcludeAll nor Filter opts the backend in fully + // under deny — "listed means allowed", so an overrides-only entry does + // not silently blank the backend out. + name: "DefaultToolVisibility deny advertises a listed backend with no filter", + aggCfg: &config.AggregationConfig{ + DefaultToolVisibility: config.DefaultToolVisibilityDeny, + Tools: []*config.WorkloadToolConfig{{Workload: backendID}}, + }, + wantAdvertised: []string{"be_tool_a", "be_tool_b"}, + }, + { + name: "DefaultToolVisibility deny with a listed backend's ExcludeAll still hides", + aggCfg: &config.AggregationConfig{ + DefaultToolVisibility: config.DefaultToolVisibilityDeny, + Tools: []*config.WorkloadToolConfig{{Workload: backendID, ExcludeAll: true}}, + }, + wantAdvertised: nil, + }, } for _, tt := range tests { diff --git a/pkg/vmcp/aggregator/aggregator.go b/pkg/vmcp/aggregator/aggregator.go index 4bc9b81d9b..6e66ba51e0 100644 --- a/pkg/vmcp/aggregator/aggregator.go +++ b/pkg/vmcp/aggregator/aggregator.go @@ -74,6 +74,9 @@ type BackendCapabilities struct { // Resources are the resources exposed by this backend. Resources []vmcp.Resource + // ResourceTemplates are the resource templates exposed by this backend. + ResourceTemplates []vmcp.ResourceTemplate + // Prompts are the prompts exposed by this backend. Prompts []vmcp.Prompt @@ -85,17 +88,30 @@ type BackendCapabilities struct { } // ResolvedCapabilities contains capabilities after conflict resolution. -// Tool names are now unique (after prefixing, priority, or manual resolution). +// Every capability identity (tool name, resource URI, resource template +// string, prompt name) is unique within its list. type ResolvedCapabilities struct { // Tools are the conflict-resolved tools. // Map key is the resolved tool name, value contains original name and backend. Tools map[string]*ResolvedTool - // Resources are passed through (conflicts rare, namespaced by URI). + // Resources are de-duplicated by URI: a URI advertised by multiple backends + // appears once, from the backend earliest in sorted-backend-ID order. URIs + // are locators the client passes back verbatim, so they are never rewritten. + // See resolveResourceConflicts. Resources []vmcp.Resource - // Prompts are passed through (conflicts rare, namespaced by name). - Prompts []vmcp.Prompt + // ResourceTemplates are de-duplicated by URI template string, with the same + // locator-identity policy as Resources. See resolveResourceTemplateConflicts. + ResourceTemplates []vmcp.ResourceTemplate + + // Prompts are conflict-resolved by name: by default every prompt is + // renamed to its backend-prefixed form; under the priority strategy, + // backends listed in priorityOrder keep their bare names. Either way the + // advertised name is a pure function of the aggregation config and + // (backendID, name) — it never shifts with group membership. See + // resolvePromptConflicts. + Prompts []ResolvedPrompt // SupportsLogging is true if any backend supports logging. SupportsLogging bool @@ -131,6 +147,20 @@ type ResolvedTool struct { ConflictResolutionApplied vmcp.ConflictResolutionStrategy } +// ResolvedPrompt represents a prompt after conflict resolution. The embedded +// Prompt is the advertised form: Name holds the resolved (client-visible) +// name. OriginalName is the name the backend itself uses; prompts/get and +// completion requests are translated back to it via +// BackendTarget.GetBackendCapabilityName, exactly like renamed tools. +type ResolvedPrompt struct { + vmcp.Prompt + + // OriginalName is the prompt's name in the backend (equal to Name only + // for priority-listed backends; otherwise Name is the backend-prefixed + // form). + OriginalName string +} + // AggregatedCapabilities is the final unified view of all backend capabilities. // This is what gets exposed to MCP clients via tools/list, resources/list, prompts/list. type AggregatedCapabilities struct { @@ -145,6 +175,9 @@ type AggregatedCapabilities struct { // Resources are the aggregated resources. Resources []vmcp.Resource + // ResourceTemplates are the aggregated resource templates. + ResourceTemplates []vmcp.ResourceTemplate + // Prompts are the aggregated prompts. Prompts []vmcp.Prompt @@ -176,6 +209,9 @@ type AggregationMetadata struct { // ResourceCount is the total number of resources. ResourceCount int + // ResourceTemplateCount is the total number of resource templates. + ResourceTemplateCount int + // PromptCount is the total number of prompts. PromptCount int @@ -203,6 +239,25 @@ type ToolOverride interface { ApplyOverrides(ctx context.Context, tools []vmcp.Tool) ([]vmcp.Tool, error) } +// CacheInvalidator is optionally implemented by an Aggregator that memoizes +// AggregateCapabilities results (see cachingAggregator). It lets a caller force +// a re-sweep of all cached entries after learning, out of band, that backend +// capabilities changed — e.g. a persistent backend connection observing +// notifications/tools/list_changed (#5748) — rather than waiting out the TTL. +// +// InvalidateAll purges the ENTIRE cache (every identity's entry), not just the +// backend that changed: the cache has no per-backend index, and coarse +// invalidation briefly de-optimizes other identities' cached views rather than +// leaving any identity looking at stale capabilities. Callers that type-assert +// an Aggregator to this interface must handle the case where it is not +// implemented (a non-caching or differently-implemented Aggregator) — see +// core.coreVMCP.InvalidateCapabilityCache for the WARN-log fallback. +type CacheInvalidator interface { + // InvalidateAll purges every cached AggregateCapabilities entry so the next + // call for any identity re-sweeps the backends. + InvalidateAll() +} + // Common aggregation errors. var ( // ErrNoBackendsFound indicates no backends were discovered. diff --git a/pkg/vmcp/aggregator/caching_aggregator.go b/pkg/vmcp/aggregator/caching_aggregator.go index 3c34973b7e..dc2ac99964 100644 --- a/pkg/vmcp/aggregator/caching_aggregator.go +++ b/pkg/vmcp/aggregator/caching_aggregator.go @@ -114,6 +114,17 @@ func (c *cachingAggregator) AggregateCapabilities( return caps, nil } +// Compile-time assertion: cachingAggregator implements CacheInvalidator. +var _ CacheInvalidator = (*cachingAggregator)(nil) + +// InvalidateAll implements CacheInvalidator by purging every cached entry, so the +// next AggregateCapabilities call for any identity re-sweeps the backends rather +// than serving a cached view until its TTL expires. See CacheInvalidator's doc +// for why this is coarse (whole-cache) rather than per-backend. +func (c *cachingAggregator) InvalidateAll() { + c.cache.Purge() +} + // cacheKey derives a collision-resistant key from the inputs that drive backend enumeration: // the caller's subject, the forwarded headers (passthrough credentials/scopes), and the // backend ID set. Hashing keeps raw credential values out of the cache keys. diff --git a/pkg/vmcp/aggregator/caching_aggregator_test.go b/pkg/vmcp/aggregator/caching_aggregator_test.go index 1fa640520f..91acc35d26 100644 --- a/pkg/vmcp/aggregator/caching_aggregator_test.go +++ b/pkg/vmcp/aggregator/caching_aggregator_test.go @@ -104,6 +104,34 @@ func TestCachingAggregator_DisabledWhenTTLNonPositive(t *testing.T) { assert.Nil(t, aggregator.NewCachingAggregator(nilAgg, time.Hour)) } +// TestCachingAggregator_InvalidateAll verifies the CacheInvalidator seam +// (#5748): the returned Aggregator satisfies aggregator.CacheInvalidator, and +// calling InvalidateAll purges the entire cache so a subsequent call for the +// SAME identity re-sweeps rather than serving the (still TTL-fresh) cached +// view. +func TestCachingAggregator_InvalidateAll(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + mock := mocks.NewMockAggregator(ctrl) + mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil).Times(2) + + c := aggregator.NewCachingAggregator(mock, time.Hour) + invalidator, ok := c.(aggregator.CacheInvalidator) + require.True(t, ok, "cachingAggregator must implement CacheInvalidator") + + ctx := ctxWithSubject("alice") + _, err := c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err) + + invalidator.InvalidateAll() + + // Well within the TTL, so a hit would prove the purge did not happen; the + // mock's Times(2) expectation is the actual re-sweep assertion. + _, err = c.AggregateCapabilities(ctx, testBackends) + require.NoError(t, err, "the second call after InvalidateAll must re-sweep, not hit a stale entry") +} + // TestCachingAggregator_ErrorNotCached: a failed sweep is not cached, so the next call retries // the wrapped aggregator. func TestCachingAggregator_ErrorNotCached(t *testing.T) { diff --git a/pkg/vmcp/aggregator/capability_conflicts.go b/pkg/vmcp/aggregator/capability_conflicts.go new file mode 100644 index 0000000000..9eeae4ad77 --- /dev/null +++ b/pkg/vmcp/aggregator/capability_conflicts.go @@ -0,0 +1,274 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package aggregator + +import ( + "log/slog" + "maps" + "slices" + "sort" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// This file resolves cross-backend identity conflicts for the list +// capabilities NOT covered by the ConflictResolver strategies: resources +// (identity: URI), resource templates (identity: URITemplate) and prompts +// (identity: Name). Tools are resolved separately through +// ConflictResolver.ResolveToolConflicts. +// +// ResolveConflicts is the ENFORCEMENT POINT for these policies: after it +// runs, every identity is unique within its list. The first-wins guards in +// the merge helpers (mergeResources/mergeResourceTemplates/mergePrompts in +// default_aggregator.go) re-check the same invariant, but only as defence in +// depth for direct MergeCapabilities callers — in the aggregation pipeline +// they can never fire. Policy changes belong here, not there. +// +// The policies deliberately differ per type: +// +// - Resource URIs and template strings are LOCATORS, not names. The client +// passes them back verbatim (resources/read, resources/subscribe, +// completion refs), backends emit them in notifications and embedded +// resource contents, and a template-matched read forwards the client's +// concrete URI untranslated (see MergeCapabilities). Rewriting them would +// break those round trips, so a duplicated URI is instead advertised +// ONCE: the backend earliest in sorted-backend-ID order wins and later +// duplicates are dropped with a warning naming the URI and both backends. +// Reads strictly improve — the routing table keys by URI, so at most one +// backend was ever served per URI, previously a nondeterministically- +// chosen one, now a stable one. Listing does regress from N +// indistinguishable entries to one, but the protocol gives a client no +// way to address the other N-1, so the extra entries promised reads that +// could never be honoured. +// +// - Prompt names are NAMES, like tool names: prompts/get is translated back +// to the backend's own name via BackendTarget.GetBackendCapabilityName, +// so renaming is lossless. By DEFAULT every prompt is renamed to its +// backend-prefixed form — the configured prefixFormat applied to the +// backend ID (default "{workload}_") — mirroring what +// PrefixConflictResolver does for every tool. Under the PRIORITY strategy +// (the same escape hatch tools have for clients that pin names), backends +// listed in priorityOrder keep their bare prompt names, a bare-name +// collision among listed backends resolves to the highest-priority one, +// and unlisted backends stay ALWAYS prefixed — deliberately stricter +// than the tool priority resolver, which lets a conflict-free unlisted +// tool keep its bare name. +// +// Under priority the losing prompt is DROPPED, never re-prefixed. That is +// deliberate and must stay that way: re-prefixing the loser to "save" it +// would re-advertise the same prompt under a name no policy mentions, so +// a forbid written against the bare name would stop matching it — a +// fail-open. Do not "improve" this by prefixing the loser instead of +// dropping it. +// +// The invariant both modes preserve, scoped precisely: the advertised +// name of a given (backendID, prompt name) PAIR is a pure function of the +// aggregation config and that pair — it NEVER shifts because an unrelated +// backend joined or left the group. That stability is a security property, +// not cosmetics: the advertised name is what authorization matches on +// (Cedar builds Prompt::"" entities — see +// pkg/vmcp/core/admission.go), so a membership-dependent rename would +// silently detach permit AND forbid policies from the prompt they were +// written for, the forbid case failing open. Names move only on an +// explicit config edit — the moment an operator reviews policy anyway. +// +// What the invariant does NOT promise is that a given advertised STRING +// keeps naming the same prompt. Under priority two listed backends can +// claim the same bare name, and the higher-ranked one takes it: a +// permit(... resource == Prompt::"review") written while b1 owned "review" +// silently begins authorizing b2's DIFFERENT prompt once a higher-ranked +// b2 advertising that name deploys, with no config edit. Cedar's resource +// identity is Prompt:: and nothing else — it carries no +// backend attribute — so whoever wins a shared name inherits every policy +// written for it. forbid still fails closed, because the loser is dropped +// rather than aliased; it is permit that gets redirected. Operators who +// write name-scoped permits must therefore review priorityOrder and their +// policy set together. +// +// A name no single backend can own is advertised by NOBODY. When two +// distinct (backendID, name) pairs compose to the same advertised string +// (backend "b1" prompt "x_y" and backend "b1_x" prompt "y" both compose to +// "b1_x_y"), EVERY colliding prompt is dropped and the collision is logged +// at ERROR. Failing the aggregation instead would take out tools/list, +// resources/list, resource-templates/list, prompts/list, Discover's +// presence flags and backend visibility for the whole group over one +// ambiguous prompt name — and that name needs no conflict-resolution +// config to reach, just an unlucky combination of operator-chosen workload +// names and backend-chosen prompt names. Dropping every claimant is the +// safe outcome rather than the convenient one: keeping a survivor would +// hand it the ambiguous name and with it any policy written for the other +// prompt, whereas an unadvertised name makes every permit and forbid on it +// vacuous, so nothing is reachable under an ambiguous identity and the +// rest of the group keeps serving. +// +// Prompt-set changes on a live backend propagate to connected sessions +// through the list_changed resync path +// (pkg/vmcp/server/serve_list_changed.go), but that path is add-only for +// prompts, so a removed prompt stays advertised on already-registered +// sessions until they reconnect (stacklok/toolhive-core#184). A backend +// joining or leaving the group emits no push signal today, so staleness +// there is bounded by the aggregator cache TTL for routing and by the +// client's next prompts/list for listing. +// +// All functions iterate backends in the caller-supplied deterministic order +// and return key-sorted slices, so the aggregated view (and therefore the +// routing table built from it) is stable run to run. + +// resolveResourceConflicts de-duplicates resources by URI across backends. +// backendIDs supplies the deterministic (sorted) iteration order. +func resolveResourceConflicts( + backendIDs []string, + capabilities map[string]*BackendCapabilities, +) []vmcp.Resource { + return dedupeByKey(backendIDs, capabilities, "resource", + func(caps *BackendCapabilities) []vmcp.Resource { return caps.Resources }, + func(resource vmcp.Resource) string { return resource.URI }) +} + +// resolveResourceTemplateConflicts de-duplicates resource templates by URI +// template string across backends. backendIDs supplies the deterministic +// (sorted) iteration order. +func resolveResourceTemplateConflicts( + backendIDs []string, + capabilities map[string]*BackendCapabilities, +) []vmcp.ResourceTemplate { + return dedupeByKey(backendIDs, capabilities, "resource_template", + func(caps *BackendCapabilities) []vmcp.ResourceTemplate { return caps.ResourceTemplates }, + func(template vmcp.ResourceTemplate) string { return template.URITemplate }) +} + +// resolvePromptConflicts forms every prompt's advertised name per the naming +// config: backends listed in naming.priorityRank keep their bare name, all +// others get the backend-prefixed form (prefixFormat applied to the backend +// ID, see applyPrefixFormat). With no priority configuration every prompt is +// prefixed. Either way the advertised name never depends on what else is +// deployed; see the file header for why that matters to authorization. +// +// An exact duplicate within one backend (same backend advertises a name +// twice) is dropped with a warning — both entries would advertise and route +// identically, so nothing reachable is lost. A bare-name collision among +// priority-listed backends resolves to the highest-priority one (lowest +// rank), dropping the others with a warning, mirroring the tool priority +// strategy. Any other collision — two advertised names composing to the same +// string (backend "b1" prompt "x_y" vs backend "b1_x" prompt "y"), or a +// prefixed name hitting a listed backend's literal bare name — has no defined +// owner, so EVERY prompt claiming that name is dropped and the collision is +// logged at ERROR. This never fails: the group keeps serving the prompts it +// can name unambiguously, and the ambiguous name is advertised by nobody. See +// the file header for why dropping all claimants beats both erroring and +// keeping a survivor. backendIDs supplies the deterministic (sorted) iteration +// order; the result is sorted by resolved name. +func resolvePromptConflicts( + naming promptNaming, + backendIDs []string, + capabilities map[string]*BackendCapabilities, +) []ResolvedPrompt { + // candidate is one backend's claim on an advertised prompt name. + type candidate struct { + prompt vmcp.Prompt + backendID string + rank int // priority rank (lower wins); meaningful only when listed + listed bool + } + + byName := make(map[string][]candidate) + for _, backendID := range backendIDs { + rank, listed := naming.priorityRank[backendID] + for _, prompt := range capabilities[backendID].Prompts { + resolvedName := prompt.Name + if !listed { + resolvedName = applyPrefixFormat(naming.prefixFormat, backendID, prompt.Name) + } + // The advertised name is injective per backend (bare, or a fixed + // prefix plus the name), so a same-backend claim on this name can + // only be an exact duplicate of the same prompt. + duplicate := slices.ContainsFunc(byName[resolvedName], func(c candidate) bool { + return c.backendID == backendID + }) + if duplicate { + slog.Warn("backend advertises the same prompt name twice, dropping later duplicate", + "backend", backendID, "prompt", prompt.Name) + continue + } + byName[resolvedName] = append(byName[resolvedName], candidate{ + prompt: prompt, backendID: backendID, rank: rank, listed: listed, + }) + } + } + + resolved := make([]ResolvedPrompt, 0, len(byName)) + for _, resolvedName := range slices.Sorted(maps.Keys(byName)) { + candidates := byName[resolvedName] + + // Distinct backends claiming one advertised name is resolvable only + // when every claimant is priority-listed: then all claims are the same + // bare name and rank decides. Anything else — prefix composition + // ambiguity, or a prefixed name hitting a listed backend's literal + // name — has no defined owner, so no claimant may keep the name. + // Promoting one would give it whatever policy was written for the + // others, since Cedar authorizes on the advertised name alone. + if len(candidates) > 1 && slices.ContainsFunc(candidates, func(c candidate) bool { return !c.listed }) { + // Index-aligned with droppedPrompts: entry i names the backend that + // claimed original prompt name i. + droppedBackends := make([]string, 0, len(candidates)) + droppedPrompts := make([]string, 0, len(candidates)) + for _, c := range candidates { + droppedBackends = append(droppedBackends, c.backendID) + droppedPrompts = append(droppedPrompts, c.prompt.Name) + } + slog.Error("advertised prompt name is ambiguous between backends, dropping every colliding prompt; "+ + "rename one of them to make the name reachable again", + "prompt", resolvedName, + "dropped_backends", droppedBackends, + "dropped_prompts", droppedPrompts) + continue + } + + winner := candidates[0] + for _, contender := range candidates[1:] { + loser := contender + if contender.rank < winner.rank { + loser, winner = winner, contender + } + // The loser is dropped, not re-prefixed; see the file header for + // why re-prefixing it would fail policy open. + slog.Warn("prompt name advertised by multiple priority-listed backends, keeping highest priority", + "prompt", resolvedName, "kept_backend", winner.backendID, "dropped_backend", loser.backendID) + } + + resolvedPrompt := ResolvedPrompt{Prompt: winner.prompt, OriginalName: winner.prompt.Name} + resolvedPrompt.Name = resolvedName + resolved = append(resolved, resolvedPrompt) + } + return resolved +} + +// dedupeByKey collects items of one capability kind across backends in the +// given deterministic order, keeping the first item seen for each key and +// dropping later duplicates with a warning. The result is sorted by key. +func dedupeByKey[T any]( + backendIDs []string, + capabilities map[string]*BackendCapabilities, + kind string, + itemsOf func(*BackendCapabilities) []T, + keyOf func(T) string, +) []T { + seen := make(map[string]string) // key -> winning backend ID + var deduped []T + for _, backendID := range backendIDs { + for _, item := range itemsOf(capabilities[backendID]) { + key := keyOf(item) + if winner, duplicate := seen[key]; duplicate { + slog.Warn("duplicate capability identity across backends, keeping first in sorted backend order", + "capability", kind, "key", key, "kept_backend", winner, "dropped_backend", backendID) + continue + } + seen[key] = backendID + deduped = append(deduped, item) + } + } + + sort.Slice(deduped, func(i, j int) bool { return keyOf(deduped[i]) < keyOf(deduped[j]) }) + return deduped +} diff --git a/pkg/vmcp/aggregator/capability_conflicts_test.go b/pkg/vmcp/aggregator/capability_conflicts_test.go new file mode 100644 index 0000000000..faacf4fdbc --- /dev/null +++ b/pkg/vmcp/aggregator/capability_conflicts_test.go @@ -0,0 +1,581 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package aggregator + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" +) + +// These tests are deliberately adversarial: the bug they pin (#6060) survived +// because every earlier fixture minted strictly unique keys and therefore +// could not express a collision. Each case here makes at least two backends +// share an identity. + +func TestDefaultAggregator_ResolveConflicts_ResourceAndTemplateDedup(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + capabilities map[string]*BackendCapabilities + // wantResources maps advertised URI -> owning backend; the advertised + // list must contain exactly these URIs, sorted. + wantResources map[string]string + // wantTemplates maps advertised URI template -> owning backend. + wantTemplates map[string]string + }{ + { + name: "two backends share a URI, first sorted backend wins", + capabilities: map[string]*BackendCapabilities{ + "beta": { + BackendID: "beta", + Resources: []vmcp.Resource{newTestResource("file:///README.md", "beta")}, + }, + "alpha": { + BackendID: "alpha", + Resources: []vmcp.Resource{ + newTestResource("file:///README.md", "alpha"), + newTestResource("res://only-alpha", "alpha"), + }, + }, + }, + wantResources: map[string]string{ + "file:///README.md": "alpha", + "res://only-alpha": "alpha", + }, + }, + { + name: "three backends share one URI", + capabilities: map[string]*BackendCapabilities{ + "b3": {BackendID: "b3", Resources: []vmcp.Resource{newTestResource("res://shared", "b3")}}, + "b1": {BackendID: "b1", Resources: []vmcp.Resource{newTestResource("res://shared", "b1")}}, + "b2": {BackendID: "b2", Resources: []vmcp.Resource{newTestResource("res://shared", "b2")}}, + }, + wantResources: map[string]string{"res://shared": "b1"}, + }, + { + name: "same backend advertises a URI twice, first occurrence wins", + capabilities: map[string]*BackendCapabilities{ + "solo": { + BackendID: "solo", + Resources: []vmcp.Resource{ + {URI: "res://twice", Name: "first", BackendID: "solo"}, + {URI: "res://twice", Name: "second", BackendID: "solo"}, + }, + }, + }, + wantResources: map[string]string{"res://twice": "solo"}, + }, + { + name: "two backends share a resource template string", + capabilities: map[string]*BackendCapabilities{ + "b2": { + BackendID: "b2", + ResourceTemplates: []vmcp.ResourceTemplate{ + {URITemplate: "file:///logs/{date}.txt", Name: "b2 logs", BackendID: "b2"}, + }, + }, + "b1": { + BackendID: "b1", + ResourceTemplates: []vmcp.ResourceTemplate{ + {URITemplate: "file:///logs/{date}.txt", Name: "b1 logs", BackendID: "b1"}, + {URITemplate: "file:///cfg/{name}.yaml", Name: "b1 cfg", BackendID: "b1"}, + }, + }, + }, + wantTemplates: map[string]string{ + "file:///logs/{date}.txt": "b1", + "file:///cfg/{name}.yaml": "b1", + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + agg := NewDefaultAggregator(nil, nil, nil, nil) + + // Repeat: map iteration order is re-randomized per call, and the + // dedup winner must not depend on it. + for range 5 { + resolved, err := agg.ResolveConflicts(context.Background(), tt.capabilities) + require.NoError(t, err) + + gotResources := make(map[string]string, len(resolved.Resources)) + var lastURI string + for _, res := range resolved.Resources { + _, dup := gotResources[res.URI] + assert.Falsef(t, dup, "URI %q advertised more than once", res.URI) + assert.LessOrEqual(t, lastURI, res.URI, "resources must be sorted by URI") + lastURI = res.URI + gotResources[res.URI] = res.BackendID + } + if tt.wantResources != nil { + assert.Equal(t, tt.wantResources, gotResources) + } + + gotTemplates := make(map[string]string, len(resolved.ResourceTemplates)) + for _, tmpl := range resolved.ResourceTemplates { + _, dup := gotTemplates[tmpl.URITemplate] + assert.Falsef(t, dup, "template %q advertised more than once", tmpl.URITemplate) + gotTemplates[tmpl.URITemplate] = tmpl.BackendID + } + if tt.wantTemplates != nil { + assert.Equal(t, tt.wantTemplates, gotTemplates) + } + } + }) + } +} + +func TestDefaultAggregator_ResolveConflicts_PromptConflicts(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + capabilities map[string]*BackendCapabilities + // want maps advertised prompt name -> {owning backend, original name}. + want map[string][2]string + // wantAbsent lists names that must NOT be advertised. + wantAbsent []string + }{ + { + // Prefixing is unconditional: the advertised name is a pure + // function of (backendID, name), so it cannot shift under a + // client (or a Cedar policy) when unrelated backends join. + name: "unique prompt names are prefixed too", + capabilities: map[string]*BackendCapabilities{ + "b1": {BackendID: "b1", Prompts: []vmcp.Prompt{newTestPrompt("review", "b1")}}, + "b2": {BackendID: "b2", Prompts: []vmcp.Prompt{newTestPrompt("summarize", "b2")}}, + }, + want: map[string][2]string{ + "b1_review": {"b1", "review"}, + "b2_summarize": {"b2", "summarize"}, + }, + wantAbsent: []string{"review", "summarize"}, + }, + { + name: "name shared by two backends stays unambiguous", + capabilities: map[string]*BackendCapabilities{ + "b2": {BackendID: "b2", Prompts: []vmcp.Prompt{newTestPrompt("review", "b2")}}, + "b1": { + BackendID: "b1", + Prompts: []vmcp.Prompt{newTestPrompt("review", "b1"), newTestPrompt("unique", "b1")}, + }, + }, + want: map[string][2]string{ + "b1_review": {"b1", "review"}, + "b2_review": {"b2", "review"}, + "b1_unique": {"b1", "unique"}, + }, + wantAbsent: []string{"review", "unique"}, + }, + { + // Under collision-only renaming this shape dropped a prompt: b2's + // literal "b1_review" collided with b1's renamed "review". + // Unconditional prefixing keeps all three reachable. + name: "literal prefixed name from another backend no longer collides", + capabilities: map[string]*BackendCapabilities{ + "b1": {BackendID: "b1", Prompts: []vmcp.Prompt{newTestPrompt("review", "b1")}}, + "b2": {BackendID: "b2", Prompts: []vmcp.Prompt{newTestPrompt("b1_review", "b2")}}, + "b3": {BackendID: "b3", Prompts: []vmcp.Prompt{newTestPrompt("review", "b3")}}, + }, + want: map[string][2]string{ + "b1_review": {"b1", "review"}, + "b2_b1_review": {"b2", "b1_review"}, + "b3_review": {"b3", "review"}, + }, + wantAbsent: []string{"review"}, + }, + { + // An exact intra-backend duplicate advertises and routes + // identically, so the later occurrence is dropped, not an error. + name: "same backend advertises a prompt name twice", + capabilities: map[string]*BackendCapabilities{ + "solo": { + BackendID: "solo", + Prompts: []vmcp.Prompt{newTestPrompt("dup", "solo"), newTestPrompt("dup", "solo")}, + }, + }, + want: map[string][2]string{"solo_dup": {"solo", "dup"}}, + wantAbsent: []string{"dup"}, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + agg := NewDefaultAggregator(nil, nil, nil, nil) + + for range 5 { + resolved, err := agg.ResolveConflicts(context.Background(), tt.capabilities) + require.NoError(t, err) + + got := make(map[string][2]string, len(resolved.Prompts)) + var lastName string + for _, prompt := range resolved.Prompts { + _, dup := got[prompt.Name] + assert.Falsef(t, dup, "prompt %q advertised more than once", prompt.Name) + assert.LessOrEqual(t, lastName, prompt.Name, "prompts must be sorted by resolved name") + lastName = prompt.Name + got[prompt.Name] = [2]string{prompt.BackendID, prompt.OriginalName} + } + assert.Equal(t, tt.want, got) + for _, absent := range tt.wantAbsent { + assert.NotContains(t, got, absent) + } + } + }) + } +} + +// TestDefaultAggregator_PromptAmbiguityDropsEveryClaimant pins the drop-all +// policy for the one residual collision unconditional prefixing cannot rule +// out: distinct (backendID, name) pairs whose advertised names compose to the +// same string. Two properties matter and both are asserted below. Aggregation +// must NOT fail — an error here would take down the group's whole aggregated +// view (tools, resources, templates, prompts, backend visibility) over one +// prompt name reachable without any conflict-resolution config. And no +// claimant may keep the name: Cedar authorizes on the advertised name alone, +// so a survivor would inherit every policy written for the prompt it collided +// with. Dropping all of them leaves the name advertised by nobody. +func TestDefaultAggregator_PromptAmbiguityDropsEveryClaimant(t *testing.T) { + t.Parallel() + + priorityCfg := &config.AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPriority, + ConflictResolutionConfig: &config.ConflictResolutionConfig{ + PriorityOrder: []string{"b1"}, + }, + } + + tests := []struct { + name string + aggCfg *config.AggregationConfig + // promptsByBackend maps backend ID -> the prompt names it serves. + promptsByBackend map[string][]string + // wantAdvertised maps every SURVIVING advertised name -> {owning + // backend, original name}. Compared for exact equality, so a claimant + // that wrongly survives the ambiguity shows up here. + wantAdvertised map[string][2]string + // wantAmbiguous is the name no backend may advertise or route. + wantAmbiguous string + }{ + { + // "b1" + "x_y" and "b1_x" + "y" both compose to "b1_x_y". + name: "two backends compose to the same prefixed name", + promptsByBackend: map[string][]string{ + "b1": {"x_y", "safe"}, + "b1_x": {"y"}, + }, + wantAdvertised: map[string][2]string{"b1_safe": {"b1", "safe"}}, + wantAmbiguous: "b1_x_y", + }, + { + // Three-way: "b1"+"x_y_z", "b1_x"+"y_z" and "b1_x_y"+"z" all + // compose to "b1_x_y_z". Dropping "the other one" is not enough. + name: "three backends compose to the same prefixed name", + promptsByBackend: map[string][]string{ + "b1": {"x_y_z"}, + "b1_x": {"y_z"}, + "b1_x_y": {"z", "safe"}, + }, + wantAdvertised: map[string][2]string{"b1_x_y_safe": {"b1_x_y", "safe"}}, + wantAmbiguous: "b1_x_y_z", + }, + { + // Priority-listed b1 advertises the literal name "ext_helper"; + // unlisted ext's "helper" prefixes to the same string. Listed + // status does not make a composition collision resolvable, so both + // go, while b1's other bare name is untouched. + name: "prefixed name hits a listed backend's literal bare name", + aggCfg: priorityCfg, + promptsByBackend: map[string][]string{ + "b1": {"ext_helper", "kept"}, + "ext": {"helper"}, + }, + wantAdvertised: map[string][2]string{"kept": {"b1", "kept"}}, + wantAmbiguous: "ext_helper", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + backends := make([]vmcp.Backend, 0, len(tt.promptsByBackend)) + capsByID := make(map[string]*vmcp.CapabilityList, len(tt.promptsByBackend)) + for backendID, names := range tt.promptsByBackend { + backends = append(backends, newTestBackend(backendID)) + prompts := make([]vmcp.Prompt, 0, len(names)) + for _, name := range names { + prompts = append(prompts, newTestPrompt(name, backendID)) + } + capsByID[backendID] = newTestCapabilityList(withPrompts(prompts...)) + } + + ctrl := gomock.NewController(t) + mockClient := mocks.NewMockBackendClient(ctrl) + expectListCapabilities(mockClient, capsByID) + + agg := NewDefaultAggregator(mockClient, nil, tt.aggCfg, nil) + result, err := agg.AggregateCapabilities(context.Background(), backends) + require.NoError(t, err, "one ambiguous prompt name must not fail the group's aggregation") + + // Survivors are read back through the routing table so the case + // also pins that each still translates to its backend's own name. + got := make(map[string][2]string, len(result.Prompts)) + for _, prompt := range result.Prompts { + target := result.RoutingTable.Prompts[prompt.Name] + require.NotNilf(t, target, "advertised prompt %q must be routable", prompt.Name) + got[prompt.Name] = [2]string{prompt.BackendID, target.GetBackendCapabilityName(prompt.Name)} + } + assert.Equal(t, tt.wantAdvertised, got) + assert.NotContains(t, result.RoutingTable.Prompts, tt.wantAmbiguous, + "an ambiguous prompt name must not be routable either") + assert.Equal(t, len(tt.wantAdvertised), result.Metadata.PromptCount) + ctrl.Finish() + }) + } +} + +// TestDefaultAggregator_ResolveConflicts_PromptPriority pins the priority +// escape hatch: backends listed in priorityOrder keep their bare prompt +// names (a bare-name collision resolves by rank), while unlisted backends +// stay always-prefixed even without a collision — the advertised name must +// be decidable from config plus (backendID, name) alone, never from what +// else happens to be deployed. +func TestDefaultAggregator_ResolveConflicts_PromptPriority(t *testing.T) { + t.Parallel() + + // b2 outranks b1 deliberately: sorted-backend-ID order would pick b1, so + // any regression from rank to encounter order flips the winner. + aggCfg := &config.AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPriority, + ConflictResolutionConfig: &config.ConflictResolutionConfig{ + PriorityOrder: []string{"b2", "b1"}, + }, + } + + capabilities := map[string]*BackendCapabilities{ + "b1": {BackendID: "b1", Prompts: []vmcp.Prompt{ + newTestPrompt("review", "b1"), newTestPrompt("unique", "b1"), + }}, + "b2": {BackendID: "b2", Prompts: []vmcp.Prompt{newTestPrompt("review", "b2")}}, + "ext": {BackendID: "ext", Prompts: []vmcp.Prompt{newTestPrompt("helper", "ext")}}, + } + + agg := NewDefaultAggregator(nil, nil, aggCfg, nil) + for range 5 { + resolved, err := agg.ResolveConflicts(context.Background(), capabilities) + require.NoError(t, err) + + got := make(map[string][2]string, len(resolved.Prompts)) + for _, prompt := range resolved.Prompts { + got[prompt.Name] = [2]string{prompt.BackendID, prompt.OriginalName} + } + assert.Equal(t, map[string][2]string{ + // "review" goes to b2: first in priorityOrder, even though b1 + // sorts first. b1's colliding prompt is dropped. + "review": {"b2", "review"}, + // Listed backends keep bare names even for unique prompts. + "unique": {"b1", "unique"}, + // Unlisted backends are ALWAYS prefixed, collision or not: + // exempting them would let a later join shift the name. + "ext_helper": {"ext", "helper"}, + }, got) + assert.NotContains(t, got, "helper") + // The priority loser is dropped, not re-prefixed: re-advertising it as + // "b1_review" would put it beyond any forbid written on "review". + assert.NotContains(t, got, "b1_review") + } + + // A prefixed name colliding with a listed backend's literal bare name is + // covered by TestDefaultAggregator_PromptAmbiguityDropsEveryClaimant, + // alongside the other composition collisions. +} + +// TestDefaultAggregator_ResolveConflicts_PromptPrefixFormat verifies prompt +// renaming honours conflictResolutionConfig.prefixFormat — the same knob the +// tool prefix strategy uses — rather than a hardcoded "_" separator. +func TestDefaultAggregator_ResolveConflicts_PromptPrefixFormat(t *testing.T) { + t.Parallel() + + capabilities := map[string]*BackendCapabilities{ + "github": {BackendID: "github", Prompts: []vmcp.Prompt{newTestPrompt("review", "github")}}, + } + aggCfg := &config.AggregationConfig{ + ConflictResolutionConfig: &config.ConflictResolutionConfig{PrefixFormat: "{workload}."}, + } + + agg := NewDefaultAggregator(nil, nil, aggCfg, nil) + resolved, err := agg.ResolveConflicts(context.Background(), capabilities) + require.NoError(t, err) + + require.Len(t, resolved.Prompts, 1) + assert.Equal(t, "github.review", resolved.Prompts[0].Name) + assert.Equal(t, "review", resolved.Prompts[0].OriginalName) +} + +// TestDefaultAggregator_AggregateCapabilities_CollisionRouting exercises the +// full pipeline with colliding identities across two backends and asserts the +// client-visible outcome: one advertised entry per resource URI routed to a +// deterministic backend, and both prompts advertised under their (always +// prefixed) names whose routing translates back to the backend's own prompt +// name. +func TestDefaultAggregator_AggregateCapabilities_CollisionRouting(t *testing.T) { + t.Parallel() + + backends := []vmcp.Backend{ + newTestBackend("backend1", withBackendName("Backend 1")), + newTestBackend("backend2", withBackendName("Backend 2")), + } + capsByID := map[string]*vmcp.CapabilityList{ + "backend1": newTestCapabilityList( + withResources(newTestResource("file:///README.md", "backend1")), + withPrompts(newTestPrompt("review", "backend1")), + ), + "backend2": newTestCapabilityList( + withResources(newTestResource("file:///README.md", "backend2")), + withPrompts(newTestPrompt("review", "backend2")), + ), + } + + // Aggregate repeatedly: the collision winner must be stable across runs, + // not whichever backend won a map write. + for run := range 5 { + ctrl := gomock.NewController(t) + mockClient := mocks.NewMockBackendClient(ctrl) + expectListCapabilities(mockClient, capsByID) + + agg := NewDefaultAggregator(mockClient, nil, nil, nil) + result, err := agg.AggregateCapabilities(context.Background(), backends) + require.NoError(t, err) + + // The duplicated URI is advertised once and reads route to backend1 + // (first in sorted backend order) -- a defined, stable outcome. + require.Lenf(t, result.Resources, 1, "run %d", run) + assert.Equal(t, "file:///README.md", result.Resources[0].URI) + assert.Equal(t, "backend1", result.Resources[0].BackendID) + target := result.RoutingTable.Resources["file:///README.md"] + require.NotNil(t, target) + assert.Equal(t, "backend1", target.WorkloadID) + + // Both colliding prompts survive under prefixed names; the bare name is + // neither advertised nor routable; prompts/get on a prefixed name + // forwards the backend's own prompt name. + require.Lenf(t, result.Prompts, 2, "run %d", run) + assert.Equal(t, "backend1_review", result.Prompts[0].Name) + assert.Equal(t, "backend2_review", result.Prompts[1].Name) + assert.NotContains(t, result.RoutingTable.Prompts, "review") + for _, advertised := range []string{"backend1_review", "backend2_review"} { + promptTarget := result.RoutingTable.Prompts[advertised] + require.NotNilf(t, promptTarget, "routing target for %q", advertised) + assert.Equal(t, "review", promptTarget.GetBackendCapabilityName(advertised), + "prompts/get on %q must forward the backend's own name", advertised) + } + + assert.Equal(t, 1, result.Metadata.ResourceCount) + assert.Equal(t, 2, result.Metadata.PromptCount) + ctrl.Finish() + } +} + +// TestDefaultAggregator_ResolveConflicts_DuplicateAtPageBoundary places the +// collision at the Modern paginator's page boundary (page size 1000): the exact +// shape that permanently dropped items before the paginator was made +// collision-safe. The aggregator must now hand the paginator a corpus with no +// duplicate keys at all, boundary included. +func TestDefaultAggregator_ResolveConflicts_DuplicateAtPageBoundary(t *testing.T) { + t.Parallel() + + const total = 1100 + const pageBoundary = 1000 // pkg/vmcp/server modernPageSize + + // backend-a advertises 1100 resources whose URIs sort to their index. + // backend-b re-advertises the URIs at sorted positions 999 and 1000, so the + // duplicates straddle the page-1/page-2 boundary. + resourcesA := make([]vmcp.Resource, 0, total) + for i := range total { + resourcesA = append(resourcesA, newTestResource(fmt.Sprintf("res://%05d", i), "backend-a")) + } + capabilities := map[string]*BackendCapabilities{ + "backend-a": {BackendID: "backend-a", Resources: resourcesA}, + "backend-b": {BackendID: "backend-b", Resources: []vmcp.Resource{ + newTestResource(fmt.Sprintf("res://%05d", pageBoundary-1), "backend-b"), + newTestResource(fmt.Sprintf("res://%05d", pageBoundary), "backend-b"), + }}, + } + + agg := NewDefaultAggregator(nil, nil, nil, nil) + resolved, err := agg.ResolveConflicts(context.Background(), capabilities) + require.NoError(t, err) + + require.Len(t, resolved.Resources, total, "duplicates must collapse, nothing else may be lost") + for i, res := range resolved.Resources { + assert.Equal(t, fmt.Sprintf("res://%05d", i), res.URI, "output must be sorted and duplicate-free") + assert.Equal(t, "backend-a", res.BackendID, "backend-a sorts first and owns every duplicated URI") + } +} + +// TestDefaultAggregator_MergeCapabilities_FirstWinsOnDuplicates pins the +// defence-in-depth guard: MergeCapabilities is independently callable, and if +// it is handed unresolved duplicates anyway, the first entry wins in BOTH the +// routing table and the advertised list -- never a silent map overwrite that +// leaves advertising and routing disagreeing. +func TestDefaultAggregator_MergeCapabilities_FirstWinsOnDuplicates(t *testing.T) { + t.Parallel() + + resolved := &ResolvedCapabilities{ + Tools: map[string]*ResolvedTool{}, + Resources: []vmcp.Resource{ + newTestResource("res://dup", "backend1"), + newTestResource("res://dup", "backend2"), + }, + ResourceTemplates: []vmcp.ResourceTemplate{ + {URITemplate: "res://tmpl/{id}", Name: "first", BackendID: "backend1"}, + {URITemplate: "res://tmpl/{id}", Name: "second", BackendID: "backend2"}, + }, + Prompts: []ResolvedPrompt{ + {Prompt: vmcp.Prompt{Name: "dup_prompt", BackendID: "backend1"}, OriginalName: "dup_prompt"}, + {Prompt: vmcp.Prompt{Name: "dup_prompt", BackendID: "backend2"}, OriginalName: "dup_prompt"}, + }, + } + registry := vmcp.NewImmutableRegistry([]vmcp.Backend{ + newTestBackend("backend1"), + newTestBackend("backend2"), + }) + + agg := NewDefaultAggregator(nil, nil, nil, nil) + aggregated, err := agg.MergeCapabilities(context.Background(), resolved, registry) + require.NoError(t, err) + + require.Len(t, aggregated.Resources, 1) + assert.Equal(t, "backend1", aggregated.Resources[0].BackendID) + assert.Equal(t, "backend1", aggregated.RoutingTable.Resources["res://dup"].WorkloadID) + + require.Len(t, aggregated.ResourceTemplates, 1) + assert.Equal(t, "backend1", aggregated.ResourceTemplates[0].BackendID) + assert.Equal(t, "backend1", aggregated.RoutingTable.ResourceTemplates["res://tmpl/{id}"].WorkloadID) + + require.Len(t, aggregated.Prompts, 1) + assert.Equal(t, "backend1", aggregated.Prompts[0].BackendID) + assert.Equal(t, "backend1", aggregated.RoutingTable.Prompts["dup_prompt"].WorkloadID) + + assert.Equal(t, 1, aggregated.Metadata.ResourceCount) + assert.Equal(t, 1, aggregated.Metadata.ResourceTemplateCount) + assert.Equal(t, 1, aggregated.Metadata.PromptCount) +} diff --git a/pkg/vmcp/aggregator/conflict_resolver.go b/pkg/vmcp/aggregator/conflict_resolver.go index aae51c2812..663b8109c2 100644 --- a/pkg/vmcp/aggregator/conflict_resolver.go +++ b/pkg/vmcp/aggregator/conflict_resolver.go @@ -10,22 +10,82 @@ package aggregator import ( "fmt" "log/slog" + "strings" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/config" ) +// defaultPrefixFormat is the prefix format applied when none is configured. +const defaultPrefixFormat = "{workload}_" + +// applyPrefixFormat expands the {workload} placeholder in prefixFormat with +// backendID and prepends the result to name. Shared by tool prefixing +// (PrefixConflictResolver) and unconditional prompt prefixing +// (resolvePromptConflicts) so both compose advertised names identically. +func applyPrefixFormat(prefixFormat, backendID, name string) string { + return strings.ReplaceAll(prefixFormat, "{workload}", backendID) + name +} + +// promptNaming captures how advertised prompt names are formed by +// resolvePromptConflicts. prefixFormat is the format backend-prefixed names +// use. priorityRank maps the backend IDs listed in +// conflictResolutionConfig.priorityOrder to their rank (lower wins); listed +// backends keep their bare prompt names. It is non-nil only under the +// priority strategy — nil means every prompt is prefixed. +type promptNaming struct { + prefixFormat string + priorityRank map[string]int +} + +// promptNamingFromConfig derives prompt naming from the aggregation config. +// Default (prefix or manual strategy, or no config): every prompt is +// prefixed with prefixFormat (conflictResolutionConfig.prefixFormat when +// set), so the advertised name is a pure function of (backendID, name). +// Under the priority strategy, backends listed in priorityOrder keep their +// bare prompt names while unlisted backends stay ALWAYS prefixed — stricter +// than the tool priority resolver, which lets a conflict-free unlisted tool +// keep its bare name. Prompts cannot afford that laxity: their advertised +// name is an authorization identity (see capability_conflicts.go), so the +// only membership-independent choices are "bare because listed" and +// "prefixed because not". +func promptNamingFromConfig(aggregationConfig *config.AggregationConfig) promptNaming { + naming := promptNaming{prefixFormat: defaultPrefixFormat} + if aggregationConfig == nil { + return naming + } + if aggregationConfig.ConflictResolutionConfig != nil && + aggregationConfig.ConflictResolutionConfig.PrefixFormat != "" { + naming.prefixFormat = aggregationConfig.ConflictResolutionConfig.PrefixFormat + } + if aggregationConfig.ConflictResolution == vmcp.ConflictStrategyPriority && + aggregationConfig.ConflictResolutionConfig != nil { + priorityOrder := aggregationConfig.ConflictResolutionConfig.PriorityOrder + if len(priorityOrder) > 0 { + naming.priorityRank = make(map[string]int, len(priorityOrder)) + for rank, backendID := range priorityOrder { + // A backend duplicated in priorityOrder keeps its first + // (highest-priority) position. + if _, seen := naming.priorityRank[backendID]; !seen { + naming.priorityRank[backendID] = rank + } + } + } + } + return naming +} + // NewConflictResolver creates the appropriate conflict resolver based on configuration. func NewConflictResolver(aggregationConfig *config.AggregationConfig) (ConflictResolver, error) { if aggregationConfig == nil { // Default to prefix strategy with default format slog.Info("no aggregation config provided, using default prefix strategy") - return NewPrefixConflictResolver("{workload}_"), nil + return NewPrefixConflictResolver(defaultPrefixFormat), nil } switch aggregationConfig.ConflictResolution { case vmcp.ConflictStrategyPrefix: - prefixFormat := "{workload}_" // Default + prefixFormat := defaultPrefixFormat if aggregationConfig.ConflictResolutionConfig != nil && aggregationConfig.ConflictResolutionConfig.PrefixFormat != "" { prefixFormat = aggregationConfig.ConflictResolutionConfig.PrefixFormat diff --git a/pkg/vmcp/aggregator/default_aggregator.go b/pkg/vmcp/aggregator/default_aggregator.go index efece920ee..581936483b 100644 --- a/pkg/vmcp/aggregator/default_aggregator.go +++ b/pkg/vmcp/aggregator/default_aggregator.go @@ -7,6 +7,8 @@ import ( "context" "fmt" "log/slog" + "maps" + "slices" "sort" "sync" "time" @@ -36,12 +38,21 @@ type defaultAggregator struct { conflictResolver ConflictResolver toolConfigMap map[string]*config.WorkloadToolConfig // Maps backend ID to tool config excludeAllTools bool // Global flag to exclude all tools - tracer trace.Tracer + // denyUnlisted advertises no tools from a backend absent from toolConfigMap + // (config.DefaultToolVisibilityDeny). False — the default — preserves the + // advertise-everything behavior predating the setting. + denyUnlisted bool + // promptNaming controls how advertised prompt names are formed (see + // resolvePromptConflicts). Derived from the aggregation config at + // construction. + promptNaming promptNaming + tracer trace.Tracer } // NewDefaultAggregator creates a new default aggregator implementation. // conflictResolver handles tool name conflicts across backends. -// aggregationConfig specifies aggregation settings including tool filtering/overrides and excludeAllTools. +// aggregationConfig specifies aggregation settings including tool filtering/overrides, +// excludeAllTools, and defaultToolVisibility. // tracerProvider is used to create a tracer for distributed tracing (pass nil for no tracing). func NewDefaultAggregator( backendClient vmcp.BackendClient, @@ -52,9 +63,13 @@ func NewDefaultAggregator( // Build tool config map for quick lookup by backend ID toolConfigMap := make(map[string]*config.WorkloadToolConfig) var excludeAllTools bool + var denyUnlisted bool if aggregationConfig != nil { excludeAllTools = aggregationConfig.ExcludeAllTools + // Only the explicit "deny" opts in; "" (unset) and "allow" both advertise, + // so a config written before this setting existed is unaffected. + denyUnlisted = aggregationConfig.DefaultToolVisibility == config.DefaultToolVisibilityDeny for _, wlConfig := range aggregationConfig.Tools { if wlConfig != nil { toolConfigMap[wlConfig.Workload] = wlConfig @@ -75,6 +90,8 @@ func NewDefaultAggregator( conflictResolver: conflictResolver, toolConfigMap: toolConfigMap, excludeAllTools: excludeAllTools, + denyUnlisted: denyUnlisted, + promptNaming: promptNamingFromConfig(aggregationConfig), tracer: tracer, } } @@ -119,12 +136,13 @@ func (a *defaultAggregator) QueryCapabilities(ctx context.Context, backend vmcp. // Convert to BackendCapabilities result := &BackendCapabilities{ - BackendID: backend.ID, - Tools: processedTools, - Resources: capabilities.Resources, - Prompts: capabilities.Prompts, - SupportsLogging: capabilities.SupportsLogging, - SupportsSampling: capabilities.SupportsSampling, + BackendID: backend.ID, + Tools: processedTools, + Resources: capabilities.Resources, + ResourceTemplates: capabilities.ResourceTemplates, + Prompts: capabilities.Prompts, + SupportsLogging: capabilities.SupportsLogging, + SupportsSampling: capabilities.SupportsSampling, } span.SetAttributes( @@ -201,10 +219,56 @@ func (a *defaultAggregator) QueryAllCapabilities( attribute.Int("successful.backends", len(capabilities)), ) + a.warnUnmatchedToolConfigs(backends) + slog.Info("successfully queried backends", "successful", len(capabilities), "total", len(backends)) return capabilities, nil } +// warnUnmatchedToolConfigs logs a Tools entry naming a workload that no backend +// in the group provides — almost always a typo in tools[].workload. +// +// Under DefaultToolVisibilityAllow a typo is nearly harmless: the entry matches +// nothing, so no filter is applied and the real backend keeps advertising. Under +// DefaultToolVisibilityDeny it inverts — the real backend has no matching entry, so +// it contributes NOTHING, and the only symptom is a short tools/list. The warning +// is what makes that debuggable, so it is loudest for deny. +func (a *defaultAggregator) warnUnmatchedToolConfigs(backends []vmcp.Backend) { + for _, workload := range a.unmatchedToolConfigWorkloads(backends) { + if a.denyUnlisted { + slog.Warn("aggregation.tools entry matches no backend in the group; "+ + "under defaultToolVisibility=deny any backend it was meant to name contributes no tools", + "workload", workload) + continue + } + slog.Warn("aggregation.tools entry matches no backend in the group; its filter/overrides are not applied", + "workload", workload) + } +} + +// unmatchedToolConfigWorkloads returns, sorted, the Tools entries naming a +// workload that no backend in the group provides. Split from the logging so the +// selection is testable without swapping the global logger. +func (a *defaultAggregator) unmatchedToolConfigWorkloads(backends []vmcp.Backend) []string { + if len(a.toolConfigMap) == 0 { + return nil + } + + present := make(map[string]struct{}, len(backends)) + for _, b := range backends { + present[b.ID] = struct{}{} + } + + var unmatched []string + for workload := range a.toolConfigMap { + if _, ok := present[workload]; !ok { + unmatched = append(unmatched, workload) + } + } + sort.Strings(unmatched) + return unmatched +} + // ResolveConflicts applies conflict resolution strategy to handle // duplicate capability names across backends. func (a *defaultAggregator) ResolveConflicts( @@ -242,11 +306,12 @@ func (a *defaultAggregator) ResolveConflicts( return nil, fmt.Errorf("conflict resolution failed: %w", err) } } else { - // Fallback: no conflict resolution (first wins, log warnings) + // Fallback: no conflict resolution (first wins in sorted backend + // order, so the winner is deterministic; log warnings) slog.Warn("no conflict resolver configured, using fallback (first wins)") resolvedTools = make(map[string]*ResolvedTool) - for backendID, tools := range toolsByBackend { - for _, tool := range tools { + for _, backendID := range slices.Sorted(maps.Keys(toolsByBackend)) { + for _, tool := range toolsByBackend[backendID] { if existing, exists := resolvedTools[tool.Name]; exists { slog.Warn("tool name conflict, keeping first", "tool", tool.Name, "existing_backend", existing.BackendID, "conflicting_backend", backendID) @@ -267,16 +332,19 @@ func (a *defaultAggregator) ResolveConflicts( // Build resolved capabilities resolved := &ResolvedCapabilities{ - Tools: resolvedTools, - Resources: []vmcp.Resource{}, - Prompts: []vmcp.Prompt{}, + Tools: resolvedTools, } - // Collect resources and prompts (no conflict resolution for these yet) - for _, caps := range capabilities { - resolved.Resources = append(resolved.Resources, caps.Resources...) - resolved.Prompts = append(resolved.Prompts, caps.Prompts...) + // Resolve conflicts for resources, resource templates, and prompts. + // Backends are iterated in sorted-ID order so that collision outcomes are + // deterministic across runs (map iteration order is not). See + // capability_conflicts.go for the per-type policies. + backendIDs := slices.Sorted(maps.Keys(capabilities)) + resolved.Resources = resolveResourceConflicts(backendIDs, capabilities) + resolved.ResourceTemplates = resolveResourceTemplateConflicts(backendIDs, capabilities) + resolved.Prompts = resolvePromptConflicts(a.promptNaming, backendIDs, capabilities) + for _, caps := range capabilities { // Aggregate logging/sampling support (OR logic - enabled if any backend supports) resolved.SupportsLogging = resolved.SupportsLogging || caps.SupportsLogging resolved.SupportsSampling = resolved.SupportsSampling || caps.SupportsSampling @@ -320,9 +388,10 @@ func (a *defaultAggregator) MergeCapabilities( // Create routing table routingTable := &vmcp.RoutingTable{ - Tools: make(map[string]*vmcp.BackendTarget), - Resources: make(map[string]*vmcp.BackendTarget), - Prompts: make(map[string]*vmcp.BackendTarget), + Tools: make(map[string]*vmcp.BackendTarget), + Resources: make(map[string]*vmcp.BackendTarget), + ResourceTemplates: make(map[string]*vmcp.BackendTarget), + Prompts: make(map[string]*vmcp.BackendTarget), } // Convert resolved tools to final vmcp.Tool format @@ -372,41 +441,16 @@ func (a *defaultAggregator) MergeCapabilities( return tools[i].Name < tools[j].Name }) - // Add resources to routing table - for _, resource := range resolved.Resources { - backend := registry.Get(ctx, resource.BackendID) - if backend == nil { - slog.Warn("backend not found in registry for resource, creating minimal target", - "backend", resource.BackendID, "resource", resource.URI) - routingTable.Resources[resource.URI] = &vmcp.BackendTarget{ - WorkloadID: resource.BackendID, - OriginalCapabilityName: resource.URI, - } - } else { - target := vmcp.BackendToTarget(backend) - // Store the original resource URI for forwarding to backend - target.OriginalCapabilityName = resource.URI - routingTable.Resources[resource.URI] = target - } - } - - // Add prompts to routing table - for _, prompt := range resolved.Prompts { - backend := registry.Get(ctx, prompt.BackendID) - if backend == nil { - slog.Warn("backend not found in registry for prompt, creating minimal target", - "backend", prompt.BackendID, "prompt", prompt.Name) - routingTable.Prompts[prompt.Name] = &vmcp.BackendTarget{ - WorkloadID: prompt.BackendID, - OriginalCapabilityName: prompt.Name, - } - } else { - target := vmcp.BackendToTarget(backend) - // Store the original prompt name for forwarding to backend - target.OriginalCapabilityName = prompt.Name - routingTable.Prompts[prompt.Name] = target - } - } + // Add resources, resource templates, and prompts to the routing table. + // ResolveConflicts is the enforcement point that makes their identities + // unique (see the capability_conflicts.go header), but Merge is + // independently callable, so as defence in depth each merge helper keeps + // a duplicate out of both the routing table and the advertised list + // (first wins, loudly) rather than silently overwriting the earlier + // routing entry. + resources := mergeResources(ctx, resolved.Resources, registry, routingTable) + templates := mergeResourceTemplates(ctx, resolved.ResourceTemplates, registry, routingTable) + prompts := mergePrompts(ctx, resolved.Prompts, registry, routingTable) // Determine conflict strategy used conflictStrategy := vmcp.ConflictStrategyPrefix // Default @@ -418,20 +462,24 @@ func (a *defaultAggregator) MergeCapabilities( } } - // Create final aggregated view + // Create final aggregated view. The advertised lists are the ones built + // alongside the routing tables above, so advertising and routing always + // agree on which entry owns a duplicated identity. aggregated := &AggregatedCapabilities{ - Tools: tools, - Resources: resolved.Resources, - Prompts: resolved.Prompts, - SupportsLogging: resolved.SupportsLogging, - SupportsSampling: resolved.SupportsSampling, - RoutingTable: routingTable, + Tools: tools, + Resources: resources, + ResourceTemplates: templates, + Prompts: prompts, + SupportsLogging: resolved.SupportsLogging, + SupportsSampling: resolved.SupportsSampling, + RoutingTable: routingTable, Metadata: &AggregationMetadata{ - BackendCount: 0, // Will be set by caller - ToolCount: len(tools), - ResourceCount: len(resolved.Resources), - PromptCount: len(resolved.Prompts), - ConflictStrategy: conflictStrategy, + BackendCount: 0, // Will be set by caller + ToolCount: len(tools), + ResourceCount: len(resources), + ResourceTemplateCount: len(templates), + PromptCount: len(prompts), + ConflictStrategy: conflictStrategy, }, } @@ -515,6 +563,120 @@ func (a *defaultAggregator) AggregateCapabilities( return aggregated, nil } +// mergeResources adds resources to the routing table (keyed by URI) and +// returns the advertised resource list. A URI already present in the routing +// table is dropped from both, with a warning (first wins), so advertising and +// routing always agree on which backend owns a URI. +func mergeResources( + ctx context.Context, + resolved []vmcp.Resource, + registry vmcp.BackendRegistry, + routingTable *vmcp.RoutingTable, +) []vmcp.Resource { + resources := make([]vmcp.Resource, 0, len(resolved)) + for _, resource := range resolved { + if existing, duplicate := routingTable.Resources[resource.URI]; duplicate { + slog.Warn("duplicate resource URI in resolved capabilities, keeping first", + "uri", resource.URI, "kept_backend", existing.WorkloadID, "dropped_backend", resource.BackendID) + continue + } + resources = append(resources, resource) + backend := registry.Get(ctx, resource.BackendID) + if backend == nil { + slog.Warn("backend not found in registry for resource, creating minimal target", + "backend", resource.BackendID, "resource", resource.URI) + routingTable.Resources[resource.URI] = &vmcp.BackendTarget{ + WorkloadID: resource.BackendID, + OriginalCapabilityName: resource.URI, + } + } else { + target := vmcp.BackendToTarget(backend) + // Store the original resource URI for forwarding to backend + target.OriginalCapabilityName = resource.URI + routingTable.Resources[resource.URI] = target + } + } + return resources +} + +// mergeResourceTemplates adds resource templates to the routing table (keyed +// by URI-template string) and returns the advertised template list. +// Pass-through, mirroring resources: no URI-template rewriting, and the same +// first-wins guard against duplicate template strings. +// +// OriginalCapabilityName is intentionally left empty. A resources/read routed +// via a template carries the client's CONCRETE, already-expanded URI (e.g. +// file:///logs/2025-01-01.txt); the backend performs its own template +// expansion, so that concrete URI must reach it verbatim. Setting +// OriginalCapabilityName to the template string would make +// GetBackendCapabilityName replace the concrete URI with the unexpanded +// template, and the backend would return unsubstituted content. vMCP does not +// rename templates, so no name translation is needed here. +func mergeResourceTemplates( + ctx context.Context, + resolved []vmcp.ResourceTemplate, + registry vmcp.BackendRegistry, + routingTable *vmcp.RoutingTable, +) []vmcp.ResourceTemplate { + templates := make([]vmcp.ResourceTemplate, 0, len(resolved)) + for _, template := range resolved { + if existing, duplicate := routingTable.ResourceTemplates[template.URITemplate]; duplicate { + slog.Warn("duplicate resource template in resolved capabilities, keeping first", + "resource_template", template.URITemplate, + "kept_backend", existing.WorkloadID, "dropped_backend", template.BackendID) + continue + } + templates = append(templates, template) + backend := registry.Get(ctx, template.BackendID) + if backend == nil { + slog.Warn("backend not found in registry for resource template, creating minimal target", + "backend", template.BackendID, "resource_template", template.URITemplate) + routingTable.ResourceTemplates[template.URITemplate] = &vmcp.BackendTarget{ + WorkloadID: template.BackendID, + } + } else { + routingTable.ResourceTemplates[template.URITemplate] = vmcp.BackendToTarget(backend) + } + } + return templates +} + +// mergePrompts adds prompts to the routing table, keyed by resolved +// (advertised) name, and returns the advertised prompt list, with the same +// first-wins guard against duplicate resolved names. +func mergePrompts( + ctx context.Context, + resolved []ResolvedPrompt, + registry vmcp.BackendRegistry, + routingTable *vmcp.RoutingTable, +) []vmcp.Prompt { + prompts := make([]vmcp.Prompt, 0, len(resolved)) + for _, prompt := range resolved { + if existing, duplicate := routingTable.Prompts[prompt.Name]; duplicate { + slog.Warn("duplicate prompt name in resolved capabilities, keeping first", + "prompt", prompt.Name, "kept_backend", existing.WorkloadID, "dropped_backend", prompt.BackendID) + continue + } + prompts = append(prompts, prompt.Prompt) + backend := registry.Get(ctx, prompt.BackendID) + if backend == nil { + slog.Warn("backend not found in registry for prompt, creating minimal target", + "backend", prompt.BackendID, "prompt", prompt.Name) + routingTable.Prompts[prompt.Name] = &vmcp.BackendTarget{ + WorkloadID: prompt.BackendID, + OriginalCapabilityName: prompt.OriginalName, + } + } else { + target := vmcp.BackendToTarget(backend) + // Store the backend's own prompt name so prompts/get on a renamed + // prompt forwards the name the backend actually knows. + target.OriginalCapabilityName = prompt.OriginalName + routingTable.Prompts[prompt.Name] = target + } + } + return prompts +} + // actualBackendCapabilityName returns the real capability name the backend uses, // reversing any per-backend override rename that processBackendTools may have applied. // @@ -541,10 +703,13 @@ func actualBackendCapabilityName(toolConfigMap map[string]*config.WorkloadToolCo // shouldAdvertiseTool returns true if a tool from the given backend should be // advertised to MCP clients (included in tools/list response). // -// ExcludeAll, Filter, and per-workload settings control advertising, not routing: -// - Tools excluded via ExcludeAll are NOT advertised to MCP clients -// - Tools not matching Filter are NOT advertised to MCP clients -// - BUT they ARE available in the routing table for composite tools to use +// ExcludeAll, Filter, DefaultToolVisibility, and per-workload settings control +// advertising, not routing: +// - Tools excluded via ExcludeAll are NOT advertised to MCP clients +// - Tools not matching Filter are NOT advertised to MCP clients +// - Under DefaultToolVisibility "deny", tools from a backend with no per-workload +// config are NOT advertised to MCP clients +// - BUT they ARE available in the routing table for composite tools to use // // This enables the use case where you want to hide raw backend tools from // direct client access while still allowing curated composite workflows to use them. @@ -561,8 +726,11 @@ func (a *defaultAggregator) shouldAdvertiseTool(backendID, originalToolName stri // Check per-workload settings wlConfig, exists := a.toolConfigMap[backendID] if !exists { - // No config for this backend, advertise the tool - return true + // No config for this backend. Under the default ("allow"), advertise the + // tool; under "deny", withhold it so only backends named in the config + // contribute tools. A backend WITH a config is opted in either way — its + // own ExcludeAll/Filter decide below. + return !a.denyUnlisted } // Check per-workload ExcludeAll setting diff --git a/pkg/vmcp/aggregator/default_aggregator_test.go b/pkg/vmcp/aggregator/default_aggregator_test.go index a25fd0343c..6f6e3d242e 100644 --- a/pkg/vmcp/aggregator/default_aggregator_test.go +++ b/pkg/vmcp/aggregator/default_aggregator_test.go @@ -218,19 +218,22 @@ func TestDefaultAggregator_ResolveConflicts(t *testing.T) { } agg := NewDefaultAggregator(nil, nil, nil, nil) - resolved, err := agg.ResolveConflicts(context.Background(), capabilities) - require.NoError(t, err) - assert.NotNil(t, resolved) - // In Phase 1, we just collect tools - conflict is detected but first one wins - assert.Contains(t, resolved.Tools, "tool1") - assert.Contains(t, resolved.Tools, "tool2") - assert.Contains(t, resolved.Tools, "shared_tool") - // Shared tool should have one backend (whichever was encountered first in map iteration) - // Map iteration order is non-deterministic, so accept either backend - sharedToolBackend := resolved.Tools["shared_tool"].BackendID - assert.True(t, sharedToolBackend == "backend1" || sharedToolBackend == "backend2", - "shared_tool should belong to either backend1 or backend2, got: %s", sharedToolBackend) + // Repeat: map iteration order is re-randomized per call, and the + // fallback winner must not depend on it. + for range 5 { + resolved, err := agg.ResolveConflicts(context.Background(), capabilities) + + require.NoError(t, err) + assert.NotNil(t, resolved) + // The no-resolver fallback keeps the first tool per name in sorted + // backend order, so the conflict winner is deterministic. + assert.Contains(t, resolved.Tools, "tool1") + assert.Contains(t, resolved.Tools, "tool2") + assert.Contains(t, resolved.Tools, "shared_tool") + assert.Equal(t, "backend1", resolved.Tools["shared_tool"].BackendID, + "fallback must keep the first backend in sorted-ID order") + } }) t.Run("no conflicts", func(t *testing.T) { @@ -283,8 +286,8 @@ func TestDefaultAggregator_MergeCapabilities(t *testing.T) { Resources: []vmcp.Resource{ {URI: "test://resource1", BackendID: "backend1"}, }, - Prompts: []vmcp.Prompt{ - {Name: "prompt1", BackendID: "backend1"}, + Prompts: []ResolvedPrompt{ + {Prompt: vmcp.Prompt{Name: "prompt1", BackendID: "backend1"}, OriginalName: "prompt1"}, }, SupportsLogging: true, SupportsSampling: false, @@ -347,6 +350,79 @@ func TestDefaultAggregator_MergeCapabilities(t *testing.T) { assert.Equal(t, 1, aggregated.Metadata.ResourceCount) assert.Equal(t, 1, aggregated.Metadata.PromptCount) }) + + t.Run("merge threads resource templates through and populates the routing table", func(t *testing.T) { + t.Parallel() + resolved := &ResolvedCapabilities{ + Tools: map[string]*ResolvedTool{}, + ResourceTemplates: []vmcp.ResourceTemplate{ + {URITemplate: "file:///logs/{date}.txt", Name: "Daily log", MimeType: "text/plain", BackendID: "backend1"}, + }, + } + + backends := []vmcp.Backend{ + { + ID: "backend1", + Name: "Backend 1", + BaseURL: "http://backend1:8080", + TransportType: "streamable-http", + HealthStatus: vmcp.BackendHealthy, + }, + } + registry := vmcp.NewImmutableRegistry(backends) + + agg := NewDefaultAggregator(nil, nil, nil, nil) + aggregated, err := agg.MergeCapabilities(context.Background(), resolved, registry) + require.NoError(t, err) + + // Pass-through: the aggregated view carries the templates unchanged. + require.Len(t, aggregated.ResourceTemplates, 1) + assert.Equal(t, "file:///logs/{date}.txt", aggregated.ResourceTemplates[0].URITemplate) + assert.Equal(t, "backend1", aggregated.ResourceTemplates[0].BackendID) + assert.Equal(t, 1, aggregated.Metadata.ResourceTemplateCount) + + // Routing table is keyed by URI template and carries full backend info. + require.Contains(t, aggregated.RoutingTable.ResourceTemplates, "file:///logs/{date}.txt") + target := aggregated.RoutingTable.ResourceTemplates["file:///logs/{date}.txt"] + require.NotNil(t, target) + assert.Equal(t, "backend1", target.WorkloadID) + assert.Equal(t, "http://backend1:8080", target.BaseURL) + // OriginalCapabilityName must be EMPTY so a resources/read routed via this + // template forwards the client's CONCRETE, already-expanded URI to the + // backend verbatim (the backend does its own expansion). If it were set to + // the template, GetBackendCapabilityName would replace the concrete URI with + // the unexpanded template and the backend would return unsubstituted content + // (conformance: "Parameter substitution not reflected in content"). + assert.Empty(t, target.OriginalCapabilityName) + assert.Equal(t, "file:///logs/2025-01-01.txt", + target.GetBackendCapabilityName("file:///logs/2025-01-01.txt"), + "a concrete expanded URI must pass through to the backend unchanged") + }) + + t.Run("registry miss still translates a renamed prompt", func(t *testing.T) { + t.Parallel() + // The minimal-target fallback (backend absent from the registry) must + // carry the backend's OWN prompt name, not the advertised one — + // otherwise prompts/get on the renamed prompt forwards a name the + // backend does not know. + resolved := &ResolvedCapabilities{ + Tools: map[string]*ResolvedTool{}, + Prompts: []ResolvedPrompt{ + {Prompt: vmcp.Prompt{Name: "b1_review", BackendID: "b1"}, OriginalName: "review"}, + }, + } + registry := vmcp.NewImmutableRegistry(nil) + + agg := NewDefaultAggregator(nil, nil, nil, nil) + aggregated, err := agg.MergeCapabilities(context.Background(), resolved, registry) + require.NoError(t, err) + + target := aggregated.RoutingTable.Prompts["b1_review"] + require.NotNil(t, target) + assert.Equal(t, "b1", target.WorkloadID) + assert.Equal(t, "review", target.GetBackendCapabilityName("b1_review"), + "prompts/get on the advertised name must forward the backend's own name") + }) } func TestDefaultAggregator_MergeCapabilities_DeterministicToolOrder(t *testing.T) { @@ -880,3 +956,88 @@ func TestDefaultAggregator_FilterPreservesRoutingTableForCompositeTools(t *testi assert.Len(t, result.RoutingTable.Tools, 2) }) } + +// TestDefaultAggregator_DefaultToolVisibilityDenyMixedBackends covers the motivating +// scenario for defaultToolVisibility (issue #6073): a group holding both a listed and +// an unlisted backend. Under "deny" the listed backend's tools are advertised and +// the unlisted backend contributes nothing — so adding a workload to the group no +// longer exposes it by default. Both backends' tools stay routable, keeping +// composite tools working over hidden ones. +func TestDefaultAggregator_DefaultToolVisibilityDenyMixedBackends(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + backends, capsByID := twoBackendCaps() + mockClient := mocks.NewMockBackendClient(ctrl) + expectListCapabilities(mockClient, capsByID) + + // backend1 is listed (opted in); backend2 is not mentioned at all. + aggCfg := &config.AggregationConfig{ + DefaultToolVisibility: config.DefaultToolVisibilityDeny, + Tools: []*config.WorkloadToolConfig{{Workload: "backend1"}}, + } + + agg := NewDefaultAggregator(mockClient, NewPrefixConflictResolver("{workload}_"), aggCfg, nil) + result, err := agg.AggregateCapabilities(context.Background(), backends) + require.NoError(t, err) + require.NotNil(t, result) + + advertised := make([]string, 0, len(result.Tools)) + for _, tool := range result.Tools { + advertised = append(advertised, tool.Name) + } + assert.ElementsMatch(t, []string{"backend1_fetch", "backend1_tool_a"}, advertised, + "only the listed backend's tools may be advertised under deny") + for _, tool := range result.Tools { + assert.Equalf(t, "backend1", tool.BackendID, + "no unlisted backend's tool may be advertised, but %q was", tool.Name) + } + + // Advertising-only: the unlisted backend stays fully routable so composite + // tools can still reach it. + require.NotNil(t, result.RoutingTable) + assert.Len(t, result.RoutingTable.Tools, 4, + "deny must not remove tools from the routing table") + assert.Contains(t, result.RoutingTable.Tools, "backend2_tool_b", + "an unlisted backend's tools remain routable for composite tools") +} + +// TestDefaultAggregator_WarnsOnUnmatchedToolConfig pins the diagnostic for a typo +// in tools[].workload. Under deny a typo means the real backend has no entry and +// so contributes nothing, and the only other symptom is a short tools/list — the +// detection is what makes that debuggable. +func TestDefaultAggregator_WarnsOnUnmatchedToolConfig(t *testing.T) { + t.Parallel() + + backends, _ := twoBackendCaps() + + t.Run("reports an entry matching no backend", func(t *testing.T) { + t.Parallel() + aggCfg := &config.AggregationConfig{ + DefaultToolVisibility: config.DefaultToolVisibilityDeny, + Tools: []*config.WorkloadToolConfig{ + {Workload: "backend1"}, + {Workload: "backend1typo"}, + }, + } + agg := NewDefaultAggregator(nil, NewPrefixConflictResolver("{workload}_"), aggCfg, nil) + + da, ok := agg.(*defaultAggregator) + require.True(t, ok) + assert.Equal(t, []string{"backend1typo"}, da.unmatchedToolConfigWorkloads(backends), + "only the entry naming no backend in the group is reported") + }) + + t.Run("reports nothing when every entry matches", func(t *testing.T) { + t.Parallel() + aggCfg := &config.AggregationConfig{ + Tools: []*config.WorkloadToolConfig{{Workload: "backend1"}, {Workload: "backend2"}}, + } + agg := NewDefaultAggregator(nil, NewPrefixConflictResolver("{workload}_"), aggCfg, nil) + + da, ok := agg.(*defaultAggregator) + require.True(t, ok) + assert.Empty(t, da.unmatchedToolConfigWorkloads(backends)) + }) +} diff --git a/pkg/vmcp/aggregator/mocks/mock_interfaces.go b/pkg/vmcp/aggregator/mocks/mock_interfaces.go index c90fe21f84..69067f8f0a 100644 --- a/pkg/vmcp/aggregator/mocks/mock_interfaces.go +++ b/pkg/vmcp/aggregator/mocks/mock_interfaces.go @@ -272,3 +272,39 @@ func (mr *MockToolOverrideMockRecorder) ApplyOverrides(ctx, tools any) *gomock.C mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ApplyOverrides", reflect.TypeOf((*MockToolOverride)(nil).ApplyOverrides), ctx, tools) } + +// MockCacheInvalidator is a mock of CacheInvalidator interface. +type MockCacheInvalidator struct { + ctrl *gomock.Controller + recorder *MockCacheInvalidatorMockRecorder + isgomock struct{} +} + +// MockCacheInvalidatorMockRecorder is the mock recorder for MockCacheInvalidator. +type MockCacheInvalidatorMockRecorder struct { + mock *MockCacheInvalidator +} + +// NewMockCacheInvalidator creates a new mock instance. +func NewMockCacheInvalidator(ctrl *gomock.Controller) *MockCacheInvalidator { + mock := &MockCacheInvalidator{ctrl: ctrl} + mock.recorder = &MockCacheInvalidatorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockCacheInvalidator) EXPECT() *MockCacheInvalidatorMockRecorder { + return m.recorder +} + +// InvalidateAll mocks base method. +func (m *MockCacheInvalidator) InvalidateAll() { + m.ctrl.T.Helper() + m.ctrl.Call(m, "InvalidateAll") +} + +// InvalidateAll indicates an expected call of InvalidateAll. +func (mr *MockCacheInvalidatorMockRecorder) InvalidateAll() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InvalidateAll", reflect.TypeOf((*MockCacheInvalidator)(nil).InvalidateAll)) +} diff --git a/pkg/vmcp/aggregator/prefix_resolver.go b/pkg/vmcp/aggregator/prefix_resolver.go index 9941bf1edb..e5a655c8e4 100644 --- a/pkg/vmcp/aggregator/prefix_resolver.go +++ b/pkg/vmcp/aggregator/prefix_resolver.go @@ -6,7 +6,6 @@ package aggregator import ( "context" "log/slog" - "strings" "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -26,7 +25,7 @@ type PrefixConflictResolver struct { // NewPrefixConflictResolver creates a new prefix-based conflict resolver. func NewPrefixConflictResolver(prefixFormat string) *PrefixConflictResolver { if prefixFormat == "" { - prefixFormat = "{workload}_" // Default format + prefixFormat = defaultPrefixFormat } return &PrefixConflictResolver{ PrefixFormat: prefixFormat, @@ -79,10 +78,5 @@ func (r *PrefixConflictResolver) ResolveToolConflicts( // applyPrefix applies the configured prefix format to a tool name. func (r *PrefixConflictResolver) applyPrefix(backendID, toolName string) string { - prefix := r.PrefixFormat - - // Replace {workload} placeholder with actual backend ID - prefix = strings.ReplaceAll(prefix, "{workload}", backendID) - - return prefix + toolName + return applyPrefixFormat(r.PrefixFormat, backendID, toolName) } diff --git a/pkg/vmcp/aggregator/priority_resolver.go b/pkg/vmcp/aggregator/priority_resolver.go index cb4d9c1308..afa0e88d3c 100644 --- a/pkg/vmcp/aggregator/priority_resolver.go +++ b/pkg/vmcp/aggregator/priority_resolver.go @@ -46,7 +46,7 @@ func NewPriorityConflictResolver(priorityOrder []string) (*PriorityConflictResol return &PriorityConflictResolver{ PriorityOrder: priorityOrder, priorityMap: priorityMap, - prefixResolver: NewPrefixConflictResolver("{workload}_"), // Fallback for unmapped backends + prefixResolver: NewPrefixConflictResolver(defaultPrefixFormat), // Fallback for unmapped backends }, nil } diff --git a/pkg/vmcp/auth/factory/authz_not_wired_test.go b/pkg/vmcp/auth/factory/authz_not_wired_test.go index 503d667ce4..08e0e6c756 100644 --- a/pkg/vmcp/auth/factory/authz_not_wired_test.go +++ b/pkg/vmcp/auth/factory/authz_not_wired_test.go @@ -120,8 +120,7 @@ func TestNewIncomingAuthMiddleware_AuthzEnforced(t *testing.T) { // Compose: auth wraps authz wraps handler wrapped := authMw(authzMw(testHandler)) - // Anonymous user has no role. This policy dereferences principal.claim_role - // directly, so Cedar raises an evaluation error instead of a clean deny. + // Anonymous user has no role, so should be denied mcpRequest := map[string]any{ "jsonrpc": "2.0", "method": "tools/call", @@ -140,12 +139,11 @@ func TestNewIncomingAuthMiddleware_AuthzEnforced(t *testing.T) { wrapped.ServeHTTP(recorder, req) - // EXPECTED: The request must not reach the handler, and the server should - // surface an explicit authz failure instead of silently degrading. + // EXPECTED: Anonymous user should be denied (no admin role) assert.False(t, handlerCalled, "handler should NOT be called - anonymous user lacks admin role") - assert.Equal(t, http.StatusInternalServerError, recorder.Code, - "response should be 500 when policy evaluation fails for missing principal attributes") + assert.Equal(t, http.StatusForbidden, recorder.Code, + "response should be 403 Forbidden for non-admin user") }) } diff --git a/pkg/vmcp/auth/factory/incoming.go b/pkg/vmcp/auth/factory/incoming.go index 400b029b04..91874cf5a7 100644 --- a/pkg/vmcp/auth/factory/incoming.go +++ b/pkg/vmcp/auth/factory/incoming.go @@ -237,6 +237,7 @@ func newOIDCAuthMiddleware( AllowPrivateIP: oidcCfg.ProtectedResourceAllowPrivateIP || oidcCfg.JwksAllowPrivateIP, InsecureAllowHTTP: oidcCfg.InsecureAllowHTTP, Scopes: oidcCfg.Scopes, + CACertPath: oidcCfg.CABundlePath, } // Wire optional dependencies from the embedded auth server so the JWT diff --git a/pkg/vmcp/auth/factory/incoming_cabundle_test.go b/pkg/vmcp/auth/factory/incoming_cabundle_test.go new file mode 100644 index 0000000000..43673f5011 --- /dev/null +++ b/pkg/vmcp/auth/factory/incoming_cabundle_test.go @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package factory + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "math/big" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp/config" +) + +// TestNewOIDCAuthMiddleware_CABundlePath verifies that the CABundlePath field on +// OIDCConfig is wired into the underlying auth.TokenValidatorConfig. A bad path +// is expected to surface as an error at middleware construction time, and a +// valid PEM bundle is expected to succeed — proving the field is actually +// consumed (not silently dropped) by the downstream HTTP client builder. +func TestNewOIDCAuthMiddleware_CABundlePath(t *testing.T) { + t.Parallel() + + // Stand up a minimal OIDC discovery server so the validator has an issuer + // to talk to during construction. The CA bundle affects the TLS trust + // store but not issuer reachability here. + server, _ := newTestOIDCServer(t) + t.Cleanup(server.Close) + + validCAPath := writeTestCAPEM(t) + + tests := []struct { + name string + caBundlePath string + wantErr bool + errContains string + }{ + { + name: "empty caBundlePath succeeds (backward compatible)", + caBundlePath: "", + wantErr: false, + }, + { + name: "valid PEM caBundlePath is loaded successfully", + caBundlePath: validCAPath, + wantErr: false, + }, + { + name: "non-existent caBundlePath surfaces as an error", + caBundlePath: "/nonexistent/ca/bundle.pem", + wantErr: true, + errContains: "CA certificate", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + oidcCfg := &config.OIDCConfig{ + Issuer: server.URL, + ClientID: "test-client", + Audience: "test-audience", + InsecureAllowHTTP: true, + JwksAllowPrivateIP: true, + CABundlePath: tt.caBundlePath, + } + + authMw, _, err := newOIDCAuthMiddleware(t.Context(), oidcCfg, nil, nil, "") + + if tt.wantErr { + require.Error(t, err) + if tt.errContains != "" { + assert.Contains(t, err.Error(), tt.errContains) + } + return + } + + require.NoError(t, err) + require.NotNil(t, authMw) + }) + } +} + +// writeTestCAPEM generates a minimal self-signed certificate and writes it as +// PEM into t.TempDir(). Returns the absolute path to the PEM file. The cert +// content does not need to match any issuer — the test only exercises that +// the bundle loads successfully into the HTTP client's trust store. +func writeTestCAPEM(t *testing.T) string { + t.Helper() + + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + + template := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "toolhive-test-ca"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + KeyUsage: x509.KeyUsageCertSign, + BasicConstraintsValid: true, + IsCA: true, + } + + der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key) + require.NoError(t, err) + + pemBytes := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}) + require.NotEmpty(t, pemBytes) + + dir := t.TempDir() + path := filepath.Join(dir, "ca.pem") + require.NoError(t, os.WriteFile(path, pemBytes, 0o600)) + return path +} diff --git a/pkg/vmcp/cli/init.go b/pkg/vmcp/cli/init.go index 8677900f64..c7bee6417b 100644 --- a/pkg/vmcp/cli/init.go +++ b/pkg/vmcp/cli/init.go @@ -73,6 +73,9 @@ outgoingAuth: # aggregation: controls how tools from multiple backends are combined. # conflictResolution: prefix prepends the backend name to each tool name. +# Prompt names are prefixed with the backend name by default; under the +# priority strategy, backends listed in priorityOrder keep their own +# prompt names. aggregation: conflictResolution: prefix conflictResolutionConfig: diff --git a/pkg/vmcp/cli/serve.go b/pkg/vmcp/cli/serve.go index a0f6bd99ca..84b49292c4 100644 --- a/pkg/vmcp/cli/serve.go +++ b/pkg/vmcp/cli/serve.go @@ -351,8 +351,13 @@ func Serve(ctx context.Context, cfg ServeConfig) error { // The factory never aggregates — the core is the single source of capability // aggregation (agg feeds it via Config.Aggregator below). The backend init // timeout follows the operational config (non-positive keeps the default). - sessionFactory := vmcpsession.NewSessionFactory(outgoingRegistry, - vmcpsession.WithBackendInitTimeout(backendInitTimeoutFromConfig(vmcpCfg))) + sessionFactoryOpts := []vmcpsession.MultiSessionFactoryOption{ + vmcpsession.WithBackendInitTimeout(backendInitTimeoutFromConfig(vmcpCfg)), + } + if revisions, ok := backendClient.(vmcp.RevisionReporter); ok { + sessionFactoryOpts = append(sessionFactoryOpts, vmcpsession.WithRevisionLookup(revisions.CachedRevision)) + } + sessionFactory := vmcpsession.NewSessionFactory(outgoingRegistry, sessionFactoryOpts...) // When the optimizer is enabled, its meta-tools are pass-through tools. // Authz uses this for optimizer-aware authorization/filtering. diff --git a/pkg/vmcp/client/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go new file mode 100644 index 0000000000..dba3dad3e7 --- /dev/null +++ b/pkg/vmcp/client/auth_error_mapping_regression_test.go @@ -0,0 +1,393 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// jsonRPCResponse is a generic JSON-RPC 2.0 response envelope. +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result,omitempty"` + Error *jsonRPCError `json:"error,omitempty"` +} + +type jsonRPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// jsonRPCRequest is a generic JSON-RPC 2.0 request envelope (for method routing). +type jsonRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Method string `json:"method"` +} + +// TestRegression_401_MapsToErrAuthenticationFailed verifies that a backend +// returning HTTP 401 on initialize is classified as ErrAuthenticationFailed. +func TestRegression_401_MapsToErrAuthenticationFailed(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Unauthorized"}, + }) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + h.clientFactory = func(ctx context.Context, target *vmcp.BackendTarget, _ bool) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-401-backend", + WorkloadName: "Test 401 Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + // Pre-seed Legacy so ListCapabilities skips the Modern discover probe (which + // needs a real registry) and exercises the Legacy error-mapping path here. + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrAuthenticationFailed), + "expected ErrAuthenticationFailed, got: %v", err) +} + +// TestRegression_403OnInitialize_LegacySSEFallback verifies that a backend +// returning HTTP 403 on initialize is classified as ErrBackendUnavailable. +// +// NOTE: a 4xx (except 401) on the initialize POST surfaces as +// transport.ErrLegacySSEServer (see wrapBackendError), which wrapBackendError maps +// to ErrBackendUnavailable while preserving the sentinel in the chain (see +// TestRegression_403OnInitialize_PreservesSentinel). Because 403 is ambiguous +// (auth rejection vs Modern-only backend), it drives a re-probe that returns the +// same revision here and surfaces this error unchanged. +func TestRegression_403OnInitialize_LegacySSEFallback(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Forbidden"}, + }) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + h.clientFactory = func(ctx context.Context, target *vmcp.BackendTarget, _ bool) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-403-backend", + WorkloadName: "Test 403 Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrBackendUnavailable), + "expected ErrBackendUnavailable, got: %v", err) + assert.Contains(t, err.Error(), "403", + "error message should reference 403 status, got: %v", err) +} + +// TestRegression_403OnInitialize_PreservesSentinel verifies that a 403 on +// initialize still classifies as ErrBackendUnavailable AND that the origin +// transport.ErrLegacySSEServer is preserved in the error chain (wrapBackendError +// now uses a second %w). The preserved sentinel is what lets dispatch detect a +// revision mismatch — a 403 is ambiguous (auth rejection vs Modern-only backend), +// so it drives a re-probe that returns the same revision here (Legacy) and +// surfaces the original error unchanged. +func TestRegression_403OnInitialize_PreservesSentinel(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Forbidden"}, + }) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + h.clientFactory = func(ctx context.Context, target *vmcp.BackendTarget, _ bool) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-403-sentinel-backend", + WorkloadName: "Test 403 Sentinel Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + + assert.ErrorIs(t, err, vmcp.ErrBackendUnavailable, + "403 on initialize must still classify as backend unavailable") + // wrapBackendError now preserves the origin via a second %w, so the sentinel + // is in the chain — this is what enables revision-mismatch detection. + assert.ErrorIs(t, err, mcptransport.ErrLegacySSEServer, + "transport.ErrLegacySSEServer must be preserved in the error chain (wrapBackendError uses %w)") +} + +// TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure verifies +// that an MCP tool error result (IsError=true on a 200 HTTP response) whose +// message contains "401 unauthorized" is NOT classified as an auth failure. +func TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, err := io.ReadAll(r.Body) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + var req jsonRPCRequest + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + switch req.Method { + case "initialize": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "test-backend", "version": "1.0.0"} + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/call": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "content": [{"type": "text", "text": "tool error: 401 unauthorized - permission denied"}], + "isError": true + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/list": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "tools": [{"name": "test-tool", "description": "A test tool", "inputSchema": {"type": "object"}}] + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + default: + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{}`), + }) + } + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + h.clientFactory = func(ctx context.Context, target *vmcp.BackendTarget, _ bool) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-tool-error-backend", + WorkloadName: "Test Tool Error Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + result, err := h.CallTool(ctx, target, "test-tool", map[string]any{"arg": "val"}, nil, nil) + + if err != nil { + if errors.Is(err, vmcp.ErrAuthenticationFailed) { + t.Fatalf("UNEXPECTED: CallTool returned ErrAuthenticationFailed for a 200 response with IsError=true containing '401 unauthorized'") + } + t.Fatalf("unexpected transport error from CallTool: %v", err) + } + + assert.NotNil(t, result, "expected non-nil result for IsError=true") + assert.True(t, result.IsError, "expected IsError=true") + assert.NotEmpty(t, result.Content, "expected non-empty content") +} + +// TestRegression_WrapBackendError_ClassificationMatrix pins the remaining +// wrapBackendError classification arms that the httptest-driven cases above do +// not exercise. wrapBackendError maps several distinct failure shapes onto vmcp +// sentinels, and health monitoring (#4935, #5223) branches on those sentinels; +// a regression that reshuffles the arms would silently mis-route recovery. Each +// case reconstructs the error shape the transport layer produces in the wild and +// asserts the sentinel via errors.Is (the chain-aware check wrapBackendError's +// %w wrapping is designed to satisfy). +func TestRegression_WrapBackendError_ClassificationMatrix(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + wantSentinel error + }{ + { + // A backend returning a 5xx surfaces as a generic HTTP status error + // (no transport sentinel). It must classify as ErrBackendUnavailable + // so the health monitor can recover once the backend is healthy again. + name: "5xx HTTP status maps to ErrBackendUnavailable", + err: fmt.Errorf("request failed with status code 503 service unavailable"), + wantSentinel: vmcp.ErrBackendUnavailable, + }, + { + // A network timeout arrives as a net.Error whose Timeout() reports true; + // wrapBackendError's errors.As(&netErr)+Timeout() arm must map it to + // ErrTimeout ahead of any string-based fallback. + name: "net.Error timeout maps to ErrTimeout", + err: &net.DNSError{Err: "i/o timeout", Name: "backend.example", IsTimeout: true}, + wantSentinel: vmcp.ErrTimeout, + }, + { + // An unexpectedly closed connection surfaces as io.EOF; it means the + // backend dropped the stream, so it maps to ErrBackendUnavailable. + name: "io.EOF maps to ErrBackendUnavailable", + err: io.EOF, + wantSentinel: vmcp.ErrBackendUnavailable, + }, + { + // A refused dial (backend down) is a net.OpError whose message contains + // "connection refused"; the Timeout() arm is skipped (not a timeout) and + // the connection-error string arm maps it to ErrBackendUnavailable. + name: "connection refused maps to ErrBackendUnavailable", + err: &net.OpError{ + Op: "dial", + Net: "tcp", + Err: errors.New("connect: connection refused"), + }, + wantSentinel: vmcp.ErrBackendUnavailable, + }, + { + // transport.ErrAuthorizationRequired is the migration-sensitive sentinel + // (mcp-go v0.49.0+) returned for 401 + WWW-Authenticate. It must map to + // ErrAuthenticationFailed so health monitoring engages the auth-aware + // branch (#4935) instead of marking the backend unhealthy (#5223). + name: "transport.ErrAuthorizationRequired maps to ErrAuthenticationFailed", + err: mcptransport.ErrAuthorizationRequired, + wantSentinel: vmcp.ErrAuthenticationFailed, + }, + { + // The production chain wraps the sentinel in *transport.Error via + // *AuthorizationRequiredError; both layers Unwrap to the sentinel so + // errors.Is must still classify it as ErrAuthenticationFailed. + name: "wrapped AuthorizationRequiredError maps to ErrAuthenticationFailed", + err: mcptransport.NewError(&mcptransport.AuthorizationRequiredError{}), + wantSentinel: vmcp.ErrAuthenticationFailed, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := wrapBackendError(tt.err, "test-backend", "initialize") + require.Error(t, err) + assert.True(t, errors.Is(err, tt.wantSentinel), + "expected %v, got: %v", tt.wantSentinel, err) + }) + } +} diff --git a/pkg/vmcp/client/auth_propagation_integration_test.go b/pkg/vmcp/client/auth_propagation_integration_test.go index 7844d1a8ae..1a69d145ac 100644 --- a/pkg/vmcp/client/auth_propagation_integration_test.go +++ b/pkg/vmcp/client/auth_propagation_integration_test.go @@ -11,10 +11,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index dc8f41a0a6..9bb13a1769 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -4,13 +4,14 @@ // Package client provides MCP protocol client implementation for communicating with backend servers. // // This package implements the BackendClient interface defined in the vmcp package, -// using the mark3labs/mcp-go SDK for protocol communication. +// using the stacklok/toolhive-core/mcpcompat SDK for protocol communication. package client import ( "context" "crypto/tls" "crypto/x509" + "encoding/json" "errors" "fmt" "io" @@ -19,17 +20,23 @@ import ( "net" "net/http" "os" + "slices" + "sync" + "sync/atomic" "syscall" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/propagation" + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/secrets" + "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/versions" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -38,6 +45,8 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/headerforward" "github.com/stacklok/toolhive/pkg/vmcp/health" healthcontext "github.com/stacklok/toolhive/pkg/vmcp/health/context" + "github.com/stacklok/toolhive/pkg/vmcp/internal/backendtelemetry" + "github.com/stacklok/toolhive/pkg/vmcp/internal/pagination" ) const ( @@ -57,6 +66,14 @@ const ( // Note: This limit is enforced per HTTP response, not per MCP request. // A tools/list response with 1000 tools would be limited to 100MB total. maxResponseSize = 100 * 1024 * 1024 // 100 MB + + // defaultRefutationTTL is the default lifetime of a modernHintRefuted + // entry. Chosen to be long enough that a persistently hint-lying backend + // (#6154) pays a negligible probe cost (one server/discover per backend + // per 5 minutes, only when calls are flowing), and short enough that a + // backend redeployed genuinely Modern self-heals its reported revision + // within minutes rather than never. + defaultRefutationTTL = 5 * time.Minute ) // Option configures an httpBackendClient. @@ -118,12 +135,30 @@ func WithRequestTimeouts(defaultTimeout time.Duration, perWorkload map[string]ti } } -// httpBackendClient implements vmcp.BackendClient using mark3labs/mcp-go HTTP client. +// WithRefutationTTL overrides how long a refuted Modern-handshake hint +// suppresses the confirming server/discover probe in legacyInit. Zero or +// negative restores the default. Production callers should not set this; it +// exists so tests can shrink the window. +func WithRefutationTTL(d time.Duration) Option { + return func(h *httpBackendClient) { + if d > 0 { + h.refutationTTL = d + } + } +} + +// httpBackendClient implements vmcp.BackendClient using stacklok/toolhive-core/mcpcompat HTTP client. // It supports streamable-HTTP and SSE transports for backend MCP servers. type httpBackendClient struct { - // clientFactory creates MCP clients for backends. + // clientFactory creates MCP clients for backends. The forwarding flag is set + // only for the tools/call path — the sole operation during which a backend + // may issue server->client requests/notifications (elicitation, sampling, + // progress, logging). It gates the continuous-listening SSE stream and the + // forwarding handlers so that aggregation/list/read/prompt/complete calls do + // NOT open a standalone GET stream (which hangs backends that don't support + // one — see the multi-backend optimizer regression). // Abstracted as a function to enable testing with mock clients. - clientFactory func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) + clientFactory func(ctx context.Context, target *vmcp.BackendTarget, forwarding bool) (*client.Client, error) // registry manages authentication strategies for outgoing requests to backend MCP servers. // Must not be nil - use UnauthenticatedStrategy for no authentication. @@ -144,6 +179,59 @@ type httpBackendClient struct { // WithRequestTimeouts; see requestTimeout. defaultRequestTimeout time.Duration perWorkloadRequestTimeouts map[string]time.Duration + + // forwarders holds the server->client forwarding requesters (elicitation, + // sampling, notifications) bound by server.New via BindForwarders. When set, + // the client factory installs handlers on each backend client so a backend's + // mid-call server->client traffic is relayed to the downstream client, and + // enables continuous listening so standalone-SSE server->client messages are + // delivered. Nil (unbound) reproduces the pre-forwarding behavior exactly, so + // direct embedders and unit tests without a bound server are unaffected. + forwarders atomic.Pointer[boundForwarders] + + // revisions caches each backend's resolved MCP revision, keyed by + // target.WorkloadID. An ABSENT key means unprobed — distinct from + // RevisionLegacy (0), a resolved result. Populated by probeRevision on the + // first ListCapabilities for a backend and read on subsequent calls to skip + // the Modern-first discover probe. + // + // NOTE: never evicted, no singleflight/CAS. Concurrent first-probes are + // last-writer-wins-safe: writes are idempotent for a deterministic backend + // (every probe agrees). A transient probe failure caches nothing + // (probeRevision returns uncached), so a blip cannot pin a revision. + // + // Self-healing now works both ways under go-sdk v1.7. A mis-cached Modern + // backend that has actually negotiated down to Legacy self-heals via + // dispatch->reclassify (errModernNegotiatedDown re-probes and flips it). A + // mis-cached Legacy backend that has actually become Modern self-heals + // in-band instead: the SDK's own client transparently negotiates Modern on + // the Legacy dispatch path, and legacyInit reads the genuinely-negotiated + // protocol version off every Initialize result and flips the cache directly + // when it is Modern — no error or reclassify round trip needed (see + // TestListCapabilities_MisCachedLegacy_SelfHealsViaSDKNegotiation in + // reclassify_test.go). + revisions sync.Map // map[string]mcpparser.Revision + + // modernHintRefuted records backends whose Legacy-handshake protocol-version + // hint was contradicted by an authoritative server/discover probe, keyed by + // target.WorkloadID. See legacyInit: a backend can negotiate 2026-07-28 on + // the Legacy initialize while its discover response is not a valid Modern + // envelope, and the hint must not override the probe in that case. Recording + // the refutation stops legacyInit re-probing on every subsequent call for a + // backend already known to lie. + // + // The refutation EXPIRES after refutationTTL: a permanently refuted flag + // would pin a redeemed backend (one that lied once but was later redeployed + // genuinely Modern) to Legacy forever, because legacyInit skips the + // confirming probe while the flag stands. With the TTL, a persistent liar + // pays one confirming probe per window instead of per call, and a redeemed + // backend self-heals within one window. + modernHintRefuted sync.Map // map[string]time.Time (refuted-at) + + // refutationTTL bounds how long a modernHintRefuted entry suppresses the + // confirming probe in legacyInit. Defaults to defaultRefutationTTL; + // configurable via WithRefutationTTL (tests). + refutationTTL time.Duration } // NewHTTPBackendClient creates a new HTTP-based backend client. @@ -166,6 +254,7 @@ func NewHTTPBackendClient(registry vmcpauth.OutgoingAuthRegistry, opts ...Option c := &httpBackendClient{ registry: registry, secretsProvider: secrets.NewEnvironmentProvider(), + refutationTTL: defaultRefutationTTL, } for _, o := range opts { o(c) @@ -359,6 +448,14 @@ func (i *identityPropagatingRoundTripper) RoundTrip(req *http.Request) (*http.Re return i.base.RoundTrip(req) } +// CloseIdleConnections forwards to the wrapped RoundTripper so it reaches the +// concrete *http.Transport at the bottom of the chain. +func (i *identityPropagatingRoundTripper) CloseIdleConnections() { + if c, ok := i.base.(interface{ CloseIdleConnections() }); ok { + c.CloseIdleConnections() + } +} + // tracePropagatingRoundTripper injects W3C Trace Context (traceparent/tracestate) and // Baggage headers into outgoing HTTP requests. This links vMCP client spans with backend // server spans in distributed traces without creating duplicate spans (unlike @@ -391,6 +488,14 @@ func (h *healthProbePropagatingRoundTripper) RoundTrip(req *http.Request) (*http return h.base.RoundTrip(clonedReq) } +// CloseIdleConnections forwards to the wrapped RoundTripper so it reaches the +// concrete *http.Transport at the bottom of the chain. +func (t *tracePropagatingRoundTripper) CloseIdleConnections() { + if c, ok := t.base.(interface{ CloseIdleConnections() }); ok { + c.CloseIdleConnections() + } +} + // authRoundTripper is an http.RoundTripper that adds authentication to backend requests. // The authentication strategy is pre-resolved and validated at client creation time, // eliminating per-request lookups and validation overhead. @@ -416,6 +521,14 @@ func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) return a.base.RoundTrip(reqClone) } +// CloseIdleConnections forwards to the wrapped RoundTripper so it reaches the +// concrete *http.Transport at the bottom of the chain. +func (a *authRoundTripper) CloseIdleConnections() { + if c, ok := a.base.(interface{ CloseIdleConnections() }); ok { + c.CloseIdleConnections() + } +} + // resolveAuthStrategy resolves the authentication strategy for a backend target. // It handles defaulting to "unauthenticated" when no auth config is specified. // This method should be called once at client creation time to enable fail-fast @@ -436,13 +549,108 @@ func (h *httpBackendClient) resolveAuthStrategy(target *vmcp.BackendTarget) (vmc return strategy, nil } -// defaultClientFactory creates mark3labs MCP clients for different transport types. -func (h *httpBackendClient) defaultClientFactory(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { - // Build transport chain (outermost to innermost, request execution order): - // size limit (response body) → trace propagation → identity propagation → authentication → HTTP - // - // Build an isolated per-call transport so each client gets its own connection pool, - // preventing stale keep-alive connections from one backend affecting others. +// defaultClientFactory creates mcpcompat MCP clients for different transport types. +// newStreamableHTTPClient builds a streamable-HTTP backend client. It bounds the +// response body size and — only when forwarding is requested (the tools/call path) +// and forwarders are bound — enables continuous listening plus the elicitation/ +// sampling handlers so a backend's mid-call server->client traffic reaches the +// downstream client. Non-forwarding calls get the plain client (no standalone GET +// stream), which is byte-for-byte the pre-forwarding construction. +func (h *httpBackendClient) newStreamableHTTPClient( + ctx context.Context, target *vmcp.BackendTarget, + baseTransport http.RoundTripper, forwarding bool, fwd *boundForwarders, +) (*client.Client, error) { + // For streamable-HTTP each MCP call is a single bounded HTTP request/response + // pair, so a per-response body size limit is safe. + sizeLimitedTransport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + resp, err := baseTransport.RoundTrip(req) + if err != nil { + return nil, err + } + resp.Body = struct { + io.Reader + io.Closer + }{ + Reader: io.LimitReader(resp.Body, maxResponseSize), + Closer: resp.Body, + } + return resp, nil + }) + requestTimeout := h.requestTimeout(target.WorkloadID) + httpClient := newBackendHTTPClient(sizeLimitedTransport, requestTimeout) + transportOpts := []transport.StreamableHTTPCOption{ + transport.WithHTTPTimeout(requestTimeout), + transport.WithHTTPBasicClient(httpClient), + } + if fwd != nil && forwarding { + // A backend's mid-call elicitation/sampling request is routed by the go-sdk + // onto the standalone SSE stream (the shim server replies with + // application/json), so the backend client must hold that stream open for the + // request to arrive and be answered. Only enabled for the tools/call path: + // opening this GET stream against a backend that does not support one hangs + // the call, so aggregation/list/read/prompt/complete must not enable it. + transportOpts = append(transportOpts, transport.WithContinuousListening()) + return client.NewStreamableHttpClientWithOpts( + target.BaseURL, transportOpts, forwardingClientOptions(ctx, fwd), + ) + } + return client.NewStreamableHttpClient(target.BaseURL, transportOpts...) +} + +// newSSEClient builds an SSE backend client. For SSE the entire session is one +// long-lived HTTP response body, so — unlike the streamable-HTTP client — no +// response-body size limit and no http.Client.Timeout apply: the former would +// silently terminate the stream after maxResponseSize CUMULATIVE bytes, the +// latter would kill the stream. When forwarding is requested (the tools/call +// path) and forwarders are bound, the elicitation/sampling handlers are +// installed so a backend's mid-call server->client traffic reaches the +// downstream client; SSE carries those requests on the already-open event +// stream, so no continuous-listening option is needed. Non-forwarding calls get +// the plain client, byte-for-byte the pre-forwarding construction. +func (*httpBackendClient) newSSEClient( + ctx context.Context, target *vmcp.BackendTarget, + baseTransport http.RoundTripper, forwarding bool, fwd *boundForwarders, +) (*client.Client, error) { + c, err := client.NewSSEMCPClient( + target.BaseURL, + // timeout 0: SSE is one long-lived stream, so no client timeout. + transport.WithHTTPClient(newBackendHTTPClient(baseTransport, 0)), + ) + if err != nil { + return nil, fmt.Errorf("failed to create SSE client: %w", err) + } + if fwd != nil && forwarding { + // NewSSEMCPClient accepts only transport options, so the client options + // are applied post-construction; both run before the caller's Initialize, + // which is what the handlers' capability declaration requires. + for _, opt := range forwardingClientOptions(ctx, fwd) { + opt(c) + } + } + return c, nil +} + +// buildBackendRoundTripper assembles the per-call backend RoundTripper chain +// shared by every backend transport (streamable-HTTP, SSE, and the raw Modern +// shim). Outermost to innermost, in request execution order: +// +// trace propagation → identity propagation → header-forward → authentication → TLS/HTTP +// +// The nesting order is load-bearing: identity MUST wrap auth so the fresh +// per-request identity is on the context before an auth strategy reads it (#5323). +// The transport is isolated per call so each client gets its own connection pool, +// preventing stale keep-alive connections from one backend affecting others. +// +// ctx is the LIVE per-call context: the header-forward and identity layers read +// forwarded headers and the fallback identity/health-check marker off it, so +// callers MUST pass the real request context, never context.Background(). +// +// This returns the CHAIN, not a wrapped *http.Client: streamable-HTTP wraps it in +// a size-limited/30s client, SSE wraps it bare (long-lived), and the Modern shim +// picks its own bound — see the callers. +func (h *httpBackendClient) buildBackendRoundTripper( + ctx context.Context, target *vmcp.BackendTarget, +) (http.RoundTripper, error) { httpTransport, err := newBackendTransport(target.CABundlePath, target.CABundleData, h.dialControl) if err != nil { return nil, fmt.Errorf("failed to create transport for backend %s: %w", target.WorkloadID, err) @@ -539,59 +747,52 @@ func (h *httpBackendClient) defaultClientFactory(ctx context.Context, target *vm propagator: otel.GetTextMapPropagator(), } + return baseTransport, nil +} + +func (h *httpBackendClient) defaultClientFactory( + ctx context.Context, target *vmcp.BackendTarget, forwarding bool, +) (*client.Client, error) { + baseTransport, err := h.buildBackendRoundTripper(ctx, target) + if err != nil { + return nil, err + } + + // Snapshot the bound server->client forwarders (nil when unbound). When set, + // the client is built with elicitation/sampling handlers and continuous + // listening so a backend's mid-call server->client traffic reaches the + // downstream client; when nil, construction is byte-for-byte the pre-forwarding + // path. + fwd := h.forwarders.Load() + var c *client.Client switch target.TransportType { case "streamable-http", "streamable": // "streamable" is a legacy alias for "streamable-http". - // - // For streamable-HTTP each MCP call is a single bounded HTTP - // request/response pair, so a per-response body size limit is safe. - sizeLimitedTransport := roundTripperFunc(func(req *http.Request) (*http.Response, error) { - resp, err := baseTransport.RoundTrip(req) - if err != nil { - return nil, err - } - resp.Body = struct { - io.Reader - io.Closer - }{ - Reader: io.LimitReader(resp.Body, maxResponseSize), - Closer: resp.Body, - } - return resp, nil - }) - requestTimeout := h.requestTimeout(target.WorkloadID) - httpClient := &http.Client{ - Transport: sizeLimitedTransport, - Timeout: requestTimeout, - } - c, err = client.NewStreamableHttpClient( - target.BaseURL, - transport.WithHTTPBasicClient(httpClient), - ) + c, err = h.newStreamableHTTPClient(ctx, target, baseTransport, forwarding, fwd) if err != nil { return nil, fmt.Errorf("failed to create streamable-http client: %w", err) } case "sse": - // For SSE the entire session is one long-lived HTTP response body. - // Applying io.LimitReader would silently terminate the stream after - // maxResponseSize cumulative bytes — not per-event — which is wrong. - // http.Client.Timeout is also omitted: it would kill the stream. - httpClient := &http.Client{Transport: baseTransport} - c, err = client.NewSSEMCPClient( - target.BaseURL, - transport.WithHTTPClient(httpClient), - ) + c, err = h.newSSEClient(ctx, target, baseTransport, forwarding, fwd) if err != nil { - return nil, fmt.Errorf("failed to create SSE client: %w", err) + return nil, err } default: return nil, fmt.Errorf("%w: %s (supported: streamable-http, sse)", vmcp.ErrUnsupportedTransport, target.TransportType) } + // Register the notification forwarder before Initialize (the caller runs it + // after this factory returns) so a backend's mid-call progress/logging + // notifications are relayed to the downstream client. OnNotification is a + // post-construction method, so it applies to both transports. + if fwd != nil && forwarding && fwd.notifier != nil { + c.OnNotification(newNotificationForwarder(ctx, fwd.notifier)) + } + // Start the client connection if err := c.Start(ctx); err != nil { return nil, fmt.Errorf("failed to start client connection: %w", err) @@ -682,13 +883,23 @@ func wrapBackendError(err error, backendID string, operation string) error { return fmt.Errorf("%w: failed to %s for backend %s: %v", vmcp.ErrAuthenticationFailed, operation, backendID, err) } - // ErrLegacySSEServer is returned for any 4xx (except 401) on initialize POST. - // This includes 403 (auth rejection) and 404/405 (endpoint not found/method not allowed). - // We cannot distinguish auth failures from routing errors without the raw status code, - // so we surface a clear message and classify as backend unavailable to allow recovery. + // ErrLegacySSEServer is toolhive-core's sentinel for a 4xx (except 401) on + // initialize POST against the "sse" transport type, where the SDK client + // cannot distinguish an auth rejection from a legacy SSE-only server without + // the raw status code. Under the streamable-http transport this arm is dead: + // mcp-go's streamable-HTTP client surfaces a generic HTTP status error for a + // 403 on initialize (e.g. "request failed with status 403"), not this + // sentinel, so that case falls through to the string-based classification + // below and still maps to ErrBackendUnavailable (see + // TestRegression_403OnInitialize_LegacySSEFallback). This arm only fires for + // the sse transport type; it is kept for backend targets still configured + // with legacy SSE. if errors.Is(err, transport.ErrLegacySSEServer) { const legacyMsg = "server rejected MCP initialize — possible auth rejection or legacy SSE-only server" - return fmt.Errorf("%w: failed to %s for backend %s (%s): %v", + // Second %w preserves ErrLegacySSEServer in the chain so a revision + // mismatch (Legacy initialize against a Modern backend) is detectable via + // errors.Is; renders identically to %v. + return fmt.Errorf("%w: failed to %s for backend %s (%s): %w", vmcp.ErrBackendUnavailable, operation, backendID, legacyMsg, err) } @@ -722,14 +933,20 @@ func wrapBackendError(err error, backendID string, operation string) error { vmcp.ErrBackendUnavailable, operation, backendID, err) } - // Default to backend unavailable for unknown errors - return fmt.Errorf("%w: failed to %s for backend %s: %v", + // Default to backend unavailable for unknown errors. Second %w preserves the + // origin error (e.g. mcp.ErrMethodNotFound from a Legacy initialize against a + // Modern backend, or errWrongEra from a Modern probe against a Legacy backend) + // so a revision mismatch is detectable via errors.Is; renders identically to %v. + return fmt.Errorf("%w: failed to %s for backend %s: %w", vmcp.ErrBackendUnavailable, operation, backendID, err) } -// initializeClient performs MCP protocol initialization handshake and returns server capabilities. -// This allows the caller to determine which optional features the server supports. -func initializeClient(ctx context.Context, c *client.Client) (*mcp.ServerCapabilities, error) { +// initializeClient performs MCP protocol initialization handshake and returns +// server capabilities plus the actually-negotiated protocol version. The +// latter comes straight from the SDK's InitializeResult, so it reflects the +// real negotiation outcome (e.g. a Modern-first Connect() that negotiated +// down) even though ProtocolVersion is requested as mcp.LATEST_PROTOCOL_VERSION. +func initializeClient(ctx context.Context, c *client.Client) (*mcp.ServerCapabilities, string, error) { result, err := c.Initialize(ctx, mcp.InitializeRequest{ Params: mcp.InitializeParams{ ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, @@ -748,128 +965,471 @@ func initializeClient(ctx context.Context, c *client.Client) (*mcp.ServerCapabil }, }) if err != nil { - return nil, err + return nil, "", err } - return &result.Capabilities, nil + return &result.Capabilities, result.ProtocolVersion, nil } // queryTools queries tools from a backend if the server advertises tool support. +// It follows MCP pagination cursors so backends that paginate (mcpcompat +// paginates at DefaultPageSize=1000) contribute their complete tool set rather +// than only the first page. func queryTools(ctx context.Context, c *client.Client, supported bool, backendID string) (*mcp.ListToolsResult, error) { if supported { - result, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + tools, err := pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Tool, mcp.Cursor, error) { + req := mcp.ListToolsRequest{} + req.Params.Cursor = cursor + result, err := c.ListTools(ctx, req) + if err != nil { + return nil, "", err + } + return result.Tools, result.NextCursor, nil + }) if err != nil { return nil, fmt.Errorf("failed to list tools from backend %s: %w", backendID, err) } - return result, nil + return &mcp.ListToolsResult{Tools: tools}, nil } slog.Debug("backend does not advertise tools capability, skipping tools query", "backend", backendID) return &mcp.ListToolsResult{Tools: []mcp.Tool{}}, nil } -// queryResources queries resources from a backend if the server advertises resource support. +// queryResources queries resources from a backend if the server advertises +// resource support. It follows MCP pagination cursors (see queryTools). func queryResources(ctx context.Context, c *client.Client, supported bool, backendID string) (*mcp.ListResourcesResult, error) { if supported { - result, err := c.ListResources(ctx, mcp.ListResourcesRequest{}) + resources, err := pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Resource, mcp.Cursor, error) { + req := mcp.ListResourcesRequest{} + req.Params.Cursor = cursor + result, err := c.ListResources(ctx, req) + if err != nil { + return nil, "", err + } + return result.Resources, result.NextCursor, nil + }) if err != nil { return nil, fmt.Errorf("failed to list resources from backend %s: %w", backendID, err) } - return result, nil + return &mcp.ListResourcesResult{Resources: resources}, nil } slog.Debug("backend does not advertise resources capability, skipping resources query", "backend", backendID) return &mcp.ListResourcesResult{Resources: []mcp.Resource{}}, nil } -// queryPrompts queries prompts from a backend if the server advertises prompt support. +// queryResourceTemplates queries resource templates from a backend if the server +// advertises resource support. Resource templates are gated on the same +// serverCaps.Resources advertisement as plain resources (there is no separate +// capability flag for templates). It follows MCP pagination cursors (see queryTools). +func queryResourceTemplates( + ctx context.Context, c *client.Client, supported bool, backendID string, +) (*mcp.ListResourceTemplatesResult, error) { + if supported { + templates, err := pagination.ListAll( + ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.ResourceTemplate, mcp.Cursor, error) { + req := mcp.ListResourceTemplatesRequest{} + req.Params.Cursor = cursor + result, err := c.ListResourceTemplates(ctx, req) + if err != nil { + return nil, "", err + } + return result.ResourceTemplates, result.NextCursor, nil + }) + // A backend that advertises the resources capability but does not + // implement resources/templates/list (JSON-RPC -32601) degrades to an + // empty template list — its other capabilities still aggregate. Other + // errors still propagate and drop the backend's capability set. + if errors.Is(err, mcp.ErrMethodNotFound) { + slog.Debug("backend does not implement resources/templates/list, treating templates as empty", + "backend", backendID) + return &mcp.ListResourceTemplatesResult{ResourceTemplates: []mcp.ResourceTemplate{}}, nil + } + if err != nil { + return nil, fmt.Errorf("failed to list resource templates from backend %s: %w", backendID, err) + } + return &mcp.ListResourceTemplatesResult{ResourceTemplates: templates}, nil + } + slog.Debug("backend does not advertise resources capability, skipping resource templates query", "backend", backendID) + return &mcp.ListResourceTemplatesResult{ResourceTemplates: []mcp.ResourceTemplate{}}, nil +} + +// queryPrompts queries prompts from a backend if the server advertises prompt +// support. It follows MCP pagination cursors (see queryTools). func queryPrompts(ctx context.Context, c *client.Client, supported bool, backendID string) (*mcp.ListPromptsResult, error) { if supported { - result, err := c.ListPrompts(ctx, mcp.ListPromptsRequest{}) + prompts, err := pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Prompt, mcp.Cursor, error) { + req := mcp.ListPromptsRequest{} + req.Params.Cursor = cursor + result, err := c.ListPrompts(ctx, req) + if err != nil { + return nil, "", err + } + return result.Prompts, result.NextCursor, nil + }) if err != nil { return nil, fmt.Errorf("failed to list prompts from backend %s: %w", backendID, err) } - return result, nil + return &mcp.ListPromptsResult{Prompts: prompts}, nil } slog.Debug("backend does not advertise prompts capability, skipping prompts query", "backend", backendID) return &mcp.ListPromptsResult{Prompts: []mcp.Prompt{}}, nil } -// ListCapabilities queries a backend for its MCP capabilities. -// Returns tools, resources, and prompts exposed by the backend. -// Only queries capabilities that the server advertises during initialization. -func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { - slog.Debug("querying capabilities from backend", "backend", target.WorkloadName, "url", target.BaseURL) +// cachedRevision returns the cached MCP revision for a backend. The second +// return is false when the backend has never been probed (distinct from a +// resolved RevisionLegacy). +func (h *httpBackendClient) cachedRevision(workloadID string) (mcpparser.Revision, bool) { + v, ok := h.revisions.Load(workloadID) + if !ok { + return 0, false + } + return v.(mcpparser.Revision), true +} - // Create a client for this backend (not yet initialized) - c, err := h.clientFactory(ctx, target) +// setRevision records a backend's resolved MCP revision. +func (h *httpBackendClient) setRevision(workloadID string, rev mcpparser.Revision) { + h.revisions.Store(workloadID, rev) +} + +// Compile-time guard that this client still satisfies the optional read-model +// every consumer type-asserts to. Without it, renaming or dropping +// CachedRevision below would compile cleanly and silently disable the telemetry +// revision label, the /status field, and the session layer's Modern-connect +// skip — each of which degrades gracefully by design when the assertion fails. +var _ vmcp.RevisionReporter = (*httpBackendClient)(nil) + +// CachedRevision reports a backend's resolved MCP revision for status/telemetry +// read-models. The second return is false when the backend has never been probed. +// Exported (not on vmcp.BackendClient) so the telemetry decorator, health +// monitor and session factory can read it via vmcp.RevisionReporter without an +// interface/mock change. +func (h *httpBackendClient) CachedRevision(workloadID string) (mcpparser.Revision, bool) { + return h.cachedRevision(workloadID) +} + +// buildModernHTTPClient wraps the shared backend RoundTripper chain (auth, +// identity, header-forward, trace, TLS/SSRF — see buildBackendRoundTripper) in an +// *http.Client for the raw Modern shim. This is a LIVE production path: the +// discover probe must carry the same security controls as every other backend +// call, so it must NOT use a bare http.Client. A 30s timeout matches the +// streamable-HTTP client; the response body is bounded inside modernCall +// (io.LimitReader), so no size-limit transport wrapper is needed here. +func (h *httpBackendClient) buildModernHTTPClient(ctx context.Context, target *vmcp.BackendTarget) (*http.Client, error) { + rt, err := h.buildBackendRoundTripper(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "create client") + return nil, err } - defer func() { - if err := c.Close(); err != nil { - slog.Debug("failed to close client", "error", err) - } - }() + return newBackendHTTPClient(rt, h.requestTimeout(target.WorkloadID)), nil +} - // Initialize the client and get server capabilities - serverCaps, err := initializeClient(ctx, c) +// newBackendHTTPClient is the single choke point for every backend *http.Client +// (Legacy streamable/SSE and Modern shim). It installs SameHostRedirectPolicy so a +// malicious backend's cross-host 30x cannot exfiltrate the auth/identity/ +// header-forward credentials the RoundTripper chain re-injects on each hop — Go's +// built-in Authorization stripping does not cover RoundTripper-injected headers. +// timeout 0 means no client timeout (long-lived SSE stream). +func newBackendHTTPClient(rt http.RoundTripper, timeout time.Duration) *http.Client { + return &http.Client{ + Transport: rt, + Timeout: timeout, + CheckRedirect: networking.SameHostRedirectPolicy(), + } +} + +// modernDiscover issues a Modern server/discover and returns the backend's +// capability flags. Used both by probeRevision and, on a Modern cache hit, by +// ListCapabilities to re-fetch the flags without re-running the fallback ladder. +// +// A clean discover response is not, by itself, proof the peer speaks Modern — +// supportedVersions is the authoritative signal. See errModernNegotiatedDown +// (modern.go) for why. Concretely: when supportedVersions does not contain +// MCPVersionModern (including when it is absent or empty), this returns +// errModernNegotiatedDown instead of the capabilities, so callers do not mistake +// a negotiated-down Legacy backend for Modern. +func (h *httpBackendClient) modernDiscover( + ctx context.Context, target *vmcp.BackendTarget, +) (*mcp.ServerCapabilities, error) { + hc, err := h.buildModernHTTPClient(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + return nil, err } + defer hc.CloseIdleConnections() + return discoverModernCapabilities(ctx, hc, target.BaseURL) +} - slog.Debug("backend capabilities", - "backend", target.WorkloadID, - "tools", serverCaps.Tools != nil, - "resources", serverCaps.Resources != nil, - "prompts", serverCaps.Prompts != nil) +// discoverModernCapabilities issues the Modern server/discover call over an +// already-built hc and applies the supportedVersions check (see modernDiscover's +// doc comment for why a clean response alone does not prove Modern). Split out +// so probeRevision can share this exact check over the *http.Client it must +// build anyway (to hard-fail on a genuine transport misconfiguration before +// classifying anything), instead of building the client twice. +func discoverModernCapabilities(ctx context.Context, hc *http.Client, endpoint string) (*mcp.ServerCapabilities, error) { + var discover struct { + Capabilities mcp.ServerCapabilities `json:"capabilities"` + SupportedVersions []string `json:"supportedVersions"` + } + if err := modernCall(ctx, hc, endpoint, "server/discover", nil, "", nil, &discover, "", nil); err != nil { + return nil, err + } + // Exact-match on MCPVersionModern (2026-07-28): vMCP's shim only speaks that + // one Modern wire shape, so a backend must advertise exactly it to be driven + // Modern. This is intentionally stricter than go-sdk's reference client + // (negotiated >= 2026-07-28). TRIPWIRE: when a newer Modern revision is + // added, this must become a set/range check (alongside the sibling + // exact-match in classifyUnsupportedProtocolVersion, modern.go) or a + // newer-only backend is wrongly classified Legacy — TestProbeRevision_RealBackends + // will catch it. + if !slices.Contains(discover.SupportedVersions, mcpparser.MCPVersionModern) { + return nil, fmt.Errorf("%w: supportedVersions=%v", errModernNegotiatedDown, discover.SupportedVersions) + } + return &discover.Capabilities, nil +} - // Query each capability type based on server advertisement - // Check for nil BEFORE passing to functions to avoid interface{} nil pointer issues - toolsResp, err := queryTools(ctx, c, serverCaps.Tools != nil, target.WorkloadID) +// probeRevision resolves a backend's MCP revision, Modern-first. +// +// It attempts a Modern server/discover (via discoverModernCapabilities, the same +// supportedVersions-checking core modernDiscover uses) and: +// - classifies MODERN (and caches it) when the discover result's +// supportedVersions contains MCPVersionModern, on a Modern-specific +// protocol error (-3202x) which proves the peer validated our Modern +// headers/_meta, or on an input_required envelope (only returned after +// decoding a valid Modern envelope, so it too proves Modern); +// - classifies LEGACY (and caches it) on a GENUINE not-Modern signal — +// errWrongEra, a -32601 (discover is mandatory for Modern), a generic +// JSON-RPC error, a bare 404/400/405, an empty/non-JSON body, a +// 200-with-Legacy-result, OR errModernNegotiatedDown (a clean discover +// envelope whose supportedVersions lacks 2026-07-28 — see +// errModernNegotiatedDown in modern.go for why a clean response alone is not +// proof of Modern); +// - returns the error UNCACHED (leaving the backend unprobed) on an +// INCONCLUSIVE outcome — an auth blip (errModernAuth, 401/403) or a transient +// failure (errModernTransient: 408/429/5xx, mid-read, or transport/timeout). +// Caching on these would flip a Modern backend to Legacy on a brief outage; +// leaving it unprobed lets the next call re-probe once the outage clears. +// +// A hard error is also returned when the backend transport cannot be built at all +// (e.g. invalid auth/CA config); that is a genuine misconfiguration. Note that +// dispatch, the sole caller, does not distinguish these error classes — it falls +// back to Legacy uncached on any error this function returns. A "sse" target +// (see below) returns before that check runs, so its transport is first +// validated on the Legacy call path instead. +// The resolved capabilities are intentionally not returned: callers that need +// them (ListCapabilities) re-fetch via modernDiscover, so probeRevision only +// classifies and caches the revision. +// +// A "sse" target is classified Legacy without a network call. TransportType == +// "sse" specifically names the deprecated 2024-11-05 two-endpoint transport +// (GET /sse + POST /messages) — see ssecommon.HTTPSSEEndpoint and the GET-only +// httpsse proxy — not "any backend that happens to use SSE" (Modern itself uses +// SSE response streams). modernCall is a single-endpoint POST, so it can never +// reach a Modern endpoint through an /sse BaseURL, even for a dual-era server +// that also hosts a Modern endpoint, because that lives at a different path. +// This also sidesteps POSTing into a GET-only stream and hanging to the client +// timeout (errModernTransient, uncached, re-probed on every call). +func (h *httpBackendClient) probeRevision( + ctx context.Context, target *vmcp.BackendTarget, +) (mcpparser.Revision, error) { + if target.TransportType == "sse" { + // The 2024-11-05 two-endpoint transport (GET /sse + POST /messages) has no + // Modern endpoint to discover at this BaseURL — see the doc comment above. + // Note: HTTP+SSE is Deprecated per the MCP spec, not removed; this gate is + // ToolHive routing, not a protocol requirement. + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + return mcpparser.RevisionLegacy, nil + } + + hc, err := h.buildModernHTTPClient(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "list tools") + return 0, fmt.Errorf("failed to build transport for backend %s: %w", target.WorkloadID, err) } + defer hc.CloseIdleConnections() - resourcesResp, err := queryResources(ctx, c, serverCaps.Resources != nil, target.WorkloadID) - if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "list resources") + _, err = discoverModernCapabilities(ctx, hc, target.BaseURL) + switch { + case err == nil: + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + return mcpparser.RevisionModern, nil + case errors.Is(err, errModernProtocolError), errors.Is(err, errModernInputRequired): + // Both prove the peer is Modern: -3202x means it validated our Modern + // protocol metadata; input_required is only returned after decoding a valid + // Modern envelope (resultType present, != "complete"). + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + return mcpparser.RevisionModern, nil + case errors.Is(err, errModernAuth), errors.Is(err, errModernTransient): + // Inconclusive: an auth blip or a transient outage tells us nothing about + // the revision. Leave the backend unprobed so the next call re-probes. + return 0, err + default: + // Includes errModernNegotiatedDown (a clean-but-negotiated-down discover + // envelope) alongside errWrongEra and every other genuine not-Modern + // signal: all fall back to Legacy here. + slog.Debug("backend is not Modern; falling back to Legacy", + "backend", target.WorkloadID, "probe_error", err) + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + return mcpparser.RevisionLegacy, nil } +} - promptsResp, err := queryPrompts(ctx, c, serverCaps.Prompts != nil, target.WorkloadID) +// modernEnumerate builds a backend's CapabilityList by enumerating each +// capability the discover flags advertise, via the Modern (2026-07-28) shim. +// Each list is gated on the matching discover flag (mirroring the Legacy +// initialize-path gating) and follows nextCursor across pages (#5851). +// +// caps may be nil: a Modern backend that rejected our discover with a -3202x +// protocol error (errModernProtocolError) is still classified Modern but yields +// no capability flags, so nothing is enumerated and an empty list is returned — +// the same outcome on the first probe and on later cache hits. +func (h *httpBackendClient) modernEnumerate( + ctx context.Context, target *vmcp.BackendTarget, caps *mcp.ServerCapabilities, +) (*vmcp.CapabilityList, error) { + hc, err := h.buildModernHTTPClient(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "list prompts") + return nil, wrapBackendError(err, target.WorkloadID, "create client") } + defer hc.CloseIdleConnections() + endpoint := target.BaseURL - // Convert MCP types to vmcp types - capabilities := &vmcp.CapabilityList{ - Tools: make([]vmcp.Tool, len(toolsResp.Tools)), - Resources: make([]vmcp.Resource, len(resourcesResp.Resources)), - Prompts: make([]vmcp.Prompt, len(promptsResp.Prompts)), + var tools []mcp.Tool + if caps != nil && caps.Tools != nil { + tools, err = modernListAll[mcp.Tool](ctx, hc, endpoint, "tools/list", "tools") + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list tools") + } + } + + var resources []mcp.Resource + var templates []mcp.ResourceTemplate + if caps != nil && caps.Resources != nil { + resources, err = modernListAll[mcp.Resource](ctx, hc, endpoint, "resources/list", "resources") + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list resources") + } + + // Resource templates share the resources capability flag. A backend that + // does not implement resources/templates/list (-32601) degrades to an + // empty template list, mirroring the Legacy queryResourceTemplates path. + templates, err = modernListAll[mcp.ResourceTemplate](ctx, hc, endpoint, "resources/templates/list", "resourceTemplates") + switch { + case errors.Is(err, mcp.ErrMethodNotFound): + templates = nil + case err != nil: + return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") + } } - // Convert tools - for i, tool := range toolsResp.Tools { - capabilities.Tools[i] = vmcp.Tool{ + var prompts []mcp.Prompt + if caps != nil && caps.Prompts != nil { + prompts, err = modernListAll[mcp.Prompt](ctx, hc, endpoint, "prompts/list", "prompts") + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list prompts") + } + } + + slog.Debug("backend capabilities queried (modern)", + "backend", target.WorkloadName, + "tools", len(tools), "resources", len(resources), + "resource_templates", len(templates), "prompts", len(prompts)) + return newCapabilityListFromMCP(target.WorkloadID, tools, resources, templates, prompts), nil +} + +// cursorParams builds the Modern list request params carrying a pagination +// cursor, or nil for the first page. +func cursorParams(cursor mcp.Cursor) map[string]any { + if cursor == "" { + return nil + } + return map[string]any{"cursor": string(cursor)} +} + +// modernListAll fetches every page of a Modern */list method, following +// nextCursor (#5851). itemsField is the result key holding the item array +// ("tools", "resources", "resourceTemplates", "prompts"); the envelope is decoded +// into a raw-message map so one helper serves all four list shapes. +func modernListAll[T any]( + ctx context.Context, hc *http.Client, endpoint, method, itemsField string, +) ([]T, error) { + return pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]T, mcp.Cursor, error) { + var page map[string]json.RawMessage + if err := modernCall(ctx, hc, endpoint, method, cursorParams(cursor), "", nil, &page, "", nil); err != nil { + return nil, "", err + } + var items []T + if raw, ok := page[itemsField]; ok { + if err := json.Unmarshal(raw, &items); err != nil { + return nil, "", fmt.Errorf("decoding %s from %s: %w", itemsField, method, err) + } + } + var next mcp.Cursor + if raw, ok := page["nextCursor"]; ok { + _ = json.Unmarshal(raw, &next) // best-effort: a malformed cursor ends pagination + } + return items, next, nil + }) +} + +// newCapabilityListFromMCP converts backend mcp types into the vmcp domain +// CapabilityList, tagging every item with backendID. Shared by the Legacy +// (initialize+enumerate) and Modern (discover+enumerate) paths so both produce +// identical domain shapes. +func newCapabilityListFromMCP( + backendID string, + tools []mcp.Tool, resources []mcp.Resource, templates []mcp.ResourceTemplate, prompts []mcp.Prompt, +) *vmcp.CapabilityList { + capabilities := &vmcp.CapabilityList{ + Tools: make([]vmcp.Tool, 0, len(tools)), + Resources: make([]vmcp.Resource, len(resources)), + ResourceTemplates: make([]vmcp.ResourceTemplate, len(templates)), + Prompts: make([]vmcp.Prompt, len(prompts)), + } + + for _, tool := range tools { + inputSchema := conversion.ConvertToolInputSchema(tool.InputSchema) + + // SEP-2243 requires a Streamable HTTP client to reject a tool definition + // whose x-mcp-header annotations violate the extension's constraints, and + // vMCP is this backend's client. Rejecting here — at the single ingestion + // seam both the Legacy and Modern paths share — also protects vMCP's own + // downstream clients: an invalid annotation republished in the aggregated + // tools/list would make a conformant downstream client reject the whole + // list, taking every other backend's tools down with it. Dropping the one + // offending tool is the narrower failure. + if err := mcpparser.ValidateParamHeaders(inputSchema); err != nil { + slog.Warn("rejecting backend tool with invalid x-mcp-header annotation", + "backend", backendID, "tool", tool.Name, "error", err) + continue + } + + capabilities.Tools = append(capabilities.Tools, vmcp.Tool{ Name: tool.Name, Description: tool.Description, - InputSchema: conversion.ConvertToolInputSchema(tool.InputSchema), + InputSchema: inputSchema, OutputSchema: conversion.ConvertToolOutputSchema(tool.OutputSchema), Annotations: conversion.ConvertToolAnnotations(tool.Annotations), - BackendID: target.WorkloadID, - } + BackendID: backendID, + }) } - // Convert resources - for i, resource := range resourcesResp.Resources { + for i, resource := range resources { capabilities.Resources[i] = vmcp.Resource{ URI: resource.URI, Name: resource.Name, Description: resource.Description, MimeType: resource.MIMEType, - BackendID: target.WorkloadID, + BackendID: backendID, } } - // Convert prompts - for i, prompt := range promptsResp.Prompts { + // Resource templates are a pass-through: no URI-template rewriting, like resources. + for i, template := range templates { + capabilities.ResourceTemplates[i] = vmcp.ResourceTemplate{ + URITemplate: template.URITemplate, + Name: template.Name, + Description: template.Description, + MimeType: template.MIMEType, + BackendID: backendID, + } + } + + for i, prompt := range prompts { args := make([]vmcp.PromptArgument, len(prompt.Arguments)) for j, arg := range prompt.Arguments { args[j] = vmcp.PromptArgument{ @@ -878,15 +1438,109 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B Required: arg.Required, } } - capabilities.Prompts[i] = vmcp.Prompt{ Name: prompt.Name, Description: prompt.Description, Arguments: args, - BackendID: target.WorkloadID, + BackendID: backendID, + } + } + + return capabilities +} + +// ListCapabilities queries a backend for its MCP capabilities, selecting the +// Legacy or Modern path by revision. Returns tools, resources, and prompts +// exposed by the backend. +// +// Like the call verbs it routes through dispatch, so a mis-cached backend (e.g. a +// first-probe blip that pinned Legacy) self-corrects: the failing path triggers a +// re-probe and one retry under the corrected revision. +func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + slog.Debug("querying capabilities from backend", "backend", target.WorkloadName, "url", target.BaseURL) + var out *vmcp.CapabilityList + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernListCapabilities(ctx, target) + } else { + out, err = h.legacyListCapabilities(ctx, target) + } + return err + }) + return out, err +} + +// modernListCapabilities resolves capabilities via Modern server/discover + +// enumeration. A -3202x protocol error is tolerated as nil caps (the backend is +// Modern but rejected our discover), yielding an empty list — consistent with +// probeRevision's classification. +func (h *httpBackendClient) modernListCapabilities( + ctx context.Context, target *vmcp.BackendTarget, +) (*vmcp.CapabilityList, error) { + caps, err := h.modernDiscover(ctx, target) + if err != nil && !errors.Is(err, errModernProtocolError) { + return nil, wrapBackendError(err, target.WorkloadID, "modern discover") + } + return h.modernEnumerate(ctx, target, caps) +} + +// legacyListCapabilities is the initialize + enumerate path (unchanged behavior). +func (h *httpBackendClient) legacyListCapabilities( + ctx context.Context, target *vmcp.BackendTarget, +) (*vmcp.CapabilityList, error) { + // Create a client for this backend (not yet initialized) + c, err := h.clientFactory(ctx, target, false) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer func() { + if err := c.Close(); err != nil { + slog.Debug("failed to close client", "error", err) } + }() + + // Initialize the client and get server capabilities + serverCaps, _, err := h.legacyInit(ctx, c, target) + if err != nil { + return nil, err } + slog.Debug("backend capabilities", + "backend", target.WorkloadID, + "tools", serverCaps.Tools != nil, + "resources", serverCaps.Resources != nil, + "prompts", serverCaps.Prompts != nil) + + // Query each capability type based on server advertisement + // Check for nil BEFORE passing to functions to avoid interface{} nil pointer issues + toolsResp, err := queryTools(ctx, c, serverCaps.Tools != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list tools") + } + + resourcesResp, err := queryResources(ctx, c, serverCaps.Resources != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list resources") + } + + // Resource templates share the same server capability advertisement as resources. + resourceTemplatesResp, err := queryResourceTemplates(ctx, c, serverCaps.Resources != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") + } + + promptsResp, err := queryPrompts(ctx, c, serverCaps.Prompts != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list prompts") + } + + // Convert MCP types to vmcp types (shared with the Modern enumeration path). + capabilities := newCapabilityListFromMCP( + target.WorkloadID, + toolsResp.Tools, resourcesResp.Resources, resourceTemplatesResp.ResourceTemplates, promptsResp.Prompts, + ) + // TODO: Query server capabilities to detect logging/sampling support // This requires additional MCP protocol support for capabilities introspection @@ -894,26 +1548,306 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B "backend", target.WorkloadName, "tools", len(capabilities.Tools), "resources", len(capabilities.Resources), + "resource_templates", len(capabilities.ResourceTemplates), "prompts", len(capabilities.Prompts)) return capabilities, nil } -// CallTool invokes a tool on the backend MCP server. -// Returns the complete tool result including _meta field. +// errLegacyInitFailed marks a Legacy initialize-step failure (see legacyInit). +// It scopes the Legacy revision-mismatch signal to the initialize step so a +// data-plane -32601 from a genuine tool/resource/prompt/completion call on a real +// Legacy backend (a legitimately unimplemented method) never triggers a re-probe. +var errLegacyInitFailed = errors.New("legacy initialize step failed") + +// legacyInit runs the Legacy initialize handshake, tagging any failure with +// errLegacyInitFailed (in addition to wrapBackendError's classification) so the +// revision-mismatch check can tell an initialize rejection apart from a +// data-plane error later in the same call. // -//nolint:gocyclo // this function is complex because it handles tool calls with various content types and error handling. +// It also self-heals a mis-cached Legacy->Modern backend in-band: under go-sdk +// v1.7 the SDK's own client transparently negotiates Modern on this path (see +// the revisions field comment), so the call succeeds and dispatch's +// reclassify-on-error trigger never fires. The negotiated protocol version +// returned by initializeClient is genuine either way (discover or plain +// initialize), so when it equals MCPVersionModern the cache is flipped here +// instead of waiting for an error that will never come. +// +// The negotiated version is a HINT, not proof, and it is only allowed to +// override an existing classification once server/discover has confirmed it. +// The two genuinely disagree in the field: github-mcp-server v1.6.0 negotiates +// 2026-07-28 on the Legacy initialize while answering server/discover with a +// body carrying no resultType, which modernCall rightly rejects as +// Legacy-shaped. Promoting on the hint alone made the two self-corrections +// fight -- the probe cached Legacy, this promoted back to Modern, the next +// Modern call failed on that same body and reclassified to Legacy -- so every +// other health check failed indefinitely (#6154). +// +// So: with no cached revision the hint is trusted outright (that is the +// uncached-Legacy fallback dispatch takes when a probe errors, and the case +// this self-heal was written for). Against a cached Legacy classification the +// hint must win a confirming probe first; probeRevision caches whatever it +// finds, so a genuinely Modern backend still self-heals in one extra round +// trip. A refuted hint is remembered (modernHintRefuted) so the confirming +// probe runs at most once per refutationTTL window per backend rather than on +// every call — and a redeemed backend (refuted once, later genuinely Modern) +// self-heals within one window instead of never. +func (h *httpBackendClient) legacyInit( + ctx context.Context, c *client.Client, target *vmcp.BackendTarget, +) (*mcp.ServerCapabilities, string, error) { + backendID := target.WorkloadID + caps, negotiatedVersion, err := initializeClient(ctx, c) + if err != nil { + return nil, "", fmt.Errorf("%w: %w", errLegacyInitFailed, wrapBackendError(err, backendID, "initialize client")) + } + if negotiatedVersion != mcpparser.MCPVersionModern { + return caps, negotiatedVersion, nil + } + + cached, isCached := h.cachedRevision(backendID) + switch { + case !isCached: + h.setRevision(backendID, mcpparser.RevisionModern) + case cached == mcpparser.RevisionLegacy: + if h.hintRefuted(backendID) { + break + } + // probeRevision caches its own verdict, so a confirmation needs no + // setRevision here. A probe error leaves the cache untouched. + if probed, perr := h.probeRevision(ctx, target); perr == nil && probed != mcpparser.RevisionModern { + slog.DebugContext(ctx, "legacy handshake negotiated Modern but discover disagrees; keeping Legacy", + "backend", backendID, "negotiated", negotiatedVersion) + h.modernHintRefuted.Store(backendID, time.Now()) + } + } + return caps, negotiatedVersion, nil +} + +// hintRefuted reports whether backendID's Modern-handshake hint is currently +// refuted: a refutation was recorded and has not yet aged past refutationTTL. +// An expired entry is deleted so the map stays bounded by live backends. +func (h *httpBackendClient) hintRefuted(backendID string) bool { + v, ok := h.modernHintRefuted.Load(backendID) + if !ok { + return false + } + refutedAt, ok := v.(time.Time) + if !ok || time.Since(refutedAt) > h.refutationTTL { + h.modernHintRefuted.Delete(backendID) + return false + } + return true +} + +// dispatch resolves the backend's MCP revision (cache hit, else probeRevision) +// and runs fn with it. Every call verb AND ListCapabilities route through this +// seam so revision selection — and self-correction — lives in one place. +// +// A probeRevision error of any kind (inconclusive probe or transport-build +// failure) falls back to Legacy — the pre-dual-era default — WITHOUT caching, +// so a genuinely Modern backend re-probes on the next call instead of being +// pinned to Legacy by one bad probe. +// +// On a revision mismatch (isRevisionMismatch), dispatch always re-probes +// authoritatively and flips the cache so future calls use the corrected +// revision. It RETRIES fn only when the failure proves the backend did NOT +// execute the request — a protocol rejection (errWrongEra) or a Legacy +// initialize-step failure. A Legacy-shaped success body (errLegacyResponseBody) +// means a lenient backend MAY have executed a side-effecting request, so the +// cache is flipped but fn is NOT re-run. The retry is unconditionally single +// (no re-check), so it can never loop. +// +// When the revision was resolved in THIS call (uncached), a mismatch is surfaced +// directly without re-probing: on a successful probe a re-probe would just repeat +// the same answer, and on the Legacy fallback above there is no cached revision to +// correct — the next call re-probes anyway once the outage clears. +func (h *httpBackendClient) dispatch( + ctx context.Context, target *vmcp.BackendTarget, + fn func(ctx context.Context, rev mcpparser.Revision) error, +) error { + rev, cached := h.cachedRevision(target.WorkloadID) + if !cached { + probed, err := h.probeRevision(ctx, target) + if err != nil { + // A failed probe (auth blip / transient 5xx / timeout / connection + // refused, or a transport that could not be built) says nothing about + // the revision and must NOT fail a backend that is reachable on the + // Legacy path — which re-validates the transport itself. Fall back to Legacy — + // the pre-dual-era default — WITHOUT caching, so a transient outage + // never pins a revision and a genuinely Modern backend is corrected on + // the next call's re-probe. + slog.DebugContext(ctx, "revision probe inconclusive; attempting Legacy uncached", + "backend", target.WorkloadID, "probe_error", err) + rev = mcpparser.RevisionLegacy + } else { + rev = probed + } + } + + err := fn(ctx, rev) + if err == nil || !isRevisionMismatch(rev, err) { + return err + } + if !cached { + // The revision was just probed authoritatively; a re-probe would agree, so + // this is a genuine error (or a lenient backend), not a stale cache. + return err + } + + corrected := h.reclassify(ctx, target, rev) + if corrected == rev { + // Re-probe agreed with the cache: the mismatch was not a revision problem. + return err + } + if errors.Is(err, errLegacyResponseBody) { + // Cache is now corrected for future calls, but the backend may have + // already executed this request — do NOT re-run it (no double-execution). + return err + } + return fn(ctx, corrected) +} + +// reclassify re-probes the backend authoritatively (overwriting the cached +// revision via probeRevision) and returns the corrected revision. On an actual +// era change it emits a WARN and increments the reclassification counter; a +// same-era re-probe (or a re-probe that can't run) is silent and returns prev. +func (h *httpBackendClient) reclassify( + ctx context.Context, target *vmcp.BackendTarget, prev mcpparser.Revision, +) mcpparser.Revision { + corrected, err := h.probeRevision(ctx, target) + if err != nil { + // Transport couldn't even be built to re-probe; keep the prior revision. + return prev + } + if corrected != prev { + slog.WarnContext(ctx, "backend MCP revision reclassified after mismatch", + "backend", target.WorkloadID, "old", prev.String(), "new", corrected.String()) + backendtelemetry.RecordRevisionReclassification(ctx) + } + return corrected +} + +// isRevisionMismatch reports whether err from an attempt made under rev signals +// that the backend actually speaks the OTHER MCP revision (triggering a cache +// reclassification). It is deliberately narrow, and keyed on rev because the same +// sentinel means different things per era: +// +// - Modern attempt: errWrongEra (protocol rejection), errLegacyResponseBody (a +// Legacy-shaped success body), or errModernNegotiatedDown (a cached-Modern +// backend's discover now reports supportedVersions without 2026-07-28 — e.g. +// redeployed stateless->stateful). A data-plane -32601 comes back as +// mcp.ErrMethodNotFound and is a real not-found on a genuine Modern backend, +// NOT a mismatch. +// - Legacy attempt: an initialize-STEP rejection only — errLegacyInitFailed +// together with mcp.ErrMethodNotFound or transport.ErrLegacySSEServer. A +// data-plane -32601 (a legitimately unimplemented method on a real Legacy +// backend) lacks the errLegacyInitFailed marker and is NOT a mismatch. +// +// Auth failures are never a mismatch: their sentinels (ErrUnauthorized, +// ErrAuthorizationRequired, ErrUpstreamTokenNotFound, ErrAuthenticationFailed) do +// not wrap the era sentinels above, so they are excluded by construction. +// +// NOTE: mismatch != safe-to-retry. dispatch reclassifies on any mismatch but only +// re-runs fn when no execution could have occurred (see dispatch). +func isRevisionMismatch(rev mcpparser.Revision, err error) bool { + if err == nil { + return false + } + if rev == mcpparser.RevisionModern { + return errors.Is(err, errWrongEra) || errors.Is(err, errLegacyResponseBody) || + errors.Is(err, errModernNegotiatedDown) + } + return errors.Is(err, errLegacyInitFailed) && + (errors.Is(err, mcp.ErrMethodNotFound) || errors.Is(err, transport.ErrLegacySSEServer)) +} + +// CallTool invokes a tool on the backend MCP server, selecting the Legacy or +// Modern (2026-07-28) path by the backend's resolved revision. Returns the +// complete tool result including _meta. func (h *httpBackendClient) CallTool( ctx context.Context, target *vmcp.BackendTarget, toolName string, arguments map[string]any, meta map[string]any, + paramHeaders map[string]string, ) (*vmcp.ToolCallResult, error) { slog.Debug("calling tool on backend", "tool", toolName, "backend", target.WorkloadName) + var out *vmcp.ToolCallResult + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernCallTool(ctx, target, toolName, arguments, meta, paramHeaders) + } else { + // paramHeaders is deliberately dropped here: SEP-2243's Mcp-Param-* + // mirroring belongs to the 2026-07-28 revision, and a Legacy backend + // neither expects the headers nor rejects their absence. + out, err = h.legacyCallTool(ctx, target, toolName, arguments, meta) + } + return err + }) + return out, err +} + +// modernCallTool invokes tools/call over the Modern (2026-07-28) shim. The +// advertised tool name is translated to the backend's capability name for BOTH +// the body identifier and the Mcp-Name header — the server rejects a mismatch +// (-32020). The caller's _meta is forwarded (modernCall strips reserved keys and +// overlays vMCP's) and the result _meta is forwarded back to core. +// +// paramHeaders carries the SEP-2243 Mcp-Param-* headers derived from the tool's +// x-mcp-header-designated arguments. A backend that designated a parameter and +// does not receive its header rejects the call with -32020, so omitting them +// makes any annotating Modern backend uncallable. +func (h *httpBackendClient) modernCallTool( + ctx context.Context, target *vmcp.BackendTarget, toolName string, arguments, meta map[string]any, + paramHeaders map[string]string, +) (*vmcp.ToolCallResult, error) { + backendToolName := target.GetBackendCapabilityName(toolName) + if backendToolName != toolName { + slog.Debug("translating tool name", "client_name", toolName, "backend_name", backendToolName) + } + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + params := map[string]any{"name": backendToolName, "arguments": arguments} + if len(meta) > 0 { + params["_meta"] = meta + } + + // Modern logging parity with the Legacy path (enableBackendLogging): when + // server->client forwarding is bound, opt the call in to the backend's + // notifications/message at debug level via the per-request logLevel _meta key + // (the Modern replacement for the removed logging/setLevel RPC), and relay any + // interleaved notifications to the downstream client. Both are no-ops when + // forwarding is unbound — no logLevel key is sent and no listener is attached. + fwd := h.forwarders.Load() + var logLevel string + var onNotification func(string, json.RawMessage) + if fwd != nil && fwd.notifier != nil { + logLevel = string(mcp.LoggingLevelDebug) + onNotification = newModernNotificationForwarder(ctx, fwd.notifier) + } + + var result mcp.CallToolResult + if err := modernCall( + ctx, hc, target.BaseURL, "tools/call", params, backendToolName, paramHeaders, &result, logLevel, onNotification, + ); err != nil { + return nil, fmt.Errorf("%w: tool call failed on backend %s: %w", vmcp.ErrBackendUnavailable, target.WorkloadID, err) + } + return toolResultFromMCP(&result, toolName, target.WorkloadID, target.OutputSchema != nil), nil +} +// legacyCallTool is the initialize + SDK CallTool path (unchanged behavior). +// +//nolint:gocyclo // this function is complex because it handles tool calls with various content types and error handling. +func (h *httpBackendClient) legacyCallTool( + ctx context.Context, target *vmcp.BackendTarget, toolName string, arguments, meta map[string]any, +) (*vmcp.ToolCallResult, error) { // Create a client for this backend - c, err := h.clientFactory(ctx, target) + c, err := h.clientFactory(ctx, target, true) if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "create client") } @@ -923,11 +1857,18 @@ func (h *httpBackendClient) CallTool( } }() - // Initialize the client - if _, err := initializeClient(ctx, c); err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + // Initialize the client and capture the backend's advertised capabilities. + serverCaps, negotiatedVersion, err := h.legacyInit(ctx, c, target) + if err != nil { + return nil, err } + // When forwarders are bound and the backend advertises logging, request debug + // level so the backend emits notifications/message during the call; the + // notification forwarder relays them to the downstream client. Best-effort: + // a failure here must not fail the tool call. + h.enableBackendLogging(ctx, c, serverCaps, negotiatedVersion, target.WorkloadID) + // Call the tool using the original capability name from the backend's perspective. // When conflict resolution renames tools (e.g., "fetch" → "fetch_fetch"), // we must use the original backend name when forwarding requests. @@ -936,17 +1877,43 @@ func (h *httpBackendClient) CallTool( slog.Debug("translating tool name", "client_name", toolName, "backend_name", backendToolName) } + // Strip the reserved io.modelcontextprotocol/* keys before forwarding to this + // Legacy (session-based, stateful) backend: a downstream Modern caller's + // _meta.protocolVersion is only valid on a stateless Modern hop, and go-sdk + // v1.7 hard-rejects ANY _meta.protocolVersion on a stateful server (HTTP 400) + // regardless of its value — see mcpparser.StripReservedMeta. + // + // conversion.ToMCPMeta below now strips too, so this is belt-and-suspenders; + // it stays because the HTTP-400 hazard is specific to this hop and deserves a + // guard at the site that knows why. (modernCallTool, which builds its _meta + // directly and never reaches ToMCPMeta, needs the helper outright.) + meta = mcpparser.StripReservedMeta(meta) result, err := c.CallTool(ctx, mcp.CallToolRequest{ Params: mcp.CallToolParams{ Name: backendToolName, Arguments: arguments, - Meta: conversion.ToMCPMeta(meta), + Meta: conversion.ToMCPMeta(telemetry.MetaWithTraceContext(ctx, meta)), }, }) if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "call tool") } + // Flush the backend's server->client stream before the deferred Close tears + // down this per-call client, so a fire-and-forget notification the backend + // emitted mid-call (progress/logging) is relayed downstream instead of being + // dropped. See drainServerToClientNotifications for the lost-notification race. + h.drainServerToClientNotifications(ctx, c) + + return toolResultFromMCP(result, toolName, target.WorkloadID, target.OutputSchema != nil), nil +} + +// toolResultFromMCP converts an mcp.CallToolResult into the vmcp domain result, +// shared by the Legacy and Modern tools/call paths so both handle IsError +// logging, structuredContent, and _meta forwarding identically. +func toolResultFromMCP( + result *mcp.CallToolResult, toolName, backendID string, hasOutputSchema bool, +) *vmcp.ToolCallResult { // Extract _meta field from backend response responseMeta := conversion.FromMCPMeta(result.Meta) @@ -966,10 +1933,10 @@ func (h *httpBackendClient) CallTool( // Log with metadata for distributed tracing if responseMeta != nil { slog.Warn("tool returned IsError=true", - "tool", toolName, "backend", target.WorkloadID, "error", errorMsg, "meta", responseMeta) + "tool", toolName, "backend", backendID, "error", errorMsg, "meta", responseMeta) } else { slog.Warn("tool returned IsError=true", - "tool", toolName, "backend", target.WorkloadID, "error", errorMsg) + "tool", toolName, "backend", backendID, "error", errorMsg) } // Continue processing - we return the result with IsError flag and metadata preserved } @@ -980,10 +1947,10 @@ func (h *httpBackendClient) CallTool( structuredContent := conversion.ToolResultStructuredContent( result.StructuredContent, contentArray, - target.OutputSchema != nil, + hasOutputSchema, ) if structuredContent != nil { - slog.Debug("resolved structured content for tool", "tool", toolName, "backend", target.WorkloadID) + slog.Debug("resolved structured content for tool", "tool", toolName, "backend", backendID) } return &vmcp.ToolCallResult{ @@ -991,18 +1958,82 @@ func (h *httpBackendClient) CallTool( StructuredContent: structuredContent, IsError: result.IsError, Meta: responseMeta, - }, nil + } } -// ReadResource retrieves a resource from the backend MCP server. -// Returns the complete resource result including _meta field. +// ReadResource retrieves a resource from the backend MCP server, selecting the +// Legacy or Modern path by revision. Returns the complete resource result +// including _meta. func (h *httpBackendClient) ReadResource( ctx context.Context, target *vmcp.BackendTarget, uri string, ) (*vmcp.ResourceReadResult, error) { slog.Debug("reading resource from backend", "resource", uri, "backend", target.WorkloadName) + var out *vmcp.ResourceReadResult + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernReadResource(ctx, target, uri) + } else { + out, err = h.legacyReadResource(ctx, target, uri) + } + return err + }) + return out, err +} +// modernReadResource invokes resources/read over the Modern shim. The URI is +// translated to the backend's capability name for both the body and the Mcp-Name +// header (server rejects a mismatch), and the result _meta is forwarded to core. +func (h *httpBackendClient) modernReadResource( + ctx context.Context, target *vmcp.BackendTarget, uri string, +) (*vmcp.ResourceReadResult, error) { + backendURI := target.GetBackendCapabilityName(uri) + if backendURI != uri { + slog.Debug("translating resource URI", "client_uri", uri, "backend_uri", backendURI) + } + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + // mcp.ResourceContents is an interface with no JSON unmarshaler, so it cannot + // decode directly. Decode the wire shape, rebuild the discriminated mcp types + // (blob takes precedence, symmetric with conversion.ToMCPResourceContents), + // then hand off to the SAME converter legacyReadResource uses so the + // content->vmcp mapping is shared, not duplicated. + var res struct { + Contents []struct { + URI string `json:"uri"` + MIMEType string `json:"mimeType"` + Text string `json:"text"` + Blob string `json:"blob"` + } `json:"contents"` + Meta map[string]any `json:"_meta"` + } + params := map[string]any{"uri": backendURI} + if err := modernCall(ctx, hc, target.BaseURL, "resources/read", params, backendURI, nil, &res, "", nil); err != nil { + return nil, fmt.Errorf("resource read failed on backend %s: %w", target.WorkloadID, err) + } + mcpContents := make([]mcp.ResourceContents, len(res.Contents)) + for i, c := range res.Contents { + if c.Blob != "" { + mcpContents[i] = mcp.BlobResourceContents{URI: c.URI, MIMEType: c.MIMEType, Blob: c.Blob} + } else { + mcpContents[i] = mcp.TextResourceContents{URI: c.URI, MIMEType: c.MIMEType, Text: c.Text} + } + } + return &vmcp.ResourceReadResult{ + Contents: conversion.ConvertMCPResourceContents(mcpContents), + Meta: res.Meta, + }, nil +} + +// legacyReadResource is the initialize + SDK ReadResource path (unchanged behavior). +func (h *httpBackendClient) legacyReadResource( + ctx context.Context, target *vmcp.BackendTarget, uri string, +) (*vmcp.ResourceReadResult, error) { // Create a client for this backend - c, err := h.clientFactory(ctx, target) + c, err := h.clientFactory(ctx, target, false) if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "create client") } @@ -1013,8 +2044,8 @@ func (h *httpBackendClient) ReadResource( }() // Initialize the client - if _, err := initializeClient(ctx, c); err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + if _, _, err := h.legacyInit(ctx, c, target); err != nil { + return nil, err } // Read the resource using the original URI from the backend's perspective. @@ -1024,9 +2055,12 @@ func (h *httpBackendClient) ReadResource( slog.Debug("translating resource URI", "client_uri", uri, "backend_uri", backendURI) } + // Forward-compat: mcpcompat drops Params.Meta on this path today (no-op on + // the wire). See TestOutboundMetaTraceContext for details/tripwire. result, err := c.ReadResource(ctx, mcp.ReadResourceRequest{ Params: mcp.ReadResourceParams{ - URI: backendURI, + URI: backendURI, + Meta: conversion.ToMCPMeta(telemetry.MetaWithTraceContext(ctx, nil)), }, }) if err != nil { @@ -1036,16 +2070,14 @@ func (h *httpBackendClient) ReadResource( // Extract _meta field from backend response meta := conversion.FromMCPMeta(result.Meta) - // Note: Due to MCP SDK limitations, the SDK's ReadResourceResult may not include Meta. - // This preserves it for future SDK improvements. return &vmcp.ResourceReadResult{ Contents: conversion.ConvertMCPResourceContents(result.Contents), Meta: meta, }, nil } -// GetPrompt retrieves a prompt from the backend MCP server. -// Returns the complete prompt result including _meta field. +// GetPrompt retrieves a prompt from the backend MCP server, selecting the Legacy +// or Modern path by revision. Returns the complete prompt result including _meta. func (h *httpBackendClient) GetPrompt( ctx context.Context, target *vmcp.BackendTarget, @@ -1053,9 +2085,71 @@ func (h *httpBackendClient) GetPrompt( arguments map[string]any, ) (*vmcp.PromptGetResult, error) { slog.Debug("getting prompt from backend", "prompt", name, "backend", target.WorkloadName) + var out *vmcp.PromptGetResult + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernGetPrompt(ctx, target, name, arguments) + } else { + out, err = h.legacyGetPrompt(ctx, target, name, arguments) + } + return err + }) + return out, err +} + +// modernGetPrompt invokes prompts/get over the Modern shim. The prompt name is +// translated to the backend's capability name for both the body and the Mcp-Name +// header (server rejects a mismatch); the result _meta is forwarded to core. +// Message content is decoded via mcp.UnmarshalContent so it goes through the same +// content conversion as the Legacy path. +func (h *httpBackendClient) modernGetPrompt( + ctx context.Context, target *vmcp.BackendTarget, name string, arguments map[string]any, +) (*vmcp.PromptGetResult, error) { + backendPromptName := target.GetBackendCapabilityName(name) + if backendPromptName != name { + slog.Debug("translating prompt name", "client_name", name, "backend_name", backendPromptName) + } + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + params := map[string]any{ + "name": backendPromptName, + "arguments": conversion.ConvertPromptArguments(arguments), + } + var res struct { + Description string `json:"description"` + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + Meta map[string]any `json:"_meta"` + } + if err := modernCall(ctx, hc, target.BaseURL, "prompts/get", params, backendPromptName, nil, &res, "", nil); err != nil { + return nil, fmt.Errorf("prompt get failed on backend %s: %w", target.WorkloadID, err) + } + messages := make([]vmcp.PromptMessage, 0, len(res.Messages)) + for _, m := range res.Messages { + content, err := mcp.UnmarshalContent(m.Content) + if err != nil { + return nil, fmt.Errorf("prompt get: decoding message content from backend %s: %w", target.WorkloadID, err) + } + messages = append(messages, vmcp.PromptMessage{ + Role: m.Role, + Content: conversion.ConvertMCPContent(content), + }) + } + return &vmcp.PromptGetResult{Messages: messages, Description: res.Description, Meta: res.Meta}, nil +} +// legacyGetPrompt is the initialize + SDK GetPrompt path (unchanged behavior). +func (h *httpBackendClient) legacyGetPrompt( + ctx context.Context, target *vmcp.BackendTarget, name string, arguments map[string]any, +) (*vmcp.PromptGetResult, error) { // Create a client for this backend - c, err := h.clientFactory(ctx, target) + c, err := h.clientFactory(ctx, target, false) if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "create client") } @@ -1066,8 +2160,8 @@ func (h *httpBackendClient) GetPrompt( }() // Initialize the client - if _, err := initializeClient(ctx, c); err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + if _, _, err := h.legacyInit(ctx, c, target); err != nil { + return nil, err } // Get the prompt using the original prompt name from the backend's perspective. @@ -1079,10 +2173,13 @@ func (h *httpBackendClient) GetPrompt( stringArgs := conversion.ConvertPromptArguments(arguments) + // Forward-compat: mcpcompat drops Params.Meta on this path today (no-op on + // the wire). See TestOutboundMetaTraceContext for details/tripwire. result, err := c.GetPrompt(ctx, mcp.GetPromptRequest{ Params: mcp.GetPromptParams{ Name: backendPromptName, Arguments: stringArgs, + Meta: conversion.ToMCPMeta(telemetry.MetaWithTraceContext(ctx, nil)), }, }) if err != nil { @@ -1095,3 +2192,177 @@ func (h *httpBackendClient) GetPrompt( Meta: conversion.FromMCPMeta(result.Meta), }, nil } + +// Complete requests argument-completion candidates from the backend MCP server, +// selecting the Legacy or Modern path by revision. Returns an empty (non-nil) +// result when the backend does not advertise completions, matching the MCP +// spec's lenient completion semantics. +func (h *httpBackendClient) Complete( + ctx context.Context, + target *vmcp.BackendTarget, + ref vmcp.CompletionRef, + argName, argValue string, + contextArgs map[string]string, +) (*vmcp.CompletionResult, error) { + slog.Debug("requesting completion from backend", "ref_type", ref.Type, "backend", target.WorkloadName) + var out *vmcp.CompletionResult + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernComplete(ctx, target, ref, argName, argValue, contextArgs) + } else { + out, err = h.legacyComplete(ctx, target, ref, argName, argValue, contextArgs) + } + return err + }) + return out, err +} + +// modernComplete invokes completion/complete over the Modern shim (not a +// name-required method, so no Mcp-Name). A prompt ref's name is translated to the +// backend capability name. A backend without completions answers -32601, which is +// treated as an empty result (lenient completion semantics). +func (h *httpBackendClient) modernComplete( + ctx context.Context, target *vmcp.BackendTarget, + ref vmcp.CompletionRef, argName, argValue string, contextArgs map[string]string, +) (*vmcp.CompletionResult, error) { + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + refMap, err := modernCompletionRef(target, ref) + if err != nil { + return nil, err + } + params := map[string]any{ + "ref": refMap, + "argument": map[string]any{"name": argName, "value": argValue}, + } + if len(contextArgs) > 0 { + params["context"] = map[string]any{"arguments": contextArgs} + } + var res struct { + Completion struct { + Values []string `json:"values"` + Total int `json:"total"` + HasMore bool `json:"hasMore"` + } `json:"completion"` + } + err = modernCall(ctx, hc, target.BaseURL, "completion/complete", params, "", nil, &res, "", nil) + if errors.Is(err, mcp.ErrMethodNotFound) { + return &vmcp.CompletionResult{Values: []string{}}, nil + } + if err != nil { + return nil, fmt.Errorf("completion failed on backend %s: %w", target.WorkloadID, err) + } + values := res.Completion.Values + if values == nil { + values = []string{} + } + return &vmcp.CompletionResult{Values: values, Total: res.Completion.Total, HasMore: res.Completion.HasMore}, nil +} + +// modernCompletionRef builds the Modern completion/complete ref params, mirroring +// buildCompletionRef's translation: a prompt ref's name is translated to the +// backend capability name; a resource ref's URI is passed through. +func modernCompletionRef(target *vmcp.BackendTarget, ref vmcp.CompletionRef) (map[string]any, error) { + switch ref.Type { + case vmcp.CompletionRefTypePrompt: + backendName := target.GetBackendCapabilityName(ref.Name) + if backendName != ref.Name { + slog.Debug("translating prompt name for completion", "client_name", ref.Name, "backend_name", backendName) + } + return map[string]any{"type": ref.Type, "name": backendName}, nil + case vmcp.CompletionRefTypeResource: + return map[string]any{"type": ref.Type, "uri": ref.URI}, nil + default: + return nil, fmt.Errorf("%w: unsupported completion ref type %q", vmcp.ErrInvalidInput, ref.Type) + } +} + +// legacyComplete is the initialize + SDK Complete path (unchanged behavior). +func (h *httpBackendClient) legacyComplete( + ctx context.Context, + target *vmcp.BackendTarget, + ref vmcp.CompletionRef, + argName, argValue string, + contextArgs map[string]string, +) (*vmcp.CompletionResult, error) { + // Create a client for this backend + c, err := h.clientFactory(ctx, target, false) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer func() { + if err := c.Close(); err != nil { + slog.Debug("failed to close client", "error", err) + } + }() + + // Initialize the client and capture the backend's advertised capabilities. + serverCaps, _, err := h.legacyInit(ctx, c, target) + if err != nil { + return nil, err + } + + // Backends that do not advertise completions cannot serve completion/complete; + // return an empty result rather than erroring (lenient completion semantics). + if serverCaps.Completions == nil { + slog.Debug("backend does not advertise completions capability, returning empty completion", + "backend", target.WorkloadID) + return &vmcp.CompletionResult{Values: []string{}}, nil + } + + mcpRef, err := buildCompletionRef(target, ref) + if err != nil { + return nil, err + } + + params := mcp.CompleteParams{ + Ref: mcpRef, + Argument: mcp.CompleteArgument{ + Name: argName, + Value: argValue, + }, + } + if len(contextArgs) > 0 { + params.Context = &mcp.CompleteContext{Arguments: contextArgs} + } + + result, err := c.Complete(ctx, mcp.CompleteRequest{Params: params}) + if err != nil { + return nil, fmt.Errorf("completion failed on backend %s: %w", target.WorkloadID, err) + } + + values := result.Completion.Values + if values == nil { + values = []string{} + } + return &vmcp.CompletionResult{ + Values: values, + Total: result.Completion.Total, + HasMore: result.Completion.HasMore, + }, nil +} + +// buildCompletionRef converts a domain CompletionRef into the mcp-go-shaped +// PromptReference/ResourceReference the shim client expects. For a prompt ref the +// name is translated to the backend's own capability name (mirroring GetPrompt); +// for a resource ref the URI is passed through unchanged (mirroring ReadResource, +// whose translation is the identity for non-renamed resources). +func buildCompletionRef(target *vmcp.BackendTarget, ref vmcp.CompletionRef) (any, error) { + switch ref.Type { + case vmcp.CompletionRefTypePrompt: + backendName := target.GetBackendCapabilityName(ref.Name) + if backendName != ref.Name { + slog.Debug("translating prompt name for completion", + "client_name", ref.Name, "backend_name", backendName) + } + return mcp.PromptReference{Type: ref.Type, Name: backendName}, nil + case vmcp.CompletionRefTypeResource: + return mcp.ResourceReference{Type: ref.Type, URI: ref.URI}, nil + default: + return nil, fmt.Errorf("%w: unsupported completion ref type %q", vmcp.ErrInvalidInput, ref.Type) + } +} diff --git a/pkg/vmcp/client/client_test.go b/pkg/vmcp/client/client_test.go index 729ef55b7c..63ae6d3b36 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -13,9 +13,11 @@ import ( "crypto/tls" "crypto/x509" "crypto/x509/pkix" + "encoding/json" "encoding/pem" "errors" "fmt" + "io" "math/big" "net" "net/http" @@ -23,14 +25,12 @@ import ( "os" "path/filepath" "strings" + gosync "sync" "sync/atomic" "syscall" "testing" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/propagation" @@ -38,7 +38,13 @@ import ( "go.opentelemetry.io/otel/trace" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" pkgauth "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/auth" authmocks "github.com/stacklok/toolhive/pkg/vmcp/auth/mocks" @@ -55,7 +61,7 @@ func TestHTTPBackendClient_ListCapabilities_WithMockFactory(t *testing.T) { t.Parallel() expectedErr := errors.New("factory error") - mockFactory := func(_ context.Context, _ *vmcp.BackendTarget) (*client.Client, error) { + mockFactory := func(_ context.Context, _ *vmcp.BackendTarget, _ bool) (*client.Client, error) { return nil, expectedErr } @@ -70,6 +76,11 @@ func TestHTTPBackendClient_ListCapabilities_WithMockFactory(t *testing.T) { TransportType: "streamable-http", } + // Pre-seed the revision cache to Legacy so ListCapabilities skips the + // Modern discover probe (which would need a real transport/registry) and + // goes straight to the injected clientFactory under test. + backendClient.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + capabilities, err := backendClient.ListCapabilities(context.Background(), target) require.Error(t, err) @@ -100,6 +111,17 @@ func TestQueryHelpers_PartialCapabilities(t *testing.T) { assert.Empty(t, result.Resources) }) + t.Run("queryResourceTemplates with unsupported capability returns empty slice", func(t *testing.T) { + t.Parallel() + + // Resource templates are gated on the resources capability; when the backend + // does not advertise resources, the query is skipped and returns empty. + result, err := queryResourceTemplates(context.Background(), nil, false, "test-backend") + require.NoError(t, err) + assert.NotNil(t, result) + assert.Empty(t, result.ResourceTemplates) + }) + t.Run("queryPrompts with unsupported capability returns empty slice", func(t *testing.T) { t.Parallel() @@ -321,7 +343,7 @@ func TestDefaultClientFactory_UnsupportedTransport(t *testing.T) { require.NoError(t, err) httpClient := backendClient.(*httpBackendClient) - _, err = httpClient.defaultClientFactory(context.Background(), target) + _, err = httpClient.defaultClientFactory(context.Background(), target, false) require.Error(t, err) assert.ErrorIs(t, err, vmcp.ErrUnsupportedTransport) @@ -337,7 +359,7 @@ func TestHTTPBackendClient_CallTool_WithMockFactory(t *testing.T) { t.Parallel() expectedErr := errors.New("connection failed") - mockFactory := func(_ context.Context, _ *vmcp.BackendTarget) (*client.Client, error) { + mockFactory := func(_ context.Context, _ *vmcp.BackendTarget, _ bool) (*client.Client, error) { return nil, expectedErr } @@ -352,7 +374,8 @@ func TestHTTPBackendClient_CallTool_WithMockFactory(t *testing.T) { TransportType: "streamable-http", } - result, err := backendClient.CallTool(context.Background(), target, "test_tool", map[string]any{}, nil) + backendClient.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + result, err := backendClient.CallTool(context.Background(), target, "test_tool", map[string]any{}, nil, nil) require.Error(t, err) assert.Nil(t, result) @@ -367,7 +390,7 @@ func TestHTTPBackendClient_ReadResource_WithMockFactory(t *testing.T) { t.Parallel() expectedErr := errors.New("connection failed") - mockFactory := func(_ context.Context, _ *vmcp.BackendTarget) (*client.Client, error) { + mockFactory := func(_ context.Context, _ *vmcp.BackendTarget, _ bool) (*client.Client, error) { return nil, expectedErr } @@ -382,6 +405,7 @@ func TestHTTPBackendClient_ReadResource_WithMockFactory(t *testing.T) { TransportType: "streamable-http", } + backendClient.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) data, err := backendClient.ReadResource(context.Background(), target, "test://resource") require.Error(t, err) @@ -397,7 +421,7 @@ func TestHTTPBackendClient_GetPrompt_WithMockFactory(t *testing.T) { t.Parallel() expectedErr := errors.New("connection failed") - mockFactory := func(_ context.Context, _ *vmcp.BackendTarget) (*client.Client, error) { + mockFactory := func(_ context.Context, _ *vmcp.BackendTarget, _ bool) (*client.Client, error) { return nil, expectedErr } @@ -412,6 +436,7 @@ func TestHTTPBackendClient_GetPrompt_WithMockFactory(t *testing.T) { TransportType: "streamable-http", } + backendClient.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) prompt, err := backendClient.GetPrompt(context.Background(), target, "test_prompt", map[string]any{"arg": "value"}) require.Error(t, err) @@ -1528,3 +1553,390 @@ func ExampleWithDialControl() { } _ = backendClient } + +// TestListCapabilities_MethodNotFoundResourceTemplates verifies FIX: a backend +// that advertises the resources capability but answers resources/templates/list +// with JSON-RPC -32601 (method not found) degrades to an EMPTY template list +// instead of dropping its entire capability set — its concrete resources still +// aggregate. +func TestListCapabilities_MethodNotFoundResourceTemplates(t *testing.T) { + t.Parallel() + + mcpServer := mcpserver.NewMCPServer("partial-backend", "1.0.0") + mcpServer.AddResource( + mcp.Resource{ + URI: "file:///readme.txt", + Name: "Readme", + MIMEType: "text/plain", + }, + func(_ context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + return []mcp.ResourceContents{ + mcp.TextResourceContents{ + URI: req.Params.URI, + MIMEType: "text/plain", + Text: "hello", + }, + }, nil + }, + ) + + // resources/templates/list is not implemented: answer it with a JSON-RPC + // -32601 (method not found) while delegating everything else to the server. + httpHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + rawMessage, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, "Failed to read request", http.StatusBadRequest) + return + } + defer r.Body.Close() + + w.Header().Set("Content-Type", "application/json") + var probe struct { + ID any `json:"id"` + Method string `json:"method"` + } + if err := json.Unmarshal(rawMessage, &probe); err == nil && probe.Method == "resources/templates/list" { + resp := map[string]any{ + "jsonrpc": "2.0", + "id": probe.ID, + "error": map[string]any{"code": -32601, "message": "method not found"}, + } + respBytes, _ := json.Marshal(resp) + _, _ = w.Write(respBytes) + return + } + + response := mcpServer.HandleMessage(r.Context(), rawMessage) + responseBytes, err := json.Marshal(response) + if err != nil { + http.Error(w, "Failed to marshal response", http.StatusInternalServerError) + return + } + _, _ = w.Write(responseBytes) + }) + + server := httptest.NewServer(httpHandler) + defer server.Close() + + registry := auth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, registry.RegisterStrategy("unauthenticated", &strategies.UnauthenticatedStrategy{})) + + backendClient, err := NewHTTPBackendClient(registry) + require.NoError(t, err) + + target := &vmcp.BackendTarget{ + WorkloadID: "partial-backend", + WorkloadName: "Partial Backend", + BaseURL: server.URL, + TransportType: "streamable-http", + } + + caps, err := backendClient.ListCapabilities(t.Context(), target) + require.NoError(t, err, "a -32601 on resources/templates/list must not drop the backend") + require.NotNil(t, caps) + + assert.Empty(t, caps.ResourceTemplates, "templates must degrade to an empty list") + require.Len(t, caps.Resources, 1, "the backend's concrete resources must still aggregate") + assert.Equal(t, "file:///readme.txt", caps.Resources[0].URI) +} + +// TestDefaultClientFactory_SSEForwarding verifies the SSE transport gets the +// same elicitation/sampling forwarding handlers as streamable-http when +// forwarding is requested and forwarders are bound, and that Initialize declares +// the corresponding client capabilities to the backend (without them a +// backend's mid-call elicitation/create or sampling/createMessage is rejected +// instead of relayed). +func TestDefaultClientFactory_SSEForwarding(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + bindForwarders bool + forwarding bool + wantElicitCap bool + wantSamplingCap bool + }{ + { + name: "bound forwarders and forwarding requested declare elicitation and sampling", + bindForwarders: true, + forwarding: true, + wantElicitCap: true, + wantSamplingCap: true, + }, + { + name: "bound forwarders without forwarding request declares neither", + bindForwarders: true, + forwarding: false, + }, + { + name: "unbound forwarders declare neither", + bindForwarders: false, + forwarding: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Minimal SSE server: one event carrying the endpoint, then the + // MCP endpoint answers initialize and reports the client's declared + // capabilities back to the test. + var ( + mu gosync.Mutex + gotElicit bool + gotSampling bool + ) + var mcpEndpoint string + // events carries JSON-RPC responses the /mcp POST handler queues for + // delivery on the SSE stream (legacy SSE transport semantics). + events := make(chan []byte, 16) + mux := http.NewServeMux() + mux.HandleFunc("/mcp", func(w http.ResponseWriter, r *http.Request) { + rawMessage, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, "Failed to read request", http.StatusBadRequest) + return + } + defer r.Body.Close() + var req struct { + Method string `json:"method"` + ID any `json:"id"` + Params struct { + Capabilities struct { + Elicitation map[string]any `json:"elicitation"` + Sampling map[string]any `json:"sampling"` + } `json:"capabilities"` + } `json:"params"` + } + if err := json.Unmarshal(rawMessage, &req); err != nil { + http.Error(w, "Bad request", http.StatusBadRequest) + return + } + if req.Method == "server/discover" { + // go-sdk v1.7's Connect() is Modern-first (SEP-2575): it always + // tries server/discover before falling back to legacy initialize. + // Answer -32601 (method not found) over the SSE message channel so + // the client falls back to initialize below (which this fake + // already handles), instead of hanging waiting for a response that + // never comes. + respBytes, _ := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "error": map[string]any{"code": -32601, "message": "method not found"}, + }) + w.WriteHeader(http.StatusAccepted) + select { + case events <- respBytes: + case <-time.After(5 * time.Second): + // t.Errorf, not t.Fatal: this runs on the server's own goroutine + // (Fatal's runtime.Goexit would kill the wrong one). + t.Errorf("timed out queuing server/discover fallback response for SSE delivery") + } + return + } + if req.Method != "initialize" { + // Notifications (initialized, cancelled) need no response body. + w.WriteHeader(http.StatusAccepted) + return + } + mu.Lock() + gotElicit = req.Params.Capabilities.Elicitation != nil + gotSampling = req.Params.Capabilities.Sampling != nil + mu.Unlock() + + respBytes, _ := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]any{ + "protocolVersion": mcp.LATEST_PROTOCOL_VERSION, + "capabilities": map[string]any{}, + "serverInfo": map[string]any{"name": "sse-backend", "version": "1.0.0"}, + }, + }) + // Per the legacy SSE transport, the server answers a client POST + // with 202 Accepted and delivers the JSON-RPC response as a + // "message" event on the client's SSE stream. + w.WriteHeader(http.StatusAccepted) + select { + case events <- respBytes: + case <-time.After(5 * time.Second): + // t.Errorf, not t.Fatal: this runs on the server's own goroutine. + t.Errorf("timed out queuing initialize response for SSE delivery") + } + }) + mux.HandleFunc("/sse", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", mcpEndpoint) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + for { + select { + case <-r.Context().Done(): + return + case msg := <-events: + _, _ = fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + } + } + }) + httpServer := httptest.NewServer(mux) + defer httpServer.Close() + mcpEndpoint = httpServer.URL + "/mcp" + + registry := auth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, registry.RegisterStrategy("unauthenticated", &strategies.UnauthenticatedStrategy{})) + + backendClient, err := NewHTTPBackendClient(registry) + require.NoError(t, err) + httpClient := backendClient.(*httpBackendClient) + + if tt.bindForwarders { + httpClient.BindForwarders( + &stubElicitationRequester{}, + &stubSamplingRequester{}, + nil, + ) + } + + target := &vmcp.BackendTarget{ + WorkloadID: "sse-backend", + WorkloadName: "SSE Backend", + BaseURL: httpServer.URL + "/sse", + TransportType: "sse", + } + + c, err := httpClient.defaultClientFactory(t.Context(), target, tt.forwarding) + require.NoError(t, err) + defer func() { _ = c.Close() }() + + // The handlers themselves live on unexported mcpcompat Client fields, so + // the observable assertion is the capability declaration: the go-sdk + // auto-declares elicitation/sampling at initialize exactly when the + // corresponding handler is installed. + initCtx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + _, _, err = initializeClient(initCtx, c) + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, tt.wantElicitCap, gotElicit, + "elicitation capability declaration must match handler installation") + assert.Equal(t, tt.wantSamplingCap, gotSampling, + "sampling capability declaration must match handler installation") + }) + } +} + +// stubElicitationRequester is a no-op ElicitationRequester used to bind +// forwarders in client-factory tests; only its non-nil-ness matters. +type stubElicitationRequester struct{} + +func (*stubElicitationRequester) RequestElicitation(context.Context, vmcp.ElicitationRequest) (*vmcp.ElicitationResult, error) { + return nil, errors.New("stub: no downstream session") +} + +// stubSamplingRequester is a no-op SamplingRequester used to bind forwarders in +// client-factory tests; only its non-nil-ness matters. +type stubSamplingRequester struct{} + +func (*stubSamplingRequester) RequestSampling(context.Context, vmcp.SamplingRequest) (*vmcp.SamplingResult, error) { + return nil, errors.New("stub: no downstream session") +} + +// idleSpy is a RoundTripper that records CloseIdleConnections calls. +type idleSpy struct{ closed int } + +func (*idleSpy) RoundTrip(*http.Request) (*http.Response, error) { + return nil, errors.New("unused") +} +func (s *idleSpy) CloseIdleConnections() { s.closed++ } + +// TestModernChain_CloseIdleConnectionsForwards verifies the trace/identity/auth +// wrapper chain forwards CloseIdleConnections down to the concrete transport at +// the bottom (a plain hc.CloseIdleConnections would otherwise be a silent no-op). +func TestModernChain_CloseIdleConnectionsForwards(t *testing.T) { + t.Parallel() + + spy := &idleSpy{} + chain := &tracePropagatingRoundTripper{ + base: &identityPropagatingRoundTripper{ + base: &authRoundTripper{base: spy}, + }, + } + + // Reached the way http.Client.CloseIdleConnections reaches its transport. + chain.CloseIdleConnections() + assert.Equal(t, 1, spy.closed, "CloseIdleConnections must reach the bottom transport") +} + +// TestNewBackendHTTPClient_RedirectPolicy verifies the shared choke point installs +// SameHostRedirectPolicy: a cross-host 30x is refused (so the RoundTripper-injected +// credential never reaches the attacker host), while a same-host redirect is +// followed. +func TestNewBackendHTTPClient_RedirectPolicy(t *testing.T) { + t.Parallel() + + // inject mimics the auth/header-forward chain: it sets a credential on EVERY hop. + inject := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + req.Header.Set("Authorization", "secret") + return http.DefaultTransport.RoundTrip(req) + }) + + t.Run("cross-host redirect refused, no credential leak", func(t *testing.T) { + t.Parallel() + + var attackerSawAuth atomic.Bool + attacker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + attackerSawAuth.Store(true) + } + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(attacker.Close) + + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Redirect(w, &http.Request{}, attacker.URL, http.StatusFound) + })) + t.Cleanup(origin.Close) + + hc := newBackendHTTPClient(inject, 5*time.Second) + resp, err := hc.Get(origin.URL) //nolint:noctx // test + if resp != nil { + _ = resp.Body.Close() + } + require.Error(t, err) + assert.ErrorIs(t, err, networking.ErrRedirectRefused) + assert.False(t, attackerSawAuth.Load(), "credential must not reach the cross-host redirect target") + }) + + t.Run("same-host redirect followed", func(t *testing.T) { + t.Parallel() + + var landed atomic.Bool + mux := http.NewServeMux() + mux.HandleFunc("/start", func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "/end", http.StatusFound) + }) + mux.HandleFunc("/end", func(w http.ResponseWriter, _ *http.Request) { + landed.Store(true) + w.WriteHeader(http.StatusOK) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + hc := newBackendHTTPClient(inject, 5*time.Second) + resp, err := hc.Get(srv.URL + "/start") //nolint:noctx // test + require.NoError(t, err) + _ = resp.Body.Close() + assert.True(t, landed.Load(), "same-host redirect must be followed") + }) +} diff --git a/pkg/vmcp/client/dispatch_legacy_fallback_test.go b/pkg/vmcp/client/dispatch_legacy_fallback_test.go new file mode 100644 index 0000000000..a22c503a49 --- /dev/null +++ b/pkg/vmcp/client/dispatch_legacy_fallback_test.go @@ -0,0 +1,94 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestDispatch_LegacyFallbackOnInconclusiveProbe pins the fix for the +// dispatch abort regressed by merge 885e8b14f: a Modern server/discover probe +// that comes back transient (HTTP 503, -> errModernTransient) must fall back +// to the Legacy initialize/tools-list path uncached, not fail the whole call. +// Caching the fallback would pin a genuinely Modern backend to Legacy past a +// transient outage, so the revision cache must remain unset afterward. +func TestDispatch_LegacyFallbackOnInconclusiveProbe(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Mcp-Method") != "" { + // Modern probe: transient failure, discovered before the body is read. + w.WriteHeader(http.StatusServiceUnavailable) + return + } + + body, err := readAll(t, r) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + var req jsonRPCRequest + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + w.Header().Set("Content-Type", "application/json") + switch req.Method { + case "initialize": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "flaky-probe", "version": "1.0.0"} + }`), + } + _ = json.NewEncoder(w).Encode(resp) + case "tools/list": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "tools": [{"name": "echo", "description": "echoes input"}] + }`), + } + _ = json.NewEncoder(w).Encode(resp) + default: + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{}`), + }) + } + })) + t.Cleanup(srv.Close) + + // newProbeClient, not setRevision: the probe must actually run so the + // Modern 503 is hit and the fallback path exercised. + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "flaky-probe", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + caps, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.Len(t, caps.Tools, 1) + assert.Equal(t, "echo", caps.Tools[0].Name) + + _, ok := h.cachedRevision("flaky-probe") + assert.False(t, ok, "an inconclusive probe must not pin a revision") +} diff --git a/pkg/vmcp/client/forwarding.go b/pkg/vmcp/client/forwarding.go new file mode 100644 index 0000000000..cffffde1b2 --- /dev/null +++ b/pkg/vmcp/client/forwarding.go @@ -0,0 +1,386 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "log/slog" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/conversion" +) + +// boundForwarders holds the server->client forwarding requesters bound onto the +// backend client after the SDK server is constructed. A nil field leaves that +// forwarding path disabled. +// +// Stored atomically on httpBackendClient so the client factory (which may run +// concurrently for different backends) reads a consistent snapshot without a +// mutex. Bound exactly once, before serving begins. +type boundForwarders struct { + elicitation vmcp.ElicitationRequester + sampling vmcp.SamplingRequester + notifier vmcp.ClientNotifier +} + +// BindForwarders implements vmcp.ClientForwarderBinder. server.New calls it once, +// after Serve builds the mcp-go server the adapters wrap and before serving +// begins, so the backend clients created per call can relay a backend's mid-call +// server->client traffic (elicitation, sampling, progress/logging) to the +// downstream client. +func (h *httpBackendClient) BindForwarders( + elicitation vmcp.ElicitationRequester, + sampling vmcp.SamplingRequester, + notifier vmcp.ClientNotifier, +) { + h.forwarders.Store(&boundForwarders{ + elicitation: elicitation, + sampling: sampling, + notifier: notifier, + }) +} + +// enableBackendLogging requests debug-level logging from the backend so it emits +// notifications/message during a tool call, which the notification forwarder +// relays to the downstream client. It is a no-op when no forwarders are bound or +// the backend does not advertise the logging capability. Best-effort: a failure +// is logged at debug and does not fail the caller. +// +// This RPC (logging/setLevel) exists only on the LEGACY revision: the 2026-07-28 +// (Modern) revision removed it — go-sdk#1116 was closed wont-fix-by-design — so +// this is only ever called from legacyCallTool. The Modern equivalent is the +// per-request io.modelcontextprotocol/logLevel _meta key, which modernCallTool +// overlays onto the request it mints (see TestModernCallTool_LogLevelGating). +func (h *httpBackendClient) enableBackendLogging( + ctx context.Context, c *client.Client, caps *mcp.ServerCapabilities, + negotiatedVersion, backendID string, +) { + if h.forwarders.Load() == nil { + return + } + if caps == nil || caps.Logging == nil { + return + } + // The 2026-07-28 revision REMOVED logging/setLevel; the per-request + // io.modelcontextprotocol/logLevel _meta key replaces it (see + // mcpparser.MetaKeyLogLevel, which the Modern path already mints). + // + // A dual-era backend can negotiate 2026-07-28 over this Legacy handshake, + // and the session then MUST carry that version on every request. go-sdk + // still sends setLevel as a Legacy-shaped call, so the request arrives with + // a Modern protocol header and no Modern _meta -- a shape ToolHive's + // classifier rejects with -32020, deliberately and by a pinned contract + // (see TestIntegration_Modern_RealBackend_LoggingContract). go-sdk treats + // that rejection as fatal, closing the session and failing the tool call + // this logging was only meant to decorate. + // + // Calling an RPC the negotiated revision removed is the actual defect, so + // skip it. Backends that negotiate a Legacy version are unaffected. + if negotiatedVersion == mcpparser.MCPVersionModern { + slog.DebugContext(ctx, "skipping logging/setLevel: removed in the negotiated revision", + "backend", backendID, "negotiated", negotiatedVersion) + return + } + if err := c.SetLoggingLevel(ctx, mcp.LoggingLevelDebug); err != nil { + slog.Debug("failed to set backend logging level", "backend", backendID, "error", err) + } +} + +// drainServerToClientNotifications flushes any in-flight backend->downstream +// notification off the per-call backend client before it is closed, closing a +// lost-notification race that otherwise drops fire-and-forget notifications +// (notifications/progress, notifications/message) under load. +// +// The race: a backend emits such a notification mid tools/call and then returns +// the result. The notification and the result travel on SEPARATE streams — the +// notification on the standalone SSE stream, the result on the tools/call +// response stream — and the client reads every stream through a single FIFO +// channel feeding one receive loop. CallTool returns as soon as the RESULT is +// read; its deferred Close then cancels the receive loop. If the standalone +// stream's notification has not yet been read and enqueued for handling when +// Close fires, it is discarded and never reaches newNotificationForwarder, so +// the downstream client never sees it (a permanent loss, not a slow delivery). +// Blocking server->client REQUESTS (elicitation/sampling) are immune: the +// backend tool blocks on the response, so the client is guaranteed to still be +// reading when they arrive. +// +// The fix is a synchronous ping used purely as a drain barrier. The backend +// flushes the notification onto the wire before returning the result, so by the +// time CallTool returns the notification bytes are already buffered on the +// client's standalone connection. A ping is a full backend round-trip; while it +// is in flight the receive loop drains the buffered notification off the shared +// FIFO channel (a local buffered read completes far ahead of the ping's network +// round-trip), enqueuing it for handling. Because the ping response arrives on +// that same FIFO channel AFTER the notification, the ping returning proves the +// notification was already read and enqueued. The subsequent Close then waits +// for enqueued handlers to finish (jsonrpc2.Connection.Close drains the handler +// queue), so the forward completes before teardown. +// +// Only runs when server->client forwarding is bound (fwd.notifier != nil, the +// same condition under which the notification forwarder is registered); other +// deployments keep the pre-forwarding fast path with no extra round-trip. +// Best-effort: a ping failure only forgoes the barrier — the tool result is +// already in hand, so it must not fail the call. +func (h *httpBackendClient) drainServerToClientNotifications(ctx context.Context, c *client.Client) { + fwd := h.forwarders.Load() + if fwd == nil || fwd.notifier == nil { + return + } + if err := c.Ping(ctx); err != nil { + slog.Debug("notification drain ping failed; forwarded notifications may be lost", "error", err) + } +} + +// forwardingClientOptions builds the client-level options that install the +// elicitation and sampling handlers on a backend client, each closing over the +// captured per-call downstream context so the handler can relay to the right +// session. Only the handlers whose requester is bound are installed. The handlers +// must be registered before Initialize, so these options are passed at +// construction (NewStreamableHttpClientWithOpts). +func forwardingClientOptions(callCtx context.Context, fwd *boundForwarders) []client.ClientOption { + var opts []client.ClientOption + if fwd.elicitation != nil { + opts = append(opts, client.WithElicitationHandler(newElicitationForwarder(callCtx, fwd.elicitation))) + } + if fwd.sampling != nil { + opts = append(opts, client.WithSamplingHandler(newSamplingForwarder(callCtx, fwd.sampling))) + } + return opts +} + +// deriveForwardCtx derives a context from base (the captured per-call downstream +// session context, which the SDK adapters need to route server->client traffic +// to the right session) that is also cancelled when handler (the receive-loop +// context the go-sdk invokes the client handler with) is cancelled. This lets a +// forwarded elicitation/sampling round-trip honor a backend-side cancellation +// while still resolving the downstream session from base. +// +// The returned cancel func must be called to release the AfterFunc watcher. +func deriveForwardCtx(base, handler context.Context) (context.Context, context.CancelFunc) { + ctx, cancel := context.WithCancel(base) + stop := context.AfterFunc(handler, cancel) + return ctx, func() { + stop() + cancel() + } +} + +// newElicitationForwarder builds a client elicitation handler that relays a +// backend's mid-call elicitation/create request to the downstream client via the +// bound requester, using the captured per-call downstream context (callCtx). The +// go-sdk invokes the handler on its receive loop with a context that does NOT +// carry the downstream session, so callCtx — not the handler's own ctx — is the +// one that resolves the session. +// +// When callCtx carries no downstream session (health probes, capability +// listing) the requester returns an error, which is relayed back to the backend +// as a clean elicitation failure rather than hanging. +func newElicitationForwarder(callCtx context.Context, req vmcp.ElicitationRequester) client.ElicitationHandlerFunc { + return func(handlerCtx context.Context, r mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + ctx, cancel := deriveForwardCtx(callCtx, handlerCtx) + defer cancel() + + res, err := req.RequestElicitation(ctx, vmcp.ElicitationRequest{ + Message: r.Params.Message, + RequestedSchema: r.Params.RequestedSchema, + Meta: conversion.FromMCPMeta(r.Params.Meta), + }) + if err != nil { + return nil, err + } + return &mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{ + Action: mcp.ElicitationResponseAction(res.Action), + Content: res.Content, + }, + }, nil + } +} + +// newSamplingForwarder builds a client sampling handler that relays a backend's +// mid-call sampling/createMessage request to the downstream client via the bound +// requester, using the captured per-call downstream context. See +// newElicitationForwarder for the captured-context rationale. +func newSamplingForwarder(callCtx context.Context, req vmcp.SamplingRequester) client.SamplingHandlerFunc { + return func(handlerCtx context.Context, r mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) { + ctx, cancel := deriveForwardCtx(callCtx, handlerCtx) + defer cancel() + + res, err := req.RequestSampling(ctx, vmcp.SamplingRequest{ + Messages: fromMCPSamplingMessages(r.Messages), + ModelPreferences: fromMCPModelPreferences(r.ModelPreferences), + SystemPrompt: r.SystemPrompt, + IncludeContext: r.IncludeContext, + Temperature: r.Temperature, + MaxTokens: r.MaxTokens, + StopSequences: r.StopSequences, + Metadata: r.Metadata, + }) + if err != nil { + return nil, err + } + return &mcp.CreateMessageResult{ + SamplingMessage: mcp.SamplingMessage{ + Role: mcp.Role(res.Role), + Content: res.Content, + }, + Model: res.Model, + StopReason: res.StopReason, + }, nil + } +} + +// newNotificationForwarder builds an OnNotification handler that relays a +// backend's mid-call notifications/progress and notifications/message to the +// downstream client via the bound notifier, using the captured per-call +// downstream context. Forwarding is best-effort: a missing downstream session +// (the notifier no-ops) or a forwarding error is logged at debug and dropped, +// never surfaced to the backend. Other notification methods are ignored (the +// go-sdk server re-emits list-changed notifications automatically). +func newNotificationForwarder(callCtx context.Context, notifier vmcp.ClientNotifier) func(mcp.JSONRPCNotification) { + // Backend notifications are delivered asynchronously and can arrive just after + // the tool call context is cancelled; keep the captured downstream-session + // values but ignore cancellation so best-effort forwarding still runs. + forwardCtx := context.WithoutCancel(callCtx) + + return func(n mcp.JSONRPCNotification) { + fields := n.Params.AdditionalFields + switch n.Method { + case vmcp.MethodProgressNotification: + err := notifier.NotifyProgress(forwardCtx, vmcp.ProgressNotification{ + ProgressToken: fields["progressToken"], + Progress: toFloat(fields["progress"]), + Total: toFloat(fields["total"]), + Message: toString(fields["message"]), + }) + if err != nil { + slog.Debug("failed to forward progress notification", "error", err) + } + case vmcp.MethodLogNotification: + err := notifier.NotifyLog(forwardCtx, vmcp.LogMessage{ + Level: toString(fields["level"]), + Logger: toString(fields["logger"]), + Data: fields["data"], + }) + if err != nil { + slog.Debug("failed to forward log notification", "error", err) + } + default: + // Other notifications (list_changed, resources/updated, ...) are not + // mid-call server->client traffic this forwarder relays. + } + } +} + +// newModernNotificationForwarder builds the onNotification callback modernCall +// invokes for each server->client notification a Modern (2026-07-28) backend +// interleaves on the SSE stream ahead of the tool-call response. It is the +// Modern analogue of newNotificationForwarder: the wire shape differs (raw +// method + params instead of a typed mcp.JSONRPCNotification), but the relay is +// the same best-effort NotifyLog/NotifyProgress to the downstream client via the +// bound notifier, using the captured per-call downstream context. Only the two +// mid-call methods vMCP opts into (log via the logLevel _meta key, progress via +// a progressToken) are relayed; anything else is ignored. +func newModernNotificationForwarder( + callCtx context.Context, notifier vmcp.ClientNotifier, +) func(method string, params json.RawMessage) { + // Same WithoutCancel rationale as newNotificationForwarder: notifications can + // arrive just after the tool call context is cancelled, but forwarding should + // still run best-effort. + forwardCtx := context.WithoutCancel(callCtx) + + return func(method string, params json.RawMessage) { + switch method { + case vmcp.MethodProgressNotification, vmcp.MethodLogNotification: + default: + return // not a mid-call method this forwarder relays + } + var fields map[string]any + if err := json.Unmarshal(params, &fields); err != nil { + slog.Debug("failed to decode modern notification params", "method", method, "error", err) + return + } + if method == vmcp.MethodProgressNotification { + if err := notifier.NotifyProgress(forwardCtx, vmcp.ProgressNotification{ + ProgressToken: fields["progressToken"], + Progress: toFloat(fields["progress"]), + Total: toFloat(fields["total"]), + Message: toString(fields["message"]), + }); err != nil { + slog.Debug("failed to forward progress notification", "error", err) + } + return + } + if err := notifier.NotifyLog(forwardCtx, vmcp.LogMessage{ + Level: toString(fields["level"]), + Logger: toString(fields["logger"]), + Data: fields["data"], + }); err != nil { + slog.Debug("failed to forward log notification", "error", err) + } + } +} + +// fromMCPSamplingMessages maps SDK sampling messages to the domain type. +func fromMCPSamplingMessages(msgs []mcp.SamplingMessage) []vmcp.SamplingMessage { + if msgs == nil { + return nil + } + out := make([]vmcp.SamplingMessage, len(msgs)) + for i, m := range msgs { + out[i] = vmcp.SamplingMessage{ + Role: string(m.Role), + Content: m.Content, + } + } + return out +} + +// fromMCPModelPreferences maps the SDK model preferences to the domain type, +// preserving a nil (unset) value. +func fromMCPModelPreferences(p *mcp.ModelPreferences) *vmcp.ModelPreferences { + if p == nil { + return nil + } + hints := make([]vmcp.ModelHint, len(p.Hints)) + for i, h := range p.Hints { + hints[i] = vmcp.ModelHint{Name: h.Name} + } + return &vmcp.ModelPreferences{ + Hints: hints, + CostPriority: p.CostPriority, + SpeedPriority: p.SpeedPriority, + IntelligencePriority: p.IntelligencePriority, + } +} + +// toFloat coerces a JSON-decoded numeric value to float64, returning 0 for a +// missing or non-numeric value. +func toFloat(v any) float64 { + switch n := v.(type) { + case float64: + return n + case float32: + return float64(n) + case int: + return float64(n) + case int64: + return float64(n) + default: + return 0 + } +} + +// toString coerces a JSON-decoded value to string, returning "" when absent or +// not a string. +func toString(v any) string { + if s, ok := v.(string); ok { + return s + } + return "" +} diff --git a/pkg/vmcp/client/forwarding_test.go b/pkg/vmcp/client/forwarding_test.go new file mode 100644 index 0000000000..b4c142222c --- /dev/null +++ b/pkg/vmcp/client/forwarding_test.go @@ -0,0 +1,303 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" +) + +// ctxKey is a private context key used to prove that the forwarder relays with +// the captured per-call downstream context (callCtx), not the receive-loop +// handler context the go-sdk supplies. +type ctxKey struct{} + +func TestNewElicitationForwarder_ForwardsWithCapturedCtx(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + req := mocks.NewMockElicitationRequester(ctrl) + + callCtx := context.WithValue(t.Context(), ctxKey{}, "downstream-session") + + req.EXPECT(). + RequestElicitation(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, r vmcp.ElicitationRequest) (*vmcp.ElicitationResult, error) { + // The captured downstream ctx is used, not the handler ctx. + assert.Equal(t, "downstream-session", ctx.Value(ctxKey{})) + // SDK -> domain conversion carried the request fields. + assert.Equal(t, "Confirm?", r.Message) + assert.Equal(t, map[string]any{"type": "object"}, r.RequestedSchema) + return &vmcp.ElicitationResult{Action: "accept", Content: map[string]any{"ok": true}}, nil + }) + + handler := newElicitationForwarder(callCtx, req) + res, err := handler.Elicit(t.Context(), mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{ + Message: "Confirm?", + RequestedSchema: map[string]any{"type": "object"}, + }, + }) + + require.NoError(t, err) + require.NotNil(t, res) + // domain -> SDK result conversion. + assert.Equal(t, mcp.ElicitationResponseActionAccept, res.Action) + assert.Equal(t, map[string]any{"ok": true}, res.Content) +} + +func TestNewElicitationForwarder_PropagatesError(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + req := mocks.NewMockElicitationRequester(ctrl) + req.EXPECT(). + RequestElicitation(gomock.Any(), gomock.Any()). + Return(nil, errors.New("no active session")) + + handler := newElicitationForwarder(t.Context(), req) + res, err := handler.Elicit(t.Context(), mcp.ElicitationRequest{}) + require.Error(t, err) + assert.Nil(t, res) +} + +func TestNewSamplingForwarder_ForwardsWithCapturedCtx(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + req := mocks.NewMockSamplingRequester(ctrl) + + callCtx := context.WithValue(t.Context(), ctxKey{}, "downstream-session") + + req.EXPECT(). + RequestSampling(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, r vmcp.SamplingRequest) (*vmcp.SamplingResult, error) { + assert.Equal(t, "downstream-session", ctx.Value(ctxKey{})) + require.Len(t, r.Messages, 1) + assert.Equal(t, "user", r.Messages[0].Role) + assert.Equal(t, 128, r.MaxTokens) + return &vmcp.SamplingResult{ + Role: "assistant", + Content: map[string]any{"type": "text", "text": "hi"}, + Model: "m", + StopReason: "endTurn", + }, nil + }) + + handler := newSamplingForwarder(callCtx, req) + res, err := handler.CreateMessage(t.Context(), mcp.CreateMessageRequest{ + CreateMessageParams: mcp.CreateMessageParams{ + Messages: []mcp.SamplingMessage{ + {Role: mcp.Role("user"), Content: map[string]any{"type": "text", "text": "hello"}}, + }, + MaxTokens: 128, + }, + }) + + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, mcp.Role("assistant"), res.Role) + assert.Equal(t, "m", res.Model) + assert.Equal(t, "endTurn", res.StopReason) +} + +func TestNewSamplingForwarder_PropagatesError(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + req := mocks.NewMockSamplingRequester(ctrl) + req.EXPECT(). + RequestSampling(gomock.Any(), gomock.Any()). + Return(nil, errors.New("session does not support sampling")) + + handler := newSamplingForwarder(t.Context(), req) + res, err := handler.CreateMessage(t.Context(), mcp.CreateMessageRequest{}) + require.Error(t, err) + assert.Nil(t, res) +} + +func TestNewNotificationForwarder_ForwardsProgress(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + notifier := mocks.NewMockClientNotifier(ctrl) + + callCtx := context.WithValue(t.Context(), ctxKey{}, "downstream-session") + + notifier.EXPECT(). + NotifyProgress(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, n vmcp.ProgressNotification) error { + assert.Equal(t, "downstream-session", ctx.Value(ctxKey{})) + assert.Equal(t, "tok-1", n.ProgressToken) + assert.Equal(t, 0.5, n.Progress) + assert.Equal(t, 1.0, n.Total) + assert.Equal(t, "halfway", n.Message) + return nil + }) + + handler := newNotificationForwarder(callCtx, notifier) + handler(progressNotification("tok-1", 0.5, 1.0, "halfway")) +} + +func TestNewNotificationForwarder_ForwardsLog(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + notifier := mocks.NewMockClientNotifier(ctrl) + + notifier.EXPECT(). + NotifyLog(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, n vmcp.LogMessage) error { + assert.Equal(t, "info", n.Level) + assert.Equal(t, "backend", n.Logger) + assert.Equal(t, "hello", n.Data) + return nil + }) + + handler := newNotificationForwarder(t.Context(), notifier) + handler(mcp.JSONRPCNotification{ + Notification: mcp.Notification{ + Method: vmcp.MethodLogNotification, + Params: mcp.NotificationParams{ + AdditionalFields: map[string]any{"level": "info", "logger": "backend", "data": "hello"}, + }, + }, + }) +} + +func TestNewNotificationForwarder_UsesWithoutCancelContext(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + notifier := mocks.NewMockClientNotifier(ctrl) + + baseCtx, cancel := context.WithCancel(context.WithValue(t.Context(), ctxKey{}, "downstream-session")) + cancel() + + notifier.EXPECT(). + NotifyLog(gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, _ vmcp.LogMessage) error { + assert.Equal(t, "downstream-session", ctx.Value(ctxKey{})) + assert.NoError(t, ctx.Err(), "notification forwarding should ignore cancellation on callCtx") + return nil + }) + + handler := newNotificationForwarder(baseCtx, notifier) + handler(mcp.JSONRPCNotification{ + Notification: mcp.Notification{ + Method: vmcp.MethodLogNotification, + Params: mcp.NotificationParams{ + AdditionalFields: map[string]any{"level": "info", "data": "hello"}, + }, + }, + }) +} + +// TestNewNotificationForwarder_IgnoresUnknownMethod verifies that a notification +// method the forwarder does not relay does not reach the notifier. +func TestNewNotificationForwarder_IgnoresUnknownMethod(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + notifier := mocks.NewMockClientNotifier(ctrl) + // No EXPECT calls: any invocation fails the test. + + handler := newNotificationForwarder(t.Context(), notifier) + handler(mcp.JSONRPCNotification{ + Notification: mcp.Notification{Method: "notifications/tools/list_changed"}, + }) +} + +// TestNewNotificationForwarder_SwallowsError verifies that a notifier error +// (e.g. best-effort no-op returning an error) does not panic the receive loop. +func TestNewNotificationForwarder_SwallowsError(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + notifier := mocks.NewMockClientNotifier(ctrl) + notifier.EXPECT(). + NotifyProgress(gomock.Any(), gomock.Any()). + Return(errors.New("transport closed")) + + handler := newNotificationForwarder(t.Context(), notifier) + assert.NotPanics(t, func() { + handler(progressNotification("tok", 1, 0, "")) + }) +} + +func TestBindForwarders_StoresSnapshot(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + elicit := mocks.NewMockElicitationRequester(ctrl) + sampling := mocks.NewMockSamplingRequester(ctrl) + notifier := mocks.NewMockClientNotifier(ctrl) + + h := &httpBackendClient{} + assert.Nil(t, h.forwarders.Load()) + + h.BindForwarders(elicit, sampling, notifier) + + fwd := h.forwarders.Load() + require.NotNil(t, fwd) + assert.Same(t, elicit, fwd.elicitation) + assert.Same(t, sampling, fwd.sampling) + assert.Same(t, notifier, fwd.notifier) +} + +// TestDeriveForwardCtx_CancelsOnHandlerCancel verifies that the derived context +// (rooted at the captured downstream ctx) is cancelled when the handler ctx is +// cancelled, so a backend-side cancellation aborts the forwarded round-trip. +func TestDeriveForwardCtx_CancelsOnHandlerCancel(t *testing.T) { + t.Parallel() + + base := context.WithValue(t.Context(), ctxKey{}, "downstream-session") + handlerCtx, cancelHandler := context.WithCancel(t.Context()) + + ctx, cancel := deriveForwardCtx(base, handlerCtx) + defer cancel() + + // The derived ctx still resolves the downstream session value from base. + assert.Equal(t, "downstream-session", ctx.Value(ctxKey{})) + + select { + case <-ctx.Done(): + t.Fatal("derived ctx cancelled prematurely") + default: + } + + cancelHandler() + select { + case <-ctx.Done(): + case <-t.Context().Done(): + t.Fatal("derived ctx not cancelled after handler cancellation") + } +} + +// progressNotification builds an mcp.JSONRPCNotification for notifications/progress +// with the wire fields the go-sdk client surfaces via OnNotification. +func progressNotification(token any, progress, total float64, message string) mcp.JSONRPCNotification { + fields := map[string]any{"progressToken": token, "progress": progress} + if total != 0 { + fields["total"] = total + } + if message != "" { + fields["message"] = message + } + return mcp.JSONRPCNotification{ + Notification: mcp.Notification{ + Method: vmcp.MethodProgressNotification, + Params: mcp.NotificationParams{AdditionalFields: fields}, + }, + } +} diff --git a/pkg/vmcp/client/header_forward_integration_test.go b/pkg/vmcp/client/header_forward_integration_test.go index 45dacc88ba..41caa2f12d 100644 --- a/pkg/vmcp/client/header_forward_integration_test.go +++ b/pkg/vmcp/client/header_forward_integration_test.go @@ -73,7 +73,7 @@ func TestHeaderForward_EndToEnd_ThroughHTTPBackendClient(t *testing.T) { // CallTool drives the full Initialize → tools/list flow through the // streamable-HTTP transport. We don't care about the call result — only // that the request reached the test server with the configured headers. - _, _ = backendClient.CallTool(ctx, target, "anything", map[string]any{}, nil) + _, _ = backendClient.CallTool(ctx, target, "anything", map[string]any{}, nil, nil) captured.mu.Lock() defer captured.mu.Unlock() diff --git a/pkg/vmcp/client/meta_integration_test.go b/pkg/vmcp/client/meta_integration_test.go index 7c65d4d7f0..8e65b9e992 100644 --- a/pkg/vmcp/client/meta_integration_test.go +++ b/pkg/vmcp/client/meta_integration_test.go @@ -10,14 +10,18 @@ import ( "io" "net" "net/http" + "sync" "testing" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/propagation" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/auth" "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" @@ -29,7 +33,7 @@ func TestMetaPreservation_CallTool(t *testing.T) { t.Parallel() // Create and start a real MCP server that returns _meta - port, cleanup := startTestMCPServer(t) + port, _, cleanup := startTestMCPServer(t) defer cleanup() // Create vMCP backend client with unauthenticated strategy @@ -54,7 +58,7 @@ func TestMetaPreservation_CallTool(t *testing.T) { // Call tool through vMCP backend client result, err := backendClient.CallTool(ctx, target, "test_tool_with_meta", map[string]any{ "input": "test-value", - }, nil) + }, nil, nil) // Verify call succeeded require.NoError(t, err) @@ -77,7 +81,7 @@ func TestMetaPreservation_CallTool(t *testing.T) { func TestMetaPreservation_CallTool_NoMeta(t *testing.T) { t.Parallel() - port, cleanup := startTestMCPServer(t) + port, _, cleanup := startTestMCPServer(t) defer cleanup() registry := auth.NewDefaultOutgoingAuthRegistry() @@ -100,7 +104,7 @@ func TestMetaPreservation_CallTool_NoMeta(t *testing.T) { // Call tool that doesn't return _meta result, err := backendClient.CallTool(ctx, target, "test_tool_no_meta", map[string]any{ "input": "test-value", - }, nil) + }, nil, nil) require.NoError(t, err) require.NotNil(t, result) @@ -124,7 +128,7 @@ func TestMetaPreservation_CallTool_NoMeta(t *testing.T) { func TestMetaPreservation_CallTool_Error(t *testing.T) { t.Parallel() - port, cleanup := startTestMCPServer(t) + port, _, cleanup := startTestMCPServer(t) defer cleanup() registry := auth.NewDefaultOutgoingAuthRegistry() @@ -147,7 +151,7 @@ func TestMetaPreservation_CallTool_Error(t *testing.T) { // Call tool that returns an error with _meta result, err := backendClient.CallTool(ctx, target, "test_tool_error", map[string]any{ "input": "trigger-error", - }, nil) + }, nil, nil) // Should return result (not a Go error) when tool returns IsError=true require.NoError(t, err, "IsError=true is not a transport error, should return result") @@ -171,7 +175,7 @@ func TestMetaPreservation_CallTool_Error(t *testing.T) { func TestMetaPreservation_GetPrompt(t *testing.T) { t.Parallel() - port, cleanup := startTestMCPServer(t) + port, _, cleanup := startTestMCPServer(t) defer cleanup() registry := auth.NewDefaultOutgoingAuthRegistry() @@ -210,24 +214,18 @@ func TestMetaPreservation_GetPrompt(t *testing.T) { assert.Equal(t, "Hello, World!", result.Messages[0].Content.Text) } -// TestMetaPreservation_ReadResource documents the SDK limitation for resource _meta. +// TestMetaPreservation_ReadResource verifies that backend resource-read _meta +// survives the vMCP backend client: the test backend registers its resource via +// mcpcompat's result-returning AddResourceWithResult (toolhive-core#194), and +// the client must extract the meta on both the Legacy and Modern paths. // -// KNOWN LIMITATION: Due to MCP SDK constraints, resource handlers return []ResourceContents -// directly, not *ReadResourceResult with _meta. This prevents backends from including _meta -// in resource responses at all. -// -// As a result: -// - Backend MCP servers cannot include _meta in resource read responses (SDK limitation) -// - vMCP client cannot extract _meta because it's not in the response -// - vMCP handler cannot forward _meta to clients -// -// This test documents the expected behavior and ensures resource reads work correctly -// even though _meta is not supported. Once the SDK adds _meta support for resource handlers, -// this test can be updated to verify _meta preservation. +// This used to document a "KNOWN LIMITATION" blaming MCP SDK constraints; that +// framing was stale — go-sdk always supported the field, and the loss was +// mcpcompat's mcp-go-shaped contents-only handler signature, fixed upstream. func TestMetaPreservation_ReadResource(t *testing.T) { t.Parallel() - port, cleanup := startTestMCPServer(t) + port, _, cleanup := startTestMCPServer(t) defer cleanup() registry := auth.NewDefaultOutgoingAuthRegistry() @@ -253,9 +251,10 @@ func TestMetaPreservation_ReadResource(t *testing.T) { require.NoError(t, err) require.NotNil(t, result) - // Verify _meta is NOT present due to SDK limitation - // The SDK handler signature doesn't support returning _meta with resources - assert.Nil(t, result.Meta, "_meta cannot be included due to SDK limitation - handler returns []ResourceContents without _meta wrapper") + // Verify _meta is preserved end to end + require.NotNil(t, result.Meta, "backend resource _meta must be preserved") + assert.Equal(t, "resource-trace-789", result.Meta["traceId"]) + assert.Equal(t, "resource-custom", result.Meta["customTag"]) // Verify resource content works correctly require.NotEmpty(t, result.Contents) @@ -263,11 +262,152 @@ func TestMetaPreservation_ReadResource(t *testing.T) { assert.Equal(t, "text/plain", result.Contents[0].MimeType) } +// TestOutboundMetaTraceContext verifies that the vMCP backend client injects the +// current W3C trace context (traceparent) into outbound params._meta (SEP-414) +// for CallTool, and that it is omitted entirely when there is no active span. +// +// ReadResource and GetPrompt are NOT asserted for on-the-wire traceparent +// delivery here: github.com/stacklok/toolhive-core's mcpcompat client (as of +// v0.0.32) builds the underlying go-sdk resources/read and prompts/get +// requests without forwarding request.Params.Meta at all (only CallTool's +// non-resume path converts and forwards Meta - see mcpcompat/client/client.go +// CallTool vs ReadResource/GetPrompt). This means the Meta we set on +// mcp.ReadResourceParams/mcp.GetPromptParams in client.go and mcp_session.go +// is silently dropped before it reaches the wire for these two operations +// today. The code changes are still correct and forward-compatible: they will +// start working the moment mcpcompat forwards Meta the same way it already +// does for CallTool, with no further change needed on our side. This is +// analogous to the inbound resource-_meta drop that existed before +// mcpcompat's result-returning resource handler (toolhive-core#211). +// +// This test mutates the global OTEL propagator, so it must NOT run in +// parallel with the other tests in this file (see propagation_test.go for the +// same convention in pkg/telemetry). +func TestOutboundMetaTraceContext(t *testing.T) { //nolint:paralleltest // Mutates global OTEL propagator + oldPropagator := otel.GetTextMapPropagator() + otel.SetTextMapPropagator(propagation.TraceContext{}) + defer otel.SetTextMapPropagator(oldPropagator) + + port, capture, cleanup := startTestMCPServer(t) + defer cleanup() + + registry := auth.NewDefaultOutgoingAuthRegistry() + err := registry.RegisterStrategy("unauthenticated", &strategies.UnauthenticatedStrategy{}) + require.NoError(t, err) + + backendClient, err := vmcpclient.NewHTTPBackendClient(registry) + require.NoError(t, err) + + target := &vmcp.BackendTarget{ + WorkloadID: "test-backend", + WorkloadName: "Test Backend", + BaseURL: "http://127.0.0.1:" + port, + TransportType: "streamable-http", + } + + tp := sdktrace.NewTracerProvider() + defer func() { _ = tp.Shutdown(context.Background()) }() + tracer := tp.Tracer("test") + + //nolint:paralleltest // sequential: shares capture/backendClient state with the sibling subtests below + t.Run("CallTool injects traceparent", func(t *testing.T) { + spanCtx, span := tracer.Start(context.Background(), "test-span") + ctx, cancel := context.WithTimeout(spanCtx, 5*time.Second) + defer cancel() + defer span.End() + + _, err := backendClient.CallTool(ctx, target, "test_tool_capture_meta", nil, nil, nil) + require.NoError(t, err) + + got := capture.get("tool") + require.NotNil(t, got, "expected params._meta to be sent") + traceparent, ok := got.AdditionalFields["traceparent"].(string) + require.True(t, ok, "expected traceparent in captured _meta") + // Prove span identity, not just presence: assert the W3C traceparent + // carries the active span's TraceID rather than a stale/global one. + assert.Contains(t, traceparent, span.SpanContext().TraceID().String(), + "traceparent must carry the active span's TraceID") + }) + + //nolint:paralleltest // sequential: shares capture/backendClient state with the sibling subtests + t.Run("ReadResource and GetPrompt: mcpcompat drops outbound Meta (documented SDK limitation)", func(t *testing.T) { + spanCtx, span := tracer.Start(context.Background(), "test-span") + ctx, cancel := context.WithTimeout(spanCtx, 5*time.Second) + defer cancel() + defer span.End() + + // TODO: tripwire for a toolhive-core/mcpcompat limitation (see this + // test's doc comment). These asserts pass ONLY because mcpcompat drops + // outbound Params.Meta on the resources/read and prompts/get paths. When + // they start failing, mcpcompat now forwards Meta: delete the + // forward-compat NOTE comments in client.go/mcp_session.go and replace + // these nil checks with traceparent assertions (as done for CallTool). + _, err := backendClient.ReadResource(ctx, target, "test://capture-resource") + require.NoError(t, err) + assert.Nil(t, capture.get("resource"), + "expected nil ONLY because mcpcompat v0.0.32 drops outbound Meta for resources/read; "+ + "if this is now non-nil, mcpcompat forwards Meta — delete the forward-compat NOTE comments "+ + "in client.go/mcp_session.go and assert traceparent here instead") + + _, err = backendClient.GetPrompt(ctx, target, "test_prompt_capture_meta", nil) + require.NoError(t, err) + assert.Nil(t, capture.get("prompt"), + "expected nil ONLY because mcpcompat v0.0.32 drops outbound Meta for prompts/get; "+ + "if this is now non-nil, mcpcompat forwards Meta — delete the forward-compat NOTE comments "+ + "in client.go/mcp_session.go and assert traceparent here instead") + }) + + //nolint:paralleltest // sequential: reuses the "tool" capture key populated above to prove _meta is now absent + t.Run("CallTool omits _meta without an active span", func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := backendClient.CallTool(ctx, target, "test_tool_capture_meta", nil, nil, nil) + require.NoError(t, err) + + got := capture.get("tool") + assert.Nil(t, got, "expected no _meta to be sent without an active trace context") + }) +} + +// metaCapture records the inbound params._meta observed by the test server's +// capture-only tool/resource/prompt handlers, keyed by capability name. It +// lets tests assert what the vMCP client actually sent on the wire (e.g. an +// injected traceparent), as opposed to what the backend echoes back. +type metaCapture struct { + mu sync.Mutex + meta map[string]*mcp.Meta +} + +func newMetaCapture() *metaCapture { + return &metaCapture{meta: make(map[string]*mcp.Meta)} +} + +func (c *metaCapture) record(key string, meta *mcp.Meta) { + c.mu.Lock() + defer c.mu.Unlock() + c.meta[key] = meta +} + +// get returns the captured _meta for key, or nil if the capability was never +// invoked or was invoked without a _meta field. +func (c *metaCapture) get(key string) *mcp.Meta { + c.mu.Lock() + defer c.mu.Unlock() + return c.meta[key] +} + // startTestMCPServer creates and starts a test MCP server with tools that return _meta. -// Returns the port and cleanup function. -func startTestMCPServer(t *testing.T) (string, func()) { +// It also registers capture-only "test_tool_capture_meta", "test_prompt_capture_meta", and +// "test://capture-resource" capabilities that record the inbound params._meta they receive +// into the returned metaCapture, so tests can assert on OUTBOUND _meta (e.g. injected W3C +// trace context) rather than only on what the backend echoes back. +// Returns the port, the meta capture, and a cleanup function. +func startTestMCPServer(t *testing.T) (string, *metaCapture, func()) { t.Helper() + capture := newMetaCapture() + // Create MCP server mcpServer := server.NewMCPServer("test-backend", "1.0.0") @@ -365,21 +505,74 @@ func startTestMCPServer(t *testing.T) (string, func()) { ) // Add resource that returns _meta - mcpServer.AddResource( + mcpServer.AddResourceWithResult( mcp.Resource{ URI: "test://resource", Name: "Test Resource", Description: "Test resource with metadata", MIMEType: "text/plain", }, - func(_ context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { - // Note: The handler returns []ResourceContents, not *ReadResourceResult - // This is why _meta cannot be forwarded - SDK limitation + func(_ context.Context, _ mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return &mcp.ReadResourceResult{ + Result: mcp.Result{ + Meta: &mcp.Meta{ + AdditionalFields: map[string]any{ + "traceId": "resource-trace-789", + "customTag": "resource-custom", + }, + }, + }, + Contents: []mcp.ResourceContents{ + mcp.TextResourceContents{ + URI: "test://resource", + MIMEType: "text/plain", + Text: "Test resource content", + }, + }, + }, nil + }, + ) + + // Add capture-only tool/resource/prompt that record inbound params._meta + // for outbound-propagation assertions (see metaCapture). + mcpServer.AddTool( + mcp.NewTool("test_tool_capture_meta", + mcp.WithDescription("Capture-only tool that records inbound _meta"), + ), + func(_ context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + capture.record("tool", request.Params.Meta) + return &mcp.CallToolResult{ + Content: []mcp.Content{mcp.NewTextContent("captured")}, + }, nil + }, + ) + mcpServer.AddPrompt( + mcp.NewPrompt("test_prompt_capture_meta", + mcp.WithPromptDescription("Capture-only prompt that records inbound _meta"), + ), + func(_ context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + capture.record("prompt", request.Params.Meta) + return &mcp.GetPromptResult{ + Messages: []mcp.PromptMessage{ + {Role: "user", Content: mcp.NewTextContent("captured")}, + }, + }, nil + }, + ) + mcpServer.AddResource( + mcp.Resource{ + URI: "test://capture-resource", + Name: "Capture Resource", + Description: "Capture-only resource that records inbound _meta", + MIMEType: "text/plain", + }, + func(_ context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + capture.record("resource", request.Params.Meta) return []mcp.ResourceContents{ mcp.TextResourceContents{ - URI: "test://resource", + URI: "test://capture-resource", MIMEType: "text/plain", - Text: "Test resource content", + Text: "captured", }, }, nil }, @@ -438,5 +631,5 @@ func startTestMCPServer(t *testing.T) (string, func()) { _ = listener.Close() } - return port, cleanup + return port, capture, cleanup } diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go new file mode 100644 index 0000000000..d397e4cf5f --- /dev/null +++ b/pkg/vmcp/client/modern.go @@ -0,0 +1,503 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "maps" + "net/http" + "slices" + "strings" + "sync/atomic" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/versions" +) + +// modernClientName is the clientInfo.name vMCP advertises to backends on Modern +// (2026-07-28) requests, matching the Legacy initialize handshake's client name +// (initializeClient in client.go). +const modernClientName = "toolhive-vmcp" + +// modernResultTypeComplete is the sole envelope resultType a single-shot Modern +// call accepts; anything else (e.g. "input_required") is a multi-round retrieval +// this shim does not drive — see errModernInputRequired. +const modernResultTypeComplete = "complete" + +// modernResultTypeInputRequired is the SEP-2322 envelope resultType announcing +// a Multi Round-Trip Request round. The only non-"complete" resultType whose +// payload survives decode (newInputRequiredError); everything else is sentinel- +// only. +const modernResultTypeInputRequired = "input_required" + +// jsonRPCCodeMethodNotFound is the JSON-RPC "method not found" code. Declared +// locally rather than imported (mcpcompat's METHOD_NOT_FOUND is the SDK's wire +// vocabulary; the Modern layer sources its codes independently — see +// pkg/vmcp/server/modern_dispatch.go for the server-side mirror). +const jsonRPCCodeMethodNotFound = -32601 + +// errWrongEra is returned when a backend's response is NOT a recognized Modern +// (2026-07-28) response: a bare 404/400 with no JSON-RPC error body, an empty or +// non-JSON body, or a 200 carrying a Legacy-shaped result (a JSON-RPC result with +// no "resultType"). It signals the peer does not speak the Modern revision at all. +// +// A -32601 (or any other) carried in a well-formed JSON-RPC error object is NOT +// wrong-era: a valid JSON-RPC error body means the backend IS Modern, so a +// -32601 there surfaces as mcp.ErrMethodNotFound and every other code as an +// ordinary call error. +var errWrongEra = errors.New("backend response is not a Modern (2026-07-28) MCP response") + +// errLegacyResponseBody is returned when a Modern request gets a well-formed 200 +// JSON-RPC SUCCESS result that is Legacy-shaped (no resultType). Unlike +// errWrongEra (a transport/protocol rejection that proves the backend did NOT +// process the request), a success body means a lenient Legacy backend MAY have +// executed the request — so the caller MUST NOT auto-retry it (double-execution +// of a side-effecting tool). The cache may still be reclassified. +var errLegacyResponseBody = errors.New("backend returned a Legacy-shaped body (no resultType); it may have executed") + +// errModernInputRequired is the classification sentinel for a Modern envelope +// whose resultType is not "complete" (e.g. "input_required"). The message +// string is frozen for compatibility — it predates the MRTR seam, and both +// probeRevision's errors.Is classification and client-visible error text pin +// it — even though "unsupported" no longer tells the whole story: an +// input_required round's SEP-2322 payload now rides the typed +// vmcp.InputRequiredError wrapping this sentinel, and MRTR consumers branch on +// vmcp.InputRequiredFromError (docs/arch/16-vmcp-mrtr.md, slice 1). +var errModernInputRequired = errors.New("modern response requires additional input (multi-round retrieval unsupported)") + +// errModernProtocolError wraps a well-formed JSON-RPC error whose code is one of +// the Modern-specific codes (-32020/-32021 unconditionally; -32022 only when +// classifyUnsupportedProtocolVersion finds the peer still advertises +// MCPVersionModern): the peer validated our Modern headers/_meta and rejected +// them, so it IS Modern even though the call failed. It is a positive Modern +// signal, distinct from errWrongEra and from a generic JSON-RPC error +// (-32600/-32603, which do not prove Modern). probeRevision classifies it as +// Modern. +var errModernProtocolError = errors.New("modern backend rejected the request with a Modern protocol error") + +// errModernAuth is returned for an HTTP 401/403 to a Modern request (auth +// rejection, often from a proxy). It is deliberately NOT errWrongEra: a transient +// auth blip must not look like a not-Modern signal, or a cached-Modern backend +// would be flipped to Legacy. The status is in the message so vmcp's string-based +// auth classification (IsAuthenticationError) still recognizes it for step-up. +var errModernAuth = errors.New("modern backend returned an auth status") + +// errModernTransient is returned for an HTTP 408/429/5xx, a mid-stream read +// failure, or a transport/network failure (connection refused, timeout, ctx +// cancel) on a Modern request. Like errModernAuth it is NOT errWrongEra: a brief +// outage must not be mistaken for a not-Modern signal. +var errModernTransient = errors.New("modern backend returned a transient error") + +// errModernNegotiatedDown has two sources, both carrying a valid Modern +// envelope whose advertised versions do NOT include mcpparser.MCPVersionModern: +// +// 1. modernDiscover / discoverModernCapabilities: a well-formed server/discover +// response (err == nil from modernCall) whose supportedVersions omits it +// (including an absent or empty list). +// 2. classifyUnsupportedProtocolVersion: a -32022 (CodeUnsupportedProtocolVersion) +// JSON-RPC error whose `data.supported` list omits it (including an absent +// or undecodable data payload). +// +// CANONICAL RATIONALE for the negotiate-down rule. Per SEP-2575, the advertised +// version list — not a clean discover response or a bare -32022 alone — is the +// authoritative signal of whether a peer actually speaks the Modern (2026-07-28) +// revision: a go-sdk v1.7 shim answers both server/discover and protocol errors +// even for a backend negotiating down to Legacy, so neither on its own is proof +// of Modern. go-sdk's own reference client applies the same rule at both +// sources above — mcp/client.go:428-444 for (1), the discover/supportedVersions +// path, and mcp/client.go:360-369 for (2), the -32022/data.supported path. +// +// modernDiscover and probeRevision (client.go) both depend on this rule and +// back-reference it here rather than restating it. Keep the explanation in one +// place: it is the surface that has to be edited whenever the exact-match +// tripwire in discoverModernCapabilities (client.go) fires for a newer Modern +// revision. Note the tripwire is in discoverModernCapabilities, which does the +// supportedVersions check; modernDiscover only builds the transport and +// delegates to it. +// +// This is a definitive Legacy signal carried in a valid Modern envelope — +// distinct from errWrongEra (peer does not speak Modern's wire shape at all) +// and from the other Modern-positive/inconclusive sentinels. Retrying under +// the corrected Legacy classification is always safe, for different reasons +// per source: (1) is inherently safe, since server/discover has no side +// effects; (2) is reached from interpretModernResult for ANY method, including +// side-effecting ones like tools/call, but go-sdk's ServerSession.handle +// (mcp/server.go) rejects an unsupported per-request protocol version BEFORE +// the switch that dispatches to a method handler, so the backend provably +// never executed the request regardless of which method was called. +var errModernNegotiatedDown = errors.New("modern backend negotiated down: advertised versions lack 2026-07-28") + +// modernRequestID supplies monotonically increasing JSON-RPC request ids. Each +// modernCall is a single request/response, so the id only has to be unique +// enough to match a response within one SSE stream. +var modernRequestID atomic.Int64 + +// modernCall issues a single MCP 2026-07-28 ("Modern") stateless JSON-RPC request +// over HTTP POST and decodes the Modern response envelope into out. +// +// It hand-rolls the Modern wire shape that mcpcompat's public Client API still +// cannot express, even now that the underlying go-sdk is v1.7.0-pre.3 (its only +// no-initialize primitive is the private, Legacy-shaped resumeCall), mirroring +// the server envelope in pkg/vmcp/server/modern_envelope.go: no initialize +// handshake, no Mcp-Session-Id, protocol metadata carried per-request in +// _meta, and a Mcp-Method header on every call. +// +// params may carry a caller "_meta"; its three reserved io.modelcontextprotocol/* +// keys are stripped and vMCP's authoritative values overlaid last (vMCP, not the +// caller, is the backend's MCP peer). name is sent as Mcp-Name only for the +// methods that require it (tools/call, resources/read, prompts/get) and only when +// non-empty. hc is the HTTP client whose transport carries the auth/identity/ +// header-forward/trace chain (see buildBackendRoundTripper); modernCall adds no +// transport concerns of its own. +// +// paramHeaders are the SEP-2243 Mcp-Param-* headers (already keyed by full header +// name and validated by pkg/mcp), set before the protocol headers below so a +// caller-derived entry can never overwrite Mcp-Method, Mcp-Name, or +// MCP-Protocol-Version. Only tools/call ever supplies them. +// +// Errors: +// - errWrongEra: the peer is not Modern (bare 4xx/5xx-free rejection, empty or +// non-JSON body, or neither result nor error). +// - errLegacyResponseBody: a 200 JSON-RPC success with no resultType — a lenient +// Legacy backend that MAY have executed the request (caller must not retry). +// - errModernAuth: an HTTP 401/403/407 auth rejection (NOT a not-Modern signal). +// - errModernTransient: an HTTP 408/429/5xx, a mid-stream read failure, or a +// transport/network/timeout failure (NOT a not-Modern signal). +// - mcp.ErrMethodNotFound: a valid -32601 error body. +// - errModernInputRequired: a non-"complete" envelope. +// - errModernProtocolError: a Modern-specific -3202x error body. +// - a wrapped call error: any other JSON-RPC error. +// +// logLevel, when non-empty, is overlaid onto the request _meta as +// io.modelcontextprotocol/logLevel (the Modern replacement for the removed +// logging/setLevel RPC): it opts the request in to the backend's +// notifications/message at that minimum level. Empty leaves the key unset so the +// backend MUST NOT emit log notifications for the request. +// +// onNotification, when non-nil, is invoked with each server->client notification +// the SSE stream interleaves ahead of the response (a nil-Method envelope is the +// response, not a notification). It is how a caller relays the log/progress +// notifications logLevel elicited; nil preserves the historical drop. It only +// ever fires on the SSE path — a single JSON 200 body carries no notifications. +func modernCall( + ctx context.Context, + hc *http.Client, + endpoint, method string, + params map[string]any, + name string, + paramHeaders map[string]string, + out any, + logLevel string, + onNotification func(method string, params json.RawMessage), +) error { + id := modernRequestID.Add(1) + + reqParams := maps.Clone(params) + if reqParams == nil { + reqParams = map[string]any{} + } + reqParams["_meta"] = mergeModernMeta(params["_meta"], logLevel) + + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": reqParams, + }) + if err != nil { + return fmt.Errorf("marshaling %s request: %w", method, err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("building %s request: %w", method, err) + } + // SEP-2243 mirrored parameter headers go on FIRST, so the protocol headers + // below win on any collision. A backend cannot name a designated parameter + // "Method" and hijack Mcp-Method, since Set overwrites. + for hdrName, hdrValue := range paramHeaders { + req.Header.Set(hdrName, hdrValue) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("MCP-Protocol-Version", mcpparser.MCPVersionModern) + // Mcp-Method is required on EVERY Modern request (ValidateHeaderConsistency). + req.Header.Set("Mcp-Method", method) + if name != "" && mcpparser.IsNameRequiredMethod(method) { + req.Header.Set("Mcp-Name", mcpparser.EncodeSentinelName(name)) + } + // Mcp-Session-Id is deliberately never set: Modern is stateless. + + resp, err := hc.Do(req) + if err != nil { + // A transport/network failure (connection refused, timeout, ctx cancel) is + // transient — not a not-Modern signal — so it must not poison the revision cache. + return fmt.Errorf("%w: sending %s request: %w", errModernTransient, method, err) + } + defer func() { + // Drain so the connection can be reused (go-style rule); the readers below + // may stop early (SSE match) or the body may be a bare error page. Bounded + // by maxResponseSize so a hostile backend can't stall us on an unbounded + // drain now that this path is live in production (probeRevision). + _, _ = io.CopyN(io.Discard, resp.Body, maxResponseSize) + _ = resp.Body.Close() + }() + + result, rpcErr, err := readModernEnvelope(resp, id, onNotification) + if err != nil { + return err + } + return interpretModernResult(result, rpcErr, method, out) +} + +// interpretModernResult maps a decoded Modern JSON-RPC response to an error or +// decodes it into out. Split from modernCall to keep each within the cyclomatic +// limit; see the modernCall doc for the error taxonomy. +func interpretModernResult(result json.RawMessage, rpcErr *modernRPCError, method string, out any) error { + if rpcErr != nil { + if rpcErr.Code == jsonRPCCodeMethodNotFound { + return fmt.Errorf("%w: %s", mcp.ErrMethodNotFound, rpcErr.Message) + } + if isModernProtocolCode(rpcErr.Code) { + return fmt.Errorf("%w: %s (rpc code %d)", errModernProtocolError, rpcErr.Message, rpcErr.Code) + } + if int64(rpcErr.Code) == mcpparser.CodeUnsupportedProtocolVersion { + return classifyUnsupportedProtocolVersion(rpcErr) + } + return fmt.Errorf("modern %s: rpc error %d: %s", method, rpcErr.Code, rpcErr.Message) + } + + // The Modern result is an envelope keyed by resultType (modern_envelope.go). + var envelope struct { + ResultType string `json:"resultType"` + } + if json.Unmarshal(result, &envelope) != nil { + return errWrongEra + } + switch envelope.ResultType { + case modernResultTypeComplete: + // proceed to decode + case "": + // A JSON-RPC success result with no resultType is a Legacy-shaped body: a + // lenient Legacy backend that ignored our Modern headers and executed the + // request. Distinct from errWrongEra so the caller does not auto-retry. + return errLegacyResponseBody + default: + // Typed so an "input_required" round's SEP-2322 payload (inputRequests, + // requestState) survives for MRTR consumers (vmcp.InputRequiredFromError) + // — gated on the three methods that may carry a round and on payload + // validity; unwraps to errModernInputRequired with an identical message, + // so classification and behavior are unchanged for everyone else. + return newInputRequiredError(method, envelope.ResultType, result) + } + + if out != nil { + if err := json.Unmarshal(result, out); err != nil { + return fmt.Errorf("decoding %s result: %w", method, err) + } + } + return nil +} + +// mergeModernMeta strips the reserved io.modelcontextprotocol/* keys from a +// caller-supplied _meta (if any) and overlays vMCP's authoritative values last. +// The caller's _meta is never mutated (StripReservedMeta clones it). +// +// logLevel, when non-empty, is overlaid as io.modelcontextprotocol/logLevel — +// vMCP, not the caller, decides whether the backend is asked to emit log +// notifications for this request (the caller's own key was already stripped as +// reserved). Empty leaves the key unset. +func mergeModernMeta(callerMeta any, logLevel string) map[string]any { + m, _ := callerMeta.(map[string]any) + meta := mcpparser.StripReservedMeta(m) + if meta == nil { + // StripReservedMeta returns nil for empty/nil input; this needs a + // non-nil map to overlay vMCP's authoritative values onto below. + meta = map[string]any{} + } + for k, v := range mcpparser.ModernRequestMeta(modernClientName, versions.Version) { + meta[k] = v + } + if logLevel != "" { + meta[mcpparser.MetaKeyLogLevel] = logLevel + } + return meta +} + +// modernRPCError is the JSON-RPC error object. Data carries the SEP-2575 +// error-specific payload (e.g. UnsupportedProtocolVersionData for -32022, +// classified by classifyUnsupportedProtocolVersion); absent for the other +// Modern-specific codes. +type modernRPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data json.RawMessage `json:"data,omitempty"` +} + +// modernRPCEnvelope is the outer JSON-RPC response envelope. Method and Params +// are set only on server->client requests/notifications interleaved on an SSE +// stream; a single-shot client matches the response (Method == "") and, when an +// onNotification listener is bound, relays Method/Params for the notifications. +type modernRPCEnvelope struct { + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *modernRPCError `json:"error"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` +} + +// readModernEnvelope reads the JSON-RPC response matching wantID, handling both +// application/json and text/event-stream bodies (mirroring mcpcompat's +// resume.go readRPCResponse). The body is bounded by maxResponseSize in both +// branches — the 100MB cap otherwise lives only inside the mcp-go client and is +// lost on this raw path. A body that is not a recognized Modern JSON-RPC response +// (empty, non-JSON, or neither result nor error) yields errWrongEra. +// +// Auth (401/403) and transient (408/429/5xx) statuses are classified BEFORE any +// body or SSE handling — even one a proxy tags as text/event-stream — so a +// transient blip is never mistaken for a not-Modern signal (which would poison +// the revision cache). A genuine Modern -32601/-32602 rides HTTP 404/400 WITH a +// JSON-RPC body and is handled by the body logic below, so those statuses are +// deliberately NOT short-circuited here. +func readModernEnvelope( + resp *http.Response, wantID int64, onNotification func(method string, params json.RawMessage), +) (json.RawMessage, *modernRPCError, error) { + switch { + case resp.StatusCode == http.StatusUnauthorized, resp.StatusCode == http.StatusForbidden, + resp.StatusCode == http.StatusProxyAuthRequired: + return nil, nil, fmt.Errorf("%w: HTTP %d", errModernAuth, resp.StatusCode) + case resp.StatusCode == http.StatusRequestTimeout, resp.StatusCode == http.StatusTooManyRequests, + resp.StatusCode >= 500: + return nil, nil, fmt.Errorf("%w: HTTP %d", errModernTransient, resp.StatusCode) + } + + body := io.LimitReader(resp.Body, maxResponseSize) + + if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") { + return readModernSSE(body, wantID, onNotification) + } + + data, err := io.ReadAll(body) + if err != nil { + return nil, nil, fmt.Errorf("%w: reading response body: %w", errModernTransient, err) + } + if len(bytes.TrimSpace(data)) == 0 { + return nil, nil, errWrongEra + } + var env modernRPCEnvelope + if json.Unmarshal(data, &env) != nil { + return nil, nil, errWrongEra + } + if env.Error == nil && len(env.Result) == 0 { + return nil, nil, errWrongEra + } + return env.Result, env.Error, nil +} + +// readModernSSE scans an SSE body for the response whose id matches wantID. +// Server->client requests/notifications interleaved on the stream (envelopes +// with a Method) are not the response: when onNotification is non-nil their +// method and params are handed to it (so a caller can relay the +// notifications/message and notifications/progress the request's logLevel / +// progressToken elicited); when nil they are dropped as before. A stream that +// ends without a matching response yields errWrongEra. +func readModernSSE( + body io.Reader, wantID int64, onNotification func(method string, params json.RawMessage), +) (json.RawMessage, *modernRPCError, error) { + sc := bufio.NewScanner(body) + // Cap the token at maxResponseSize (the doc-promised bound) so a valid single + // data: event up to that size decodes; the outer io.LimitReader already bounds + // the total, so this cannot over-allocate. + sc.Buffer(make([]byte, 0, 64*1024), maxResponseSize) + for sc.Scan() { + data, ok := strings.CutPrefix(sc.Text(), "data:") + if !ok { + continue + } + var env modernRPCEnvelope + if json.Unmarshal([]byte(strings.TrimSpace(data)), &env) != nil { + continue + } + if env.Method != "" { + // server->client request/notification; not our response. Relay it when a + // listener is bound, otherwise drop it (historical behavior). + if onNotification != nil { + onNotification(env.Method, env.Params) + } + continue + } + if !modernIDMatches(env.ID, wantID) { + continue + } + if env.Error == nil && len(env.Result) == 0 { + return nil, nil, errWrongEra + } + return env.Result, env.Error, nil + } + if err := sc.Err(); err != nil { + return nil, nil, fmt.Errorf("%w: reading SSE stream: %w", errModernTransient, err) + } + return nil, nil, errWrongEra +} + +// modernIDMatches reports whether the raw JSON id equals wantID. +func modernIDMatches(raw json.RawMessage, wantID int64) bool { + if len(raw) == 0 { + return false + } + var n int64 + return json.Unmarshal(raw, &n) == nil && n == wantID +} + +// isModernProtocolCode reports whether code is one of the two Modern-specific +// JSON-RPC error codes that unconditionally prove the peer is Modern +// (-32020/-32021: header/meta validation). Generic JSON-RPC codes +// (-32600/-32601/-32603) do not, since a Legacy backend also returns them. +// +// CodeUnsupportedProtocolVersion (-32022) is deliberately excluded: unlike the +// other two, it does not by itself prove Modern — it means "I don't support +// the version you asked for" and can equally mean the peer negotiated down to +// Legacy. See classifyUnsupportedProtocolVersion, which resolves it by +// decoding the error's advertised supported-versions list. +func isModernProtocolCode(code int) bool { + switch int64(code) { + case mcpparser.CodeHeaderMismatch, mcpparser.CodeMissingClientCapability: + return true + default: + return false + } +} + +// classifyUnsupportedProtocolVersion resolves a -32022 +// (CodeUnsupportedProtocolVersion) error into a Modern-positive or +// negotiated-down signal. go-sdk v1.7's reference client (mcp/client.go:360-369) +// decodes the error's `data.supported` list and retries Modern when it finds a +// mutually-supported version >= 2026-07-28 (a range check), falling back to +// Legacy initialize otherwise. This diverges from that: it exact-matches +// MCPVersionModern instead of a range, same as discoverModernCapabilities' +// exact-match (client.go) and safe for the same reason — vMCP's shim only +// speaks that one Modern wire shape, and a dual-era peer offering a future +// Modern revision also lists 2025-11-25 for us to fall back to. TRIPWIRE: when +// a newer Modern revision is added, this must become a set/range check +// alongside discoverModernCapabilities' — this is the second exact-match site. +// +// Otherwise — including an absent or undecodable data payload — it is +// errModernNegotiatedDown. +func classifyUnsupportedProtocolVersion(rpcErr *modernRPCError) error { + var data struct { + Supported []string `json:"supported"` + } + _ = json.Unmarshal(rpcErr.Data, &data) // best-effort; empty Supported on any decode failure + if slices.Contains(data.Supported, mcpparser.MCPVersionModern) { + return fmt.Errorf("%w: %s (rpc code %d)", errModernProtocolError, rpcErr.Message, rpcErr.Code) + } + return fmt.Errorf("%w: data.supported=%v", errModernNegotiatedDown, data.Supported) +} diff --git a/pkg/vmcp/client/modern_calls_test.go b/pkg/vmcp/client/modern_calls_test.go new file mode 100644 index 0000000000..5b53500ab8 --- /dev/null +++ b/pkg/vmcp/client/modern_calls_test.go @@ -0,0 +1,248 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// bodyRecordingServer stands up a fake Modern backend that records the last +// request's headers and decoded params (so tests can assert the body identifier +// matches the Mcp-Name header), and replies with result. Callers pre-seed the +// revision cache Modern so the verb skips the discover probe and hits this server. +func bodyRecordingServer(t *testing.T, result map[string]any) (*httptest.Server, *http.Header, *map[string]any) { + t.Helper() + var gotHeader http.Header + gotBody := map[string]any{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHeader = r.Header.Clone() + body, _ := readAll(t, r) + var req struct { + ID any `json:"id"` + Params map[string]any `json:"params"` + } + require.NoError(t, json.Unmarshal(body, &req)) + for k, v := range req.Params { + gotBody[k] = v + } + writeModernResult(t, w, req.ID, result) + })) + t.Cleanup(srv.Close) + return srv, &gotHeader, &gotBody +} + +func modernClient(t *testing.T, url string) (*httpBackendClient, *vmcp.BackendTarget) { + t.Helper() + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: url, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) // skip the probe + return h, target +} + +// TestModernCallTool verifies tools/call request shaping (Mcp-Method, translated +// Mcp-Name matching the body name, no session id) and result decode (isError, +// structuredContent, _meta forwarded). +func TestModernCallTool(t *testing.T) { + t.Parallel() + + srv, hdr, body := bodyRecordingServer(t, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "out"}}, + "structuredContent": map[string]any{"k": "v"}, + "isError": true, + "_meta": map[string]any{"trace": "x"}, + }) + h, target := modernClient(t, srv.URL) + target.OriginalCapabilityName = "backend_echo" // advertised "echo" -> backend "backend_echo" + + res, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hi"}, map[string]any{"caller": "meta"}, nil) + require.NoError(t, err) + + assert.Equal(t, "tools/call", hdr.Get("Mcp-Method")) + assert.Equal(t, "backend_echo", hdr.Get("Mcp-Name"), "Mcp-Name must be the translated name") + assert.Equal(t, "backend_echo", (*body)["name"], "body name must match Mcp-Name") + assert.Empty(t, hdr.Get("Mcp-Session-Id")) + + require.Len(t, res.Content, 1) + assert.Equal(t, "out", res.Content[0].Text) + assert.Equal(t, "v", res.StructuredContent["k"]) + assert.True(t, res.IsError) + assert.Equal(t, "x", res.Meta["trace"], "result _meta must be forwarded to core") +} + +// TestModernCallTool_LogLevelGating pins when the Modern tools/call opts into +// backend log notifications: the io.modelcontextprotocol/logLevel _meta key is +// injected at debug level ONLY when server->client forwarding is bound (the +// Modern analogue of enableBackendLogging on the Legacy path). Unbound, the key +// must be absent so a conformant backend does not emit notifications/message. +func TestModernCallTool_LogLevelGating(t *testing.T) { + t.Parallel() + + logLevelOf := func(body *map[string]any) (any, bool) { + meta, _ := (*body)["_meta"].(map[string]any) + v, ok := meta["io.modelcontextprotocol/logLevel"] + return v, ok + } + + t.Run("bound forwarders inject debug logLevel", func(t *testing.T) { + t.Parallel() + srv, _, body := bodyRecordingServer(t, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + }) + h, target := modernClient(t, srv.URL) + h.BindForwarders(&stubElicitationRequester{}, &stubSamplingRequester{}, stubClientNotifier{}) + + _, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hi"}, nil, nil) + require.NoError(t, err) + + v, ok := logLevelOf(body) + require.True(t, ok, "bound forwarders must opt the call into log notifications") + assert.Equal(t, "debug", v) + }) + + t.Run("unbound forwarders send no logLevel", func(t *testing.T) { + t.Parallel() + srv, _, body := bodyRecordingServer(t, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + }) + h, target := modernClient(t, srv.URL) + + _, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hi"}, nil, nil) + require.NoError(t, err) + + _, ok := logLevelOf(body) + assert.False(t, ok, "without forwarding the backend must not be asked to emit logs") + }) +} + +// TestModernReadResource verifies resources/read shaping (Mcp-Name mirrors the +// translated uri) and text/blob content decode. +func TestModernReadResource(t *testing.T) { + t.Parallel() + + srv, hdr, body := bodyRecordingServer(t, map[string]any{ + "contents": []any{ + map[string]any{"uri": "file:///x", "mimeType": "text/plain", "text": "hello"}, + map[string]any{"uri": "file:///y", "mimeType": "application/octet-stream", "blob": "AAAA"}, + }, + "_meta": map[string]any{"trace": "r"}, + }) + h, target := modernClient(t, srv.URL) + target.OriginalCapabilityName = "file:///backend" + + res, err := h.ReadResource(context.Background(), target, "file:///advertised") + require.NoError(t, err) + + assert.Equal(t, "resources/read", hdr.Get("Mcp-Method")) + assert.Equal(t, "file:///backend", hdr.Get("Mcp-Name")) + assert.Equal(t, "file:///backend", (*body)["uri"], "body uri must match Mcp-Name") + require.Len(t, res.Contents, 2) + assert.Equal(t, "hello", res.Contents[0].Text) + assert.Equal(t, "AAAA", res.Contents[1].Blob) + assert.Equal(t, "r", res.Meta["trace"]) +} + +// TestModernGetPrompt verifies prompts/get shaping (translated Mcp-Name matching +// the body name) and message content decode. +func TestModernGetPrompt(t *testing.T) { + t.Parallel() + + srv, hdr, body := bodyRecordingServer(t, map[string]any{ + "description": "d", + "messages": []any{ + map[string]any{"role": "user", "content": map[string]any{"type": "text", "text": "hi"}}, + }, + "_meta": map[string]any{"trace": "p"}, + }) + h, target := modernClient(t, srv.URL) + target.OriginalCapabilityName = "backend_prompt" + + res, err := h.GetPrompt(context.Background(), target, "advertised_prompt", map[string]any{"a": "b"}) + require.NoError(t, err) + + assert.Equal(t, "prompts/get", hdr.Get("Mcp-Method")) + assert.Equal(t, "backend_prompt", hdr.Get("Mcp-Name")) + assert.Equal(t, "backend_prompt", (*body)["name"]) + assert.Equal(t, "d", res.Description) + require.Len(t, res.Messages, 1) + assert.Equal(t, "user", res.Messages[0].Role) + assert.Equal(t, "hi", res.Messages[0].Content.Text) + assert.Equal(t, "p", res.Meta["trace"]) +} + +// TestModernComplete verifies completion/complete shaping (NOT name-required, so +// no Mcp-Name) and value decode, plus the -32601 -> empty leniency. +func TestModernComplete(t *testing.T) { + t.Parallel() + + t.Run("returns values, no Mcp-Name", func(t *testing.T) { + t.Parallel() + srv, hdr, _ := bodyRecordingServer(t, map[string]any{ + "completion": map[string]any{"values": []any{"a", "b"}, "total": 2, "hasMore": false}, + }) + h, target := modernClient(t, srv.URL) + + res, err := h.Complete(context.Background(), target, + vmcp.CompletionRef{Type: vmcp.CompletionRefTypeResource, URI: "file:///x"}, "arg", "va", nil) + require.NoError(t, err) + assert.Equal(t, "completion/complete", hdr.Get("Mcp-Method")) + assert.Empty(t, hdr.Get("Mcp-Name"), "completion/complete is not name-required") + assert.Equal(t, []string{"a", "b"}, res.Values) + }) + + t.Run("method not found yields empty result", func(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"method not found"}}`)) + })) + t.Cleanup(srv.Close) + h, target := modernClient(t, srv.URL) + + res, err := h.Complete(context.Background(), target, + vmcp.CompletionRef{Type: vmcp.CompletionRefTypeResource, URI: "file:///x"}, "arg", "", nil) + require.NoError(t, err) + assert.Equal(t, []string{}, res.Values, "a backend without completions degrades to empty, not an error") + }) +} + +// TestIntegration_ModernCallTool_NameRequired proves a name-required verb +// round-trips end-to-end against the real dispatchModern server: Mcp-Name is +// mirrored in header + body, no Mcp-Session-Id on the wire, and the echo tool +// result comes back. +func TestIntegration_ModernCallTool_NameRequired(t *testing.T) { + t.Parallel() + + backendURL := startEchoBackend(t) + vmcpSrv := newModernVMCPServer(t, backendURL) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "vmcp-modern", + WorkloadName: "vMCP Modern", + BaseURL: vmcpSrv.URL + "/mcp", + TransportType: "streamable-http", + } + + res, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hello modern"}, nil, nil) + require.NoError(t, err) + + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev) + + require.Len(t, res.Content, 1) + assert.Equal(t, "hello modern", res.Content[0].Text) + assert.False(t, res.IsError) +} diff --git a/pkg/vmcp/client/modern_enumerate_test.go b/pkg/vmcp/client/modern_enumerate_test.go new file mode 100644 index 0000000000..14c25c1da6 --- /dev/null +++ b/pkg/vmcp/client/modern_enumerate_test.go @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// modernReq decodes a Modern JSON-RPC request, reading the body exactly once +// (callers must not read r.Body separately) and returning the id and cursor. +func modernReq(t *testing.T, r *http.Request) (id any, cursor string) { + t.Helper() + body, err := readAll(t, r) + require.NoError(t, err) + var req struct { + ID any `json:"id"` + Params struct { + Cursor string `json:"cursor"` + } `json:"params"` + } + require.NoError(t, json.Unmarshal(body, &req)) + return req.ID, req.Params.Cursor +} + +// writeModernResult writes a Modern JSON-RPC success envelope for the given +// request id, wrapping result under a "complete" resultType. +func writeModernResult(t *testing.T, w http.ResponseWriter, id any, result map[string]any) { + t.Helper() + if result["resultType"] == nil { + result["resultType"] = "complete" + } + out, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + require.NoError(t, err) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(out) +} + +// TestModernEnumerate_Pagination verifies tools/list follows nextCursor across +// pages and aggregates every tool (no page-1 truncation, #5851). +func TestModernEnumerate_Pagination(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "tools/list", r.Header.Get("Mcp-Method")) + id, cursor := modernReq(t, r) + if cursor == "" { + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "t1", "inputSchema": map[string]any{"type": "object"}}}, + "nextCursor": "page2", + }) + return + } + require.Equal(t, "page2", cursor) + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "t2", "inputSchema": map[string]any{"type": "object"}}}, + }) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + caps := &mcpmcp.ServerCapabilities{Tools: &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{}} + + list, err := h.modernEnumerate(context.Background(), target, caps) + require.NoError(t, err) + require.Len(t, list.Tools, 2, "both pages must aggregate") + assert.Equal(t, "t1", list.Tools[0].Name) + assert.Equal(t, "t2", list.Tools[1].Name) + assert.Equal(t, "b", list.Tools[0].BackendID) +} + +// TestModernEnumerate_GatesOnFlags verifies only advertised capabilities are +// enumerated: a backend advertising only Tools is never asked for resources or +// prompts. +func TestModernEnumerate_GatesOnFlags(t *testing.T) { + t.Parallel() + + var toolsList, otherList atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Mcp-Method") == "tools/list" { + toolsList.Add(1) + } else { + otherList.Add(1) + } + id, _ := modernReq(t, r) + writeModernResult(t, w, id, map[string]any{"tools": []any{}}) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + caps := &mcpmcp.ServerCapabilities{Tools: &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{}} + + _, err := h.modernEnumerate(context.Background(), target, caps) + require.NoError(t, err) + + assert.Equal(t, int32(1), toolsList.Load(), "tools advertised: listed once") + assert.Zero(t, otherList.Load(), "resources/prompts not advertised: must not be listed") +} + +// TestModernEnumerate_NilCapsEmpty verifies nil caps (a -3202x backend) yields +// an empty list with no list calls at all. +func TestModernEnumerate_NilCapsEmpty(t *testing.T) { + t.Parallel() + + var called atomic.Bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + called.Store(true) + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + + list, err := h.modernEnumerate(context.Background(), target, nil) + require.NoError(t, err) + assert.Empty(t, list.Tools) + assert.Empty(t, list.Resources) + assert.Empty(t, list.Prompts) + assert.False(t, called.Load(), "nil caps must not issue any list request") +} diff --git a/pkg/vmcp/client/modern_integration_test.go b/pkg/vmcp/client/modern_integration_test.go new file mode 100644 index 0000000000..a58d4ab238 --- /dev/null +++ b/pkg/vmcp/client/modern_integration_test.go @@ -0,0 +1,193 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// startEchoBackend stands up a real in-process MCP server over streamable-HTTP +// exposing a single "echo" tool, returning the URL of its /mcp endpoint. Mirrors +// the server package's startRealMCPBackend test utility (not importable across +// the package boundary). +func startEchoBackend(t *testing.T) string { + t.Helper() + + mcpSrv := mcpserver.NewMCPServer("real-backend", "1.0.0") + mcpSrv.AddTool( + mcpmcp.NewTool("echo", + mcpmcp.WithDescription("Echoes the input back"), + mcpmcp.WithString("input", mcpmcp.Required()), + ), + func(_ context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + args, _ := req.Params.Arguments.(map[string]any) + input, _ := args["input"].(string) + return &mcpmcp.CallToolResult{Content: []mcpmcp.Content{mcpmcp.NewTextContent(input)}}, nil + }, + ) + + mux := http.NewServeMux() + mux.Handle("/mcp", mcpserver.NewStreamableHTTPServer(mcpSrv)) + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + return ts.URL + "/mcp" +} + +// newModernVMCPServer stands up a real vMCP server with Modern (2026-07-28) +// dispatch enabled, aggregating the echo backend at backendURL. A well-formed +// Modern request routes through classifyingHandler -> dispatchModern (the Phase-2 +// server side), so this is the real target modernCall must satisfy end-to-end. +func newModernVMCPServer(t *testing.T, backendURL string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: backendURL, + TransportType: "streamable-http", + } + mockRegistry := mocks.NewMockBackendRegistry(ctrl) + mockRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, authReg.RegisterStrategy( + authtypes.StrategyTypeUnauthenticated, strategies.NewUnauthenticatedStrategy(), + )) + + backendClient, err := NewHTTPBackendClient(authReg) + require.NoError(t, err) + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + srv, err := server.New( + context.Background(), + &server.Config{ + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: vmcpsession.NewSessionFactory(authReg), + Aggregator: agg, + }, + router.NewSessionRouter(&vmcp.RoutingTable{}), + backendClient, + mockRegistry, + nil, + ) + require.NoError(t, err) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// TestIntegration_ModernCall_Discover proves the shim end-to-end: modernCall's +// hand-rolled Modern request satisfies the real Phase-2 dispatchModern server +// (the headers and _meta it validates), and the server's discover envelope +// decodes back through the shim with capability flags reflecting the echo +// backend's actual tool set. +func TestIntegration_ModernCall_Discover(t *testing.T) { + t.Parallel() + + backendURL := startEchoBackend(t) + vmcpSrv := newModernVMCPServer(t, backendURL) + + // Capture the exact headers the shim sent by recording them in the client's + // transport, then delegating to the server's real client transport. + var gotHeaders http.Header + hc := &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) { + gotHeaders = req.Header.Clone() + return http.DefaultTransport.RoundTrip(req) + })} + + var out struct { + ResultType string `json:"resultType"` + Capabilities struct { + Tools json.RawMessage `json:"tools"` + Completions json.RawMessage `json:"completions"` + Resources json.RawMessage `json:"resources"` + Prompts json.RawMessage `json:"prompts"` + } `json:"capabilities"` + SupportedVersions []string `json:"supportedVersions"` + } + err := modernCall(context.Background(), hc, vmcpSrv.URL+"/mcp", "server/discover", nil, "", nil, &out, "", nil) + require.NoError(t, err) + + // Request shaping reached the real server intact. + assert.Equal(t, "server/discover", gotHeaders.Get("Mcp-Method")) + assert.Equal(t, "2026-07-28", gotHeaders.Get("MCP-Protocol-Version")) + assert.Contains(t, gotHeaders.Get("Accept"), "text/event-stream") + assert.Empty(t, gotHeaders.Get("Mcp-Session-Id"), "Modern responses/requests carry no session id") + + // Decoded discover envelope reflects the echo backend: a tool present, + // completions unconditional, resources/prompts absent. + assert.Equal(t, "complete", out.ResultType) + assert.Contains(t, out.SupportedVersions, "2026-07-28") + assert.NotEmpty(t, out.Capabilities.Tools, "echo backend has a tool") + assert.NotEmpty(t, out.Capabilities.Completions, "completions is advertised unconditionally") + assert.Empty(t, out.Capabilities.Resources, "echo backend exposes no resources") + assert.Empty(t, out.Capabilities.Prompts, "echo backend exposes no prompts") +} + +// TestIntegration_ModernListCapabilities proves the full Modern enumeration path +// end-to-end: ListCapabilities probes the Phase-2 dispatchModern server as +// Modern, then enumerates its real capabilities (tools/list) via the shim — no +// initialize handshake, no Mcp-Session-Id — and returns the echo backend's tool. +func TestIntegration_ModernListCapabilities(t *testing.T) { + t.Parallel() + + backendURL := startEchoBackend(t) + vmcpSrv := newModernVMCPServer(t, backendURL) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "vmcp-modern", + WorkloadName: "vMCP Modern", + BaseURL: vmcpSrv.URL + "/mcp", + TransportType: "streamable-http", + } + + caps, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + + // Classified Modern (not Legacy initialize) and the echo tool enumerated. + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev) + + require.Len(t, caps.Tools, 1, "the echo backend's tool must enumerate through Modern") + assert.Equal(t, "echo", caps.Tools[0].Name) + assert.Equal(t, target.WorkloadID, caps.Tools[0].BackendID) + assert.Empty(t, caps.Resources, "echo backend exposes no resources") + assert.Empty(t, caps.Prompts, "echo backend exposes no prompts") +} diff --git a/pkg/vmcp/client/modern_mrtr.go b/pkg/vmcp/client/modern_mrtr.go new file mode 100644 index 0000000000..b16061da86 --- /dev/null +++ b/pkg/vmcp/client/modern_mrtr.go @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "encoding/json" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// mrtrMethods are the only client requests that may carry an input_required +// round: "Servers MUST NOT send InputRequiredResult responses on any other +// client requests" (spec, basic/patterns/mrtr). The set coincides with +// pkg/mcp's nameRequiredMethods today, but the concepts are independent — +// Mcp-Name validation is a header rule, this is the MRTR method allow-list, +// and the final 2026-07-28 cut may grow this one alone (SEP-2322 Final also +// allowed GetTaskPayloadRequest before tasks moved to an extension; see the +// drift list in docs/arch/16-vmcp-mrtr.md) — so they are deliberately not +// shared. +var mrtrMethods = map[string]bool{ + "tools/call": true, + "resources/read": true, + "prompts/get": true, +} + +// newInputRequiredError builds the typed error (vmcp.InputRequiredError, +// unwrapping to errModernInputRequired) for a non-"complete" Modern result +// envelope. The error always classifies and renders identically to the +// previous fmt.Errorf wrapping; whether it also carries an extractable round +// (vmcp.InputRequiredFromError reporting ok=true) is gated fail-closed: +// +// - resultType must be exactly "input_required" — the spec says an +// unrecognized resultType "MUST be considered invalid", so anything else +// (including the Tasks extension's "task", which vMCP never solicits — +// it advertises no extensions) keeps pure sentinel semantics; +// - method must be one of mrtrMethods — a round on any other request is a +// backend protocol violation, and relaying it would make vMCP the +// non-conformant server; +// - the payload must decode strictly and be drivable: server requirement 6 +// mandates at least one of inputRequests or requestState, and an envelope +// violating it (or one with a wrong-typed field) must not surface as a +// round — a consumer following client requirement 1 would retry it having +// gathered nothing, which the backend may answer with the same envelope +// again, a loop with no exit. Fail closed to the sentinel instead. +func newInputRequiredError(method, resultType string, result json.RawMessage) error { + e := &vmcp.InputRequiredError{ResultType: resultType, Sentinel: errModernInputRequired} + if resultType != modernResultTypeInputRequired || !mrtrMethods[method] { + return e + } + var payload struct { + InputRequests map[string]json.RawMessage `json:"inputRequests"` + RequestState *string `json:"requestState"` + } + if err := json.Unmarshal(result, &payload); err != nil { + // Wrong-typed inputRequests/requestState: a schema violation, not a + // drivable round. The sentinel classification survives. + return e + } + if len(payload.InputRequests) == 0 && payload.RequestState == nil { + // Server requirement 6 violation ("at least one of inputRequests or + // requestState"): nothing to fulfill and nothing to echo. + return e + } + e.Result = &vmcp.InputRequiredResult{ + InputRequests: payload.InputRequests, + RequestState: payload.RequestState, + } + return e +} diff --git a/pkg/vmcp/client/modern_mrtr_test.go b/pkg/vmcp/client/modern_mrtr_test.go new file mode 100644 index 0000000000..8b89f6479b --- /dev/null +++ b/pkg/vmcp/client/modern_mrtr_test.go @@ -0,0 +1,235 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestInterpretModernResultInputRequired pins the egress MRTR seam +// (docs/arch/16-vmcp-mrtr.md, slice 1): an input_required envelope must +// surface as an error that (a) still satisfies every existing +// errors.Is(err, errModernInputRequired) classification, (b) renders the +// exact message the previous fmt.Errorf wrapping produced (no client-visible +// drift), and (c) carries an extractable SEP-2322 round — but only for a +// valid, drivable round on a method that may carry one; protocol violations +// fail closed to sentinel-only semantics. +func TestInterpretModernResultInputRequired(t *testing.T) { + t.Parallel() + + strPtr := func(s string) *string { return &s } + + tests := []struct { + name string + method string + result string + wantMsg string // "" means derive from the resultType template + wantPayload bool + wantKeys map[string]string // request key -> method, per Methods() + wantState *string + }{ + { + name: "input_required with elicitation and state", + method: "tools/call", + result: `{ + "resultType": "input_required", + "inputRequests": { + "github_login": {"method": "elicitation/create", "params": {"message": "user?"}}, + "summary": {"method": "sampling/createMessage", "params": {"maxTokens": 10}} + }, + "requestState": "opaque-blob" + }`, + wantMsg: `modern response requires additional input (multi-round retrieval unsupported): ` + + `resultType="input_required"`, + wantPayload: true, + wantKeys: map[string]string{ + "github_login": "elicitation/create", + "summary": "sampling/createMessage", + }, + wantState: strPtr("opaque-blob"), + }, + { + // SEP-2322's load-shedding shape: requestState only, no inputRequests. + // The client may retry immediately; the payload must still decode. + name: "load-shedding round (requestState only)", + method: "resources/read", + result: `{"resultType": "input_required", "requestState": "resume-here"}`, + wantPayload: true, + wantKeys: nil, + wantState: strPtr("resume-here"), + }, + { + // requestState?: string has no minimum length: present-and-empty is + // legal, distinct from absent, and must survive as such — client + // requirement 2 makes the retry echo the exact value. + name: "present-and-empty requestState is a drivable round", + method: "prompts/get", + result: `{"resultType": "input_required", "requestState": ""}`, + wantPayload: true, + wantKeys: nil, + wantState: strPtr(""), + }, + { + // Wrong-typed inputRequests: a schema violation. The typed error and + // sentinel classification survive, but no round is extractable — a + // consumer must fail closed on a protocol violation, not relay an + // empty envelope. + name: "malformed inputRequests fails closed to sentinel-only", + method: "tools/call", + result: `{"resultType": "input_required", "inputRequests": "not-a-map"}`, + wantPayload: false, + }, + { + // Server requirement 6 violation: neither inputRequests nor + // requestState. Indistinguishable-from-empty must not become a + // retry-forever round. + name: "empty envelope (server requirement 6 violation) fails closed", + method: "tools/call", + result: `{"resultType": "input_required"}`, + wantPayload: false, + }, + { + // A wrong-typed requestState (numeric) must not silently decode as + // absent or empty. + name: "wrong-typed requestState fails closed", + method: "tools/call", + result: `{"resultType": "input_required", "requestState": 42}`, + wantPayload: false, + }, + { + // Only tools/call, resources/read, and prompts/get may carry a round + // ("Servers MUST NOT send InputRequiredResult responses on any other + // client requests"). A valid-looking round on another method is a + // backend protocol violation; relaying it would make vMCP the + // non-conformant server. + name: "input_required on a disallowed method carries no round", + method: "completion/complete", + result: `{"resultType": "input_required", "requestState": "x"}`, + wantPayload: false, + }, + { + // An unrecognized resultType "MUST be considered invalid" per the + // spec; it keeps pure sentinel semantics with NO extractable round. + name: "unrecognized resultType carries no payload", + method: "tools/call", + result: `{"resultType": "partial"}`, + wantPayload: false, + }, + { + // "task" is the Tasks extension's CreateTaskResult discriminator, not + // an input round. vMCP advertises no extensions, so rejecting it here + // is correct — it lands on the same sentinel as a genuinely + // unrecognized value, even though the sentinel's "requires additional + // input" text names the wrong reason (message frozen for compat). + name: "Tasks-extension resultType carries no payload", + method: "tools/call", + result: `{"resultType": "task"}`, + wantPayload: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := interpretModernResult(json.RawMessage(tt.result), nil, tt.method, nil) + require.Error(t, err) + + // (a) classification is unchanged. + require.ErrorIs(t, err, errModernInputRequired) + + // (b) message is byte-identical to the previous wrapping. The + // expected text is hardcoded for the first case so the production + // template and this assertion cannot drift together; the rest derive + // it from the envelope like the old wrapping did. + if tt.wantMsg != "" { + assert.Equal(t, tt.wantMsg, err.Error()) + } else { + var envelope struct { + ResultType string `json:"resultType"` + } + require.NoError(t, json.Unmarshal([]byte(tt.result), &envelope)) + assert.Equal(t, + fmt.Errorf("%w: resultType=%q", errModernInputRequired, envelope.ResultType).Error(), + err.Error()) + } + + // (c) payload extraction, through an extra wrapping layer to mirror + // how modernCallTool wraps backend errors before they reach dispatch. + wrapped := fmt.Errorf("tool call failed on backend %s: %w", "b1", err) + round, ok := vmcp.InputRequiredFromError(wrapped) + require.Equal(t, tt.wantPayload, ok) + if !tt.wantPayload { + assert.Nil(t, round) + return + } + require.NotNil(t, round) + assert.Equal(t, tt.wantState, round.RequestState) + assert.Equal(t, tt.wantKeys, round.Methods()) + }) + } +} + +// TestInputRequiredErrorTruncatesHostileResultType pins that a +// backend-controlled resultType does not ride unbounded into the error text +// (which reaches the downstream client via writeModernDispatchError): beyond +// 64 bytes the message carries a prefix and the length, per the #6079/#6066 +// report-the-length rule. +func TestInputRequiredErrorTruncatesHostileResultType(t *testing.T) { + t.Parallel() + + hostile := strings.Repeat("A", 500) + result := fmt.Sprintf(`{"resultType": %q}`, hostile) + + err := interpretModernResult(json.RawMessage(result), nil, "tools/call", nil) + require.Error(t, err) + require.ErrorIs(t, err, errModernInputRequired) + assert.Equal(t, + "modern response requires additional input (multi-round retrieval unsupported): "+ + fmt.Sprintf("resultType=%q... (500 bytes)", strings.Repeat("A", 64)), + err.Error()) +} + +// TestInputRequiredFromErrorRejectsForeignErrors pins that the extractor is a +// safe single branch point: unrelated errors — including the OTHER Modern +// sentinels — report ok=false. +func TestInputRequiredFromErrorRejectsForeignErrors(t *testing.T) { + t.Parallel() + + for _, err := range []error{ + nil, + errWrongEra, + errLegacyResponseBody, + fmt.Errorf("wrapped: %w", errModernProtocolError), + } { + round, ok := vmcp.InputRequiredFromError(err) + assert.False(t, ok, "error %v must not extract as input_required", err) + assert.Nil(t, round) + } +} + +// TestInputRequiredRequestsAreRelayedVerbatim pins the pass-through +// invariant: the raw bytes of each input request survive the decode +// untouched (json.RawMessage preserves them exactly), so the relay (slice 2) +// can forward them without reinterpretation. +func TestInputRequiredRequestsAreRelayedVerbatim(t *testing.T) { + t.Parallel() + + const rawReq = `{"method":"elicitation/create","params":{"mode":"form","message":"x","extra":[1,2,3]}}` + result := `{"resultType":"input_required","inputRequests":{"k1":` + rawReq + `}}` + + err := interpretModernResult(json.RawMessage(result), nil, "tools/call", nil) + round, ok := vmcp.InputRequiredFromError(err) + require.True(t, ok) + require.Contains(t, round.InputRequests, "k1") + assert.Equal(t, rawReq, string(round.InputRequests["k1"])) +} diff --git a/pkg/vmcp/client/modern_test.go b/pkg/vmcp/client/modern_test.go new file mode 100644 index 0000000000..5aef8b7d24 --- /dev/null +++ b/pkg/vmcp/client/modern_test.go @@ -0,0 +1,452 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// completeEnvelope is a minimal valid Modern success body for a given method's +// result payload merged with the resultType envelope key. The id is always 1: +// these fakes serve the JSON (not SSE) path, which does not match by id. +func completeEnvelope(t *testing.T, payload map[string]any) []byte { + t.Helper() + result := map[string]any{"resultType": "complete"} + for k, v := range payload { + result[k] = v + } + body, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "result": result}) + require.NoError(t, err) + return body +} + +// TestModernCall_RequestShaping verifies the wire shape the shim produces: +// mandatory headers, conditional Mcp-Name, the reserved _meta keys, and the +// absence of any session header. +func TestModernCall_RequestShaping(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + method string + mcpName string + wantNameHdr string // expected Mcp-Name header value ("" = must be absent) + callerMeta map[string]any + }{ + { + name: "name-required method sets Mcp-Name", + method: "tools/call", + mcpName: "echo", + wantNameHdr: "echo", + }, + { + name: "non-name-required method omits Mcp-Name even when provided", + method: "tools/list", + mcpName: "ignored", + wantNameHdr: "", + }, + { + name: "empty name omits Mcp-Name", + method: "tools/call", + mcpName: "", + wantNameHdr: "", + }, + { + name: "caller reserved _meta keys are overridden by vMCP's", + method: "tools/list", + mcpName: "", + wantNameHdr: "", + callerMeta: map[string]any{ + "io.modelcontextprotocol/protocolVersion": "1999-01-01", + "io.modelcontextprotocol/clientCapabilities": "not-an-object", + "userKey": "preserved", + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var gotReq *http.Request + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotReq = r + gotBody, _ = readAll(t, r) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(completeEnvelope(t, map[string]any{})) + })) + t.Cleanup(srv.Close) + + params := map[string]any{} + if tt.callerMeta != nil { + params["_meta"] = tt.callerMeta + } + err := modernCall(context.Background(), srv.Client(), srv.URL, tt.method, params, tt.mcpName, nil, nil, "", nil) + require.NoError(t, err) + + assert.Equal(t, "application/json", gotReq.Header.Get("Content-Type")) + assert.Equal(t, "application/json, text/event-stream", gotReq.Header.Get("Accept")) + assert.Equal(t, "2026-07-28", gotReq.Header.Get("MCP-Protocol-Version")) + assert.Equal(t, tt.method, gotReq.Header.Get("Mcp-Method")) + assert.Equal(t, tt.wantNameHdr, gotReq.Header.Get("Mcp-Name")) + assert.Empty(t, gotReq.Header.Get("Mcp-Session-Id"), "Modern is stateless: never send a session id") + + // Verify the reserved _meta keys are present and authoritative. + var decoded struct { + Params struct { + Meta map[string]any `json:"_meta"` + } `json:"params"` + } + require.NoError(t, json.Unmarshal(gotBody, &decoded)) + meta := decoded.Params.Meta + assert.Equal(t, "2026-07-28", meta["io.modelcontextprotocol/protocolVersion"]) + assert.IsType(t, map[string]any{}, meta["io.modelcontextprotocol/clientCapabilities"]) + assert.Contains(t, meta, "io.modelcontextprotocol/clientInfo") + if tt.callerMeta != nil { + assert.Equal(t, "preserved", meta["userKey"], "non-reserved caller _meta keys survive") + } + }) + } +} + +// TestModernCall_CallerMetaNotMutated verifies the shim never writes through the +// caller's params or _meta map (go-style copy-before-mutate rule). +func TestModernCall_CallerMetaNotMutated(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(completeEnvelope(t, map[string]any{})) + })) + t.Cleanup(srv.Close) + + callerMeta := map[string]any{"userKey": "v"} + params := map[string]any{"_meta": callerMeta, "name": "x"} + require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "tools/list", params, "", nil, nil, "", nil)) + + assert.Equal(t, map[string]any{"userKey": "v"}, callerMeta, "caller _meta must be untouched") + assert.NotContains(t, params, "does-not-add-keys") + _, addedMeta := params["_meta"].(map[string]any) + assert.True(t, addedMeta) + assert.Len(t, callerMeta, 1, "no reserved keys leaked into caller's _meta") +} + +// TestModernCall_Decode verifies a complete envelope decodes into out. +func TestModernCall_Decode(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(completeEnvelope(t, map[string]any{ + "supportedVersions": []string{"2026-07-28"}, + })) + })) + t.Cleanup(srv.Close) + + var out struct { + ResultType string `json:"resultType"` + SupportedVersions []string `json:"supportedVersions"` + } + require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil, &out, "", nil)) + assert.Equal(t, "complete", out.ResultType) + assert.Equal(t, []string{"2026-07-28"}, out.SupportedVersions) +} + +// TestModernCall_SSEResponse verifies the dual-body reader handles a +// text/event-stream response, ignoring interleaved notifications and returning +// the final matching JSON-RPC response. +func TestModernCall_SSEResponse(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Echo the request's JSON-RPC id: modernCall's SSE reader matches the + // response frame by id (unlike the JSON path), and the id comes from a + // shared counter, so it cannot be hardcoded. + var req struct { + ID json.RawMessage `json:"id"` + } + body, _ := io.ReadAll(r.Body) + require.NoError(t, json.Unmarshal(body, &req)) + + w.Header().Set("Content-Type", "text/event-stream") + // A progress notification (has "method", no id match) then the response. + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{}}\n\n")) + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"id\":" + string(req.ID) + + ",\"result\":{\"resultType\":\"complete\",\"ok\":true}}\n\n")) + })) + t.Cleanup(srv.Close) + + var out map[string]any + require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil, &out, "", nil)) + assert.Equal(t, "complete", out["resultType"]) + assert.Equal(t, true, out["ok"]) +} + +// TestModernCall_LogLevelMeta verifies the logLevel argument is overlaid onto +// the minted request _meta as io.modelcontextprotocol/logLevel — the Modern +// (2026-07-28) replacement for the removed logging/setLevel RPC — and that an +// empty logLevel leaves the key unset (so a conformant backend MUST NOT emit +// notifications/message for the request). +func TestModernCall_LogLevelMeta(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + logLevel string + wantKey bool + wantLogLevel any + }{ + {name: "non-empty logLevel is injected", logLevel: "debug", wantKey: true, wantLogLevel: "debug"}, + {name: "empty logLevel omits the key", logLevel: "", wantKey: false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = readAll(t, r) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(completeEnvelope(t, map[string]any{})) + })) + t.Cleanup(srv.Close) + + require.NoError(t, modernCall( + context.Background(), srv.Client(), srv.URL, "tools/call", + map[string]any{"name": "x"}, "x", nil, nil, tt.logLevel, nil, + )) + + var decoded struct { + Params struct { + Meta map[string]any `json:"_meta"` + } `json:"params"` + } + require.NoError(t, json.Unmarshal(gotBody, &decoded)) + if tt.wantKey { + assert.Equal(t, tt.wantLogLevel, decoded.Params.Meta["io.modelcontextprotocol/logLevel"]) + } else { + assert.NotContains(t, decoded.Params.Meta, "io.modelcontextprotocol/logLevel") + } + }) + } +} + +// TestModernCall_SSENotificationRelay verifies that when an onNotification +// listener is bound, an interleaved server->client notification is handed to it +// (method + params) AND the matching response is still returned — the listener +// must not swallow the response the caller is waiting on. +func TestModernCall_SSENotificationRelay(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req struct { + ID json.RawMessage `json:"id"` + } + body, _ := io.ReadAll(r.Body) + require.NoError(t, json.Unmarshal(body, &req)) + + w.Header().Set("Content-Type", "text/event-stream") + // A log notification (method, no id) then the matching response. + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/message\"," + + "\"params\":{\"level\":\"info\",\"data\":\"hello\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"id\":" + string(req.ID) + + ",\"result\":{\"resultType\":\"complete\",\"ok\":true}}\n\n")) + })) + t.Cleanup(srv.Close) + + var gotMethod string + var gotParams map[string]any + onNotification := func(method string, params json.RawMessage) { + gotMethod = method + _ = json.Unmarshal(params, &gotParams) + } + + var out map[string]any + require.NoError(t, modernCall( + context.Background(), srv.Client(), srv.URL, "tools/call", + map[string]any{"name": "x"}, "x", nil, &out, "debug", onNotification, + )) + + assert.Equal(t, "notifications/message", gotMethod) + assert.Equal(t, "info", gotParams["level"]) + assert.Equal(t, "hello", gotParams["data"]) + assert.Equal(t, "complete", out["resultType"], "the response must still be delivered to the caller") +} + +// TestModernCall_ErrorMapping verifies the era/error classification: a valid +// -32601 body is method-not-found (backend IS Modern), a non-"complete" envelope +// is input-required, and non-Modern responses are wrong-era. +func TestModernCall_ErrorMapping(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + status int + contentType string + body string + wantErr error // errors.Is target; nil means "some error, but not a known sentinel" + wantMsg string // substring the error message must contain (checked when wantErr is nil) + }{ + { + name: "valid -32601 error body is method-not-found", + status: http.StatusNotFound, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"method not found"}}`, + wantErr: mcp.ErrMethodNotFound, + }, + { + name: "non-complete envelope is input-required", + status: http.StatusOK, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"result":{"resultType":"input_required"}}`, + wantErr: errModernInputRequired, + }, + { + name: "bare 404 with no JSON-RPC body is wrong-era", + status: http.StatusNotFound, + contentType: "text/plain", + body: "not found", + wantErr: errWrongEra, + }, + { + name: "empty body is wrong-era", + status: http.StatusOK, + contentType: "application/json", + body: "", + wantErr: errWrongEra, + }, + { + name: "200 with Legacy-shaped result (no resultType) is a legacy-response-body", + status: http.StatusOK, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}`, + wantErr: errLegacyResponseBody, + }, + { + name: "other JSON-RPC error surfaces as a call error, not wrong-era", + status: http.StatusOK, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"error":{"code":-32603,"message":"boom"}}`, + wantMsg: "boom", + }, + { + name: "401 is auth, not wrong-era", + status: http.StatusUnauthorized, + contentType: "text/plain", + body: "", + wantErr: errModernAuth, + }, + { + name: "403 is auth, not wrong-era", + status: http.StatusForbidden, + contentType: "application/json", + body: "", + wantErr: errModernAuth, + }, + { + name: "407 proxy-auth is auth, not wrong-era", + status: http.StatusProxyAuthRequired, + contentType: "text/plain", + body: "", + wantErr: errModernAuth, + }, + { + name: "503 is transient, not wrong-era", + status: http.StatusServiceUnavailable, + contentType: "text/plain", + body: "", + wantErr: errModernTransient, + }, + { + name: "429 is transient, not wrong-era", + status: http.StatusTooManyRequests, + contentType: "application/json", + body: "", + wantErr: errModernTransient, + }, + { + name: "401 tagged text/event-stream is still auth, not wrong-era", + status: http.StatusUnauthorized, + contentType: "text/event-stream", + body: "", + wantErr: errModernAuth, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", tt.contentType) + w.WriteHeader(tt.status) + _, _ = w.Write([]byte(tt.body)) + })) + t.Cleanup(srv.Close) + + err := modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil, nil, "", nil) + require.Error(t, err) + if tt.wantErr != nil { + assert.ErrorIs(t, err, tt.wantErr) + // Auth/transient statuses must never masquerade as the not-Modern + // signal (that is what would poison a cached-Modern backend). + if tt.wantErr == errModernAuth || tt.wantErr == errModernTransient { + assert.NotErrorIs(t, err, errWrongEra) + assert.NotErrorIs(t, err, errLegacyResponseBody) + } + return + } + // A valid JSON-RPC error body means the backend IS Modern: the error + // must surface (not be reclassified as wrong-era or method-not-found). + assert.NotErrorIs(t, err, errWrongEra) + assert.NotErrorIs(t, err, mcp.ErrMethodNotFound) + assert.Contains(t, err.Error(), tt.wantMsg) + }) + } +} + +// TestModernCall_LargeSSEEvent verifies a single SSE data: event larger than the +// old 4 MiB scanner cap (but under maxResponseSize) decodes successfully. +func TestModernCall_LargeSSEEvent(t *testing.T) { + t.Parallel() + + big := strings.Repeat("x", 5*1024*1024) // 5 MiB > old 4 MiB token cap + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + w.Header().Set("Content-Type", "text/event-stream") + out, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, + "result": map[string]any{"resultType": "complete", "blob": big}}) + require.NoError(t, err) + _, _ = w.Write([]byte("data: " + string(out) + "\n\n")) + })) + t.Cleanup(srv.Close) + + var got struct { + Blob string `json:"blob"` + } + require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil, &got, "", nil)) + assert.Len(t, got.Blob, len(big)) +} + +// readAll returns the request body bytes for assertions. +func readAll(t *testing.T, r *http.Request) ([]byte, error) { + t.Helper() + return io.ReadAll(r.Body) +} diff --git a/pkg/vmcp/client/reclassify_test.go b/pkg/vmcp/client/reclassify_test.go new file mode 100644 index 0000000000..5f3bf392c3 --- /dev/null +++ b/pkg/vmcp/client/reclassify_test.go @@ -0,0 +1,363 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" +) + +// TestIsRevisionMismatch is the narrow-predicate truth table. +func TestIsRevisionMismatch(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + rev mcpparser.Revision + err error + want bool + }{ + {"nil err", mcpparser.RevisionLegacy, nil, false}, + {"legacy INITIALIZE-step method-not-found -> mismatch", mcpparser.RevisionLegacy, + fmt.Errorf("%w: %w", errLegacyInitFailed, mcp.ErrMethodNotFound), true}, + {"legacy INITIALIZE-step 4xx (ErrLegacySSEServer) -> mismatch", mcpparser.RevisionLegacy, + fmt.Errorf("%w: %w", errLegacyInitFailed, transport.ErrLegacySSEServer), true}, + {"legacy DATA-PLANE method-not-found (no init marker) -> NOT mismatch", mcpparser.RevisionLegacy, + fmt.Errorf("completion failed: %w", mcp.ErrMethodNotFound), false}, + {"modern wrong-era -> mismatch", mcpparser.RevisionModern, + fmt.Errorf("wrap: %w", errWrongEra), true}, + {"modern legacy-response-body -> mismatch", mcpparser.RevisionModern, + fmt.Errorf("wrap: %w", errLegacyResponseBody), true}, + {"modern data-plane method-not-found -> NOT mismatch", mcpparser.RevisionModern, + fmt.Errorf("wrap: %w", mcp.ErrMethodNotFound), false}, + {"modern errLegacySSE (not a modern signal) -> NOT mismatch", mcpparser.RevisionModern, + fmt.Errorf("wrap: %w", transport.ErrLegacySSEServer), false}, + {"modern negotiated-down -> mismatch", mcpparser.RevisionModern, + fmt.Errorf("wrap: %w", errModernNegotiatedDown), true}, + {"legacy negotiated-down (not a legacy signal) -> NOT mismatch", mcpparser.RevisionLegacy, + fmt.Errorf("wrap: %w", errModernNegotiatedDown), false}, + {"auth: ErrUnauthorized -> NOT mismatch", mcpparser.RevisionLegacy, transport.ErrUnauthorized, false}, + {"auth: ErrAuthorizationRequired -> NOT mismatch", mcpparser.RevisionLegacy, transport.ErrAuthorizationRequired, false}, + {"auth: ErrUpstreamTokenNotFound -> NOT mismatch", mcpparser.RevisionLegacy, authtypes.ErrUpstreamTokenNotFound, false}, + {"auth: ErrAuthenticationFailed -> NOT mismatch", mcpparser.RevisionLegacy, vmcp.ErrAuthenticationFailed, false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, isRevisionMismatch(tt.rev, tt.err)) + }) + } +} + +// modernDiscoverServer serves Modern server/discover (advertising tools and +// supportedVersions containing 2026-07-28, the classification signal added by +// Fix 1) so a re-probe classifies Modern, and rejects Legacy requests (no +// Mcp-Method header) with 404 so a Legacy initialize surfaces +// transport.ErrLegacySSEServer. It also acks subscriptions/listen (see the +// case below) since go-sdk v1.7's own client — used by the Legacy code path, +// not just this repo's raw shim — is itself Modern-first and would otherwise +// die on that unconditional post-discover subscription. +func modernDiscoverServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + method := r.Header.Get("Mcp-Method") + if method == "" { + w.WriteHeader(http.StatusNotFound) // Legacy initialize -> ErrLegacySSEServer + return + } + id, _ := modernReq(t, r) + switch method { + case "server/discover": + writeModernResult(t, w, id, map[string]any{ + "capabilities": map[string]any{"tools": map[string]any{}}, + "supportedVersions": []string{"2026-07-28"}, + }) + case "tools/list": + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "echo", "inputSchema": map[string]any{"type": "object"}}}, + }) + case "subscriptions/listen": + // go-sdk v1.7's Connect() unconditionally subscribes to list-changed + // notifications once discover negotiates Modern (mcpcompat always + // installs a ToolListChangedHandler), independent of our own shim's + // raw discover probe. Ack it immediately so the real underlying + // client session used by legacyInit doesn't die. + writeModernResult(t, w, id, map[string]any{}) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + t.Cleanup(srv.Close) + return srv +} + +// TestDispatch_ReclassifyAndRetry: a backend mis-cached Legacy but actually +// Modern flips after one re-probe and the retry succeeds. +func TestDispatch_ReclassifyAndRetry(t *testing.T) { + t.Parallel() + + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) // mis-cached + + var attempts []mcpparser.Revision + err := h.dispatch(context.Background(), target, func(_ context.Context, rev mcpparser.Revision) error { + attempts = append(attempts, rev) + if rev == mcpparser.RevisionLegacy { + // Legacy initialize-step rejection (carries the init marker). + return fmt.Errorf("%w: %w", errLegacyInitFailed, + wrapBackendError(transport.ErrLegacySSEServer, target.WorkloadID, "initialize client")) + } + return nil + }) + require.NoError(t, err) + assert.Equal(t, []mcpparser.Revision{mcpparser.RevisionLegacy, mcpparser.RevisionModern}, attempts, + "exactly one retry, under the flipped revision") + + rev, _ := h.cachedRevision(target.WorkloadID) + assert.Equal(t, mcpparser.RevisionModern, rev, "cache updated to the corrected revision") +} + +// TestDispatch_NoRetryOnDataPlaneMethodNotFound: a genuine method-not-found on a +// correctly-classified Modern backend must not trigger a re-probe. +func TestDispatch_NoRetryOnDataPlaneMethodNotFound(t *testing.T) { + t.Parallel() + + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + + attempts := 0 + err := h.dispatch(context.Background(), target, func(_ context.Context, _ mcpparser.Revision) error { + attempts++ + return fmt.Errorf("tool missing: %w", mcp.ErrMethodNotFound) + }) + require.Error(t, err) + assert.Equal(t, 1, attempts, "data-plane method-not-found must not retry") +} + +// TestDispatch_NoInfiniteLoop: even if fn always fails with a mismatch and the +// re-probe flips the era, fn runs at most twice. +func TestDispatch_NoInfiniteLoop(t *testing.T) { + t.Parallel() + + // This server rejects everything (no Modern discover), so a re-probe from a + // Modern cache classifies Legacy — a flip that would loop without the guard. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + + attempts := 0 + err := h.dispatch(context.Background(), target, func(_ context.Context, _ mcpparser.Revision) error { + attempts++ + return fmt.Errorf("always wrong era: %w", errWrongEra) + }) + require.Error(t, err) + assert.Equal(t, 2, attempts, "one original + at most one retry, never a loop") +} + +// TestDispatch_NoDoubleExecOnLegacyBody: a lenient Legacy backend that EXECUTES a +// Modern tools/call and returns a Legacy-shaped 200 body must NOT be retried +// (no double-execution), but the cache is still flipped to the corrected era. +func TestDispatch_NoDoubleExecOnLegacyBody(t *testing.T) { + t.Parallel() + + var toolCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + if r.Header.Get("Mcp-Method") == "tools/call" { + toolCalls.Add(1) + // Legacy-shaped success: JSON-RPC result WITHOUT resultType (backend executed it). + w.Header().Set("Content-Type", "application/json") + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, + "result": map[string]any{"content": []any{map[string]any{"type": "text", "text": "done"}}}}) + _, _ = w.Write(out) + return + } + // server/discover during re-probe: reject so the backend classifies Legacy. + w.WriteHeader(http.StatusNotFound) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) // mis-cached Modern + + _, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "x"}, nil, nil) + require.Error(t, err, "a Legacy-shaped body must surface as an error, not a blank success") + assert.EqualValues(t, 1, toolCalls.Load(), "the side-effecting tool must run exactly once (no double-exec)") + + rev, _ := h.cachedRevision(target.WorkloadID) + assert.Equal(t, mcpparser.RevisionLegacy, rev, "cache flipped for future calls despite no retry") +} + +// TestDispatch_NoReprobeOnLegacyDataPlaneMethodNotFound: a data-plane -32601 from +// a genuine Legacy backend (a legitimately unimplemented method) must not +// reclassify or re-probe. +func TestDispatch_NoReprobeOnLegacyDataPlaneMethodNotFound(t *testing.T) { + t.Parallel() + + // This server WOULD classify Modern if probed — so if a re-probe wrongly + // fired, the cache would flip to Modern. It must stay Legacy. + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + attempts := 0 + err := h.dispatch(context.Background(), target, func(_ context.Context, _ mcpparser.Revision) error { + attempts++ + // Data-plane -32601 (no errLegacyInitFailed marker), e.g. completion/complete. + return fmt.Errorf("completion failed: %w", mcp.ErrMethodNotFound) + }) + require.Error(t, err) + assert.Equal(t, 1, attempts, "data-plane -32601 must not retry") + + rev, _ := h.cachedRevision(target.WorkloadID) + assert.Equal(t, mcpparser.RevisionLegacy, rev, "data-plane -32601 must not re-probe/reclassify") +} + +// TestDispatch_TransientDoesNotReclassify verifies an auth blip or transient +// outage on a cached-Modern backend neither retries nor flips the cache to Legacy +// (F1: status-blind errWrongEra used to poison the cache here). +func TestDispatch_TransientDoesNotReclassify(t *testing.T) { + t.Parallel() + + for _, blip := range []error{errModernAuth, errModernTransient} { + blip := blip + t.Run(blip.Error(), func(t *testing.T) { + t.Parallel() + + // This server WOULD classify Modern if re-probed — proving no re-probe fired. + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + + attempts := 0 + err := h.dispatch(context.Background(), target, func(_ context.Context, _ mcpparser.Revision) error { + attempts++ + return fmt.Errorf("%w: HTTP blip", blip) + }) + require.Error(t, err) + assert.Equal(t, 1, attempts, "a transient/auth blip must not retry") + + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev, "a blip must not flip a Modern backend to Legacy") + }) + } +} + +// TestReclassify_WarnsOnlyOnActualChange captures slog to confirm the WARN (which +// gates the reclassification counter in the same branch) fires only when the +// revision actually changes. +// +//nolint:paralleltest // swaps the global slog default; must not run concurrently with other tests +func TestReclassify_WarnsOnlyOnActualChange(t *testing.T) { + // Not parallel: swaps the global slog default. + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + srv := modernDiscoverServer(t) // re-probe classifies Modern + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + + // prev=Legacy -> re-probe Modern: change, must WARN. + got := h.reclassify(context.Background(), target, mcpparser.RevisionLegacy) + assert.Equal(t, mcpparser.RevisionModern, got) + assert.Contains(t, buf.String(), "reclassified", "a real change must WARN") + + // prev=Modern -> re-probe Modern: no change, must be silent. + buf.Reset() + got = h.reclassify(context.Background(), target, mcpparser.RevisionModern) + assert.Equal(t, mcpparser.RevisionModern, got) + assert.Empty(t, buf.String(), "a no-op re-probe must not WARN") +} + +// TestListCapabilities_MisCachedLegacy_SelfHealsViaSDKNegotiation: a backend +// pinned Legacy by a stale cache is still queried successfully, because the +// underlying go-sdk v1.7 client (used by the Legacy/session-based code path, +// legacyListCapabilities) is itself Modern-first (SEP-2575): its Connect() tries +// server/discover before any legacy initialize, so a genuinely Modern backend +// answers correctly even when dispatch hands it the stale Legacy revision. The +// call succeeds on the first attempt, so dispatch's reclassify-on-mismatch path +// (which only fires on error; see dispatch's doc comment) is never entered — +// but legacyInit reads the genuinely-negotiated protocol version off the +// Initialize result and flips the cache to Modern in-band regardless (see the +// revisions field doc in client.go). +// +// The staleness this closes is not purely internal bookkeeping: CachedRevision +// surfaces it to health status (pkg/vmcp/health/status.go:260), which +// republishes it on every health tick into the mcpRevision CRD status field — +// so before this self-heal, an operator reading that field would see "legacy" +// for a backend actually being driven Modern on every call. +func TestListCapabilities_MisCachedLegacy_SelfHealsViaSDKNegotiation(t *testing.T) { + t.Parallel() + + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) // mis-cached + + caps, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err, "the Legacy code path's go-sdk client negotiates Modern transparently") + require.Len(t, caps.Tools, 1) + assert.Equal(t, "echo", caps.Tools[0].Name) + + // legacyInit flips the cache directly off the negotiated Initialize result, + // without needing an error to trigger dispatch's reclassify path. + rev, _ := h.cachedRevision(target.WorkloadID) + assert.Equal(t, mcpparser.RevisionModern, rev) +} + +// TestLegacyInit_SelfHealsRevisionCache exercises legacyInit directly (rather +// than through the full ListCapabilities/query pipeline) to isolate the +// self-heal contract: a mis-cached Legacy backend that is genuinely Modern +// ends up cached Modern after a single Legacy-path Initialize call. +func TestLegacyInit_SelfHealsRevisionCache(t *testing.T) { + t.Parallel() + + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) // mis-cached + + c, err := h.clientFactory(context.Background(), target, false) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, _, err = h.legacyInit(context.Background(), c, target) + require.NoError(t, err, "the SDK client negotiates Modern transparently on the Legacy path") + + rev, _ := h.cachedRevision(target.WorkloadID) + assert.Equal(t, mcpparser.RevisionModern, rev) +} diff --git a/pkg/vmcp/client/request_timeout_test.go b/pkg/vmcp/client/request_timeout_test.go index 3c11fafe0f..9f087f3588 100644 --- a/pkg/vmcp/client/request_timeout_test.go +++ b/pkg/vmcp/client/request_timeout_test.go @@ -89,7 +89,7 @@ func TestWithRequestTimeouts_CutsSlowBackendCall(t *testing.T) { } start := time.Now() - _, err = bc.CallTool(t.Context(), target, "any", nil, nil) + _, err = bc.CallTool(t.Context(), target, "any", nil, nil, nil) require.Error(t, err) assert.Less(t, time.Since(start), 10*time.Second, "call must end at the per-workload timeout") } diff --git a/pkg/vmcp/client/resource_templates_integration_test.go b/pkg/vmcp/client/resource_templates_integration_test.go new file mode 100644 index 0000000000..2b94f9aafd --- /dev/null +++ b/pkg/vmcp/client/resource_templates_integration_test.go @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client_test + +import ( + "context" + "encoding/json" + "io" + "net" + "net/http" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" +) + +// TestListCapabilities_IncludesResourceTemplates verifies that a backend +// advertising a resource template surfaces it in the aggregated CapabilityList, +// as a pass-through (no URI-template rewriting) tagged with the backend ID. +func TestListCapabilities_IncludesResourceTemplates(t *testing.T) { + t.Parallel() + + port, cleanup := startResourceTemplateMCPServer(t) + defer cleanup() + + registry := auth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, registry.RegisterStrategy("unauthenticated", &strategies.UnauthenticatedStrategy{})) + + backendClient, err := vmcpclient.NewHTTPBackendClient(registry) + require.NoError(t, err) + + target := &vmcp.BackendTarget{ + WorkloadID: "docs-backend", + WorkloadName: "Docs Backend", + BaseURL: "http://127.0.0.1:" + port, + TransportType: "streamable-http", + } + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + caps, err := backendClient.ListCapabilities(ctx, target) + require.NoError(t, err) + require.NotNil(t, caps) + + require.Len(t, caps.ResourceTemplates, 1, "the advertised resource template must be surfaced") + tmpl := caps.ResourceTemplates[0] + assert.Equal(t, "file:///logs/{date}.txt", tmpl.URITemplate) + assert.Equal(t, "Daily log", tmpl.Name) + assert.Equal(t, "text/plain", tmpl.MimeType) + assert.Equal(t, "docs-backend", tmpl.BackendID, "the template must be tagged with the backend ID") +} + +// TestComplete_NoCompletionsCapability verifies the spec-lenient completion path: +// a backend that does NOT advertise the completions capability (this fixture only +// registers a resource template, no completion handler) yields a non-nil, EMPTY +// completion result and no error. The client must not surface a hard failure for a +// backend that simply cannot complete. +func TestComplete_NoCompletionsCapability(t *testing.T) { + t.Parallel() + + port, cleanup := startResourceTemplateMCPServer(t) + defer cleanup() + + registry := auth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, registry.RegisterStrategy("unauthenticated", &strategies.UnauthenticatedStrategy{})) + + backendClient, err := vmcpclient.NewHTTPBackendClient(registry) + require.NoError(t, err) + + target := &vmcp.BackendTarget{ + WorkloadID: "docs-backend", + WorkloadName: "Docs Backend", + BaseURL: "http://127.0.0.1:" + port, + TransportType: "streamable-http", + } + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + res, err := backendClient.Complete(ctx, target, + vmcp.CompletionRef{Type: vmcp.CompletionRefTypePrompt, Name: "greeting"}, + "name", "wor", nil) + require.NoError(t, err, "a backend without completions must not error (spec-lenient)") + require.NotNil(t, res, "the completion result must be non-nil") + assert.Empty(t, res.Values, "the completion result must carry an empty value set") +} + +// startResourceTemplateMCPServer starts a test MCP server advertising a single +// resource template and serving its expanded URIs via the template handler. +func startResourceTemplateMCPServer(t *testing.T) (string, func()) { + t.Helper() + + mcpServer := server.NewMCPServer("docs-backend", "1.0.0") + + mcpServer.AddResourceTemplate( + mcp.ResourceTemplate{ + URITemplate: "file:///logs/{date}.txt", + Name: "Daily log", + Description: "A day's log file", + MIMEType: "text/plain", + }, + func(_ context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + return []mcp.ResourceContents{ + mcp.TextResourceContents{ + URI: req.Params.URI, + MIMEType: "text/plain", + Text: "log contents for " + req.Params.URI, + }, + }, nil + }, + ) + + httpHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + rawMessage, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, "Failed to read request", http.StatusBadRequest) + return + } + defer r.Body.Close() + response := mcpServer.HandleMessage(r.Context(), rawMessage) + responseBytes, err := json.Marshal(response) + if err != nil { + http.Error(w, "Failed to marshal response", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(responseBytes) + }) + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + port := listener.Addr().(*net.TCPAddr).Port + httpServer := &http.Server{Handler: httpHandler} + go func() { _ = httpServer.Serve(listener) }() + + cleanup := func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = httpServer.Shutdown(ctx) + } + return strconv.Itoa(port), cleanup +} diff --git a/pkg/vmcp/client/revision_realbackend_test.go b/pkg/vmcp/client/revision_realbackend_test.go new file mode 100644 index 0000000000..f6ecf29d17 --- /dev/null +++ b/pkg/vmcp/client/revision_realbackend_test.go @@ -0,0 +1,538 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" +) + +// newRealEchoServer stands up a real go-sdk v1.7-backed streamable-HTTP MCP +// server (via mcpcompat, the same SDK backends run in production) exposing a +// single "echo" tool that echoes back its "input" argument. stateless controls +// whether the server is built with mcpserver.WithStateless(true) — the same +// knob that determines whether the backend answers server/discover with +// 2026-07-28 in supportedVersions (Modern) or negotiates it away (Legacy); see +// TestProbeRevision_RealBackends. +func newRealEchoServer( + t *testing.T, stateless bool, onCallTool func(mcpmcp.CallToolRequest), srvOpts ...mcpserver.ServerOption, +) *httptest.Server { + t.Helper() + + mcpSrv := mcpserver.NewMCPServer("real-backend", "1.0.0", srvOpts...) + mcpSrv.AddTool( + mcpmcp.NewTool("echo", + mcpmcp.WithDescription("Echoes the input back"), + mcpmcp.WithString("input", mcpmcp.Required()), + ), + func(_ context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + if onCallTool != nil { + onCallTool(req) + } + args, _ := req.Params.Arguments.(map[string]any) + input, _ := args["input"].(string) + return &mcpmcp.CallToolResult{Content: []mcpmcp.Content{mcpmcp.NewTextContent(input)}}, nil + }, + ) + + var opts []mcpserver.StreamableHTTPOption + if stateless { + opts = append(opts, mcpserver.WithStateless(true)) + } + + mux := http.NewServeMux() + mux.Handle("/mcp", mcpserver.NewStreamableHTTPServer(mcpSrv, opts...)) + // httptest.NewServer binds an OS-assigned (random) port. + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + return ts +} + +// TestProbeRevision_RealBackends is the anti-regression pin for Fix 1: it +// exercises probeRevision against a REAL go-sdk v1.7 backend (not a hand-rolled +// httptest fake), stateful and stateless, so a future go-sdk change to +// server/discover's supportedVersions semantics fails HERE with a clear signal +// instead of surfacing as dozens of mystery failures elsewhere. +// +// A stateful streamable-HTTP server negotiates DOWN on server/discover (its +// supportedVersions excludes 2026-07-28 — see go-sdk's discover(), which +// requires WithStateless(true) to advertise the Modern revision), so it must +// classify Legacy. A stateless server advertises 2026-07-28 and must classify +// Modern. +func TestProbeRevision_RealBackends(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + stateless bool + wantRev mcpparser.Revision + }{ + {name: "stateful streamable-HTTP backend negotiates down -> Legacy", stateless: false, wantRev: mcpparser.RevisionLegacy}, + {name: "stateless streamable-HTTP backend advertises 2026-07-28 -> Modern", stateless: true, wantRev: mcpparser.RevisionModern}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := newRealEchoServer(t, tt.stateless, nil) + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "real-backend", + WorkloadName: "Real Backend", + BaseURL: srv.URL + "/mcp", + TransportType: "streamable-http", + } + + rev, err := h.probeRevision(context.Background(), target) + require.NoError(t, err) + assert.Equal(t, tt.wantRev, rev) + + cached, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, tt.wantRev, cached) + }) + } +} + +// TestLegacyCallTool_StripsReservedMeta_RealBackend pins Fix 5 end-to-end +// against a real go-sdk v1.7 stateful (Legacy) backend: a legacy CallTool +// carrying both reserved io.modelcontextprotocol/* _meta keys (as a downstream +// Modern caller's request would) and a custom caller key must succeed — the +// reserved keys would otherwise make the real backend reject the request +// outright (HTTP 400: "protocol version ... is only supported on stateless +// HTTP servers") — and the backend must actually receive the custom key but +// NOT the reserved ones. +func TestLegacyCallTool_StripsReservedMeta_RealBackend(t *testing.T) { + t.Parallel() + + var ( + mu sync.Mutex + gotMeta map[string]any + sawCall bool + ) + srv := newRealEchoServer(t, false, func(req mcpmcp.CallToolRequest) { + mu.Lock() + defer mu.Unlock() + sawCall = true + if req.Params.Meta != nil { + gotMeta = req.Params.Meta.AdditionalFields + } + }) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "real-backend", + WorkloadName: "Real Backend", + BaseURL: srv.URL + "/mcp", + TransportType: "streamable-http", + } + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + callerMeta := map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": map[string]any{"name": "downstream-modern-client", "version": "1.0.0"}, + "io.modelcontextprotocol/clientCapabilities": map[string]any{}, + "io.modelcontextprotocol/logLevel": "debug", + "custom-caller-key": "custom-value", + } + + res, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hello legacy"}, callerMeta, nil) + require.NoError(t, err, "reserved Modern _meta must not leak onto the Legacy backend hop") + require.Len(t, res.Content, 1) + assert.Equal(t, "hello legacy", res.Content[0].Text) + + mu.Lock() + defer mu.Unlock() + require.True(t, sawCall, "the backend's tool handler must have been invoked") + // Assert on the namespace, not a fixed list: this catches any reserved key, + // including ones added to the fixture later. No passthroughMetaKeys entry is + // in play here, so a blanket prefix check is exact. + for k := range gotMeta { + assert.False(t, strings.HasPrefix(k, mcpparser.ReservedMetaPrefix), + "reserved _meta key %q must be stripped before the Legacy hop", k) + } + assert.Equal(t, "custom-value", gotMeta["custom-caller-key"], "non-reserved caller _meta must survive") +} + +// stubClientNotifier is a minimal ClientNotifier for BindForwarders: the +// tools/call path only reads the bound requesters to decide WHICH forwarding +// handlers to install — the real backend here never notifies mid-call, so it is +// never invoked. (stubElicitationRequester/stubSamplingRequester live in +// client_test.go.) +type stubClientNotifier struct{} + +func (stubClientNotifier) NotifyProgress(context.Context, vmcp.ProgressNotification) error { + return nil +} +func (stubClientNotifier) NotifyLog(context.Context, vmcp.LogMessage) error { return nil } + +// TestCallTool_MisCachedLegacy_ForwardingAgainstStatelessBackend pins the +// forwarding tools/call arm of the "stale Legacy cache is harmless" property +// (issue #5992), previously covered only for ListCapabilities +// (forwarding=false). A backend mis-cached Legacy that is genuinely +// stateless-Modern is queried through the FORWARDING path: forwarders are +// bound, so newStreamableHTTPClient enables transport.WithContinuousListening +// — the standalone server->client GET stream that, opened against a session-less +// backend, is exactly what made the POOLED session-factory connect fail in +// #6006 (subscriptions/listen -> "session not found"). +// +// The per-call path survives because of WHERE the skip lives: go-sdk's +// streamableClientConn.sessionUpdated opens the standalone stream ONLY when the +// negotiated protocol version is < 2026-07-28 (mcp/streamable.go). On the +// mis-cached-Legacy hop the SDK's Modern-first client negotiates Modern off the +// backend, so the version gate suppresses the standalone stream entirely — no +// subscriptions/listen is ever attempted, Connect succeeds, and the call +// round-trips. legacyInit then flips the cache to Modern in-band, so the next +// call skips the legacy handshake. +// +// This is the behavior a regression in the version gate (or a shim change that +// opens the stream unconditionally) would silently break: the call would start +// failing only on the forwarding path, invisible to the ListCapabilities pin. +func TestCallTool_MisCachedLegacy_ForwardingAgainstStatelessBackend(t *testing.T) { + t.Parallel() + + srv := newRealEchoServer(t, true, nil) // stateless => Modern + + h := newProbeClient(t) + h.BindForwarders(&stubElicitationRequester{}, &stubSamplingRequester{}, stubClientNotifier{}) + + target := &vmcp.BackendTarget{ + WorkloadID: "real-backend", + WorkloadName: "Real Backend", + BaseURL: srv.URL + "/mcp", + TransportType: "streamable-http", + } + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) // mis-cached + + res, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hello forwarding"}, nil, nil) + require.NoError(t, err, + "the negotiated-Modern version gate must suppress the standalone stream so the per-call forwarding path survives") + require.Len(t, res.Content, 1) + assert.Equal(t, "hello forwarding", res.Content[0].Text) + + // legacyInit flips the cache off the genuinely-negotiated Initialize result. + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev, + "a successful mis-cached-Legacy call self-heals the cache to Modern") +} + +// newRecordingEchoServer is a real go-sdk backend that advertises the logging +// capability and reports every JSON-RPC method it receives. +func newRecordingEchoServer(t *testing.T, onMethod func(method string)) *httptest.Server { + t.Helper() + + backend := newRealEchoServer(t, false, nil, mcpserver.WithLogging()) + backendURL, err := url.Parse(backend.URL) + require.NoError(t, err) + rp := httputil.NewSingleHostReverseProxy(backendURL) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, readErr := io.ReadAll(r.Body) + require.NoError(t, readErr) + r.Body = io.NopCloser(bytes.NewReader(body)) + + var probe struct { + Method string `json:"method"` + } + _ = json.Unmarshal(body, &probe) + if probe.Method != "" && onMethod != nil { + onMethod(probe.Method) + } + rp.ServeHTTP(w, r) + })) + t.Cleanup(ts.Close) + return ts +} + +// newHintLyingServer emulates github-mcp-server v1.6.0's dual-era behaviour: a +// backend that negotiates 2026-07-28 on the Legacy initialize handshake while +// answering server/discover with a body that is NOT a valid Modern envelope +// (no resultType), which modernCall rejects as Legacy-shaped. +// +// It delegates everything except server/discover to a real stateless go-sdk +// backend -- which is what makes the handshake genuinely negotiate 2026-07-28 +// rather than a hand-rolled approximation -- and substitutes the resultType-less +// discover body that produces the mismatch. +func newHintLyingServer(t *testing.T) *httptest.Server { + t.Helper() + return newHintLyingServerRecording(t, nil) +} + +// newHintLyingServerRecording is newHintLyingServer with a hook invoked for +// every JSON-RPC method the backend receives. +func newHintLyingServerRecording(t *testing.T, onMethod func(method string)) *httptest.Server { + t.Helper() + + // WithLogging so the backend advertises the logging capability: without it + // enableBackendLogging returns before ever reaching the version check, and + // a test asserting setLevel is not sent would pass vacuously. + backend := newRealEchoServer(t, true, nil, mcpserver.WithLogging()) // stateless => negotiates 2026-07-28 + backendURL, err := url.Parse(backend.URL) + require.NoError(t, err) + rp := httputil.NewSingleHostReverseProxy(backendURL) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, readErr := io.ReadAll(r.Body) + require.NoError(t, readErr) + r.Body = io.NopCloser(bytes.NewReader(body)) + + if onMethod != nil { + var probe struct { + Method string `json:"method"` + } + _ = json.Unmarshal(body, &probe) + if probe.Method != "" { + onMethod(probe.Method) + } + } + + if !bytes.Contains(body, []byte(`"server/discover"`)) { + rp.ServeHTTP(w, r) + return + } + + var req struct { + ID json.RawMessage `json:"id"` + } + _ = json.Unmarshal(body, &req) + id := req.ID + if len(id) == 0 { + id = json.RawMessage(`null`) + } + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"jsonrpc":"2.0","id":%s,"result":{`+ + `"capabilities":{"tools":{}},`+ + `"supportedVersions":["2026-07-28","2025-11-25"],`+ + `"serverInfo":{"name":"hint-liar","version":"1.0.0"}}}`, id) + })) + t.Cleanup(ts.Close) + return ts +} + +// TestLegacyInit_DoesNotPromoteOnHandshakeHintAlone pins #6154. +// +// The backend negotiates 2026-07-28 on the Legacy initialize -- the SDK client +// offers its own latest and the dual-era server accepts it -- while its +// server/discover answer is not a valid Modern envelope. github-mcp-server +// v1.6.0 behaves exactly this way in production. +// +// Promoting the cache to Modern on that handshake hint alone made the two +// self-corrections fight: probeRevision cached Legacy, legacyInit promoted back +// to Modern, the next Modern call failed on the non-Modern discover body and +// reclassified to Legacy, and so on. Every other health check failed, forever. +// The cached revision must therefore stay Legacy across repeated calls. +func TestLegacyInit_DoesNotPromoteOnHandshakeHintAlone(t *testing.T) { + t.Parallel() + + srv := newHintLyingServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "b", + BaseURL: srv.URL + "/mcp", + TransportType: "streamable-http", + } + + for i := range 3 { + _, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err, "ListCapabilities call %d", i) + + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok, "revision should be cached after call %d", i) + assert.Equal(t, mcpparser.RevisionLegacy, rev, + "cached revision must stay Legacy after call %d; flipping to Modern is the #6154 oscillation", i) + } +} + +// TestEnableBackendLogging_SkipsRemovedRPCOnModernSession pins the fix for +// tool calls dying against a dual-era backend. +// +// A backend that negotiates 2026-07-28 over the LEGACY initialize handshake +// leaves the session obliged to carry that version on every request. go-sdk +// still sends logging/setLevel as a Legacy-shaped call, so it arrives with a +// Modern protocol header and no Modern _meta -- a shape ToolHive's classifier +// rejects with -32020 by a deliberate, pinned contract (see +// TestIntegration_Modern_RealBackend_LoggingContract in pkg/vmcp/server). +// go-sdk treats that rejection as fatal, so the session closes and the tool +// call the logging was only meant to decorate dies with it. +// +// The RPC was removed in 2026-07-28, so issuing it on a session that +// negotiated that revision is the defect. Legacy sessions must still get it. +func TestEnableBackendLogging_SkipsRemovedRPCOnModernSession(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + negotiated string + wantSetLevelSent bool + }{ + { + name: "modern negotiated: removed RPC is skipped", + negotiated: mcpparser.MCPVersionModern, + wantSetLevelSent: false, + }, + { + name: "legacy negotiated: RPC is still sent", + negotiated: mcpparser.MCPVersionLegacy, + wantSetLevelSent: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var sawSetLevel bool + srv := newRecordingEchoServer(t, func(method string) { + if method == "logging/setLevel" { + mu.Lock() + sawSetLevel = true + mu.Unlock() + } + }) + + h := newProbeClient(t) + h.BindForwarders(&stubElicitationRequester{}, &stubSamplingRequester{}, stubClientNotifier{}) + target := &vmcp.BackendTarget{ + WorkloadID: "b", + BaseURL: srv.URL + "/mcp", + TransportType: "streamable-http", + } + + c, err := h.clientFactory(context.Background(), target, true) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + // The session must be live before SetLoggingLevel can reach the + // wire; an unconnected client fails inside the SDK instead. + _, _, err = h.legacyInit(context.Background(), c, target) + require.NoError(t, err) + + caps := &mcpmcp.ServerCapabilities{Logging: &struct{}{}} + h.enableBackendLogging(context.Background(), c, caps, tt.negotiated, target.WorkloadID) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, tt.wantSetLevelSent, sawSetLevel) + }) + } +} + +// TestLegacyInit_RefutedHintExpires covers #5992: a backend that lied once +// (Modern on the Legacy initialize, Legacy-shaped discover) is recorded in +// modernHintRefuted so the confirming probe does not run on every call. That +// memory must NOT be permanent: if the backend is later redeployed genuinely +// Modern (hint and discover now agree), the refutation expires after +// refutationTTL and the next call's confirming probe flips the cache. +// +// The redeeming server flips its discover answer mid-test: Legacy-shaped until +// redeemed.Store(true), then a valid Modern envelope. +func TestLegacyInit_RefutedHintExpires(t *testing.T) { + t.Parallel() + + var redeemed atomic.Bool + backend := newRealEchoServer(t, true, nil) // stateless => negotiates 2026-07-28 + backendURL, err := url.Parse(backend.URL) + require.NoError(t, err) + rp := httputil.NewSingleHostReverseProxy(backendURL) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, readErr := io.ReadAll(r.Body) + require.NoError(t, readErr) + r.Body = io.NopCloser(bytes.NewReader(body)) + + if !bytes.Contains(body, []byte(`"server/discover"`)) { + rp.ServeHTTP(w, r) + return + } + if redeemed.Load() { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(discoverEnvelope(t, r)) + return + } + + var req struct { + ID json.RawMessage `json:"id"` + } + _ = json.Unmarshal(body, &req) + id := req.ID + if len(id) == 0 { + id = json.RawMessage(`null`) + } + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"jsonrpc":"2.0","id":%s,"result":{`+ + `"capabilities":{"tools":{}},`+ + `"supportedVersions":["2026-07-28","2025-11-25"],`+ + `"serverInfo":{"name":"hint-liar","version":"1.0.0"}}}`, id) + })) + t.Cleanup(ts.Close) + + reg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, reg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, &strategies.UnauthenticatedStrategy{})) + c, err := NewHTTPBackendClient(reg, WithRefutationTTL(50*time.Millisecond)) + require.NoError(t, err) + h := c.(*httpBackendClient) + + target := &vmcp.BackendTarget{ + WorkloadID: "b", + BaseURL: ts.URL + "/mcp", + TransportType: "streamable-http", + } + + // Phase 1: the backend lies. The refutation is recorded and the cache + // stays Legacy (#6154 behavior is preserved within the window). + _, err = h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + require.Equal(t, mcpparser.RevisionLegacy, rev) + _, refuted := h.modernHintRefuted.Load(target.WorkloadID) + require.True(t, refuted, "the refutation should be recorded after the confirming probe disagrees") + + // Phase 2: the backend is "redeployed" genuinely Modern. Within the TTL + // the refutation still suppresses the confirming probe... + _, err = h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + rev, _ = h.cachedRevision(target.WorkloadID) + require.Equal(t, mcpparser.RevisionLegacy, rev, + "within the TTL the refutation still suppresses re-probing") + + // Phase 3: ...but once the window lapses the next call re-probes and the + // cache self-heals to Modern. + redeemed.Store(true) + time.Sleep(2 * 50 * time.Millisecond) + + _, err = h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + rev, ok = h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev, + "after the refutation expires, a redeemed backend must self-heal to Modern") +} diff --git a/pkg/vmcp/client/revision_test.go b/pkg/vmcp/client/revision_test.go new file mode 100644 index 0000000000..2f8e50603a --- /dev/null +++ b/pkg/vmcp/client/revision_test.go @@ -0,0 +1,323 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" +) + +// newProbeClient builds a real httpBackendClient with an unauthenticated +// registry so probeRevision's buildBackendRoundTripper succeeds. +func newProbeClient(t *testing.T) *httpBackendClient { + t.Helper() + reg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, reg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, &strategies.UnauthenticatedStrategy{})) + c, err := NewHTTPBackendClient(reg) + require.NoError(t, err) + return c.(*httpBackendClient) +} + +// discoverEnvelope is a valid Modern server/discover success body echoing the +// request id, advertising supportedVersions that include 2026-07-28 (the +// authoritative Modern signal — see errModernNegotiatedDown). +func discoverEnvelope(t *testing.T, r *http.Request) []byte { + t.Helper() + body, _ := readAll(t, r) + var req struct { + ID any `json:"id"` + } + require.NoError(t, json.Unmarshal(body, &req)) + out, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]any{ + "resultType": "complete", + "capabilities": map[string]any{"tools": map[string]any{}, "completions": map[string]any{}}, + "supportedVersions": []string{"2026-07-28", "2025-11-25"}, + }, + }) + require.NoError(t, err) + return out +} + +// TestProbeRevision_TruthTable exercises the Modern-first probe's classification: +// only a clean discover or a Modern-specific protocol error (-3202x) yields +// Modern; every other backend response falls back to Legacy. +func TestProbeRevision_TruthTable(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + handler http.HandlerFunc + wantRev mcpparser.Revision + }{ + { + name: "clean 2xx discover, supportedVersions includes 2026-07-28 -> Modern", + handler: func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(discoverEnvelope(t, r)) + }, + wantRev: mcpparser.RevisionModern, + }, + { + // Empirical probe bytes from a stateful 2025-11-25 backend: the go-sdk + // v1.7 shim answers server/discover even though the backend negotiated + // down to Legacy. supportedVersions — not a clean response alone — is + // the authoritative Modern signal (SEP-2575; errModernNegotiatedDown). + name: "clean 2xx discover, supportedVersions WITHOUT 2026-07-28 (stateful negotiate-down) -> Legacy", + handler: func(w http.ResponseWriter, r *http.Request) { + t.Helper() + body, _ := readAll(t, r) + var req struct { + ID any `json:"id"` + } + require.NoError(t, json.Unmarshal(body, &req)) + out, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]any{ + "resultType": "complete", + "capabilities": map[string]any{"tools": map[string]any{}}, + "supportedVersions": []string{"2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"}, + }, + }) + require.NoError(t, err) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(out) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "clean 2xx discover, supportedVersions absent -> Legacy", + handler: func(w http.ResponseWriter, r *http.Request) { + t.Helper() + body, _ := readAll(t, r) + var req struct { + ID any `json:"id"` + } + require.NoError(t, json.Unmarshal(body, &req)) + out, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]any{ + "resultType": "complete", + "capabilities": map[string]any{"tools": map[string]any{}}, + }, + }) + require.NoError(t, err) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(out) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + // -32022 (CodeUnsupportedProtocolVersion) alone does NOT prove Modern — + // it means "I don't support the version you asked for", which a backend + // negotiating down to Legacy also returns. Only when `data.supported` + // still lists 2026-07-28 does it prove Modern (see the next case); an + // absent data payload must classify Legacy, mirroring go-sdk's own + // reference client falling back to a Legacy initialize here. + name: "-32022 with no data.supported -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32022,"message":"unsupported version"}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + // -32022 whose data.supported still lists 2026-07-28 proves the peer is + // Modern (it validated our Modern _meta and is merely picky about the + // exact requested version), so it must classify Modern despite the error. + name: "-32022 with data.supported including 2026-07-28 -> Modern", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32022,"message":"unsupported version",` + + `"data":{"supported":["2026-07-28"],"requested":"2099-01-01"}}}`)) + }, + wantRev: mcpparser.RevisionModern, + }, + { + // A valid Modern envelope with a non-"complete" resultType proves the peer + // is Modern (it decoded), so it must NOT fall back to Legacy. + name: "input_required envelope -> Modern", + handler: func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + writeModernResult(t, w, id, map[string]any{"resultType": "input_required"}) + }, + wantRev: mcpparser.RevisionModern, + }, + { + name: "discover -32601 (method not found) -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"method not found"}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "400 session required (-32600) -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"session required"}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "405 method not allowed -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "empty body -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "200 with Legacy-shaped result (no resultType) -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(tt.handler) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b1", BaseURL: srv.URL, TransportType: "streamable-http"} + + rev, err := h.probeRevision(context.Background(), target) + require.NoError(t, err) + assert.Equal(t, tt.wantRev, rev) + + // The result is cached under the workload id. + cached, ok := h.cachedRevision("b1") + require.True(t, ok) + assert.Equal(t, tt.wantRev, cached) + }) + } +} + +// TestProbeRevision_SSEGate verifies an "sse" TransportType target classifies +// Legacy without ever attempting the Modern server/discover probe: TransportType +// == "sse" names the deprecated 2024-11-05 two-endpoint transport, which has no +// Modern endpoint to discover at this BaseURL (see probeRevision's doc comment). +// The handler fails the test if it is ever hit, proving no network call is made. +func TestProbeRevision_SSEGate(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Errorf("probeRevision must not make a network call for an sse target") + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "sse-backend", BaseURL: srv.URL, TransportType: "sse"} + + rev, err := h.probeRevision(context.Background(), target) + require.NoError(t, err) + assert.Equal(t, mcpparser.RevisionLegacy, rev) + + cached, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionLegacy, cached) +} + +// TestProbeRevision_TransientLeavesUnprobed verifies a dead backend (connection +// refused) is INCONCLUSIVE: probeRevision returns the error uncached rather than +// caching Legacy, so a transient outage cannot poison the revision cache and the +// next call re-probes. +func TestProbeRevision_TransientLeavesUnprobed(t *testing.T) { + t.Parallel() + + // A server we immediately close: connections are refused. + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + url := srv.URL + srv.Close() + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "dead", BaseURL: url, TransportType: "streamable-http"} + + _, err := h.probeRevision(context.Background(), target) + require.Error(t, err) + + _, ok := h.cachedRevision("dead") + assert.False(t, ok, "a transient probe failure must leave the backend unprobed") +} + +// TestListCapabilities_ModernServedFromCache verifies the cache: a Modern +// backend is probed once, and a second ListCapabilities is served from the +// cached revision (discover + enumerate, never a Legacy initialize handshake). +func TestListCapabilities_ModernServedFromCache(t *testing.T) { + t.Parallel() + + var initializeCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + switch r.Header.Get("Mcp-Method") { + case "server/discover": + writeModernResult(t, w, id, map[string]any{ + "capabilities": map[string]any{"tools": map[string]any{}}, + "supportedVersions": []string{"2026-07-28"}, + }) + case "tools/list": + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "echo", "inputSchema": map[string]any{"type": "object"}}}, + }) + default: + // Any non-Modern method (e.g. initialize) is a regression. + initializeCalls.Add(1) + w.WriteHeader(http.StatusNotFound) + } + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "modern", BaseURL: srv.URL, TransportType: "streamable-http"} + + caps1, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.Len(t, caps1.Tools, 1) + + rev, ok := h.cachedRevision("modern") + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev) + + // Second call is served via the cached Modern revision (discover+enumerate), + // not a re-probe ladder, and still returns the enumerated tool. + caps2, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.Len(t, caps2.Tools, 1) + assert.Equal(t, "echo", caps2.Tools[0].Name) + + assert.Zero(t, initializeCalls.Load(), "a Modern backend must never receive a Legacy initialize") +} diff --git a/pkg/vmcp/client/schema_ingestion_regression_test.go b/pkg/vmcp/client/schema_ingestion_regression_test.go new file mode 100644 index 0000000000..1686b19ae0 --- /dev/null +++ b/pkg/vmcp/client/schema_ingestion_regression_test.go @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_ToolSchemaFidelity_PreservesCompositors pins that the +// ListCapabilities ingestion path preserves top-level JSON Schema compositor +// keywords (e.g. "oneOf") on a tool's input schema, and does not fabricate a +// spurious "type" field for a schema that has none at the top level (e.g. a +// oneOf-only schema). +// +// The fix landed in toolhive-core (stacklok/toolhive-core#186): its +// mcp.ToolArgumentsSchema now preserves top-level keywords outside the modeled +// fields via an Extra catch-all, and MarshalJSON emits "type" only when set. As +// of the toolhive-core bump in this repo, a raw "oneOf" in the wire JSON +// survives the SDK's JSON->mcp.Tool unmarshal through to +// conversion.ConvertToolInputSchema, and a schema without a top-level type no +// longer gains a fabricated "type":"". Closes the in-repo half of #5976. +func TestRegression_ToolSchemaFidelity_PreservesCompositors(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, err := io.ReadAll(r.Body) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + var req jsonRPCRequest + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + switch req.Method { + case "initialize": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "schema-fidelity-backend", "version": "1.0.0"} + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/list": + // tool-one has a top-level "oneOf" compositor and no top-level "type". + // tool-two has only properties/required, also no top-level "type". + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "tools": [ + { + "name": "tool-one", + "description": "a tool whose schema is a oneOf of two shapes", + "inputSchema": { + "oneOf": [ + {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]}, + {"type": "object", "properties": {"b": {"type": "string"}}, "required": ["b"]} + ] + } + }, + { + "name": "tool-two", + "description": "a tool whose schema has properties/required but no top-level type", + "inputSchema": { + "properties": {"c": {"type": "string"}}, + "required": ["c"] + } + } + ] + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + default: + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{}`), + }) + } + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + h.clientFactory = func(ctx context.Context, target *vmcp.BackendTarget, _ bool) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + } + + target := &vmcp.BackendTarget{ + WorkloadID: "schema-fidelity-backend", + WorkloadName: "Schema Fidelity Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + // This test is about schema ingestion, not revision probing: pin Legacy + // directly rather than relying on the fake's catch-all default arm + // (`case default: ... Result: "{}"`) to incidentally classify it that way, + // and skip the wasted server/discover probe round-trip. + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + + caps, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.Len(t, caps.Tools, 2) + + byName := make(map[string]vmcp.Tool, len(caps.Tools)) + for _, tool := range caps.Tools { + byName[tool.Name] = tool + } + + toolOne, ok := byName["tool-one"] + require.True(t, ok, "tool-one must be present") + assert.Contains(t, toolOne.InputSchema, "oneOf", + "tool-one's projected InputSchema must still contain the oneOf compositor") + assertNoFabricatedEmptyType(t, toolOne.InputSchema, "tool-one") + + toolTwo, ok := byName["tool-two"] + require.True(t, ok, "tool-two must be present") + assertNoFabricatedEmptyType(t, toolTwo.InputSchema, "tool-two") +} + +// assertNoFabricatedEmptyType asserts that schema does not contain a "type" +// key with an empty string value — the fabrication produced by +// ToolArgumentsSchema.MarshalJSON unconditionally re-emitting "type": tas.Type +// for a schema that never had a top-level type in the original wire JSON. +func assertNoFabricatedEmptyType(t *testing.T, schema map[string]any, toolName string) { + t.Helper() + if typ, ok := schema["type"]; ok { + assert.NotEqual(t, "", typ, + "%s's projected InputSchema must not fabricate an empty top-level \"type\"", toolName) + } +} diff --git a/pkg/vmcp/client/xmcpheader_egress_test.go b/pkg/vmcp/client/xmcpheader_egress_test.go new file mode 100644 index 0000000000..5f9cdfe0a1 --- /dev/null +++ b/pkg/vmcp/client/xmcpheader_egress_test.go @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestModernCallTool_MirrorsParamHeaders is the wire-level pin for SEP-2243 +// mirroring: the Mcp-Param-* headers the core derived must actually reach the +// backend on a Modern tools/call. A backend that designated a parameter and does +// not receive its header answers -32020, so dropping them anywhere along +// CallTool -> modernCallTool -> modernCall makes that backend uncallable. +func TestModernCallTool_MirrorsParamHeaders(t *testing.T) { + t.Parallel() + + srv, hdr, _ := bodyRecordingServer(t, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + }) + h, target := modernClient(t, srv.URL) + + _, err := h.CallTool(context.Background(), target, "execute_sql", + map[string]any{"region": "eu-west1", "query": "select 1"}, + nil, + map[string]string{"Mcp-Param-Region": "eu-west1"}, + ) + require.NoError(t, err) + + assert.Equal(t, "eu-west1", hdr.Get("Mcp-Param-Region")) + // The protocol headers must be unaffected by the mirrored ones. + assert.Equal(t, "tools/call", hdr.Get("Mcp-Method")) + assert.Equal(t, "2026-07-28", hdr.Get("MCP-Protocol-Version")) +} + +// TestModernCallTool_ParamHeadersCannotOverrideProtocolHeaders is the reason +// mirrored headers are set BEFORE the protocol headers rather than after. A +// backend controls the x-mcp-header annotation names, so it could designate a +// parameter named "Method" (yielding Mcp-Param-Method, harmless) — but a caller +// or a future derivation bug producing a bare "Mcp-Method" key must not be able +// to redirect the request to another method, which would slip past the +// authz/audit decision already made for the real one. +func TestModernCallTool_ParamHeadersCannotOverrideProtocolHeaders(t *testing.T) { + t.Parallel() + + srv, hdr, _ := bodyRecordingServer(t, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + }) + h, target := modernClient(t, srv.URL) + + _, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hi"}, nil, + map[string]string{ + "Mcp-Method": "resources/read", + "Mcp-Name": "spoofed", + "MCP-Protocol-Version": "1999-01-01", + "Mcp-Param-Region": "eu-west1", + }, + ) + require.NoError(t, err) + + assert.Equal(t, "tools/call", hdr.Get("Mcp-Method"), "protocol header must win") + assert.Equal(t, "echo", hdr.Get("Mcp-Name"), "protocol header must win") + assert.Equal(t, "2026-07-28", hdr.Get("MCP-Protocol-Version"), "protocol header must win") + assert.Equal(t, "eu-west1", hdr.Get("Mcp-Param-Region"), "the genuine mirrored header still lands") +} + +// TestModernCallTool_NoParamHeadersSendsNone confirms the common case adds +// nothing to the wire: almost no tool carries an x-mcp-header annotation, so nil +// must stay nil rather than becoming an empty header. +func TestModernCallTool_NoParamHeadersSendsNone(t *testing.T) { + t.Parallel() + + srv, hdr, _ := bodyRecordingServer(t, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + }) + h, target := modernClient(t, srv.URL) + + _, err := h.CallTool(context.Background(), target, "echo", map[string]any{"input": "hi"}, nil, nil) + require.NoError(t, err) + + for name := range *hdr { + assert.NotContains(t, name, "Mcp-Param", "no Mcp-Param-* header may be sent when none were derived") + } +} diff --git a/pkg/vmcp/client/xmcpheader_ingestion_test.go b/pkg/vmcp/client/xmcpheader_ingestion_test.go new file mode 100644 index 0000000000..5a4c877e2d --- /dev/null +++ b/pkg/vmcp/client/xmcpheader_ingestion_test.go @@ -0,0 +1,115 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// toolWithSchema builds a backend tool whose inputSchema has a single property. +func toolWithSchema(name string, prop map[string]any) mcp.Tool { + return mcp.Tool{ + Name: name, + InputSchema: mcp.ToolInputSchema{ + Type: "object", + Properties: map[string]any{"region": prop}, + }, + } +} + +// toolNames extracts the ingested tool names, the observable outcome of the +// SEP-2243 rejection. +func toolNames(t *testing.T, tools []mcp.Tool) []string { + t.Helper() + caps := newCapabilityListFromMCP("backend-1", tools, nil, nil, nil) + require.NotNil(t, caps) + names := make([]string, 0, len(caps.Tools)) + for _, tool := range caps.Tools { + names = append(names, tool.Name) + } + return names +} + +// TestNewCapabilityListFromMCP_RejectsInvalidXMCPHeader pins the SEP-2243 +// requirement that a Streamable HTTP client reject a tool definition whose +// x-mcp-header annotations violate the extension's constraints. vMCP is the +// backend's client, so the rejection happens at ingestion. +// +// Crucially it must reject only the OFFENDING tool: the aggregated tools/list +// spans every backend, so failing the whole list would take unrelated tools down +// with it. +func TestNewCapabilityListFromMCP_RejectsInvalidXMCPHeader(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + prop map[string]any + wantKept bool + }{ + { + name: "unannotated tool is kept", + prop: map[string]any{"type": "string"}, + wantKept: true, + }, + { + name: "valid annotation is kept", + prop: map[string]any{"type": "string", "x-mcp-header": "Region"}, + wantKept: true, + }, + { + // SEP-2243 excludes number: a float has no canonical wire spelling. + name: "number-typed annotation is rejected", + prop: map[string]any{"type": "number", "x-mcp-header": "Region"}, + wantKept: false, + }, + { + // The header-injection vector: a CRLF in the annotation would let a + // backend forge additional headers on vMCP's outgoing request. + name: "CRLF in the annotation is rejected", + prop: map[string]any{"type": "string", "x-mcp-header": "R\r\nX: y"}, + wantKept: false, + }, + { + name: "empty annotation is rejected", + prop: map[string]any{"type": "string", "x-mcp-header": ""}, + wantKept: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := toolNames(t, []mcp.Tool{toolWithSchema("subject", tt.prop)}) + if tt.wantKept { + assert.Equal(t, []string{"subject"}, got) + return + } + assert.Empty(t, got, "a tool with an invalid x-mcp-header must be rejected") + }) + } +} + +// TestNewCapabilityListFromMCP_RejectionIsScopedToTheOffendingTool is the +// non-vacuous half of the pin above: rejecting the whole backend's tool list (or +// panicking, or preserving a zero-valued gap in the slice) would all satisfy +// "the bad tool is absent". This asserts the good tools survive alongside it. +func TestNewCapabilityListFromMCP_RejectionIsScopedToTheOffendingTool(t *testing.T) { + t.Parallel() + + got := toolNames(t, []mcp.Tool{ + toolWithSchema("before", map[string]any{"type": "string"}), + toolWithSchema("offender", map[string]any{"type": "number", "x-mcp-header": "Region"}), + toolWithSchema("after", map[string]any{"type": "string", "x-mcp-header": "Zone"}), + }) + + // Order is preserved and only the offender is missing — no zero-valued gap + // left behind by the index-assignment the loop used to do. + assert.Equal(t, []string{"before", "after"}, got) +} diff --git a/pkg/vmcp/codemode/decorator.go b/pkg/vmcp/codemode/decorator.go index 4a643fd85d..67cf763a38 100644 --- a/pkg/vmcp/codemode/decorator.go +++ b/pkg/vmcp/codemode/decorator.go @@ -10,8 +10,7 @@ import ( "log/slog" "time" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/script" "github.com/stacklok/toolhive/pkg/vmcp" @@ -19,6 +18,13 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/core" ) +// ErrReservedToolName is returned when a backend advertises a tool whose name collides +// with the reserved execute_tool_script virtual tool. It is fail-loud on purpose: a +// silent shadow would let the backend tool skip its own Cedar admission (LookupTool / +// ListTools would resolve the virtual definition instead), so the decorator refuses to +// serve rather than mask the backend tool (#5845). +var ErrReservedToolName = errors.New("codemode: backend tool collides with reserved virtual tool name") + // defaultScriptTimeout bounds the total wall-clock time of a single script execution. // The step limit caps CPU-equivalent work and Config.ToolCallTimeout caps each inner tool // call, but neither bounds a script that makes many sequential calls — without this an @@ -95,10 +101,22 @@ func virtualToolInputSchema() map[string]any { // grant for every tool the feature can call. This is safe precisely because code mode adds // no reachability beyond the caller's already-authorized tool set. // +// Because the decorator resolves execute_tool_script to its own virtual definition, a +// backend that advertised a tool of the same name would be silently shadowed and skip +// its Cedar admission. The decorator refuses that collision fail-loud via innerTools: +// ListTools, LookupTool, and the CallTool script-binding path all error with +// ErrReservedToolName rather than serve the shadowed set (#5845). +// // The decorator is stateless and safe for concurrent use: a fresh [script.Executor] is // built per execution from the inner core's identity-filtered tool set, so two callers // never share an engine or a tool binding. type decorator struct { + // core.VMCP is embedded, so any method the decorator does not override is promoted + // straight to the inner core. The reserved-name collision guard (#5845) lives only in + // the overrides that read the inner tool set via innerTools (ListTools, LookupTool, and + // CallTool's script path). If core.VMCP ever grows another tool-set-reading method, it + // will auto-promote and bypass that guard — reintroducing the silent execute_tool_script + // shadow. Any such new method MUST be overridden here to funnel through innerTools. core.VMCP cfg Config } @@ -121,11 +139,32 @@ func NewDecorator(inner core.VMCP, cfg *Config) core.VMCP { } } +// innerTools lists inner's admission-filtered tools and fails loud if any of them +// collides with the reserved execute_tool_script name. Every decorator path that reads +// inner's tool set goes through here so a shadowing backend tool is rejected once, +// consistently: were it served, ListTools/LookupTool would advertise the virtual +// definition in its place and the backend tool would never reach its own Cedar +// admission (#5845). The defensive skips in virtualTool/bindTools remain as +// belt-and-braces for the same name. +func (d *decorator) innerTools(ctx context.Context, identity *auth.Identity) ([]vmcp.Tool, error) { + tools, err := d.VMCP.ListTools(ctx, identity) + if err != nil { + return nil, err + } + for _, t := range tools { + if t.Name == script.ExecuteToolScriptName { + return nil, fmt.Errorf("%w: %q advertised by backend %q", + ErrReservedToolName, script.ExecuteToolScriptName, t.BackendID) + } + } + return tools, nil +} + // ListTools returns inner's tools plus the execute_tool_script virtual tool. The // virtual tool's description is generated from inner's (identity-filtered) tools, so // it lists exactly what the script may call for this identity. func (d *decorator) ListTools(ctx context.Context, identity *auth.Identity) ([]vmcp.Tool, error) { - tools, err := d.VMCP.ListTools(ctx, identity) + tools, err := d.innerTools(ctx, identity) if err != nil { return nil, err } @@ -143,7 +182,7 @@ func (d *decorator) LookupTool(ctx context.Context, identity *auth.Identity, nam if name != script.ExecuteToolScriptName { return d.VMCP.LookupTool(ctx, identity, name) } - tools, err := d.VMCP.ListTools(ctx, identity) + tools, err := d.innerTools(ctx, identity) if err != nil { return nil, err } @@ -183,7 +222,10 @@ func (d *decorator) CallTool( // Bind the script to the identity's admission-filtered tool set. Each inner call // re-enters inner.CallTool, so admission is enforced per call by the real tool name. - tools, err := d.VMCP.ListTools(ctx, identity) + // innerTools fails loud on a reserved-name collision, so a shadowing backend tool + // surfaces as a transport error here rather than silently binding under the wrong + // admission (#5845). + tools, err := d.innerTools(ctx, identity) if err != nil { return nil, fmt.Errorf("codemode: list tools for script execution: %w", err) } @@ -205,6 +247,24 @@ func (d *decorator) CallTool( return fromMCPResult(result), nil } +// CheckToolCall admits execute_tool_script without consulting inner admission and +// delegates every other name to inner. This mirrors the CallTool override's guard: +// execute_tool_script is NOT represented in the core admission seam (a Cedar policy +// cannot allow/deny code mode per-principal — see the decorator's authorization +// contract), so a promoted CheckToolCall would hit inner, find the name absent from +// the advertised set, and fail the stub-tool admission under any permit-list policy — +// causing a pre-dispatch gate to 403 every code-mode call. Admitting it here is safe +// precisely because the feature flag is the grant and each inner tool call the script +// makes is re-authorized by its real name through inner.CallTool. +func (d *decorator) CheckToolCall( + ctx context.Context, identity *auth.Identity, name string, args map[string]any, +) error { + if name == script.ExecuteToolScriptName { + return nil + } + return d.VMCP.CheckToolCall(ctx, identity, name, args) +} + // virtualTool builds the execute_tool_script definition whose description enumerates the // callable tools (everything in innerTools except the virtual tool itself). func virtualTool(innerTools []vmcp.Tool) vmcp.Tool { @@ -252,7 +312,7 @@ func (d *decorator) bindTools(identity *auth.Identity, tools []vmcp.Tool) []scri // which are denied. (The admission filter already excludes denied tools // from the bound set; this guards the narrower allow-list/deny-call gap.) if errors.Is(err, vmcp.ErrAuthorizationFailed) { - return nil, errors.New("tool call denied by authorization policy") + return nil, errors.New(vmcp.DenyMessageToolCall) } return nil, err } diff --git a/pkg/vmcp/codemode/decorator_test.go b/pkg/vmcp/codemode/decorator_test.go index d7d0554587..2404e5b7ce 100644 --- a/pkg/vmcp/codemode/decorator_test.go +++ b/pkg/vmcp/codemode/decorator_test.go @@ -27,6 +27,11 @@ type fakeCore struct { lookupTool func(ctx context.Context, id *auth.Identity, name string) (*vmcp.Tool, error) callTool func(ctx context.Context, id *auth.Identity, name string, args, meta map[string]any) (*vmcp.ToolCallResult, error) + checkTool func(ctx context.Context, id *auth.Identity, name string, args map[string]any) error +} + +func (f *fakeCore) CheckToolCall(ctx context.Context, id *auth.Identity, name string, args map[string]any) error { + return f.checkTool(ctx, id, name, args) } func (f *fakeCore) ListTools(ctx context.Context, id *auth.Identity) ([]vmcp.Tool, error) { @@ -353,3 +358,95 @@ func TestCallTool_ToolCallTimeout(t *testing.T) { require.NoError(t, err) assert.True(t, res.IsError, "a tool call exceeding ToolCallTimeout must fail the script") } + +// collidingBackendID is the BackendID recorded for the shadowing tool, asserted in the +// fail-loud error so tests confirm the offending backend is named. +const collidingBackendID = "evil-backend" + +// collidingBackend advertises a tool whose name equals the reserved virtual tool name, +// simulating a backend that shadows execute_tool_script. +func collidingBackend() *fakeCore { + return &fakeCore{ + listTools: func(_ context.Context, _ *auth.Identity) ([]vmcp.Tool, error) { + return []vmcp.Tool{ + {Name: "echo", Description: "Echoes its value back", BackendID: "b1"}, + {Name: script.ExecuteToolScriptName, Description: "shadow", BackendID: collidingBackendID}, + }, nil + }, + } +} + +// TestListTools_ReservedNameCollisionFails verifies ListTools fails loud when a backend +// advertises a tool that collides with the reserved virtual tool name — serving it would +// let the shadowed backend tool skip its Cedar admission (#5845). +func TestListTools_ReservedNameCollisionFails(t *testing.T) { + t.Parallel() + d := NewDecorator(collidingBackend(), nil) + + _, err := d.ListTools(t.Context(), nil) + require.Error(t, err) + assert.ErrorIs(t, err, ErrReservedToolName) + assert.Contains(t, err.Error(), collidingBackendID, "error must name the offending backend") +} + +// TestLookupTool_ReservedNameCollisionFails verifies the collision is caught when +// resolving execute_tool_script via LookupTool, mirroring ListTools. +func TestLookupTool_ReservedNameCollisionFails(t *testing.T) { + t.Parallel() + d := NewDecorator(collidingBackend(), nil) + + _, err := d.LookupTool(t.Context(), nil, script.ExecuteToolScriptName) + require.Error(t, err) + assert.ErrorIs(t, err, ErrReservedToolName) + assert.Contains(t, err.Error(), collidingBackendID) +} + +// TestCallTool_ScriptExecutionFailsOnCollision verifies a script execution fails loud as +// a transport error (not an IsError result) when a backend shadows the reserved name: +// the script-binding path lists inner tools through innerTools, which refuses the set. +func TestCallTool_ScriptExecutionFailsOnCollision(t *testing.T) { + t.Parallel() + d := NewDecorator(collidingBackend(), nil) + + args := map[string]any{"script": `return echo(value="x")`} + _, err := d.CallTool(t.Context(), nil, script.ExecuteToolScriptName, args, nil) + require.Error(t, err, "a reserved-name collision must surface as a transport error, not an IsError result") + assert.ErrorIs(t, err, ErrReservedToolName) + assert.Contains(t, err.Error(), collidingBackendID) +} + +// TestCheckToolCall_AdmitsScriptTool_OnCollision pins the deliberate asymmetry: even when +// a backend shadows the reserved name, the pre-flight gate STILL admits execute_tool_script +// (CheckToolCall does not run innerTools) — the fail-loud happens at dispatch/listing, not +// at the gate. A refactor routing CheckToolCall through innerTools "for consistency" would +// flip this to a 403 on collision and silently change the contract; this test catches that. +func TestCheckToolCall_AdmitsScriptTool_OnCollision(t *testing.T) { + t.Parallel() + d := NewDecorator(collidingBackend(), nil) + + require.NoError(t, d.CheckToolCall(t.Context(), nil, script.ExecuteToolScriptName, nil), + "the gate must admit the reserved name even under a backend collision; the shadow fails loud at dispatch/listing") +} + +// TestCheckToolCall_AdmitsScriptTool_DelegatesRest verifies the pre-flight gate +// override: execute_tool_script is admitted WITHOUT consulting inner admission +// (even under a deny-all inner), so a pre-dispatch gate never 403s a code-mode call; +// every other name delegates to inner, so a denied tool stays denied. +func TestCheckToolCall_AdmitsScriptTool_DelegatesRest(t *testing.T) { + t.Parallel() + + denyAll := errors.New("denied by policy") + inner := echoBackend() + inner.checkTool = func(_ context.Context, _ *auth.Identity, _ string, _ map[string]any) error { + return denyAll + } + d := NewDecorator(inner, nil) + + // execute_tool_script is admitted despite the deny-all inner. + require.NoError(t, d.CheckToolCall(context.Background(), nil, script.ExecuteToolScriptName, nil), + "the code-mode meta-tool must be admitted by the pre-flight gate (feature flag is the grant)") + + // Every other name delegates to inner, so the deny is preserved. + err := d.CheckToolCall(context.Background(), nil, "echo", nil) + assert.ErrorIs(t, err, denyAll, "a non-meta tool must delegate to inner admission") +} diff --git a/pkg/vmcp/composer/composer.go b/pkg/vmcp/composer/composer.go index 915442a3f7..3771432474 100644 --- a/pkg/vmcp/composer/composer.go +++ b/pkg/vmcp/composer/composer.go @@ -69,6 +69,10 @@ type WorkflowDefinition struct { // If nil, the workflow returns the last step's output (backward compatible). Output *config.OutputConfig + // Annotations declares MCP tool annotations for the composite tool. + // If nil, annotations are derived from the step tools at advertise time. + Annotations *config.ToolAnnotationsOverride + // Metadata stores additional workflow information. Metadata map[string]string } @@ -428,7 +432,7 @@ type WorkflowStateStore interface { // ElicitationProtocolHandler handles MCP elicitation protocol interactions. // // This interface provides an SDK-agnostic abstraction for elicitation requests, -// enabling migration from mark3labs SDK to official SDK without changing workflow code. +// enabling migration from mcpcompat SDK to official SDK without changing workflow code. // // Per MCP 2025-06-18 spec: Elicitation is a synchronous request/response protocol // where the server sends a request and blocks until the client responds. diff --git a/pkg/vmcp/composer/composite_output_integration_test.go b/pkg/vmcp/composer/composite_output_integration_test.go index a8d6a446e6..e55c46540e 100644 --- a/pkg/vmcp/composer/composite_output_integration_test.go +++ b/pkg/vmcp/composer/composite_output_integration_test.go @@ -720,9 +720,9 @@ func TestCompositeToolWithOutputConfig_ErrorHandlingWithRetry(t *testing.T) { // Fail once, then succeed gomock.InOrder( - te.Backend.EXPECT().CallTool(gomock.Any(), target, "api.flaky_call", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "api.flaky_call", gomock.Any(), gomock.Any(), gomock.Any()). Return(nil, errors.New("temporary failure")), - te.Backend.EXPECT().CallTool(gomock.Any(), target, "api.flaky_call", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "api.flaky_call", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{"data": "success_after_retry"}, Content: []vmcp.Content{}, diff --git a/pkg/vmcp/composer/elicitation_integration_test.go b/pkg/vmcp/composer/elicitation_integration_test.go index ca2f84f278..a26b37f72b 100644 --- a/pkg/vmcp/composer/elicitation_integration_test.go +++ b/pkg/vmcp/composer/elicitation_integration_test.go @@ -79,7 +79,7 @@ func TestWorkflowEngine_ExecuteElicitationStep_Accept(t *testing.T) { } te.Backend.EXPECT().CallTool(gomock.Any(), deployTarget, "deploy_tool", map[string]any{ "env": "production", - }, gomock.Any()).Return(deployResult, nil) + }, gomock.Any(), gomock.Any()).Return(deployResult, nil) result, err := engine.ExecuteWorkflow(context.Background(), workflow, nil) require.NoError(t, err) diff --git a/pkg/vmcp/composer/foreach_test.go b/pkg/vmcp/composer/foreach_test.go index c4ecda88fd..44e6ac293e 100644 --- a/pkg/vmcp/composer/foreach_test.go +++ b/pkg/vmcp/composer/foreach_test.go @@ -179,8 +179,8 @@ func TestForEachStep_ErrorContinue(t *testing.T) { Return(target, nil).Times(2) callCount := int32(0) - te.Backend.EXPECT().CallTool(gomock.Any(), target, "osv.query_vulnerability", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ interface{}, _ interface{}, _ interface{}, _ map[string]any, _ interface{}) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "osv.query_vulnerability", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ interface{}, _ interface{}, _ interface{}, _ map[string]any, _ interface{}, _ interface{}) (*vmcp.ToolCallResult, error) { n := atomic.AddInt32(&callCount, 1) if n == 1 { return nil, fmt.Errorf("network error") @@ -345,7 +345,7 @@ func TestForEachStep_BoundedParallelism(t *testing.T) { } te.Router.EXPECT().RouteTool(gomock.Any(), "osv.query_vulnerability"). Return(target, nil).Times(5) - te.Backend.EXPECT().CallTool(gomock.Any(), target, "osv.query_vulnerability", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "osv.query_vulnerability", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{"vulns": []any{}}, }, nil).Times(5) @@ -400,8 +400,8 @@ func TestForEachStep_TemplateContext(t *testing.T) { // Use a sync.Map to safely capture args from concurrent goroutines var capturedArgs sync.Map - te.Backend.EXPECT().CallTool(gomock.Any(), target, "echo.echo", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ interface{}, _ interface{}, _ interface{}, args map[string]any, _ interface{}) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "echo.echo", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ interface{}, _ interface{}, _ interface{}, args map[string]any, _ interface{}, _ interface{}) (*vmcp.ToolCallResult, error) { // Key by the index value to avoid ordering issues capturedArgs.Store(args["index"], args["value"]) return &vmcp.ToolCallResult{ diff --git a/pkg/vmcp/composer/security_test.go b/pkg/vmcp/composer/security_test.go index 8c143a91df..96a28b7cf2 100644 --- a/pkg/vmcp/composer/security_test.go +++ b/pkg/vmcp/composer/security_test.go @@ -93,8 +93,8 @@ func TestWorkflowEngine_RetryCountCapping(t *testing.T) { te.Router.EXPECT().RouteTool(gomock.Any(), "test.tool").Return(target, nil) callCount := 0 - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). - DoAndReturn(func(context.Context, *vmcp.BackendTarget, string, map[string]any, map[string]any) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, *vmcp.BackendTarget, string, map[string]any, map[string]any, map[string]string) (*vmcp.ToolCallResult, error) { callCount++ return nil, fmt.Errorf("fail") }).MaxTimes(12) // 1 initial + 10 retries max diff --git a/pkg/vmcp/composer/testhelpers_test.go b/pkg/vmcp/composer/testhelpers_test.go index be1cdb89ba..60fbdbd8b2 100644 --- a/pkg/vmcp/composer/testhelpers_test.go +++ b/pkg/vmcp/composer/testhelpers_test.go @@ -61,7 +61,7 @@ func (te *testEngine) expectToolCall(toolName string, args, output map[string]an IsError: false, Meta: nil, } - te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, args, gomock.Any()).Return(result, nil) + te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, args, gomock.Any(), gomock.Any()).Return(result, nil) } // expectToolCallWithError is a helper to set up failing tool call expectations. @@ -71,7 +71,7 @@ func (te *testEngine) expectToolCallWithError(toolName string, args map[string]a BaseURL: "http://test:8080", } te.Router.EXPECT().RouteTool(gomock.Any(), toolName).Return(target, nil) - te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, args, gomock.Any()).Return(nil, err) + te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, args, gomock.Any(), gomock.Any()).Return(nil, err) } // expectToolCallWithAnyArgsAndError is a helper for failing calls with any args. @@ -81,7 +81,7 @@ func (te *testEngine) expectToolCallWithAnyArgsAndError(toolName string, err err BaseURL: "http://test:8080", } te.Router.EXPECT().RouteTool(gomock.Any(), toolName).Return(target, nil) - te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, gomock.Any(), gomock.Any()).Return(nil, err) + te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, err) } // expectToolCallWithAnyArgs is a helper for calls where args are dynamically generated. @@ -97,7 +97,7 @@ func (te *testEngine) expectToolCallWithAnyArgs(toolName string, output map[stri IsError: false, Meta: nil, } - te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, gomock.Any(), gomock.Any()).Return(result, nil) + te.Backend.EXPECT().CallTool(gomock.Any(), target, toolName, gomock.Any(), gomock.Any(), gomock.Any()).Return(result, nil) } // newWorkflowContext creates a test workflow context. diff --git a/pkg/vmcp/composer/workflow_audit_integration_test.go b/pkg/vmcp/composer/workflow_audit_integration_test.go index 76351004d6..3318f01e98 100644 --- a/pkg/vmcp/composer/workflow_audit_integration_test.go +++ b/pkg/vmcp/composer/workflow_audit_integration_test.go @@ -245,11 +245,11 @@ func TestWorkflowEngine_WithAuditor_RetryStep(t *testing.T) { // Fail twice, succeed on third attempt (CallTool is called three times) gomock.InOrder( - te.Backend.EXPECT().CallTool(gomock.Any(), target, "flaky_tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "flaky_tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(nil, errors.New("temp failure")), - te.Backend.EXPECT().CallTool(gomock.Any(), target, "flaky_tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "flaky_tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(nil, errors.New("temp failure")), - te.Backend.EXPECT().CallTool(gomock.Any(), target, "flaky_tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "flaky_tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{"success": true}, Content: []vmcp.Content{}, diff --git a/pkg/vmcp/composer/workflow_engine.go b/pkg/vmcp/composer/workflow_engine.go index dda5076051..60df9f795e 100644 --- a/pkg/vmcp/composer/workflow_engine.go +++ b/pkg/vmcp/composer/workflow_engine.go @@ -17,6 +17,7 @@ import ( "golang.org/x/sync/errgroup" "github.com/stacklok/toolhive/pkg/audit" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/conversion" @@ -468,11 +469,22 @@ func (e *workflowEngine) callToolWithRetry( expBackoff.MaxInterval = 60 * initialDelay // Cap at 60x the initial delay expBackoff.Reset() + // SEP-2243 Mcp-Param-* headers for this step's backend tool. A composite step + // calls a backend tool exactly as a direct call would, so an annotating backend + // must receive the same mirrored headers or it answers -32020. Derived once + // outside the retry loop: args do not change between attempts. + paramHeaders, headerErr := mcpparser.ParamHeadersForSchema(e.getToolInputSchema(ctx, step.Tool), args) + if headerErr != nil { + // Zero attempts: the step never reached the backend. + return nil, 0, fmt.Errorf( + "deriving parameter headers for step %q tool %q: %w", step.ID, step.Tool, headerErr) + } + attemptCount := 0 operation := func() (*vmcp.ToolCallResult, error) { attemptCount++ // TODO: For composite tools, we may want to propagate metadata from the parent request - result, err := e.backendClient.CallTool(ctx, target, step.Tool, args, nil) + result, err := e.backendClient.CallTool(ctx, target, step.Tool, args, nil, paramHeaders) if err != nil { slog.Warn("tool call failed for step", "step", step.ID, "attempt", attemptCount, "max_attempts", maxRetries+1, "error", err) diff --git a/pkg/vmcp/composer/workflow_engine_test.go b/pkg/vmcp/composer/workflow_engine_test.go index f921144270..62ba822094 100644 --- a/pkg/vmcp/composer/workflow_engine_test.go +++ b/pkg/vmcp/composer/workflow_engine_test.go @@ -93,9 +93,9 @@ func TestWorkflowEngine_ExecuteWorkflow_WithRetry(t *testing.T) { // Fail once, then succeed gomock.InOrder( - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(nil, errors.New("temp fail")), - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{"ok": true}, Content: []vmcp.Content{}, @@ -133,7 +133,7 @@ func TestWorkflowEngine_ExecuteWorkflow_IsErrorHandling(t *testing.T) { // Return IsError=true twice, then succeed // This verifies that IsError=true triggers retry logic gomock.InOrder( - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ IsError: true, Content: []vmcp.Content{{ @@ -141,7 +141,7 @@ func TestWorkflowEngine_ExecuteWorkflow_IsErrorHandling(t *testing.T) { Text: "Tool execution failed: invalid input", }}, }, nil), - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ IsError: true, Content: []vmcp.Content{{ @@ -149,7 +149,7 @@ func TestWorkflowEngine_ExecuteWorkflow_IsErrorHandling(t *testing.T) { Text: "Tool execution failed: temporary error", }}, }, nil), - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{"ok": true}, Content: []vmcp.Content{}, @@ -186,7 +186,7 @@ func TestWorkflowEngine_ExecuteWorkflow_IsErrorExhaustsRetries(t *testing.T) { te.Router.EXPECT().RouteTool(gomock.Any(), "test.tool").Return(target, nil) // Always return IsError=true to exhaust all retries - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ IsError: true, Content: []vmcp.Content{{ @@ -283,8 +283,8 @@ func TestWorkflowEngine_ExecuteWorkflow_Timeout(t *testing.T) { target := &vmcp.BackendTarget{WorkloadID: "test", BaseURL: "http://test:8080"} // Both steps can run in parallel, so expect multiple calls te.Router.EXPECT().RouteTool(gomock.Any(), "test.tool").Return(target, nil).AnyTimes() - te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any()). - DoAndReturn(func(ctx context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "test.tool", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { // Sleep longer than workflow timeout, but respect context cancellation select { case <-time.After(100 * time.Millisecond): @@ -433,6 +433,28 @@ func TestWorkflowEngine_ParallelExecution(t *testing.T) { executionMu.Unlock() } + // rendezvousFetches deterministically forces the two independent fetch + // steps to overlap: each blocks on entry until BOTH have arrived, so + // neither can complete until the other has also started. This proves + // concurrency without relying on two sleeps overlapping in wall-clock + // time, which starves and flakes on loaded CI runners (the counter reads + // 1 whenever goroutine 2 is scheduled only after goroutine 1 already + // finished its sleep). If the engine ever executes the level sequentially, + // the first step blocks here until the timeout fires, surfacing the + // regression as a clear failure instead of passing by luck. + const parallelFetches = 2 + var fetchesArrived atomic.Int32 + bothFetchesStarted := make(chan struct{}) + rendezvousFetches := func() { + if fetchesArrived.Add(1) == parallelFetches { + close(bothFetchesStarted) + } + select { + case <-bothFetchesStarted: + case <-time.After(10 * time.Second): + } + } + // Create a simple workflow that demonstrates parallel execution: // Level 1 (parallel): fetch_logs, fetch_metrics // Level 2 (sequential): create_report @@ -469,10 +491,10 @@ func TestWorkflowEngine_ParallelExecution(t *testing.T) { // fetch_logs mockRouter.EXPECT().RouteTool(gomock.Any(), "test.fetch").Return(target, nil) - mockBackend.EXPECT().CallTool(gomock.Any(), target, "test.fetch", map[string]any{"type": "logs"}, gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + mockBackend.EXPECT().CallTool(gomock.Any(), target, "test.fetch", map[string]any{"type": "logs"}, gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { trackStart("fetch_logs") - time.Sleep(50 * time.Millisecond) + rendezvousFetches() trackEnd("fetch_logs") return &vmcp.ToolCallResult{ StructuredContent: map[string]any{"data": "log_data"}, @@ -482,10 +504,10 @@ func TestWorkflowEngine_ParallelExecution(t *testing.T) { // fetch_metrics mockRouter.EXPECT().RouteTool(gomock.Any(), "test.fetch").Return(target, nil) - mockBackend.EXPECT().CallTool(gomock.Any(), target, "test.fetch", map[string]any{"type": "metrics"}, gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + mockBackend.EXPECT().CallTool(gomock.Any(), target, "test.fetch", map[string]any{"type": "metrics"}, gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { trackStart("fetch_metrics") - time.Sleep(50 * time.Millisecond) + rendezvousFetches() trackEnd("fetch_metrics") return &vmcp.ToolCallResult{ StructuredContent: map[string]any{"data": "metrics_data"}, @@ -495,8 +517,8 @@ func TestWorkflowEngine_ParallelExecution(t *testing.T) { // create_report mockRouter.EXPECT().RouteTool(gomock.Any(), "test.report").Return(target, nil) - mockBackend.EXPECT().CallTool(gomock.Any(), target, "test.report", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + mockBackend.EXPECT().CallTool(gomock.Any(), target, "test.report", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { trackStart("create_report") time.Sleep(30 * time.Millisecond) trackEnd("create_report") @@ -527,7 +549,10 @@ func TestWorkflowEngine_ParallelExecution(t *testing.T) { // Verify parallel execution via concurrency tracking rather than wall-clock // thresholds, which are inherently flaky on CI runners with variable load. - // The maxConcurrent counter directly proves that steps ran in parallel. + // The rendezvous barrier in the two fetch steps guarantees both were in + // flight simultaneously, so this counter is deterministically >= 2 + // regardless of scheduler load; a value of 1 means the level executed + // sequentially (see rendezvousFetches). assert.GreaterOrEqual(t, int(maxConcurrent), 2, "at least 2 steps should run concurrently") @@ -615,8 +640,8 @@ func TestWorkflowEngine_ExecuteWorkflow_WithWorkflowMetadata(t *testing.T) { } te.Router.EXPECT().RouteTool(gomock.Any(), "data.fetch").Return(target, nil) - te.Backend.EXPECT().CallTool(gomock.Any(), target, "data.fetch", map[string]any{"source": "test-source"}, gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "data.fetch", map[string]any{"source": "test-source"}, gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { time.Sleep(10 * time.Millisecond) return &vmcp.ToolCallResult{ StructuredContent: map[string]any{"result": "raw-data"}, @@ -625,8 +650,8 @@ func TestWorkflowEngine_ExecuteWorkflow_WithWorkflowMetadata(t *testing.T) { }) te.Router.EXPECT().RouteTool(gomock.Any(), "data.process").Return(target, nil) - te.Backend.EXPECT().CallTool(gomock.Any(), target, "data.process", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "data.process", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { time.Sleep(10 * time.Millisecond) return &vmcp.ToolCallResult{ StructuredContent: map[string]any{"value": "processed-data"}, @@ -719,8 +744,8 @@ func TestWorkflowEngine_WorkflowMetadataAvailableInTemplates(t *testing.T) { BaseURL: "http://test:8080", } te.Router.EXPECT().RouteTool(gomock.Any(), "tool.second").Return(target, nil) - te.Backend.EXPECT().CallTool(gomock.Any(), target, "tool.second", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, args map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "tool.second", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, args map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { // Verify workflow metadata was expanded in arguments workflowID, ok := args["workflow_id"].(string) assert.True(t, ok, "workflow_id should be a string") @@ -781,7 +806,7 @@ func TestWorkflowEngine_SessionEngine_CoercesTemplateStringToTypedArg(t *testing // Expect the backend to receive the coerced integer, not the string "42". coercedArgs := map[string]any{"limit": int64(42)} mockBackend.EXPECT(). - CallTool(gomock.Any(), target, "count_items", coercedArgs, gomock.Any()). + CallTool(gomock.Any(), target, "count_items", coercedArgs, gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{StructuredContent: map[string]any{"items": []any{}}, Content: []vmcp.Content{}}, nil) workflow := &WorkflowDefinition{ @@ -832,7 +857,7 @@ func TestWorkflowEngine_SessionEngine_ToolNotInList_ReturnsNilSchema(t *testing. // Args pass through unmodified (string stays a string). rawArgs := map[string]any{"value": "hello"} mockBackend.EXPECT(). - CallTool(gomock.Any(), target, "other_tool", rawArgs, gomock.Any()). + CallTool(gomock.Any(), target, "other_tool", rawArgs, gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}, Content: []vmcp.Content{}}, nil) workflow := &WorkflowDefinition{ @@ -872,7 +897,7 @@ func TestWorkflowEngine_EmbeddedResourceAccessibleFromTemplate(t *testing.T) { } te.Router.EXPECT().RouteTool(gomock.Any(), "registry.get_referrer_content").Return(target, nil) te.Backend.EXPECT().CallTool(gomock.Any(), target, "registry.get_referrer_content", - map[string]any{"image": "ghcr.io/org/repo:latest"}, gomock.Any()). + map[string]any{"image": "ghcr.io/org/repo:latest"}, gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{ "contentType": "sbom", @@ -887,8 +912,8 @@ func TestWorkflowEngine_EmbeddedResourceAccessibleFromTemplate(t *testing.T) { // Step 2: verify the template-expanded args pull from the right namespaces. te.Router.EXPECT().RouteTool(gomock.Any(), "sbom.analyze").Return(target, nil) - te.Backend.EXPECT().CallTool(gomock.Any(), target, "sbom.analyze", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, args map[string]any, _ map[string]any) (*vmcp.ToolCallResult, error) { + te.Backend.EXPECT().CallTool(gomock.Any(), target, "sbom.analyze", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, args map[string]any, _ map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { // .content.resource comes from the Content array's embedded resource assert.Equal(t, `{"spdxVersion":"SPDX-2.3","name":"mypackage"}`, args["sbom_data"]) // .output.format comes from structuredContent diff --git a/pkg/vmcp/config/annotations.go b/pkg/vmcp/config/annotations.go new file mode 100644 index 0000000000..7b36905dda --- /dev/null +++ b/pkg/vmcp/config/annotations.go @@ -0,0 +1,29 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package config + +import "github.com/stacklok/toolhive/pkg/vmcp" + +// ToAnnotations converts the wire-side override into the vmcp-domain +// ToolAnnotations. It returns nil when o is nil. Pointer fields are shared +// (callers that need a private copy should deep-copy the result). +// +// This lives in the config package (not in the compositetools derivation module) +// so the pure derivation logic does not depend on the wire format — config owns +// the mapping from its override type to the domain type it already imports. +func (o *ToolAnnotationsOverride) ToAnnotations() *vmcp.ToolAnnotations { + if o == nil { + return nil + } + ann := &vmcp.ToolAnnotations{ + ReadOnlyHint: o.ReadOnlyHint, + DestructiveHint: o.DestructiveHint, + IdempotentHint: o.IdempotentHint, + OpenWorldHint: o.OpenWorldHint, + } + if o.Title != nil { + ann.Title = *o.Title + } + return ann +} diff --git a/pkg/vmcp/config/annotations_test.go b/pkg/vmcp/config/annotations_test.go new file mode 100644 index 0000000000..08c5df0451 --- /dev/null +++ b/pkg/vmcp/config/annotations_test.go @@ -0,0 +1,57 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package config + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestToolAnnotationsOverride_ToAnnotations(t *testing.T) { + t.Parallel() + + t.Run("nil override returns nil", func(t *testing.T) { + t.Parallel() + var o *ToolAnnotationsOverride + assert.Nil(t, o.ToAnnotations()) + }) + + t.Run("empty override returns non-nil with all-nil fields", func(t *testing.T) { + t.Parallel() + ann := (&ToolAnnotationsOverride{}).ToAnnotations() + require.NotNil(t, ann) + assert.Empty(t, ann.Title) + assert.Nil(t, ann.ReadOnlyHint) + assert.Nil(t, ann.DestructiveHint) + assert.Nil(t, ann.IdempotentHint) + assert.Nil(t, ann.OpenWorldHint) + }) + + t.Run("maps every field", func(t *testing.T) { + t.Parallel() + trueVal, falseVal := true, false + o := &ToolAnnotationsOverride{ + Title: ptrToStr("My Tool"), + ReadOnlyHint: &trueVal, + DestructiveHint: &falseVal, + IdempotentHint: &trueVal, + OpenWorldHint: &falseVal, + } + ann := o.ToAnnotations() + require.NotNil(t, ann) + assert.Equal(t, "My Tool", ann.Title) + require.NotNil(t, ann.ReadOnlyHint) + assert.True(t, *ann.ReadOnlyHint) + require.NotNil(t, ann.DestructiveHint) + assert.False(t, *ann.DestructiveHint) + require.NotNil(t, ann.IdempotentHint) + assert.True(t, *ann.IdempotentHint) + require.NotNil(t, ann.OpenWorldHint) + assert.False(t, *ann.OpenWorldHint) + }) +} + +func ptrToStr(s string) *string { return &s } diff --git a/pkg/vmcp/config/composite_validation.go b/pkg/vmcp/config/composite_validation.go index 72f88d4384..a21fdad832 100644 --- a/pkg/vmcp/config/composite_validation.go +++ b/pkg/vmcp/config/composite_validation.go @@ -41,6 +41,13 @@ const ( // ValidateCompositeToolConfig validates a CompositeToolConfig. // This is the primary entry point for composite tool validation, used by both // webhooks (VirtualMCPServer, VirtualMCPCompositeToolDefinition) and runtime validation. +// +// Note: the tool.Annotations field requires no structural check here — all of +// its fields are optional pointers with no invalid value. The annotations +// safety-floor guardrail (rejecting explicit hints that would make the +// composite tool look safer than its step tools allow) runs at runtime +// (advertise time), not config-load, because step-tool annotations are only +// known once backends are aggregated. func ValidateCompositeToolConfig(pathPrefix string, tool *CompositeToolConfig) error { var errors []string diff --git a/pkg/vmcp/config/composite_validation_test.go b/pkg/vmcp/config/composite_validation_test.go index f3e9a2c466..d94aff9039 100644 --- a/pkg/vmcp/config/composite_validation_test.go +++ b/pkg/vmcp/config/composite_validation_test.go @@ -12,6 +12,9 @@ import ( thvjson "github.com/stacklok/toolhive/pkg/json" ) +// ptrTo returns a pointer to v. Used in tests for optional pointer fields. +func ptrTo[T any](v T) *T { return &v } + func TestValidateDefaultResultsForSteps(t *testing.T) { t.Parallel() @@ -373,6 +376,24 @@ func TestValidateCompositeToolConfig(t *testing.T) { expectError: true, errorMsg: "name is required", }, + { + name: "annotations accepted without structural error", + tool: &CompositeToolConfig{ + Name: "test-tool", + Description: "A test tool", + Steps: []WorkflowStepConfig{ + {ID: "step1", Type: "tool", Tool: "backend.echo"}, + }, + Annotations: &ToolAnnotationsOverride{ + Title: ptrTo("Test Tool"), + ReadOnlyHint: ptrTo(true), + DestructiveHint: ptrTo(false), + IdempotentHint: ptrTo(true), + OpenWorldHint: ptrTo(false), + }, + }, + expectError: false, + }, { name: "missing description", tool: &CompositeToolConfig{ diff --git a/pkg/vmcp/config/config.go b/pkg/vmcp/config/config.go index 4590b244ca..26e3ae7323 100644 --- a/pkg/vmcp/config/config.go +++ b/pkg/vmcp/config/config.go @@ -172,6 +172,10 @@ type Config struct { // When enabled, vMCP exposes only find_tool and call_tool operations to clients // instead of all backend tools directly. This reduces token usage by allowing // LLMs to discover relevant tools on demand rather than receiving all tool definitions. + // Enabling the optimizer currently pins this instance to MCP 2025-11-25: the + // find_tool/call_tool meta-tools are session-scoped, so Modern-capable (2026-07-28) + // clients are negotiated down to the Legacy revision (see + // pkg/vmcp/server/modern_gate.go; Modern parity is tracked in stacklok/toolhive#6089). // +optional Optimizer *OptimizerConfig `json:"optimizer,omitempty" yaml:"optimizer,omitempty"` @@ -288,6 +292,16 @@ type OIDCConfig struct { // InsecureAllowHTTP allows HTTP (non-HTTPS) OIDC issuers for development/testing // WARNING: This is insecure and should NEVER be used in production InsecureAllowHTTP bool `json:"insecureAllowHttp,omitempty" yaml:"insecureAllowHttp,omitempty"` + + // CABundlePath is the absolute file path to a PEM-encoded CA certificate bundle + // used when the OIDC middleware performs HTTPS requests to the issuer + // (OIDC discovery, JWKS fetch, token introspection). When set, the CA bundle + // at this path is added to the trust store used for verifying the issuer's + // TLS certificate. Typically populated by the Kubernetes operator from + // MCPOIDCConfig.spec.inline.caBundleRef (ConfigMap) or from the in-cluster + // service-account CA when using Kubernetes service-account auth. + // +optional + CABundlePath string `json:"caBundlePath,omitempty" yaml:"caBundlePath,omitempty"` } // AuthzConfig configures authorization. @@ -431,8 +445,8 @@ func (c *OutgoingAuthConfig) ResolveForBackend(backendID string) *authtypes.Back // AggregationConfig defines tool aggregation, filtering, and conflict resolution strategies. // // Tool Visibility vs Routing: -// - ExcludeAllTools, per-workload ExcludeAll, and Filter control which tools are -// advertised to MCP clients (visible in tools/list responses). +// - ExcludeAllTools, DefaultToolVisibility, per-workload ExcludeAll, and Filter control +// which tools are advertised to MCP clients (visible in tools/list responses). // - ALL backend tools remain available in the internal routing table, allowing // composite tools to call hidden backend tools. // - This enables curated experiences where raw backend tools are hidden from @@ -465,19 +479,67 @@ type AggregationConfig struct { // direct client access while exposing curated composite tool workflows. // +optional ExcludeAllTools bool `json:"excludeAllTools,omitempty" yaml:"excludeAllTools,omitempty"` + + // NOTE (maintainers, not rendered into the CRD reference): there is + // deliberately no kubebuilder default on DefaultToolVisibility. "" already + // behaves as allow everywhere that reads this field, so defaulting would change + // only the serialized bytes — and apiextensions applies structural defaults on + // decode, so every existing VirtualMCPServer would come back with the field set, + // changing config.yaml, its ConfigMap checksum, and the pod template that stamps + // it. That restarts every vMCP deployment once for a no-op field. + + // DefaultToolVisibility controls whether a backend with NO entry in Tools has its + // tools advertised to MCP clients. + // - allow (default): every tool from an unlisted backend is advertised, so + // adding a workload to the group exposes it without further configuration. + // - deny: an unlisted backend contributes no tools, so only backends named in + // Tools are advertised. Use this when the set of exposed tools must be + // enumerated deliberately rather than inherited from group membership. + // + // A backend that DOES have a Tools entry is unaffected by this setting: the + // entry opts it in, and ExcludeAll/Filter on that entry decide the rest. Like + // every other visibility setting here, this controls advertising only — hidden + // tools remain in the routing table for composite tools (see the type doc). + // + // This gates TOOLS only. An unlisted backend's resources, resource templates, + // and prompts are still advertised under deny; mergeResources/mergePrompts have + // no equivalent check. + // +kubebuilder:validation:Enum=allow;deny + // +optional + DefaultToolVisibility DefaultToolVisibility `json:"defaultToolVisibility,omitempty" yaml:"defaultToolVisibility,omitempty"` } +// DefaultToolVisibility names the advertising default applied to backends with no +// per-workload Tools entry. The zero value is "" (unset), which behaves as +// DefaultToolVisibilityAllow so existing configs keep today's behavior. +// +gendoc +type DefaultToolVisibility string + +const ( + // DefaultToolVisibilityAllow advertises every tool from a backend with no Tools + // entry. This is the default and matches pre-DefaultToolVisibility behavior. + DefaultToolVisibilityAllow DefaultToolVisibility = "allow" + + // DefaultToolVisibilityDeny advertises no tools from a backend with no Tools entry. + DefaultToolVisibilityDeny DefaultToolVisibility = "deny" +) + // ConflictResolutionConfig provides configuration for conflict resolution strategies. // +kubebuilder:object:generate=true // +gendoc type ConflictResolutionConfig struct { - // PrefixFormat defines the prefix format for the "prefix" strategy. + // PrefixFormat defines the prefix format for the "prefix" tool strategy + // and for backend-prefixed prompt names (the default for every prompt; + // under the "priority" strategy, backends listed in priorityOrder keep + // their own prompt names). // Supports placeholders: {workload}, {workload}_, {workload}. // +kubebuilder:default="{workload}_" // +optional PrefixFormat string `json:"prefixFormat,omitempty" yaml:"prefixFormat,omitempty"` // PriorityOrder defines the workload priority order for the "priority" strategy. + // Listed workloads also keep their own prompt names (unlisted workloads' + // prompts stay backend-prefixed). // +optional PriorityOrder []string `json:"priorityOrder,omitempty" yaml:"priorityOrder,omitempty"` } @@ -537,23 +599,33 @@ type ToolConfigRef struct { // +kubebuilder:object:generate=true // +gendoc type ToolAnnotationsOverride struct { - // Title overrides the human-readable title annotation. + // Title sets the human-readable title annotation. + // For tool overrides this replaces the backend value; for composite tools + // it is an explicit declaration merged over the derived safety floor. // +optional Title *string `json:"title,omitempty" yaml:"title,omitempty"` - // ReadOnlyHint overrides the read-only hint annotation. + // ReadOnlyHint sets the read-only hint annotation. + // For tool overrides this replaces the backend value; for composite tools + // it is an explicit declaration merged over the derived safety floor. // +optional ReadOnlyHint *bool `json:"readOnlyHint,omitempty" yaml:"readOnlyHint,omitempty"` - // DestructiveHint overrides the destructive hint annotation. + // DestructiveHint sets the destructive hint annotation. + // For tool overrides this replaces the backend value; for composite tools + // it is an explicit declaration merged over the derived safety floor. // +optional DestructiveHint *bool `json:"destructiveHint,omitempty" yaml:"destructiveHint,omitempty"` - // IdempotentHint overrides the idempotent hint annotation. + // IdempotentHint sets the idempotent hint annotation. + // For tool overrides this replaces the backend value; for composite tools + // it is an explicit declaration merged over the derived safety floor. // +optional IdempotentHint *bool `json:"idempotentHint,omitempty" yaml:"idempotentHint,omitempty"` - // OpenWorldHint overrides the open-world hint annotation. + // OpenWorldHint sets the open-world hint annotation. + // For tool overrides this replaces the backend value; for composite tools + // it is an explicit declaration merged over the derived safety floor. // +optional OpenWorldHint *bool `json:"openWorldHint,omitempty" yaml:"openWorldHint,omitempty"` } @@ -716,6 +788,26 @@ type CompositeToolConfig struct { // If not specified, the workflow returns the last step's output (backward compatible). // +optional Output *OutputConfig `json:"output,omitempty" yaml:"output,omitempty"` + + // Annotations declares MCP tool annotations for the composite tool. + // + // Annotation derivation runs at ADVERTISE TIME (when tools/list is served + // and the backend tools are aggregated), not at CRD admission — thv vmcp + // validate does NOT check annotation contradictions. The derived floor is + // fail-closed: when the workflow has one or more tool steps the floor is + // always non-nil, and any step whose annotations are nil/unknown taints the + // floor conservatively (readOnly=false, destructive=true, openWorld=true). + // A workflow with no tool steps (e.g. only elicitation) has no floor. + // + // When nil, annotations are derived from the annotations of the backend tools + // referenced by the workflow's steps (e.g. readOnlyHint is true only when every + // step tool is read-only). When set, the values are an explicit author + // declaration merged over the derived floor — subject to a safety-floor + // guardrail that drops the composite tool (with a warning naming the offending + // step tools) if an explicit hint would make the tool look safer than its + // steps allow. + // +optional + Annotations *ToolAnnotationsOverride `json:"annotations,omitempty" yaml:"annotations,omitempty"` } // CompositeToolRef defines a reference to a VirtualMCPCompositeToolDefinition resource. diff --git a/pkg/vmcp/config/crd_cli_roundtrip_test.go b/pkg/vmcp/config/crd_cli_roundtrip_test.go index 1ad76d3064..30d12c9a81 100644 --- a/pkg/vmcp/config/crd_cli_roundtrip_test.go +++ b/pkg/vmcp/config/crd_cli_roundtrip_test.go @@ -489,6 +489,11 @@ func TestConfigRoundtrip(t *testing.T) { Tool: "github-mcp.search_repos", }, }, + Annotations: &ToolAnnotationsOverride{ + Title: ptrTo("Test Tool"), + ReadOnlyHint: ptrTo(true), + OpenWorldHint: ptrTo(true), + }, }, }, } @@ -526,6 +531,24 @@ func TestConfigRoundtrip(t *testing.T) { if parsedConfig.CompositeTools[0].Name != "test-tool" { t.Errorf("CompositeTools[0].Name = %q, want %q", parsedConfig.CompositeTools[0].Name, "test-tool") } + // Annotations must survive the YAML round-trip (guards the yaml tags). + wantAnn := originalConfig.CompositeTools[0].Annotations + gotAnn := parsedConfig.CompositeTools[0].Annotations + if gotAnn == nil { + t.Fatal("CompositeTools[0].Annotations is nil after round-trip") + } + if gotAnn.Title == nil || *gotAnn.Title != *wantAnn.Title { + t.Errorf("Annotations.Title = %v, want %v", gotAnn.Title, *wantAnn.Title) + } + if gotAnn.ReadOnlyHint == nil || *gotAnn.ReadOnlyHint != *wantAnn.ReadOnlyHint { + t.Errorf("Annotations.ReadOnlyHint = %v, want %v", gotAnn.ReadOnlyHint, *wantAnn.ReadOnlyHint) + } + if gotAnn.OpenWorldHint == nil || *gotAnn.OpenWorldHint != *wantAnn.OpenWorldHint { + t.Errorf("Annotations.OpenWorldHint = %v, want %v", gotAnn.OpenWorldHint, *wantAnn.OpenWorldHint) + } + if gotAnn.DestructiveHint != nil || gotAnn.IdempotentHint != nil { + t.Errorf("unset annotations hints = (%v, %v), want nil", gotAnn.DestructiveHint, gotAnn.IdempotentHint) + } } // containsString checks if s contains substr. diff --git a/pkg/vmcp/config/validator.go b/pkg/vmcp/config/validator.go index 486a68b0f6..461aa993b2 100644 --- a/pkg/vmcp/config/validator.go +++ b/pkg/vmcp/config/validator.go @@ -177,6 +177,16 @@ func (v *DefaultValidator) validateIncomingAuth(auth *IncomingAuthConfig) error // - PKCE flows (public clients) // - Token validation without introspection // - Kubernetes service account token validation + + // Validate CA bundle path: reject null bytes, path traversal, and relative paths. + if auth.OIDC.CABundlePath != "" { + if strings.ContainsRune(auth.OIDC.CABundlePath, 0) || strings.Contains(auth.OIDC.CABundlePath, "..") { + return fmt.Errorf("incomingAuth.oidc.caBundlePath contains invalid path characters") + } + if !filepath.IsAbs(auth.OIDC.CABundlePath) { + return fmt.Errorf("incomingAuth.oidc.caBundlePath must be an absolute path") + } + } } // Validate authorization configuration @@ -344,9 +354,63 @@ func (v *DefaultValidator) validateAggregation(agg *AggregationConfig) error { return err } + // Mirrors the CRD's enum for the CLI path, which has no API-server admission + // to reject a bad value. "" is accepted and means allow, so configs written + // before this setting existed stay valid. + switch agg.DefaultToolVisibility { + case "", DefaultToolVisibilityAllow, DefaultToolVisibilityDeny: + default: + return fmt.Errorf("defaultToolVisibility must be one of: allow, deny") + } + + if err := validateDenyVisibilityPriorityOrder(agg); err != nil { + return err + } + return v.validateToolConfigurations(agg.Tools) } +// validateDenyVisibilityPriorityOrder rejects a priorityOrder entry that has no +// corresponding Tools entry when DefaultToolVisibility is deny. +// +// Conflict resolution runs BEFORE the advertising filter, and priorityOrder is +// independent of Tools. So an unlisted backend can win a name conflict — dropping +// the listed backend's version of that tool at resolution — and then be withheld +// by deny, leaving the tool advertised by nobody. That is strictly worse than +// allow, where the winner at least stays visible. +// +// Rejecting the combination up front is cheaper than reordering the pipeline, +// which would mean resolving conflicts twice: once over the visible candidates +// and once over the complete set that composite tools route against. +// +// validateAggregation guarantees ConflictResolutionConfig is non-nil before +// calling this, so it is safe to dereference below. +func validateDenyVisibilityPriorityOrder(agg *AggregationConfig) error { + if agg.DefaultToolVisibility != DefaultToolVisibilityDeny || + agg.ConflictResolution != vmcp.ConflictStrategyPriority { + return nil + } + + listed := make(map[string]struct{}, len(agg.Tools)) + for _, wl := range agg.Tools { + if wl != nil { + listed[wl.Workload] = struct{}{} + } + } + + for _, workload := range agg.ConflictResolutionConfig.PriorityOrder { + if _, ok := listed[workload]; !ok { + return fmt.Errorf( + "defaultToolVisibility deny with priority strategy: priorityOrder entry %q has no tools entry; "+ + "it could win a name conflict and then be withheld, hiding the tool entirely. "+ + "Add a tools entry for %q or remove it from priorityOrder", + workload, workload) + } + } + + return nil +} + // validateConflictStrategy validates strategy-specific configuration func (*DefaultValidator) validateConflictStrategy(agg *AggregationConfig) error { switch agg.ConflictResolution { diff --git a/pkg/vmcp/config/validator_test.go b/pkg/vmcp/config/validator_test.go index 7176a9ec5d..810eb05410 100644 --- a/pkg/vmcp/config/validator_test.go +++ b/pkg/vmcp/config/validator_test.go @@ -199,6 +199,57 @@ func TestValidator_ValidateIncomingAuth(t *testing.T) { wantErr: true, errMsg: "issuer is required", }, + { + name: "valid OIDC auth with absolute caBundlePath", + auth: &IncomingAuthConfig{ + Type: "oidc", + OIDC: &OIDCConfig{ + Issuer: "https://example.com", + Audience: "vmcp", + CABundlePath: "/config/certs/example-ca/ca.crt", + }, + }, + wantErr: false, + }, + { + name: "OIDC rejects relative caBundlePath", + auth: &IncomingAuthConfig{ + Type: "oidc", + OIDC: &OIDCConfig{ + Issuer: "https://example.com", + Audience: "vmcp", + CABundlePath: "certs/ca.crt", + }, + }, + wantErr: true, + errMsg: "caBundlePath must be an absolute path", + }, + { + name: "OIDC rejects caBundlePath with path traversal", + auth: &IncomingAuthConfig{ + Type: "oidc", + OIDC: &OIDCConfig{ + Issuer: "https://example.com", + Audience: "vmcp", + CABundlePath: "/config/certs/../../etc/passwd", + }, + }, + wantErr: true, + errMsg: "caBundlePath contains invalid path characters", + }, + { + name: "OIDC rejects caBundlePath with null byte", + auth: &IncomingAuthConfig{ + Type: "oidc", + OIDC: &OIDCConfig{ + Issuer: "https://example.com", + Audience: "vmcp", + CABundlePath: "/config/certs/ca\x00.crt", + }, + }, + wantErr: true, + errMsg: "caBundlePath contains invalid path characters", + }, } for _, tt := range tests { @@ -552,6 +603,89 @@ func TestValidator_ValidateAggregation(t *testing.T) { wantErr: true, errMsg: "tool overrides are required", }, + { + // Omitted defaultToolVisibility keeps pre-existing configs valid. + name: "unset defaultToolVisibility is valid", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPrefix, + ConflictResolutionConfig: &ConflictResolutionConfig{PrefixFormat: "{workload}_"}, + }, + wantErr: false, + }, + { + name: "defaultToolVisibility deny is valid", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPrefix, + ConflictResolutionConfig: &ConflictResolutionConfig{PrefixFormat: "{workload}_"}, + DefaultToolVisibility: DefaultToolVisibilityDeny, + }, + wantErr: false, + }, + { + name: "defaultToolVisibility allow is valid", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPrefix, + ConflictResolutionConfig: &ConflictResolutionConfig{PrefixFormat: "{workload}_"}, + DefaultToolVisibility: DefaultToolVisibilityAllow, + }, + wantErr: false, + }, + { + // The CLI path has no admission webhook, so a typo must be rejected here + // rather than silently falling back to advertise-everything. + name: "defaultToolVisibility rejects an unknown value", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPrefix, + ConflictResolutionConfig: &ConflictResolutionConfig{PrefixFormat: "{workload}_"}, + DefaultToolVisibility: "denied", + }, + wantErr: true, + errMsg: "defaultToolVisibility must be one of", + }, + { + // An unlisted priority backend can win a conflict and then be withheld, + // hiding the tool from every backend that offered it. + name: "deny with priority rejects an unlisted priorityOrder entry", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPriority, + ConflictResolutionConfig: &ConflictResolutionConfig{ + PriorityOrder: []string{"github", "jira"}, + }, + Tools: []*WorkloadToolConfig{{Workload: "github"}}, + DefaultToolVisibility: DefaultToolVisibilityDeny, + }, + wantErr: true, + errMsg: `priorityOrder entry "jira" has no tools entry`, + }, + { + name: "deny with priority accepts a fully listed priorityOrder", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPriority, + ConflictResolutionConfig: &ConflictResolutionConfig{ + PriorityOrder: []string{"github", "jira"}, + }, + Tools: []*WorkloadToolConfig{ + {Workload: "github"}, + {Workload: "jira"}, + }, + DefaultToolVisibility: DefaultToolVisibilityDeny, + }, + wantErr: false, + }, + { + // Under allow the unlisted winner is still advertised, so the + // combination is harmless and must stay valid. + name: "allow with priority permits an unlisted priorityOrder entry", + agg: &AggregationConfig{ + ConflictResolution: vmcp.ConflictStrategyPriority, + ConflictResolutionConfig: &ConflictResolutionConfig{ + PriorityOrder: []string{"github", "jira"}, + }, + Tools: []*WorkloadToolConfig{{Workload: "github"}}, + DefaultToolVisibility: DefaultToolVisibilityAllow, + }, + wantErr: false, + }, } for _, tt := range tests { @@ -605,6 +739,28 @@ func TestValidator_ValidateCompositeTools(t *testing.T) { }, wantErr: false, }, + { + name: "valid composite tool with annotations", + tools: []CompositeToolConfig{ + { + Name: "report_workflow", + Description: "Read-only report workflow", + Timeout: Duration(30 * time.Minute), + Steps: []WorkflowStepConfig{ + { + ID: "fetch", + Type: "tool", + Tool: "github.get_pr", + }, + }, + Annotations: &ToolAnnotationsOverride{ + Title: ptrTo("Report Workflow"), + ReadOnlyHint: ptrTo(true), + }, + }, + }, + wantErr: false, + }, { name: "missing tool name", tools: []CompositeToolConfig{ diff --git a/pkg/vmcp/config/zz_generated.deepcopy.go b/pkg/vmcp/config/zz_generated.deepcopy.go index a541abdd6f..187d321e44 100644 --- a/pkg/vmcp/config/zz_generated.deepcopy.go +++ b/pkg/vmcp/config/zz_generated.deepcopy.go @@ -124,6 +124,11 @@ func (in *CompositeToolConfig) DeepCopyInto(out *CompositeToolConfig) { *out = new(OutputConfig) (*in).DeepCopyInto(*out) } + if in.Annotations != nil { + in, out := &in.Annotations, &out.Annotations + *out = new(ToolAnnotationsOverride) + (*in).DeepCopyInto(*out) + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CompositeToolConfig. diff --git a/pkg/vmcp/conversion/content.go b/pkg/vmcp/conversion/content.go index e22cb8427a..b4f20a6573 100644 --- a/pkg/vmcp/conversion/content.go +++ b/pkg/vmcp/conversion/content.go @@ -10,8 +10,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) diff --git a/pkg/vmcp/conversion/content_test.go b/pkg/vmcp/conversion/content_test.go index c8a3f4cc56..7a89c768ef 100644 --- a/pkg/vmcp/conversion/content_test.go +++ b/pkg/vmcp/conversion/content_test.go @@ -6,10 +6,10 @@ package conversion import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) diff --git a/pkg/vmcp/conversion/conversion_test.go b/pkg/vmcp/conversion/conversion_test.go index 41044e47b5..0d9a95a07c 100644 --- a/pkg/vmcp/conversion/conversion_test.go +++ b/pkg/vmcp/conversion/conversion_test.go @@ -6,10 +6,10 @@ package conversion_test import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" ) @@ -953,6 +953,60 @@ func TestToMCPMeta(t *testing.T) { }, }, }, + { + // #5986: vMCP, not the backend, is the client's MCP peer, so a backend + // must not be able to set reserved io.modelcontextprotocol/* keys on a + // response. Only serverInfo is even schema-legal on a result; the + // request-only keys arriving on one are a backend fabricating the + // client's own identity. + name: "reserved keys are stripped, non-reserved survive", + input: map[string]any{ + "io.modelcontextprotocol/serverInfo": map[string]any{"name": "attacker"}, + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": map[string]any{}, + "io.modelcontextprotocol/futureThing": "whatever", + "traceId": "trace-keep", + }, + expected: &mcp.Meta{ + AdditionalFields: map[string]any{"traceId": "trace-keep"}, + }, + }, + { + // Collapses to nil rather than an empty _meta object on the wire. + name: "map of only reserved keys returns nil", + input: map[string]any{ + "io.modelcontextprotocol/serverInfo": map[string]any{"name": "attacker"}, + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + }, + expected: nil, + }, + { + // The reserved namespace is not stripped wholesale: end-to-end keys + // (task correlation) ride through while per-hop control keys go. + name: "passthrough reserved key survives alongside a stripped one", + input: map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/related-task": map[string]any{"taskId": "t-1"}, + }, + expected: &mcp.Meta{ + AdditionalFields: map[string]any{ + "io.modelcontextprotocol/related-task": map[string]any{"taskId": "t-1"}, + }, + }, + }, + { + // Pins the ordering: the strip must run BEFORE the progressToken split, + // or a reserved key would land in AdditionalFields. + name: "reserved keys stripped alongside a progressToken", + input: map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "progressToken": "token-xyz", + }, + expected: &mcp.Meta{ + ProgressToken: "token-xyz", + AdditionalFields: map[string]any{}, + }, + }, } for _, tt := range tests { diff --git a/pkg/vmcp/conversion/errors.go b/pkg/vmcp/conversion/errors.go index 83058a514f..c3f3a057e7 100644 --- a/pkg/vmcp/conversion/errors.go +++ b/pkg/vmcp/conversion/errors.go @@ -6,8 +6,7 @@ package conversion import ( "errors" - sdkmcp "github.com/mark3labs/mcp-go/mcp" - + sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvmcp "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -18,13 +17,17 @@ import ( // Go errors: the SDK maps them to generic internal errors. For domain errors // that opt into CodedError, preserve the code/data in StructuredContent instead. func ErrorToToolResult(err error) *sdkmcp.CallToolResult { + // Denial first (matches writeModernDispatchError): an error that is both a + // CodedError and wraps ErrAuthorizationFailed must render as the denial, + // never as retry-shaped coded data (the sets are disjoint today; this + // ordering is the invariant). + if errors.Is(err, vmcp.ErrAuthorizationFailed) { + return sdkmcp.NewToolResultError(vmcp.DenyMessageToolCall) + } var coded thvmcp.CodedError if errors.As(err, &coded) { return CodedErrorResult(err, coded) } - if errors.Is(err, vmcp.ErrAuthorizationFailed) { - return sdkmcp.NewToolResultError("call denied by authorization policy") - } return sdkmcp.NewToolResultError(err.Error()) } @@ -37,7 +40,7 @@ func CodedErrorResult(err error, coded thvmcp.CodedError) *sdkmcp.CallToolResult "code": coded.Code(), "message": err.Error(), } - if data := coded.Data(); data != nil { + if data := coded.Data(); len(data) > 0 { structured["data"] = data } result.StructuredContent = structured diff --git a/pkg/vmcp/conversion/meta.go b/pkg/vmcp/conversion/meta.go index c7b6d03b0b..17932aaad8 100644 --- a/pkg/vmcp/conversion/meta.go +++ b/pkg/vmcp/conversion/meta.go @@ -6,7 +6,8 @@ package conversion import ( "maps" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" ) // FromMCPMeta converts MCP SDK meta to map[string]any for vmcp wrapper types. @@ -41,9 +42,23 @@ func FromMCPMeta(meta *mcp.Meta) map[string]any { // ToMCPMeta converts vmcp meta map to MCP SDK meta for forwarding to clients. // This reconstructs the _meta field when sending responses back through the MCP protocol. // +// Reserved io.modelcontextprotocol/* keys are stripped first: vMCP, not the +// backend, is the client's MCP peer, so a backend must not speak for it. This is +// the single chokepoint for every Legacy egress that carries backend _meta +// (serve_handlers, sessionmanager, the elicitation adapter); the Modern path's +// mirror is newModernResultMeta, and both call the same mcpparser helper so the +// two revisions cannot drift. +// +// Note the Legacy/Modern asymmetry: Modern re-stamps its own +// io.modelcontextprotocol/serverInfo after stripping, and Legacy deliberately +// does not — the 2025-11-25 revision has no serverInfo _meta key at all. The +// missing Legacy stamp is correct, not an oversight. +// // Returns nil if meta is nil or empty, following the MCP specification that -// _meta is optional and should be omitted when empty. +// _meta is optional and should be omitted when empty. A map consisting only of +// reserved keys therefore collapses to nil rather than an empty _meta object. func ToMCPMeta(meta map[string]any) *mcp.Meta { + meta = mcpparser.StripReservedMeta(meta) if len(meta) == 0 { return nil } diff --git a/pkg/vmcp/core/admission_test.go b/pkg/vmcp/core/admission_test.go index 8a0064d1d1..cba2edafb8 100644 --- a/pkg/vmcp/core/admission_test.go +++ b/pkg/vmcp/core/admission_test.go @@ -545,7 +545,7 @@ func TestAdmission_ListCallLookupEnforceSameDecision(t *testing.T) { // Only the permitted tool ever reaches the backend. want := &vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}} - m.client.EXPECT().CallTool(gomock.Any(), gomock.Any(), "weather", gomock.Any(), gomock.Any()).Return(want, nil) + m.client.EXPECT().CallTool(gomock.Any(), gomock.Any(), "weather", gomock.Any(), gomock.Any(), gomock.Any()).Return(want, nil) c, err := New(cfg) require.NoError(t, err) @@ -601,7 +601,7 @@ func TestAdmission_AnnotationGatedDecisionMatchesListAndCall(t *testing.T) { // Only the read-only tool is callable, so only it reaches the backend. want := &vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}} - m.client.EXPECT().CallTool(gomock.Any(), gomock.Any(), "ro", gomock.Any(), gomock.Any()).Return(want, nil) + m.client.EXPECT().CallTool(gomock.Any(), gomock.Any(), "ro", gomock.Any(), gomock.Any(), gomock.Any()).Return(want, nil) c, err := New(cfg) require.NoError(t, err) @@ -684,7 +684,7 @@ func TestAdmission_CallToolForwardsArgsToAuthorizer(t *testing.T) { // Only the args-satisfying call reaches the backend. want := &vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}} - m.client.EXPECT().CallTool(gomock.Any(), gomock.Any(), "deploy", gomock.Any(), gomock.Any()).Return(want, nil) + m.client.EXPECT().CallTool(gomock.Any(), gomock.Any(), "deploy", gomock.Any(), gomock.Any(), gomock.Any()).Return(want, nil) c, err := New(cfg) require.NoError(t, err) diff --git a/pkg/vmcp/core/core.go b/pkg/vmcp/core/core.go index 1bf4111fb6..4c2a9be74a 100644 --- a/pkg/vmcp/core/core.go +++ b/pkg/vmcp/core/core.go @@ -32,6 +32,17 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/router" ) +// DiscoverCapabilities summarizes, for one identity, whether each capability +// kind has at least one admission-filtered entry. It carries no descriptor +// arrays -- only presence flags -- unlike the []vmcp.Tool/Resource/etc. slices +// ListTools/ListResources/ListResourceTemplates/ListPrompts return. +type DiscoverCapabilities struct { + HasTools bool + HasResources bool + HasResourceTemplates bool + HasPrompts bool +} + // VMCP is the core Virtual MCP domain object. // // Contract: @@ -74,6 +85,14 @@ type VMCP interface { // ListResources returns the resources advertised to identity. ListResources(ctx context.Context, identity *auth.Identity) ([]vmcp.Resource, error) + // ListResourceTemplates returns the resource templates advertised to identity. + // Each template's URI template is treated as the resource id for the admission + // filter (the same read-side decision ListResources applies), so a template the + // identity may not read is withheld. Reading an expanded URI reuses ReadResource + // (the router matches the URI against these templates); there is no dedicated + // template read method. + ListResourceTemplates(ctx context.Context, identity *auth.Identity) ([]vmcp.ResourceTemplate, error) + // ReadResource reads the resource at uri on behalf of identity. The // returned result's Meta may be nil (see the interface contract). ReadResource(ctx context.Context, identity *auth.Identity, uri string) (*vmcp.ResourceReadResult, error) @@ -89,6 +108,27 @@ type VMCP interface { GetPrompt(ctx context.Context, identity *auth.Identity, name string, args map[string]any) (*vmcp.PromptGetResult, error) + // Complete resolves argument-completion candidates for a prompt argument or a + // resource-template variable on behalf of identity (the completion/complete + // method). + // + // ref selects the target: a CompletionRefTypePrompt ref resolves the backend via + // the prompts routing table (by ref.Name); a CompletionRefTypeResource ref resolves + // it via the resource-templates routing table (matching ref.URI against templates) + // with an exact concrete-resource fallback. The referenced capability is + // admission-checked (the same get/read decision GetPrompt/ReadResource enforce) + // before the backend is queried. + // + // argName/argValue are the argument being completed; contextArgs carries + // previously-resolved argument values for multi-argument completion (may be nil). + // + // An unroutable ref, or a backend that does not advertise completions, yields an + // empty (non-nil) CompletionResult rather than an error — matching the MCP spec's + // lenient completion semantics. Admission denial still returns an error wrapping + // vmcp.ErrAuthorizationFailed. See ListTools for the nil/anonymous identity semantics. + Complete(ctx context.Context, identity *auth.Identity, ref vmcp.CompletionRef, + argName, argValue string, contextArgs map[string]string) (*vmcp.CompletionResult, error) + // LookupTool resolves an advertised tool name to its capability (incl. // BackendID) WITHOUT invoking it. Returns an error for an unknown or // unadvertised name — the validation seam for the call path. Lookups apply @@ -106,6 +146,29 @@ type VMCP interface { // unadvertised name. Applies the same admission filter as ListPrompts. LookupPrompt(ctx context.Context, identity *auth.Identity, name string) (*vmcp.Prompt, error) + // CheckToolCall runs the CALL-side admission decision for the named tool + // WITHOUT invoking it: it returns nil when identity is allowed to call name + // with args, and an error wrapping [vmcp.ErrAuthorizationFailed] on deny AND on + // an authorizer error (fail closed). It shares the exact admission block CallTool + // runs, so a pre-flight check and the call can never drift. + // + // Unlike LookupTool (list-side, no args), this is the call-side decision: it + // evaluates argument-conditional policies with the real args. identity is never + // logged. See ListTools for the nil/anonymous semantics. + CheckToolCall(ctx context.Context, identity *auth.Identity, name string, args map[string]any) error + + // CheckResourceRead runs the read-side admission decision for uri WITHOUT + // reading it. Same contract as CheckToolCall: nil when allowed, an error + // wrapping [vmcp.ErrAuthorizationFailed] on deny or authorizer error (fail + // closed). Shares ReadResource's admission block. + CheckResourceRead(ctx context.Context, identity *auth.Identity, uri string) error + + // CheckPromptGet runs the get-side admission decision for the named prompt + // WITHOUT retrieving it. Same contract as CheckToolCall: nil when allowed, an + // error wrapping [vmcp.ErrAuthorizationFailed] on deny or authorizer error (fail + // closed). Shares GetPrompt's admission block. + CheckPromptGet(ctx context.Context, identity *auth.Identity, name string) error + // ListBackends returns the backends visible to identity, scoped to the gateway's // group (groupRef always applies; the BackendRegistry is the group-scoped source). // @@ -145,12 +208,35 @@ type VMCP interface { // corresponding authorized list would not show. LookupBackend(ctx context.Context, identity *auth.Identity, backendID string) (*vmcp.Backend, error) + // Discover returns identity's capability-presence flags -- whether it is + // admitted to at least one tool, resource, resource template, and prompt -- + // from a SINGLE aggregation of backend capabilities, for server/discover. + // + // It applies the exact same admission-filtered code paths + // ListTools/ListResources/ListResourceTemplates/ListPrompts use against one + // shared aggregated view, so a flag is true iff the ADMISSION-FILTERED set + // for that capability is non-empty -- never derived from the raw aggregate, + // which would leak capabilities identity cannot reach. This mirrors the + // post-admission-summary precedent ListBackends(filterUnauthorized=true) + // established, applied to presence flags instead of a backend list. See + // ListTools for the nil/anonymous identity semantics. + Discover(ctx context.Context, identity *auth.Identity) (DiscoverCapabilities, error) + // BackendHealth returns the backend health reporter the core owns, or nil when health // monitoring is disabled. The core builds, starts, and (via Close) stops the monitor and // filters capabilities with it; the transport layer uses this only to report on or sync // backend health (e.g. the /health route and periodic status reporting). BackendHealth() health.Reporter + // InvalidateCapabilityCache forces the next List/Lookup/Call on every identity + // to re-sweep backend capabilities rather than serve a cached view. Serve's + // Server calls it after a backend reports notifications/tools/list_changed + // (#5748), before re-deriving the session's advertised tool set — otherwise + // the resync would immediately re-read the now-stale cached aggregation. It + // is a no-op (WARN-logged, not silent) when the configured Aggregator does + // not memoize results — see aggregator.CacheInvalidator. + InvalidateCapabilityCache() + // Close releases core-held resources (backend connections, the health monitor, etc.). // Implementations must be idempotent: calling Close multiple times returns nil. Close() error diff --git a/pkg/vmcp/core/core_calls.go b/pkg/vmcp/core/core_calls.go index d928c4e464..452ad77245 100644 --- a/pkg/vmcp/core/core_calls.go +++ b/pkg/vmcp/core/core_calls.go @@ -10,8 +10,10 @@ import ( "fmt" "log/slog" "maps" + "slices" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/composer" "github.com/stacklok/toolhive/pkg/vmcp/router" @@ -23,6 +25,13 @@ import ( // from the same table. Returns vmcp.ErrNotFound for an unadvertised name and // vmcp.ErrAuthorizationFailed when admission denies identity the call. // +// "Unadvertised" here means the AGGREGATION view that ListTools filters — NOT +// the routing table, which intentionally holds more (see the advertised-view +// check below). Admission narrowing on top of that view is enforced separately, +// by authorizeToolCall. So a tool hidden by excludeAllTools, per-workload +// excludeAll, or filter is not directly callable, while composite workflow steps +// may still reach it. +// // args and meta are treated as read-only and copied before being forwarded // (go-style: copy before mutating caller input). The admission decision enforces // the same policy ListTools filters on. identity is never logged. See ListTools @@ -42,28 +51,43 @@ func (c *coreVMCP) CallTool( return nil, err } - // Admission deny: enforce the same decision ListTools filters on, sourcing the - // tool's annotations from the advertised set (mirroring the annotation cache the - // HTTP middleware populates from tools/list) so annotation-gated policies - // evaluate. advertisedTools includes composites, so their annotations are sourced - // too. A name absent from the advertised set carries no annotations, so an - // annotation-gated decision evaluates with no hints (and may deny). In normal - // operation the advertised set and the routing table are derived from the same - // aggregation, so this only arises if they diverge. - tool := findAdvertisedTool(c.advertisedTools(agg), name) - if tool == nil { - tool = &vmcp.Tool{Name: name} + // Resolve the advertised view ONCE and thread it through the three consumers + // below (admission, the not-found guard, the composite branch). Each of + // accessibleComposites and advertisedTools re-runs FilterWorkflowDefsForSession, + // ValidateNoToolConflicts and ConvertWorkflowDefsToTools, so computing them per + // consumer meant three passes and two concatenated-slice allocations per call. + // Sharing one result also removes any chance of the admission lookup and the + // guard below disagreeing about what is advertised. + composites := c.accessibleComposites(agg) + advertised := advertisedToolsWith(agg, composites) + tool := findAdvertisedTool(advertised, name) + + if err := c.authorizeToolCall(ctx, identity, name, argsCopy, tool); err != nil { + return nil, err } - // An authorizer ERROR is treated as a denial (fail closed), classified as - // ErrAuthorizationFailed so the Serve adapter can distinguish it from a transport - // failure via errors.Is — mirroring the live authorizeAndServe, which routes both - // err and !authorized to the unauthorized response. The underlying error is - // preserved in the chain for server-side diagnostics; Serve decides what reaches - // the client. - if allowed, err := c.admission.AllowToolCall(ctx, identity, tool, argsCopy); err != nil { - return nil, fmt.Errorf("%w: tool %q: %w", vmcp.ErrAuthorizationFailed, name, err) - } else if !allowed { - return nil, fmt.Errorf("%w: tool %q", vmcp.ErrAuthorizationFailed, name) + + // Hold the call to the ADVERTISED view, not the routing table. The routing + // table deliberately carries every backend tool — including those hidden from + // tools/list by excludeAllTools, per-workload excludeAll, or filter — so that + // composite-tool workflow STEPS can reach them (#3636, + // aggregator/default_aggregator.go:349). Workflow steps never come through + // here: a composite enters CallTool once under its own (advertised) name and + // its steps then run composer -> router.RouteTool -> backendClient.CallTool, + // bypassing this method. So narrowing to the advertised set closes direct + // invocation of a hidden tool without weakening #3636 at all. + // + // This is checked AFTER authorizeToolCall on purpose: a tool that admission + // denies must keep returning ErrAuthorizationFailed, never ErrNotFound, or the + // two errors together would let a caller probe which denied tools exist. + // + // advertised includes accessible composites, so an advertised workflow still + // resolves here and falls through to the composite branch below. Note this + // also rejects RouteTool's "{workloadID}.{toolName}" alias + // (router/session_router.go:105) for a direct call: an alias is never an + // advertised name. That alias exists for workflow step definitions and keeps + // working there, inside the composer. + if tool == nil { + return nil, fmt.Errorf("%w: tool %q", vmcp.ErrNotFound, name) } // Composite tool: execute only when the workflow is actually advertised in the @@ -71,7 +95,7 @@ func (c *coreVMCP) CallTool( // uses the same gate as ListTools (accessibleComposites), so advertised equals // executed. A name that collides with a backend tool is NOT in the set and falls // through to backend routing, matching the legacy decorator. - if def, ok := c.accessibleComposites(agg)[name]; ok { + if def, ok := composites[name]; ok { engine := c.composerFactory(agg.RoutingTable, agg.Tools) return executeComposite(ctx, engine, def, argsCopy) } @@ -86,7 +110,23 @@ func (c *coreVMCP) CallTool( } return nil, fmt.Errorf("routing tool %q: %w", name, err) } - return c.backendClient.CallTool(ctx, target, name, argsCopy, metaCopy) + // SEP-2243 Mcp-Param-* mirroring: a backend MAY designate tool parameters, + // via x-mcp-header in its inputSchema, whose values it also expects as HTTP + // headers — and rejects the call with -32020 if they are missing. Deriving + // that here is what makes an annotating backend callable at all. The core is + // the right place because the aggregated view already holds the tool's + // schema, so the backend client needs no schema cache of its own. + paramHeaders, err := paramHeadersFor(agg.Tools, name, argsCopy) + if err != nil { + return nil, err + } + + result, err := c.backendClient.CallTool(ctx, target, name, argsCopy, metaCopy, paramHeaders) + if err != nil { + return nil, err + } + result.BackendID = target.WorkloadID + return result, nil } // ReadResource reads the resource at uri from its backend. Returns @@ -102,13 +142,8 @@ func (c *coreVMCP) ReadResource( return nil, err } - // Admission deny: mirror ListResources' filter for the single read. Resources - // carry no annotations, so the URI alone identifies the decision. An authorizer - // error fails closed, classified as ErrAuthorizationFailed (see CallTool). - if allowed, err := c.admission.AllowResourceRead(ctx, identity, &vmcp.Resource{URI: uri}); err != nil { - return nil, fmt.Errorf("%w: resource %q: %w", vmcp.ErrAuthorizationFailed, uri, err) - } else if !allowed { - return nil, fmt.Errorf("%w: resource %q", vmcp.ErrAuthorizationFailed, uri) + if err := c.authorizeResourceRead(ctx, identity, uri); err != nil { + return nil, err } target, err := router.NewSessionRouter(agg.RoutingTable).RouteResource(ctx, uri) @@ -120,7 +155,12 @@ func (c *coreVMCP) ReadResource( } // Pass the advertised URI; the backend client owns the single translation to // the backend's capability name (client.go:874), matching CallTool. - return c.backendClient.ReadResource(ctx, target, uri) + result, err := c.backendClient.ReadResource(ctx, target, uri) + if err != nil { + return nil, err + } + result.BackendID = target.WorkloadID + return result, nil } // GetPrompt retrieves the named prompt from its backend. args is treated as @@ -138,13 +178,8 @@ func (c *coreVMCP) GetPrompt( return nil, err } - // Admission deny: mirror ListPrompts' filter for the single get. Prompts carry - // no annotations, so the name alone identifies the decision. An authorizer error - // fails closed, classified as ErrAuthorizationFailed (see CallTool). - if allowed, err := c.admission.AllowPromptGet(ctx, identity, &vmcp.Prompt{Name: name}); err != nil { - return nil, fmt.Errorf("%w: prompt %q: %w", vmcp.ErrAuthorizationFailed, name, err) - } else if !allowed { - return nil, fmt.Errorf("%w: prompt %q", vmcp.ErrAuthorizationFailed, name) + if err := c.authorizePromptGet(ctx, identity, name); err != nil { + return nil, err } target, err := router.NewSessionRouter(agg.RoutingTable).RoutePrompt(ctx, name) @@ -156,7 +191,79 @@ func (c *coreVMCP) GetPrompt( } // Pass the advertised name; the backend client owns the single translation to // the backend's capability name (client.go:927), matching CallTool. - return c.backendClient.GetPrompt(ctx, target, name, maps.Clone(args)) + result, err := c.backendClient.GetPrompt(ctx, target, name, maps.Clone(args)) + if err != nil { + return nil, err + } + result.BackendID = target.WorkloadID + return result, nil +} + +// Complete resolves argument-completion candidates for the referenced prompt or +// resource template. It resolves the backend from the freshly aggregated routing +// table (prompts table for a prompt ref, resource-templates table with a concrete +// fallback for a resource ref), admission-checks the referenced capability (the same +// get/read decision GetPrompt/ReadResource enforce), and forwards to the backend. +// +// An unroutable ref returns an empty (non-nil) result rather than an error, matching +// the MCP spec's lenient completion semantics (a client asking for completions on an +// unknown ref should get no candidates, not a protocol error). Admission denial +// returns an error wrapping vmcp.ErrAuthorizationFailed. See ListTools for identity +// semantics; identity is never logged. +func (c *coreVMCP) Complete( + ctx context.Context, + identity *auth.Identity, + ref vmcp.CompletionRef, + argName, argValue string, + contextArgs map[string]string, +) (*vmcp.CompletionResult, error) { + agg, err := c.aggregatedView(ctx) + if err != nil { + return nil, err + } + + sessionRouter := router.NewSessionRouter(agg.RoutingTable) + + switch ref.Type { + case vmcp.CompletionRefTypePrompt: + if err := c.authorizePromptGet(ctx, identity, ref.Name); err != nil { + return nil, err + } + target, err := sessionRouter.RoutePrompt(ctx, ref.Name) + if err != nil { + if errors.Is(err, router.ErrPromptNotFound) { + return emptyCompletion(), nil + } + return nil, fmt.Errorf("routing prompt %q for completion: %w", ref.Name, err) + } + return c.backendClient.Complete(ctx, target, ref, argName, argValue, contextArgs) + + case vmcp.CompletionRefTypeResource: + if err := c.authorizeResourceRead(ctx, identity, ref.URI); err != nil { + return nil, err + } + // RouteResource matches the URI against concrete resources first, then the + // resource-template table (the same fallback ReadResource uses). + target, err := sessionRouter.RouteResource(ctx, ref.URI) + if err != nil { + if errors.Is(err, router.ErrResourceNotFound) { + return emptyCompletion(), nil + } + return nil, fmt.Errorf("routing resource %q for completion: %w", ref.URI, err) + } + return c.backendClient.Complete(ctx, target, ref, argName, argValue, contextArgs) + + default: + // Unknown ref type: no candidates, not a hard error (lenient completion). + slog.Debug("unknown completion ref type, returning empty completion", "ref_type", ref.Type) + return emptyCompletion(), nil + } +} + +// emptyCompletion returns a non-nil, empty completion result. It is the lenient +// answer for an unroutable or unknown completion ref. +func emptyCompletion() *vmcp.CompletionResult { + return &vmcp.CompletionResult{Values: []string{}} } // executeComposite runs a composite-tool workflow and converts the result to a @@ -197,6 +304,36 @@ func executeComposite( }, nil } +// paramHeadersFor derives the SEP-2243 Mcp-Param-* headers for a call to the +// named tool, from the tool's own x-mcp-header annotations and the call's +// arguments. Returns nil when the tool is not in the view or declares no +// annotations — the common case, and cheap: ParamHeaders exits immediately on a +// schema with no annotation. +// +// A malformed annotation cannot reach here: pkg/vmcp/client rejects such a tool +// at ingestion, so it is never advertised. An UNMIRRORABLE ARGUMENT can, though — +// the caller supplies the values — and is surfaced as an invalid-parameters error +// rather than dropped, because silently omitting the header would make the +// backend answer -32020 and turn a caller mistake into an opaque backend failure. +func paramHeadersFor(tools []vmcp.Tool, name string, args map[string]any) (map[string]string, error) { + idx := slices.IndexFunc(tools, func(t vmcp.Tool) bool { return t.Name == name }) + if idx < 0 { + return nil, nil + } + headers, err := mcpparser.ParamHeadersForSchema(tools[idx].InputSchema, args) + if err == nil { + return headers, nil + } + if errors.Is(err, mcpparser.ErrUnmirrorableValue) { + // The caller's argument value is at fault (a control character, a + // non-integral integer), so name it as invalid input. + return nil, fmt.Errorf("%w: tool %q: %w", vmcp.ErrInvalidInput, name, err) + } + // A malformed annotation. Defensive: ingestion already rejected this shape, so + // reaching here is an internal inconsistency, not the caller's mistake. + return nil, fmt.Errorf("tool %q has an invalid x-mcp-header annotation: %w", name, err) +} + // compositeErrorResult builds a tool-level error result for a failed workflow. func compositeErrorResult(msg string) *vmcp.ToolCallResult { return &vmcp.ToolCallResult{ diff --git a/pkg/vmcp/core/core_calls_test.go b/pkg/vmcp/core/core_calls_test.go index c456d1e28a..de11aff205 100644 --- a/pkg/vmcp/core/core_calls_test.go +++ b/pkg/vmcp/core/core_calls_test.go @@ -15,6 +15,7 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/aggregator" "github.com/stacklok/toolhive/pkg/vmcp/composer" + "github.com/stacklok/toolhive/pkg/vmcp/config" ) // expectAggregation wires reg.List + agg.AggregateCapabilities to return agg once. @@ -35,7 +36,7 @@ func TestCallTool_RoutesToBackend(t *testing.T) { want := &vmcp.ToolCallResult{StructuredContent: map[string]any{"result": "ok"}} m.client.EXPECT(). - CallTool(gomock.Any(), gomock.Any(), "tool_a", gomock.Any(), gomock.Any()). + CallTool(gomock.Any(), gomock.Any(), "tool_a", gomock.Any(), gomock.Any(), gomock.Any()). Return(want, nil) c, err := New(cfg) @@ -45,6 +46,7 @@ func TestCallTool_RoutesToBackend(t *testing.T) { got, err := c.CallTool(context.Background(), nil, "tool_a", map[string]any{"a": 1}, nil) require.NoError(t, err) assert.Equal(t, want, got) + assert.Equal(t, testBackendID, got.BackendID, "CallTool must stamp the routed target's backend onto the result") } func TestCallTool_NotFound(t *testing.T) { @@ -61,20 +63,62 @@ func TestCallTool_NotFound(t *testing.T) { assert.ErrorIs(t, err, vmcp.ErrNotFound) } +// TestCallTool_RejectsUnadvertisedRoutableTool is the hidden-tool regression +// guard at the core boundary. The routing table intentionally holds EVERY backend +// tool, including ones hidden from tools/list by excludeAllTools, per-workload +// excludeAll, or filter, so that composite workflow steps can still reach them +// (#3636, aggregator/default_aggregator.go:349). Before this guard, CallTool +// resolved a name straight against that table, so a hidden tool was directly +// callable by name — its ErrNotFound contract was documented but unenforced. +// +// The Times(0) on the backend client is the real security assertion: an +// ErrNotFound alone would still pass if the call had already been forwarded to +// the backend and only the response discarded. +func TestCallTool_RejectsUnadvertisedRoutableTool(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + target := backendTarget() + expectAggregation(m, &aggregator.AggregatedCapabilities{ + // "hidden" is routable but NOT advertised — exactly the shape excludeAll + // and filter produce. + Tools: []vmcp.Tool{backendTool("visible")}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{ + "visible": target, + "hidden": target, + }}, + }) + + m.client.EXPECT(). + CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Times(0) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.CallTool(context.Background(), nil, "hidden", nil, nil) + assert.ErrorIs(t, err, vmcp.ErrNotFound, + "a routable-but-unadvertised tool must not be directly callable") +} + func TestCallTool_CopyBeforeMutate(t *testing.T) { t.Parallel() cfg, m := baseConfig(t) target := backendTarget() expectAggregation(m, &aggregator.AggregatedCapabilities{ + // Advertised as well as routable: CallTool holds a direct call to the + // advertised view, so a routing-table-only entry would not reach the backend. + Tools: []vmcp.Tool{backendTool("tool_a")}, RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{"tool_a": target}}, }) // The backend client mutates the maps it receives; the caller's originals // must be untouched because CallTool forwards clones. m.client.EXPECT(). - CallTool(gomock.Any(), gomock.Any(), "tool_a", gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, args, meta map[string]any) (*vmcp.ToolCallResult, error) { + CallTool(gomock.Any(), gomock.Any(), "tool_a", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget, _ string, args, meta map[string]any, _ map[string]string) (*vmcp.ToolCallResult, error) { args["injected"] = true meta["injected"] = true return &vmcp.ToolCallResult{}, nil @@ -109,7 +153,7 @@ func TestCallTool_CompositeWorkflow(t *testing.T) { // The composite workflow's single tool step routes to the backend through the // per-call composer built from the aggregated routing table. m.client.EXPECT(). - CallTool(gomock.Any(), gomock.Any(), "be1.echo", gomock.Any(), gomock.Any()). + CallTool(gomock.Any(), gomock.Any(), "be1.echo", gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}}, nil) c, err := New(cfg) @@ -121,6 +165,7 @@ func TestCallTool_CompositeWorkflow(t *testing.T) { require.NotNil(t, got) assert.False(t, got.IsError) assert.Equal(t, true, got.StructuredContent["ok"]) + assert.Empty(t, got.BackendID, "a composite tool has no single serving backend") } func TestCallTool_CompositeNotAccessible(t *testing.T) { @@ -141,6 +186,51 @@ func TestCallTool_CompositeNotAccessible(t *testing.T) { assert.ErrorIs(t, err, vmcp.ErrNotFound) } +// TestStepAnnotationResolver_DottedNameResolution (Q5) verifies that a +// composite-tool step reference using the "{workloadID}.{toolName}" dot +// convention resolves to the correct backend annotations through the routing +// table's resolved (prefixed) key. With prefix conflict resolution the routing +// table stores "be1_echo" while the step still references "be1.echo"; the +// resolver must match via WorkloadID + original capability name and return the +// annotations registered under the resolved name. +func TestStepAnnotationResolver_DottedNameResolution(t *testing.T) { + t.Parallel() + + trueVal, falseVal := true, false + backendAnn := &vmcp.ToolAnnotations{ + ReadOnlyHint: &trueVal, + DestructiveHint: &falseVal, + OpenWorldHint: &falseVal, + } + + agg := &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{ + {Name: "be1_echo", BackendID: testBackendID, Annotations: backendAnn}, + }, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{ + "be1_echo": { + WorkloadID: testBackendID, + BaseURL: "http://" + testBackendID + ":8080", + OriginalCapabilityName: "echo", + }, + }}, + } + + resolver := stepAnnotationResolver(agg) + require.NotNil(t, resolver) + + // Dotted reference resolves to the prefixed routing-table key's annotations. + got := resolver("be1.echo") + require.NotNil(t, got, "dotted step ref must resolve through the routing table") + assert.Same(t, backendAnn, got, "resolved annotations must come from the resolved (prefixed) tool name") + + // The resolved key itself also resolves (exact-match fast path). + assert.Same(t, backendAnn, resolver("be1_echo")) + + // An unknown step tool resolves to nil (treated conservatively upstream). + assert.Nil(t, resolver("be1.unknown")) +} + func TestReadResource(t *testing.T) { t.Parallel() cfg, m := baseConfig(t) @@ -160,6 +250,7 @@ func TestReadResource(t *testing.T) { got, err := c.ReadResource(context.Background(), nil, "file://a") require.NoError(t, err) assert.Equal(t, want, got) + assert.Equal(t, testBackendID, got.BackendID, "ReadResource must stamp the routed target's backend onto the result") } func TestReadResource_NotFound(t *testing.T) { @@ -194,6 +285,7 @@ func TestGetPrompt(t *testing.T) { got, err := c.GetPrompt(context.Background(), nil, "p1", map[string]any{"x": 1}) require.NoError(t, err) assert.Equal(t, want, got) + assert.Equal(t, testBackendID, got.BackendID, "GetPrompt must stamp the routed target's backend onto the result") } func TestGetPrompt_CopyBeforeMutate(t *testing.T) { @@ -238,31 +330,92 @@ func TestGetPrompt_NotFound(t *testing.T) { } // TestCallTool_ResolvesRenamedTool exercises the conflict-resolution name path: -// the advertised name uses the dot convention ("be1.echo") while the routing -// table is keyed by the prefixed resolved name ("be1_echo") whose target carries -// a non-empty OriginalCapabilityName. The session router resolves it via -// GetBackendCapabilityName, and the core forwards the advertised name to the -// client (the client owns the translation to the backend's capability name). +// the advertised and routing-table name is the prefixed resolved name +// ("be1_echo") whose target carries a non-empty OriginalCapabilityName. The core +// forwards the advertised name to the client, which owns the translation to the +// backend's capability name ("echo"). +// +// The second leg pins the narrowing this test used to contradict: RouteTool also +// accepts the "{workloadID}.{toolName}" alias ("be1.echo", +// router/session_router.go:105), but an alias is not an ADVERTISED name, so a +// direct CallTool on it is ErrNotFound. The alias exists for composite workflow +// step definitions and still resolves there, inside the composer — the step path +// never enters CallTool. Note that TestCallTool_CompositeWorkflow does NOT cover +// the alias fallback: it registers "be1.echo" as an exact routing-table key, so +// RouteTool's fast path answers and the fallback never runs. The third leg below +// is what actually drives it. func TestCallTool_ResolvesRenamedTool(t *testing.T) { t.Parallel() - cfg, m := baseConfig(t) target := &vmcp.BackendTarget{WorkloadID: "be1", OriginalCapabilityName: "echo", BaseURL: "http://be1:8080"} - expectAggregation(m, &aggregator.AggregatedCapabilities{ - Tools: []vmcp.Tool{{Name: "be1_echo", BackendID: "be1"}}, - RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{"be1_echo": target}}, + caps := func() *aggregator.AggregatedCapabilities { + return &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{{Name: "be1_echo", BackendID: "be1"}}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{"be1_echo": target}}, + } + } + + t.Run("advertised resolved name routes to the backend", func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + expectAggregation(m, caps()) + + m.client.EXPECT(). + CallTool(gomock.Any(), target, "be1_echo", gomock.Any(), gomock.Any(), gomock.Any()). + Return(&vmcp.ToolCallResult{}, nil) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.CallTool(context.Background(), nil, "be1_echo", nil, nil) + require.NoError(t, err) }) - m.client.EXPECT(). - CallTool(gomock.Any(), target, "be1.echo", gomock.Any(), gomock.Any()). - Return(&vmcp.ToolCallResult{}, nil) + t.Run("dot alias of an advertised tool is not directly callable", func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + expectAggregation(m, caps()) - c, err := New(cfg) - require.NoError(t, err) - t.Cleanup(func() { _ = c.Close() }) + // No client EXPECT: the alias must be rejected before any backend call. + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) - _, err = c.CallTool(context.Background(), nil, "be1.echo", nil, nil) - require.NoError(t, err) + _, err = c.CallTool(context.Background(), nil, "be1.echo", nil, nil) + assert.ErrorIs(t, err, vmcp.ErrNotFound, + "the routing-table alias is not an advertised name, so a direct call must not resolve it") + }) + + // The guarantee the alias actually exists for, and the only test in the suite + // that drives RouteTool's dot-convention FALLBACK (session_router.go:110-118) + // rather than its exact-key fast path: a workflow step written "be1.echo" + // still resolves after conflict resolution rekeyed the table to "be1_echo". + // Without this leg, rejecting the alias for direct calls would leave the + // fallback with no coverage at all. + t.Run("composite step reaches the backend through the dot-alias fallback", func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + cfg.WorkflowDefs = map[string]*composer.WorkflowDefinition{ + "wf": {Name: "wf", Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "be1.echo"}, + }}, + } + expectAggregation(m, caps()) + + // The step's "be1.echo" is NOT a routing-table key (the table holds only + // "be1_echo"), so reaching the backend here proves the fallback resolved it. + m.client.EXPECT(). + CallTool(gomock.Any(), target, "be1.echo", gomock.Any(), gomock.Any(), gomock.Any()). + Return(&vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}}, nil) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.CallTool(context.Background(), nil, "wf", nil, nil) + require.NoError(t, err) + }) } // TestCompositeNameConflict_AdvertisedEqualsExecuted is the F1 parity guard: when a @@ -304,13 +457,284 @@ func TestCompositeNameConflict_AdvertisedEqualsExecuted(t *testing.T) { // CallTool("shared") must route to the backend, not execute the composite. want := &vmcp.ToolCallResult{StructuredContent: map[string]any{"from": "backend"}} - m.client.EXPECT().CallTool(gomock.Any(), beTarget, "shared", gomock.Any(), gomock.Any()).Return(want, nil) + m.client.EXPECT().CallTool(gomock.Any(), beTarget, "shared", gomock.Any(), gomock.Any(), gomock.Any()).Return(want, nil) got, err := c.CallTool(context.Background(), nil, "shared", nil, nil) require.NoError(t, err) assert.Equal(t, want, got, "conflicting name must resolve to the backend tool, not the composite") } +// TestCompositeAnnotationContradiction_AdvertisedEqualsExecuted is the parity +// guard for annotation drop: a composite whose explicit annotations contradict +// the derived floor is omitted from ListTools and must also be uncallable via +// CallTool (ErrNotFound). No backend CallTool expectation is set — if the +// composite incorrectly remains executable, CallTool would hit the composer and +// invoke the backend, failing the mock. +func TestCompositeAnnotationContradiction_AdvertisedEqualsExecuted(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + cfg.WorkflowDefs = map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + // Optimistic readOnlyHint against a silent backend (nil annotations → + // fail-closed floor with readOnly=false) is a contradiction → drop. + Annotations: &config.ToolAnnotationsOverride{ReadOnlyHint: boolPtr(true)}, + Steps: []composer.WorkflowStep{{ID: "s1", Type: composer.StepTypeTool, Tool: "be1.echo"}}, + }, + } + + beTarget := backendTarget() + // Backend tool has nil Annotations → silent backend → conservative floor. + agg := &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{{Name: "be1.echo", BackendID: testBackendID}}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{ + "be1.echo": beTarget, + }}, + } + // Two aggregations: one for ListTools, one for CallTool. + m.reg.EXPECT().List(gomock.Any()).Return(nil).Times(2) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(agg, nil).Times(2) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + tools, err := c.ListTools(context.Background(), nil) + require.NoError(t, err) + require.Len(t, tools, 1) + assert.Equal(t, "be1.echo", tools[0].Name) + assert.Equal(t, testBackendID, tools[0].BackendID) + + // No m.client.EXPECT().CallTool — a regression that executes the workflow + // would call the backend and fail the mock controller. + _, err = c.CallTool(context.Background(), nil, "wf", nil, nil) + assert.ErrorIs(t, err, vmcp.ErrNotFound, + "contradicting composite must be uncallable; advertised equals executed") +} + +// TestCompositeNameConflict_WithOptimisticAnnotations_RoutesToBackend is the +// bug-2 regression: a composite named like a backend tool and carrying +// optimistic annotations (readOnlyHint:true) must still be detected as a name +// conflict. Pre-fix, ConvertWorkflowDefsToTools + noop resolver dropped the +// composite from conversion (annotation contradiction), so ValidateNoToolConflicts +// never saw the name and CallTool executed the workflow. Post-fix, +// CompositeToolNames feeds the conflict check → ALL composites dropped → +// CallTool routes to the backend. +func TestCompositeNameConflict_WithOptimisticAnnotations_RoutesToBackend(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + cfg.WorkflowDefs = map[string]*composer.WorkflowDefinition{ + "be1.echo": { + Name: "be1.echo", + Annotations: &config.ToolAnnotationsOverride{ReadOnlyHint: boolPtr(true)}, + Steps: []composer.WorkflowStep{{ID: "s1", Type: composer.StepTypeTool, Tool: "be1.echo"}}, + }, + } + + beTarget := backendTarget() + agg := &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{{Name: "be1.echo", BackendID: testBackendID}}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{ + "be1.echo": beTarget, + }}, + } + // Two aggregations: one for ListTools, one for CallTool. + m.reg.EXPECT().List(gomock.Any()).Return(nil).Times(2) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(agg, nil).Times(2) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + tools, err := c.ListTools(context.Background(), nil) + require.NoError(t, err) + require.Len(t, tools, 1) + assert.Equal(t, "be1.echo", tools[0].Name) + assert.Equal(t, testBackendID, tools[0].BackendID) + + want := &vmcp.ToolCallResult{StructuredContent: map[string]any{"from": "backend"}} + m.client.EXPECT().CallTool(gomock.Any(), beTarget, "be1.echo", gomock.Any(), gomock.Any(), gomock.Any()).Return(want, nil) + + got, err := c.CallTool(context.Background(), nil, "be1.echo", nil, nil) + require.NoError(t, err) + assert.Equal(t, want, got, "conflicting composite with optimistic annotations must route to the backend") +} + +// TestComplete_RoutesPromptRef verifies a ref/prompt completion resolves the +// backend through the prompts routing table and forwards to the backend client. +func TestComplete_RoutesPromptRef(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + target := backendTarget() + expectAggregation(m, &aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{Prompts: map[string]*vmcp.BackendTarget{"p1": target}}, + }) + + want := &vmcp.CompletionResult{Values: []string{"alpha", "beta"}, Total: 2} + ref := vmcp.CompletionRef{Type: vmcp.CompletionRefTypePrompt, Name: "p1"} + m.client.EXPECT(). + Complete(gomock.Any(), target, ref, "arg", "a", gomock.Any()). + Return(want, nil) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + got, err := c.Complete(t.Context(), nil, ref, "arg", "a", nil) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +// TestComplete_RoutesResourceTemplateRef verifies a ref/resource completion resolves +// the backend through the resource-templates routing table and forwards to the +// backend client. Per the MCP spec a ref/resource carries the URI TEMPLATE STRING +// itself (ResourceTemplateReference's uri is "@format uri-template"), which a +// template does not match by expansion — the router must route the exact template +// string to its backend. An expanded URI still routes via template matching. +func TestComplete_RoutesResourceTemplateRef(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + uri string + }{ + { + name: "template string (spec form)", + uri: "file:///logs/{date}.txt", + }, + { + name: "expanded URI (template expansion match)", + uri: "file:///logs/2025-01-01.txt", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + target := backendTarget() + expectAggregation(m, &aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{ + ResourceTemplates: map[string]*vmcp.BackendTarget{"file:///logs/{date}.txt": target}, + }, + }) + + want := &vmcp.CompletionResult{Values: []string{"2025-01-01.txt"}} + ref := vmcp.CompletionRef{Type: vmcp.CompletionRefTypeResource, URI: tc.uri} + m.client.EXPECT(). + Complete(gomock.Any(), target, ref, "date", "2025", gomock.Any()). + Return(want, nil) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + got, err := c.Complete(t.Context(), nil, ref, "date", "2025", nil) + require.NoError(t, err) + assert.Equal(t, want, got) + }) + } +} + +// TestComplete_UnroutableRefReturnsEmpty covers the lenient-completion contract: +// an unroutable prompt ref, an unroutable resource ref, and an unknown ref type all +// yield a non-nil empty result rather than an error, and never reach the backend. +func TestComplete_UnroutableRefReturnsEmpty(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + ref vmcp.CompletionRef + }{ + { + name: "unknown prompt", + ref: vmcp.CompletionRef{Type: vmcp.CompletionRefTypePrompt, Name: "missing"}, + }, + { + name: "unmatched resource", + ref: vmcp.CompletionRef{Type: vmcp.CompletionRefTypeResource, URI: "file:///nope"}, + }, + { + name: "unknown ref type", + ref: vmcp.CompletionRef{Type: "ref/bogus", Name: "x"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + // Empty routing table: nothing routes. No client.Complete expectation, so + // a forwarded call would fail the test. + expectAggregation(m, &aggregator.AggregatedCapabilities{RoutingTable: &vmcp.RoutingTable{}}) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + got, err := c.Complete(t.Context(), nil, tc.ref, "arg", "", nil) + require.NoError(t, err) + require.NotNil(t, got) + assert.Empty(t, got.Values) + }) + } +} + +// TestComplete_AdmissionDenied verifies a completion ref whose referenced capability +// is denied by admission returns ErrAuthorizationFailed without reaching the backend, +// for both a prompt ref (get-side decision) and a resource ref (read-side decision). +func TestComplete_AdmissionDenied(t *testing.T) { + t.Parallel() + + t.Run("prompt ref denied", func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + cfg.ServerName = "test-vmcp" + cfg.Authz = cedarAuthzConfig(t, + `permit(principal, action == Action::"get_prompt", resource == Prompt::"allowed");`) + + m.reg.EXPECT().List(gomock.Any()). + Return([]vmcp.Backend{{ID: testBackendID, HealthStatus: vmcp.BackendHealthy}}).AnyTimes() + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(&aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{Prompts: map[string]*vmcp.BackendTarget{"denied": backendTarget()}}, + }, nil).AnyTimes() + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + ref := vmcp.CompletionRef{Type: vmcp.CompletionRefTypePrompt, Name: "denied"} + _, err = c.Complete(t.Context(), cedarIdentity(), ref, "arg", "", nil) + assert.ErrorIs(t, err, vmcp.ErrAuthorizationFailed) + }) + + t.Run("resource ref denied", func(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + cfg.ServerName = "test-vmcp" + cfg.Authz = cedarAuthzConfig(t, + `permit(principal, action == Action::"read_resource", resource == Resource::"file:///ok");`) + + m.reg.EXPECT().List(gomock.Any()). + Return([]vmcp.Backend{{ID: testBackendID, HealthStatus: vmcp.BackendHealthy}}).AnyTimes() + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(&aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{ + ResourceTemplates: map[string]*vmcp.BackendTarget{"file:///{name}": backendTarget()}, + }, + }, nil).AnyTimes() + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + ref := vmcp.CompletionRef{Type: vmcp.CompletionRefTypeResource, URI: "file:///secret"} + _, err = c.Complete(t.Context(), cedarIdentity(), ref, "name", "", nil) + assert.ErrorIs(t, err, vmcp.ErrAuthorizationFailed) + }) +} + // stubComposer is a configurable composer.Composer for unit-testing // executeComposite's result/error conversion without the real workflow engine. type stubComposer struct { diff --git a/pkg/vmcp/core/core_checks.go b/pkg/vmcp/core/core_checks.go new file mode 100644 index 0000000000..134b6400b8 --- /dev/null +++ b/pkg/vmcp/core/core_checks.go @@ -0,0 +1,115 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package core + +import ( + "context" + "fmt" + "maps" + + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// CheckToolCall runs the CallTool admission decision without invoking the tool. +// It fetches its own aggregated view (the tool decision sources annotations from +// the advertised set, and composites must be included) and delegates to the shared +// authorizeToolCall helper, so a pre-flight gate and the call path enforce the same +// decision from one source. +// +// The gated allowed-call flow thus aggregates twice (once here, once in CallTool); +// this is intentional and cheap in practice — the Serve-layer CachingAggregator +// (30s TTL) turns the second view into a cache hit (vmcp anti-patterns #8/#9: the +// core does not cache, Serve does). identity is never logged. +func (c *coreVMCP) CheckToolCall(ctx context.Context, identity *auth.Identity, name string, args map[string]any) error { + // Clone args before handing them to the admission seam, matching CallTool's + // copy-before-mutating-caller-input discipline (the caller's map — here the + // parsed request's Arguments, shared via context — must not be mutated). + args = maps.Clone(args) + agg, err := c.aggregatedView(ctx) + if err != nil { + // An aggregation failure is a transport/plumbing error, NOT an authorization + // denial — return it unwrapped so the gate admits and the call path surfaces + // it through existing mapping, rather than converting infra faults into 403s. + return err + } + return c.authorizeToolCall(ctx, identity, name, args, findAdvertisedTool(c.advertisedTools(agg), name)) +} + +// CheckResourceRead runs the ReadResource admission decision without reading the +// resource. It needs no aggregated view (the URI alone identifies the decision), +// so it works even when the aggregator is failing. identity is never logged. +func (c *coreVMCP) CheckResourceRead(ctx context.Context, identity *auth.Identity, uri string) error { + return c.authorizeResourceRead(ctx, identity, uri) +} + +// CheckPromptGet runs the GetPrompt admission decision without retrieving the +// prompt. It needs no aggregated view (the name alone identifies the decision), +// so it works even when the aggregator is failing. identity is never logged. +func (c *coreVMCP) CheckPromptGet(ctx context.Context, identity *auth.Identity, name string) error { + return c.authorizePromptGet(ctx, identity, name) +} + +// authorizeToolCall runs the CALL-side admission decision for name with args, +// sourcing the tool's annotations from the caller-supplied advertised entry so +// annotation-gated policies evaluate. It returns nil when allowed and an error +// wrapping vmcp.ErrAuthorizationFailed on deny AND on an authorizer error (fail +// closed), classified so the Serve adapter can distinguish it from a transport +// failure via errors.Is — mirroring the live authorizeAndServe. The underlying +// error is preserved in the chain for server-side diagnostics. +// +// This is the single source of truth for the tool-call decision: both CallTool +// (before dispatch) and CheckToolCall (pre-flight) resolve the tool the same way +// (findAdvertisedTool over the advertised view) and hand it here, so the two +// cannot drift. args is treated as read-only. +// +// tool is nil when name is absent from the advertised set; the decision then runs +// against a bare stub carrying only the name, so an annotation-gated policy +// evaluates with no hints (and may deny). CallTool rejects an absent name outright +// but only AFTER this call: the denial must win over not-found (see CallTool). +func (c *coreVMCP) authorizeToolCall( + ctx context.Context, + identity *auth.Identity, + name string, + args map[string]any, + tool *vmcp.Tool, +) error { + if tool == nil { + tool = &vmcp.Tool{Name: name} + } + if allowed, err := c.admission.AllowToolCall(ctx, identity, tool, args); err != nil { + return fmt.Errorf("%w: tool %q: %w", vmcp.ErrAuthorizationFailed, name, err) + } else if !allowed { + return fmt.Errorf("%w: tool %q", vmcp.ErrAuthorizationFailed, name) + } + return nil +} + +// authorizeResourceRead runs the read-side admission decision for uri, mirroring +// ListResources' filter for the single read. Resources carry no annotations, so +// the URI alone identifies the decision — it needs no aggregated view. An +// authorizer error fails closed, classified as ErrAuthorizationFailed (see +// authorizeToolCall). Shared by ReadResource and CheckResourceRead. +func (c *coreVMCP) authorizeResourceRead(ctx context.Context, identity *auth.Identity, uri string) error { + if allowed, err := c.admission.AllowResourceRead(ctx, identity, &vmcp.Resource{URI: uri}); err != nil { + return fmt.Errorf("%w: resource %q: %w", vmcp.ErrAuthorizationFailed, uri, err) + } else if !allowed { + return fmt.Errorf("%w: resource %q", vmcp.ErrAuthorizationFailed, uri) + } + return nil +} + +// authorizePromptGet runs the get-side admission decision for name, mirroring +// ListPrompts' filter for the single get. Prompts carry no annotations, so the +// name alone identifies the decision — it needs no aggregated view. An authorizer +// error fails closed, classified as ErrAuthorizationFailed (see authorizeToolCall). +// Shared by GetPrompt and CheckPromptGet. +func (c *coreVMCP) authorizePromptGet(ctx context.Context, identity *auth.Identity, name string) error { + if allowed, err := c.admission.AllowPromptGet(ctx, identity, &vmcp.Prompt{Name: name}); err != nil { + return fmt.Errorf("%w: prompt %q: %w", vmcp.ErrAuthorizationFailed, name, err) + } else if !allowed { + return fmt.Errorf("%w: prompt %q", vmcp.ErrAuthorizationFailed, name) + } + return nil +} diff --git a/pkg/vmcp/core/core_checks_test.go b/pkg/vmcp/core/core_checks_test.go new file mode 100644 index 0000000000..9adb7c20b0 --- /dev/null +++ b/pkg/vmcp/core/core_checks_test.go @@ -0,0 +1,287 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package core + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" +) + +// checkCore builds a *coreVMCP wired with the given admission and the base mocks, +// bypassing New so tests can inject a specific admission seam directly. The tool +// checks call aggregatedView; resource/prompt checks do not. +func checkCore(m *coreMocks, adm Admission) *coreVMCP { + return &coreVMCP{ + aggregator: m.agg, + backendRegistry: m.reg, + backendClient: m.client, + admission: adm, + } +} + +// expectAggregationAnyTimes wires reg.List + agg.AggregateCapabilities to return +// agg on every call (the check/call pairing aggregates more than once). +func expectAggregationAnyTimes(m *coreMocks, agg *aggregator.AggregatedCapabilities) { + m.reg.EXPECT().List(gomock.Any()). + Return([]vmcp.Backend{{ID: testBackendID, HealthStatus: vmcp.BackendHealthy}}).AnyTimes() + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(agg, nil).AnyTimes() +} + +// isAuthErr reports whether err is an authorization denial (nil is not). +func isAuthErr(err error) bool { return errors.Is(err, vmcp.ErrAuthorizationFailed) } + +func TestCheckToolCall_AllowDenyFailClosed(t *testing.T) { + t.Parallel() + _, m := baseConfig(t) + + authorizer := &mockAuthorizer{results: map[string]mockResult{ + "echo": {authorized: true}, + "boom": {err: errors.New("authorizer exploded")}, + "denied": {authorized: false}, + }} + c := checkCore(m, newCedarAdmission(authorizer)) + expectAggregationAnyTimes(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{backendTool("echo"), backendTool("boom"), backendTool("denied")}, + RoutingTable: &vmcp.RoutingTable{}, + }) + + require.NoError(t, c.CheckToolCall(t.Context(), cedarIdentity(), "echo", nil), + "an allowed tool call must pass the pre-flight check") + + err := c.CheckToolCall(t.Context(), cedarIdentity(), "denied", nil) + assert.ErrorIs(t, err, vmcp.ErrAuthorizationFailed, "a denied tool call must fail the check") + + err = c.CheckToolCall(t.Context(), cedarIdentity(), "boom", nil) + assert.ErrorIs(t, err, vmcp.ErrAuthorizationFailed, "an authorizer error must fail closed as an authz denial") +} + +// TestCheckToolCall_AggregationError_NotAuthzDenial pins the infra-error branch: +// when aggregatedView fails (a transport/plumbing fault, e.g. the backend registry +// or aggregator is unavailable), CheckToolCall must return that error UNWRAPPED — +// not classified as ErrAuthorizationFailed. The gate admits on a non-authz error, +// so wrapping a transient aggregation failure as an authz denial would silently +// turn it into a hard HTTP 403 (fail-open infra handling must not become fail-closed +// authz). This is the classification the fake-core gate test cannot exercise. +func TestCheckToolCall_AggregationError_NotAuthzDenial(t *testing.T) { + t.Parallel() + _, m := baseConfig(t) + + // Admission is never reached; the aggregation failure short-circuits first. + c := checkCore(m, newCedarAdmission(&mockAuthorizer{results: map[string]mockResult{}})) + + infraErr := errors.New("aggregator unavailable") + m.reg.EXPECT().List(gomock.Any()). + Return([]vmcp.Backend{{ID: testBackendID, HealthStatus: vmcp.BackendHealthy}}).AnyTimes() + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(nil, infraErr).AnyTimes() + + err := c.CheckToolCall(t.Context(), cedarIdentity(), "echo", nil) + require.Error(t, err, "an aggregation failure must surface as an error") + assert.False(t, isAuthErr(err), + "a transient aggregation failure must NOT be classified as ErrAuthorizationFailed "+ + "(the gate would otherwise convert an infra fault into a 403)") +} + +func TestCheckResourceRead_AllowDenyFailClosed_NoAggregation(t *testing.T) { + t.Parallel() + // No aggregation expectations are set: if CheckResourceRead ever aggregated, the + // gomock controller would fail on the unexpected reg.List/AggregateCapabilities + // call. This proves the read check has no aggregation dependency. + _, m := baseConfig(t) + + authorizer := &mockAuthorizer{results: map[string]mockResult{ + "res://ok": {authorized: true}, + "res://boom": {err: errors.New("authorizer exploded")}, + }} + c := checkCore(m, newCedarAdmission(authorizer)) + + require.NoError(t, c.CheckResourceRead(t.Context(), cedarIdentity(), "res://ok")) + assert.ErrorIs(t, c.CheckResourceRead(t.Context(), cedarIdentity(), "res://denied"), + vmcp.ErrAuthorizationFailed) + assert.ErrorIs(t, c.CheckResourceRead(t.Context(), cedarIdentity(), "res://boom"), + vmcp.ErrAuthorizationFailed, "authorizer error must fail closed") +} + +// TestCallTool_DenialPrecedesNotFound pins the ORDER of CallTool's two rejections. +// A tool that is both hidden from tools/list and denied by policy must report the +// DENIAL, not ErrNotFound: if the advertised-view check ran first, the pair of +// answers would become an oracle — ErrNotFound for "hidden and denied" vs. +// ErrAuthorizationFailed for "advertised but denied" would tell an unauthorized +// caller which hidden tools exist. +// +// "hidden" is routable but unadvertised, and the authorizer denies it. The +// advertised "visible" tool proves the same authorizer produces a denial (not a +// not-found) for a name that IS in the view, so the two legs differ only in +// advertising. +func TestCallTool_DenialPrecedesNotFound(t *testing.T) { + t.Parallel() + _, m := baseConfig(t) + + // Neither name is authorized: mockAuthorizer denies anything absent from results. + c := checkCore(m, newCedarAdmission(&mockAuthorizer{results: map[string]mockResult{}})) + target := backendTarget() + expectAggregationAnyTimes(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{backendTool("visible")}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{ + "visible": target, + "hidden": target, + }}, + }) + + _, hiddenErr := c.CallTool(t.Context(), cedarIdentity(), "hidden", nil, nil) + assert.ErrorIs(t, hiddenErr, vmcp.ErrAuthorizationFailed, + "a denied tool must report the denial even when it is also unadvertised") + assert.NotErrorIs(t, hiddenErr, vmcp.ErrNotFound, + "reporting not-found for a denied hidden tool would leak which hidden tools exist") + + _, visibleErr := c.CallTool(t.Context(), cedarIdentity(), "visible", nil, nil) + assert.ErrorIs(t, visibleErr, vmcp.ErrAuthorizationFailed, + "an advertised-but-denied tool must be indistinguishable from a hidden denied one") +} + +func TestCheckPromptGet_AllowDenyFailClosed_NoAggregation(t *testing.T) { + t.Parallel() + _, m := baseConfig(t) + + authorizer := &mockAuthorizer{results: map[string]mockResult{ + "greet": {authorized: true}, + "boom": {err: errors.New("authorizer exploded")}, + }} + c := checkCore(m, newCedarAdmission(authorizer)) + + require.NoError(t, c.CheckPromptGet(t.Context(), cedarIdentity(), "greet")) + assert.ErrorIs(t, c.CheckPromptGet(t.Context(), cedarIdentity(), "secret"), + vmcp.ErrAuthorizationFailed) + assert.ErrorIs(t, c.CheckPromptGet(t.Context(), cedarIdentity(), "boom"), + vmcp.ErrAuthorizationFailed, "authorizer error must fail closed") +} + +// TestCheckToolCall_AntiDrift is the invariant guard: for the same (policy, identity, +// name, args), CheckToolCall must fail with ErrAuthorizationFailed if and only if +// CallTool does. The shared authorizeToolCall helper makes this structural, but the +// test pins it — including an argument-conditional Cedar policy so the args flow +// through both routes is proven. An allowed call falls through to routing, which +// returns ErrNotFound (empty routing table) — NOT an authz error — so the invariant +// compares only the authorization outcome, not the transport outcome. +func TestCheckToolCall_AntiDrift(t *testing.T) { + t.Parallel() + + const permitEcho = `permit(principal, action == Action::"call_tool", resource == Tool::"echo");` + const forbidBlocked = `forbid(principal, action == Action::"call_tool", resource == Tool::"echo") ` + + `when { context has arg_input && context.arg_input == "blocked" };` + + tests := []struct { + name string + policies []string + toolName string + args map[string]any + // wantAuthzDenied is the ABSOLUTE expected outcome. Asserting it (not just that + // both paths agree) is what proves the args/name are actually consulted: an + // agreement-only check would still pass if BOTH paths dropped args entirely. + wantAuthzDenied bool + }{ + {"unconditional permit", []string{permitEcho}, "echo", nil, false}, + {"default deny unpermitted name", []string{permitEcho}, "other", nil, true}, + {"arg-gated allow", []string{permitEcho, forbidBlocked}, "echo", map[string]any{"input": "hello"}, false}, + {"arg-gated deny", []string{permitEcho, forbidBlocked}, "echo", map[string]any{"input": "blocked"}, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, m := baseConfig(t) + c := checkCore(m, cedarAdmissionWith(t, tt.policies...)) + expectAggregationAnyTimes(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{backendTool("echo")}, + RoutingTable: &vmcp.RoutingTable{}, + }) + + checkErr := c.CheckToolCall(t.Context(), cedarIdentity(), tt.toolName, tt.args) + _, callErr := c.CallTool(t.Context(), cedarIdentity(), tt.toolName, tt.args, nil) + + // Absolute outcome: proves args are consulted (arg-gated deny/allow differ + // only by the "input" value), not merely that the two paths agree. + assert.Equal(t, tt.wantAuthzDenied, isAuthErr(checkErr), + "CheckToolCall authz outcome mismatch (err=%v)", checkErr) + // Agreement: the shared helper keeps the two routes from drifting. + assert.Equal(t, isAuthErr(checkErr), isAuthErr(callErr), + "CheckToolCall and CallTool must agree on the authorization outcome (check=%v call=%v)", + checkErr, callErr) + }) + } +} + +// TestCheckResourceRead_AntiDrift pairs CheckResourceRead with ReadResource: the +// check errors with ErrAuthorizationFailed iff the read does. ReadResource +// aggregates before authorizing (allowed reads fall through to routing → +// ErrNotFound); the check does not aggregate, so aggregation is wired AnyTimes. +func TestCheckResourceRead_AntiDrift(t *testing.T) { + t.Parallel() + + cases := []struct { + uri string + wantAuthzDenied bool + }{ + {"res://ok", false}, + {"res://denied", true}, + } + for _, tc := range cases { + t.Run(tc.uri, func(t *testing.T) { + t.Parallel() + // A fresh authorizer per subtest: mockAuthorizer records calls into a shared + // slice, so a single instance shared across parallel subtests would race. + authorizer := &mockAuthorizer{results: map[string]mockResult{"res://ok": {authorized: true}}} + _, m := baseConfig(t) + c := checkCore(m, newCedarAdmission(authorizer)) + expectAggregationAnyTimes(m, &aggregator.AggregatedCapabilities{RoutingTable: &vmcp.RoutingTable{}}) + + checkErr := c.CheckResourceRead(t.Context(), cedarIdentity(), tc.uri) + _, readErr := c.ReadResource(t.Context(), cedarIdentity(), tc.uri) + + // Absolute outcome (proves the URI is consulted) plus agreement (no drift). + assert.Equal(t, tc.wantAuthzDenied, isAuthErr(checkErr), + "CheckResourceRead authz outcome mismatch (err=%v)", checkErr) + assert.Equal(t, isAuthErr(checkErr), isAuthErr(readErr), + "CheckResourceRead and ReadResource must agree (check=%v read=%v)", checkErr, readErr) + }) + } +} + +// TestCheckPromptGet_AntiDrift pairs CheckPromptGet with GetPrompt (same invariant). +func TestCheckPromptGet_AntiDrift(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + wantAuthzDenied bool + }{ + {"greet", false}, + {"secret", true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + // Fresh authorizer per subtest (shared mockAuthorizer records calls → would race). + authorizer := &mockAuthorizer{results: map[string]mockResult{"greet": {authorized: true}}} + _, m := baseConfig(t) + c := checkCore(m, newCedarAdmission(authorizer)) + expectAggregationAnyTimes(m, &aggregator.AggregatedCapabilities{RoutingTable: &vmcp.RoutingTable{}}) + + checkErr := c.CheckPromptGet(t.Context(), cedarIdentity(), tc.name) + _, getErr := c.GetPrompt(t.Context(), cedarIdentity(), tc.name, nil) + + // Absolute outcome (proves the name is consulted) plus agreement (no drift). + assert.Equal(t, tc.wantAuthzDenied, isAuthErr(checkErr), + "CheckPromptGet authz outcome mismatch (err=%v)", checkErr) + assert.Equal(t, isAuthErr(checkErr), isAuthErr(getErr), + "CheckPromptGet and GetPrompt must agree (check=%v get=%v)", checkErr, getErr) + }) + } +} diff --git a/pkg/vmcp/core/core_vmcp.go b/pkg/vmcp/core/core_vmcp.go index ac72694fd5..e5730a4b21 100644 --- a/pkg/vmcp/core/core_vmcp.go +++ b/pkg/vmcp/core/core_vmcp.go @@ -10,6 +10,8 @@ import ( "sync" "time" + "github.com/yosida95/uritemplate/v3" + "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" @@ -289,6 +291,54 @@ func (c *coreVMCP) ListResources(ctx context.Context, identity *auth.Identity) ( return c.admission.FilterResources(ctx, identity, agg.Resources) } +// ListResourceTemplates health-filters the registry, aggregates resource +// templates on demand, and returns only those identity is admitted to read. +// +// Resource templates carry no admission entity of their own: the URI template is +// treated as the resource id, reusing the exact resource-read filter (the same +// per-item skip-on-error, log-and-continue decision ListResources applies). A +// template whose URI-template string is denied is withheld. +func (c *coreVMCP) ListResourceTemplates( + ctx context.Context, identity *auth.Identity, +) ([]vmcp.ResourceTemplate, error) { + agg, err := c.aggregatedView(ctx) + if err != nil { + return nil, err + } + return c.filterResourceTemplates(ctx, identity, agg.ResourceTemplates) +} + +// filterResourceTemplates applies the resource-read admission filter to templates +// by projecting each template's URI template onto a vmcp.Resource and reusing +// admission.FilterResources, so templates and resources share one decision path. +// The surviving URI templates select which templates are returned. +func (c *coreVMCP) filterResourceTemplates( + ctx context.Context, identity *auth.Identity, templates []vmcp.ResourceTemplate, +) ([]vmcp.ResourceTemplate, error) { + if len(templates) == 0 { + return []vmcp.ResourceTemplate{}, nil + } + proxies := make([]vmcp.Resource, len(templates)) + for i := range templates { + proxies[i] = vmcp.Resource{URI: templates[i].URITemplate, BackendID: templates[i].BackendID} + } + allowedProxies, err := c.admission.FilterResources(ctx, identity, proxies) + if err != nil { + return nil, err + } + allowed := make(map[string]struct{}, len(allowedProxies)) + for i := range allowedProxies { + allowed[allowedProxies[i].URI] = struct{}{} + } + filtered := make([]vmcp.ResourceTemplate, 0, len(templates)) + for i := range templates { + if _, ok := allowed[templates[i].URITemplate]; ok { + filtered = append(filtered, templates[i]) + } + } + return filtered, nil +} + // ListPrompts health-filters the registry, aggregates prompts on demand, and // returns only those identity is admitted to get. func (c *coreVMCP) ListPrompts(ctx context.Context, identity *auth.Identity) ([]vmcp.Prompt, error) { @@ -299,6 +349,45 @@ func (c *coreVMCP) ListPrompts(ctx context.Context, identity *auth.Identity) ([] return c.admission.FilterPrompts(ctx, identity, agg.Prompts) } +// Discover aggregates backend capabilities ONCE and derives all four +// capability-presence flags from that single view, applying the exact same +// admission-filter code paths ListTools/ListResources/ListResourceTemplates/ +// ListPrompts each apply independently (advertisedTools+FilterTools, +// FilterResources, filterResourceTemplates, FilterPrompts). Sharing those +// helpers against one aggregatedView call -- rather than reimplementing the +// filtering -- is what guarantees a flag can never drift from what the +// corresponding List* verb would show for the same identity. +func (c *coreVMCP) Discover(ctx context.Context, identity *auth.Identity) (DiscoverCapabilities, error) { + agg, err := c.aggregatedView(ctx) + if err != nil { + return DiscoverCapabilities{}, err + } + + tools, err := c.admission.FilterTools(ctx, identity, c.advertisedTools(agg)) + if err != nil { + return DiscoverCapabilities{}, err + } + resources, err := c.admission.FilterResources(ctx, identity, agg.Resources) + if err != nil { + return DiscoverCapabilities{}, err + } + templates, err := c.filterResourceTemplates(ctx, identity, agg.ResourceTemplates) + if err != nil { + return DiscoverCapabilities{}, err + } + prompts, err := c.admission.FilterPrompts(ctx, identity, agg.Prompts) + if err != nil { + return DiscoverCapabilities{}, err + } + + return DiscoverCapabilities{ + HasTools: len(tools) > 0, + HasResources: len(resources) > 0, + HasResourceTemplates: len(templates) > 0, + HasPrompts: len(prompts) > 0, + }, nil +} + // LookupTool resolves an advertised tool name (incl. composite tools) to its // capability without invoking it. It delegates to ListTools, so it applies the // same health/advertising AND admission view: a name that is unknown, unadvertised, @@ -320,6 +409,13 @@ func (c *coreVMCP) LookupTool(ctx context.Context, identity *auth.Identity, name // LookupResource resolves an advertised resource URI to its capability without // reading it. Delegates to ListResources, so a URI that is unknown, unadvertised, // or denied to identity returns vmcp.ErrNotFound. +// +// A URI backed by a resource TEMPLATE resolves too: after the concrete-resource +// miss, the URI is checked against the admission-filtered resource templates +// (the same view ListResourceTemplates advertises) — an exact template-string +// match (the form a resources/subscribe or completion ref/resource carries) or +// a template-expansion match. This keeps subscribe's admission aligned with +// read: any URI ReadResource would serve is a valid subscription target. func (c *coreVMCP) LookupResource(ctx context.Context, identity *auth.Identity, uri string) (*vmcp.Resource, error) { resources, err := c.ListResources(ctx, identity) if err != nil { @@ -331,9 +427,39 @@ func (c *coreVMCP) LookupResource(ctx context.Context, identity *auth.Identity, return &resource, nil } } + + templates, err := c.ListResourceTemplates(ctx, identity) + if err != nil { + return nil, err + } + for i := range templates { + if templates[i].URITemplate == uri || templateMatches(templates[i].URITemplate, uri) { + // Project the template onto a resource (the same projection + // filterResourceTemplates uses) so lookup and list agree on the shape. + return &vmcp.Resource{ + URI: templates[i].URITemplate, + Name: templates[i].Name, + Description: templates[i].Description, + MimeType: templates[i].MimeType, + BackendID: templates[i].BackendID, + }, nil + } + } + return nil, fmt.Errorf("%w: resource %q", vmcp.ErrNotFound, uri) } +// templateMatches reports whether the RFC 6570 URI template matches uri. An +// invalid template string reports no match (it was already logged and skipped +// at routing-table construction). +func templateMatches(tmplStr, uri string) bool { + tmpl, err := uritemplate.New(tmplStr) + if err != nil { + return false + } + return tmpl.Match(uri) != nil +} + // LookupPrompt resolves an advertised prompt name to its capability without // retrieving it. Delegates to ListPrompts, so a name that is unknown, unadvertised, // or denied to identity returns vmcp.ErrNotFound. @@ -395,6 +521,22 @@ func (c *coreVMCP) LookupBackend( return nil, fmt.Errorf("%w: backend %q", vmcp.ErrNotFound, backendID) } +// InvalidateCapabilityCache implements VMCP.InvalidateCapabilityCache. It +// type-asserts the configured aggregator to aggregator.CacheInvalidator and +// delegates; when the aggregator does not implement it, the call is a WARN-logged +// no-op — not a silent one — because a caller relying on invalidation (the +// list_changed resync path) needs to know the request had no effect (go-style: +// no silent no-ops). +func (c *coreVMCP) InvalidateCapabilityCache() { + invalidator, ok := c.aggregator.(aggregator.CacheInvalidator) + if !ok { + slog.Warn("capability cache invalidation requested but the configured aggregator does not support it", + "aggregatorType", fmt.Sprintf("%T", c.aggregator)) + return + } + invalidator.InvalidateAll() +} + // Close stops the workflow state store's cleanup goroutine. It is idempotent: // the underlying Stop closes a channel that cannot be closed twice, so the work // is guarded by sync.Once and subsequent calls return nil. @@ -491,12 +633,22 @@ func (c *coreVMCP) authorizedBackends( // may only SUBTRACT reachability ([VMCP] contract), so the core is the only layer // that can advertise them. func (c *coreVMCP) advertisedTools(agg *aggregator.AggregatedCapabilities) []vmcp.Tool { - defs := c.accessibleComposites(agg) + return advertisedToolsWith(agg, c.accessibleComposites(agg)) +} + +// advertisedToolsWith is advertisedTools against an ALREADY-RESOLVED composite set. +// CallTool needs the same composites map for its own dispatch, and resolving it is +// not free (FilterWorkflowDefsForSession + ValidateNoToolConflicts + +// ConvertWorkflowDefsToTools), so it resolves once and passes the result here +// rather than having this recompute it. +func advertisedToolsWith( + agg *aggregator.AggregatedCapabilities, defs map[string]*composer.WorkflowDefinition, +) []vmcp.Tool { if len(defs) == 0 { return agg.Tools } - composite := compositetools.ConvertWorkflowDefsToTools(defs) + composite := compositetools.ConvertWorkflowDefsToTools(defs, stepAnnotationResolver(agg)) out := make([]vmcp.Tool, 0, len(agg.Tools)+len(composite)) out = append(out, agg.Tools...) out = append(out, composite...) @@ -505,9 +657,12 @@ func (c *coreVMCP) advertisedTools(agg *aggregator.AggregatedCapabilities) []vmc // accessibleComposites returns the composite-tool definitions the core advertises // (and therefore executes) for agg's view: those whose every tool step is reachable -// in the routing table AND whose names do not collide with a backend tool. On a name +// in the routing table, whose names do not collide with a backend tool, AND whose +// explicit annotations do not contradict the derived safety floor. On a name // collision ALL composites are dropped so the backend tool wins, matching the legacy // compositeToolsDecorator (sessionmanager/factory.go:158-168, decorator.go:83-86). +// On an annotation contradiction the offending composite is dropped (with a warning) +// so CallTool and ListTools share the same set. // // This is the single source of truth shared by advertisedTools (what ListTools shows) // and CallTool (what executes), so a withheld composite is never executed — advertised @@ -523,11 +678,39 @@ func (c *coreVMCP) accessibleComposites( if len(defs) == 0 { return nil } - if err := compositetools.ValidateNoToolConflicts(agg.Tools, compositetools.ConvertWorkflowDefsToTools(defs)); err != nil { + // Name-only conflict check — must not run annotation conversion, which can + // drop composites carrying optimistic hints and hide a colliding name. + if err := compositetools.ValidateNoToolConflicts( + agg.Tools, compositetools.CompositeToolNames(defs)); err != nil { slog.Warn("composite tool name conflict detected; omitting composite tools", "error", err) return nil } - return defs + // Annotation guardrail: drop contradicting composites so CallTool and + // ListTools share the same set (advertised equals executed). + return compositetools.FilterWorkflowDefsByAnnotations(defs, stepAnnotationResolver(agg)) +} + +// stepAnnotationResolver returns a StepAnnotationResolver that maps a composite +// tool's step reference ("{workloadID}.{toolName}") to the backend tool's +// annotations, using agg's routing table and aggregated tools. Resolution goes +// through the shared router.ResolveToolRef primitive (the same path +// isToolStepAccessible uses), so accessibility filtering and annotation +// resolution cannot drift. Unknown step tools resolve to nil annotations +// (treated conservatively by the derivation). +func stepAnnotationResolver(agg *aggregator.AggregatedCapabilities) compositetools.StepAnnotationResolver { + annByName := make(map[string]*vmcp.ToolAnnotations, len(agg.Tools)) + for i := range agg.Tools { + if agg.Tools[i].Annotations != nil { + annByName[agg.Tools[i].Name] = agg.Tools[i].Annotations + } + } + return func(stepTool string) *vmcp.ToolAnnotations { + resolvedName, ok := router.ResolveToolRef(agg.RoutingTable, stepTool) + if !ok { + return nil + } + return annByName[resolvedName] + } } // validateConfig checks New's required inputs and the elicitation contract, @@ -653,9 +836,9 @@ func filterHealthyBackends(backends []vmcp.Backend, healthStatusProvider health. // Include healthy, degraded, and empty/zero-value (assume healthy) backends. // Explicitly exclude unhealthy, unknown, and unauthenticated backends. - if healthStatus == "" || - healthStatus == vmcp.BackendHealthy || - healthStatus == vmcp.BackendDegraded { + // The predicate is shared with session establishment's stricter + // counterpart — see health.ShouldAdvertise / health.ShouldOpenSession. + if health.ShouldAdvertise(healthStatus) { healthy = append(healthy, *backend) } else { excluded++ diff --git a/pkg/vmcp/core/core_vmcp_test.go b/pkg/vmcp/core/core_vmcp_test.go index 69f6346d2e..5506635fba 100644 --- a/pkg/vmcp/core/core_vmcp_test.go +++ b/pkg/vmcp/core/core_vmcp_test.go @@ -283,6 +283,145 @@ func TestListResourcesAndPrompts(t *testing.T) { assert.Equal(t, "p1", prompts[0].Name) } +// denyResourceAdmission is a test Admission that drops resources whose URI is in +// deny. It embeds allowAllAdmission so the tool/prompt methods are satisfied, and +// overrides only FilterResources — the sole method ListResourceTemplates exercises +// (the URI template is projected onto a resource URI for admission). +type denyResourceAdmission struct { + allowAllAdmission + deny map[string]struct{} +} + +func (d denyResourceAdmission) FilterResources( + _ context.Context, _ *auth.Identity, resources []vmcp.Resource, +) ([]vmcp.Resource, error) { + out := make([]vmcp.Resource, 0, len(resources)) + for i := range resources { + if _, blocked := d.deny[resources[i].URI]; !blocked { + out = append(out, resources[i]) + } + } + return out, nil +} + +func TestListResourceTemplates(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + m.reg.EXPECT().List(gomock.Any()).Return(nil).Times(2) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(&aggregator.AggregatedCapabilities{ + ResourceTemplates: []vmcp.ResourceTemplate{ + {URITemplate: "file:///logs/{date}.txt", BackendID: "be1"}, + {URITemplate: "secret:///{id}", BackendID: "be1"}, + }, + }, nil).Times(2) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + // Default allow-all admission: both templates surface, tagged with their backend. + templates, err := c.ListResourceTemplates(t.Context(), nil) + require.NoError(t, err) + require.Len(t, templates, 2) + + // Inject a deny-admission that filters templates by URI template (treated as the + // resource id): the denied template is withheld, mirroring ListResources. + c.(*coreVMCP).admission = denyResourceAdmission{deny: map[string]struct{}{"secret:///{id}": {}}} + templates, err = c.ListResourceTemplates(t.Context(), nil) + require.NoError(t, err) + require.Len(t, templates, 1) + assert.Equal(t, "file:///logs/{date}.txt", templates[0].URITemplate) + assert.Equal(t, "be1", templates[0].BackendID) +} + +func TestListResourceTemplates_Empty(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + m.reg.EXPECT().List(gomock.Any()).Return(nil) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return( + &aggregator.AggregatedCapabilities{}, nil) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + templates, err := c.ListResourceTemplates(t.Context(), nil) + require.NoError(t, err) + assert.Empty(t, templates) +} + +// TestDiscover_SingleAggregation verifies Discover derives all four +// capability-presence flags from ONE aggregation call (reg.List + +// AggregateCapabilities each Times(1)) rather than the four independent +// fan-outs ListTools/ListResources/ListResourceTemplates/ListPrompts would +// cost if called separately. +func TestDiscover_SingleAggregation(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + backends := []vmcp.Backend{{ID: testBackendID, HealthStatus: vmcp.BackendHealthy}} + m.reg.EXPECT().List(gomock.Any()).Return(backends).Times(1) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), backends).Return(&aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{backendTool("echo")}, + Resources: []vmcp.Resource{{URI: "file://a", BackendID: testBackendID}}, + ResourceTemplates: []vmcp.ResourceTemplate{{URITemplate: "file:///logs/{date}.txt", BackendID: testBackendID}}, + Prompts: []vmcp.Prompt{{Name: "p1", BackendID: testBackendID}}, + RoutingTable: &vmcp.RoutingTable{}, + }, nil).Times(1) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + caps, err := c.Discover(context.Background(), nil) + require.NoError(t, err) + assert.Equal(t, DiscoverCapabilities{ + HasTools: true, HasResources: true, HasResourceTemplates: true, HasPrompts: true, + }, caps) +} + +// TestDiscover_EmptyAggregate verifies an empty aggregated view yields every +// flag false, not an error. +func TestDiscover_EmptyAggregate(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + m.reg.EXPECT().List(gomock.Any()).Return(nil) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(&aggregator.AggregatedCapabilities{}, nil) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + caps, err := c.Discover(context.Background(), nil) + require.NoError(t, err) + assert.Equal(t, DiscoverCapabilities{}, caps) +} + +// TestDiscover_DeniedIdentityHidesTools is the security-critical case: a +// backend advertising a tool must NOT set HasTools for an identity the +// admission seam denies that tool to. Deriving the flag from the raw +// aggregate instead of the admission-filtered set would leak the tool's +// existence to an identity that cannot call it. +func TestDiscover_DeniedIdentityHidesTools(t *testing.T) { + t.Parallel() + _, m := baseConfig(t) + + authorizer := &mockAuthorizer{results: map[string]mockResult{"echo": {authorized: false}}} + c := checkCore(m, newCedarAdmission(authorizer)) + expectAggregationAnyTimes(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{backendTool("echo")}, + RoutingTable: &vmcp.RoutingTable{}, + }) + + caps, err := c.Discover(t.Context(), cedarIdentity()) + require.NoError(t, err) + assert.False(t, caps.HasTools, + "a denied identity must not see HasTools=true even though the backend advertises a tool") +} + func TestListTools_AggregationError(t *testing.T) { t.Parallel() cfg, m := baseConfig(t) @@ -396,12 +535,15 @@ func TestLookup(t *testing.T) { t.Parallel() cfg, m := baseConfig(t) - m.reg.EXPECT().List(gomock.Any()).Return(nil).Times(6) + // Seven aggregations: each Lookup aggregates once, except the failing + // LookupResource, which also consults the resource-template view on the + // concrete-resource miss. + m.reg.EXPECT().List(gomock.Any()).Return(nil).Times(7) m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(&aggregator.AggregatedCapabilities{ Tools: []vmcp.Tool{backendTool("tool_a")}, Resources: []vmcp.Resource{{URI: "file://a", BackendID: "be1"}}, Prompts: []vmcp.Prompt{{Name: "p1", BackendID: "be1"}}, - }, nil).Times(6) + }, nil).Times(7) c, err := New(cfg) require.NoError(t, err) @@ -427,6 +569,50 @@ func TestLookup(t *testing.T) { assert.ErrorIs(t, err, vmcp.ErrNotFound) } +// TestLookupResource_TemplateBackedURI verifies LookupResource accepts the URIs a +// templated read serves — both the template string itself (the form a +// resources/subscribe or completion ref/resource carries) and an expanded URI +// matching an advertised template — while an unmatched URI still misses. This +// keeps the subscribe admission (coreSubscribeHandler) aligned with read. +func TestLookupResource_TemplateBackedURI(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + m.reg.EXPECT().List(gomock.Any()).Return(nil).Times(6) + m.agg.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(&aggregator.AggregatedCapabilities{ + Resources: []vmcp.Resource{{URI: "file://a", BackendID: "be1"}}, + ResourceTemplates: []vmcp.ResourceTemplate{ + { + URITemplate: "file:///logs/{date}.txt", + Name: "Daily log", + Description: "A day's log file", + MimeType: "text/plain", + BackendID: "be1", + }, + }, + }, nil).Times(6) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + ctx := context.Background() + + // The exact template string resolves to the template's capability. + res, err := c.LookupResource(ctx, nil, "file:///logs/{date}.txt") + require.NoError(t, err) + assert.Equal(t, "file:///logs/{date}.txt", res.URI) + assert.Equal(t, "be1", res.BackendID) + + // An expanded URI matching the template resolves to the same capability. + res, err = c.LookupResource(ctx, nil, "file:///logs/2025-01-01.txt") + require.NoError(t, err) + assert.Equal(t, "be1", res.BackendID) + + // A URI matching no concrete resource or template still misses. + _, err = c.LookupResource(ctx, nil, "db:///users/42") + assert.ErrorIs(t, err, vmcp.ErrNotFound) +} + func TestClose_Idempotent(t *testing.T) { t.Parallel() cfg, _ := baseConfig(t) @@ -441,6 +627,59 @@ func TestClose_Idempotent(t *testing.T) { }) } +// invalidatorAggregator is a minimal aggregator.Aggregator that also +// implements aggregator.CacheInvalidator, so TestInvalidateCapabilityCache_* +// can assert coreVMCP.InvalidateCapabilityCache's delegation without depending +// on the real cachingAggregator (covered separately in the aggregator package). +type invalidatorAggregator struct { + aggregator.Aggregator + invalidateCalls int +} + +func (a *invalidatorAggregator) InvalidateAll() { a.invalidateCalls++ } + +var _ aggregator.CacheInvalidator = (*invalidatorAggregator)(nil) + +// TestInvalidateCapabilityCache_DelegatesWhenSupported verifies (#5748) that +// InvalidateCapabilityCache type-asserts the configured aggregator to +// aggregator.CacheInvalidator and calls InvalidateAll when it is implemented. +func TestInvalidateCapabilityCache_DelegatesWhenSupported(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + inv := &invalidatorAggregator{Aggregator: m.agg} + cfg.Aggregator = inv + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + c.InvalidateCapabilityCache() + assert.Equal(t, 1, inv.invalidateCalls) +} + +// TestInvalidateCapabilityCache_WarnsWhenUnsupported verifies (#5748) that +// InvalidateCapabilityCache is a WARN-logged no-op — not a silent one — when +// the configured aggregator does not implement aggregator.CacheInvalidator +// (baseConfig's plain MockAggregator does not). +// +//nolint:paralleltest // installs a global slog default + non-thread-safe buffer; must not run in parallel +func TestInvalidateCapabilityCache_WarnsWhenUnsupported(t *testing.T) { + cfg, _ := baseConfig(t) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, nil))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + assert.NotPanics(t, func() { c.InvalidateCapabilityCache() }) + assert.Contains(t, buf.String(), "does not support it", + "unsupported aggregator must WARN-log, not silently no-op") +} + // Must run serially: it swaps the global slog default to capture output into a // non-thread-safe buffer, which is unsafe if other tests log concurrently. // diff --git a/pkg/vmcp/core/xmcpheader_derivation_test.go b/pkg/vmcp/core/xmcpheader_derivation_test.go new file mode 100644 index 0000000000..9cb3825246 --- /dev/null +++ b/pkg/vmcp/core/xmcpheader_derivation_test.go @@ -0,0 +1,133 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package core + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" +) + +// annotatedTool is backendTool with an x-mcp-header annotation on one parameter, +// the SEP-2243 worked example (execute_sql / Region). +func annotatedTool(name, header string) vmcp.Tool { + t := backendTool(name) + t.InputSchema = map[string]any{ + "type": "object", + "properties": map[string]any{ + "region": map[string]any{"type": "string", "x-mcp-header": header}, + "query": map[string]any{"type": "string"}, + }, + } + return t +} + +// TestCallTool_DerivesParamHeadersFromSchema is the pin that makes the whole +// feature non-vacuous: the core is the only layer holding both the tool's +// inputSchema and the call's arguments, so if it fails to derive the Mcp-Param-* +// headers nothing downstream can. Asserted on the exact map handed to the backend +// client rather than on a wire capture, because that hand-off is the seam this +// change introduces. +func TestCallTool_DerivesParamHeadersFromSchema(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + target := backendTarget() + expectAggregation(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{annotatedTool("execute_sql", "Region")}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{"execute_sql": target}}, + }) + + var gotHeaders map[string]string + m.client.EXPECT(). + CallTool(gomock.Any(), gomock.Any(), "execute_sql", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func( + _ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, + paramHeaders map[string]string, + ) (*vmcp.ToolCallResult, error) { + gotHeaders = paramHeaders + return &vmcp.ToolCallResult{}, nil + }) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.CallTool(context.Background(), nil, "execute_sql", + map[string]any{"region": "eu-west1", "query": "select 1"}, nil) + require.NoError(t, err) + + // Only the designated parameter is mirrored; "query" carries no annotation. + assert.Equal(t, map[string]string{"Mcp-Param-Region": "eu-west1"}, gotHeaders) +} + +// TestCallTool_NoAnnotationsDerivesNoHeaders covers the common case: a tool with +// no x-mcp-header annotation must hand the backend client nil, not an empty map, +// so nothing is added to the wire. +func TestCallTool_NoAnnotationsDerivesNoHeaders(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + target := backendTarget() + expectAggregation(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{backendTool("tool_a")}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{"tool_a": target}}, + }) + + var gotHeaders map[string]string + called := false + m.client.EXPECT(). + CallTool(gomock.Any(), gomock.Any(), "tool_a", gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func( + _ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any, _ map[string]any, + paramHeaders map[string]string, + ) (*vmcp.ToolCallResult, error) { + gotHeaders = paramHeaders + called = true + return &vmcp.ToolCallResult{}, nil + }) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.CallTool(context.Background(), nil, "tool_a", map[string]any{"a": 1}, nil) + require.NoError(t, err) + assert.True(t, called) + assert.Nil(t, gotHeaders) +} + +// TestCallTool_UnmirrorableArgumentFailsBeforeDispatch pins the fail-closed +// direction. A caller-supplied value carrying a CRLF must abort the call rather +// than be dropped: dropping it would send no header, the backend would answer +// -32020, and a caller mistake would surface as an opaque backend failure. The +// backend client must therefore never be reached — asserted by setting no EXPECT +// on it, so any call fails the gomock controller. +func TestCallTool_UnmirrorableArgumentFailsBeforeDispatch(t *testing.T) { + t.Parallel() + cfg, m := baseConfig(t) + + target := backendTarget() + expectAggregation(m, &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{annotatedTool("execute_sql", "Region")}, + RoutingTable: &vmcp.RoutingTable{Tools: map[string]*vmcp.BackendTarget{"execute_sql": target}}, + }) + + c, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.CallTool(context.Background(), nil, "execute_sql", + map[string]any{"region": "eu\r\nX-Evil: 1"}, nil) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrInvalidInput), + "a bad argument value must be reported as invalid input, got %v", err) +} diff --git a/pkg/vmcp/errors.go b/pkg/vmcp/errors.go index 1458f3a5ce..9db4a0ad47 100644 --- a/pkg/vmcp/errors.go +++ b/pkg/vmcp/errors.go @@ -68,6 +68,25 @@ var ( ErrToolNameConflict = errors.New("tool name conflict") ) +// Authorization-denial messages, one per admission kind. +// +// These are kind-only: each names the capability class, never the specific +// tool/resource/prompt or whether it is advertised vs. nonexistent, so a denial +// can never be used to enumerate what exists (the "no oracle" property). They +// live here at the domain root and are referenced by every site that emits a +// denial — conversion.ErrorToToolResult, the Serve resources/read handler, the +// codemode script closure, and the pre-dispatch call gate — so a denial has one +// wording per kind by reference, not by copied literals that can silently drift +// apart. Test assertions deliberately keep the literal to pin the exact wording. +const ( + // DenyMessageToolCall is emitted when a tools/call is denied by policy. + DenyMessageToolCall = "call denied by authorization policy" + // DenyMessageResourceRead is emitted when a resources/read is denied by policy. + DenyMessageResourceRead = "read denied by authorization policy" + // DenyMessagePromptGet is emitted when a prompts/get is denied by policy. + DenyMessagePromptGet = "prompt denied by authorization policy" +) + // Error Categorization Helpers // // These functions categorize errors by examining error message strings. diff --git a/pkg/vmcp/errors_test.go b/pkg/vmcp/errors_test.go new file mode 100644 index 0000000000..55517c19b7 --- /dev/null +++ b/pkg/vmcp/errors_test.go @@ -0,0 +1,164 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestIsAuthenticationError pins IsAuthenticationError's exact matching +// boundary (§13): it must recognize the specific phrase patterns the +// function checks for (e.g. "401 unauthorized", "unauthorized (401)", +// "403 forbidden", "authorization required", "access denied") but must NOT +// fire on a bare status code or a bare keyword with no surrounding context — +// otherwise unrelated errors that merely mention "401" or "unauthorized" in +// passing would be misclassified as authentication failures. +func TestIsAuthenticationError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + want bool + }{ + // --- negatives: bare substrings and unrelated errors must not match --- + { + name: "nil error", + err: nil, + want: false, + }, + { + name: "bare 401 with no surrounding phrase", + err: errors.New("401"), + want: false, + }, + { + name: "bare unauthorized with no surrounding phrase", + err: errors.New("unauthorized"), + want: false, + }, + { + name: "unrelated error that merely contains 401 as a numeric substring", + err: errors.New("error 4010 occurred"), + want: false, + }, + { + name: "hostname that merely contains 401", + err: errors.New("http://backend401.example.com"), + want: false, + }, + { + name: "connection refused", + err: errors.New("connection refused"), + want: false, + }, + { + name: "request timeout", + err: errors.New("request timeout"), + want: false, + }, + { + name: "404 not found", + err: errors.New("404 not found"), + want: false, + }, + { + name: "500 internal server error", + err: errors.New("500 internal server error"), + want: false, + }, + { + // Pins against accidental loosening of the "authorization required" + // matcher: "field 'authorization' required" must NOT match. The + // matcher looks for the contiguous substring "authorization + // required"; a future change allowing arbitrary whitespace between + // the words would silently regress this. + name: "validation message with 'authorization' and 'required' separated", + err: errors.New("field 'authorization' required"), + want: false, + }, + + // --- positives: recognized phrase patterns --- + { + name: "authentication failed", + err: errors.New("authentication failed"), + want: true, + }, + { + name: "Authentication Failed (case-insensitive)", + err: errors.New("Authentication Failed"), + want: true, + }, + { + name: "authentication error phrase", + err: errors.New("authentication error: bad token"), + want: true, + }, + { + name: "401 Unauthorized phrase", + err: errors.New("401 Unauthorized"), + want: true, + }, + { + name: "unauthorized (401) reversed phrase (mcp-go ErrUnauthorized form)", + err: errors.New("unauthorized (401)"), + want: true, + }, + { + name: "403 forbidden phrase", + err: errors.New("403 forbidden"), + want: true, + }, + { + name: "HTTP 401 phrase", + err: errors.New("HTTP 401"), + want: true, + }, + { + name: "HTTP 403 phrase", + err: errors.New("HTTP 403"), + want: true, + }, + { + name: "status code 401 phrase", + err: errors.New("status code 401"), + want: true, + }, + { + name: "status code 403 phrase", + err: errors.New("status code 403"), + want: true, + }, + { + name: "request unauthenticated phrase", + err: errors.New("request unauthenticated"), + want: true, + }, + { + name: "request unauthorized phrase", + err: errors.New("request unauthorized"), + want: true, + }, + { + name: "authorization required phrase", + err: errors.New("authorization required"), + want: true, + }, + { + name: "access denied phrase", + err: errors.New("access denied"), + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, IsAuthenticationError(tt.err)) + }) + } +} diff --git a/pkg/vmcp/forwarding.go b/pkg/vmcp/forwarding.go new file mode 100644 index 0000000000..06ccade2b8 --- /dev/null +++ b/pkg/vmcp/forwarding.go @@ -0,0 +1,214 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp + +import "context" + +//go:generate mockgen -destination=mocks/mock_forwarding.go -package=mocks -source=forwarding.go SamplingRequester ClientNotifier + +// MCP wire method names for the server->client notifications the vMCP relay +// forwards. These are the on-the-wire JSON-RPC method strings and MUST stay in +// lockstep between the backend client's OnNotification handler (which matches on +// them) and the server's notifier adapter (which emits them). The mcpcompat mcp +// package does not export these, so they are defined here in the domain package +// as the single source of truth. +const ( + // MethodProgressNotification is the notifications/progress wire method. + MethodProgressNotification = "notifications/progress" + + // MethodLogNotification is the notifications/message (logging) wire method. + MethodLogNotification = "notifications/message" + + // MethodToolsListChangedNotification is the notifications/tools/list_changed + // wire method. Used by the persistent session-backed backend connection + // (pkg/vmcp/session/internal/backend) to recognize a backend's asynchronous + // tool-set change and trigger cache invalidation + session resync (#5748). + MethodToolsListChangedNotification = "notifications/tools/list_changed" + + // MethodResourcesListChangedNotification is the notifications/resources/list_changed + // wire method. Per MCP 2025-11-25 there is no separate wire method for resource + // TEMPLATE changes, so this method also covers a backend's resource-template-set + // change. Used by the persistent session-backed backend connection + // (pkg/vmcp/session/internal/backend) to recognize a backend's asynchronous + // resource (and resource-template) set change and trigger cache invalidation + + // session resync (#5969, mirroring #5748's tools handling). + MethodResourcesListChangedNotification = "notifications/resources/list_changed" + + // MethodPromptsListChangedNotification is the notifications/prompts/list_changed + // wire method. Used by the persistent session-backed backend connection + // (pkg/vmcp/session/internal/backend) to recognize a backend's asynchronous + // prompt-set change and trigger cache invalidation + session resync (#5969, + // mirroring #5748's tools handling). + MethodPromptsListChangedNotification = "notifications/prompts/list_changed" +) + +// SamplingRequester sends an MCP sampling (sampling/createMessage) request to the +// client and blocks for the response. +// +// This is the server->client forwarding seam for MCP sampling, mirroring +// ElicitationRequester: a backend that, mid tools/call, asks its caller to sample +// an LLM must have that request relayed to the downstream vMCP client. Callers +// (the backend client's sampling handler) depend only on this interface and the +// domain SamplingRequest/SamplingResult value types, never on the underlying SDK. +// The transport adapter (sdkSamplingAdapter in pkg/vmcp/server) is the sole point +// that translates to/from mcp-go types. +// +// The underlying SDK handles JSON-RPC ID correlation and session routing +// internally, so implementations do not track request IDs. +type SamplingRequester interface { + // RequestSampling sends a sampling request and blocks until the client + // responds or the context is cancelled. Implementations return an error when + // the context carries no downstream session or the client did not advertise + // the sampling capability. + RequestSampling(ctx context.Context, req SamplingRequest) (*SamplingResult, error) +} + +// ClientNotifier forwards a backend's mid-call server->client notifications +// (progress, logging) to the downstream vMCP client on the same session. +// +// This is the server->client forwarding seam for notifications. Callers (the +// backend client's OnNotification handler) depend only on this interface and the +// domain notification value types, never on the underlying SDK. The transport +// adapter (sdkNotifierAdapter in pkg/vmcp/server) wraps the SDK's +// SendNotificationToClient. Forwarding is best-effort: a missing downstream +// session is not an error worth surfacing to the backend, so implementations +// treat it as a no-op. +type ClientNotifier interface { + // NotifyProgress forwards a notifications/progress message to the downstream + // client. It is best-effort; a nil error does not guarantee delivery. + NotifyProgress(ctx context.Context, n ProgressNotification) error + + // NotifyLog forwards a notifications/message (logging) message to the + // downstream client. It is best-effort. + NotifyLog(ctx context.Context, n LogMessage) error +} + +// ClientForwarderBinder is implemented by a BackendClient that can have the +// server->client forwarding requesters bound after the SDK server is +// constructed. +// +// server.New evaluates the BackendClient before Serve builds the mcp-go server +// that the SDK adapters wrap (a construction-order inversion, the same one the +// late-bound elicitation requester solves). New type-asserts the concrete +// backend client to this interface and binds the real adapters once, before the +// server begins serving. A BackendClient that does not implement it simply does +// not forward server->client traffic. +type ClientForwarderBinder interface { + // BindForwarders installs the elicitation, sampling, and notification + // requesters used to relay a backend's mid-call server->client traffic to the + // downstream client. It is called exactly once, before serving begins. Any of + // the arguments may be nil to leave that forwarding path disabled. + BindForwarders(elicitation ElicitationRequester, sampling SamplingRequester, notifier ClientNotifier) +} + +// SamplingRequest is the domain-typed sampling (sampling/createMessage) request. +// +// It mirrors, one-to-one, the mcp-go CreateMessageParams fields the vMCP relay +// forwards. The adapter in pkg/vmcp/server translates this to the SDK request +// type; the backend client's sampling handler builds it from the SDK request. +type SamplingRequest struct { + // Messages is the conversation history to sample from. + Messages []SamplingMessage + + // ModelPreferences expresses the backend's model-selection preferences. The + // client may ignore them. Nil when the backend set none. + ModelPreferences *ModelPreferences + + // SystemPrompt is an optional system prompt the backend wants to use. + SystemPrompt string + + // IncludeContext requests context be attached ("none", "thisServer", + // "allServers"). The client may ignore it. + IncludeContext string + + // Temperature is the sampling temperature. + Temperature float64 + + // MaxTokens is the maximum number of tokens to sample. + MaxTokens int + + // StopSequences are sequences that, when generated, stop sampling. + StopSequences []string + + // Metadata is optional provider-specific metadata passed through to the LLM. + Metadata any +} + +// SamplingMessage is a single message in a sampling conversation. Content stays +// typed as any (a text/image/audio content block) matching the SDK's wire shape, +// which the vMCP relay passes through without inspection. +type SamplingMessage struct { + // Role is the message role ("user" or "assistant"). + Role string + + // Content is the message content block (text/image/audio). + Content any +} + +// ModelPreferences mirrors the mcp-go model-selection preferences. All fields +// are advisory. +type ModelPreferences struct { + // Hints are ordered model-selection hints. + Hints []ModelHint + + // CostPriority weights cost in model selection (0..1). + CostPriority float64 + + // SpeedPriority weights latency in model selection (0..1). + SpeedPriority float64 + + // IntelligencePriority weights capability in model selection (0..1). + IntelligencePriority float64 +} + +// ModelHint is a single model-selection hint. +type ModelHint struct { + // Name is treated by the client as a substring of a model name. + Name string +} + +// SamplingResult is the domain-typed sampling response. Content stays typed as +// any, matching the SDK's wire shape. +type SamplingResult struct { + // Role is the role of the sampled message ("assistant"). + Role string + + // Content is the sampled content block (text/image/audio). + Content any + + // Model is the name of the model that generated the message. + Model string + + // StopReason is the reason sampling stopped, if known. + StopReason string +} + +// ProgressNotification is the domain-typed notifications/progress message. +type ProgressNotification struct { + // ProgressToken correlates the notification with the request that requested + // progress. It is an opaque token (string or number) passed through unchanged. + ProgressToken any + + // Progress is the current progress value. + Progress float64 + + // Total is the optional total against which Progress is measured. Zero means + // unset. + Total float64 + + // Message is an optional human-readable progress description. + Message string +} + +// LogMessage is the domain-typed notifications/message (logging) message. +type LogMessage struct { + // Level is the syslog-style severity ("debug", "info", ..., "emergency"). + Level string + + // Logger is an optional logger name. + Logger string + + // Data is the log payload (any JSON value). + Data any +} diff --git a/pkg/vmcp/headerforward/transport.go b/pkg/vmcp/headerforward/transport.go index cf2c9cfcb3..e87710f197 100644 --- a/pkg/vmcp/headerforward/transport.go +++ b/pkg/vmcp/headerforward/transport.go @@ -64,6 +64,15 @@ func (h *headerForwardRoundTripper) RoundTrip(req *http.Request) (*http.Response return h.base.RoundTrip(reqCopy) } +// CloseIdleConnections forwards to the wrapped RoundTripper so http.Client's +// CloseIdleConnections reaches the concrete *http.Transport at the bottom of the +// chain (this wrapper would otherwise silently swallow the call). +func (h *headerForwardRoundTripper) CloseIdleConnections() { + if c, ok := h.base.(interface{ CloseIdleConnections() }); ok { + c.CloseIdleConnections() + } +} + // BuildHeaderForwardTripper constructs a headerForwardRoundTripper for the // backend's pre-resolved HeaderForwardConfig. Returns base unchanged when no // header injection is configured or the effective header set is empty. diff --git a/pkg/vmcp/headerforward/transport_test.go b/pkg/vmcp/headerforward/transport_test.go index c8f2b42766..464026709c 100644 --- a/pkg/vmcp/headerforward/transport_test.go +++ b/pkg/vmcp/headerforward/transport_test.go @@ -398,3 +398,22 @@ func TestMergeForwardedHeaders_RestrictedHeadersList(t *testing.T) { } } } + +// closeIdleSpy is a RoundTripper that records CloseIdleConnections calls. +type closeIdleSpy struct{ closed int } + +func (*closeIdleSpy) RoundTrip(*http.Request) (*http.Response, error) { + return nil, errors.New("unused") +} +func (s *closeIdleSpy) CloseIdleConnections() { s.closed++ } + +// TestHeaderForwardRoundTripper_CloseIdleConnections verifies the wrapper forwards +// CloseIdleConnections to its base rather than swallowing it. +func TestHeaderForwardRoundTripper_CloseIdleConnections(t *testing.T) { + t.Parallel() + + spy := &closeIdleSpy{} + rt := &headerForwardRoundTripper{base: spy} + rt.CloseIdleConnections() + assert.Equal(t, 1, spy.closed) +} diff --git a/pkg/vmcp/health/checker_test.go b/pkg/vmcp/health/checker_test.go index c8d2a20cf1..75949209d7 100644 --- a/pkg/vmcp/health/checker_test.go +++ b/pkg/vmcp/health/checker_test.go @@ -10,11 +10,11 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client/transport" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" "github.com/stacklok/toolhive/pkg/vmcp" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" "github.com/stacklok/toolhive/pkg/vmcp/mocks" @@ -521,58 +521,11 @@ func TestCategorizeError(t *testing.T) { } } -func TestIsAuthenticationError(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - err error - expectErr bool - }{ - // Positive cases - {name: "authentication failed", err: errors.New("authentication failed"), expectErr: true}, - {name: "Authentication Failed (uppercase)", err: errors.New("Authentication Failed"), expectErr: true}, - {name: "authentication error", err: errors.New("authentication error: bad token"), expectErr: true}, - {name: "401 unauthorized", err: errors.New("401 unauthorized"), expectErr: true}, - {name: "403 forbidden", err: errors.New("403 forbidden"), expectErr: true}, - {name: "HTTP 401", err: errors.New("HTTP 401"), expectErr: true}, - {name: "HTTP 403", err: errors.New("HTTP 403"), expectErr: true}, - {name: "status code 401", err: errors.New("status code 401"), expectErr: true}, - {name: "status code 403", err: errors.New("status code 403"), expectErr: true}, - {name: "request unauthenticated", err: errors.New("request unauthenticated"), expectErr: true}, - {name: "request unauthorized", err: errors.New("request unauthorized"), expectErr: true}, - {name: "access denied", err: errors.New("access denied"), expectErr: true}, - - // mcp-go ErrUnauthorized format: "unauthorized (401)" (reversed order vs "401 unauthorized") - {name: "unauthorized (401) - mcp-go ErrUnauthorized format", err: errors.New("unauthorized (401)"), expectErr: true}, - - // Negative cases - should NOT be detected as auth errors - {name: "connection refused", err: errors.New("connection refused"), expectErr: false}, - {name: "timeout", err: errors.New("request timeout"), expectErr: false}, - {name: "generic error", err: errors.New("something went wrong"), expectErr: false}, - {name: "404 not found", err: errors.New("404 not found"), expectErr: false}, - {name: "500 internal server error", err: errors.New("500 internal server error"), expectErr: false}, - {name: "hostname with 401", err: errors.New("http://backend401.example.com"), expectErr: false}, - // Pin against accidental loosening of the "authorization required" - // substring matcher: a validation message of the form "field - // 'authorization' required" must not be misclassified as an auth - // failure. The current matcher uses the contiguous substring - // "authorization required" (one space, no punctuation), so this - // string does not match — but a future loosening (e.g. allowing - // any whitespace) would silently regress. - {name: "validation message containing 'authorization' and 'required'", err: errors.New("field 'authorization' required"), expectErr: false}, - {name: "nil error", err: nil, expectErr: false}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - result := vmcp.IsAuthenticationError(tt.err) - assert.Equal(t, tt.expectErr, result) - }) - } -} +// NOTE: IsAuthenticationError is exhaustively tested in its owning package at +// pkg/vmcp/errors_test.go (TestIsAuthenticationError). It was previously +// re-tested here, but a classifier owned by pkg/vmcp belongs under test there, +// not in the health package that merely consumes it (see .claude/rules/testing.md +// "Test Scope"). func TestIsTimeoutError(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/health/exposure.go b/pkg/vmcp/health/exposure.go deleted file mode 100644 index 561f4cf7d7..0000000000 --- a/pkg/vmcp/health/exposure.go +++ /dev/null @@ -1,87 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package health - -import ( - "log/slog" - "reflect" - - "github.com/stacklok/toolhive/pkg/vmcp" -) - -// IsStatusEligibleForExposure returns true when a backend should remain visible -// in vMCP discovery/session capability surfaces. -func IsStatusEligibleForExposure(status vmcp.BackendHealthStatus) bool { - return status == "" || - status == vmcp.BackendHealthy || - status == vmcp.BackendDegraded -} - -// FilterBackendsForExposure filters backends to only include those that are -// currently eligible to be exposed to sessions and discovery. -// -// Health status filtering: -// - healthy: included -// - degraded: included -// - empty/zero-value: included (assume healthy when health monitoring is disabled) -// - unhealthy: excluded -// - unknown: excluded -// - unauthenticated: excluded -// -// When healthStatusProvider is provided, the current health status from the -// health monitor is used. When nil, falls back to the backend registry status. -func FilterBackendsForExposure(backends []vmcp.Backend, healthStatusProvider StatusProvider) []vmcp.Backend { - if len(backends) == 0 { - return backends - } - - eligible := make([]vmcp.Backend, 0, len(backends)) - excluded := 0 - - for i := range backends { - backend := &backends[i] - - healthStatus := backend.HealthStatus - if !isNilStatusProvider(healthStatusProvider) { - if status, exists := healthStatusProvider.QueryBackendStatus(backend.ID); exists { - healthStatus = status - } - } - - if IsStatusEligibleForExposure(healthStatus) { - eligible = append(eligible, *backend) - continue - } - - excluded++ - slog.Debug("excluding backend from exposure due to health status", - "backend_name", backend.Name, - "backend_id", backend.ID, - "health_status", healthStatus, - ) - } - - if excluded > 0 { - slog.Debug("filtered backends for exposure", - "total_backends", len(backends), - "eligible_backends", len(eligible), - "excluded_backends", excluded, - ) - } - - return eligible -} - -func isNilStatusProvider(provider StatusProvider) bool { - if provider == nil { - return true - } - - value := reflect.ValueOf(provider) - if kind := value.Kind(); kind == reflect.Chan || kind == reflect.Func || kind == reflect.Interface || - kind == reflect.Map || kind == reflect.Pointer || kind == reflect.Slice { - return value.IsNil() - } - return false -} diff --git a/pkg/vmcp/health/monitor.go b/pkg/vmcp/health/monitor.go index f7c8afbed8..907a802272 100644 --- a/pkg/vmcp/health/monitor.go +++ b/pkg/vmcp/health/monitor.go @@ -121,6 +121,10 @@ type Monitor struct { // checker performs health checks on backends. checker vmcp.HealthChecker + // revisions reads each backend's negotiated MCP revision for the status + // read-model. Nil when the client does not implement vmcp.RevisionReporter. + revisions vmcp.RevisionReporter + // statusTracker tracks health status for all backends. statusTracker *statusTracker @@ -265,8 +269,13 @@ func NewMonitor( // The status tracker will lazily initialize circuit breakers as needed statusTracker := newStatusTracker(config.UnhealthyThreshold, config.CircuitBreaker) + // The client (directly or via the telemetry decorator) optionally reports the + // negotiated MCP revision for the status read-model; nil when unsupported. + revisions, _ := client.(vmcp.RevisionReporter) + return &Monitor{ checker: checker, + revisions: revisions, statusTracker: statusTracker, checkInterval: config.CheckInterval, backends: backends, @@ -514,6 +523,14 @@ func (m *Monitor) performHealthCheck(ctx context.Context, backend *vmcp.Backend) m.statusTracker.RecordSuccess(backend.ID, backend.Name, status) } + // Refresh the MCP revision read-model from the client's cache (empty until the + // backend is probed). Read-only; a no-op when the client doesn't report it. + if m.revisions != nil { + if rev, ok := m.revisions.CachedRevision(backend.ID); ok { + m.statusTracker.RecordRevision(backend.ID, rev.String()) + } + } + currentStatus, currentExists := m.statusTracker.GetStatus(backend.ID) if previousExists && currentExists && previousStatus != currentStatus { m.notifyStatusChange(StatusChangeEvent{ @@ -781,6 +798,7 @@ func (m *Monitor) convertToDiscoveredBackends(allStates map[string]*State) []vmc CircuitBreakerState: string(state.CircuitState), CircuitLastChanged: metav1.NewTime(state.CircuitLastChanged), ConsecutiveFailures: state.ConsecutiveFailures, + MCPRevision: state.MCPRevision, }) continue } @@ -798,6 +816,7 @@ func (m *Monitor) convertToDiscoveredBackends(allStates map[string]*State) []vmc CircuitBreakerState: string(state.CircuitState), CircuitLastChanged: metav1.NewTime(state.CircuitLastChanged), ConsecutiveFailures: state.ConsecutiveFailures, + MCPRevision: state.MCPRevision, }) } diff --git a/pkg/vmcp/health/monitor_test.go b/pkg/vmcp/health/monitor_test.go index f59059e21b..f0c87ec14b 100644 --- a/pkg/vmcp/health/monitor_test.go +++ b/pkg/vmcp/health/monitor_test.go @@ -13,6 +13,7 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/mocks" ) @@ -1362,3 +1363,68 @@ func TestMonitor_CircuitBreakerStatusReporting(t *testing.T) { err = monitor.Stop() require.NoError(t, err) } + +// TestConvertToDiscoveredBackends_MCPRevision verifies the negotiated revision is +// surfaced in DiscoveredBackend: Modern, Legacy, and empty-when-unprobed, across +// both the in-list and fallback (not-in-list) construction branches. +func TestConvertToDiscoveredBackends_MCPRevision(t *testing.T) { + t.Parallel() + + m := &Monitor{backends: []vmcp.Backend{{ID: "modern", Name: "modern-b", BaseURL: "u"}}} + states := map[string]*State{ + "modern": {Status: vmcp.BackendHealthy, MCPRevision: "2026-07-28"}, + "legacy": {Status: vmcp.BackendHealthy, MCPRevision: "2025-11-25"}, // not in backends -> fallback branch + "unprobed": {Status: vmcp.BackendHealthy}, + } + + byName := make(map[string]string) + for _, b := range m.convertToDiscoveredBackends(states) { + byName[b.Name] = b.MCPRevision + } + + assert.Equal(t, "2026-07-28", byName["modern-b"], "in-list backend surfaces Modern") + assert.Equal(t, "2025-11-25", byName["legacy"], "fallback backend surfaces Legacy") + assert.Equal(t, "", byName["unprobed"], "unprobed backend has empty revision") +} + +// revClientStub wraps a MockBackendClient (for the health-check ListCapabilities +// calls) and adds the optional CachedRevision accessor the monitor reads. +type revClientStub struct { + *mocks.MockBackendClient + rev mcpparser.Revision +} + +func (s revClientStub) CachedRevision(string) (mcpparser.Revision, bool) { return s.rev, true } + +// TestMonitor_SurfacesMCPRevision verifies the performHealthCheck wiring: the +// NewMonitor revisionReporter type-assertion plus the read-and-RecordRevision +// after a check surface the client's revision in GetState. +func TestMonitor_SurfacesMCPRevision(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + client := revClientStub{MockBackendClient: mockClient, rev: mcpparser.RevisionModern} + + backends := []vmcp.Backend{{ID: "backend-1", Name: "Backend 1", BaseURL: "http://localhost:8080", TransportType: "sse"}} + monitor, err := NewMonitor(client, backends, MonitorConfig{ + CheckInterval: 100 * time.Millisecond, + UnhealthyThreshold: 3, + Timeout: 50 * time.Millisecond, + }) + require.NoError(t, err) + + require.NoError(t, monitor.Start(context.Background())) + t.Cleanup(func() { _ = monitor.Stop() }) + + require.Eventually(t, func() bool { + state, err := monitor.GetBackendState("backend-1") + return err == nil && state != nil && state.MCPRevision == "2026-07-28" + }, 500*time.Millisecond, 10*time.Millisecond, "the client's negotiated revision must surface in health state") +} diff --git a/pkg/vmcp/health/policy.go b/pkg/vmcp/health/policy.go new file mode 100644 index 0000000000..aab8f99227 --- /dev/null +++ b/pkg/vmcp/health/policy.go @@ -0,0 +1,78 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import "github.com/stacklok/toolhive/pkg/vmcp" + +// ShouldAdvertise reports whether a backend in this status may contribute +// capabilities to the advertised view (tools/list and friends). +// +// Degraded backends are included: they are slow but working, and hiding their +// tools would be a worse outcome for the caller than serving them. An empty +// status means health monitoring is disabled, which is treated as healthy so a +// deployment without a monitor behaves as it did before monitoring existed. +// +// Excluded: unhealthy (not responding), unknown (not yet probed), and +// unauthenticated (operator misconfiguration). +func ShouldAdvertise(status vmcp.BackendHealthStatus) bool { + return status == "" || + status == vmcp.BackendHealthy || + status == vmcp.BackendDegraded +} + +// ShouldOpenSession reports whether a session should attempt to open a +// connection to a backend in this status. +// +// It skips only statuses that positively establish the backend is a bad bet and +// admits everything else, including not-yet-classified. This is the fix for +// #5861: session establishment waits for every backend it attempts +// (session.makeBaseSession's wg.Wait), so a backend the monitor already knows is +// bad sets the floor for the whole tenant's session-establishment latency. Worse, +// the handshake makes several sequential round trips, so the cost is a multiple +// of the backend's per-request latency, not one unit of it. The reported +// backend's 10-25s latency exceeds the 10s probe timeout, so its checks fail and +// it reaches Unhealthy after UnhealthyThreshold consecutive failures — which is +// the status this predicate skips. +// +// Degraded is deliberately ADMITTED. The tempting reading — "degraded means slow, +// so don't block on it" — does not survive the status's three producers +// (see vmcp.BackendDegraded): +// +// - Slow probe (healthChecker.CheckHealth): a successful check slower than +// DegradedThreshold. Genuinely slow, and the only producer where skipping +// would buy latency. +// - Recovering (statusTracker.RecordSuccess): ANY success recorded while +// consecutiveFailures > 0 is forced to degraded, overriding the check's own +// verdict. The backend just answered, possibly in microseconds, and stays +// labelled degraded for up to one CheckInterval (30s default). +// - Auth retrying (workloads.mapWorkloadStatusToVMCPHealth): a transient +// OAuth-refresh failure, with no latency component at all. +// +// Skipping degraded would exclude a fast, working, just-recovered backend from +// every session created in the ~30s after it recovers, making recovery slower to +// take effect — a worse and more surprising failure than the one #5861 reports. +// Telling the producers apart needs a degradation reason the enum does not carry; +// until it does, the safe reading of degraded at session-open time is "still +// worth attempting". +// +// The residual cost is accepted knowingly: a backend slow enough to be degraded +// but faster than the probe timeout (DegradedThreshold..Timeout) stays on the +// session-establishment path. That window is bounded by the probe timeout, where +// #5861's backend was not. +// +// Unknown is admitted for a separate reason: serving is not gated on the first +// health check completing (only the status reporter calls +// WaitForInitialHealthChecks), so sessions are routinely created while backends +// are still Unknown — during pod startup, and for a backend whose first check +// failed below the unhealthy threshold, which the monitor records as Unknown with +// a non-zero failure count (statusTracker.RecordFailure). Skipping those would +// connect a session to zero backends during the startup window: both a +// regression against the pre-#5861 behaviour and a worse failure than the one +// being fixed. "Not yet known to be bad" must fail open. +func ShouldOpenSession(status vmcp.BackendHealthStatus) bool { + // Skip only confirmed-bad statuses; everything else — including Degraded, + // Unknown, and the empty zero value — is attempted. + return status != vmcp.BackendUnhealthy && + status != vmcp.BackendUnauthenticated +} diff --git a/pkg/vmcp/health/policy_test.go b/pkg/vmcp/health/policy_test.go new file mode 100644 index 0000000000..bb2a24f791 --- /dev/null +++ b/pkg/vmcp/health/policy_test.go @@ -0,0 +1,108 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package health + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestShouldAdvertiseAndShouldOpenSession pins both health predicates against +// every BackendHealthStatus constant, plus the empty zero value and an +// unrecognized value. +// +// The two are asserted together because their relationship is the contract that +// matters, and it is not a simple ordering. ShouldOpenSession is LOOSER for +// unknown (session establishment must not fail closed before the first health +// check completes), and the two now COINCIDE on degraded — the deliberate outcome +// of #5861's review: that status conflates "slow", "recovering" and "auth +// retrying", so it cannot be read as "slow" at session-open time. Testing them +// side by side makes an accidental change to either one visible as a change in +// the pairing. +// +// The unrecognized-value row pins a subtlety worth stating explicitly: the two +// predicates have opposite defaults for a status neither knows about. +// ShouldAdvertise is an allow-list, so an unrecognized status is NOT advertised +// (fails closed — a capability that may not be servable is withheld). +// ShouldOpenSession is a deny-list, so it IS attempted (fails open — better to +// connect to a backend of uncertain health than to strand a session with none). +// Each default is the conservative choice for its own question, but they point +// in opposite directions, so anyone adding a status must consider both. +func TestShouldAdvertiseAndShouldOpenSession(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + status vmcp.BackendHealthStatus + wantAdvertise bool + wantOpenSession bool + }{ + { + name: "empty means health monitoring disabled: assume usable", + status: "", + wantAdvertise: true, + wantOpenSession: true, + }, + { + name: "healthy", + status: vmcp.BackendHealthy, + wantAdvertise: true, + wantOpenSession: true, + }, + { + // Degraded is attempted, not skipped. Only one of its three producers + // is latency; the "recovering" producer forces degraded onto a backend + // that just answered successfully, so skipping it would sideline a + // fast, working backend for up to one check interval. + name: "degraded is advertised AND attempted", + status: vmcp.BackendDegraded, + wantAdvertise: true, + wantOpenSession: true, + }, + { + name: "unhealthy", + status: vmcp.BackendUnhealthy, + wantAdvertise: false, + wantOpenSession: false, + }, + { + // The asymmetry that remains: aggregation waits for confirmation, + // session establishment must not, or a cold monitor connects sessions + // to zero backends during pod startup. + name: "unknown is not advertised but is still attempted", + status: vmcp.BackendUnknown, + wantAdvertise: false, + wantOpenSession: true, + }, + { + name: "unauthenticated (operator misconfiguration)", + status: vmcp.BackendUnauthenticated, + wantAdvertise: false, + wantOpenSession: false, + }, + { + // ShouldAdvertise is an allow-list (fails closed); ShouldOpenSession is + // a deny-list (fails open). Opposite defaults, deliberately — see the + // doc comment above. + name: "unrecognized status: not advertised, but still attempted", + status: vmcp.BackendHealthStatus("some-future-status"), + wantAdvertise: false, + wantOpenSession: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + assert.Equal(t, tt.wantAdvertise, ShouldAdvertise(tt.status), + "ShouldAdvertise(%q)", tt.status) + assert.Equal(t, tt.wantOpenSession, ShouldOpenSession(tt.status), + "ShouldOpenSession(%q)", tt.status) + }) + } +} diff --git a/pkg/vmcp/health/status.go b/pkg/vmcp/health/status.go index f0056315e1..6dbd03c586 100644 --- a/pkg/vmcp/health/status.go +++ b/pkg/vmcp/health/status.go @@ -32,6 +32,12 @@ type backendHealthState struct { // circuitBreaker manages circuit breaker state for this backend. // Always non-nil; uses alwaysClosedCircuit when circuit breaker is disabled. circuitBreaker CircuitBreaker + + // mcpRevision is the backend's negotiated MCP revision as a read-model copy + // (e.g. "2026-07-28"/"2025-11-25"), sourced from the client's cache during the + // health check. Empty when the backend has not been probed. May lag the client + // cache slightly — fine for status. + mcpRevision string } // statusTracker tracks health status for multiple backends. @@ -171,6 +177,7 @@ func (*statusTracker) copyState(state *backendHealthState) *State { LastErrorCategory: sanitizeError(state.lastError), LastError: state.lastError, LastTransitionTime: state.lastTransitionTime, + MCPRevision: state.mcpRevision, } // Include circuit breaker state @@ -242,6 +249,18 @@ func (t *statusTracker) RecordSuccess(backendID string, backendName string, stat state.circuitBreaker.RecordSuccess() } +// RecordRevision stores the backend's negotiated MCP revision read-model. It is +// a no-op for an untracked or removed backend (RecordSuccess/RecordFailure runs +// first and creates the state, so a tracked backend always exists here). +func (t *statusTracker) RecordRevision(backendID, revision string) { + t.mu.Lock() + defer t.mu.Unlock() + + if state, exists := t.states[backendID]; exists { + state.mcpRevision = revision + } +} + // RecordFailure records a failed health check for a backend. // This increments the consecutive failure count and may transition the backend to unhealthy // if the threshold is exceeded. Status transitions are logged. @@ -514,4 +533,9 @@ type State struct { // CircuitLastChanged is when the circuit breaker state last changed. // When circuit breaker is disabled, this will be zero time (via alwaysClosedCircuit). CircuitLastChanged time.Time + + // MCPRevision is the backend's negotiated MCP revision ("2026-07-28" or + // "2025-11-25"), or empty when the backend has not been probed. This is a + // read-model copy of the client's cached revision. + MCPRevision string } diff --git a/pkg/vmcp/health/status_test.go b/pkg/vmcp/health/status_test.go index 4ddc7b7ebd..9d34ac8681 100644 --- a/pkg/vmcp/health/status_test.go +++ b/pkg/vmcp/health/status_test.go @@ -746,3 +746,22 @@ func TestSanitizeError(t *testing.T) { }) } } + +// TestStatusTracker_RecordRevision verifies the MCP revision read-model is stored +// on a tracked backend and is a no-op for an untracked one. +func TestStatusTracker_RecordRevision(t *testing.T) { + t.Parallel() + + tr := newStatusTracker(3, nil) + tr.RecordSuccess("b1", "n1", vmcp.BackendHealthy) // creates the state + tr.RecordRevision("b1", "2026-07-28") + + st, ok := tr.GetState("b1") + require.True(t, ok) + assert.Equal(t, "2026-07-28", st.MCPRevision) + + // No-op for an untracked backend (must not create state). + tr.RecordRevision("missing", "2025-11-25") + _, ok = tr.GetState("missing") + assert.False(t, ok) +} diff --git a/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go b/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go index 57ae2c9691..b804d95c8a 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go @@ -13,14 +13,17 @@ package backendtelemetry import ( "context" "fmt" + "sync" "time" + "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/telemetry" transporttypes "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/pkg/vmcp" @@ -30,6 +33,33 @@ const ( instrumentationName = "github.com/stacklok/toolhive/pkg/vmcp" ) +var ( + reclassCounterOnce sync.Once + reclassCounter metric.Int64Counter +) + +// RecordRevisionReclassification increments the count of backends whose MCP +// revision was reclassified after a call revealed the cached revision was wrong. +// +// The backend client resolves revisions below the telemetry decorator, so this +// is a free function backed by the global meter provider rather than the injected +// one used by MonitorBackends. +// +// NOTE: no labels yet — old/new revision labels (and the CRD status surface) +// are deferred. If the global provider ever diverges from the injected one, thread +// the meter down instead. +func RecordRevisionReclassification(ctx context.Context) { + reclassCounterOnce.Do(func() { + reclassCounter, _ = otel.GetMeterProvider().Meter(instrumentationName).Int64Counter( + "toolhive_vmcp_backend_revision_reclassifications", + metric.WithDescription("Number of times a backend's MCP revision was reclassified after a mismatch"), + ) + }) + if reclassCounter != nil { + reclassCounter.Add(ctx, 1) + } +} + // MonitorBackends decorates the backend client so it records telemetry on each method call. // It also emits a gauge for the number of backends discovered once, since the number of backends is static. func MonitorBackends( @@ -103,6 +133,26 @@ type telemetryBackendClient struct { var _ vmcp.BackendClient = telemetryBackendClient{} +// CachedRevision forwards to the wrapped client's optional vmcp.RevisionReporter so +// callers reaching the client THROUGH this decorator (e.g. the health monitor) +// can still read the negotiated revision. Returns (0, false) when the wrapped +// client does not report revisions. +func (t telemetryBackendClient) CachedRevision(workloadID string) (mcpparser.Revision, bool) { + if r, ok := t.backendClient.(vmcp.RevisionReporter); ok { + return r.CachedRevision(workloadID) + } + return 0, false +} + +// revisionLabel returns the backend's negotiated MCP revision as a metric label +// value, or "" when unprobed/unknown (low cardinality: 2 values + empty). +func (t telemetryBackendClient) revisionLabel(workloadID string) string { + if rev, ok := t.CachedRevision(workloadID); ok { + return rev.String() + } + return "" +} + // mapActionToMCPMethod maps internal action names to MCP method names per the OTEL MCP spec. func mapActionToMCPMethod(action string) string { switch action { @@ -153,6 +203,8 @@ func (t telemetryBackendClient) record( attribute.String("action", action), // OTEL MCP spec-required attributes attribute.String("mcp.method.name", mcpMethod), + // Negotiated MCP revision (low cardinality: 2 values + empty when unprobed). + attribute.String("mcp.protocol.revision", t.revisionLabel(target.WorkloadID)), } commonAttrs = append(commonAttrs, attrs...) @@ -205,6 +257,7 @@ func (t telemetryBackendClient) CallTool( toolName string, arguments map[string]any, meta map[string]any, + paramHeaders map[string]string, ) (_ *vmcp.ToolCallResult, retErr error) { attrs := []attribute.KeyValue{ attribute.String("tool_name", toolName), // backward compat @@ -216,7 +269,7 @@ func (t telemetryBackendClient) CallTool( } ctx, done := t.record(ctx, target, "call_tool", toolName, &retErr, attrs...) defer done() - return t.backendClient.CallTool(ctx, target, toolName, arguments, meta) + return t.backendClient.CallTool(ctx, target, toolName, arguments, meta, paramHeaders) } func (t telemetryBackendClient) ReadResource( @@ -253,6 +306,26 @@ func (t telemetryBackendClient) GetPrompt( return t.backendClient.GetPrompt(ctx, target, name, arguments) } +func (t telemetryBackendClient) Complete( + ctx context.Context, + target *vmcp.BackendTarget, + ref vmcp.CompletionRef, + argName, argValue string, + contextArgs map[string]string, +) (_ *vmcp.CompletionResult, retErr error) { + attrs := []attribute.KeyValue{ + attribute.String("completion.ref_type", ref.Type), + attribute.String("completion.argument_name", argName), + } + // Check if caller is authenticated (extract from context) + if caller, _ := auth.IdentityFromContext(ctx); caller != nil && caller.Subject != "" { + attrs = append(attrs, attribute.Bool("auth.authenticated", true)) + } + ctx, done := t.record(ctx, target, "complete", "", &retErr, attrs...) + defer done() + return t.backendClient.Complete(ctx, target, ref, argName, argValue, contextArgs) +} + func (t telemetryBackendClient) ListCapabilities( ctx context.Context, target *vmcp.BackendTarget, ) (_ *vmcp.CapabilityList, retErr error) { diff --git a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go index 0c96135357..5a43f7605e 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go @@ -4,9 +4,60 @@ package backendtelemetry import ( + "context" "testing" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" ) +// fakeRevClient embeds vmcp.BackendClient (nil — its methods are never called +// here) and adds the optional CachedRevision accessor. +type fakeRevClient struct { + vmcp.BackendClient + rev mcpparser.Revision + ok bool +} + +func (f fakeRevClient) CachedRevision(string) (mcpparser.Revision, bool) { return f.rev, f.ok } + +// fakeNoRevClient embeds vmcp.BackendClient but does NOT implement revisionReporter. +type fakeNoRevClient struct{ vmcp.BackendClient } + +// TestTelemetryBackendClient_CachedRevisionForwarding verifies the decorator +// forwards CachedRevision to a client that reports it, and reports nothing for a +// client that doesn't. +func TestTelemetryBackendClient_CachedRevisionForwarding(t *testing.T) { + t.Parallel() + + d := telemetryBackendClient{backendClient: fakeRevClient{rev: mcpparser.RevisionModern, ok: true}} + rev, ok := d.CachedRevision("b") + if !ok || rev != mcpparser.RevisionModern { + t.Fatalf("CachedRevision = (%v, %v), want (Modern, true)", rev, ok) + } + if got := d.revisionLabel("b"); got != "2026-07-28" { + t.Errorf("revisionLabel = %q, want 2026-07-28", got) + } + + dn := telemetryBackendClient{backendClient: fakeNoRevClient{}} + if _, ok := dn.CachedRevision("b"); ok { + t.Error("CachedRevision should report false for a client without the accessor") + } + if got := dn.revisionLabel("b"); got != "" { + t.Errorf("revisionLabel = %q, want empty for unprobed/unsupported", got) + } +} + +// TestRecordRevisionReclassification is a smoke test: the counter lazily binds to +// the global meter provider and increments without panicking (the noop provider +// makes the value unobservable here — the WARN in the same reclassify branch is +// asserted in the client package's reclassify test). +func TestRecordRevisionReclassification(t *testing.T) { + t.Parallel() + RecordRevisionReclassification(context.Background()) + RecordRevisionReclassification(context.Background()) +} + func TestMapActionToMCPMethod(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/internal/compositetools/annotations.go b/pkg/vmcp/internal/compositetools/annotations.go new file mode 100644 index 0000000000..9a0e85ebc4 --- /dev/null +++ b/pkg/vmcp/internal/compositetools/annotations.go @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package compositetools + +import ( + "fmt" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// StepAnnotationResolver resolves a composite-tool step's tool reference +// ("{workloadID}.{toolName}") to the backend tool's annotations. It returns nil +// when the step tool is unknown or the backend declares no annotations. +type StepAnnotationResolver func(stepTool string) *vmcp.ToolAnnotations + +// DeriveCompositeAnnotations computes the safety-floor annotations for a +// composite tool from the annotations of its step tools. +// +// Fail-closed contract (issue #6192): +// - A workflow with NO tool steps (e.g. only elicitation steps) yields an +// empty/nil stepAnn slice, and Derive returns nil — there is nothing to +// derive, so no floor is advertised. +// - A workflow with one or more tool steps yields a non-nil stepAnn slice, +// and Derive ALWAYS returns a conservative floor, even when every step +// tool's annotations are nil/unknown. An unknown step never makes the tool +// look safer — it taints destructive/openWorld to true and readOnly to +// false. This means an explicit readOnlyHint:true over backends that +// declare no annotations (the common case) CONTRADICTS the floor and is +// dropped by CheckAnnotationContradiction. +// +// Derivation rules per hint: +// - ReadOnlyHint: AND across steps — true only when every step declares +// readOnlyHint=true; any step that is nil or false makes it non-read-only. +// - DestructiveHint: OR across steps — true when any step is nil or declares +// destructiveHint=true (an unknown step is tainted conservatively). +// - OpenWorldHint: OR across steps — true when any step is nil or declares +// openWorldHint=true. +// - IdempotentHint: never derived (always nil) — idempotency does not +// compose across multi-step workflows. +func DeriveCompositeAnnotations(stepAnn []*vmcp.ToolAnnotations) *vmcp.ToolAnnotations { + // No tool steps at all: nothing to derive. This is distinct from "≥1 tool + // step exists but none declares annotations", which produces a conservative + // floor below. + if len(stepAnn) == 0 { + return nil + } + + readOnly := true + destructive := false + openWorld := false + for _, ann := range stepAnn { + if ann == nil { + readOnly = false + destructive = true + openWorld = true + continue + } + if ann.ReadOnlyHint == nil || !*ann.ReadOnlyHint { + readOnly = false + } + if ann.DestructiveHint == nil || *ann.DestructiveHint { + destructive = true + } + if ann.OpenWorldHint == nil || *ann.OpenWorldHint { + openWorld = true + } + } + + return &vmcp.ToolAnnotations{ + ReadOnlyHint: &readOnly, + DestructiveHint: &destructive, + OpenWorldHint: &openWorld, + } +} + +// CheckAnnotationContradiction reports whether explicit annotations would make +// a composite tool look SAFER than its derived safety floor allows: +// - readOnlyHint=true while the floor is not true +// - destructiveHint=false while the floor is true +// - openWorldHint=false while the floor is true +// +// idempotentHint never contradicts, and an explicit hint MORE conservative than +// the floor (e.g. readOnlyHint=false when the floor is true) is allowed. +// +// DeriveCompositeAnnotations always populates ReadOnlyHint, DestructiveHint, +// and OpenWorldHint when it returns a non-nil floor, so the nil-hint branches +// below are defensive (they would only matter for a hand-built floor). A nil +// floor (no tool steps) or nil explicit means there is nothing to guard and +// the check returns nil. +func CheckAnnotationContradiction(explicit, floor *vmcp.ToolAnnotations) error { + if explicit == nil || floor == nil { + return nil + } + + if explicit.ReadOnlyHint != nil && *explicit.ReadOnlyHint && + (floor.ReadOnlyHint == nil || !*floor.ReadOnlyHint) { + return fmt.Errorf( + "annotations declare readOnlyHint=true, but not every step tool is read-only; " + + "the composite tool cannot be advertised as read-only") + } + if explicit.DestructiveHint != nil && !*explicit.DestructiveHint && + floor.DestructiveHint != nil && *floor.DestructiveHint { + return fmt.Errorf( + "annotations declare destructiveHint=false, but at least one step tool is destructive or unknown; " + + "the composite tool cannot be advertised as non-destructive") + } + if explicit.OpenWorldHint != nil && !*explicit.OpenWorldHint && + floor.OpenWorldHint != nil && *floor.OpenWorldHint { + return fmt.Errorf( + "annotations declare openWorldHint=false, but at least one step tool is open-world or unknown; " + + "the composite tool cannot be advertised as closed-world") + } + return nil +} + +// MergeAnnotations merges explicit annotations over the derived floor. Per +// hint, an explicit non-nil value wins over the floor; Title behaves the same +// (a non-empty explicit Title wins). Returns nil when both inputs are nil. +// +// The returned value is a deep copy: *bool pointer fields and the Title string +// are copied so the advertised tool does not alias the author's config +// pointers. The inputs are not modified. +func MergeAnnotations(floor, explicit *vmcp.ToolAnnotations) *vmcp.ToolAnnotations { + if floor == nil && explicit == nil { + return nil + } + + merged := &vmcp.ToolAnnotations{} + if floor != nil { + copyBoolPtr(&merged.ReadOnlyHint, floor.ReadOnlyHint) + copyBoolPtr(&merged.DestructiveHint, floor.DestructiveHint) + copyBoolPtr(&merged.IdempotentHint, floor.IdempotentHint) + copyBoolPtr(&merged.OpenWorldHint, floor.OpenWorldHint) + merged.Title = floor.Title + } + if explicit != nil { + if explicit.Title != "" { + merged.Title = explicit.Title + } + copyBoolPtr(&merged.ReadOnlyHint, explicit.ReadOnlyHint) + copyBoolPtr(&merged.DestructiveHint, explicit.DestructiveHint) + copyBoolPtr(&merged.IdempotentHint, explicit.IdempotentHint) + copyBoolPtr(&merged.OpenWorldHint, explicit.OpenWorldHint) + } + return merged +} + +// copyBoolPtr overwrites dst with a fresh copy of src when src is non-nil, so +// dst does not alias src's memory. A nil src leaves dst untouched (preserving +// the prior value, if any). +func copyBoolPtr(dst **bool, src *bool) { + if src == nil { + return + } + v := *src + *dst = &v +} diff --git a/pkg/vmcp/internal/compositetools/workflow_converter.go b/pkg/vmcp/internal/compositetools/workflow_converter.go index 1fb41f075d..a1290d81be 100644 --- a/pkg/vmcp/internal/compositetools/workflow_converter.go +++ b/pkg/vmcp/internal/compositetools/workflow_converter.go @@ -5,11 +5,12 @@ package compositetools import ( "fmt" - "strings" + "log/slog" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/composer" "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/pkg/vmcp/router" ) // FilterWorkflowDefsForSession returns only the workflow definitions whose every @@ -40,11 +41,6 @@ func FilterWorkflowDefsForSession( // references a backend tool that is present in the session routing table. // Returns false if rt is nil and the workflow contains any tool steps, // since a nil routing table means no tools are routable in this session. -// -// Composite tool step names use the convention "{workloadID}.{toolName}" where -// workloadID is a Kubernetes resource name (no dots). The routing table may store -// tools under resolved/prefixed names (e.g. "{workloadID}_echo" with prefix strategy), -// so we look up by BackendTarget.WorkloadID rather than the resolved key directly. func allToolStepsAccessible(def *composer.WorkflowDefinition, rt *vmcp.RoutingTable) bool { for _, step := range def.Steps { if step.Type == composer.StepTypeTool { @@ -71,38 +67,12 @@ func allToolStepsAccessible(def *composer.WorkflowDefinition, rt *vmcp.RoutingTa } // isToolStepAccessible reports whether a composite tool step's tool name can be -// resolved to an accessible backend tool in the given routing table. -// -// Step tool names use the "{workloadID}.{toolName}" convention. Since conflict -// resolution strategies (e.g. prefix) may rename tools in the routing table -// (e.g. "echo" → "yardstick-backend_echo"), we check for accessibility by -// matching on WorkloadID and the original backend capability name rather than -// the resolved routing table key. +// resolved to an accessible backend tool in the given routing table. It is a +// thin nil-safe wrapper around the shared router.ResolveToolRef primitive so +// accessibility filtering and annotation resolution cannot drift. func isToolStepAccessible(stepTool string, rt *vmcp.RoutingTable) bool { - // Fast path: exact match in the routing table. - if _, ok := rt.Tools[stepTool]; ok { - return true - } - - // Parse "{workloadID}.{toolName}" convention. - // Workload IDs are Kubernetes resource names and cannot contain dots, - // so the first dot separates the workload ID from the tool name. - dotIdx := strings.Index(stepTool, ".") - if dotIdx <= 0 { - return false - } - workloadID := stepTool[:dotIdx] - originalName := stepTool[dotIdx+1:] - - for resolvedName, target := range rt.Tools { - if target.WorkloadID != workloadID { - continue - } - if target.GetBackendCapabilityName(resolvedName) == originalName { - return true - } - } - return false + _, ok := router.ResolveToolRef(rt, stepTool) + return ok } // ConvertWorkflowDefsToTools converts workflow definitions to vmcp.Tool format. @@ -116,9 +86,24 @@ func isToolStepAccessible(stepTool string, rt *vmcp.RoutingTable) bool { // - Description: workflow.Description // - InputSchema: workflow.Parameters (JSON Schema format) // - OutputSchema: workflow.Output (JSON Schema format, if defined) +// - Annotations: the safety floor derived from the step tools' annotations, +// merged with the workflow's explicit annotations, if any. When a workflow +// has at least one tool step the floor is always non-nil (fail-closed): +// backends that declare no annotations produce a conservative floor rather +// than no floor. +// +// stepResolver may be nil (or return nil for every step), in which case each +// tool step's annotations are treated as unknown and taint the floor +// conservatively. When a workflow's explicit annotations contradict the derived +// safety floor, the composite tool is DROPPED (not advertised) with a warning +// that names the offending step tool(s) — an explicit declaration must never +// make a tool look safer than its steps allow. // // Returns a slice of vmcp.Tool ready for aggregation and exposure to clients. -func ConvertWorkflowDefsToTools(defs map[string]*composer.WorkflowDefinition) []vmcp.Tool { +func ConvertWorkflowDefsToTools( + defs map[string]*composer.WorkflowDefinition, + stepResolver StepAnnotationResolver, +) []vmcp.Tool { if len(defs) == 0 { return nil // Idiomatic Go: nil slice for empty result } @@ -136,12 +121,120 @@ func ConvertWorkflowDefsToTools(defs map[string]*composer.WorkflowDefinition) [] tool.OutputSchema = buildOutputSchema(def.Output) } + // Derive the safety floor from the step tools, merge explicit + // annotations over it, and drop the tool on a contradiction. + if ann, ok := resolveCompositeAnnotations(def, stepResolver); ok { + tool.Annotations = ann + } else { + continue + } + tools = append(tools, tool) } return tools } +// FilterWorkflowDefsByAnnotations returns only the workflow definitions whose +// explicit annotations do not contradict the derived safety floor. Definitions +// that contradict are omitted (with the same warning resolveCompositeAnnotations +// already emits). This is the CallTool/ListTools shared gate — callers that +// execute composites must use this filtered set so a dropped tool is never callable. +func FilterWorkflowDefsByAnnotations( + defs map[string]*composer.WorkflowDefinition, + stepResolver StepAnnotationResolver, +) map[string]*composer.WorkflowDefinition { + if len(defs) == 0 { + return defs + } + + filtered := make(map[string]*composer.WorkflowDefinition, len(defs)) + for name, def := range defs { + if _, ok := resolveCompositeAnnotations(def, stepResolver); ok { + filtered[name] = def + } + } + return filtered +} + +// resolveCompositeAnnotations computes the advertised annotations for a +// composite tool: the floor derived from the step tools merged with the +// workflow's explicit annotations. It returns ok=false when the explicit +// annotations contradict the floor, in which case the caller must not +// advertise the tool. +// +// stepRefs (the tool references of the steps that produced the floor) is +// threaded into the drop warning so the author can locate the offending step. +func resolveCompositeAnnotations( + def *composer.WorkflowDefinition, + stepResolver StepAnnotationResolver, +) (ann *vmcp.ToolAnnotations, ok bool) { + // No resolver and no explicit annotations: nothing to derive or merge. + // (A workflow whose backends declare no annotations but has explicit + // annotations still runs through Derive + the contradiction guard below.) + if stepResolver == nil && def.Annotations == nil { + return nil, true + } + + stepAnn, stepRefs := resolveStepAnnotations(def, stepResolver) + floor := DeriveCompositeAnnotations(stepAnn) + explicit := def.Annotations.ToAnnotations() + if err := CheckAnnotationContradiction(explicit, floor); err != nil { + slog.Warn("composite tool annotations contradict the safety floor; omitting composite tool", + "tool", def.Name, "step_tools", stepRefs, "error", err) + return nil, false + } + return MergeAnnotations(floor, explicit), true +} + +// resolveStepAnnotations collects the annotations of every tool step in the +// workflow, including forEach inner steps, mirroring the traversal in +// allToolStepsAccessible. It returns the annotations and, in parallel, the +// step tool references that produced them — so a contradiction warning can name +// the offending step(s). A nil resolver (or an unknown step tool) yields a nil +// entry, which the derivation treats conservatively as "unknown". A forEach +// step with a nil InnerStep contributes nothing and never panics. +func resolveStepAnnotations( + def *composer.WorkflowDefinition, + stepResolver StepAnnotationResolver, +) (anns []*vmcp.ToolAnnotations, refs []string) { + resolve := func(stepTool string) { + var ann *vmcp.ToolAnnotations + if stepResolver != nil { + ann = stepResolver(stepTool) + } + anns = append(anns, ann) + refs = append(refs, stepTool) + } + for i := range def.Steps { + step := &def.Steps[i] + if step.Type == composer.StepTypeTool { + resolve(step.Tool) + } + // A forEach step whose InnerStep is nil is structurally invalid (caught + // earlier by validation); guard against a nil deref here regardless. + if step.Type == composer.StepTypeForEach && step.InnerStep != nil && + step.InnerStep.Type == composer.StepTypeTool { + resolve(step.InnerStep.Tool) + } + } + return anns, refs +} + +// CompositeToolNames returns the names of the composite tools in defs. +// Order is undefined (map iteration). Used for name-conflict detection so +// annotation policy cannot drop a colliding name from the conflict check. +func CompositeToolNames(defs map[string]*composer.WorkflowDefinition) []string { + if len(defs) == 0 { + return nil + } + names := make([]string, 0, len(defs)) + for name := range defs { + names = append(names, name) + } + return names +} + // ValidateNoToolConflicts validates that composite tool names don't conflict with backend tool names. // // Tool name conflicts would cause ambiguity in routing/execution: @@ -150,7 +243,10 @@ func ConvertWorkflowDefsToTools(defs map[string]*composer.WorkflowDefinition) [] // // This validation ensures clear separation and prevents runtime confusion. // Returns an error listing all conflicting tool names if any conflicts are found. -func ValidateNoToolConflicts(backendTools, compositeTools []vmcp.Tool) error { +// +// Prefer CompositeToolNames(defs) for the compositeNames argument so conflict +// detection never depends on annotation conversion (which can drop tools). +func ValidateNoToolConflicts(backendTools []vmcp.Tool, compositeNames []string) error { // Build set of backend tool names for O(1) lookups backendNames := make(map[string]bool, len(backendTools)) for _, tool := range backendTools { @@ -159,9 +255,9 @@ func ValidateNoToolConflicts(backendTools, compositeTools []vmcp.Tool) error { // Check for conflicts var conflicts []string - for _, compTool := range compositeTools { - if backendNames[compTool.Name] { - conflicts = append(conflicts, compTool.Name) + for _, name := range compositeNames { + if backendNames[name] { + conflicts = append(conflicts, name) } } diff --git a/pkg/vmcp/internal/compositetools/workflow_converter_test.go b/pkg/vmcp/internal/compositetools/workflow_converter_test.go index 795a448c52..3fd64d18e3 100644 --- a/pkg/vmcp/internal/compositetools/workflow_converter_test.go +++ b/pkg/vmcp/internal/compositetools/workflow_converter_test.go @@ -4,9 +4,12 @@ package compositetools import ( + "errors" "testing" "github.com/google/go-cmp/cmp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/composer" @@ -348,7 +351,7 @@ func TestConvertWorkflowDefsToToolsWithOutputSchema(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - tools := ConvertWorkflowDefsToTools(tt.defs) + tools := ConvertWorkflowDefsToTools(tt.defs, nil) if len(tools) != tt.want { t.Errorf("ConvertWorkflowDefsToTools() returned %d tools, want %d", len(tools), tt.want) @@ -536,6 +539,692 @@ func TestFilterWorkflowDefsForSession(t *testing.T) { } } +func TestDeriveCompositeAnnotations(t *testing.T) { + t.Parallel() + + trueVal, falseVal := true, false + + tests := []struct { + name string + stepAnn []*vmcp.ToolAnnotations + want *vmcp.ToolAnnotations + }{ + { + name: "empty input (no tool steps) yields nil floor", + stepAnn: nil, + want: nil, + }, + { + // Fail-closed (issue #6192): ≥1 tool step exists but none declares + // annotations → conservative floor, NOT nil. An explicit readOnlyHint:true + // against this floor is dropped by CheckAnnotationContradiction. + name: "all nil annotations (tool steps exist, none declare) yields conservative floor", + stepAnn: []*vmcp.ToolAnnotations{nil, nil}, + want: &vmcp.ToolAnnotations{ + ReadOnlyHint: &falseVal, + DestructiveHint: &trueVal, + OpenWorldHint: &trueVal, + }, + }, + { + name: "all steps read-only", + stepAnn: []*vmcp.ToolAnnotations{ + {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &trueVal}, + }, + want: &vmcp.ToolAnnotations{ + ReadOnlyHint: &trueVal, + DestructiveHint: &falseVal, + OpenWorldHint: &trueVal, + }, + }, + { + name: "one non-read-only step makes floor not read-only", + stepAnn: []*vmcp.ToolAnnotations{ + {ReadOnlyHint: &trueVal}, + {ReadOnlyHint: &falseVal}, + }, + want: &vmcp.ToolAnnotations{ + ReadOnlyHint: &falseVal, + DestructiveHint: &trueVal, // nil hints taint conservatively + OpenWorldHint: &trueVal, + }, + }, + { + name: "nil step annotations taint destructive and open-world", + stepAnn: []*vmcp.ToolAnnotations{ + {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + nil, + }, + want: &vmcp.ToolAnnotations{ + ReadOnlyHint: &falseVal, + DestructiveHint: &trueVal, + OpenWorldHint: &trueVal, + }, + }, + { + name: "destructive OR across steps", + stepAnn: []*vmcp.ToolAnnotations{ + {ReadOnlyHint: &falseVal, DestructiveHint: &trueVal, OpenWorldHint: &falseVal}, + {ReadOnlyHint: &falseVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }, + want: &vmcp.ToolAnnotations{ + ReadOnlyHint: &falseVal, + DestructiveHint: &trueVal, + OpenWorldHint: &falseVal, + }, + }, + { + name: "idempotent hint is never derived", + stepAnn: []*vmcp.ToolAnnotations{ + {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal, IdempotentHint: &trueVal}, + }, + want: &vmcp.ToolAnnotations{ + ReadOnlyHint: &trueVal, + DestructiveHint: &falseVal, + OpenWorldHint: &falseVal, + IdempotentHint: nil, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := DeriveCompositeAnnotations(tt.stepAnn) + + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("DeriveCompositeAnnotations() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestCheckAnnotationContradiction(t *testing.T) { + t.Parallel() + + trueVal, falseVal := true, false + + tests := []struct { + name string + explicit *vmcp.ToolAnnotations + floor *vmcp.ToolAnnotations + wantErr bool + }{ + { + name: "nil explicit", + explicit: nil, + floor: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + wantErr: false, + }, + { + name: "nil floor", + explicit: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + floor: nil, + wantErr: false, + }, + { + name: "readOnly true against non-read-only floor", + explicit: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + floor: &vmcp.ToolAnnotations{ReadOnlyHint: &falseVal}, + wantErr: true, + }, + { + name: "readOnly true against unknown read-only floor", + explicit: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + floor: &vmcp.ToolAnnotations{ReadOnlyHint: nil}, + wantErr: true, + }, + { + name: "destructive false against destructive floor", + explicit: &vmcp.ToolAnnotations{DestructiveHint: &falseVal}, + floor: &vmcp.ToolAnnotations{DestructiveHint: &trueVal}, + wantErr: true, + }, + { + name: "openWorld false against open-world floor", + explicit: &vmcp.ToolAnnotations{OpenWorldHint: &falseVal}, + floor: &vmcp.ToolAnnotations{OpenWorldHint: &trueVal}, + wantErr: true, + }, + { + name: "equal annotations", + explicit: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal, DestructiveHint: &falseVal}, + floor: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal, DestructiveHint: &falseVal}, + wantErr: false, + }, + { + name: "more conservative readOnly false when floor is true", + explicit: &vmcp.ToolAnnotations{ReadOnlyHint: &falseVal}, + floor: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + wantErr: false, + }, + { + name: "more conservative destructive true when floor is false", + explicit: &vmcp.ToolAnnotations{DestructiveHint: &trueVal}, + floor: &vmcp.ToolAnnotations{DestructiveHint: &falseVal}, + wantErr: false, + }, + { + name: "idempotent never contradicts", + explicit: &vmcp.ToolAnnotations{IdempotentHint: &trueVal}, + floor: &vmcp.ToolAnnotations{IdempotentHint: &falseVal}, + wantErr: false, + }, + { + // Q1: explicit destructiveHint:false against a floor whose destructiveHint + // is nil (unknown) is allowed — the floor does not assert the tool is + // destructive, so the explicit claim does not contradict it. + name: "destructiveHint:false against nil-floor destructiveHint (nil-floor-hint guard)", + explicit: &vmcp.ToolAnnotations{DestructiveHint: &falseVal}, + floor: &vmcp.ToolAnnotations{DestructiveHint: nil}, + wantErr: false, + }, + { + // Q2: explicit openWorldHint:false against a floor whose openWorldHint is + // nil (unknown) is allowed for the same reason. + name: "openWorldHint:false against nil-floor openWorldHint (nil-floor-hint guard)", + explicit: &vmcp.ToolAnnotations{OpenWorldHint: &falseVal}, + floor: &vmcp.ToolAnnotations{OpenWorldHint: nil}, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := CheckAnnotationContradiction(tt.explicit, tt.floor) + + if (err != nil) != tt.wantErr { + t.Errorf("CheckAnnotationContradiction() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestMergeAnnotations(t *testing.T) { + t.Parallel() + + trueVal, falseVal := true, false + + tests := []struct { + name string + floor *vmcp.ToolAnnotations + explicit *vmcp.ToolAnnotations + want *vmcp.ToolAnnotations + }{ + { + name: "both nil", + floor: nil, + explicit: nil, + want: nil, + }, + { + name: "floor only", + floor: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + explicit: nil, + want: &vmcp.ToolAnnotations{ReadOnlyHint: &trueVal}, + }, + { + name: "explicit only", + floor: nil, + explicit: &vmcp.ToolAnnotations{Title: "T", IdempotentHint: &trueVal}, + want: &vmcp.ToolAnnotations{Title: "T", IdempotentHint: &trueVal}, + }, + { + name: "explicit non-nil fields win per hint", + floor: &vmcp.ToolAnnotations{ + Title: "Floor", + ReadOnlyHint: &trueVal, + DestructiveHint: &trueVal, + OpenWorldHint: &trueVal, + }, + explicit: &vmcp.ToolAnnotations{ + Title: "Explicit", + ReadOnlyHint: &falseVal, + IdempotentHint: &trueVal, + DestructiveHint: nil, // nil keeps floor + }, + want: &vmcp.ToolAnnotations{ + Title: "Explicit", + ReadOnlyHint: &falseVal, + DestructiveHint: &trueVal, + IdempotentHint: &trueVal, + OpenWorldHint: &trueVal, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := MergeAnnotations(tt.floor, tt.explicit) + + if diff := cmp.Diff(tt.want, got); diff != "" { + t.Errorf("MergeAnnotations() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestConvertWorkflowDefsToToolsAnnotations(t *testing.T) { + t.Parallel() + + trueVal, falseVal := true, false + + resolver := func(ann map[string]*vmcp.ToolAnnotations) StepAnnotationResolver { + return func(stepTool string) *vmcp.ToolAnnotations { return ann[stepTool] } + } + + tests := []struct { + name string + defs map[string]*composer.WorkflowDefinition + stepResolver StepAnnotationResolver + want map[string]*vmcp.ToolAnnotations // tool name -> expected annotations (absent = dropped) + }{ + { + name: "derives annotations from step tools", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + {ID: "s2", Type: composer.StepTypeTool, Tool: "backend.list"}, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + "backend.list": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }), + want: map[string]*vmcp.ToolAnnotations{ + "wf": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }, + }, + { + name: "forEach inner step contributes to derivation", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + { + ID: "s2", + Type: composer.StepTypeForEach, + InnerStep: &composer.WorkflowStep{ + ID: "inner", Type: composer.StepTypeTool, Tool: "backend.write", + }, + }, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + "backend.write": {ReadOnlyHint: &falseVal, DestructiveHint: &trueVal, OpenWorldHint: &trueVal}, + }), + want: map[string]*vmcp.ToolAnnotations{ + "wf": {ReadOnlyHint: &falseVal, DestructiveHint: &trueVal, OpenWorldHint: &trueVal}, + }, + }, + { + name: "explicit annotations merge over the derived floor", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + }, + Annotations: &config.ToolAnnotationsOverride{ + Title: ptr("My Tool"), + IdempotentHint: &trueVal, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }), + want: map[string]*vmcp.ToolAnnotations{ + "wf": { + Title: "My Tool", + ReadOnlyHint: &trueVal, + DestructiveHint: &falseVal, + OpenWorldHint: &falseVal, + IdempotentHint: &trueVal, + }, + }, + }, + { + name: "contradicting tool is dropped, others kept", + defs: map[string]*composer.WorkflowDefinition{ + "wf_bad": { + Name: "wf_bad", + Description: "contradicting workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.write"}, + }, + Annotations: &config.ToolAnnotationsOverride{ReadOnlyHint: &trueVal}, + }, + "wf_ok": { + Name: "wf_ok", + Description: "valid workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + "backend.write": {ReadOnlyHint: &falseVal, DestructiveHint: &trueVal, OpenWorldHint: &trueVal}, + }), + want: map[string]*vmcp.ToolAnnotations{ + "wf_ok": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }, + }, + { + name: "unknown step tool taints the floor", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + {ID: "s2", Type: composer.StepTypeTool, Tool: "backend.unknown"}, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }), + want: map[string]*vmcp.ToolAnnotations{ + "wf": {ReadOnlyHint: &falseVal, DestructiveHint: &trueVal, OpenWorldHint: &trueVal}, + }, + }, + { + name: "nil resolver with no explicit annotations yields no annotations", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + }, + }, + }, + stepResolver: nil, + want: map[string]*vmcp.ToolAnnotations{"wf": nil}, + }, + { + // Q4: with a nil stepResolver the floor cannot be derived (all steps are + // unknown), but an explicit annotation that does not tighten safety + // against the conservative floor still passes through. + name: "nil resolver with explicit annotations passes explicit through", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + }, + Annotations: &config.ToolAnnotationsOverride{ + Title: ptr("Explicit Tool"), + IdempotentHint: &trueVal, + }, + }, + }, + stepResolver: nil, + // Conservative floor (readOnly=false, destructive=true, openWorld=true) + // merged with the explicit Title + IdempotentHint. + want: map[string]*vmcp.ToolAnnotations{ + "wf": { + Title: "Explicit Tool", + ReadOnlyHint: &falseVal, + DestructiveHint: &trueVal, + OpenWorldHint: &trueVal, + IdempotentHint: &trueVal, + }, + }, + }, + { + // Q3: a forEach step with a nil InnerStep is structurally invalid but must + // be skipped (no panic) during annotation resolution. + name: "forEach with nil InnerStep is skipped without panic", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Description: "workflow", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + {ID: "s2", Type: composer.StepTypeForEach, InnerStep: nil}, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }), + want: map[string]*vmcp.ToolAnnotations{ + "wf": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tools := ConvertWorkflowDefsToTools(tt.defs, tt.stepResolver) + + if len(tools) != len(tt.want) { + t.Fatalf("ConvertWorkflowDefsToTools() returned %d tools, want %d", len(tools), len(tt.want)) + } + for _, tool := range tools { + wantAnn, ok := tt.want[tool.Name] + if !ok { + t.Errorf("unexpected tool %q in result", tool.Name) + continue + } + if diff := cmp.Diff(wantAnn, tool.Annotations); diff != "" { + t.Errorf("tool %q annotations mismatch (-want +got):\n%s", tool.Name, diff) + } + } + }) + } +} + +func ptr[T any](v T) *T { return &v } + +func TestCompositeToolNames(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + defs map[string]*composer.WorkflowDefinition + want []string + }{ + { + name: "nil map", + defs: nil, + want: nil, + }, + { + name: "empty map", + defs: map[string]*composer.WorkflowDefinition{}, + want: nil, + }, + { + name: "returns all definition names", + defs: map[string]*composer.WorkflowDefinition{ + "wf_a": {Name: "wf_a"}, + "wf_b": {Name: "wf_b"}, + "wf_c": {Name: "wf_c"}, + }, + want: []string{"wf_a", "wf_b", "wf_c"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := CompositeToolNames(tt.defs) + require.ElementsMatch(t, tt.want, got) + }) + } +} + +func TestValidateNoToolConflicts(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + backendTools []vmcp.Tool + compositeNames []string + wantConflict bool + }{ + { + name: "no conflict", + backendTools: []vmcp.Tool{{Name: "be1.echo"}, {Name: "be1.list"}}, + compositeNames: []string{"wf", "deploy"}, + wantConflict: false, + }, + { + name: "single name conflict", + backendTools: []vmcp.Tool{{Name: "be1.echo"}, {Name: "shared"}}, + compositeNames: []string{"wf", "shared"}, + wantConflict: true, + }, + { + name: "empty composites", + backendTools: []vmcp.Tool{{Name: "be1.echo"}}, + compositeNames: nil, + wantConflict: false, + }, + { + name: "empty backends", + backendTools: nil, + compositeNames: []string{"wf"}, + wantConflict: false, + }, + { + // Name-only check: even an optimistic/contradicting composite name is + // still visible to conflict detection (bug-2 regression surface). + name: "conflict with optimistic annotation name still detected", + backendTools: []vmcp.Tool{{Name: "be1.echo"}}, + compositeNames: []string{"be1.echo"}, + wantConflict: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := ValidateNoToolConflicts(tt.backendTools, tt.compositeNames) + if tt.wantConflict { + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrToolNameConflict), + "expected ErrToolNameConflict, got %v", err) + return + } + require.NoError(t, err) + }) + } +} + +func TestFilterWorkflowDefsByAnnotations(t *testing.T) { + t.Parallel() + + trueVal, falseVal := true, false + resolver := func(ann map[string]*vmcp.ToolAnnotations) StepAnnotationResolver { + return func(stepTool string) *vmcp.ToolAnnotations { return ann[stepTool] } + } + + tests := []struct { + name string + defs map[string]*composer.WorkflowDefinition + stepResolver StepAnnotationResolver + wantNames []string + }{ + { + name: "empty defs", + defs: map[string]*composer.WorkflowDefinition{}, + wantNames: nil, + }, + { + name: "keeps non-contradicting, drops contradicting", + defs: map[string]*composer.WorkflowDefinition{ + "wf_bad": { + Name: "wf_bad", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.write"}, + }, + Annotations: &config.ToolAnnotationsOverride{ReadOnlyHint: &trueVal}, + }, + "wf_ok": { + Name: "wf_ok", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + "backend.write": {ReadOnlyHint: &falseVal, DestructiveHint: &trueVal, OpenWorldHint: &trueVal}, + }), + wantNames: []string{"wf_ok"}, + }, + { + // Silent backend (nil step annotations) → fail-closed floor; explicit + // readOnlyHint:true contradicts and the def is dropped. + name: "drops optimistic readOnly against silent backend floor", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.echo"}, + }, + Annotations: &config.ToolAnnotationsOverride{ReadOnlyHint: &trueVal}, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.echo": nil, + }), + wantNames: []string{}, + }, + { + name: "keeps workflow with no explicit annotations", + defs: map[string]*composer.WorkflowDefinition{ + "wf": { + Name: "wf", + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: "backend.read"}, + }, + }, + }, + stepResolver: resolver(map[string]*vmcp.ToolAnnotations{ + "backend.read": {ReadOnlyHint: &trueVal, DestructiveHint: &falseVal, OpenWorldHint: &falseVal}, + }), + wantNames: []string{"wf"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := FilterWorkflowDefsByAnnotations(tt.defs, tt.stepResolver) + gotNames := make([]string, 0, len(got)) + for name := range got { + gotNames = append(gotNames, name) + } + require.ElementsMatch(t, tt.wantNames, gotNames) + }) + } +} + func TestBuildOutputPropertySchema(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/internal/pagination/pagination.go b/pkg/vmcp/internal/pagination/pagination.go new file mode 100644 index 0000000000..95e25c9704 --- /dev/null +++ b/pkg/vmcp/internal/pagination/pagination.go @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package pagination provides a helper for following MCP list-pagination +// cursors when querying backends. +// +// MCP list operations (tools/list, resources/list, prompts/list) are +// paginated: a result may carry a NextCursor pointing at the next page, and a +// client that issues a single request without following the cursor silently +// drops every item beyond the first page. Backends served by the +// toolhive-core mcpcompat server paginate at DefaultPageSize=1000, so any +// backend advertising more than 1000 tools/resources/prompts would lose the +// tail without cursor following. +// +// This package lives under pkg/vmcp/internal so both pkg/vmcp/client and +// pkg/vmcp/session can depend on it without introducing an import cycle +// (pkg/vmcp/client already imports pkg/vmcp/session). +package pagination + +import ( + "context" + "fmt" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// maxIterations caps the number of pages ListAll will fetch. It is a generous +// upper bound (well above the ~O(items/pageSize) pages any legitimate backend +// needs) that exists solely to prevent an infinite loop if a backend keeps +// returning a non-empty cursor forever. +const maxIterations = 10000 + +// ListAll repeatedly invokes fetch, following the MCP pagination cursor, until +// the backend returns an empty next cursor. It accumulates and returns every +// item across all pages. +// +// fetch is called with the cursor for the page to retrieve (empty on the first +// call) and must return that page's items, the cursor for the next page (empty +// when no more pages remain), and any error. A fetch error is returned to the +// caller immediately. +// +// ListAll guards against a misbehaving backend in two ways: it returns an error +// if the same non-empty cursor is seen twice (a cycle), and it caps the total +// number of pages at maxIterations. +func ListAll[T any]( + ctx context.Context, + fetch func(ctx context.Context, cursor mcp.Cursor) (items []T, next mcp.Cursor, err error), +) ([]T, error) { + var all []T + var cursor mcp.Cursor + seen := make(map[mcp.Cursor]struct{}) + + for i := 0; i < maxIterations; i++ { + items, next, err := fetch(ctx, cursor) + if err != nil { + return nil, err + } + all = append(all, items...) + + // An empty next cursor signals the last page. + if next == "" { + return all, nil + } + + // A repeated cursor means the backend is not advancing; abort rather + // than loop forever. + if _, ok := seen[next]; ok { + return nil, fmt.Errorf("pagination cursor %q repeated; backend is not advancing pages", next) + } + seen[next] = struct{}{} + cursor = next + } + + return nil, fmt.Errorf("pagination exceeded %d pages; aborting to avoid an unbounded loop", maxIterations) +} diff --git a/pkg/vmcp/internal/pagination/pagination_test.go b/pkg/vmcp/internal/pagination/pagination_test.go new file mode 100644 index 0000000000..757aa33e42 --- /dev/null +++ b/pkg/vmcp/internal/pagination/pagination_test.go @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package pagination + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +func TestListAll(t *testing.T) { + t.Parallel() + + errFetch := errors.New("boom") + + tests := []struct { + name string + // pages maps the requested cursor to the page it returns. + pages map[mcp.Cursor]page + fetchErr error + wantItems []int + wantErr bool + }{ + { + name: "single page with empty next cursor terminates after one call", + pages: map[mcp.Cursor]page{ + "": {items: []int{1, 2, 3}, next: ""}, + }, + wantItems: []int{1, 2, 3}, + }, + { + name: "multiple pages accumulate in order", + pages: map[mcp.Cursor]page{ + "": {items: []int{1, 2}, next: "p1"}, + "p1": {items: []int{3, 4}, next: "p2"}, + "p2": {items: []int{5}, next: ""}, + }, + wantItems: []int{1, 2, 3, 4, 5}, + }, + { + name: "empty result returns no items", + pages: map[mcp.Cursor]page{ + "": {items: nil, next: ""}, + }, + wantItems: nil, + }, + { + name: "repeated cursor aborts with error", + pages: map[mcp.Cursor]page{ + "": {items: []int{1}, next: "loop"}, + "loop": {items: []int{2}, next: "loop"}, + }, + wantErr: true, + }, + { + name: "fetch error is propagated", + fetchErr: errFetch, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + fetch := func(_ context.Context, cursor mcp.Cursor) ([]int, mcp.Cursor, error) { + if tt.fetchErr != nil { + return nil, "", tt.fetchErr + } + p, ok := tt.pages[cursor] + require.True(t, ok, "unexpected cursor requested: %q", cursor) + return p.items, p.next, nil + } + + got, err := ListAll(t.Context(), fetch) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.wantItems, got) + }) + } +} + +// page is a canned pagination response used by the table test. +type page struct { + items []int + next mcp.Cursor +} diff --git a/pkg/vmcp/k8s/manager.go b/pkg/vmcp/k8s/manager.go index 97c3df1bda..881b34d72d 100644 --- a/pkg/vmcp/k8s/manager.go +++ b/pkg/vmcp/k8s/manager.go @@ -309,7 +309,7 @@ func (w *BackendWatcher) addBackendWatchController(ctx context.Context) error { } // Register reconciler with manager - // This sets up watches on MCPServer, MCPRemoteProxy, MCPServerEntry, ExternalAuthConfig, and ConfigMap. + // This sets up watches on MCPServer, MCPRemoteProxy, and ExternalAuthConfig if err := reconciler.SetupWithManager(w.ctrlManager); err != nil { return fmt.Errorf("failed to setup backend reconciler: %w", err) } diff --git a/pkg/vmcp/mocks/mock_backend_client.go b/pkg/vmcp/mocks/mock_backend_client.go index 1a187df3f8..eb3e2cf8bf 100644 --- a/pkg/vmcp/mocks/mock_backend_client.go +++ b/pkg/vmcp/mocks/mock_backend_client.go @@ -13,6 +13,7 @@ import ( context "context" reflect "reflect" + mcp "github.com/stacklok/toolhive/pkg/mcp" vmcp "github.com/stacklok/toolhive/pkg/vmcp" gomock "go.uber.org/mock/gomock" ) @@ -81,18 +82,33 @@ func (m *MockBackendClient) EXPECT() *MockBackendClientMockRecorder { } // CallTool mocks base method. -func (m *MockBackendClient) CallTool(ctx context.Context, target *vmcp.BackendTarget, toolName string, arguments, meta map[string]any) (*vmcp.ToolCallResult, error) { +func (m *MockBackendClient) CallTool(ctx context.Context, target *vmcp.BackendTarget, toolName string, arguments, meta map[string]any, paramHeaders map[string]string) (*vmcp.ToolCallResult, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "CallTool", ctx, target, toolName, arguments, meta) + ret := m.ctrl.Call(m, "CallTool", ctx, target, toolName, arguments, meta, paramHeaders) ret0, _ := ret[0].(*vmcp.ToolCallResult) ret1, _ := ret[1].(error) return ret0, ret1 } // CallTool indicates an expected call of CallTool. -func (mr *MockBackendClientMockRecorder) CallTool(ctx, target, toolName, arguments, meta any) *gomock.Call { +func (mr *MockBackendClientMockRecorder) CallTool(ctx, target, toolName, arguments, meta, paramHeaders any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CallTool", reflect.TypeOf((*MockBackendClient)(nil).CallTool), ctx, target, toolName, arguments, meta) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CallTool", reflect.TypeOf((*MockBackendClient)(nil).CallTool), ctx, target, toolName, arguments, meta, paramHeaders) +} + +// Complete mocks base method. +func (m *MockBackendClient) Complete(ctx context.Context, target *vmcp.BackendTarget, ref vmcp.CompletionRef, argName, argValue string, contextArgs map[string]string) (*vmcp.CompletionResult, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Complete", ctx, target, ref, argName, argValue, contextArgs) + ret0, _ := ret[0].(*vmcp.CompletionResult) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Complete indicates an expected call of Complete. +func (mr *MockBackendClientMockRecorder) Complete(ctx, target, ref, argName, argValue, contextArgs any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Complete", reflect.TypeOf((*MockBackendClient)(nil).Complete), ctx, target, ref, argName, argValue, contextArgs) } // GetPrompt mocks base method. @@ -139,3 +155,42 @@ func (mr *MockBackendClientMockRecorder) ReadResource(ctx, target, uri any) *gom mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadResource", reflect.TypeOf((*MockBackendClient)(nil).ReadResource), ctx, target, uri) } + +// MockRevisionReporter is a mock of RevisionReporter interface. +type MockRevisionReporter struct { + ctrl *gomock.Controller + recorder *MockRevisionReporterMockRecorder + isgomock struct{} +} + +// MockRevisionReporterMockRecorder is the mock recorder for MockRevisionReporter. +type MockRevisionReporterMockRecorder struct { + mock *MockRevisionReporter +} + +// NewMockRevisionReporter creates a new mock instance. +func NewMockRevisionReporter(ctrl *gomock.Controller) *MockRevisionReporter { + mock := &MockRevisionReporter{ctrl: ctrl} + mock.recorder = &MockRevisionReporterMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRevisionReporter) EXPECT() *MockRevisionReporterMockRecorder { + return m.recorder +} + +// CachedRevision mocks base method. +func (m *MockRevisionReporter) CachedRevision(workloadID string) (mcp.Revision, bool) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CachedRevision", workloadID) + ret0, _ := ret[0].(mcp.Revision) + ret1, _ := ret[1].(bool) + return ret0, ret1 +} + +// CachedRevision indicates an expected call of CachedRevision. +func (mr *MockRevisionReporterMockRecorder) CachedRevision(workloadID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CachedRevision", reflect.TypeOf((*MockRevisionReporter)(nil).CachedRevision), workloadID) +} diff --git a/pkg/vmcp/mocks/mock_forwarding.go b/pkg/vmcp/mocks/mock_forwarding.go new file mode 100644 index 0000000000..726bf2d1b7 --- /dev/null +++ b/pkg/vmcp/mocks/mock_forwarding.go @@ -0,0 +1,145 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: forwarding.go +// +// Generated by this command: +// +// mockgen -destination=mocks/mock_forwarding.go -package=mocks -source=forwarding.go SamplingRequester ClientNotifier +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + reflect "reflect" + + vmcp "github.com/stacklok/toolhive/pkg/vmcp" + gomock "go.uber.org/mock/gomock" +) + +// MockSamplingRequester is a mock of SamplingRequester interface. +type MockSamplingRequester struct { + ctrl *gomock.Controller + recorder *MockSamplingRequesterMockRecorder + isgomock struct{} +} + +// MockSamplingRequesterMockRecorder is the mock recorder for MockSamplingRequester. +type MockSamplingRequesterMockRecorder struct { + mock *MockSamplingRequester +} + +// NewMockSamplingRequester creates a new mock instance. +func NewMockSamplingRequester(ctrl *gomock.Controller) *MockSamplingRequester { + mock := &MockSamplingRequester{ctrl: ctrl} + mock.recorder = &MockSamplingRequesterMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSamplingRequester) EXPECT() *MockSamplingRequesterMockRecorder { + return m.recorder +} + +// RequestSampling mocks base method. +func (m *MockSamplingRequester) RequestSampling(ctx context.Context, req vmcp.SamplingRequest) (*vmcp.SamplingResult, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RequestSampling", ctx, req) + ret0, _ := ret[0].(*vmcp.SamplingResult) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RequestSampling indicates an expected call of RequestSampling. +func (mr *MockSamplingRequesterMockRecorder) RequestSampling(ctx, req any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RequestSampling", reflect.TypeOf((*MockSamplingRequester)(nil).RequestSampling), ctx, req) +} + +// MockClientNotifier is a mock of ClientNotifier interface. +type MockClientNotifier struct { + ctrl *gomock.Controller + recorder *MockClientNotifierMockRecorder + isgomock struct{} +} + +// MockClientNotifierMockRecorder is the mock recorder for MockClientNotifier. +type MockClientNotifierMockRecorder struct { + mock *MockClientNotifier +} + +// NewMockClientNotifier creates a new mock instance. +func NewMockClientNotifier(ctrl *gomock.Controller) *MockClientNotifier { + mock := &MockClientNotifier{ctrl: ctrl} + mock.recorder = &MockClientNotifierMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockClientNotifier) EXPECT() *MockClientNotifierMockRecorder { + return m.recorder +} + +// NotifyLog mocks base method. +func (m *MockClientNotifier) NotifyLog(ctx context.Context, n vmcp.LogMessage) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "NotifyLog", ctx, n) + ret0, _ := ret[0].(error) + return ret0 +} + +// NotifyLog indicates an expected call of NotifyLog. +func (mr *MockClientNotifierMockRecorder) NotifyLog(ctx, n any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NotifyLog", reflect.TypeOf((*MockClientNotifier)(nil).NotifyLog), ctx, n) +} + +// NotifyProgress mocks base method. +func (m *MockClientNotifier) NotifyProgress(ctx context.Context, n vmcp.ProgressNotification) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "NotifyProgress", ctx, n) + ret0, _ := ret[0].(error) + return ret0 +} + +// NotifyProgress indicates an expected call of NotifyProgress. +func (mr *MockClientNotifierMockRecorder) NotifyProgress(ctx, n any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "NotifyProgress", reflect.TypeOf((*MockClientNotifier)(nil).NotifyProgress), ctx, n) +} + +// MockClientForwarderBinder is a mock of ClientForwarderBinder interface. +type MockClientForwarderBinder struct { + ctrl *gomock.Controller + recorder *MockClientForwarderBinderMockRecorder + isgomock struct{} +} + +// MockClientForwarderBinderMockRecorder is the mock recorder for MockClientForwarderBinder. +type MockClientForwarderBinderMockRecorder struct { + mock *MockClientForwarderBinder +} + +// NewMockClientForwarderBinder creates a new mock instance. +func NewMockClientForwarderBinder(ctrl *gomock.Controller) *MockClientForwarderBinder { + mock := &MockClientForwarderBinder{ctrl: ctrl} + mock.recorder = &MockClientForwarderBinderMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockClientForwarderBinder) EXPECT() *MockClientForwarderBinderMockRecorder { + return m.recorder +} + +// BindForwarders mocks base method. +func (m *MockClientForwarderBinder) BindForwarders(elicitation vmcp.ElicitationRequester, sampling vmcp.SamplingRequester, notifier vmcp.ClientNotifier) { + m.ctrl.T.Helper() + m.ctrl.Call(m, "BindForwarders", elicitation, sampling, notifier) +} + +// BindForwarders indicates an expected call of BindForwarders. +func (mr *MockClientForwarderBinderMockRecorder) BindForwarders(elicitation, sampling, notifier any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BindForwarders", reflect.TypeOf((*MockClientForwarderBinder)(nil).BindForwarders), elicitation, sampling, notifier) +} diff --git a/pkg/vmcp/mrtr.go b/pkg/vmcp/mrtr.go new file mode 100644 index 0000000000..231f2416d2 --- /dev/null +++ b/pkg/vmcp/mrtr.go @@ -0,0 +1,143 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp + +import ( + "encoding/json" + "errors" + "fmt" +) + +// InputRequiredResult carries a Modern (2026-07-28) backend's +// resultType:"input_required" envelope — a Multi Round-Trip Request round +// (SEP-2322): the backend cannot complete the call until the caller fulfills +// InputRequests and retries the original request echoing RequestState. +// +// This is the domain-typed egress surface of MRTR (the seam the client-edge +// limitation in docs/arch/10-virtual-mcp-architecture.md names): the backend +// client decodes the envelope into this type and surfaces it through +// InputRequiredError, so upper layers can relay a round to a Modern +// downstream client or fulfill it in-request for a Legacy one. Full flow +// design in docs/arch/16-vmcp-mrtr.md. +type InputRequiredResult struct { + // InputRequests maps the backend's server-assigned keys to raw request + // objects (ElicitRequest, CreateMessageRequest, or ListRootsRequest — + // schema/draft's InputRequest union). Values are deliberately opaque + // json.RawMessage: on the pass-through path vMCP relays them verbatim and + // must not reinterpret them. Use Methods to inspect only the wire method + // names (e.g. for capability gating). Nil on a load-shedding round + // (RequestState only). + InputRequests map[string]json.RawMessage + + // RequestState is the backend's opaque round-trip state. Per SEP-2322 the + // client "MUST echo back the exact value" on retry — vMCP relays it + // verbatim and never interprets it. Nil when the backend's envelope had no + // requestState field, which the retry must preserve: client requirement 2 + // says that if the result does not contain the field, the client "MUST NOT + // include one in the retry". The schema puts no minimum length on the + // field, so present-and-empty ("") is legal, distinct from absent, and + // must be echoed as the empty string. + RequestState *string +} + +// InputRequiredError is the typed error that carries an input_required round +// through the backend client's error return, so every existing +// complete-result path stays untouched. Result is non-nil only for a valid, +// drivable round (see the extraction rules on InputRequiredFromError); for +// any other non-"complete" envelope only the sentinel classification remains. +// +// Sentinel is the classification error this unwraps to — the backend client +// sets its errModernInputRequired so every existing errors.Is check keeps +// working unchanged. The fields are exported so tests (e.g. against +// mocks.MockBackendClient) can construct the error without going through the +// real client's envelope decode. +type InputRequiredError struct { + // ResultType is the envelope's raw resultType. Backend-controlled text: + // Error() truncates it, and it must never be interpolated elsewhere + // unbounded. + ResultType string + + // Result is the decoded round, nil unless the envelope was a valid + // input_required round on a method that may carry one. + Result *InputRequiredResult + + // Sentinel is the error Unwrap returns; required. It leads the Error() + // message, so the rendered text stays byte-identical to a plain + // fmt.Errorf("%w: ...", Sentinel, ...) wrapping. Because the field is + // exported for direct construction, nothing enforces "required" at compile + // time: a literal that omits it unwraps to nil and renders + // missingSentinelText in place of the classification. + Sentinel error +} + +// maxResultTypeInError bounds how much of the backend-controlled resultType +// Error() interpolates. Real resultTypes are short tokens ("complete", +// "input_required", "task"); anything longer is hostile or broken, and this +// error's text reaches the downstream client via writeModernDispatchError — +// report a prefix and the length instead of splicing the payload (the #6079 / +// #6066 rule). +const maxResultTypeInError = 64 + +// missingSentinelText leads Error() when Sentinel is nil, which a direct +// construction that omits the required field can produce. Naming the mistake +// beats letting fmt splice "%!s()" into a message that reaches the +// downstream client. +const missingSentinelText = "input_required error with no classification sentinel" + +func (e *InputRequiredError) Error() string { + sentinel := func() string { + if e.Sentinel == nil { + return missingSentinelText + } + return e.Sentinel.Error() + }() + if len(e.ResultType) > maxResultTypeInError { + return fmt.Sprintf("%s: resultType=%q... (%d bytes)", + sentinel, e.ResultType[:maxResultTypeInError], len(e.ResultType)) + } + return fmt.Sprintf("%s: resultType=%q", sentinel, e.ResultType) +} + +// Unwrap returns the classification sentinel, keeping every existing +// errors.Is check on the untyped sentinel working unchanged. +func (e *InputRequiredError) Unwrap() error { return e.Sentinel } + +// InputRequiredFromError extracts the SEP-2322 input_required payload from an +// error chain, however deeply the backend client wrapped it. It reports false +// for errors that are not a drivable input_required round — including the +// unrecognized-resultType, disallowed-method, and invalid-payload variants of +// the same sentinel — so callers can use it as the single MRTR branch point: +// +// if round, ok := vmcp.InputRequiredFromError(err); ok { ... relay round ... } +// +// This is the seam the ingress half (dispatchModern's input_required envelope +// and the Legacy-client bridge; docs/arch/16-vmcp-mrtr.md slices 2-4) consumes. +func InputRequiredFromError(err error) (*InputRequiredResult, bool) { + var ire *InputRequiredError + if errors.As(err, &ire) && ire.Result != nil { + return ire.Result, true + } + return nil, false +} + +// Methods returns the JSON-RPC method name of each input request, keyed by +// the backend's request key. An entry whose value does not decode as an +// object with a "method" field maps to the empty string, so callers gating on +// method (capability checks) fail closed on malformed entries instead of +// skipping them. +func (r *InputRequiredResult) Methods() map[string]string { + if len(r.InputRequests) == 0 { + return nil + } + methods := make(map[string]string, len(r.InputRequests)) + for key, raw := range r.InputRequests { + var probe struct { + Method string `json:"method"` + } + // A failed decode leaves probe.Method empty — the fail-closed value. + _ = json.Unmarshal(raw, &probe) + methods[key] = probe.Method + } + return methods +} diff --git a/pkg/vmcp/mrtr_test.go b/pkg/vmcp/mrtr_test.go new file mode 100644 index 0000000000..cec6723ed1 --- /dev/null +++ b/pkg/vmcp/mrtr_test.go @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp + +import ( + "encoding/json" + "errors" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestInputRequiredError pins the carrier's contract for consumers that +// construct it directly (e.g. tests against mocks.MockBackendClient, which +// cannot reach the backend client's unexported sentinel): it unwraps to +// whatever Sentinel it was built with, renders the sentinel-prefixed message, +// and InputRequiredFromError extracts the round only when Result is non-nil. +func TestInputRequiredError(t *testing.T) { + t.Parallel() + + sentinel := errors.New("some classification") + state := "s1" + withRound := &InputRequiredError{ + ResultType: "input_required", + Result: &InputRequiredResult{RequestState: &state}, + Sentinel: sentinel, + } + require.ErrorIs(t, withRound, sentinel) + assert.Equal(t, `some classification: resultType="input_required"`, withRound.Error()) + + round, ok := InputRequiredFromError(fmt.Errorf("wrapped: %w", withRound)) + require.True(t, ok) + assert.Equal(t, &state, round.RequestState) + + // Result nil (unrecognized resultType, disallowed method, invalid + // payload): classification survives, extraction reports false. + sentinelOnly := &InputRequiredError{ResultType: "task", Sentinel: sentinel} + require.ErrorIs(t, sentinelOnly, sentinel) + round, ok = InputRequiredFromError(sentinelOnly) + assert.False(t, ok) + assert.Nil(t, round) +} + +// TestInputRequiredErrorBoundsResultType pins that Error() never interpolates +// more than 64 bytes of the backend-controlled resultType. +func TestInputRequiredErrorBoundsResultType(t *testing.T) { + t.Parallel() + + e := &InputRequiredError{ + ResultType: strings.Repeat("x", 200), + Sentinel: errors.New("sentinel"), + } + assert.Equal(t, + fmt.Sprintf("sentinel: resultType=%q... (200 bytes)", strings.Repeat("x", 64)), + e.Error()) +} + +// TestInputRequiredErrorNilSentinel pins what a construction that omits the +// required Sentinel renders: nothing enforces the field, so Error() must name +// the missing classification rather than splicing fmt's "%!s()" into text +// that reaches the downstream client. Unwrap still reports nil — the guard +// makes the mistake legible, it does not invent a classification. +func TestInputRequiredErrorNilSentinel(t *testing.T) { + t.Parallel() + + e := &InputRequiredError{ResultType: "input_required"} + assert.Equal(t, + missingSentinelText+`: resultType="input_required"`, + e.Error()) + assert.NotContains(t, e.Error(), "%!s") + assert.NoError(t, e.Unwrap()) + + // The bounded branch guards the sentinel too. + long := &InputRequiredError{ResultType: strings.Repeat("x", 200)} + assert.Equal(t, + fmt.Sprintf("%s: resultType=%q... (200 bytes)", missingSentinelText, strings.Repeat("x", 64)), + long.Error()) +} + +func TestInputRequiredResultMethods(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + requests map[string]json.RawMessage + want map[string]string + }{ + { + name: "methods extracted per key", + requests: map[string]json.RawMessage{ + "a": json.RawMessage(`{"method":"elicitation/create","params":{}}`), + "b": json.RawMessage(`{"method":"sampling/createMessage"}`), + }, + want: map[string]string{"a": "elicitation/create", "b": "sampling/createMessage"}, + }, + { + // A malformed entry maps to "" — the fail-closed value for a + // capability gate — rather than being silently dropped. + name: "malformed entry fails closed as empty method", + requests: map[string]json.RawMessage{ + "bad": json.RawMessage(`[1,2]`), + "good": json.RawMessage(`{"method":"roots/list"}`), + }, + want: map[string]string{"bad": "", "good": "roots/list"}, + }, + { + name: "empty requests yield nil", + requests: nil, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + r := &InputRequiredResult{InputRequests: tt.requests} + assert.Equal(t, tt.want, r.Methods()) + }) + } +} diff --git a/pkg/vmcp/optimizer/internal/tokencounter/counter.go b/pkg/vmcp/optimizer/internal/tokencounter/counter.go index f0f6495f73..606414324b 100644 --- a/pkg/vmcp/optimizer/internal/tokencounter/counter.go +++ b/pkg/vmcp/optimizer/internal/tokencounter/counter.go @@ -7,7 +7,7 @@ package tokencounter import ( "encoding/json" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // Counter estimates the number of tokens a tool definition would consume diff --git a/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go b/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go index d3abd6cfb5..9f720a03d5 100644 --- a/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go +++ b/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go @@ -7,8 +7,9 @@ import ( "encoding/json" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestJSONByteDivisionCounter(t *testing.T) { diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go index a3c6f67b28..acb8a154a4 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go @@ -19,11 +19,11 @@ import ( "sort" "strings" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "golang.org/x/sync/errgroup" _ "modernc.org/sqlite" // registers the "sqlite" database/sql driver + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" ) @@ -50,7 +50,29 @@ var schemaSQL string // and optional vector embedding-based semantic search. // It satisfies the types.ToolStore interface. type sqliteToolStore struct { - db *sql.DB + db *sql.DB + + // keepAlive is a dedicated connection held open for the store's whole lifetime. + // + // The database is in-memory with a shared cache, so it exists only while at + // least one connection to it is open. If the database/sql pool ever drops to + // zero connections, SQLite discards the entire database — table, FTS5 index + // and triggers — and nothing recreates it, because the schema is executed + // only in the constructor. Every later operation then fails with + // "no such table: llm_capabilities" until the process restarts. + // + // The pool does reach zero in practice. modernc.org/sqlite implements context + // cancellation with sqlite3_interrupt and reports an interrupted connection as + // unusable from conn.IsValid, so database/sql discards a connection whose + // statement was canceled instead of returning it to the pool. A single + // canceled request while the pool holds one connection is enough to destroy + // the database for the life of the process (#5889). + // + // This connection is acquired with a background context and never used to run + // statements, so it can never be interrupted and never discarded. It exists + // solely to keep the database alive. + keepAlive *sql.Conn + embeddingClient types.EmbeddingClient // nil = FTS5-only maxToolsToReturn int hybridSemanticRatio float64 @@ -77,8 +99,17 @@ func newSQLiteToolStore( return sqliteToolStore{}, fmt.Errorf("failed to open sqlite database: %w", err) } + // Pin the keep-alive connection before anything else, so the database cannot + // be discarded from this point on. See the keepAlive field for why. + keepAlive, err := db.Conn(context.Background()) + if err != nil { + _ = db.Close() + return sqliteToolStore{}, fmt.Errorf("failed to acquire keep-alive connection: %w", err) + } + // Execute schema if _, err := db.Exec(schemaSQL); err != nil { + _ = keepAlive.Close() _ = db.Close() return sqliteToolStore{}, fmt.Errorf("failed to initialize schema: %w", err) } @@ -100,6 +131,7 @@ func newSQLiteToolStore( store := sqliteToolStore{ db: db, + keepAlive: keepAlive, embeddingClient: embeddingClient, maxToolsToReturn: maxTools, hybridSemanticRatio: hybridRatio, @@ -176,30 +208,40 @@ func (s sqliteToolStore) generateEmbeddings(ctx context.Context, tools []server. return blobs, nil } -// Search finds tools matching the query string using FTS5 full-text search -// and optional semantic search when an embedding client is configured. +// Search finds tools matching q using FTS5 full-text search and optional +// semantic search when an embedding client is configured. The lexical arm +// prefers q.Keywords, falling back to q.Description when Keywords is empty; +// the semantic arm always embeds q.Description. // The allowedTools parameter limits results to only tools with names in the given set. // If allowedTools is empty, no results are returned (empty = no access). // Returns matches ranked by relevance. -func (s sqliteToolStore) Search(ctx context.Context, query string, allowedTools []string) ([]mcp.Tool, error) { +func (s sqliteToolStore) Search(ctx context.Context, q types.SearchQuery, allowedTools []string) ([]mcp.Tool, error) { if len(allowedTools) == 0 { slog.Debug("search skipped, no allowed tools") return nil, nil } - ftsExpr := sanitizeFTS5Query(query) + // The lexical arm prefers explicit keywords and falls back to the words of + // the description, either when no keywords were supplied or when every + // keyword was dropped as too common to discriminate. The semantic arm + // always embeds the description. + ftsExpr := sanitizeFTS5Terms(q.Keywords) + if ftsExpr == "" { + ftsExpr = sanitizeFTS5Terms(strings.Fields(q.Description)) + } // FTS5-only path (no embedding client) if s.embeddingClient == nil { if ftsExpr == "" { - slog.Debug("search skipped, empty FTS5 expression", "query", query) + slog.Debug("search skipped, empty FTS5 expression", "query", q.Description, "keywords", q.Keywords) return nil, nil } results, err := s.searchFTS5(ctx, ftsExpr, allowedTools, s.maxToolsToReturn) if err != nil { return nil, err } - slog.Debug("search completed (FTS5-only)", "query", query, "results", len(results), "matched_tools", matchNames(results)) + slog.Debug("search completed (FTS5-only)", + "query", q.Description, "keywords", q.Keywords, "results", len(results), "matched_tools", matchNames(results)) return results, nil } @@ -221,7 +263,7 @@ func (s sqliteToolStore) Search(ctx context.Context, query string, allowedTools if semanticLimit > 0 { g.Go(func() error { var err error - semanticResults, err = s.searchSemantic(gCtx, query, allowedTools, semanticLimit) + semanticResults, err = s.searchSemantic(gCtx, q.Description, allowedTools, semanticLimit) return err }) } @@ -233,7 +275,8 @@ func (s sqliteToolStore) Search(ctx context.Context, query string, allowedTools merged := mergeResults(ftsResults, semanticResults, s.maxToolsToReturn) slog.Debug("search completed (hybrid)", - "query", query, + "query", q.Description, + "keywords", q.Keywords, "fts5_results", len(ftsResults), "semantic_results", len(semanticResults), "merged_results", len(merged), @@ -243,14 +286,27 @@ func (s sqliteToolStore) Search(ctx context.Context, query string, allowedTools return merged, nil } -// Close releases the underlying database connection. +// Close releases the underlying database connections. Releasing the keep-alive +// connection drops the in-memory database, so it happens only here, on the way +// to closing the pool itself. +// +// Close may be called more than once: a keep-alive connection that was already +// released reports sql.ErrConnDone, which is not a failure. func (s sqliteToolStore) Close() error { var embErr error if s.embeddingClient != nil { embErr = s.embeddingClient.Close() } + + var keepAliveErr error + if s.keepAlive != nil { + if err := s.keepAlive.Close(); err != nil && !errors.Is(err, sql.ErrConnDone) { + keepAliveErr = err + } + } + dbErr := s.db.Close() - return errors.Join(embErr, dbErr) + return errors.Join(embErr, keepAliveErr, dbErr) } // searchFTS5 performs a full-text search using FTS5 MATCH with BM25 ranking. @@ -263,7 +319,7 @@ func (s sqliteToolStore) Close() error { // it has high cosine similarity, because the semantic query runs separately // and will surface it. // -// The ftsExpr is produced by sanitizeFTS5Query and is always passed as a +// The ftsExpr is produced by sanitizeFTS5Terms and is always passed as a // parameterized ? value, never interpolated into SQL. func (s sqliteToolStore) searchFTS5( ctx context.Context, ftsExpr string, allowedTools []string, limit int, @@ -458,8 +514,12 @@ func matchNames(matches []mcp.Tool) []string { } // problematicWords contains words that FTS5 interprets as operators or that -// are too common in tool metadata to be useful search terms. This set aligns -// with Python mcp_optimizer's DEFAULT_FTS_PROBLEMATIC_WORDS. +// are too common in tool metadata to be useful search terms. +// +// sanitizeFTS5Terms drops them. The alternative, falling back to a phrase +// search over the whole query when one of these appears, demands exact token +// adjacency and so reliably matches nothing; dropping the noise word and +// OR-joining the rest avoids the result flood without going empty-handed. var problematicWords = map[string]struct{}{ "name": {}, "description": {}, "schema": {}, "input": {}, "output": {}, "type": {}, "properties": {}, "required": {}, @@ -469,13 +529,16 @@ var problematicWords = map[string]struct{}{ "index": {}, "key": {}, "primary": {}, } -// sanitizeFTS5Query prepares a user query string for use with FTS5 MATCH. +// sanitizeFTS5Terms builds an FTS5 MATCH expression from a list of search terms. // -// The returned string is designed to be passed as a single ? parameter to -// QueryContext. It cannot cause SQL injection because it is always bound via ?. +// Each entry is split on whitespace (a single term may arrive multi-word), words +// in problematicWords are dropped, the remainder is deduplicated case-insensitively, +// and what is left is OR-joined. Returns "" when no usable term remains, letting +// the caller fall back to another source of terms. // -// FTS5 MATCH requires a single string operand containing the full query -// expression (e.g., "read" OR "write"). Individual terms cannot be separate +// The returned string is designed to be passed as a single ? parameter to +// QueryContext. FTS5 MATCH requires a single string operand containing the full +// query expression (e.g., "read" OR "write"). Individual terms cannot be separate // ? SQL parameters because the OR/AND operators are part of the FTS5 query // language, not SQL. // See: https://sqlite.org/fts5.html#full_text_query_syntax @@ -484,32 +547,24 @@ var problematicWords = map[string]struct{}{ // - SQL injection is prevented because the expression is always bound via ?. // - FTS5 operator injection is prevented by double-quoting each term and // escaping embedded double-quotes (standard FTS5 escaping). -func sanitizeFTS5Query(query string) string { - words := strings.Fields(strings.TrimSpace(query)) - if len(words) == 0 { - return "" - } - - hasProblematic := false - for _, word := range words { - if _, ok := problematicWords[strings.ToLower(word)]; ok { - hasProblematic = true - break +func sanitizeFTS5Terms(terms []string) string { + seen := make(map[string]struct{}, len(terms)) + var quoted []string + for _, term := range terms { + for _, word := range strings.Fields(term) { + lower := strings.ToLower(word) + if _, ok := problematicWords[lower]; ok { + continue + } + if _, ok := seen[lower]; ok { + continue + } + seen[lower] = struct{}{} + escaped := strings.ReplaceAll(word, `"`, `""`) + quoted = append(quoted, `"`+escaped+`"`) } } - // Single word or any problematic word present: use phrase search - if len(words) == 1 || hasProblematic { - escaped := strings.ReplaceAll(strings.Join(words, " "), `"`, `""`) - return `"` + escaped + `"` - } - - // Multi-word with no problematic words: join with OR - quoted := make([]string, len(words)) - for i, word := range words { - escaped := strings.ReplaceAll(word, `"`, `""`) - quoted[i] = `"` + escaped + `"` - } return strings.Join(quoted, " OR ") } diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go index 8e5bc2728b..908774bed6 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go @@ -12,10 +12,10 @@ import ( "fmt" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" ) @@ -56,7 +56,7 @@ func BenchmarkSearch_FTS5Only_1000Tools(b *testing.B) { b.ResetTimer() b.ReportAllocs() for b.Loop() { - _, _ = store.Search(ctx, "task operation", names) + _, _ = store.Search(ctx, types.SearchQuery{Description: "task operation"}, names) } } @@ -86,7 +86,7 @@ func BenchmarkSearch_Hybrid_1000Tools(b *testing.B) { b.ResetTimer() b.ReportAllocs() for b.Loop() { - _, _ = store.Search(ctx, "task operation", names) + _, _ = store.Search(ctx, types.SearchQuery{Description: "task operation"}, names) } } diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go index 4894a45258..e2a8e2caa4 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go @@ -6,14 +6,16 @@ package toolstore import ( "context" "fmt" + "slices" "sync" "sync/atomic" "testing" + "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" ) @@ -54,6 +56,7 @@ func TestNewSQLiteToolStore(t *testing.T) { t.Parallel() store := newTestStore(t, nil, nil) require.NotNil(t, store.db) + require.NotNil(t, store.keepAlive) require.Nil(t, store.embeddingClient) }) @@ -65,6 +68,45 @@ func TestNewSQLiteToolStore(t *testing.T) { }) } +// TestSQLiteToolStore_SurvivesCanceledStatement is a regression test for #5889. +// +// Canceling a statement mid-flight makes modernc.org/sqlite interrupt the +// connection, which database/sql then discards rather than returning to the +// pool. Without the store's keep-alive connection that empties the pool, and +// emptying the pool destroys the shared in-memory database: every later +// operation fails with "no such table: llm_capabilities" until the process +// restarts. +func TestSQLiteToolStore_SurvivesCanceledStatement(t *testing.T) { + t.Parallel() + + store := newTestStore(t, nil, nil) + tools := makeTools(mcp.NewTool("read_file", mcp.WithDescription("Read a file from disk"))) + require.NoError(t, store.UpsertTools(context.Background(), tools)) + + // Counting to 200 million takes far longer than the deadline below, so the + // statement is always canceled while it is still executing. + const slowQuery = `WITH RECURSIVE counter(x) AS ( + SELECT 1 UNION ALL SELECT x + 1 FROM counter WHERE x < 200000000 + ) SELECT count(*) FROM counter` + + for range 3 { + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + var count int + require.Error(t, store.db.QueryRowContext(ctx, slowQuery).Scan(&count)) + cancel() + } + + // The store still works: schema, FTS5 index and rows all intact. + results, err := store.Search( + context.Background(), + types.SearchQuery{Description: "read a file"}, + toolNames(tools), + ) + require.NoError(t, err) + require.Len(t, results, 1) + require.Equal(t, "read_file", results[0].Name) +} + func TestSQLiteToolStore_UpsertTools(t *testing.T) { t.Parallel() @@ -113,7 +155,7 @@ func TestSQLiteToolStore_UpsertTools(t *testing.T) { } require.NoError(t, store.UpsertTools(ctx, tc.upsert)) - results, err := store.Search(ctx, tc.searchQuery, tc.allowedTools) + results, err := store.Search(ctx, types.SearchQuery{Description: tc.searchQuery}, tc.allowedTools) require.NoError(t, err) require.Len(t, results, tc.wantLen) if tc.wantDesc != "" && len(results) > 0 { @@ -149,6 +191,7 @@ func TestSQLiteToolStore_Search(t *testing.T) { name string tools []server.ServerTool query string + keywords []string allowedTools []string wantNames []string wantNonEmpty bool // just assert results are non-empty (when exact names vary) @@ -241,6 +284,39 @@ func TestSQLiteToolStore_Search(t *testing.T) { allowedTools: []string{"generic_tool", "search_tool"}, wantNonEmpty: true, }, + { + name: "keywords surface a tool the description alone misses", + tools: makeTools( + mcp.NewTool("run_query", mcp.WithDescription("Executes arbitrary commands against a Postgres datastore")), + ), + query: "do something with data", + keywords: []string{"postgres"}, + allowedTools: []string{"run_query"}, + wantNames: []string{"run_query"}, + }, + { + // Negative twin of the case above. Without it, that case would keep + // passing if keywords were ignored but the description happened to + // start matching (say, if "data" left problematicWords). + name: "same description without keywords finds nothing", + tools: makeTools( + mcp.NewTool("run_query", mcp.WithDescription("Executes arbitrary commands against a Postgres datastore")), + ), + query: "do something with data", + keywords: nil, + allowedTools: []string{"run_query"}, + wantNames: nil, + }, + { + name: "keywords all problematic falls back to description", + tools: makeTools( + mcp.NewTool("read_file", mcp.WithDescription("Read a file from disk")), + ), + query: "read disk", + keywords: []string{"table", "schema"}, + allowedTools: []string{"read_file"}, + wantNames: []string{"read_file"}, + }, } for _, tc := range tests { @@ -251,7 +327,7 @@ func TestSQLiteToolStore_Search(t *testing.T) { require.NoError(t, store.UpsertTools(ctx, tc.tools)) - results, err := store.Search(ctx, tc.query, tc.allowedTools) + results, err := store.Search(ctx, types.SearchQuery{Description: tc.query, Keywords: tc.keywords}, tc.allowedTools) require.NoError(t, err) if tc.wantNonEmpty { @@ -307,7 +383,7 @@ func TestSQLiteToolStore_Search_ResultsCapped(t *testing.T) { ) require.NoError(t, store.UpsertTools(ctx, tools)) - results, err := store.Search(ctx, "file", toolNames(tools)) + results, err := store.Search(ctx, types.SearchQuery{Description: "file"}, toolNames(tools)) require.NoError(t, err) require.LessOrEqual(t, len(results), tc.wantMax, "results should be capped at %d", tc.wantMax) @@ -372,7 +448,7 @@ func TestSQLiteToolStore_Concurrent(t *testing.T) { go func(idx int) { defer wg.Done() // Pass a known tool name so we don't hit the empty-allowedTools shortcut - _, err := store.Search(ctx, "tool", []string{"tool_0"}) + _, err := store.Search(ctx, types.SearchQuery{Description: "tool"}, []string{"tool_0"}) if err != nil { t.Errorf("concurrent search failed for goroutine %d: %v", idx, err) } @@ -415,12 +491,84 @@ func TestSQLiteToolStore_HybridSearch(t *testing.T) { require.NoError(t, store.UpsertTools(ctx, tools)) // Hybrid search should return results from both FTS5 and semantic - results, err := store.Search(ctx, "file", toolNames(tools)) + results, err := store.Search(ctx, types.SearchQuery{Description: "file"}, toolNames(tools)) require.NoError(t, err) require.NotEmpty(t, results) require.LessOrEqual(t, len(results), DefaultMaxToolsToReturn) } +// TestSQLiteToolStore_SemanticArmEmbedsDescriptionOnly pins half of the hybrid +// design: whatever the keywords are, the semantic arm embeds Description and +// nothing else. A regression that embedded the keywords, or a concatenation of +// both fields, would go unnoticed through result assertions alone, because the +// merged set can look identical either way. +// +// The other half, keywords driving the lexical arm, is covered by the +// "keywords surface a tool the description alone misses" case in +// TestSQLiteToolStore_Search, which runs FTS5-only so the semantic arm cannot +// mask the result. +func TestSQLiteToolStore_SemanticArmEmbedsDescriptionOnly(t *testing.T) { + t.Parallel() + client := newFakeEmbeddingClient(384) + store := newTestStore(t, client, nil) + ctx := context.Background() + + tools := makeTools( + mcp.NewTool("run_query", mcp.WithDescription("Executes arbitrary commands against a Postgres datastore")), + mcp.NewTool("send_email", mcp.WithDescription("Send an email message")), + ) + require.NoError(t, store.UpsertTools(ctx, tools)) + + _, err := store.Search(ctx, types.SearchQuery{ + Description: "do something with data", + Keywords: []string{"postgres", "sql"}, + }, toolNames(tools)) + require.NoError(t, err) + + require.Equal(t, []string{"do something with data"}, client.embeddedTexts(), + "semantic arm must embed Description verbatim, never the keywords") +} + +// TestSQLiteToolStore_KeywordsOnly covers a SearchQuery with no description. +// FindTool rejects that upstream, but the store is a separate unit. +func TestSQLiteToolStore_KeywordsOnly(t *testing.T) { + t.Parallel() + + tools := makeTools( + mcp.NewTool("github_create_issue", mcp.WithDescription("Create a GitHub issue")), + mcp.NewTool("slack_send_message", mcp.WithDescription("Send a Slack message")), + ) + + t.Run("keywords alone drive the lexical arm", func(t *testing.T) { + t.Parallel() + // FTS5-only, so the semantic arm cannot supply the match. + store := newTestStore(t, nil, nil) + ctx := context.Background() + require.NoError(t, store.UpsertTools(ctx, tools)) + + results, err := store.Search(ctx, types.SearchQuery{ + Keywords: []string{"github"}, + }, toolNames(tools)) + require.NoError(t, err) + require.Equal(t, []string{"github_create_issue"}, matchNames(results)) + }) + + t.Run("hybrid tolerates an empty description", func(t *testing.T) { + t.Parallel() + client := newFakeEmbeddingClient(384) + store := newTestStore(t, client, nil) + ctx := context.Background() + require.NoError(t, store.UpsertTools(ctx, tools)) + + _, err := store.Search(ctx, types.SearchQuery{ + Keywords: []string{"github"}, + }, toolNames(tools)) + require.NoError(t, err) + // The semantic arm still runs and embeds the empty description. + require.Equal(t, []string{""}, client.embeddedTexts()) + }) +} + func TestSQLiteToolStore_ConcurrentSemantic(t *testing.T) { t.Parallel() client := newFakeEmbeddingClient(384) @@ -440,7 +588,7 @@ func TestSQLiteToolStore_ConcurrentSemantic(t *testing.T) { wg.Add(1) go func(idx int) { defer wg.Done() - _, err := store.Search(ctx, "file", toolNames(tools)) + _, err := store.Search(ctx, types.SearchQuery{Description: "file"}, toolNames(tools)) if err != nil { t.Errorf("concurrent semantic search failed for goroutine %d: %v", idx, err) } @@ -460,41 +608,55 @@ func TestSQLiteToolStore_EmbeddingRoundTrip(t *testing.T) { require.Equal(t, original, decoded) } -func TestSanitizeFTS5Query(t *testing.T) { +// TestSanitizeFTS5Terms covers both callers of the sanitizer: an explicit +// keyword list, and the words of a description. Both converge on the same +// function, so a single-element slice containing spaces exercises the same +// whitespace split that strings.Fields(description) produces. +func TestSanitizeFTS5Terms(t *testing.T) { t.Parallel() tests := []struct { - input string + name string + terms []string wantExpr string }{ - {input: "simple", wantExpr: `"simple"`}, - {input: "two words", wantExpr: `"two" OR "words"`}, - {input: "hello world foo", wantExpr: `"hello" OR "world" OR "foo"`}, - {input: "", wantExpr: ""}, - {input: " ", wantExpr: ""}, - - // Special chars are NOT stripped (unlike previous behavior) - {input: "key:value", wantExpr: `"key:value"`}, - {input: `"quoted"`, wantExpr: `"""quoted"""`}, - {input: "read*", wantExpr: `"read*"`}, - {input: "***", wantExpr: `"***"`}, - {input: "read + file", wantExpr: `"read" OR "+" OR "file"`}, - - // Problematic words trigger phrase search - {input: "name value", wantExpr: `"name value"`}, - {input: "search description fast", wantExpr: `"search description fast"`}, - {input: "read tool write", wantExpr: `"read tool write"`}, - {input: "schema definition", wantExpr: `"schema definition"`}, - - // Non-problematic multi-word queries use OR - {input: "read write", wantExpr: `"read" OR "write"`}, - {input: "github slack", wantExpr: `"github" OR "slack"`}, + {name: "single term", terms: []string{"simple"}, wantExpr: `"simple"`}, + {name: "two terms", terms: []string{"two", "words"}, wantExpr: `"two" OR "words"`}, + {name: "three terms", terms: []string{"hello", "world", "foo"}, wantExpr: `"hello" OR "world" OR "foo"`}, + {name: "nil list", terms: nil, wantExpr: ""}, + {name: "whitespace-only entry", terms: []string{" "}, wantExpr: ""}, + {name: "empty string entry", terms: []string{""}, wantExpr: ""}, + + // Special chars are NOT stripped; quoting makes them inert to FTS5. + {name: "colon", terms: []string{"key:value"}, wantExpr: `"key:value"`}, + {name: "embedded quotes", terms: []string{`"quoted"`}, wantExpr: `"""quoted"""`}, + {name: "trailing star", terms: []string{"read*"}, wantExpr: `"read*"`}, + {name: "all stars", terms: []string{"***"}, wantExpr: `"***"`}, + {name: "plus operator", terms: []string{"read", "+", "file"}, wantExpr: `"read" OR "+" OR "file"`}, + + // Problematic words are dropped, not phrase-searched. A phrase search + // would demand exact adjacency and reliably match nothing. + {name: "problematic word dropped", terms: []string{"name", "lookup"}, wantExpr: `"lookup"`}, + {name: "problematic word in middle", terms: []string{"search", "description", "fast"}, wantExpr: `"search" OR "fast"`}, + {name: "problematic word leaves two", terms: []string{"read", "tool", "write"}, wantExpr: `"read" OR "write"`}, + {name: "problematic word leaves one", terms: []string{"schema", "definition"}, wantExpr: `"definition"`}, + {name: "all problematic", terms: []string{"table", "schema"}, wantExpr: ""}, + {name: "mixed problematic and valid", terms: []string{"table", "postgres"}, wantExpr: `"postgres"`}, + + // A single entry containing spaces splits, matching the description path. + {name: "multi-word entry splits", terms: []string{"list issues"}, wantExpr: `"list" OR "issues"`}, + {name: "multi-word entry drops problematic", terms: []string{"name lookup"}, wantExpr: `"lookup"`}, + {name: "every word problematic yields empty", terms: []string{"name value"}, wantExpr: ""}, + + {name: "duplicates deduped", terms: []string{"github", "github"}, wantExpr: `"github"`}, + {name: "dedup is case-insensitive", terms: []string{"GitHub", "github"}, wantExpr: `"GitHub"`}, + {name: "escapes embedded quote", terms: []string{`db"name`}, wantExpr: `"db""name"`}, } for _, tt := range tests { - t.Run(tt.input, func(t *testing.T) { + t.Run(tt.name, func(t *testing.T) { t.Parallel() - gotExpr := sanitizeFTS5Query(tt.input) + gotExpr := sanitizeFTS5Terms(tt.terms) require.Equal(t, tt.wantExpr, gotExpr) }) } @@ -713,14 +875,45 @@ func TestSQLiteToolStore_SemanticDistanceThreshold(t *testing.T) { // import cycles. type fakeEmbeddingClient struct { dim int + + // mu guards embedded, which Search writes from an errgroup goroutine. + mu sync.Mutex + // embedded records every text passed to Embed, so tests can assert which + // SearchQuery field reaches the semantic arm. + embedded []string } func newFakeEmbeddingClient(dim int) *fakeEmbeddingClient { return &fakeEmbeddingClient{dim: dim} } +// embeddedTexts returns the texts passed to Embed so far. EmbedBatch, used on +// the write path by UpsertTools, is deliberately not recorded. +func (f *fakeEmbeddingClient) embeddedTexts() []string { + f.mu.Lock() + defer f.mu.Unlock() + return slices.Clone(f.embedded) +} + func (f *fakeEmbeddingClient) Embed(_ context.Context, text string) ([]float32, error) { - // Simple deterministic hash: use string bytes as seed + f.mu.Lock() + f.embedded = append(f.embedded, text) + f.mu.Unlock() + + return f.embedVector(text), nil +} + +func (f *fakeEmbeddingClient) EmbedBatch(_ context.Context, texts []string) ([][]float32, error) { + result := make([][]float32, len(texts)) + for i, text := range texts { + result[i] = f.embedVector(text) + } + return result, nil +} + +// embedVector produces the deterministic vector without recording the text, +// keeping embeddedTexts limited to the read path. +func (f *fakeEmbeddingClient) embedVector(text string) []float32 { vec := make([]float32, f.dim) for i := range vec { // Use text bytes to generate deterministic values @@ -730,19 +923,7 @@ func (f *fakeEmbeddingClient) Embed(_ context.Context, text string) ([]float32, } vec[i] = float32(b)/128.0 - 1.0 + float32(i)*0.001 } - return vec, nil -} - -func (f *fakeEmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { - result := make([][]float32, len(texts)) - for i, text := range texts { - vec, err := f.Embed(ctx, text) - if err != nil { - return nil, err - } - result[i] = vec - } - return result, nil + return vec } func (*fakeEmbeddingClient) Close() error { return nil } diff --git a/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go b/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go index 1a344ee73e..6ac9e8e9e4 100644 --- a/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go +++ b/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go @@ -13,8 +13,9 @@ import ( context "context" reflect "reflect" - mcp "github.com/mark3labs/mcp-go/mcp" - server "github.com/mark3labs/mcp-go/server" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + server "github.com/stacklok/toolhive-core/mcpcompat/server" + types "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" gomock "go.uber.org/mock/gomock" ) @@ -57,18 +58,18 @@ func (mr *MockToolStoreMockRecorder) Close() *gomock.Call { } // Search mocks base method. -func (m *MockToolStore) Search(ctx context.Context, query string, allowedTools []string) ([]mcp.Tool, error) { +func (m *MockToolStore) Search(ctx context.Context, q types.SearchQuery, allowedTools []string) ([]mcp.Tool, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Search", ctx, query, allowedTools) + ret := m.ctrl.Call(m, "Search", ctx, q, allowedTools) ret0, _ := ret[0].([]mcp.Tool) ret1, _ := ret[1].(error) return ret0, ret1 } // Search indicates an expected call of Search. -func (mr *MockToolStoreMockRecorder) Search(ctx, query, allowedTools any) *gomock.Call { +func (mr *MockToolStoreMockRecorder) Search(ctx, q, allowedTools any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Search", reflect.TypeOf((*MockToolStore)(nil).Search), ctx, query, allowedTools) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Search", reflect.TypeOf((*MockToolStore)(nil).Search), ctx, q, allowedTools) } // UpsertTools mocks base method. diff --git a/pkg/vmcp/optimizer/internal/types/types.go b/pkg/vmcp/optimizer/internal/types/types.go index 9c54f5a92c..3b064d5c8d 100644 --- a/pkg/vmcp/optimizer/internal/types/types.go +++ b/pkg/vmcp/optimizer/internal/types/types.go @@ -10,10 +10,23 @@ import ( "context" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" ) +// SearchQuery describes a tool search request. +// +// The two fields feed different arms of hybrid search: Description is embedded +// for semantic matching, Keywords build the FTS5/BM25 expression. When Keywords +// is empty the lexical arm falls back to Description. +type SearchQuery struct { + // Description is a natural language description of the desired capability. + Description string + + // Keywords are optional discrete terms for the keyword-search arm. + Keywords []string +} + // ToolStore defines the interface for storing and searching tools. // Implementations may use in-memory maps, SQLite FTS5, or other backends. // @@ -24,12 +37,14 @@ type ToolStore interface { // Tools are identified by name; duplicate names are overwritten. UpsertTools(ctx context.Context, tools []server.ServerTool) error - // Search finds tools matching the query string. + // Search finds tools matching q. Description is embedded for semantic + // matching; Keywords build the FTS5/BM25 expression, falling back to + // Description when empty. // The allowedTools parameter limits results to only tools with names in the given set. // If allowedTools is empty, no results are returned (empty = no access). // Returns matches ranked by relevance. The returned mcp.Tool values contain // only Name and Description; the caller is responsible for enriching with schemas. - Search(ctx context.Context, query string, allowedTools []string) ([]mcp.Tool, error) + Search(ctx context.Context, q SearchQuery, allowedTools []string) ([]mcp.Tool, error) // Close releases any resources held by the store (e.g., database connections). // For in-memory stores this is a no-op. diff --git a/pkg/vmcp/optimizer/optimizer.go b/pkg/vmcp/optimizer/optimizer.go index 35389cf0f2..41c538a264 100644 --- a/pkg/vmcp/optimizer/optimizer.go +++ b/pkg/vmcp/optimizer/optimizer.go @@ -14,16 +14,18 @@ package optimizer import ( "context" + "encoding/json" "fmt" "log/slog" + "maps" "os" + "slices" "strconv" "strings" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" httpval "github.com/stacklok/toolhive-core/validation/http" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity" @@ -202,7 +204,7 @@ type FindToolInput struct { // ToolKeywords is an optional list of keywords to narrow the search. //nolint:lll // Long description tag provides essential context for LLM tool usage. - ToolKeywords []string `json:"tool_keywords,omitempty" description:"Optional keywords for BM25 text search to narrow results (e.g. ['list', 'issues', 'github'] or ['SQL', 'query', 'postgres']). Combined with tool_description for hybrid search."` + ToolKeywords []string `json:"tool_keywords,omitempty" description:"Optional keywords driving the BM25 keyword-search arm (e.g. ['list', 'issues', 'github'] or ['SQL', 'query', 'postgres']). Semantic matching always uses tool_description, so provide a complete description even when supplying keywords."` } // FindToolOutput contains the results of a tool search. @@ -230,6 +232,50 @@ type CallToolInput struct { Parameters map[string]any `json:"parameters" description:"Dictionary of arguments required by the tool. The structure must match the tool's input schema as returned by find_tool."` } +// callToolArgToolName is the parameters key resolveCallToolTarget hoists a nested +// name out of. It must match the json tag on CallToolInput.ToolName, since the +// nested lookup is a map index while the top-level one goes through encoding/json; +// TestCallToolArgToolNameMatchesStructTag guards the pair against drift. +const callToolArgToolName = "tool_name" + +// resolveCallToolTarget resolves the tool a call_tool invocation targets, +// accepting the common LLM malformation where tool_name is nested inside +// parameters instead of sitting alongside it. A top-level name always wins, so a +// backend tool with its own tool_name argument still works. +// +// It is deliberately unexported: decoding a payload into a CallToolInput is the +// only supported way to learn which tool a call_tool request names. Reading the +// name out of a raw arguments map instead misses the case-variant keys +// encoding/json accepts, and a target that authorization and dispatch disagree on +// is a tool executing under a policy decision made for a different name. +// +// params is never modified; a copy is returned when a nested name is hoisted. +func resolveCallToolTarget(name string, params map[string]any) (string, map[string]any) { + if name != "" { + return name, params + } + nested, ok := params[callToolArgToolName].(string) + if !ok || nested == "" { + return name, params + } + hoisted := maps.Clone(params) + delete(hoisted, callToolArgToolName) + return nested, hoisted +} + +// UnmarshalJSON hoists a nested tool_name so dispatch targets the same tool +// authorization approved. See resolveCallToolTarget. +func (in *CallToolInput) UnmarshalJSON(data []byte) error { + type rawCallToolInput CallToolInput // drops the method set to avoid recursion + var raw rawCallToolInput + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + raw.ToolName, raw.Parameters = resolveCallToolTarget(raw.ToolName, raw.Parameters) + *in = CallToolInput(raw) + return nil +} + // NewOptimizerFactory creates the embedding client and SQLite tool store from // the given OptimizerConfig, then returns an OptimizerFactory and a cleanup // function that closes the store. The caller must invoke the cleanup function @@ -337,7 +383,10 @@ func (d *toolOptimizer) FindTool(ctx context.Context, input FindToolInput) (*Fin return nil, fmt.Errorf("tool_description is required") } - matches, err := d.store.Search(ctx, input.ToolDescription, d.toolNames) + matches, err := d.store.Search(ctx, types.SearchQuery{ + Description: input.ToolDescription, + Keywords: input.ToolKeywords, + }, d.toolNames) if err != nil { return nil, fmt.Errorf("tool search failed: %w", err) } @@ -378,7 +427,10 @@ func (d *toolOptimizer) FindTool(ctx context.Context, input FindToolInput) (*Fin // is invoked directly with the given parameters. func (d *toolOptimizer) CallTool(ctx context.Context, input CallToolInput) (*mcp.CallToolResult, error) { if input.ToolName == "" { - return nil, fmt.Errorf("tool_name is required") + return nil, fmt.Errorf( + `tool_name is required: call_tool expects {"tool_name": "", `+ + `"parameters": {}}, got parameters keys %v`, + slices.Sorted(maps.Keys(input.Parameters))) } // Verify the tool exists diff --git a/pkg/vmcp/optimizer/optimizer_test.go b/pkg/vmcp/optimizer/optimizer_test.go index 1c59228507..8ae1457999 100644 --- a/pkg/vmcp/optimizer/optimizer_test.go +++ b/pkg/vmcp/optimizer/optimizer_test.go @@ -7,19 +7,22 @@ import ( "context" "encoding/json" "fmt" + "maps" + "reflect" "strings" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/tokencounter" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/schema" ) func TestGetAndValidateConfig(t *testing.T) { @@ -437,11 +440,11 @@ func newMockStoreWithSubstringSearch(ctrl *gomock.Controller) *mocks.MockToolSto ).AnyTimes() store.EXPECT().Search(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, query string, allowedTools []string) ([]mcp.Tool, error) { + func(_ context.Context, q types.SearchQuery, allowedTools []string) ([]mcp.Tool, error) { if len(allowedTools) == 0 { return nil, nil } - searchTerm := strings.ToLower(query) + searchTerm := strings.ToLower(q.Description) allowedSet := make(map[string]struct{}, len(allowedTools)) for _, name := range allowedTools { allowedSet[name] = struct{}{} @@ -483,8 +486,8 @@ func TestOptimizer_SearchDelegation(t *testing.T) { } store.EXPECT().UpsertTools(gomock.Any(), gomock.Any()).Return(nil) - store.EXPECT().Search(gomock.Any(), "query", gomock.Any()).DoAndReturn( - func(_ context.Context, _ string, allowedTools []string) ([]mcp.Tool, error) { + store.EXPECT().Search(gomock.Any(), types.SearchQuery{Description: "query"}, gomock.Any()).DoAndReturn( + func(_ context.Context, _ types.SearchQuery, allowedTools []string) ([]mcp.Tool, error) { require.ElementsMatch(t, []string{"tool_a", "tool_b"}, allowedTools) return []mcp.Tool{ {Name: "tool_a", Description: "Tool A"}, @@ -688,6 +691,35 @@ func TestOptimizer_FindTool(t *testing.T) { } } +// TestOptimizer_FindToolPassesKeywords is the regression guard for the bug +// where ToolKeywords was decoded off the wire, logged, and then dropped +// instead of reaching the store's Search call. +func TestOptimizer_FindToolPassesKeywords(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + store := mocks.NewMockToolStore(ctrl) + + tools := []server.ServerTool{ + {Tool: mcp.Tool{Name: "tool_a", Description: "Tool A"}}, + } + + store.EXPECT().UpsertTools(gomock.Any(), gomock.Any()).Return(nil) + store.EXPECT().Search(gomock.Any(), types.SearchQuery{ + Description: "query", + Keywords: []string{"list", "issues", "github"}, + }, gomock.Any()).Return(nil, nil) + + opt, err := newToolOptimizer(context.Background(), store, tokencounter.NewJSONByteCounter(), tools) + require.NoError(t, err) + + _, err = opt.FindTool(context.Background(), FindToolInput{ + ToolDescription: "query", + ToolKeywords: []string{"list", "issues", "github"}, + }) + require.NoError(t, err) +} + func TestOptimizerFactoryWithStore(t *testing.T) { t.Parallel() @@ -816,7 +848,7 @@ func TestOptimizer_CallTool(t *testing.T) { Parameters: map[string]any{}, }, expectedError: true, - errorContains: "tool_name is required", + errorContains: `call_tool expects {"tool_name"`, }, } @@ -848,3 +880,151 @@ func TestOptimizer_CallTool(t *testing.T) { }) } } + +func TestResolveCallToolTarget(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + toolName string + params map[string]any + expectedName string + expectedParams map[string]any + }{ + { + name: "top-level name is returned unchanged", + toolName: "search", + params: map[string]any{"query": "weather"}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "nested name is hoisted and removed from params", + params: map[string]any{"tool_name": "search", "query": "weather"}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "top-level name wins and params are left untouched", + toolName: "search", + params: map[string]any{"tool_name": "other", "query": "weather"}, + expectedName: "search", + expectedParams: map[string]any{"tool_name": "other", "query": "weather"}, + }, + { + name: "nested empty name is not hoisted", + params: map[string]any{"tool_name": ""}, + expectedName: "", + expectedParams: map[string]any{"tool_name": ""}, + }, + { + name: "nested non-string name is not hoisted", + params: map[string]any{"tool_name": 123}, + expectedName: "", + expectedParams: map[string]any{"tool_name": 123}, + }, + { + name: "no name anywhere", + params: map[string]any{"query": "weather"}, + expectedName: "", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "nil params", + expectedName: "", + expectedParams: nil, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Callers share this map with downstream consumers, so it must survive intact. + original := maps.Clone(tc.params) + + gotName, gotParams := resolveCallToolTarget(tc.toolName, tc.params) + require.Equal(t, tc.expectedName, gotName) + require.Equal(t, tc.expectedParams, gotParams) + require.Equal(t, original, tc.params, "input params must not be mutated") + }) + } +} + +// Both call_tool handlers and the authz middleware resolve the target with +// schema.Translate, so every shape that reaches a tool through that call must +// resolve to the same name for all three. A reader that indexes the arguments map +// instead sees a different target for the case-variant keys below. +func TestCallToolInput_TranslateResolvesTarget(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + args map[string]any + expectedName string + expectedParams map[string]any + }{ + { + name: "top-level name", + args: map[string]any{"tool_name": "search", "parameters": map[string]any{"query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "nested name is hoisted", + args: map[string]any{"parameters": map[string]any{"tool_name": "search", "query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + // encoding/json prefers an exact field match but falls back to a + // case-insensitive one, so these name a tool just as the lowercase key does. + name: "case-variant top-level key still names the tool", + args: map[string]any{"Tool_Name": "search", "parameters": map[string]any{"query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "case-variant parameters key still carries the arguments", + args: map[string]any{"tool_name": "search", "PARAMETERS": map[string]any{"query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + // The nested lookup is a map index, not a struct field match, so it is + // exact. Dispatch names no tool here and neither may authorization. + name: "case-variant nested key is not hoisted", + args: map[string]any{"parameters": map[string]any{"Tool_Name": "search"}}, + expectedName: "", + expectedParams: map[string]any{"Tool_Name": "search"}, + }, + { + name: "no name anywhere", + args: map[string]any{"parameters": map[string]any{"query": "weather"}}, + expectedName: "", + expectedParams: map[string]any{"query": "weather"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, err := schema.Translate[CallToolInput](tc.args) + require.NoError(t, err) + require.Equal(t, tc.expectedName, got.ToolName) + require.Equal(t, tc.expectedParams, got.Parameters) + }) + } +} + +// TestCallToolArgToolNameMatchesStructTag guards the one place the target is +// resolved by indexing a map rather than decoding a struct. If the json tag is +// renamed and the constant is not, a nested name silently stops being hoisted. +func TestCallToolArgToolNameMatchesStructTag(t *testing.T) { + t.Parallel() + + f, ok := reflect.TypeOf(CallToolInput{}).FieldByName("ToolName") + require.True(t, ok, "CallToolInput must have a ToolName field") + assert.Equal(t, callToolArgToolName, f.Tag.Get("json")) +} diff --git a/pkg/vmcp/ratelimit/decorator.go b/pkg/vmcp/ratelimit/decorator.go index ff4f30f236..cd472cd7f6 100644 --- a/pkg/vmcp/ratelimit/decorator.go +++ b/pkg/vmcp/ratelimit/decorator.go @@ -18,6 +18,10 @@ import ( // decorator wraps a [core.VMCP] to rate-limit tool calls. Every method except // CallTool is promoted from the embedded inner core unchanged. The decorator sits // below the session optimizer, so name is already the resolved backend tool name. +// +// CheckToolCall is deliberately NOT overridden: promoting it to inner is correct +// because a pre-flight admission check must not consume a rate-limit token — the +// limiter applies only when the call is actually dispatched via CallTool. type decorator struct { core.VMCP limiter baseratelimit.Limiter diff --git a/pkg/vmcp/ratelimit/decorator_test.go b/pkg/vmcp/ratelimit/decorator_test.go index cb0832ff31..fb9bbe1337 100644 --- a/pkg/vmcp/ratelimit/decorator_test.go +++ b/pkg/vmcp/ratelimit/decorator_test.go @@ -39,13 +39,19 @@ func (l *recordingLimiter) Allow(_ context.Context, toolName, userID string) (*b type recordingCore struct { core.VMCP - called bool - identity *auth.Identity - name string - args map[string]any - meta map[string]any - result *vmcp.ToolCallResult - err error + called bool + checkCalled bool + identity *auth.Identity + name string + args map[string]any + meta map[string]any + result *vmcp.ToolCallResult + err error +} + +func (c *recordingCore) CheckToolCall(_ context.Context, _ *auth.Identity, _ string, _ map[string]any) error { + c.checkCalled = true + return nil } func (c *recordingCore) CallTool( @@ -157,3 +163,21 @@ func TestCallToolNilIdentityUsesEmptyUserID(t *testing.T) { assert.True(t, inner.called) assert.Empty(t, limiter.userID) } + +// TestCheckToolCallConsumesNoTokens verifies the deliberate non-override of +// CheckToolCall: a pre-flight admission check promotes straight to inner without +// touching the limiter, so it never consumes a rate-limit token (the limiter +// applies only at CallTool dispatch). +func TestCheckToolCallConsumesNoTokens(t *testing.T) { + t.Parallel() + + limiter := &recordingLimiter{} + inner := &recordingCore{} + decorated := NewDecorator(inner, limiter) + + err := decorated.CheckToolCall(t.Context(), nil, "backend_a_echo", nil) + + require.NoError(t, err) + assert.True(t, inner.checkCalled, "the check must promote to inner") + assert.Empty(t, limiter.toolName, "a pre-flight check must not consume a rate-limit token") +} diff --git a/pkg/vmcp/router/session_router.go b/pkg/vmcp/router/session_router.go index 436cc18c70..b357807b5c 100644 --- a/pkg/vmcp/router/session_router.go +++ b/pkg/vmcp/router/session_router.go @@ -6,8 +6,12 @@ package router import ( "context" "fmt" + "log/slog" + "sort" "strings" + "github.com/yosida95/uritemplate/v3" + "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -17,6 +21,19 @@ import ( // on the discovery middleware injecting DiscoveredCapabilities into the context. type sessionRouter struct { routingTable *vmcp.RoutingTable + + // resourceTemplates is the aggregated resource-template table pre-sorted by + // template string with each template pre-parsed. These are invariant for a + // given routing table, so they are computed once at construction rather than + // re-sorted and re-parsed on every RouteResource miss. + resourceTemplates []compiledResourceTemplate +} + +// compiledResourceTemplate pairs an aggregated resource-template entry with its +// pre-parsed RFC 6570 template for matchResourceTemplate. +type compiledResourceTemplate struct { + tmpl *uritemplate.Template + target *vmcp.BackendTarget } // NewSessionRouter creates a Router that routes from the provided RoutingTable @@ -24,11 +41,80 @@ type sessionRouter struct { // composite tool workflow engines because it couples routing to the session // rather than to middleware-managed context values. func NewSessionRouter(rt *vmcp.RoutingTable) Router { - return &sessionRouter{routingTable: rt} + return &sessionRouter{routingTable: rt, resourceTemplates: compileResourceTemplates(rt)} +} + +// compileResourceTemplates builds the sorted, pre-parsed resource-template list +// for a routing table. Template keys are sorted so that when overlapping +// templates match the same URI (e.g. a greedy "{+path}" template alongside a +// more specific one) the first-match winner is deterministic and stable across +// runs. A template string that fails to parse is dropped here (logged once) +// rather than on every routing miss. +func compileResourceTemplates(rt *vmcp.RoutingTable) []compiledResourceTemplate { + if rt == nil || len(rt.ResourceTemplates) == 0 { + return nil + } + tmplStrs := make([]string, 0, len(rt.ResourceTemplates)) + for tmplStr := range rt.ResourceTemplates { + tmplStrs = append(tmplStrs, tmplStr) + } + sort.Strings(tmplStrs) + + compiled := make([]compiledResourceTemplate, 0, len(tmplStrs)) + for _, tmplStr := range tmplStrs { + tmpl, err := uritemplate.New(tmplStr) + if err != nil { + slog.Warn("skipping invalid resource URI template during routing", + "template", tmplStr, "error", err) + continue + } + compiled = append(compiled, compiledResourceTemplate{ + tmpl: tmpl, + target: rt.ResourceTemplates[tmplStr], + }) + } + return compiled +} + +// ResolveToolRef resolves a composite-tool step reference ("{workloadID}.{toolName}") +// to its routing-table key, returning ok=false when the reference does not +// resolve. It is the shared primitive behind both composite-tool accessibility +// filtering (compositetools.isToolStepAccessible) and per-step annotation +// resolution (core.stepAnnotationResolver), so the two paths cannot drift. +// +// rt may be nil (returns "", false). Resolution order mirrors ResolveToolName: +// 1. Exact key: the resolved/conflict-resolved name stored in rt.Tools. +// 2. Dot convention "{workloadID}.{originalCapabilityName}": workload IDs +// are Kubernetes resource names (no dots), so the first dot separates the +// workload ID from the original backend capability name. A leading dot +// (dotIdx == 0) is rejected so an empty workload ID never matches. +func ResolveToolRef(rt *vmcp.RoutingTable, stepTool string) (resolvedName string, ok bool) { + if rt == nil || rt.Tools == nil || stepTool == "" { + return "", false + } + + // Fast path: exact key match. + if _, exists := rt.Tools[stepTool]; exists { + return stepTool, true + } + + // Fallback: dot convention "{workloadID}.{toolName}". + if dotIdx := strings.Index(stepTool, "."); dotIdx > 0 { + workloadID := stepTool[:dotIdx] + capName := stepTool[dotIdx+1:] + for resolvedName, target := range rt.Tools { + if target.WorkloadID == workloadID && target.GetBackendCapabilityName(resolvedName) == capName { + return resolvedName, true + } + } + } + + return "", false } // RouteTool resolves a tool name to its backend target using the session's -// routing table directly. +// routing table directly. Resolution is delegated to ResolveToolRef so this +// path cannot drift from composite-tool accessibility filtering. // // Two naming conventions are supported: // @@ -45,73 +131,75 @@ func NewSessionRouter(rt *vmcp.RoutingTable) Router { // tools by their pre-conflict-resolution name (e.g. "my-backend.echo"), while // the routing table may store them under a prefixed key ("my-backend_echo"). func (r *sessionRouter) RouteTool(_ context.Context, toolName string) (*vmcp.BackendTarget, error) { - if r.routingTable == nil || r.routingTable.Tools == nil { + resolvedName, ok := ResolveToolRef(r.routingTable, toolName) + if !ok { return nil, fmt.Errorf("%w: %s", ErrToolNotFound, toolName) } - - // Fast path: exact key match. - if target, exists := r.routingTable.Tools[toolName]; exists { - return target, nil - } - - // Fallback: dot convention "{workloadID}.{toolName}". - // Workload IDs are Kubernetes resource names and cannot contain dots, - // so the first dot unambiguously separates the workload ID from the - // original backend capability name. - if dotIdx := strings.Index(toolName, "."); dotIdx > 0 { - workloadID := toolName[:dotIdx] - capName := toolName[dotIdx+1:] - for resolvedName, target := range r.routingTable.Tools { - if target.WorkloadID == workloadID && target.GetBackendCapabilityName(resolvedName) == capName { - return target, nil - } - } - } - - return nil, fmt.Errorf("%w: %s", ErrToolNotFound, toolName) + return r.routingTable.Tools[resolvedName], nil } // ResolveToolName returns the routing table key (conflict-resolved name) for -// toolName. If toolName is an exact key it is returned unchanged. If it uses -// the dot convention "{workloadID}.{originalCapabilityName}", the matching -// routing table key is returned. Falls back to returning toolName unchanged -// when the routing table is absent or the name cannot be resolved (pass-through -// semantics, consistent with the Router interface contract). +// toolName via ResolveToolRef. If toolName is an exact key it is returned +// unchanged. If it uses the dot convention "{workloadID}.{originalCapabilityName}", +// the matching routing table key is returned. Falls back to returning toolName +// unchanged when the routing table is absent or the name cannot be resolved +// (pass-through semantics, consistent with the Router interface contract). func (r *sessionRouter) ResolveToolName(_ context.Context, toolName string) string { - if r.routingTable == nil || r.routingTable.Tools == nil { - return toolName - } - - // Fast path: exact key match. - if _, exists := r.routingTable.Tools[toolName]; exists { - return toolName - } - - // Fallback: dot convention "{workloadID}.{toolName}". - if dotIdx := strings.Index(toolName, "."); dotIdx > 0 { - workloadID := toolName[:dotIdx] - capName := toolName[dotIdx+1:] - for resolvedName, target := range r.routingTable.Tools { - if target.WorkloadID == workloadID && target.GetBackendCapabilityName(resolvedName) == capName { - return resolvedName - } - } + if n, ok := ResolveToolRef(r.routingTable, toolName); ok { + return n } - return toolName } // RouteResource resolves a resource URI to its backend target using the // session's routing table directly. +// +// Resolution order: +// +// 1. Exact match against the aggregated concrete resources (the fast path, +// covering resources/list entries). +// +// 2. Template match: when no concrete resource matches, the URI is tested +// against the aggregated resource TEMPLATES (RFC 6570) and routed to the +// first template whose expansion matches. This lets a client read a +// templated resource (e.g. "file:///logs/2025-01-01.txt" matching +// "file:///logs/{date}.txt") through the ordinary resources/read path +// without a dedicated template read method. func (r *sessionRouter) RouteResource(_ context.Context, uri string) (*vmcp.BackendTarget, error) { - if r.routingTable == nil || r.routingTable.Resources == nil { + if r.routingTable == nil { return nil, fmt.Errorf("%w: %s", ErrResourceNotFound, uri) } - target, exists := r.routingTable.Resources[uri] - if !exists { - return nil, fmt.Errorf("%w: %s", ErrResourceNotFound, uri) + // Fast path: exact concrete-resource match. + if target, exists := r.routingTable.Resources[uri]; exists { + return target, nil } - return target, nil + // Exact template-string match: per the MCP spec a completion/complete with a + // ref/resource (and resources/subscribe on a templated resource) carries the + // URI TEMPLATE string itself (e.g. "file:///logs/{date}.txt"), not an + // expanded URI. A template does not match its own template string, so look + // the string up directly before falling back to template expansion. + if target, exists := r.routingTable.ResourceTemplates[uri]; exists { + return target, nil + } + // Fallback: match the URI against the aggregated resource templates. + if target := r.matchResourceTemplate(uri); target != nil { + return target, nil + } + return nil, fmt.Errorf("%w: %s", ErrResourceNotFound, uri) +} + +// matchResourceTemplate returns the backend target for the first precompiled +// resource template whose RFC 6570 expansion matches uri, or nil when none +// match. First-match wins over the sorted template keys (see +// compileResourceTemplates). +func (r *sessionRouter) matchResourceTemplate(uri string) *vmcp.BackendTarget { + for i := range r.resourceTemplates { + // Match returns non-nil Values on a match, nil otherwise. + if r.resourceTemplates[i].tmpl.Match(uri) != nil { + return r.resourceTemplates[i].target + } + } + return nil } // RoutePrompt resolves a prompt name to its backend target using the session's diff --git a/pkg/vmcp/router/session_router_test.go b/pkg/vmcp/router/session_router_test.go index 0b1d520891..fae275f420 100644 --- a/pkg/vmcp/router/session_router_test.go +++ b/pkg/vmcp/router/session_router_test.go @@ -265,6 +265,149 @@ func TestSessionRouter_RouteResource(t *testing.T) { } } +// TestSessionRouter_RouteResource_TemplateFallback exercises the template-match +// fallback: when the exact-URI lookup misses, the URI is matched against the +// aggregated resource templates (RFC 6570), routing to the matching template's +// backend. It covers a single match, no match, and disjoint multiple templates +// where exactly one matches (first-match selection). +func TestSessionRouter_RouteResource_TemplateFallback(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + routingTable *vmcp.RoutingTable + uri string + expectedID string + expectError bool + errorContains string + }{ + { + name: "URI matches a single template", + routingTable: &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{}, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///logs/{date}.txt": {WorkloadID: "logs-backend"}, + }, + }, + uri: "file:///logs/2025-01-01.txt", + expectedID: "logs-backend", + }, + { + // Per the MCP spec, a completion/complete with a ref/resource carries + // the URI TEMPLATE string itself, and a template does not match its + // own template string. The exact template-string key must route. + name: "exact template string routes to its backend", + routingTable: &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{}, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///logs/{date}.txt": {WorkloadID: "logs-backend"}, + }, + }, + uri: "file:///logs/{date}.txt", + expectedID: "logs-backend", + }, + { + name: "exact resource wins over a matching template", + routingTable: &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{ + "file:///logs/2025-01-01.txt": {WorkloadID: "exact-backend"}, + }, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///logs/{date}.txt": {WorkloadID: "template-backend"}, + }, + }, + uri: "file:///logs/2025-01-01.txt", + expectedID: "exact-backend", + }, + { + name: "URI matches no template", + routingTable: &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{}, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///logs/{date}.txt": {WorkloadID: "logs-backend"}, + }, + }, + uri: "db:///users/42", + expectError: true, + errorContains: "resource not found", + }, + { + name: "disjoint multiple templates: only the matching one is selected", + routingTable: &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{}, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///logs/{date}.txt": {WorkloadID: "logs-backend"}, + "db:///users/{id}": {WorkloadID: "users-backend"}, + }, + }, + uri: "db:///users/42", + expectedID: "users-backend", + }, + { + name: "invalid template is skipped, valid one still matches", + routingTable: &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{}, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///logs/{date}.txt": {WorkloadID: "logs-backend"}, + "://{bad": {WorkloadID: "broken-backend"}, + }, + }, + uri: "file:///logs/2025-01-01.txt", + expectedID: "logs-backend", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + r := router.NewSessionRouter(tt.routingTable) + target, err := r.RouteResource(t.Context(), tt.uri) + + if tt.expectError { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errorContains) + assert.Nil(t, target) + } else { + require.NoError(t, err) + require.NotNil(t, target) + assert.Equal(t, tt.expectedID, target.WorkloadID) + } + }) + } +} + +// TestSessionRouter_RouteResource_TemplateDeterminism locks in deterministic +// resolution when two OVERLAPPING templates match the same URI. Map iteration +// order is randomized per range, so before the sorted-key fix the winner could +// flip between runs; matchResourceTemplate now iterates template keys in sorted +// order, so the first-match winner is stable. "file:///{+path}" sorts before +// "file:///{a}/{b}" ('+' 0x2B < 'a' 0x61), so the greedy template wins every run. +func TestSessionRouter_RouteResource_TemplateDeterminism(t *testing.T) { + t.Parallel() + + const uri = "file:///logs/foo" + + // A single shared map is sufficient: Go randomizes iteration order on every + // range, so each RouteResource call below exercises a fresh order. + rt := &vmcp.RoutingTable{ + Resources: map[string]*vmcp.BackendTarget{}, + ResourceTemplates: map[string]*vmcp.BackendTarget{ + "file:///{+path}": {WorkloadID: "greedy-backend"}, + "file:///{a}/{b}": {WorkloadID: "specific-backend"}, + }, + } + + for i := range 50 { + r := router.NewSessionRouter(rt) + target, err := r.RouteResource(t.Context(), uri) + require.NoError(t, err) + require.NotNil(t, target) + assert.Equal(t, "greedy-backend", target.WorkloadID, + "overlapping templates must resolve to the sorted-first key deterministically (run %d)", i) + } +} + func TestSessionRouter_RoutePrompt(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/schema/reflect_test.go b/pkg/vmcp/schema/reflect_test.go index 7194b4f55a..65fc43af63 100644 --- a/pkg/vmcp/schema/reflect_test.go +++ b/pkg/vmcp/schema/reflect_test.go @@ -25,7 +25,7 @@ func TestGenerateSchema_FindToolInput(t *testing.T) { "tool_keywords": map[string]any{ "type": "array", "items": map[string]any{"type": "string"}, - "description": "Optional keywords for BM25 text search to narrow results (e.g. ['list', 'issues', 'github'] or ['SQL', 'query', 'postgres']). Combined with tool_description for hybrid search.", + "description": "Optional keywords driving the BM25 keyword-search arm (e.g. ['list', 'issues', 'github'] or ['SQL', 'query', 'postgres']). Semantic matching always uses tool_description, so provide a complete description even when supplying keywords.", }, }, "required": []string{"tool_description"}, diff --git a/pkg/vmcp/server/adapter/capability_adapter.go b/pkg/vmcp/server/adapter/capability_adapter.go index d02b2d3c6c..68463f9c06 100644 --- a/pkg/vmcp/server/adapter/capability_adapter.go +++ b/pkg/vmcp/server/adapter/capability_adapter.go @@ -8,9 +8,8 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" ) @@ -19,7 +18,7 @@ import ( // // This is the Anti-Corruption Layer between: // - Domain model (aggregator.AggregatedCapabilities) -// - External library (mark3labs/mcp-go SDK types) +// - External library (stacklok/toolhive-core/mcpcompat SDK types) // // The adapter: // 1. Converts aggregator types to SDK types diff --git a/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go b/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go index 1b85496f15..263fcea5fc 100644 --- a/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go +++ b/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go @@ -7,12 +7,12 @@ import ( "context" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter/mocks" diff --git a/pkg/vmcp/server/adapter/capability_adapter_test.go b/pkg/vmcp/server/adapter/capability_adapter_test.go index 71e0bf13d9..4150e0239c 100644 --- a/pkg/vmcp/server/adapter/capability_adapter_test.go +++ b/pkg/vmcp/server/adapter/capability_adapter_test.go @@ -7,12 +7,12 @@ import ( "context" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter/mocks" diff --git a/pkg/vmcp/server/adapter/handler_factory.go b/pkg/vmcp/server/adapter/handler_factory.go index ace09654d4..7581b68684 100644 --- a/pkg/vmcp/server/adapter/handler_factory.go +++ b/pkg/vmcp/server/adapter/handler_factory.go @@ -13,8 +13,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" "github.com/stacklok/toolhive/pkg/vmcp/internal/compositetools" @@ -93,8 +92,14 @@ func (f *DefaultHandlerFactory) CreateToolHandler( // Extract metadata from request to forward to backend meta := conversion.FromMCPMeta(request.Params.Meta) - // Call the backend tool - the backend client handles name translation and metadata forwarding - result, err := f.backendClient.CallTool(ctx, target, toolName, args, meta) + // Call the backend tool - the backend client handles name translation and metadata forwarding. + // + // No SEP-2243 Mcp-Param-* headers are mirrored on this path: deriving them + // needs the tool's inputSchema, and this factory holds only a router and a + // backend client. That is a limitation of this legacy adapter, not a policy + // choice — the live Serve path routes tools/call through the core, which has + // the aggregated schema and does mirror (see core_calls.go paramHeadersFor). + result, err := f.backendClient.CallTool(ctx, target, toolName, args, meta, nil) if err != nil { // Only actual network/transport errors reach here now (IsError=true is handled in result) if errors.Is(err, vmcp.ErrBackendUnavailable) { diff --git a/pkg/vmcp/server/adapter/handler_factory_test.go b/pkg/vmcp/server/adapter/handler_factory_test.go index f6a0b599c7..81b3c95b44 100644 --- a/pkg/vmcp/server/adapter/handler_factory_test.go +++ b/pkg/vmcp/server/adapter/handler_factory_test.go @@ -8,11 +8,11 @@ import ( "errors" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" vmcpmocks "github.com/stacklok/toolhive/pkg/vmcp/mocks" "github.com/stacklok/toolhive/pkg/vmcp/router" @@ -67,7 +67,7 @@ func TestDefaultHandlerFactory_CreateToolHandler(t *testing.T) { CallTool(gomock.Any(), target, "test_tool", map[string]any{ "input": "test", "count": 42, - }, gomock.Any()). + }, gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{StructuredContent: expectedResult}, nil) }, request: mcp.CallToolRequest{ @@ -172,7 +172,7 @@ func TestDefaultHandlerFactory_CreateToolHandler(t *testing.T) { Return(target, nil) mockClient.EXPECT(). - CallTool(gomock.Any(), target, "test_tool", map[string]any{"input": "test"}, gomock.Any()). + CallTool(gomock.Any(), target, "test_tool", map[string]any{"input": "test"}, gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ Content: []vmcp.Content{ {Type: vmcp.ContentTypeText, Text: "tool execution failed"}, @@ -206,7 +206,7 @@ func TestDefaultHandlerFactory_CreateToolHandler(t *testing.T) { Return(target, nil) mockClient.EXPECT(). - CallTool(gomock.Any(), target, "test_tool", map[string]any{"input": "test"}, gomock.Any()). + CallTool(gomock.Any(), target, "test_tool", map[string]any{"input": "test"}, gomock.Any(), gomock.Any()). Return(nil, vmcp.ErrBackendUnavailable) }, request: mcp.CallToolRequest{ @@ -235,7 +235,7 @@ func TestDefaultHandlerFactory_CreateToolHandler(t *testing.T) { Return(target, nil) mockClient.EXPECT(). - CallTool(gomock.Any(), target, "test_tool", map[string]any{"input": "test"}, gomock.Any()). + CallTool(gomock.Any(), target, "test_tool", map[string]any{"input": "test"}, gomock.Any(), gomock.Any()). Return(nil, errors.New("unknown backend error")) }, request: mcp.CallToolRequest{ @@ -269,7 +269,7 @@ func TestDefaultHandlerFactory_CreateToolHandler(t *testing.T) { // Handler factory now passes the client-facing name (backend1_fetch) // Backend client handles translation to original name (fetch) mockClient.EXPECT(). - CallTool(gomock.Any(), target, "backend1_fetch", map[string]any{"url": "https://example.com"}, gomock.Any()). + CallTool(gomock.Any(), target, "backend1_fetch", map[string]any{"url": "https://example.com"}, gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{StructuredContent: expectedResult}, nil) }, request: mcp.CallToolRequest{ diff --git a/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go b/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go index 5d516c2177..53862acf60 100644 --- a/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go +++ b/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go @@ -13,7 +13,7 @@ import ( context "context" reflect "reflect" - mcp "github.com/mark3labs/mcp-go/mcp" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" adapter "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" gomock "go.uber.org/mock/gomock" ) diff --git a/pkg/vmcp/server/authz_integration_test.go b/pkg/vmcp/server/authz_integration_test.go index 20a8e18405..a7580e8dd1 100644 --- a/pkg/vmcp/server/authz_integration_test.go +++ b/pkg/vmcp/server/authz_integration_test.go @@ -4,11 +4,13 @@ package server_test import ( - "bytes" - "context" + "encoding/json" "io" "net/http" "net/http/httptest" + "os" + "path/filepath" + "strings" "testing" "time" @@ -16,21 +18,48 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/aggregator" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" + "github.com/stacklok/toolhive/pkg/vmcp/codemode" "github.com/stacklok/toolhive/pkg/vmcp/mocks" "github.com/stacklok/toolhive/pkg/vmcp/router" "github.com/stacklok/toolhive/pkg/vmcp/server" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) +// rpcErrorFields is the subset of a JSON-RPC error envelope the authz tests assert on. +type rpcErrorFields struct { + code int + message string + present bool +} + +// parseRPCError extracts the JSON-RPC error code/message from a response body. It +// reports present=false when the body carries no top-level "error" object (e.g. a +// successful result). +func parseRPCError(t *testing.T, body []byte) rpcErrorFields { + t.Helper() + var env struct { + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(body, &env); err != nil || env.Error == nil { + return rpcErrorFields{} + } + return rpcErrorFields{code: env.Error.Code, message: env.Error.Message, present: true} +} + // newCedarAuthzTestServer builds a vMCP server via the rerouted server.New (→ core.New + // Serve) backed by the real "echo" MCP backend at backendURL. An AuthMiddleware injects a // fixed authenticated principal (the Cedar authorizer must resolve one — a nil identity @@ -39,6 +68,28 @@ import ( // chain that replaced the legacy HTTP authz middleware on the Serve path. func newCedarAuthzTestServer(t *testing.T, backendURL string, policies ...string) *httptest.Server { t.Helper() + return buildCedarAuthzServer(t, backendURL, nil, nil, nil, policies...) +} + +// newCedarAuthzCodeModeServer is newCedarAuthzTestServer with code mode enabled, so +// the core is wrapped by the codemode decorator. It proves the pre-dispatch gate's +// codemode carve-out: execute_tool_script is admitted (the feature flag is the grant) +// while a directly-denied tool still 403s. +func newCedarAuthzCodeModeServer(t *testing.T, backendURL string, policies ...string) *httptest.Server { + t.Helper() + return buildCedarAuthzServer(t, backendURL, &codemode.Config{}, nil, nil, policies...) +} + +// buildCedarAuthzServer builds the vMCP test server. A non-nil codeModeCfg enables +// the codemode decorator; a non-nil auditCfg enables the audit middleware; a +// non-nil authMw replaces the default identity-injecting auth middleware (used by +// the auth-failure audit test); a nil policies slice leaves Authz unset +// (allow-all, gate not installed) — used by the no-Authz parity guard. +func buildCedarAuthzServer( + t *testing.T, backendURL string, codeModeCfg *codemode.Config, auditCfg *audit.Config, + authMw func(http.Handler) http.Handler, policies ...string, +) *httptest.Server { + t.Helper() ctrl := gomock.NewController(t) t.Cleanup(ctrl.Finish) @@ -69,28 +120,50 @@ func newCedarAuthzTestServer(t *testing.T, backendURL string, policies ...string require.NoError(t, err) agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) - // Inject a fixed authenticated identity on every request so the session binds to it at + // Inject an authenticated identity on every request so the session binds to it at // initialize and the Cedar authorizer can resolve the principal on subsequent calls. - identityMiddleware := func(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ - Subject: "user-123", - Name: "Test User", - Claims: map[string]any{"sub": "user-123", "name": "Test User"}, - }} - next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) - }) + // The principal defaults to the fixed "user-123" (all existing authz tests rely on + // this), but a request carrying X-Test-Principal derives the identity from that header + // instead — letting a single running server authenticate different callers as distinct + // principals (see TestRegression_PerSessionToolProjection_DivergesByPrincipal), without + // changing behavior for any test that never sets the header. The MCP parser is composed + // inside it, mirroring the production incoming-auth factory (see pkg/vmcp/auth/factory): + // audit and authz read parsed MCP data from the request context, and the audit + // middleware sits between auth and the parser applied in Handler. + identityMiddleware := authMw + if identityMiddleware == nil { + identityMiddleware = func(next http.Handler) http.Handler { + withParser := mcpparser.ParsingMiddleware(next) + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + principal := r.Header.Get("X-Test-Principal") + if principal == "" { + principal = "user-123" + } + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: principal, + Name: "Test User", + Claims: map[string]any{"sub": principal, "name": "Test User"}, + }} + withParser.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) + } } - authzCfg, err := authorizers.NewConfig(cedar.Config{ - Version: "1.0", - Type: cedar.ConfigType, - Options: &cedar.ConfigOptions{Policies: policies, EntitiesJSON: "[]"}, - }) - require.NoError(t, err) + // A nil policies slice means "no authz": leave Config.Authz nil so the gate is not + // installed (the no-Authz parity guard exercises this). Otherwise compile the Cedar + // policies into the config the core admission seam consumes. + var authzCfg *authorizers.Config + if policies != nil { + authzCfg, err = authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{Policies: policies, EntitiesJSON: "[]"}, + }) + require.NoError(t, err) + } srv, err := server.New( - context.Background(), + t.Context(), &server.Config{ // Name is non-empty: deriveCoreConfig forwards the raw server name to the core, // and the Cedar admission seam requires it (resource entities are scoped to @@ -103,6 +176,8 @@ func newCedarAuthzTestServer(t *testing.T, backendURL string, policies ...string Aggregator: agg, AuthMiddleware: identityMiddleware, Authz: authzCfg, + AuditConfig: auditCfg, + CodeModeConfig: codeModeCfg, }, router.NewSessionRouter(&vmcp.RoutingTable{}), backendClient, @@ -111,7 +186,7 @@ func newCedarAuthzTestServer(t *testing.T, backendURL string, policies ...string ) require.NoError(t, err) - handler, err := srv.Handler(context.Background()) + handler, err := srv.Handler(t.Context()) require.NoError(t, err) ts := httptest.NewServer(handler) @@ -139,19 +214,21 @@ func TestIntegration_RealBackend_CedarAuthzGatesToolCall(t *testing.T) { client := NewMCPTestClient(t, ts.URL) client.InitializeSession() - // "echo" is filtered out, so there is no list signal to wait on; poll the call until - // the session is established and the SDK rejects the unregistered tool (transient - // pre-registration 404s are tolerated and retried). - var lastBody string - require.Eventually(t, func() bool { - resp := client.CallTool("echo", map[string]any{"input": "hello"}) - defer resp.Body.Close() - b, _ := io.ReadAll(resp.Body) - lastBody = string(b) - return resp.StatusCode == http.StatusOK && bytes.Contains(b, []byte("not found")) - }, 5*time.Second, 50*time.Millisecond, - "a tool filtered out by the Cedar policy must be rejected as an unknown tool") - assert.NotContains(t, lastBody, "hello", "a filtered tool must never reach the backend") + // #5827: a policy-denied direct tools/call is now rejected by the pre-dispatch + // gate as HTTP 403 + JSON-RPC error code 403 — NOT the old 200 + SDK "not found" + // that leaked filtered-vs-nonexistent. The gate decision is deterministic (it + // re-runs core admission), so no polling is needed. + resp := client.CallTool("echo", map[string]any{"input": "hello"}) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, resp.StatusCode, + "a tool denied by policy must be rejected with HTTP 403, body: %s", string(body)) + rpcErr := parseRPCError(t, body) + require.True(t, rpcErr.present, "the 403 must carry a JSON-RPC error envelope, body: %s", string(body)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, rpcErr.code, "the JSON-RPC error code must be 403") + assert.Equal(t, "call denied by authorization policy", rpcErr.message) + assert.NotContains(t, string(body), "hello", "a filtered tool must never reach the backend") // And it must be absent from tools/list. listResp := client.ListTools() @@ -187,16 +264,469 @@ func TestIntegration_RealBackend_CedarAuthzGatesToolCall(t *testing.T) { require.Equal(t, http.StatusOK, okResp.StatusCode, "body: %s", string(okBody)) assert.Contains(t, string(okBody), "hello", "a permitted call must return the backend result") - // Denied call: the core admission seam rejects it with a generic message that never - // leaks the underlying authorizer/policy detail. + // Denied call: #5827 — the pre-dispatch gate rejects the arg-gated denial as + // HTTP 403 + JSON-RPC error code 403 (NOT the old 200 + IsError), with a generic + // message that never leaks the authorizer/policy detail. denyResp := client.CallTool("echo", map[string]any{"input": "blocked"}) defer denyResp.Body.Close() denyBody, err := io.ReadAll(denyResp.Body) require.NoError(t, err) - require.Equal(t, http.StatusOK, denyResp.StatusCode, "body: %s", string(denyBody)) - assert.Contains(t, string(denyBody), "call denied by authorization policy", + require.Equal(t, http.StatusForbidden, denyResp.StatusCode, "body: %s", string(denyBody)) + rpcErr := parseRPCError(t, denyBody) + require.True(t, rpcErr.present, "the 403 must carry a JSON-RPC error envelope, body: %s", string(denyBody)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, rpcErr.code, "the JSON-RPC error code must be 403") + assert.Equal(t, "call denied by authorization policy", rpcErr.message, "an arg-gated policy denial must surface the genericized authorization message") assert.NotContains(t, string(denyBody), "cedar", "the authorizer implementation detail must not leak") assert.NotContains(t, string(denyBody), "forbid", "the policy text must not leak") }) } + +// TestIntegration_RealBackend_CedarAuthz_NoEnumerationOracle proves the denial carries +// no information about whether the named tool exists: under a default-deny policy a +// nonexistent tool is rejected with the SAME 403 + message as a real-but-denied tool. +// A separate permissive policy documents the flip side — when the name IS permitted, +// the gate admits and the SDK's own -32602 "not found" surfaces at 200, so authz never +// masks nonexistence for an allowed name. +func TestIntegration_RealBackend_CedarAuthz_NoEnumerationOracle(t *testing.T) { + t.Parallel() + + t.Run("nonexistent tool under default-deny is an identical 403", func(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newCedarAuthzTestServer(t, backendURL, + `permit(principal, action == Action::"call_tool", resource == Tool::"unrelated");`) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + resp := client.CallTool("does-not-exist", map[string]any{"input": "x"}) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, resp.StatusCode, "body: %s", string(body)) + rpcErr := parseRPCError(t, body) + require.True(t, rpcErr.present, "body: %s", string(body)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, rpcErr.code) + assert.Equal(t, "call denied by authorization policy", rpcErr.message, + "a nonexistent tool must be denied identically to a real one (no enumeration oracle)") + assert.NotContains(t, string(body), "does-not-exist", "the tool name must not leak in the denial") + }) + + t.Run("permitted nonexistent name yields the SDK's -32602 at 200", func(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + // Permit call_tool on any resource: the gate admits, so a nonexistent name reaches + // the SDK and gets the standard invalid-params "not found" — authz does not mask it. + ts := newCedarAuthzTestServer(t, backendURL, + `permit(principal, action == Action::"call_tool", resource);`) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + resp := client.CallTool("ghost-tool", map[string]any{"input": "x"}) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, + "a permitted-but-nonexistent name must not be gated, body: %s", string(body)) + rpcErr := parseRPCError(t, body) + require.True(t, rpcErr.present, "body: %s", string(body)) + assert.Equal(t, -32602, rpcErr.code, "the SDK's INVALID_PARAMS code must surface for an unknown tool") + }) +} + +// TestIntegration_RealBackend_CedarAuthz_ResourceAndPromptDenied proves the gate covers +// resources/read and prompts/get, not just tools/call: a denied read/get is rejected as +// HTTP 403 + JSON-RPC 403 with the kind-specific message. +func TestIntegration_RealBackend_CedarAuthz_ResourceAndPromptDenied(t *testing.T) { + t.Parallel() + + // Permit only an unrelated tool so every resource/prompt is default-denied. + const denyAllReadsAndPrompts = `permit(principal, action == Action::"call_tool", resource == Tool::"unrelated");` + + t.Run("resources/read denied", func(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newCedarAuthzTestServer(t, backendURL, denyAllReadsAndPrompts) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + resp := client.ReadResource("file://secret") + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, resp.StatusCode, "body: %s", string(body)) + rpcErr := parseRPCError(t, body) + require.True(t, rpcErr.present, "body: %s", string(body)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, rpcErr.code) + assert.Equal(t, "read denied by authorization policy", rpcErr.message) + }) + + t.Run("prompts/get denied", func(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newCedarAuthzTestServer(t, backendURL, denyAllReadsAndPrompts) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + resp := client.GetPrompt("secret-prompt") + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, resp.StatusCode, "body: %s", string(body)) + rpcErr := parseRPCError(t, body) + require.True(t, rpcErr.present, "body: %s", string(body)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, rpcErr.code) + assert.Equal(t, "prompt denied by authorization policy", rpcErr.message) + }) +} + +// TestIntegration_RealBackend_NoAuthz_UnknownToolIsSDKError is the byte-parity guard: +// with no Authz configured the gate is NOT installed, so an unknown tool falls through +// to the SDK and gets -32602 at 200 — exactly today's behavior, unchanged by the gate. +func TestIntegration_RealBackend_NoAuthz_UnknownToolIsSDKError(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + // nil policies (no variadic args) ⇒ Config.Authz stays nil ⇒ no gate. + ts := newCedarAuthzTestServer(t, backendURL) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + resp := client.CallTool("unknown-tool", map[string]any{"input": "x"}) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, + "without Authz the gate is absent; the SDK handles unknown tools at 200, body: %s", string(body)) + rpcErr := parseRPCError(t, body) + require.True(t, rpcErr.present, "body: %s", string(body)) + assert.Equal(t, -32602, rpcErr.code, "the SDK's INVALID_PARAMS code must surface unchanged") +} + +// TestIntegration_RealBackend_CodeMode_Authz is the regression guard for the codemode +// carve-out: with both code mode and Authz enabled, the execute_tool_script meta-tool is +// admitted by the gate (200) — it is not in the admission seam and would otherwise be +// denied under a permit-list policy — while a directly-called denied tool still 403s. +func TestIntegration_RealBackend_CodeMode_Authz(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + // Permit "echo" (so a script can call it) but nothing grants the execute_tool_script + // meta-tool itself — the gate must admit it via the codemode carve-out, not policy. + ts := newCedarAuthzCodeModeServer(t, backendURL, + `permit(principal, action == Action::"call_tool", resource == Tool::"echo");`) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + waitForEchoTool(t, ts.URL, client.SessionID()) + + // execute_tool_script is admitted by the gate and runs the script (which calls the + // permitted "echo"), returning 200 — NOT a 403. + scriptResp := client.CallTool("execute_tool_script", map[string]any{ + "script": `return echo(input="hi from script")`, + }) + defer scriptResp.Body.Close() + scriptBody, err := io.ReadAll(scriptResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, scriptResp.StatusCode, + "execute_tool_script must be admitted by the gate's codemode carve-out, body: %s", string(scriptBody)) + assert.NotEqual(t, mcpparser.JSONRPCCodeDenied, parseRPCError(t, scriptBody).code, + "the code-mode meta-tool must never be 403'd") + // Prove the script actually RAN (not a 200 + IsError): the echo backend returns its + // input verbatim, so the script's return value must carry "hi from script". + assert.Contains(t, string(scriptBody), "hi from script", + "the admitted script must execute and return the permitted tool's output") + + // A directly-called denied tool is still gated to 403. + denyResp := client.CallTool("unrelated", map[string]any{"input": "x"}) + defer denyResp.Body.Close() + denyBody, err := io.ReadAll(denyResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, denyResp.StatusCode, "body: %s", string(denyBody)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, parseRPCError(t, denyBody).code) +} + +// TestIntegration_CedarAuthzDenialIsAudited proves that on the vMCP Serve path a +// policy-denied tools/call still produces an audit event with outcome "denied". +// The pre-dispatch authorization gate runs inside the audit middleware, so the +// 403 it writes must be captured as the event outcome. This is the vMCP +// counterpart of the proxyrunner-path guard in pkg/runner +// (TestAuthzDecisionIsAudited). +func TestIntegration_CedarAuthzDenialIsAudited(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + auditLogPath := filepath.Join(t.TempDir(), "audit.log") + // Permit only an unrelated tool: "echo" is default-denied. + ts := buildCedarAuthzServer(t, backendURL, nil, + &audit.Config{Component: "vmcp-server", LogFile: auditLogPath}, + nil, + `permit(principal, action == Action::"call_tool", resource == Tool::"unrelated");`) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + resp := client.CallTool("echo", map[string]any{"input": "hello"}) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, resp.StatusCode, "body: %s", string(body)) + + // The audit event is written after the response is flushed, so poll briefly. + require.Eventually(t, func() bool { + return findToolCallAuditEvent(t, auditLogPath) != nil + }, 5*time.Second, 50*time.Millisecond, "a tools/call audit event must be emitted for the denied call") + + event := findToolCallAuditEvent(t, auditLogPath) + assert.Equal(t, "denied", event["outcome"], + "a policy-denied tools/call must be audited with outcome denied") + subjects, ok := event["subjects"].(map[string]any) + require.True(t, ok, "the event must carry subjects") + assert.Equal(t, "user-123", subjects["user_id"], + "audit wraps auth, so the identity must flow back via the auth.IdentityHolder carrier") +} + +// TestIntegration_AuthFailureIsAudited proves that an authentication failure +// (401 from the auth middleware) still produces an audit event: audit wraps +// auth on the vMCP Serve path, so rejected requests are recorded with outcome +// "denied" and an anonymous subject. +func TestIntegration_AuthFailureIsAudited(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + auditLogPath := filepath.Join(t.TempDir(), "audit.log") + rejectAll := func(http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "invalid token", http.StatusUnauthorized) + }) + } + ts := buildCedarAuthzServer(t, backendURL, nil, + &audit.Config{Component: "vmcp-server", LogFile: auditLogPath}, + rejectAll) + + resp, err := http.Post(ts.URL+"/mcp", "application/json", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusUnauthorized, resp.StatusCode) + + require.Eventually(t, func() bool { + return findAuditEvent(t, auditLogPath, "http_request") != nil + }, 5*time.Second, 50*time.Millisecond, "an authentication failure must be audited") + + event := findAuditEvent(t, auditLogPath, "http_request") + assert.Equal(t, "denied", event["outcome"]) + subjects, ok := event["subjects"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "anonymous", subjects["user"], + "no identity exists when authentication fails") +} + +// TestIntegration_CedarAuthzDenial_ModernPath_IsAudited is the Modern (2026-07-28) +// counterpart of TestIntegration_CedarAuthzDenialIsAudited. The Modern stateless +// dispatcher bypasses the SDK server (and its CallGate), so it re-homes the +// pre-dispatch authorization gate itself; this proves that gate end-to-end through +// the fully assembled server rather than only in the dispatchModern unit tests. A +// policy-denied Modern tools/call must surface as HTTP 403 + JSON-RPC 403 AND be +// audited with outcome "denied" (audit -> parsing -> classification -> +// dispatchModern), confirming the audit middleware wraps the re-homed gate exactly +// as it wraps the SDK CallGate on the Legacy path. +func TestIntegration_CedarAuthzDenial_ModernPath_IsAudited(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + auditLogPath := filepath.Join(t.TempDir(), "audit.log") + // Permit only an unrelated tool: "echo" is default-denied, so the re-homed gate + // in dispatchModern rejects the call before it reaches the backend. + ts := buildCedarAuthzServer(t, backendURL, nil, + &audit.Config{Component: "vmcp-server", LogFile: auditLogPath}, nil, + `permit(principal, action == Action::"call_tool", resource == Tool::"unrelated");`) + + resp, decoded := postModern(t, ts.URL, "tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "hello modern"}, + }, 1, "echo") + defer resp.Body.Close() + + require.Equal(t, http.StatusForbidden, resp.StatusCode, "decoded: %+v", decoded) + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "the 403 must carry a JSON-RPC error envelope: %+v", decoded) + assert.EqualValues(t, mcpparser.JSONRPCCodeDenied, errObj["code"], + "a policy-denied Modern tools/call must surface JSON-RPC 403") + + // The audit event is written after the response is flushed, so poll briefly. + require.Eventually(t, func() bool { + return findToolCallAuditEvent(t, auditLogPath) != nil + }, 5*time.Second, 50*time.Millisecond, + "a tools/call audit event must be emitted for the denied Modern call") + + event := findToolCallAuditEvent(t, auditLogPath) + assert.Equal(t, "denied", event["outcome"], + "a policy-denied Modern tools/call must be audited with outcome denied") +} + +// waitForClientTool polls c.ListTools until toolName appears in the advertised +// set, or the deadline elapses. The poll condition is best-effort: any +// transport/read error returns false so the poll simply retries, and it never +// calls require/FailNow inside the closure (testify runs the condition on its +// own goroutine, where FailNow would Goexit off the test goroutine, hang the +// poll to its ceiling, and misreport the failure). +func waitForClientTool(t *testing.T, c *MCPTestClient, toolName string) { + t.Helper() + require.Eventually(t, func() bool { + resp := c.ListTools() + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return false + } + return resp.StatusCode == http.StatusOK && strings.Contains(string(body), `"`+toolName+`"`) + }, 5*time.Second, 50*time.Millisecond, + "client's tools/list must advertise "+toolName) +} + +// listToolsBody returns the client's current tools/list response body as a +// string. Unlike waitForClientTool's poll closure, this runs on the test +// goroutine, so requiring a clean read/200 here is safe. +func listToolsBody(t *testing.T, c *MCPTestClient) string { + t.Helper() + resp := c.ListTools() + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "tools/list must succeed") + return string(body) +} + +// TestRegression_PerSessionToolProjection_DivergesByPrincipal pins the fixed +// per-principal Cedar admission behind the mcp-go -> go-sdk migration +// (toolhive-core v0.0.32, #5742): two clients authenticated as DIFFERENT +// principals, against the SAME running server instance, must get divergent +// tools/list projections when a Cedar policy permits a tool for one principal +// only. Before the fix, the server's identity middleware stamped every request +// with a single fixed principal, so per-session admission could not diverge by +// caller — this test would have failed (or passed vacuously) against that +// regression. +func TestRegression_PerSessionToolProjection_DivergesByPrincipal(t *testing.T) { + t.Parallel() + + backendURL := startEchoPingBackend(t) + // The Cedar principal for a JWT-claims-derived caller is Client::"" + // (cedar authorizeToolCall). Give each principal a DISTINCT positively- + // permitted tool -- alice may call echo, bob may call ping -- so the two + // sessions' tools/list projections (gated by these call_tool permissions) + // must diverge in BOTH directions on the same running server. + ts := newCedarAuthzTestServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"echo");`, + `permit(principal == Client::"bob", action == Action::"call_tool", resource == Tool::"ping");`) + + alice := NewMCPTestClient(t, ts.URL).WithPrincipal("alice") + alice.InitializeSession() + bob := NewMCPTestClient(t, ts.URL).WithPrincipal("bob") + bob.InitializeSession() + + // Wait until EACH session has registered its OWN permitted tool. This is + // the anti-vacuity guard: it proves both sessions finished their + // (asynchronous) registration and each principal resolved to its own + // identity. Without it, an empty tools/list (registration still pending) + // would be indistinguishable from a correctly-filtered one, so a real + // cross-principal projection leak could pass unnoticed. + waitForClientTool(t, alice, "echo") + waitForClientTool(t, bob, "ping") + + // Divergence in both directions: alice sees echo but NOT bob's ping; bob + // sees ping but NOT alice's echo. A regression that collapsed all sessions + // onto one fixed principal, or shared a cached projection across principals, + // would leak the other principal's tool here. + aliceTools := listToolsBody(t, alice) + assert.Contains(t, aliceTools, `"echo"`, "alice must see her permitted tool echo") + assert.NotContains(t, aliceTools, `"ping"`, + "alice must NOT see bob's tool ping on the SAME server instance") + + bobTools := listToolsBody(t, bob) + assert.Contains(t, bobTools, `"ping"`, "bob must see his permitted tool ping") + assert.NotContains(t, bobTools, `"echo"`, + "bob must NOT see alice's tool echo on the SAME server instance") +} + +// TestRegression_PerSessionToolCall_DeniedForUnprivilegedPrincipal is the +// call-path counterpart of TestRegression_PerSessionToolProjection_DivergesByPrincipal: +// on the SAME running server, a tools/call for a tool permitted to one +// principal (alice) must succeed, while the identical call from a different +// principal (bob) on their own concurrently-registered session must be denied +// by the pre-dispatch authorization gate (403), not silently allowed through a +// shared/fixed identity. +func TestRegression_PerSessionToolCall_DeniedForUnprivilegedPrincipal(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newCedarAuthzTestServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"echo");`) + + alice := NewMCPTestClient(t, ts.URL).WithPrincipal("alice") + alice.InitializeSession() + bob := NewMCPTestClient(t, ts.URL).WithPrincipal("bob") + bob.InitializeSession() + + // Wait for alice's session to register echo before calling it. + waitForClientTool(t, alice, "echo") + + aliceResp := alice.CallTool("echo", map[string]any{"input": "hello"}) + defer aliceResp.Body.Close() + aliceBody, err := io.ReadAll(aliceResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, aliceResp.StatusCode, + "alice is permitted to call echo; body: %s", string(aliceBody)) + assert.Contains(t, string(aliceBody), "hello", "alice's permitted call must reach the real backend") + + bobResp := bob.CallTool("echo", map[string]any{"input": "hello"}) + defer bobResp.Body.Close() + bobBody, err := io.ReadAll(bobResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusForbidden, bobResp.StatusCode, + "bob is not permitted to call echo on the SAME running server; body: %s", string(bobBody)) + rpcErr := parseRPCError(t, bobBody) + require.True(t, rpcErr.present, "the 403 must carry a JSON-RPC error envelope, body: %s", string(bobBody)) + assert.Equal(t, mcpparser.JSONRPCCodeDenied, rpcErr.code, + "bob's denied call must surface the genericized authorization JSON-RPC error code") +} + +// findToolCallAuditEvent reads the newline-delimited JSON audit log at path and +// returns the first "mcp_tool_call" event, or nil if none is present yet. +func findToolCallAuditEvent(t *testing.T, path string) map[string]any { + t.Helper() + return findAuditEvent(t, path, "mcp_tool_call") +} + +// findAuditEvent reads the newline-delimited JSON audit log at path and returns +// the first event whose "type" matches eventType, or nil if none is present yet. +func findAuditEvent(t *testing.T, path string, eventType string) map[string]any { + t.Helper() + + data, err := os.ReadFile(path) + if err != nil { + return nil + } + for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") { + if line == "" { + continue + } + var event map[string]any + if err := json.Unmarshal([]byte(line), &event); err != nil { + continue + } + if event["type"] == eventType { + return event + } + } + return nil +} diff --git a/pkg/vmcp/server/bridge_regression_test.go b/pkg/vmcp/server/bridge_regression_test.go new file mode 100644 index 0000000000..93a13b6d2a --- /dev/null +++ b/pkg/vmcp/server/bridge_regression_test.go @@ -0,0 +1,811 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// --------------------------------------------------------------------------- +// This file pins the two cross-era "bridge cells" of vMCP's dual-protocol +// dispatch: Modern client -> Legacy backend (Cell A) and Legacy client -> +// Modern backend (Cell B). Both already work; nothing here changes production +// behavior. RFC THV-0083 D6 plans a shared backend connection pool, which +// would silently break Cell A's per-principal isolation if it ever collapsed +// distinct backend sessions -- these tests exist so that change cannot land +// without noticing. +// +// Scope: these pins cover the TRANSPORT/SESSION dimension of isolation only +// (distinct backend connections, no cross-request credential bleed on the +// wire). Credential *derivation* caching (e.g. token-exchange result caching, +// pkg/vmcp/auth/strategies/tokenexchange.go) is a different layer and is not +// exercised here. +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Cell A: Modern client -> Legacy backend +// --------------------------------------------------------------------------- + +// testIdentityEchoStrategyName is the outgoing-auth strategy name registered +// by identityEchoStrategy below. +const testIdentityEchoStrategyName = "test-identity-echo" + +// identityEchoStrategy stamps the calling principal's identity onto every +// outgoing backend request as an Authorization header, so a test can observe +// -- on the wire, via newRecordingBackendProxy -- which principal's +// credential reached the backend on each individual request. The default +// UnauthenticatedStrategy (used by every other real-backend test in this +// package) injects nothing, so there would be nothing to assert; this is the +// minimal seam for exercising the real property under test (the *calling* +// principal's identity reaches outgoing auth on every re-originated call) +// without needing real IdP infrastructure. +type identityEchoStrategy struct{} + +func (identityEchoStrategy) Name() string { return testIdentityEchoStrategyName } + +func (identityEchoStrategy) Authenticate(ctx context.Context, req *http.Request, _ *authtypes.BackendAuthStrategy) error { + if identity, ok := auth.IdentityFromContext(ctx); ok && identity != nil { + req.Header.Set("Authorization", "Bearer test-"+identity.Subject) + } + return nil +} + +func (identityEchoStrategy) Validate(_ *authtypes.BackendAuthStrategy) error { return nil } + +// bridgePrincipalMiddleware stamps a per-request auth.Identity derived from +// the X-Test-Principal header. It duplicates the small stamping technique in +// principalIdentityMiddleware (context_isolation_regression_test.go) rather +// than importing it: that helper is unexported in package server, and this +// file is package server_test so it can reach the real-backend helpers in +// testutil_test.go. +func bridgePrincipalMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + principal := r.Header.Get("X-Test-Principal") + if principal == "" { + next.ServeHTTP(w, r) + return + } + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: principal, Claims: map[string]any{"sub": principal}}} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) +} + +// recordedBackendRequest captures one HTTP request newRecordingBackendProxy +// forwarded to the real backend, plus the backend's response session id (when +// present), so a test can correlate an initialize call with the tools/call +// that reused its backend session via Mcp-Session-Id. +type recordedBackendRequest struct { + jsonRPCMethod string + authorization string + reqSessionID string + respSessionID string + body map[string]any +} + +// newRecordingBackendProxy starts an httptest.Server that transparently +// forwards every request to backendURL while recording it. This gives the +// test real MCP semantics plus full request observability, without +// hand-rolling a fake MCP server (per the task's design suggestion). Returns +// the proxy's /mcp URL and an accessor for a snapshot of everything recorded +// so far. +// flushWriter flushes the underlying http.ResponseWriter after every Write, +// so a streamed response (the standalone SSE GET below) reaches the client +// as it arrives instead of sitting in net/http's default write buffer. +type flushWriter struct{ w http.ResponseWriter } + +func (fw flushWriter) Write(p []byte) (int, error) { + n, err := fw.w.Write(p) + if f, ok := fw.w.(http.Flusher); ok { + f.Flush() + } + return n, err +} + +func newRecordingBackendProxy(t *testing.T, backendURL string) (string, func() []recordedBackendRequest) { + t.Helper() + + var mu sync.Mutex + var records []recordedBackendRequest + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + bodyBytes, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + _ = r.Body.Close() + + var decoded map[string]any + _ = json.Unmarshal(bodyBytes, &decoded) + + fwdReq, err := http.NewRequestWithContext(r.Context(), r.Method, backendURL, bytes.NewReader(bodyBytes)) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + fwdReq.Header = r.Header.Clone() + + resp, err := http.DefaultClient.Do(fwdReq) + if err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + + method, _ := decoded["method"].(string) + mu.Lock() + records = append(records, recordedBackendRequest{ + jsonRPCMethod: method, + authorization: r.Header.Get("Authorization"), + reqSessionID: r.Header.Get("Mcp-Session-Id"), + respSessionID: resp.Header.Get("Mcp-Session-Id"), + body: decoded, + }) + mu.Unlock() + + for k, vs := range resp.Header { + for _, v := range vs { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + // Flush the status line/headers immediately: the standalone SSE GET below + // can sit with headers sent but no body for a while (streamable.go's + // hangResponse), and the mcpcompat client's connectSSE blocks until it + // receives the response headers -- net/http's default buffered writer + // would otherwise hold them back indefinitely. + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + // Stream and flush every chunk, don't buffer: after `initialize`, the + // mcpcompat client opens that persistent standalone SSE GET, which the + // real backend intentionally never closes. Buffering the full response + // body before writing anything back would block forever on it, and a + // plain io.Copy alone would still hold each event back in the same + // buffered writer. flushWriter flushes after every write. + _, _ = io.Copy(flushWriter{w}, resp.Body) + }) + + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + + getRecords := func() []recordedBackendRequest { + mu.Lock() + defer mu.Unlock() + out := make([]recordedBackendRequest, len(records)) + copy(out, records) + return out + } + return ts.URL + "/mcp", getRecords +} + +// newBridgeCellAServer builds a vMCP server routing a single Legacy backend +// (at proxyURL) through identityEchoStrategy so the credential that reaches +// the backend on each call is observable. Mirrors newRealTestHandler's construction +// (session_management_realbackend_integration_test.go), diverging only in +// the backend's AuthConfig -- that helper always uses "unauthenticated", +// which would leave nothing for this file's tests to assert. +func newBridgeCellAServer(t *testing.T, proxyURL string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: proxyURL, + TransportType: "streamable-http", + AuthConfig: &authtypes.BackendAuthStrategy{Type: testIdentityEchoStrategyName}, + } + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, authReg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, strategies.NewUnauthenticatedStrategy())) + require.NoError(t, authReg.RegisterStrategy(testIdentityEchoStrategyName, identityEchoStrategy{})) + + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + rt := router.NewSessionRouter(&vmcp.RoutingTable{}) + srv, err := server.New( + context.Background(), + &server.Config{ + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: vmcpsession.NewSessionFactory(authReg), + Aggregator: agg, + }, + rt, + backendClient, + mockBackendRegistry, + nil, + ) + require.NoError(t, err) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + + ts := httptest.NewServer(bridgePrincipalMiddleware(handler)) + t.Cleanup(ts.Close) + return ts +} + +// postModernToolCallAsPrincipal sends a Modern (2026-07-28) tools/call for +// echo, stamped with X-Test-Principal so bridgePrincipalMiddleware +// authenticates the caller as principal, and carrying the reserved +// io.modelcontextprotocol/* _meta keys a real Modern client sends. Mirrors +// postModern (modern_realbackend_integration_test.go), which has no hook for +// extra headers -- needed here for the principal header. +func postModernToolCallAsPrincipal(t *testing.T, baseURL, principal, input string) *http.Response { + t.Helper() + + payload, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": input}, + "_meta": map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": map[string]any{}, + }, + }, + }) + require.NoError(t, err) + + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, baseURL+"/mcp", bytes.NewReader(payload)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", "2026-07-28") + req.Header.Set("Mcp-Method", "tools/call") + req.Header.Set("Mcp-Name", "echo") + req.Header.Set("X-Test-Principal", principal) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// TestRegression_BridgeCellA_PerPrincipalSessionIsolation pins A1: two +// different principals issuing Modern tools/calls, bridged to a Legacy +// backend via dispatchModern -> core.CallTool -> legacyCallTool, must each +// get their OWN backend initialize/session on EVERY call, and each +// synthesized initialize/tools/call must carry that principal's own +// credential -- nothing shared, cached, or reused across principals OR +// across repeat calls from the same principal. legacyCallTool creates a +// fresh MCP client (and runs a fresh initialize) per request and +// `defer c.Close()`s it, so this holds today by construction; RFC THV-0083 +// D6's planned shared backend connection pool is the future change that +// could silently break it. +// +// TWO SEQUENTIAL rounds of the same two principals, not one call each: with +// only one call per principal, per-request client creation is never actually +// exercised -- a pool keyed by (backend, principal) looks identical to +// per-request creation on a single call, since there's nothing yet for it to +// reuse. Round 2 gives such a pool its one chance to show up. The nonce is +// decoupled from the principal name ("a1"/"a2" vs "alice") specifically so +// the same principal calling twice is independently verifiable. +func TestRegression_BridgeCellA_PerPrincipalSessionIsolation(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + proxyURL, getRecords := newRecordingBackendProxy(t, backendURL) + ts := newBridgeCellAServer(t, proxyURL) + + type call struct{ principal, nonce string } + round1 := []call{{"alice", "a1"}, {"bob", "b1"}} // pool cold + round2 := []call{{"alice", "a2"}, {"bob", "b2"}} // pool warm, if there were one + calls := append(append([]call{}, round1...), round2...) + wantCred := map[string]string{ + "a1": "Bearer test-alice", "b1": "Bearer test-bob", + "a2": "Bearer test-alice", "b2": "Bearer test-bob", + } + + doRound := func(round []call) { + t.Helper() + var wg sync.WaitGroup + wg.Add(len(round)) + errs := make([]error, len(round)) + statusCodes := make([]int, len(round)) + for i, c := range round { + i, c := i, c + go func() { + defer wg.Done() + // No require/assert here: FailNow from a worker goroutine only runs + // Goexit off the test goroutine and misreports (see testing.md). + resp, err := func() (*http.Response, error) { + resp := postModernToolCallAsPrincipal(t, ts.URL, c.principal, c.nonce) + defer resp.Body.Close() + _, readErr := io.ReadAll(resp.Body) + return resp, readErr + }() + errs[i] = err + if resp != nil { + statusCodes[i] = resp.StatusCode + } + }() + } + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent Modern tool calls") + } + for i, err := range errs { + require.NoError(t, err, "principal %q", round[i].principal) + require.Equal(t, http.StatusOK, statusCodes[i], "principal %q", round[i].principal) + } + } + + doRound(round1) + doRound(round2) + + records := getRecords() + var initRecords, callRecords []recordedBackendRequest + for _, r := range records { + switch r.jsonRPCMethod { + case "initialize": + initRecords = append(initRecords, r) + case "tools/call": + callRecords = append(callRecords, r) + } + } + // mcpcompat's client library deterministically performs TWO initialize + // round-trips per legacyCallTool call (a capability pre-flight with empty + // capabilities, then the real one with elicitation/sampling added) -- so + // each call produces one SURVIVING session (the one its tools/call + // reuses) plus one abandoned pre-flight session. Assert "at least one", + // not "exactly one": the pre-flight count is a library implementation + // detail this pin does not depend on. + require.GreaterOrEqual(t, len(initRecords), len(calls), "expected at least one backend initialize per call") + require.Len(t, callRecords, len(calls), "expected one backend tools/call per call") + + // (a) + (c): every backend session id is distinct -- none reused, across + // principals OR across repeat calls from the same principal (covers + // surviving AND abandoned pre-flight sessions alike). A pool keyed by + // (backend, principal) would reuse alice's round-1 session in round 2 and + // get caught right here. + seen := make(map[string]bool, len(initRecords)) + for _, r := range initRecords { + require.NotEmpty(t, r.respSessionID, "backend initialize response must carry Mcp-Session-Id") + require.False(t, seen[r.respSessionID], "backend session id %q was reused", r.respSessionID) + seen[r.respSessionID] = true + } + + // callByNonce keys each SURVIVING tools/call by its own nonce (not by + // session id, which a pool could legitimately vary): every call must + // appear exactly once on the wire. + callByNonce := make(map[string]recordedBackendRequest, len(callRecords)) + for _, r := range callRecords { + params, _ := r.body["params"].(map[string]any) + args, _ := params["arguments"].(map[string]any) + nonce, _ := args["input"].(string) + callByNonce[nonce] = r + } + require.Len(t, callByNonce, len(calls), "every call's nonce must appear exactly once on the wire") + + // (b) + the primary data-plane check: assert the credential on the + // tools/call itself (not only the initialize step), keyed by the call's + // own nonce so the same principal's repeat call in round 2 is checked + // independently of round 1. This is what actually distinguishes + // per-request client creation from a (backend, principal)-keyed pool -- + // with only one call per principal there would be nothing to compare a + // repeat call against. + for _, c := range calls { + r := callByNonce[c.nonce] + assert.Equal(t, wantCred[c.nonce], r.authorization, + "the tools/call for nonce %q must carry principal %q's own credential", c.nonce, c.principal) + } + + // Correlate each SURVIVING initialize's Authorization credential with the + // principal named by the tools/call that reused its backend session id + // (via Mcp-Session-Id). This catches a SYMMETRIC SWAP at the initialize + // step too: if two calls' credentials were exchanged, the session-id + // cardinality check above still passes, but the credential recorded on + // session X's initialize would not match the principal named by session + // X's own tools/call. An abandoned pre-flight session (no matching + // tools/call) has nothing to correlate against and is skipped -- its + // uniqueness was already checked above. + nonceBySession := make(map[string]string, len(callByNonce)) + for nonce, r := range callByNonce { + nonceBySession[r.reqSessionID] = nonce + } + checked := 0 + for _, r := range initRecords { + nonce, ok := nonceBySession[r.respSessionID] + if !ok { + continue + } + checked++ + assert.Equal(t, wantCred[nonce], r.authorization, + "backend session %q's initialize must carry the credential of the principal behind nonce %q, not a bled/swapped one", + r.respSessionID, nonce) + } + require.Equal(t, len(calls), checked, "every call's surviving session must have been credential-checked at the initialize step too") +} + +// TestRegression_BridgeCellA_NoReservedMetaLeakToLegacyBackend pins A2: a +// Modern client's request carries reserved io.modelcontextprotocol/* _meta +// keys, but vMCP -- not the downstream caller -- is the Legacy backend's MCP +// peer on this hop, and go-sdk v1.7 hard-rejects (HTTP 400) ANY +// _meta.protocolVersion on a stateful server. legacyCallTool strips them via +// mcpparser.StripReservedMeta before forwarding +// (pkg/vmcp/client/client.go); this asserts the Legacy backend receives NO +// io.modelcontextprotocol/* key in the tools/call request body it actually +// gets. Request path only -- response _meta is issue #5986's scope. +// +// Also carries a non-reserved _meta key ("vmcp.test/nonce") and asserts it +// SURVIVES: proving absence of the reserved keys alone can't distinguish a +// surgical strip from a blanket _meta drop (or the dispatcher simply ceasing +// to forward _meta at all) -- either would also pass the reserved-key-absence +// checks below on a zero value, silently killing real per-request metadata +// (progressToken, trace context) on the Legacy hop. +func TestRegression_BridgeCellA_NoReservedMetaLeakToLegacyBackend(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + proxyURL, getRecords := newRecordingBackendProxy(t, backendURL) + ts := newBridgeCellAServer(t, proxyURL) + + resp, decoded := postModern(t, ts.URL, "tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "hello modern"}, + "_meta": map[string]any{"vmcp.test/nonce": "n1"}, + }, 1, "echo") + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + + records := getRecords() + var sawToolCall bool + for _, r := range records { + if r.jsonRPCMethod != "tools/call" { + continue + } + sawToolCall = true + params, _ := r.body["params"].(map[string]any) + meta, _ := params["_meta"].(map[string]any) + // Assert on the namespace, not a fixed list: this catches any reserved + // key, including ones the spec adds later. No passthroughMetaKeys entry + // is in play on a tools/call, so a blanket prefix check is exact. + for k := range meta { + assert.False(t, strings.HasPrefix(k, mcpparser.ReservedMetaPrefix), + "Legacy backend must never see reserved _meta key %q on the request path; got %+v", k, meta) + } + assert.Equal(t, "n1", meta["vmcp.test/nonce"], + "the strip must be surgical, not a blanket _meta drop; got %+v", meta) + } + require.True(t, sawToolCall, "expected the recording proxy to observe a tools/call request") +} + +// --------------------------------------------------------------------------- +// Cell B: Legacy client -> Modern backend +// --------------------------------------------------------------------------- + +// bridgeModernBackend is a minimal REAL Modern (2026-07-28) MCP backend used +// to pin Cell B: it speaks the raw Modern wire protocol directly +// (server/discover, tools/list, tools/call), giving real MCP semantics plus +// full request observability -- mirroring +// pkg/vmcp/client/reclassify_test.go's modernDiscoverServer technique (that +// helper, and the writeModernResult/modernReq helpers it uses, are +// unexported in package client and unusable from server_test). +type bridgeModernBackend struct { + srv *httptest.Server + + mu sync.Mutex + requests []string // Mcp-Method header value of every request received, in order +} + +func newBridgeModernBackend(t *testing.T) *bridgeModernBackend { + t.Helper() + b := &bridgeModernBackend{} + b.srv = httptest.NewServer(http.HandlerFunc(b.handle)) + t.Cleanup(b.srv.Close) + return b +} + +func (b *bridgeModernBackend) handle(w http.ResponseWriter, r *http.Request) { + method := r.Header.Get("Mcp-Method") + + b.mu.Lock() + b.requests = append(b.requests, method) + b.mu.Unlock() + + if method == "" { + // No Mcp-Method header: a Legacy `initialize` handshake attempt (or any + // other Legacy request) against this Modern-only backend. Reject with + // 404, same as reclassify_test.go's modernDiscoverServer, rather than + // silently answering it. + w.WriteHeader(http.StatusNotFound) + return + } + + var body struct { + ID any `json:"id"` + Params json.RawMessage `json:"params"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + + switch method { + case "server/discover": + b.writeResult(w, body.ID, map[string]any{ + "capabilities": map[string]any{"tools": map[string]any{}}, + "supportedVersions": []string{mcpparser.MCPVersionModern}, + }) + case "tools/list": + b.writeResult(w, body.ID, map[string]any{ + "tools": []any{map[string]any{"name": "ping", "inputSchema": map[string]any{"type": "object"}}}, + }) + case "tools/call": + b.writeResult(w, body.ID, map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "pong"}}, + }) + default: + w.WriteHeader(http.StatusNotFound) + } +} + +// writeResult writes a Modern JSON-RPC success envelope, wrapping result +// under resultType "complete" -- modernCall (pkg/vmcp/client/modern.go) +// treats any other resultType, or its absence, as NOT this backend's genuine +// response. +func (*bridgeModernBackend) writeResult(w http.ResponseWriter, id any, result map[string]any) { + if result["resultType"] == nil { + result["resultType"] = "complete" + } + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(out) +} + +func (b *bridgeModernBackend) requestMethods() []string { + b.mu.Lock() + defer b.mu.Unlock() + out := make([]string, len(b.requests)) + copy(out, b.requests) + return out +} + +func (b *bridgeModernBackend) url() string { return b.srv.URL } + +// revisionReporter mirrors the unexported interface of the same name in +// pkg/vmcp/cli (serve.go) that wires WithRevisionLookup in production; +// redeclared here because it is unexported there and this test needs the +// same type assertion against the *vmcpclient.NewHTTPBackendClient it built +// directly, without going through cli.Serve. +type revisionReporter interface { + CachedRevision(workloadID string) (mcpparser.Revision, bool) +} + +// newBridgeCellBServer builds a vMCP server routing a single Modern backend +// through backendClient/sessionFactory, shared by both Cell B tests below. +func newBridgeCellBServer( + t *testing.T, backend vmcp.Backend, backendClient vmcp.BackendClient, sessionFactory vmcpsession.MultiSessionFactory, +) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + rt := router.NewSessionRouter(&vmcp.RoutingTable{}) + + srv, err := server.New( + context.Background(), + &server.Config{ + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: sessionFactory, + Aggregator: agg, + }, + rt, + backendClient, + mockBackendRegistry, + nil, + ) + require.NoError(t, err) + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// TestRegression_BridgeCellB_SubscriptionsListenFingerprintIsValid is a +// positive control for the fingerprint TestRegression_BridgeCellB_ +// ModernBackendSetSkipsHandshakeAndStillWorks relies on below: it builds the +// SAME Modern backend and session factory but WITHOUT WithRevisionLookup, so +// the connector is never skipped and always runs go-sdk's Client.Connect(). +// Self-validates that subscriptions/listen really is observed against the +// currently-pinned go-sdk version -- if a future go-sdk bump changes that +// (e.g. gating the subscribe on registered handlers, or failing Connect +// before discover completes), this test starts failing and flags the other +// test's fingerprint as stale, instead of both silently passing forever. +func TestRegression_BridgeCellB_SubscriptionsListenFingerprintIsValid(t *testing.T) { + t.Parallel() + + modernBackend := newBridgeModernBackend(t) + backend := vmcp.Backend{ + ID: "modern-backend", + Name: "modern-backend", + BaseURL: modernBackend.url(), + TransportType: "streamable-http", + } + + authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, authReg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, strategies.NewUnauthenticatedStrategy())) + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + + // No WithRevisionLookup: the connector always attempts to connect, + // skip or no skip. + sessionFactory := vmcpsession.NewSessionFactory(authReg) + ts := newBridgeCellBServer(t, backend, backendClient, sessionFactory) + + client := NewMCPTestClient(t, ts.URL) + client.InitializeSession() + + require.Eventually(t, func() bool { + return slices.Contains(modernBackend.requestMethods(), "subscriptions/listen") + }, 5*time.Second, 50*time.Millisecond, + "go-sdk's Connect() must still open subscriptions/listen against a Modern backend when the connector is not skipped") +} + +// TestRegression_BridgeCellB_ModernBackendSetSkipsHandshakeAndStillWorks pins +// both Cell B properties against ONE session, since B2 (capabilities/calls +// working) is the direct observable consequence of B1 (the connector being +// skipped) actually mattering -- asserting them separately would just +// duplicate this setup: +// +// - B1: initOneBackend (pkg/vmcp/session/factory.go) must not attempt to +// connect against a backend whose cached MCP revision is known Modern. +// Asserted on the fake backend's own recorded request log two ways: (1) +// no request ever carries an empty Mcp-Method header (safe even though +// it never fires -- probeRevision is Modern-first and never sends a raw +// Legacy initialize, see reclassify_test.go); (2) no subscriptions/listen +// request is ever observed -- go-sdk's Client.Connect() issues that +// UNCONDITIONALLY once server/discover negotiates Modern, and nothing +// else in this path (this test's own priming probe, or the core's +// per-call Modern client) ever uses a go-sdk Client, so its presence is +// the fingerprint that the connector ran a full Connect() at all -- +// validated against the currently-pinned SDK by the sibling +// TestRegression_BridgeCellB_SubscriptionsListenFingerprintIsValid. +// - B2: despite the connector being skipped (so this session holds no +// connection for the backend at all), a Legacy client session must still +// see the backend's tools via tools/list and successfully invoke one via +// tools/call -- because capability advertisement and calls flow through +// core.* (the revision-aware per-call client), never through the +// session's per-backend connections (serve_handlers.go). +func TestRegression_BridgeCellB_ModernBackendSetSkipsHandshakeAndStillWorks(t *testing.T) { + t.Parallel() + + modernBackend := newBridgeModernBackend(t) + backend := vmcp.Backend{ + ID: "modern-backend", + Name: "modern-backend", + BaseURL: modernBackend.url(), + TransportType: "streamable-http", + } + + authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, authReg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, strategies.NewUnauthenticatedStrategy())) + + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + + // Prime the revision cache to "known Modern", the way the health monitor's + // first probe (or any earlier call) would in production -- initOneBackend's + // Modern-skip only fires for an ALREADY-classified backend (see its #5992 + // cold-start-race doc comment). ListCapabilities dispatches through + // probeRevision on an unprobed backend, which is exactly that probe. + target := vmcp.BackendToTarget(&backend) + _, err = backendClient.ListCapabilities(context.Background(), target) + require.NoError(t, err, "priming probe against the real Modern backend must succeed") + + revisions, ok := backendClient.(revisionReporter) + require.True(t, ok, "httpBackendClient must expose CachedRevision") + rev, known := revisions.CachedRevision(target.WorkloadID) + require.True(t, known && rev == mcpparser.RevisionModern, "backend must be cached Modern before session creation") + + sessionFactory := vmcpsession.NewSessionFactory(authReg, vmcpsession.WithRevisionLookup(revisions.CachedRevision)) + ts := newBridgeCellBServer(t, backend, backendClient, sessionFactory) + + // Legacy client session (real initialize handshake against vMCP). + client := NewMCPTestClient(t, ts.URL) + sessionID := client.InitializeSession() + + // B2a: tools/list must expose the Modern backend's tool once session + // registration (injectCoreSessionCapabilities -> core.ListTools) settles. + require.Eventually(t, func() bool { + names := listToolNames(t, ts.URL, sessionID) + return len(names) == 1 && names[0] == "ping" + }, 5*time.Second, 50*time.Millisecond, "tools/list should expose the Modern backend's 'ping' tool") + + // B2b: tools/call must reach the Modern backend and succeed. + toolResp := client.CallTool("ping", map[string]any{}) + defer toolResp.Body.Close() + toolBody, err := io.ReadAll(toolResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, toolResp.StatusCode, "body: %s", string(toolBody)) + + var rpc struct { + Result struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + IsError bool `json:"isError"` + } `json:"result"` + } + require.NoError(t, json.Unmarshal(toolBody, &rpc), "body: %s", string(toolBody)) + assert.False(t, rpc.Result.IsError) + require.Len(t, rpc.Result.Content, 1) + assert.Equal(t, "pong", rpc.Result.Content[0].Text) + + // B1: the session-level connector (backend.NewHTTPConnector) must never + // have run its own go-sdk Client.Connect() against this backend. Checked + // two ways: + // 1. No request ever carries an empty Mcp-Method header (the wire + // signature of a raw Legacy request) -- safe even though it should + // never fire either way: go-sdk v1.7's client is itself Modern-first + // (SEP-2575), so even a "Legacy handshake attempt" against a Modern + // backend negotiates via server/discover, not a raw Legacy initialize + // (see pkg/vmcp/client/reclassify_test.go's identical observation for + // the call-path client). + // 2. No subscriptions/listen request is ever observed. go-sdk's + // Client.Connect() issues that call UNCONDITIONALLY once + // server/discover negotiates Modern, and nothing else in this path + // (this test's own priming probe, or the core's per-call Modern + // client -- both raw modernCall, no go-sdk Client) ever makes it, so + // its presence is the fingerprint that the connector ran a full + // Connect() at all, skip or no skip. Validated against the + // currently-pinned SDK by the sibling + // TestRegression_BridgeCellB_SubscriptionsListenFingerprintIsValid. + for _, m := range modernBackend.requestMethods() { + assert.NotEmpty(t, m, "backend must never receive a request with no Mcp-Method header") + assert.NotEqual(t, "subscriptions/listen", m, + "the session connector must never run go-sdk Client.Connect() against a backend classified Modern") + } +} diff --git a/pkg/vmcp/server/call_gate.go b/pkg/vmcp/server/call_gate.go new file mode 100644 index 0000000000..0c86045b65 --- /dev/null +++ b/pkg/vmcp/server/call_gate.go @@ -0,0 +1,115 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "log/slog" + "net/http" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// authzCallGate builds the pre-dispatch CallGate consulted once per POST on the +// Streamable HTTP transport, BEFORE session validation and BEFORE the message +// reaches the SDK. It re-runs the core admission decision for the gated methods so +// a Cedar-denied direct tools/call, resources/read, or prompts/get is rejected on +// the wire as HTTP 403 + JSON-RPC error code 403 — instead of the SDK's 200/"not +// found" (a filtered tool) or the tool-result 200/IsError fallback. +// +// It reads only context the host middleware already populated (the parsed request +// and the authenticated identity), matching the seam's contract; it never mutates +// context and never logs the identity. The gate is installed only when authz is +// configured (see Handler), so this closure runs solely on authorized deployments. +// +// The gate and the subsequent dispatch (core.CallTool/ReadResource/GetPrompt) both +// authorize, and they invoke the SAME core admission helper (Check* and Call*/Read*/ +// Get* share one extracted decision — see pkg/vmcp/core), so the DECISION LOGIC +// cannot diverge. The gate is the wire-level representation of that decision (it +// turns a denial into 403 before dispatch); the call-path check remains as +// defense-in-depth for embedders that bypass this transport. +// +// For argument-conditional policies (tools/call), the gated decision and the enforced +// decision derive from the SAME parse by construction: the gate authorizes on +// parsed.Arguments (pkg/mcp's transport parse) and dispatch prefers that same map via +// gateParsedArgs before re-authorizing and forwarding (see serve_handlers.go). Where +// no matching parse exists (batch, embedders bypassing this transport, method/tool +// mismatch), dispatch falls back to the SDK decode and makes a single decision on that +// single map — so no path can produce an allow-then-deny split between gate and call. +func (s *Server) authzCallGate() server.CallGate { + return func(ctx context.Context, _ *http.Request) *server.Denial { + // An unparsable body or a batch leaves no ParsedMCPRequest: admit and let the + // SDK handle it. The batch blind spot matches the pre-existing single-server + // parity gap (cross-ref #5745); it is not widened here. + parsed := mcpparser.GetParsedMCPRequest(ctx) + if parsed == nil { + return nil + } + + // Sanctioned transport-boundary identity read (same as the Serve handlers): + // the core takes identity as an explicit parameter, so the gate resolves it + // once here from the context the auth middleware populated. A nil identity is + // anonymous — the core admission seam decides what that may do. + identity, _ := auth.IdentityFromContext(ctx) + + var err error + var message string + switch parsed.Method { + case "tools/call": + err = s.core.CheckToolCall(ctx, identity, parsed.ResourceID, parsed.Arguments) + message = vmcp.DenyMessageToolCall + case "resources/read": + err = s.core.CheckResourceRead(ctx, identity, parsed.ResourceID) + message = vmcp.DenyMessageResourceRead + case "prompts/get": + err = s.core.CheckPromptGet(ctx, identity, parsed.ResourceID) + message = vmcp.DenyMessagePromptGet + default: + // Non-gated method (initialize, tools/list, ping, ...): admit unchanged. + // + // This default is deliberately fail-OPEN, unlike pkg/authz's + // MCPMethodToFeatureOperation ("methods not in this map are denied by + // default"). The gate mirrors the core admission seam. For methods with + // NO admission decision at all (elicitation/create, sampling/createMessage, + // tasks/*), denying here would diverge from the core rather than enforce a + // real policy. + // + // Two methods DO carry a core admission decision but are intentionally + // NOT wire-gated here — this is a conscious choice, not an oversight: + // - completion/complete: core.Complete authorizes the underlying prompt/ + // resource ref at dispatch. + // - resources/subscribe: core validates the URI via LookupResource at + // dispatch. + // Both are enforced at dispatch instead of on the wire because they are not + // argument-conditional in the way the gate's parsed-args fast path assumes, + // and their denials already surface through the core's error mapping. To + // gate a NEW verb on the wire: add a Check* method on core.VMCP for it and a + // case above — do NOT flip this default to deny, which would reject protocol + // methods (initialize/ping/list/notifications) that must always pass. + return nil + } + + if err == nil { + return nil + } + + // Only an authorization denial (or a fail-closed authorizer error, already + // classified as ErrAuthorizationFailed inside Check*) becomes a 403. Any other + // error is infrastructure (aggregation/backend plumbing): admit so the call + // path hits the same failure and surfaces it through the existing mapping — + // the gate must not convert infra faults into 403s. + if !errors.Is(err, vmcp.ErrAuthorizationFailed) { + slog.WarnContext(ctx, "vmcp authz gate: non-authorization error, admitting request", + "method", parsed.Method, "error", err) + return nil + } + + // HTTPStatus left zero ⇒ the shim writes HTTP 403 (see server.Denial). + return &server.Denial{Code: mcpparser.JSONRPCCodeDenied, Message: message} + } +} diff --git a/pkg/vmcp/server/call_gate_test.go b/pkg/vmcp/server/call_gate_test.go new file mode 100644 index 0000000000..4b8fb7c3cb --- /dev/null +++ b/pkg/vmcp/server/call_gate_test.go @@ -0,0 +1,176 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/core" +) + +// gateFakeCore is a core.VMCP whose only exercised methods are the Check* triple. +// The embedded nil interface satisfies the rest and panics if the gate ever calls +// a method it should not. +type gateFakeCore struct { + core.VMCP + toolErr error + resErr error + promptErr error + + gotMethod string + gotResourceID string + gotArgs map[string]any +} + +func (f *gateFakeCore) CheckToolCall(_ context.Context, _ *auth.Identity, name string, args map[string]any) error { + f.gotMethod = "tools/call" + f.gotResourceID = name + f.gotArgs = args + return f.toolErr +} + +func (f *gateFakeCore) CheckResourceRead(_ context.Context, _ *auth.Identity, uri string) error { + f.gotMethod = "resources/read" + f.gotResourceID = uri + return f.resErr +} + +func (f *gateFakeCore) CheckPromptGet(_ context.Context, _ *auth.Identity, name string) error { + f.gotMethod = "prompts/get" + f.gotResourceID = name + return f.promptErr +} + +func TestAuthzCallGate(t *testing.T) { + t.Parallel() + + deny := fmt.Errorf("%w: policy said no", vmcp.ErrAuthorizationFailed) + infra := errors.New("aggregation exploded") + + tests := []struct { + name string + // parsed is the request the parsing middleware would have populated; nil means + // no ParsedMCPRequest in context (unparsable body / batch). + parsed *mcpparser.ParsedMCPRequest + core *gateFakeCore + wantDenial bool + wantMessage string + wantMethod string // Check* method the gate must have invoked ("" = none) + wantResource string + wantArgs map[string]any + forbidInMsg string // substring that must NOT appear in the denial message + }{ + { + name: "no parsed request admits", + core: &gateFakeCore{}, + }, + { + name: "non-gated method admits", + parsed: &mcpparser.ParsedMCPRequest{Method: "tools/list"}, + core: &gateFakeCore{}, + }, + { + name: "tools/call allowed admits and forwards name+args", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", ResourceID: "echo", Arguments: map[string]any{"input": "hi"}, + }, + core: &gateFakeCore{}, + wantMethod: "tools/call", + wantResource: "echo", + wantArgs: map[string]any{"input": "hi"}, + }, + { + name: "tools/call denied returns 403 denial", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", ResourceID: "secret-tool", Arguments: map[string]any{"input": "hi"}, + }, + core: &gateFakeCore{toolErr: deny}, + wantDenial: true, + wantMessage: vmcp.DenyMessageToolCall, + wantMethod: "tools/call", + wantResource: "secret-tool", + forbidInMsg: "secret-tool", + }, + { + name: "tools/call non-authz error admits (no 403 for infra faults)", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", ResourceID: "echo", + }, + core: &gateFakeCore{toolErr: infra}, + wantMethod: "tools/call", + }, + { + name: "resources/read denied returns 403 denial", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "resources/read", ResourceID: "file://secret", + }, + core: &gateFakeCore{resErr: deny}, + wantDenial: true, + wantMessage: vmcp.DenyMessageResourceRead, + wantMethod: "resources/read", + wantResource: "file://secret", + forbidInMsg: "file://secret", + }, + { + name: "prompts/get denied returns 403 denial", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "prompts/get", ResourceID: "secret-prompt", + }, + core: &gateFakeCore{promptErr: deny}, + wantDenial: true, + wantMessage: vmcp.DenyMessagePromptGet, + wantMethod: "prompts/get", + wantResource: "secret-prompt", + forbidInMsg: "secret-prompt", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Build the context per-subtest from t.Context() (our test convention), + // injecting the parsed request as the parsing middleware would. + ctx := t.Context() + if tt.parsed != nil { + ctx = context.WithValue(ctx, mcpparser.MCPRequestContextKey, tt.parsed) + } + + s := &Server{core: tt.core} + gate := s.authzCallGate() + + denial := gate(ctx, nil) + + if tt.wantDenial { + require.NotNil(t, denial, "expected a denial") + assert.Equal(t, mcpparser.JSONRPCCodeDenied, denial.Code, "denial code must be 403") + assert.Equal(t, 0, denial.HTTPStatus, "HTTPStatus must be zero so the shim writes 403") + assert.Equal(t, tt.wantMessage, denial.Message) + if tt.forbidInMsg != "" { + assert.NotContains(t, denial.Message, tt.forbidInMsg, + "the denial message must not leak the capability name (no enumeration oracle)") + } + } else { + assert.Nil(t, denial, "expected the request to be admitted") + } + + assert.Equal(t, tt.wantMethod, tt.core.gotMethod, "unexpected Check* method invoked") + if tt.wantResource != "" { + assert.Equal(t, tt.wantResource, tt.core.gotResourceID) + } + if tt.wantArgs != nil { + assert.Equal(t, tt.wantArgs, tt.core.gotArgs) + } + }) + } +} diff --git a/pkg/vmcp/server/capability_regression_test.go b/pkg/vmcp/server/capability_regression_test.go new file mode 100644 index 0000000000..7288c73408 --- /dev/null +++ b/pkg/vmcp/server/capability_regression_test.go @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_InitializeAdvertisesToolsAndResourcesCapabilities pins the +// capabilities advertised in the initialize response on the Serve path. +// +// BEHAVIOR (go-sdk bridge, toolhive-core v0.0.28): the initialize response +// advertises tools and resources capabilities alongside logging, matching +// mcp-go. Earlier releases advertised only {"logging":{}} on the Serve path; +// v0.0.28 surfaces the tool/resource capabilities in the initialize response. +// +// This test pins that behavior so a future regression is a deliberate, visible +// flip rather than a silent drift. It asserts on the RAW initialize response +// body parsed with encoding/json — not tools/list. +func TestRegression_InitializeAdvertisesToolsAndResourcesCapabilities(t *testing.T) { + t.Parallel() + + fc := &fakeCore{ + tools: []vmcp.Tool{{Name: "cap-tool", Description: "a capability test tool"}}, + resources: []vmcp.Resource{{Name: "cap-doc", URI: "file:///cap.txt"}}, + } + _, _, baseURL := registerServeSession(t, fc) + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, 200, initResp.StatusCode, "initialize should succeed") + + env, _ := readServeJSONRPC(t, initResp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "initialize response must have a result object; env: %v", env) + + capabilities, ok := result["capabilities"].(map[string]any) + require.True(t, ok, "result.capabilities must be present; result: %v", result) + + // As of toolhive-core v0.0.28 the go-sdk bridge advertises tools and + // resources capabilities in the initialize response on the Serve path + // (alongside logging). This pins that behavior so a future regression is a + // deliberate, visible flip. + assert.Contains(t, capabilities, "logging", + "logging capability must be advertised; got %v", capabilities) + assert.Contains(t, capabilities, "tools", + "tools capability must be advertised in initialize on the Serve path; got %v", capabilities) + assert.Contains(t, capabilities, "resources", + "resources capability must be advertised in initialize on the Serve path; got %v", capabilities) + assert.Contains(t, capabilities, "prompts", + "prompts capability must be advertised in initialize on the Serve path (#5969 fixes a latent bug"+ + " where it was never advertised even though per-session prompts were served); got %v", capabilities) +} + +// TestRegression_InitializeAdvertisesListChanged pins the #5748 (tools) and +// #5969 (resources/prompts) capability flips: tools.listChanged, +// resources.listChanged, and prompts.listChanged must all be true in the +// initialize response (they were false — and therefore omitted, per the +// SDK's omitempty — before these changes). resources.subscribe must remain +// true throughout (unrelated to listChanged, pinned pre-existing behavior). +func TestRegression_InitializeAdvertisesListChanged(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "cap-tool"}}} + _, _, baseURL := registerServeSession(t, fc) + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, 200, initResp.StatusCode, "initialize should succeed") + + env, _ := readServeJSONRPC(t, initResp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "initialize response must have a result object; env: %v", env) + capabilities, ok := result["capabilities"].(map[string]any) + require.True(t, ok, "result.capabilities must be present; result: %v", result) + + toolsCaps, ok := capabilities["tools"].(map[string]any) + require.True(t, ok, "tools capability must be an object; capabilities: %v", capabilities) + assert.Equal(t, true, toolsCaps["listChanged"], + "tools.listChanged must be true so downstream clients trust the resync notification (#5748)") + + resourcesCaps, ok := capabilities["resources"].(map[string]any) + require.True(t, ok, "resources capability must be an object; capabilities: %v", capabilities) + assert.Equal(t, true, resourcesCaps["subscribe"], "resources.subscribe must remain true") + assert.Equal(t, true, resourcesCaps["listChanged"], + "resources.listChanged must be true so downstream clients trust the resync notification (#5969)") + + promptsCaps, ok := capabilities["prompts"].(map[string]any) + require.True(t, ok, "prompts capability must be an object; capabilities: %v", capabilities) + assert.Equal(t, true, promptsCaps["listChanged"], + "prompts.listChanged must be true so downstream clients trust the resync notification (#5969)") +} + +// TestRegression_SessionRegistrationDoesNotSpuriouslyInvalidateCache verifies +// the #5748 R1 concern at the layer this PR controls: registering a session — +// with no backend list_changed notification ever firing — must NOT invoke +// InvalidateCapabilityCache or the resync path. Those only run from the sink +// built in buildListChangedSink, which is invoked solely by a persistent +// backend connection's OnNotification handler (mcp_session.go), never by +// session registration itself. +// +// (go-sdk's OWN early/broadcast-scoped auto-notify behavior at registration — +// documented in serve.go next to WithToolCapabilities(true) — is an upstream +// SDK characteristic this PR cannot suppress; it is unrelated to whether OUR +// code spuriously self-triggers, which is what this test pins.) +func TestRegression_SessionRegistrationDoesNotSpuriouslyInvalidateCache(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "cap-tool"}}} + registerServeSession(t, fc) + + assert.Equal(t, int32(0), fc.invalidateCacheCalls.Load(), + "session registration alone must never invalidate the capability cache") +} diff --git a/pkg/vmcp/server/classification.go b/pkg/vmcp/server/classification.go new file mode 100644 index 0000000000..d4a7d15a94 --- /dev/null +++ b/pkg/vmcp/server/classification.go @@ -0,0 +1,157 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "net/http" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" +) + +// methodServerDiscover is the Modern (2026-07-28) capability-probe method — how +// a client learns which revisions a server supports. It is special-cased in two +// places — the capability-gate branch in classifyingHandler and dispatchModern's +// method switch (dispatchModernDiscover) — because it must never be refused on +// version grounds: it is the one request a client needs answered to negotiate +// down. +const methodServerDiscover = "server/discover" + +// classifyingHandler classifies a parsed MCP request as Legacy (2025-11-25) or +// Modern (2026-07-28) at the decode seam, rejects a malformed Modern request +// with the correct JSON-RPC error before it reaches dispatch, and routes a +// well-formed Modern request to dispatchModern instead of the SDK. Legacy +// traffic always falls through to next unchanged. +// +// Classification is exact-match on both the MCP-Protocol-Version header and the +// reserved io.modelcontextprotocol/* _meta keys (ClassifyRevision), so nothing +// that is not unambiguously Modern reaches dispatchModern: Legacy wire behavior +// is unaffected by this handler. +// +// Whether vMCP serves the Modern revision at all is a per-instance capability +// question, not a global switch: #5959's temporary env-var kill-switch is +// replaced by modernDispatchBlockers (modern_gate.go), which enumerates the +// enabled features the stateless dispatch path cannot serve. When that list is +// empty (the common case) a well-formed Modern request dispatches; when it is +// not, the gate branch below refuses Modern with a conformant -32022 listing +// Legacy so capable clients negotiate down and reach the SDK path, where every +// enabled feature works. A request whose _meta names some OTHER protocol +// version is refused with the same -32022 by ClassifyRevision below, which +// owns every other version-grounds rejection. +// +// ValidateHeaderConsistency (Mcp-Method/Mcp-Name) only applies to Modern +// requests: a Legacy request carrying a stray Mcp-Method/Mcp-Name header +// (e.g. from a misbehaving proxy) must not be rejected for it, since Legacy +// clients never send these headers and have no obligation to omit them. +// +// This handler makes no authentication/authorization decision of its own and +// confers no elevated trust on requests that pass it — it only validates +// protocol shape and routes. It must run after ParsingMiddleware (so +// GetParsedMCPRequest is populated) and is expected to run after any auth +// middleware in the chain, so a Modern dispatch that gets gated 403 is still +// audited as "denied". +func (s *Server) classifyingHandler(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + parsed := mcpparser.GetParsedMCPRequest(r.Context()) + if parsed == nil { + next.ServeHTTP(w, r) + return + } + + protoHeader := r.Header.Get("MCP-Protocol-Version") + rev, err := mcpparser.ClassifyRevision(parsed.Method, parsed.Meta, protoHeader) + if err != nil { + mcpparser.WriteClassificationError(w, parsed.ID, err) + return + } + + if rev != mcpparser.RevisionModern { + next.ServeHTTP(w, r) + return + } + + // A Modern (2026-07-28) request over HTTP MUST carry the + // MCP-Protocol-Version header (draft Streamable HTTP "Server Validation"). + // ClassifyRevision is transport-agnostic -- it also serves header-less + // stdio -- so it admits a Modern request signaled only by the reserved + // io.modelcontextprotocol/* _meta keys and defers the header-presence rule + // to the HTTP layer (see the TODO in pkg/mcp/revision.go). Without this + // check a well-formed Modern _meta with no header would dispatch and return + // 200 instead of the mandated -32020 rejection. + if protoHeader == "" { + mcpparser.WriteClassificationError(w, parsed.ID, errMissingProtocolVersionHeader) + return + } + + if err := mcpparser.ValidateHeaderConsistency(parsed); err != nil { + mcpparser.WriteClassificationError(w, parsed.ID, err) + return + } + + // Capability gate: when this instance has enabled features the stateless + // dispatch path cannot serve (modernDispatchBlockers, modern_gate.go), + // vMCP does not serve the Modern revision. + // + // Answer that conformantly here rather than letting the request reach the + // SDK. The draft's Streamable HTTP "Protocol Version Header" section + // requires a server that does not implement a requested version -- "whether + // the version is unknown to the server, or is a known version the server has + // chosen not to support" -- to reply 400 with an UnsupportedProtocolVersionError + // listing the versions it does support. A gated instance is exactly the + // second case. Falling through instead yields go-sdk's stateful-server + // rejection, which is a 400 with a PLAIN-TEXT body whose text is Go-API advice + // for the server author ("set StreamableHTTPOptions.Stateless = true") -- not + // parseable as a protocol error and carrying no version list. + // + // server/discover is deliberately exempt: it is how a client learns which + // revisions a server supports, so rejecting it on version grounds would leave + // the client no way to negotiate down. go-sdk exempts it for the same reason, + // and its stateful path answers discover with the transport-filtered version + // list -- which excludes 2026-07-28 for a stateful transport -- so the + // fall-through is not just tolerable but the correct advertisement: a + // Modern-first client (go-sdk v1.7+ probes discover before initialize) + // reads it and lands on the Legacy handshake without ever seeing an error. + if blocked := s.modernDispatchBlockers(); len(blocked) > 0 { + if parsed.Method == methodServerDiscover { + next.ServeHTTP(w, r) + return + } + mcpparser.WriteClassificationError(w, parsed.ID, &mcpparser.UnsupportedVersionError{ + Requested: mcpparser.MCPVersionModern, + Supported: []string{mcpparser.MCPVersionLegacy}, + }) + return + } + + s.dispatchModern(w, r, parsed) + }) +} + +// missingProtocolVersionHeaderError is returned when a request classifies Modern +// (2026-07-28) over HTTP but omits the mandatory MCP-Protocol-Version header. The +// draft Streamable HTTP "Server Validation" rules make a missing required standard +// header a -32020 (HeaderMismatch) condition, so this maps to the same wire code +// and HTTP 400 as the header/body mismatch mcp.ClassifyRevision already produces +// when the header is present but wrong. +// +// The enforcement lives here, at the HTTP layer, rather than in the +// transport-agnostic mcp.ClassifyRevision (which also serves header-less stdio and +// cannot know the header was required); see the TODO in pkg/mcp/revision.go. It +// implements mcp.CodedError so mcp.WriteClassificationError renders it correctly. +type missingProtocolVersionHeaderError struct{} + +func (missingProtocolVersionHeaderError) Error() string { + return "MCP-Protocol-Version header is required for Modern (2026-07-28) requests" +} + +// Code implements mcp.CodedError. +func (missingProtocolVersionHeaderError) Code() int64 { return mcpparser.CodeHeaderMismatch } + +// Data implements mcp.CodedError. +func (missingProtocolVersionHeaderError) Data() map[string]any { + return map[string]any{"header": "MCP-Protocol-Version"} +} + +// errMissingProtocolVersionHeader is the singleton rejection for a Modern HTTP +// request that omits the mandatory MCP-Protocol-Version header. +var errMissingProtocolVersionHeader = missingProtocolVersionHeaderError{} diff --git a/pkg/vmcp/server/classification_test.go b/pkg/vmcp/server/classification_test.go new file mode 100644 index 0000000000..adcb82a0ec --- /dev/null +++ b/pkg/vmcp/server/classification_test.go @@ -0,0 +1,436 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" +) + +// Reserved Modern _meta keys, mirrored from pkg/mcp/revision.go's unexported +// constants since classification_test.go cannot import them directly. +const ( + metaKeyProtocolVersion = "io.modelcontextprotocol/protocolVersion" + metaKeyClientInfo = "io.modelcontextprotocol/clientInfo" + metaKeyClientCapabilities = "io.modelcontextprotocol/clientCapabilities" +) + +func sentinelEncode(v string) string { + return "=?base64?" + base64.StdEncoding.EncodeToString([]byte(v)) + "?=" +} + +type classificationErrorBody struct { + Error struct { + Code int64 `json:"code"` + Message string `json:"message"` + } `json:"error"` +} + +// classifyingHandlerTestServer builds a minimal *Server for driving +// classifyingHandler in isolation, carrying only what the handler reads beyond +// config scalars: the core a Modern dispatch routes to. +func classifyingHandlerTestServer() *Server { + return &Server{ + config: &Config{ + Name: testServerName, + Version: testServerVersion, + }, + core: &modernFakeCore{tools: []vmcp.Tool{{Name: "echo", InputSchema: map[string]any{"type": "object"}}}}, + } +} + +func TestClassifyingHandler(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + parsed *mcpparser.ParsedMCPRequest + protocolHeader string + wantPassthrough bool + wantDispatched bool + wantCode int64 + }{ + { + name: "nil parsed request passes through", + parsed: nil, + wantPassthrough: true, + }, + { + name: "legacy body with no modern signal passes through", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + }, + wantPassthrough: true, + }, + { + // A non-Modern MCP-Protocol-Version header, with no reserved _meta key, + // is not a Modern signal (ClassifyRevision requires an exact match on + // MCPVersionModern): this must still reach next, never dispatchModern, + // now that Modern dispatch is unconditional for well-formed requests. + name: "legacy request with an old protocol version header still passes through", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + }, + protocolHeader: "2025-11-25", + wantPassthrough: true, + }, + { + // tools/list is deliberately not in the Mcp-Name-required set, so this + // case only needs Mcp-Method (required on every Modern request) to pass + // ValidateHeaderConsistency. This test server has no enabled feature + // the stateless path cannot serve (modernDispatchBlockers is empty), + // so a well-formed Modern request dispatches to the core rather than + // falling through to next. The gated counterpart is + // TestClassifyingHandler_ModernCapabilityGate. + name: "well-formed modern request dispatches to the core", + parsed: wellFormedModernToolsList(), + protocolHeader: mcpparser.MCPVersionModern, + wantDispatched: true, + }, + { + // With no capability-gate blockers, server/discover is dispatched like + // any other Modern method: dispatchModern answers it directly. When + // the gate is active it instead falls through to the SDK so the client + // can negotiate down — see TestClassifyingHandler_ModernCapabilityGate. + name: "server/discover dispatches to the core", + parsed: wellFormedModernDiscover(), + protocolHeader: mcpparser.MCPVersionModern, + wantDispatched: true, + }, + { + // A body that is otherwise a well-formed Modern request (valid _meta + // protocolVersion + clientCapabilities) but omits the mandatory + // MCP-Protocol-Version header MUST be rejected with -32020, not + // dispatched: the draft Streamable HTTP spec requires the header on + // every Modern POST. Without the header-presence check in + // classifyingHandler this would classify Modern with a nil error and + // return 200. + name: "well-formed modern body missing the protocol version header is rejected", + parsed: wellFormedModernToolsList(), + protocolHeader: "", + wantCode: mcpparser.CodeHeaderMismatch, + }, + { + // initialize is forced Legacy unconditionally (ClassifyRevision), even + // with a full spoofed Modern signal on both header and _meta -- mirrors + // revision_test.go's "legacy: initialize wins over spoofed modern meta + // and header" at the classifyingHandler boundary, so a future change + // ahead of the ClassifyRevision call can't silently route this to + // dispatchModern. + name: "initialize with spoofed modern signal still passes through", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "initialize", + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + }, + protocolHeader: mcpparser.MCPVersionModern, + wantPassthrough: true, + }, + { + name: "modern signal via reserved key with no protocolVersion and no header is invalid params", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + Meta: map[string]any{ + metaKeyClientInfo: map[string]any{}, + }, + }, + wantCode: mcpparser.CodeInvalidParams, + }, + { + name: "modern header/body protocolVersion mismatch is a header mismatch", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + }, + protocolHeader: "2099-01-01", + wantCode: mcpparser.CodeHeaderMismatch, + }, + { + name: "modern unsupported protocolVersion is rejected", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + Meta: map[string]any{ + metaKeyProtocolVersion: "1.0", + }, + }, + wantCode: mcpparser.CodeUnsupportedProtocolVersion, + }, + { + name: "modern missing clientCapabilities is rejected", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + }, + }, + protocolHeader: mcpparser.MCPVersionModern, + wantCode: mcpparser.CodeMissingClientCapability, + }, + { + // No Modern signal anywhere (no header, no reserved _meta key), so this + // classifies Legacy and ValidateHeaderConsistency must not run at all — + // a stray/mismatched Mcp-Method header on Legacy traffic is not an error. + name: "legacy request carrying a stray Mcp-Method header passes through unchanged", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + MCPMethodHeader: "resources/read", + }, + wantPassthrough: true, + }, + { + name: "sentinel-encoded Mcp-Name header mismatched against ResourceID is a header mismatch", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "echo", + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + MCPMethodHeader: "tools/call", + MCPNameHeader: sentinelEncode("other-tool"), + }, + protocolHeader: mcpparser.MCPVersionModern, + wantCode: mcpparser.CodeHeaderMismatch, + }, + { + name: "modern request missing required Mcp-Method header is rejected", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/list", + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + }, + protocolHeader: mcpparser.MCPVersionModern, + wantCode: mcpparser.CodeHeaderMismatch, + }, + { + name: "modern tools/call request missing required Mcp-Name header is rejected", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + ResourceID: "echo", + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + MCPMethodHeader: "tools/call", + }, + protocolHeader: mcpparser.MCPVersionModern, + wantCode: mcpparser.CodeHeaderMismatch, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + ctx := t.Context() + if tt.parsed != nil { + ctx = context.WithValue(ctx, mcpparser.MCPRequestContextKey, tt.parsed) + } + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(ctx) + if tt.protocolHeader != "" { + req.Header.Set("MCP-Protocol-Version", tt.protocolHeader) + } + + nextCalled := false + next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + }) + + rec := httptest.NewRecorder() + classifyingHandlerTestServer().classifyingHandler(next).ServeHTTP(rec, req) + + if tt.wantPassthrough { + assert.True(t, nextCalled, "expected the request to fall through to next") + return + } + assert.False(t, nextCalled, "expected classification to short-circuit before next") + + if tt.wantDispatched { + var body map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "expected a Modern result envelope, got %v", body) + assert.Equal(t, "complete", result["resultType"]) + return + } + + var body classificationErrorBody + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + assert.Equal(t, tt.wantCode, body.Error.Code) + }) + } +} + +// wellFormedModernToolsList returns a well-formed Modern tools/list request: +// tools/list is deliberately not in the Mcp-Name-required set, so only +// Mcp-Method (required on every Modern request) is needed for it to pass +// ValidateHeaderConsistency. +func wellFormedModernToolsList() *mcpparser.ParsedMCPRequest { + return &mcpparser.ParsedMCPRequest{ + Method: "tools/list", + ID: "1", + IsRequest: true, + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + MCPMethodHeader: "tools/list", + } +} + +// wellFormedModernDiscover is wellFormedModernToolsList for server/discover, the +// Modern capability probe. +func wellFormedModernDiscover() *mcpparser.ParsedMCPRequest { + return &mcpparser.ParsedMCPRequest{ + Method: methodServerDiscover, + ID: "1", + IsRequest: true, + Meta: map[string]any{ + metaKeyProtocolVersion: mcpparser.MCPVersionModern, + metaKeyClientCapabilities: map[string]any{}, + }, + MCPMethodHeader: methodServerDiscover, + } +} + +// stubOptimizerFactory marks the optimizer as enabled for gate tests. The gate +// (modernDispatchBlockers) only checks the factory for nil-ness — no gated +// request ever builds an optimizer — so returning an error is safe and keeps +// the stub honest about never being invoked. +func stubOptimizerFactory(context.Context, []mcpserver.ServerTool) (optimizer.Optimizer, error) { + return nil, errors.New("stub optimizer factory must not be invoked by the capability gate") +} + +// TestModernDispatchBlockers pins the gate's contents: exactly which enabled +// features keep an instance off the Modern revision, by name. A parity change +// (deleting an entry in modern_gate.go) must flip a case here. +func TestModernDispatchBlockers(t *testing.T) { + t.Parallel() + + plain := classifyingHandlerTestServer() + assert.Empty(t, plain.modernDispatchBlockers(), + "an instance with no Serve-layer-only features must serve Modern") + + withOptimizer := classifyingHandlerTestServer() + withOptimizer.optimizerFactory = stubOptimizerFactory + assert.Equal(t, []string{blockerOptimizer}, withOptimizer.modernDispatchBlockers(), + "the session-scoped optimizer is not servable by the stateless Modern path") +} + +// TestClassifyingHandler_ModernCapabilityGate pins the wire behavior of the +// capability gate for an instance with a feature the stateless Modern path +// cannot serve (the optimizer): +// +// - A well-formed Modern request is refused with a conformant 400 + -32022 +// UnsupportedProtocolVersionError whose data lists the Legacy version, so a +// Modern-first client (go-sdk v1.7+ probes before initialize) negotiates +// down cleanly instead of hitting go-sdk's unparsable plain-text +// stateful rejection. +// - server/discover is exempt and falls through to the SDK path, whose +// stateful transport advertises the Legacy-only version list — the one +// request a client needs answered to negotiate down. +// - Legacy traffic falls through untouched: the gate only ever refuses +// requests that classified Modern. +// +// The ungated counterparts (same requests dispatching to the core) live in +// TestClassifyingHandler; together the two pin that the gate cannot be +// "simplified" away in either direction. +func TestClassifyingHandler_ModernCapabilityGate(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + parsed *mcpparser.ParsedMCPRequest + protocolHeader string + wantPassthrough bool + }{ + { + name: "modern tools/list is refused with -32022 listing Legacy", + parsed: wellFormedModernToolsList(), + protocolHeader: mcpparser.MCPVersionModern, + }, + { + name: "server/discover falls through to the SDK so the client can negotiate down", + parsed: wellFormedModernDiscover(), + protocolHeader: mcpparser.MCPVersionModern, + wantPassthrough: true, + }, + { + name: "legacy request falls through untouched", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "tools/list", + ID: "1", + IsRequest: true, + }, + protocolHeader: mcpparser.MCPVersionLegacy, + wantPassthrough: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := classifyingHandlerTestServer() + srv.optimizerFactory = stubOptimizerFactory + + ctx := context.WithValue(t.Context(), mcpparser.MCPRequestContextKey, tt.parsed) + req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(ctx) + req.Header.Set("MCP-Protocol-Version", tt.protocolHeader) + + nextCalled := false + next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + nextCalled = true + w.WriteHeader(http.StatusOK) + }) + + rec := httptest.NewRecorder() + srv.classifyingHandler(next).ServeHTTP(rec, req) + + if tt.wantPassthrough { + assert.True(t, nextCalled, "expected the request to fall through to next") + return + } + assert.False(t, nextCalled, "expected the gate to short-circuit before next") + assert.Equal(t, http.StatusBadRequest, rec.Code) + + var body struct { + Error struct { + Code int64 `json:"code"` + Data struct { + Requested string `json:"requested"` + Supported []string `json:"supported"` + } `json:"data"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + assert.Equal(t, mcpparser.CodeUnsupportedProtocolVersion, body.Error.Code) + assert.Equal(t, mcpparser.MCPVersionModern, body.Error.Data.Requested) + assert.Equal(t, []string{mcpparser.MCPVersionLegacy}, body.Error.Data.Supported, + "the refusal must list the Legacy version so the client can negotiate down") + }) + } +} diff --git a/pkg/vmcp/server/client_sessions.go b/pkg/vmcp/server/client_sessions.go deleted file mode 100644 index ce3844045d..0000000000 --- a/pkg/vmcp/server/client_sessions.go +++ /dev/null @@ -1,85 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "sync" - - "github.com/mark3labs/mcp-go/server" -) - -// clientSessionSet tracks the SDK sessions registered on this replica. It feeds -// the session cardinality gauges and selects the sessions a capability refresh -// clears (serve_refresh.go). -// -// A registration runs for up to defaultSessionRegistrationTimeout, and mcp-go -// can unregister the session meanwhile (a DELETE, the end of a GET, the idle -// sweep). So a registration is begun when the hook starts and committed when it -// succeeds, and remove cancels a registration still in progress: a session that -// was unregistered before its registration finished is never added. The zero -// value is ready to use. -type clientSessionSet struct { - mu sync.Mutex - pending map[string]struct{} - active map[string]server.ClientSession -} - -// begin records that the registration of sessionID has started. -func (c *clientSessionSet) begin(sessionID string) { - c.mu.Lock() - defer c.mu.Unlock() - if c.pending == nil { - c.pending = make(map[string]struct{}) - } - c.pending[sessionID] = struct{}{} -} - -// commit adds session if its registration is still in progress (begin was -// called and remove was not), and ends the registration. -func (c *clientSessionSet) commit(session server.ClientSession) { - c.mu.Lock() - defer c.mu.Unlock() - sessionID := session.SessionID() - if _, ok := c.pending[sessionID]; !ok { - return - } - delete(c.pending, sessionID) - if c.active == nil { - c.active = make(map[string]server.ClientSession) - } - c.active[sessionID] = session -} - -// abort ends a registration that failed, without adding the session. -func (c *clientSessionSet) abort(sessionID string) { - c.mu.Lock() - defer c.mu.Unlock() - delete(c.pending, sessionID) -} - -// remove drops sessionID and cancels its registration if one is in progress. -func (c *clientSessionSet) remove(sessionID string) { - c.mu.Lock() - defer c.mu.Unlock() - delete(c.pending, sessionID) - delete(c.active, sessionID) -} - -// count returns the number of registered sessions. -func (c *clientSessionSet) count() int64 { - c.mu.Lock() - defer c.mu.Unlock() - return int64(len(c.active)) -} - -// snapshot returns the registered sessions. -func (c *clientSessionSet) snapshot() []server.ClientSession { - c.mu.Lock() - defer c.mu.Unlock() - sessions := make([]server.ClientSession, 0, len(c.active)) - for _, session := range c.active { - sessions = append(sessions, session) - } - return sessions -} diff --git a/pkg/vmcp/server/context_isolation_regression_test.go b/pkg/vmcp/server/context_isolation_regression_test.go new file mode 100644 index 0000000000..46e90f0877 --- /dev/null +++ b/pkg/vmcp/server/context_isolation_regression_test.go @@ -0,0 +1,406 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/auth" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + sessionfactorymocks "github.com/stacklok/toolhive/pkg/vmcp/session/mocks" + sessionmocks "github.com/stacklok/toolhive/pkg/vmcp/session/types/mocks" +) + +// perToolCore is a fakeCore variant whose CallTool returns a result text equal +// to the called tool's name. This lets the context-isolation test detect +// cross-contamination: each concurrent call must observe its own tool name. +type perToolCore struct { + fakeCore +} + +func (p *perToolCore) CallTool( + _ context.Context, _ *auth.Identity, name string, _ map[string]any, _ map[string]any, +) (*vmcp.ToolCallResult, error) { + p.callToolCalls.Add(1) + p.lastCallToolName.Store(name) + return &vmcp.ToolCallResult{ + Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: name}}, + }, nil +} + +// newPerToolSessionFactory mirrors newToolSessionFactory but wires a per-tool +// CallTool result on the mock MultiSession so the SDK tool handler routes +// through perToolCore (whose result carries the tool name). The mock session's +// CallTool is only a fallback; the Serve path routes through coreToolHandler → +// core.CallTool, so perToolCore.CallTool is the one that runs. +func newPerToolSessionFactory( + t *testing.T, ctrl *gomock.Controller, tools []vmcp.Tool, +) *sessionfactorymocks.MockMultiSessionFactory { + t.Helper() + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { + mock := sessionmocks.NewMockMultiSession(ctrl) + mock.EXPECT().ID().Return(id).AnyTimes() + mock.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes() + mock.EXPECT().CreatedAt().Return(time.Time{}).AnyTimes() + mock.EXPECT().Type().Return(transportsession.SessionType("")).AnyTimes() + mock.EXPECT().GetData().Return(nil).AnyTimes() + mock.EXPECT().SetData(gomock.Any()).AnyTimes() + mock.EXPECT().GetMetadata().Return(map[string]string{ + vmcpsession.MetadataKeyIdentityBinding: "unauthenticated", + }).AnyTimes() + mock.EXPECT().GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding). + Return("unauthenticated", true).AnyTimes() + mock.EXPECT().SetMetadata(gomock.Any(), gomock.Any()).AnyTimes() + toolsCopy := make([]vmcp.Tool, len(tools)) + copy(toolsCopy, tools) + mock.EXPECT().Tools().Return(toolsCopy).AnyTimes() + mock.EXPECT().AllTools().Return(toolsCopy).AnyTimes() + mock.EXPECT().Resources().Return(nil).AnyTimes() + mock.EXPECT().Prompts().Return(nil).AnyTimes() + mock.EXPECT().BackendSessions().Return(nil).AnyTimes() + rt := &vmcp.RoutingTable{Tools: make(map[string]*vmcp.BackendTarget, len(tools))} + for _, tool := range tools { + rt.Tools[tool.Name] = &vmcp.BackendTarget{WorkloadID: tool.Name} + } + mock.EXPECT().GetRoutingTable().Return(rt).AnyTimes() + mock.EXPECT().ReadResource(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mock.EXPECT().GetPrompt(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mock.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(&vmcp.ToolCallResult{Content: []vmcp.Content{{Type: "text", Text: "ok"}}}, nil).AnyTimes() + mock.EXPECT().Close().Return(nil).AnyTimes() + return mock, nil + }).AnyTimes() + return factory +} + +// principalIdentityMiddleware stamps a per-request auth.Identity onto the +// request context, derived from the X-Test-Principal header, mirroring the +// identity-stamping pattern in buildCedarAuthzServer's identityMiddleware +// (authz_integration_test.go). It deliberately leaves Identity.Token empty: +// the mock session factory below binds the session as "unauthenticated", and +// validateCallerBinding's anonymous-session upgrade-attack guard only rejects +// a caller that presents a non-empty Token -- so distinct Subjects with no +// token pass the binding check while still giving each request its own +// identity for coreToolHandler/core.CallTool to observe. A request with no +// header is left unauthenticated, matching every other test in this file that +// never sets this middleware. +func principalIdentityMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + principal := r.Header.Get("X-Test-Principal") + if principal == "" { + next.ServeHTTP(w, r) + return + } + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: principal, + Claims: map[string]any{"sub": principal}, + }} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) +} + +// TestRegression_ConcurrentSameSession_NoIdentityBleed pins the fix for issue +// #156 item U3 (toolhive-core v0.0.32): go-sdk handles messages for an +// existing session on that session's own connection goroutine using the +// initialize-time context, so a per-POST request context (identity included) +// would be lost unless bridged per-request. The shim bridges it via a nonce +// keyed to the individual POST -- NOT the session ID -- specifically so two +// concurrent POSTs on the SAME session cannot clobber each other's identity. +// TestRegression_ConcurrentToolCalls_NoAuditBleed (above) already covers +// concurrent same-session calls, but both requests carry the SAME (anonymous) +// principal, so a regression that collapsed the per-POST bridge back to the +// old session-keyed (or session-wide) context would not be caught there. This +// test fires two concurrent tools/call POSTs on ONE session with DISTINCT +// principals and asserts core.CallTool observed each request's own identity, +// correlated by principal (a stable key), not by arrival order. +func TestRegression_ConcurrentSameSession_NoIdentityBleed(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + testTool := vmcp.Tool{Name: "t"} + fc := &fakeCore{tools: []vmcp.Tool{testTool}} + fc.enableCallIdentityRecording() + factory, _ := newToolSessionFactory(t, ctrl, []vmcp.Tool{testTool}) + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(principalIdentityMiddleware(streamable)) + t.Cleanup(ts.Close) + baseURL := ts.URL + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-11-25", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return ok + }, 2*time.Second, 10*time.Millisecond, "session should be registered") + + principals := []string{"p1", "p2"} + var wg sync.WaitGroup + wg.Add(len(principals)) + callErrs := make([]error, len(principals)) + for i, principal := range principals { + i, principal := i, principal + go func() { + defer wg.Done() + // No require/assert here: FailNow from a worker goroutine only runs + // Goexit off the test goroutine and misreports. Collect the error and + // assert on the test goroutine after wg.Wait() below. + resp, doErr := doServeMCPWithHeaders(baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2 + i, + "method": "tools/call", + "params": map[string]any{ + "name": "t", + // "caller" travels in the request BODY and is echoed to + // fakeCore.CallTool via args, giving a correlator that is + // independent of the context-bridged identity. The recorder + // keys on it so the assertion below can detect a swap, not + // just a collapse. + "arguments": map[string]any{"caller": principal}, + }, + }, sessionID, map[string]string{"X-Test-Principal": principal}) + if doErr != nil { + callErrs[i] = doErr + return + } + defer resp.Body.Close() + _, callErrs[i] = io.ReadAll(resp.Body) + }() + } + + // Fail-fast wait: never block indefinitely on a WaitGroup. + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent same-session calls to complete") + } + for i, callErr := range callErrs { + require.NoError(t, callErr, "call for principal %q must succeed", principals[i]) + } + + require.Equal(t, int32(len(principals)), fc.callToolCalls.Load(), + "core.CallTool must be called exactly once per concurrent request") + + // For each request (identified by its body-supplied "caller"), the identity + // core.CallTool observed must be that request's OWN identity. The recorder + // keys on the caller argument (from the request body) and stores the + // observed identity (from the context bridge) as the value; asserting + // value.Subject == caller therefore correlates the two independent sources. + // This catches BOTH failure modes of the U3 regression: a collapse (one + // caller key never recorded -> NotNil fails) AND a symmetric swap (request + // caller=p1 handled with p2's identity -> callIdentityFor("p1").Subject is + // "p2" -> Equal fails), which a Subject-keyed recorder could not detect. + for _, principal := range principals { + got := fc.callIdentityFor(principal) + require.NotNil(t, got, "no identity was recorded for caller %q; did it bleed into another request?", principal) + assert.Equal(t, principal, got.Subject, + "the request with caller=%q must be handled with its OWN identity, not a concurrent request's", principal) + } +} + +// TestRegression_ConcurrentToolCalls_NoAuditBleed fires two tools/call POSTs +// concurrently against the same session and asserts each response carries its +// OWN tool's result text — no cross-contamination. The Serve path sets the +// audit BackendInfo per-request via the handler closure (coreToolHandler writes +// the pre-resolved backend name), so concurrent calls must not bleed a backend +// name or result from one request into another. +// +// The audit BackendInfo is attached per-request by the audit middleware (a fresh +// *BackendInfo per request, see pkg/audit/auditor.go Middleware), and the Serve +// handler writes to it from the per-request context — so the isolation surface +// under test is the result text and the observed backend label. We assert each +// call's result text matches the tool it invoked. +func TestRegression_ConcurrentToolCalls_NoAuditBleed(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + tools := []vmcp.Tool{{Name: "tool-a"}, {Name: "tool-b"}} + fc := &perToolCore{fakeCore: fakeCore{tools: tools}} + factory := newPerToolSessionFactory(t, ctrl, tools) + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + baseURL := ts.URL + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return ok + }, 2*time.Second, 10*time.Millisecond, "session should be registered") + + type callOutcome struct { + toolName string + body string + err error + } + + var wg sync.WaitGroup + outcomes := make([]callOutcome, len(tools)) + wg.Add(len(tools)) + for i, tool := range tools { + i, tool := i, tool + go func() { + defer wg.Done() + // doServeMCP returns an error instead of calling require: FailNow from a + // worker goroutine only runs Goexit off the test goroutine and misreports. + // All assertions happen on the test goroutine after wg.Wait() below. + resp, doErr := doServeMCP(baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2 + i, + "method": "tools/call", + "params": map[string]any{ + "name": tool.Name, + "arguments": map[string]any{}, + }, + }, sessionID) + if doErr != nil { + outcomes[i] = callOutcome{toolName: tool.Name, err: doErr} + return + } + defer resp.Body.Close() + body, readErr := io.ReadAll(resp.Body) + outcomes[i] = callOutcome{toolName: tool.Name, body: string(body), err: readErr} + }() + } + + // Fail-fast wait: never block indefinitely on a WaitGroup. + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent tool calls to complete") + } + + // Each call must have returned its OWN tool name in the result text, proving + // no result cross-contamination between the concurrent requests. + for _, o := range outcomes { + require.NoError(t, o.err) + assert.Contains(t, o.body, o.toolName, + "concurrent call for %q must carry its own tool's result; got %s", o.toolName, o.body) + } + + // The core must have been reached exactly once per tool (2 total) — no + // duplicated or dropped calls. + assert.Equal(t, int32(len(tools)), fc.callToolCalls.Load(), + "core.CallTool must be called exactly once per concurrent request") + + // Verify coreToolHandler labels the request-scoped audit BackendInfo with the + // backend name it was constructed with. Each goroutine owns a freshly-allocated + // BackendInfo and passes its own backend name, so this is NOT a shared-state + // isolation test — it asserts the handler copies the resolved backend name into + // the per-request BackendInfo (rather than dropping or swapping it). Running it + // concurrently additionally lets the race detector flag any accidental sharing + // of the labelling surface across invocations. + type labelOutcome struct { + want string + got string + hasRes bool + err error + } + var bgWG sync.WaitGroup + bgWG.Add(len(tools)) + labelOutcomes := make([]labelOutcome, len(tools)) + for i, tool := range tools { + i, tool := i, tool + go func() { + defer bgWG.Done() + // No require/assert here: collect results and assert on the test + // goroutine after Wait, since FailNow is only safe there. + bi := &audit.BackendInfo{} + ctx := audit.WithBackendInfo(context.Background(), bi) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: tool.Name, Arguments: map[string]any{}}} + res, err := srv.coreToolHandler(sessionID, tool.Name, tool.Name)(ctx, req) + labelOutcomes[i] = labelOutcome{want: tool.Name, got: bi.BackendName, hasRes: res != nil, err: err} + }() + } + bgDone := make(chan struct{}) + go func() { bgWG.Wait(); close(bgDone) }() + select { + case <-bgDone: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for backend-label goroutines") + } + + // Assert on the test goroutine (require/FailNow is only safe here). Each + // handler invocation must have labelled its own BackendInfo with the backend + // name it was given. + for _, o := range labelOutcomes { + require.NoError(t, o.err) + require.True(t, o.hasRes, "coreToolHandler must return a non-nil result") + assert.Equal(t, o.want, o.got, + "coreToolHandler must copy the resolved backend name %q into the request-scoped audit BackendInfo", + o.want) + } +} diff --git a/pkg/vmcp/server/derive.go b/pkg/vmcp/server/derive.go index 450926c637..ec2e2dcdde 100644 --- a/pkg/vmcp/server/derive.go +++ b/pkg/vmcp/server/derive.go @@ -77,6 +77,7 @@ func deriveServerConfig( EndpointPath: cfg.EndpointPath, SessionTTL: cfg.SessionTTL, WriteTimeout: cfg.WriteTimeout, + HeartbeatInterval: cfg.HeartbeatInterval, AuthMiddleware: cfg.AuthMiddleware, AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, diff --git a/pkg/vmcp/server/derive_test.go b/pkg/vmcp/server/derive_test.go index e2543f24ac..4e35e8ff5b 100644 --- a/pkg/vmcp/server/derive_test.go +++ b/pkg/vmcp/server/derive_test.go @@ -40,6 +40,7 @@ func populatedLegacyConfig() *Config { EndpointPath: "/custom", SessionTTL: 17 * time.Minute, WriteTimeout: 95 * time.Second, + HeartbeatInterval: 5 * time.Second, AuthMiddleware: passthrough, AuthzMiddleware: passthrough, AuthInfoHandler: http.NewServeMux(), @@ -72,6 +73,7 @@ func TestDeriveServerConfigProjectsTransportFields(t *testing.T) { assert.Equal(t, "/custom", got.EndpointPath) assert.Equal(t, 17*time.Minute, got.SessionTTL) assert.Equal(t, 95*time.Second, got.WriteTimeout) + assert.Equal(t, 5*time.Second, got.HeartbeatInterval) assert.Equal(t, 11*time.Second, got.StatusReportingInterval) // Func/handler/pointer fields projected by reference. diff --git a/pkg/vmcp/server/dns_rebinding_regression_test.go b/pkg/vmcp/server/dns_rebinding_regression_test.go new file mode 100644 index 0000000000..b043ab8cb1 --- /dev/null +++ b/pkg/vmcp/server/dns_rebinding_regression_test.go @@ -0,0 +1,114 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" +) + +// TestRegression_DNSRebinding_RejectsForeignHostByDefault pins that the vMCP +// server does NOT pass go-sdk's WithDisableLocalhostProtection option, so a +// Streamable HTTP server bound to a loopback listener retains go-sdk's default +// DNS-rebinding protection: a POST whose Host header names a non-localhost +// value is rejected with 403 before the request reaches the MCP dispatcher. +// +// Only toolhive-core (mcpcompat/server) exercises this behaviour directly +// today; this test pins it at the vMCP integration point so a future change +// that starts threading WithDisableLocalhostProtection(true) into the Serve +// path's streamableOpts (server.go Handler) regresses loudly instead of +// silently reopening the DNS-rebinding hole. +// +// The positive control (httptest's own loopback Host, which the client sets +// automatically) must NOT be rejected, so the negative case cannot pass +// vacuously (e.g. if the whole POST path were broken). +func TestRegression_DNSRebinding_RejectsForeignHostByDefault(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + factory, _ := newToolSessionFactory(t, ctrl, nil) + fc := &fakeCore{} + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + + initBody, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "dns-rebinding-test", "version": "1.0"}, + }, + }) + require.NoError(t, err) + + tests := []struct { + name string + host string // empty means "leave the client-assigned default (loopback) Host" + wantStatus int + }{ + { + name: "foreign Host header is rejected (DNS-rebinding protection)", + host: "evil.example.com", + wantStatus: http.StatusForbidden, + }, + { + name: "loopback Host header is accepted (positive control)", + // Leave req.Host at its client-assigned default (the loopback + // listener's own address) so this case cannot pass vacuously if the + // POST path itself were broken. + host: "", + wantStatus: http.StatusOK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, ts.URL+"/mcp", bytes.NewReader(initBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if tt.host != "" { + // Set the Host FIELD (not a header) — go-sdk's DNS-rebinding check + // inspects http.Request.Host, which net/http populates from the + // request line / Host header on the wire. Setting Header.Set("Host", + // ...) would not exercise the same code path. + req.Host = tt.host + } + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, tt.wantStatus, resp.StatusCode, + "POST with Host=%q must yield status %d", tt.host, tt.wantStatus) + }) + } +} diff --git a/pkg/vmcp/server/forwarding_realbackend_integration_test.go b/pkg/vmcp/server/forwarding_realbackend_integration_test.go new file mode 100644 index 0000000000..000caea625 --- /dev/null +++ b/pkg/vmcp/server/forwarding_realbackend_integration_test.go @@ -0,0 +1,770 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// Forwarding integration fixtures. These exercise the server->client forwarding +// cluster end-to-end: a real in-process backend, mid tools/call, drives +// elicitation/sampling/progress/logging back at its caller; vMCP must relay that +// traffic to the real downstream client on the same session. That session is +// necessarily Legacy (2025-11-25): the Modern revision removed server-initiated +// requests, so the downstream clients here are explicitly pinned to Legacy — +// see legacyPinningRoundTripper. The Modern-client behavior for the same +// eliciting backend is asserted separately in +// TestIntegration_Modern_RealBackend_ElicitingToolFailsCleanly. +const ( + fwdElicitTool = "elicit_tool" + fwdSampleTool = "sample_tool" + fwdProgressTool = "progress_tool" + fwdLogTool = "log_tool" + + fwdSampledSummary = "a short summary" + fwdSampleModel = "test-model" + fwdProgressToken = "tok-1" + fwdLogData = "hello-from-backend" +) + +// forwardingRealBackendTimeout bounds each real-backend forwarding test. These +// exercise async, server-initiated traffic relayed backend -> vMCP -> downstream +// over real in-process HTTP clients; under the full-suite parallel `-race` load +// on CI a single relay can take many seconds. The timeout is deliberately +// generous so a slow-but-working relay is not mistaken for a hang: it is the +// single deadline the tests wait against (waitNotification derives its own +// deadline from the per-test context), so a genuine hang still fails with a +// clear error rather than blocking to the `go test` global timeout. See #5962. +const forwardingRealBackendTimeout = 60 * time.Second + +// startForwardingBackend starts a real in-process MCP backend whose tools drive +// server->client traffic (elicitation, sampling, progress, logging) during a +// tools/call. Returns the backend's /mcp URL. +func startForwardingBackend(t *testing.T) string { + t.Helper() + + srv := mcpserver.NewMCPServer("forwarding-backend", "1.0.0", + mcpserver.WithToolCapabilities(false), + mcpserver.WithLogging(), + ) + + srv.AddTool( + mcpmcp.NewTool(fwdElicitTool, mcpmcp.WithDescription("ask the client to confirm")), + func(ctx context.Context, _ mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + res, err := srv.RequestElicitation(ctx, mcpmcp.ElicitationRequest{ + Params: mcpmcp.ElicitationParams{ + Message: "Confirm?", + RequestedSchema: map[string]any{"type": "object"}, + }, + }) + if err != nil { + return nil, err + } + return mcpmcp.NewToolResultText("action=" + string(res.Action)), nil + }, + ) + + srv.AddTool( + mcpmcp.NewTool(fwdSampleTool, mcpmcp.WithDescription("ask the client to sample")), + func(ctx context.Context, _ mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + res, err := srv.RequestSampling(ctx, mcpmcp.CreateMessageRequest{ + CreateMessageParams: mcpmcp.CreateMessageParams{ + MaxTokens: 100, + Messages: []mcpmcp.SamplingMessage{{ + Role: mcpmcp.RoleUser, + Content: mcpmcp.NewTextContent("summarize this"), + }}, + }, + }) + if err != nil { + return nil, err + } + text, _ := res.Content.(map[string]any)["text"].(string) + return mcpmcp.NewToolResultText("sampled=" + text + " model=" + res.Model), nil + }, + ) + + srv.AddTool( + mcpmcp.NewTool(fwdProgressTool, mcpmcp.WithDescription("emit progress")), + func(ctx context.Context, _ mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + if err := srv.SendNotificationToClient(ctx, "notifications/progress", map[string]any{ + "progressToken": fwdProgressToken, + "progress": 0.5, + "total": 1.0, + "message": "halfway", + }); err != nil { + return nil, err + } + return mcpmcp.NewToolResultText("done"), nil + }, + ) + + srv.AddTool( + mcpmcp.NewTool(fwdLogTool, mcpmcp.WithDescription("emit a log message")), + func(ctx context.Context, _ mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + if err := srv.SendNotificationToClient(ctx, "notifications/message", map[string]any{ + "level": "info", + "data": fwdLogData, + }); err != nil { + return nil, err + } + return mcpmcp.NewToolResultText("logged"), nil + }, + ) + + streamableSrv := mcpserver.NewStreamableHTTPServer(srv) + mux := http.NewServeMux() + mux.Handle("/mcp", streamableSrv) + + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + return ts.URL + "/mcp" +} + +// legacyPinningRoundTripper pins a downstream mcpcompat client to the Legacy +// (2025-11-25) revision by answering its Modern-first server/discover probe +// with a successful DiscoverResult whose supportedVersions lists ONLY +// 2025-11-25 — literally what a Legacy-only server sends, and the same shape +// vMCP's SDK path answers today: classification.go deliberately EXEMPTS +// server/discover from Modern-dispatch gating (rejecting the probe would +// leave a client no way to negotiate down), so the probe reaches the SDK +// stateful path and gets a 200 whose version list omits 2026-07-28. go-sdk's +// Connect finds no mutually supported Modern version in that list and falls +// back to the Legacy initialize handshake (mcp/client.go's discover +// negotiation; any non-nil discover error takes the same fallback). This +// RoundTripper reproduces that version-omitting success answer +// deterministically, independent of what the server under test advertises. +// Every other request passes through to base untouched. +// +// WHY PIN: these fixtures assert the Legacy-only server-initiated surface; +// see the file header above and the client-edge limitation in +// docs/arch/10-virtual-mcp-architecture.md for the full disposition. Modern +// dispatch is served whenever the capability gate is open +// (modernDispatchBlockers, modern_gate.go; #5959 removed the kill-switch), and +// these fixtures configure no blockers, so without +// this pin the go-sdk-based client would negotiate Modern and the surface +// under test would vanish mid-test — failing at connect rather than at the +// behavior being asserted. Before that, the server's own version-omitting +// discover answer pinned these clients to Legacy incidentally; the pin makes +// the dependency explicit instead of relying on that accident. +// +// The pin lives in the transport because mcpcompat documents it cannot set a +// protocol version (go-sdk's ClientSessionOptions.protocolVersion is +// unexported — see the LIMITATION note in mcpcompat/client and #5911). +// Replace this with a real client option if #5911 lands. +// +// LOAD-BEARING after #6033: with the capability gate open (modern_gate.go) -- +// and these fixtures configure no blockers -- this +// RoundTripper is the ONLY thing keeping these downstream clients on Legacy. +// It is not leftover scaffolding — deleting it silently flips every test in +// this file to Modern and voids what they assert. The intercepted counter +// (asserted after every connect) exists to catch exactly that: if go-sdk ever +// stops probing server/discover first, the pin becomes dead code and the +// counter assertion fails instead of the tests silently changing meaning. +type legacyPinningRoundTripper struct { + base http.RoundTripper + // intercepted counts server/discover probes answered with the Legacy-only + // DiscoverResult. Asserted > 0 after connect — proof the pin FIRED, not + // merely that it exists. + intercepted atomic.Int32 +} + +func (rt *legacyPinningRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if req.Method != http.MethodPost || req.Body == nil { + return rt.base.RoundTrip(req) + } + // Clone before touching anything: RoundTrippers must not modify the + // caller's request beyond consuming the body (net/http contract; see also + // the copy-before-mutating rule). The body is the one thing a routing + // interceptor must read, so consume it from the original and give the + // clone a fresh reader. + body, err := io.ReadAll(req.Body) + _ = req.Body.Close() + if err != nil { + return nil, err + } + req = req.Clone(req.Context()) + req.Body = io.NopCloser(bytes.NewReader(body)) + var probe struct { + ID any `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal(body, &probe) == nil && probe.Method == "server/discover" { + rt.intercepted.Add(1) + return legacyOnlyDiscoverResponse(req, probe.ID) + } + return rt.base.RoundTrip(req) +} + +// legacyOnlyDiscoverResponse builds the HTTP 200 server/discover answer of a +// Legacy-only server: a successful JSON-RPC result whose supportedVersions +// lists only 2025-11-25. go-sdk's Connect negotiates against that list, finds +// no Modern version, and falls back to the Legacy initialize handshake — the +// same path vMCP's own version-omitting discover answer drives today. +func legacyOnlyDiscoverResponse(req *http.Request, id any) (*http.Response, error) { + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "result": map[string]any{ + "resultType": "complete", + "supportedVersions": []string{"2025-11-25"}, + "capabilities": map[string]any{}, + }, + }) + if err != nil { + return nil, err + } + hdr := make(http.Header) + hdr.Set("Content-Type", "application/json") + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: hdr, + ContentLength: int64(len(body)), + Body: io.NopCloser(bytes.NewReader(body)), + Request: req, + }, nil +} + +// newLegacyPinnedHTTPClient returns an *http.Client for +// transport.WithHTTPBasicClient that pins the mcpcompat client to Legacy, +// plus the RoundTripper itself so callers can assert the pin actually fired — +// see legacyPinningRoundTripper. +func newLegacyPinnedHTTPClient() (*http.Client, *legacyPinningRoundTripper) { + rt := &legacyPinningRoundTripper{base: http.DefaultTransport} + return &http.Client{Transport: rt}, rt +} + +// requirePinFired asserts the Legacy pin intercepted at least one +// server/discover probe during connect. Without this, the pin is +// indistinguishable from dead code: with Modern dispatch disabled (today's +// default) the tests pass even with the RoundTripper deleted (the server's +// own version-omitting discover answer pins the client), and if go-sdk ever +// stopped probing +// server/discover first the tests would silently stop testing what they +// claim. +func requirePinFired(t *testing.T, rt *legacyPinningRoundTripper) { + t.Helper() + require.Positive(t, rt.intercepted.Load(), + "Legacy pin never fired: the client did not probe server/discover; the pin may be dead code") +} + +// downstreamClient builds a real mcpcompat client against the vMCP endpoint, +// wired for server->client traffic (elicitation/sampling handlers, continuous +// listening) and collecting forwarded notifications on the returned channel. +// It is pinned to the Legacy revision (see legacyPinningRoundTripper): the +// forwarded traffic it collects exists only on a Legacy session. +type downstreamClient struct { + c *client.Client + notifCh chan mcpmcp.JSONRPCNotification + elicited chan struct{} +} + +// newDownstreamClient connects a downstream client to vmcpURL. When +// withHandlers is true it advertises elicitation and sampling and answers those +// requests; when false it advertises neither (the negative path). It always +// registers an OnNotification collector. +func newDownstreamClient(ctx context.Context, t *testing.T, vmcpURL string, withHandlers bool) *downstreamClient { + t.Helper() + return newDownstreamClientOpts(ctx, t, vmcpURL, withHandlers, true) +} + +// newDownstreamClientOpts is newDownstreamClient with explicit control over +// the standalone SSE stream (listen): pass listen=false to model a client that +// advertises elicitation/sampling but never opens the standalone stream. +func newDownstreamClientOpts( + ctx context.Context, t *testing.T, vmcpURL string, withHandlers, listen bool, +) *downstreamClient { + t.Helper() + + dc := &downstreamClient{ + notifCh: make(chan mcpmcp.JSONRPCNotification, 8), + elicited: make(chan struct{}, 1), + } + + var clientOpts []client.ClientOption + if withHandlers { + clientOpts = append(clientOpts, + client.WithElicitationHandler(client.ElicitationHandlerFunc( + func(_ context.Context, _ mcpmcp.ElicitationRequest) (*mcpmcp.ElicitationResult, error) { + select { + case dc.elicited <- struct{}{}: + default: + } + return &mcpmcp.ElicitationResult{ + ElicitationResponse: mcpmcp.ElicitationResponse{ + Action: mcpmcp.ElicitationResponseActionAccept, + Content: map[string]any{"confirmed": true}, + }, + }, nil + }, + )), + client.WithSamplingHandler(client.SamplingHandlerFunc( + func(_ context.Context, _ mcpmcp.CreateMessageRequest) (*mcpmcp.CreateMessageResult, error) { + return &mcpmcp.CreateMessageResult{ + SamplingMessage: mcpmcp.SamplingMessage{ + Role: mcpmcp.RoleAssistant, + Content: mcpmcp.NewTextContent(fwdSampledSummary), + }, + Model: fwdSampleModel, + StopReason: "endTurn", + }, nil + }, + )), + ) + } + + hc, pinRT := newLegacyPinnedHTTPClient() + transportOpts := []transport.StreamableHTTPCOption{ + transport.WithHTTPBasicClient(hc), + } + if listen { + transportOpts = append(transportOpts, transport.WithContinuousListening()) + } + c, err := client.NewStreamableHttpClientWithOpts( + vmcpURL, + transportOpts, + clientOpts, + ) + require.NoError(t, err) + dc.c = c + + c.OnNotification(func(n mcpmcp.JSONRPCNotification) { + select { + case dc.notifCh <- n: + default: + } + }) + + require.NoError(t, c.Start(ctx)) + if !listen { + // Force-close this downstream client's idle keep-alive connections at + // teardown. Registration order matters: t.Cleanup runs + // last-added-first, and c.Close() sends the session-terminating DELETE + // over hc, opening a fresh connection that goes idle — so + // CloseIdleConnections must be registered FIRST to run AFTER c.Close(). + // Note the vMCP server holds its own backend-client connection, so the + // suite's ~30s httptest.Server.Close stall persists regardless; this + // only covers the connection this helper owns. + t.Cleanup(hc.CloseIdleConnections) + } + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcpmcp.InitializeRequest{ + Params: mcpmcp.InitializeParams{ + ProtocolVersion: mcpmcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcpmcp.Implementation{Name: "downstream", Version: "1.0"}, + }, + }) + require.NoError(t, err) + requirePinFired(t, pinRT) + + return dc +} + +// waitNotification blocks for a forwarded notification with the given method, +// or until ctx is done. It waits against the caller's context rather than an +// independent hardcoded timer so there is a single deadline for the test (see +// forwardingRealBackendTimeout): a fixed timer shorter than the context flaked +// under CI `-race` load — the async backend -> vMCP -> downstream relay can take +// well over a second — while a timer longer than the context would mask a hang. +func (dc *downstreamClient) waitNotification(ctx context.Context, t *testing.T, method string) mcpmcp.JSONRPCNotification { + t.Helper() + for { + select { + case n := <-dc.notifCh: + if n.Method == method { + return n + } + case <-ctx.Done(): + t.Fatalf("timed out waiting for %s notification: %v", method, ctx.Err()) + } + } +} + +// TestForwarding_Elicitation_RealBackend verifies a backend's mid-call +// elicitation/create is relayed to the downstream client and its response +// carried back into the tool result. +// +// Legacy-pinned: see legacyPinningRoundTripper and the client-edge limitation +// in docs/arch/10-virtual-mcp-architecture.md. +func TestForwarding_Elicitation_RealBackend(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + dc := newDownstreamClient(ctx, t, vmcpTS.URL+"/mcp", true) + + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdElicitTool}, + }) + require.NoError(t, err) + require.False(t, res.IsError, "elicitation must round-trip to the downstream client") + require.Len(t, res.Content, 1) + txt, ok := mcpmcp.AsTextContent(res.Content[0]) + require.True(t, ok) + assert.Equal(t, "action=accept", txt.Text) + + // The downstream elicitation handler must have actually fired. + select { + case <-dc.elicited: + default: + t.Fatal("downstream elicitation handler was not invoked") + } +} + +// TestForwarding_Sampling_RealBackend verifies a backend's mid-call +// sampling/createMessage is relayed to the downstream client and the sampled +// message carried back into the tool result. +// +// Legacy-pinned: see legacyPinningRoundTripper and the client-edge limitation +// in docs/arch/10-virtual-mcp-architecture.md (which also covers SEP-2577's +// deprecation of sampling itself). +func TestForwarding_Sampling_RealBackend(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + dc := newDownstreamClient(ctx, t, vmcpTS.URL+"/mcp", true) + + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdSampleTool}, + }) + require.NoError(t, err) + require.False(t, res.IsError, "sampling must round-trip to the downstream client") + require.Len(t, res.Content, 1) + txt, ok := mcpmcp.AsTextContent(res.Content[0]) + require.True(t, ok) + assert.Equal(t, "sampled="+fwdSampledSummary+" model="+fwdSampleModel, txt.Text) +} + +// TestForwarding_Progress_RealBackend verifies a backend's mid-call +// notifications/progress is relayed to the downstream client, arriving before +// the tool result is read. +// +// Legacy-pinned because dispatchModern cannot stream a response, NOT because +// Modern lacks a progress channel — it has one (request-scoped notifications +// on the POST SSE response stream, SEP-2260). Full disposition in the +// client-edge limitation in docs/arch/10-virtual-mcp-architecture.md; today's +// Modern behavior (notification dropped, call completes) is pinned by +// TestIntegration_Modern_RealBackend_ProgressDropped. +func TestForwarding_Progress_RealBackend(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + dc := newDownstreamClient(ctx, t, vmcpTS.URL+"/mcp", true) + + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdProgressTool}, + }) + require.NoError(t, err) + require.False(t, res.IsError) + + n := dc.waitNotification(ctx, t, "notifications/progress") + assert.Equal(t, fwdProgressToken, n.Params.AdditionalFields["progressToken"]) + assert.InDelta(t, 0.5, n.Params.AdditionalFields["progress"], 1e-9) + assert.Equal(t, "halfway", n.Params.AdditionalFields["message"]) +} + +// TestForwarding_Logging_RealBackend verifies that vMCP requests debug logging +// from the backend (so it emits) and relays the backend's notifications/message +// to the downstream client, which has itself set a logging level. +// +// Legacy-pinned. The one fact worth keeping AT this test so nobody "fixes" +// vMCP to make it pass on Modern: go-sdk's SetLoggingLevel omits the per-request +// _meta injection every other Modern-aware method performs, so the Modern +// request it sends is MALFORMED (header without _meta.protocolVersion) and +// vMCP's -32020 rejection is CORRECT — accepting it would be worse than the +// failing call. This is a PERMANENT fixture of go-sdk v1.7.x, not a bug to wait +// out: modelcontextprotocol/go-sdk#1116 was closed wont-fix-by-design because +// the 2026-07-28 revision REMOVED the logging/setLevel RPC (the maintainer's +// answer: use the per-request logLevel _meta key instead). The rest of the +// disposition (the RPC removal, the logLevel _meta replacement, SEP-2577's +// deprecation of logging) lives in the client-edge limitation in +// docs/arch/10-virtual-mcp-architecture.md; today's Modern contract is pinned +// by TestIntegration_Modern_RealBackend_LoggingContract, and the Modern log +// opt-in + relay is exercised by TestModernCallTool_LogLevelGating. +func TestForwarding_Logging_RealBackend(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + dc := newDownstreamClient(ctx, t, vmcpTS.URL+"/mcp", true) + + // notifications/message is delivered downstream only once the downstream + // client has set a logging level. + require.NoError(t, dc.c.SetLoggingLevel(ctx, mcpmcp.LoggingLevelDebug)) + + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdLogTool}, + }) + require.NoError(t, err) + require.False(t, res.IsError) + + n := dc.waitNotification(ctx, t, "notifications/message") + assert.Equal(t, "info", n.Params.AdditionalFields["level"]) + assert.Equal(t, fwdLogData, n.Params.AdditionalFields["data"]) +} + +// TestForwarding_Sampling_NoDownstreamCapability verifies the negative path: +// when the downstream client did not advertise sampling, the backend's +// sampling request fails cleanly (a failed tools/call) rather than hanging until +// the deadline. +// +// Legacy-pinned even though it happens to pass on Modern: capability-gated +// failure on a live session only exists on Legacy. On Modern the same call +// fails with the sessionless error REGARDLESS of what the client advertises +// (withHandlers makes no difference), so a lenient assertion here would be +// satisfied for a reason unrelated to the gating this test exists to +// exercise. The pin keeps it testing what its name says. +func TestForwarding_Sampling_NoDownstreamCapability(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + dc := newDownstreamClient(ctx, t, vmcpTS.URL+"/mcp", false) + + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdSampleTool}, + }) + + // Structural, not substring: the call must ROUND-TRIP on the live Legacy + // session (err == nil proves it was not a connect failure — the vacuity + // route) and fail as a tool error. The failure text is deliberately not + // asserted: its identifying substrings are go-sdk's error wrapping, not + // ToolHive's, and testing a dependency's strings breaks on upstream + // rewords (testing.md, test scope). + require.NoError(t, err, + "the call must round-trip on the live session, not die at transport level") + require.NotNil(t, res) + assert.True(t, res.IsError, "backend sampling must fail when downstream lacks the capability") +} + +// TestForwarding_Elicitation_NoDownstreamCapability is the elicitation twin of +// TestForwarding_Sampling_NoDownstreamCapability: when the downstream client did +// not advertise elicitation, the backend's mid-call elicitation/create fails +// cleanly (a failed tools/call) rather than hanging until the deadline. +// Legacy-pinned for the same vacuous-pass reason as its sampling twin above. +func TestForwarding_Elicitation_NoDownstreamCapability(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + dc := newDownstreamClient(ctx, t, vmcpTS.URL+"/mcp", false) + + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdElicitTool}, + }) + + // Same structural shape as the sampling twin — see the rationale there. + require.NoError(t, err, + "the call must round-trip on the live session, not die at transport level") + require.NotNil(t, res) + assert.True(t, res.IsError, "backend elicitation must fail when downstream lacks the capability") +} + +// TestForwarding_Elicitation_AdvertisedButNoStream_FastFails pins the runtime +// twin of TestForwarding_Elicitation_NoDownstreamCapability (#5975): a client +// that ADVERTISED the elicitation capability but holds NO open standalone SSE +// stream passes go-sdk's capability gate, yet the elicitation cannot be +// delivered — under JSONResponse the go-sdk routes server->client requests to +// the standalone stream, and a missing stream rejects the write +// ("rejected by transport: stream not connected or already closed"). +// +// Legacy-pinned like its siblings, for the same vacuous-pass reason: on Modern +// the call fails with the sessionless error regardless of stream state, so a +// Modern run would satisfy the assertions without exercising the delivery +// path this test exists for. +// +// The assertion is timing-structural, not string-matching: with a generous +// outer deadline, the call must fail as a tool error FAR below it (a hang-to- +// timeout regression blows the full deadline instead). This documents and pins +// the fail-fast until an upstream mcpcompat stream-presence accessor lets vMCP +// fail before dispatch with a cleaner error. +func TestForwarding_Elicitation_AdvertisedButNoStream_FastFails(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + // withHandlers=true (capability advertised) but listen=false (no stream). + dc := newDownstreamClientOpts(ctx, t, vmcpTS.URL+"/mcp", true, false) + + start := time.Now() + res, err := dc.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdElicitTool}, + }) + elapsed := time.Since(start) + + // Same structural shape as the capability-gate twin — see the rationale + // there — plus the timing assertion that distinguishes fail-fast from hang. + require.NoError(t, err, + "the call must round-trip on the live session, not die at transport level") + require.NotNil(t, res) + assert.True(t, res.IsError, + "backend elicitation must fail when the downstream holds no standalone stream") + assert.Less(t, elapsed, forwardingRealBackendTimeout/4, + "elicitation without a standalone stream must fail fast, not hang to the deadline") +} + +// samplingClient is a downstream client whose sampling handler returns a +// DISTINGUISHABLE summary+model and counts its own invocations, used to prove +// per-session isolation of forwarded server->client sampling. +type samplingClient struct { + c *client.Client + sampleCalls atomic.Int32 +} + +// newSamplingClient connects a downstream client to vmcpURL whose sampling +// handler always answers with the given summary and model (so the tool result +// distinguishes which client's handler served the request) and increments a +// per-client counter. Pinned to Legacy like newDownstreamClient — see +// legacyPinningRoundTripper. +func newSamplingClient(ctx context.Context, t *testing.T, vmcpURL, summary, model string) *samplingClient { + t.Helper() + + sc := &samplingClient{} + hc, pinRT := newLegacyPinnedHTTPClient() + c, err := client.NewStreamableHttpClientWithOpts( + vmcpURL, + []transport.StreamableHTTPCOption{ + transport.WithContinuousListening(), + transport.WithHTTPBasicClient(hc), + }, + []client.ClientOption{ + client.WithSamplingHandler(client.SamplingHandlerFunc( + func(_ context.Context, _ mcpmcp.CreateMessageRequest) (*mcpmcp.CreateMessageResult, error) { + sc.sampleCalls.Add(1) + return &mcpmcp.CreateMessageResult{ + SamplingMessage: mcpmcp.SamplingMessage{ + Role: mcpmcp.RoleAssistant, + Content: mcpmcp.NewTextContent(summary), + }, + Model: model, + StopReason: "endTurn", + }, nil + }, + )), + }, + ) + require.NoError(t, err) + sc.c = c + + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcpmcp.InitializeRequest{ + Params: mcpmcp.InitializeParams{ + ProtocolVersion: mcpmcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcpmcp.Implementation{Name: "downstream", Version: "1.0"}, + }, + }) + require.NoError(t, err) + requirePinFired(t, pinRT) + return sc +} + +// TestForwarding_Sampling_RealBackend_SessionIsolation proves a backend's +// forwarded sampling request reaches the CALLING session, never another session +// on the same vMCP server. Two downstream clients A and B, each with a +// distinguishable sampling response, call the sampling tool on their OWN session; +// each tool result must reflect THAT client's own response, and each client's +// sampling handler must fire exactly once (for its own call). +// +// This test would FAIL if forwarding routed to the wrong session: if A's tool +// call's sampling request were relayed to B's session, A's tool result would show +// "summary-B" (mismatching the "summary-A" assertion) and the handler counters +// would be lopsided (A=0, B=2 instead of 1/1) — either check trips. +// +// Legacy-pinned: see legacyPinningRoundTripper — per-session routing of a +// server-initiated request presupposes sessions, which Modern removed. +func TestForwarding_Sampling_RealBackend_SessionIsolation(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(t.Context(), forwardingRealBackendTimeout) + defer cancel() + + backendURL := startForwardingBackend(t) + vmcpTS := newRealTestServer(t, backendURL) + + clientA := newSamplingClient(ctx, t, vmcpTS.URL+"/mcp", "summary-A", "model-A") + clientB := newSamplingClient(ctx, t, vmcpTS.URL+"/mcp", "summary-B", "model-B") + + // Call concurrently so the two forwarded sampling round-trips are in flight at + // the same time — the strongest check that each resolves to its own session. + type callResult struct { + res *mcpmcp.CallToolResult + err error + } + resA := make(chan callResult, 1) + resB := make(chan callResult, 1) + go func() { + r, err := clientA.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdSampleTool}, + }) + resA <- callResult{r, err} + }() + go func() { + r, err := clientB.c.CallTool(ctx, mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: fwdSampleTool}, + }) + resB <- callResult{r, err} + }() + + outA := <-resA + outB := <-resB + require.NoError(t, outA.err) + require.NoError(t, outB.err) + require.False(t, outA.res.IsError) + require.False(t, outB.res.IsError) + + txtA, ok := mcpmcp.AsTextContent(outA.res.Content[0]) + require.True(t, ok) + txtB, ok := mcpmcp.AsTextContent(outB.res.Content[0]) + require.True(t, ok) + + // Each tool result must carry the CALLING client's own sampling response. + assert.Equal(t, "sampled=summary-A model=model-A", txtA.Text, + "client A's tool result must reflect A's own sampling response") + assert.Equal(t, "sampled=summary-B model=model-B", txtB.Text, + "client B's tool result must reflect B's own sampling response") + + // Each client's handler fired exactly once — for its own call only. + assert.Equal(t, int32(1), clientA.sampleCalls.Load(), + "client A's sampling handler must fire once, for A's own call") + assert.Equal(t, int32(1), clientB.sampleCalls.Load(), + "client B's sampling handler must fire once, for B's own call") +} diff --git a/pkg/vmcp/server/hidden_tool_regression_test.go b/pkg/vmcp/server/hidden_tool_regression_test.go new file mode 100644 index 0000000000..67fa4233f6 --- /dev/null +++ b/pkg/vmcp/server/hidden_tool_regression_test.go @@ -0,0 +1,280 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "encoding/json" + "io" + "net/http" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/composer" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// backendCallRecorder records the tool name of every call that reaches +// BackendClient.CallTool. Asserting a name is ABSENT here is what proves the +// request was rejected before dispatch: a JSON-RPC error alone would also be +// produced if vMCP had forwarded the call to the backend and then discarded the +// response. +type backendCallRecorder struct { + mu sync.Mutex + names []string +} + +func (r *backendCallRecorder) record(name string) { + r.mu.Lock() + defer r.mu.Unlock() + r.names = append(r.names, name) +} + +func (r *backendCallRecorder) got() []string { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.names...) +} + +// TestRegression_HiddenToolNotDirectlyCallable is the two-directional guard for +// the tool-visibility invariant: +// +// ADVERTISED == DIRECTLY CALLABLE, while the routing table stays deliberately +// WIDER so composite workflow steps can reach hidden tools. +// +// The aggregator puts EVERY backend tool in the routing table, including those +// hidden from tools/list by excludeAllTools, per-workload excludeAll, or filter +// (default_aggregator.go:349) — on purpose, so a composite tool can call a tool +// its operator chose not to expose (#3636). core.CallTool used to resolve a +// direct call against that same table, so a hidden tool was callable by name; on +// the Modern (2026-07-28) path, where tools/call goes straight to core.CallTool, +// that was reachable straight off the wire. +// +// BOTH directions are asserted because they must move in opposite directions, and +// a single-direction test is exactly what let the bug ship: +// +// - "hidden_tool is rejected AND never reaches the backend" fails if the +// advertised-view check in core.CallTool is reverted. +// - "a composite whose step targets hidden_tool still succeeds, and the backend +// DOES receive hidden_tool" fails if someone instead 'fixes' this by pruning +// hidden tools from the routing table — which would silently re-break #3636. +// +// Both eras are covered: Legacy is protected only by the SDK registering +// advertised names on the session, so pinning it here keeps the two paths from +// drifting apart again. +func TestRegression_HiddenToolNotDirectlyCallable(t *testing.T) { + t.Parallel() + + const ( + visibleTool = "visible_tool" + hiddenTool = "hidden_tool" + compositeWF = "wf_using_hidden" + schemaObject = "object" + ) + + // newFixture builds a server whose advertised set is {visible_tool} while the + // routing table also holds hidden_tool, plus one composite workflow whose only + // step targets the hidden tool. factory is the session factory: the Legacy leg + // needs one that stubs GetMetadataValue (enforceSessionBinding reads the + // identity binding through it on every tools/call), which newNoopMockFactory + // does not. + newFixture := func(t *testing.T, factory vmcpsession.MultiSessionFactory) (string, *backendCallRecorder) { + t.Helper() + rec := &backendCallRecorder{} + ts := buildTestServerWithOptions(t, factory, serverOptions{ + tools: []vmcp.Tool{{ + Name: visibleTool, + Description: "advertised", + InputSchema: map[string]any{"type": schemaObject}, + }}, + hiddenTools: []vmcp.Tool{{ + Name: hiddenTool, + Description: "routable but withheld from tools/list", + InputSchema: map[string]any{"type": schemaObject}, + }}, + workflowDefs: map[string]*composer.WorkflowDefinition{ + compositeWF: { + Name: compositeWF, + Steps: []composer.WorkflowStep{ + {ID: "s1", Type: composer.StepTypeTool, Tool: hiddenTool}, + }, + }, + }, + onBackendCall: func(name string) { rec.record(name) }, + }) + return ts.URL, rec + } + + t.Run("modern", func(t *testing.T) { + t.Parallel() + baseURL, rec := newFixture(t, newNoopMockFactory(t)) + + // tools/list must not leak the hidden tool. + listResp, listBody := postModern(t, baseURL, "tools/list", nil, 1, "") + defer listResp.Body.Close() + require.Equal(t, http.StatusOK, listResp.StatusCode, "decoded: %+v", listBody) + advertised := modernToolNames(t, listBody) + assert.Contains(t, advertised, visibleTool) + assert.NotContains(t, advertised, hiddenTool, + "a tool hidden by excludeAll/filter must not appear in tools/list") + + // The advertised tool is callable, proving the fixture's happy path works and + // the rejection below is about visibility, not a broken harness. + okResp, okBody := postModern(t, baseURL, "tools/call", + map[string]any{"name": visibleTool, "arguments": map[string]any{}}, 2, visibleTool) + defer okResp.Body.Close() + assert.Equal(t, http.StatusOK, okResp.StatusCode, "decoded: %+v", okBody) + assert.NotContains(t, okBody, "error", "an advertised tool must remain callable: %+v", okBody) + + // THE REGRESSION: calling the hidden tool by name must be refused with + // -32602 (HTTP 400), matching what the Legacy SDK path answers for a tool it + // never registered — not -32603, and certainly not a success. + hiddenResp, hiddenBody := postModern(t, baseURL, "tools/call", + map[string]any{"name": hiddenTool, "arguments": map[string]any{}}, 3, hiddenTool) + defer hiddenResp.Body.Close() + assert.Equal(t, http.StatusBadRequest, hiddenResp.StatusCode, + "an unadvertised tool name is caller input, so it must map to HTTP 400: %+v", hiddenBody) + errObj, ok := hiddenBody["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope, got %+v", hiddenBody) + assert.Equal(t, float64(-32602), errObj["code"], + "a hidden tool must be answered as an unknown tool, not laundered into -32603") + + // The security assertion: the call was refused BEFORE reaching the backend. + assert.NotContains(t, rec.got(), hiddenTool, + "vMCP must never forward a direct call to a tool hidden from tools/list") + + // THE COUNTERWEIGHT: a composite tool whose step targets the hidden tool must + // still work, and the backend must actually receive the hidden tool name. + // This is what fails if the routing table is pruned instead (#3636). + wfResp, wfBody := postModern(t, baseURL, "tools/call", + map[string]any{"name": compositeWF, "arguments": map[string]any{}}, 4, compositeWF) + defer wfResp.Body.Close() + require.Equal(t, http.StatusOK, wfResp.StatusCode, "decoded: %+v", wfBody) + assert.Contains(t, rec.got(), hiddenTool, + "a composite workflow step must still reach a hidden backend tool (#3636)") + }) + + t.Run("legacy", func(t *testing.T) { + t.Parallel() + // The Serve path sources the advertised set from the aggregator, not the + // factory, so newMockFactory's tool argument only shapes its own routing + // table; it is used here for its GetMetadataValue stub. + factory, _ := newMockFactory(t, gomock.NewController(t), nil) + baseURL, rec := newFixture(t, factory) + sessionID := legacyInitialize(t, baseURL) + + names := listToolNames(t, baseURL, sessionID) + assert.Contains(t, names, visibleTool) + assert.NotContains(t, names, hiddenTool, + "a tool hidden by excludeAll/filter must not appear in tools/list") + + // Legacy has no dispatcher-level classification: the SDK simply never + // registered the hidden tool on the session, so it answers -32602. Asserting + // the decoded error.code (not a substring of the body) is what keeps the two + // eras from drifting apart on the same input -- this is the same assertion + // the Modern leg makes above, so a divergence in either direction fails. + // The message text does differ by era: toolhive-core's mcpcompat rewrites + // go-sdk's `unknown tool "X"` to `tool "X" not found`, so only the code is + // compared. + hiddenErr := legacyCallError(t, baseURL, sessionID, 3, hiddenTool) + assert.Equal(t, float64(-32602), hiddenErr["code"], + "Legacy must also refuse a hidden tool as an unknown tool: %+v", hiddenErr) + assert.NotContains(t, rec.got(), hiddenTool, + "vMCP must never forward a direct call to a tool hidden from tools/list") + + // Same counterweight as the Modern leg. + legacyCall(t, baseURL, sessionID, 4, "tools/call", + map[string]any{"name": compositeWF, "arguments": map[string]any{}}) + assert.Contains(t, rec.got(), hiddenTool, + "a composite workflow step must still reach a hidden backend tool (#3636)") + }) +} + +// modernToolNames extracts result.tools[].name from a Modern tools/list envelope. +func modernToolNames(t *testing.T, decoded map[string]any) []string { + t.Helper() + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "expected a result object, got %+v", decoded) + raw, ok := result["tools"].([]any) + require.True(t, ok, "expected result.tools, got %+v", result) + names := make([]string, 0, len(raw)) + for _, entry := range raw { + tool, ok := entry.(map[string]any) + require.True(t, ok, "expected a tool object, got %T", entry) + name, ok := tool["name"].(string) + require.True(t, ok, "expected tool.name, got %+v", tool) + names = append(names, name) + } + return names +} + +// legacyInitialize performs the Legacy (session-based) handshake and returns the +// session ID every subsequent request must carry. +func legacyInitialize(t *testing.T, baseURL string) string { + t.Helper() + resp := postMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "initialize should succeed") + sessionID := resp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID, "initialize must return a session ID") + return sessionID +} + +// legacyCall sends a Legacy JSON-RPC request on an established session and returns +// the raw response body. Used where only "the call was answered" matters (the +// composite counterweight); assertions on an error envelope go through +// legacyCallError, which decodes. +func legacyCall(t *testing.T, baseURL, sessionID string, id int, method string, params map[string]any) string { + t.Helper() + body := map[string]any{"jsonrpc": "2.0", "id": id, "method": method} + if params != nil { + body["params"] = params + } + resp := postMCP(t, baseURL, body, sessionID) + defer resp.Body.Close() + raw, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "%s should be answered; body: %s", method, string(raw)) + // Sanity: the body must be a JSON-RPC payload (possibly SSE-framed), not empty. + require.NotEmpty(t, raw, "%s returned an empty body", method) + return string(raw) +} + +// legacyCallError issues a Legacy tools/call for toolName and returns the decoded +// JSON-RPC `error` object, failing the test if the response carries no error. +// Decoding (rather than substring-matching the body) is what makes this assertion +// as strong as the Modern leg's: it pins error.code for THIS request rather than +// accepting the digits appearing anywhere in the payload. +func legacyCallError(t *testing.T, baseURL, sessionID string, id int, toolName string) map[string]any { + t.Helper() + resp := postMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "tools/call", + "params": map[string]any{"name": toolName, "arguments": map[string]any{}}, + }, sessionID) + defer resp.Body.Close() + raw, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + var env struct { + Error map[string]any `json:"error"` + } + require.NoError(t, json.Unmarshal(raw, &env), "tools/call response was not JSON: %s", string(raw)) + require.NotNil(t, env.Error, "expected a JSON-RPC error envelope, got: %s", string(raw)) + return env.Error +} diff --git a/pkg/vmcp/server/integration_test.go b/pkg/vmcp/server/integration_test.go index 9b441b9795..f73b6bfb13 100644 --- a/pkg/vmcp/server/integration_test.go +++ b/pkg/vmcp/server/integration_test.go @@ -395,7 +395,7 @@ func TestIntegration_AuditLogging(t *testing.T) { AnyTimes() mockBackendClient.EXPECT(). - CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{ "result": "Sunny, 72°F", @@ -463,8 +463,8 @@ func TestIntegration_AuditLogging(t *testing.T) { // table needed for tool calls and resource reads to be audit-logged correctly. auditSessionFactory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) auditSessionFactory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { mock := sessionmocks.NewMockMultiSession(ctrl) mock.EXPECT().ID().Return(id).AnyTimes() mock.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes() diff --git a/pkg/vmcp/server/modern_capability_refusal.go b/pkg/vmcp/server/modern_capability_refusal.go new file mode 100644 index 0000000000..aca5d127f8 --- /dev/null +++ b/pkg/vmcp/server/modern_capability_refusal.go @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "sync" +) + +// Client capability names as they appear in the draft schema's +// ClientCapabilities object — the keys a Modern client declares under +// _meta's io.modelcontextprotocol/clientCapabilities and the keys a +// MissingRequiredClientCapabilityError reports back in +// data.requiredCapabilities. +const ( + capabilityElicitation = "elicitation" + capabilitySampling = "sampling" +) + +// modernClientCapabilitiesKey is pkg/mcp's metaKeyClientCapabilities, +// reproduced by hand because that constant is unexported — the same precedent +// as modernServerInfoKey in modern_envelope.go. Keep the two in sync. +const modernClientCapabilitiesKey = "io.modelcontextprotocol/clientCapabilities" + +// modernClientDeclaredCapability reports whether the request's _meta declared +// the named client capability. Presence of the key is the declaration (its +// value is a per-capability options object, typically {}), matching how +// ClassifyRevision treats clientCapabilities itself: presence, not content. +func modernClientDeclaredCapability(meta map[string]any, capability string) bool { + caps, _ := meta[modernClientCapabilitiesKey].(map[string]any) + _, ok := caps[capability] + return ok +} + +// capabilityRefusalRecorder captures, in-process, which client capability a +// backend's mid-call server-initiated request needed when there was no +// downstream session to forward it to (Modern ingress). dispatchModern's call +// verbs install one before dispatching; the SDK requester adapters +// (sdk_elicitation_adapter.go, sdk_sampling_adapter.go) record into it at the +// exact point of refusal (mcpcompat's ErrNoActiveSession); the dispatcher +// reads it after a failed call to classify the failure as the draft schema's +// MissingRequiredClientCapabilityError (-32021) when the client did not +// declare the capability. +// +// Why this rides the context rather than a parameter or the error chain: the +// refusal evidence cannot travel back through the return path. The forwarder +// answers the BACKEND's elicitation/sampling request; the backend then fails +// its own tool, and what the dispatcher receives is the backend's error +// STRING — the typed sentinel is laundered away at the wire boundary. A +// Modern request has no MultiSession to hang per-call state on (that absence +// is the very thing being recorded), and the requesters are bound once at +// startup (BindForwarders), so a per-call explicit parameter cannot reach +// them either. This mirrors the audit.BackendInfoFromContext pattern the same +// dispatch verbs already use: a transport-boundary observation channel +// written by a lower layer, not domain data flowing between middleware +// (vmcp-anti-patterns #1 targets the latter). +type capabilityRefusalRecorder struct { + mu sync.Mutex + capability string +} + +// capabilityRefusalKey is the context key for the recorder. +type capabilityRefusalKey struct{} + +// withCapabilityRefusalRecorder returns ctx carrying a fresh recorder, plus +// the recorder itself for the caller to read after dispatch. Installed only +// by dispatchModern's call verbs; on the Legacy/SDK path no recorder exists +// and recordCapabilityRefusal is a benign no-op (Legacy calls have a session, +// so the refusal it observes cannot occur there in the first place). +func withCapabilityRefusalRecorder(ctx context.Context) (context.Context, *capabilityRefusalRecorder) { + rec := &capabilityRefusalRecorder{} + return context.WithValue(ctx, capabilityRefusalKey{}, rec), rec +} + +// recordCapabilityRefusal notes that a mid-call server-initiated request +// needing the named capability was refused for lack of a downstream session. +// The first recorded capability wins — one call verb produces one wire error, +// and the first refusal is what failed the backend tool. No-op when ctx +// carries no recorder. +func recordCapabilityRefusal(ctx context.Context, capability string) { + rec, _ := ctx.Value(capabilityRefusalKey{}).(*capabilityRefusalRecorder) + if rec == nil { + return + } + rec.mu.Lock() + defer rec.mu.Unlock() + if rec.capability == "" { + rec.capability = capability + } +} + +// refused returns the recorded capability name, or "" if no refusal occurred. +func (r *capabilityRefusalRecorder) refused() string { + r.mu.Lock() + defer r.mu.Unlock() + return r.capability +} diff --git a/pkg/vmcp/server/modern_capability_refusal_test.go b/pkg/vmcp/server/modern_capability_refusal_test.go new file mode 100644 index 0000000000..88fa3769fc --- /dev/null +++ b/pkg/vmcp/server/modern_capability_refusal_test.go @@ -0,0 +1,155 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +func TestModernClientDeclaredCapability(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + meta map[string]any + capability string + want bool + }{ + { + name: "nil meta", + meta: nil, + capability: capabilityElicitation, + want: false, + }, + { + name: "no clientCapabilities key", + meta: map[string]any{"other": "x"}, + capability: capabilityElicitation, + want: false, + }, + { + name: "clientCapabilities not an object", + meta: map[string]any{modernClientCapabilitiesKey: "elicitation"}, + capability: capabilityElicitation, + want: false, + }, + { + name: "empty declaration", + meta: map[string]any{modernClientCapabilitiesKey: map[string]any{}}, + capability: capabilityElicitation, + want: false, + }, + { + name: "declared with empty options object", + meta: map[string]any{modernClientCapabilitiesKey: map[string]any{ + capabilityElicitation: map[string]any{}, + }}, + capability: capabilityElicitation, + want: true, + }, + { + name: "different capability declared", + meta: map[string]any{modernClientCapabilitiesKey: map[string]any{ + capabilitySampling: map[string]any{}, + }}, + capability: capabilityElicitation, + want: false, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tc.want, modernClientDeclaredCapability(tc.meta, tc.capability)) + }) + } +} + +func TestCapabilityRefusalRecorder(t *testing.T) { + t.Parallel() + + t.Run("no recorder in context is a no-op", func(t *testing.T) { + t.Parallel() + // Must not panic; there is simply nothing to observe. + recordCapabilityRefusal(context.Background(), capabilityElicitation) + }) + + t.Run("records and reads back", func(t *testing.T) { + t.Parallel() + ctx, rec := withCapabilityRefusalRecorder(context.Background()) + assert.Empty(t, rec.refused()) + recordCapabilityRefusal(ctx, capabilitySampling) + assert.Equal(t, capabilitySampling, rec.refused()) + }) + + t.Run("first recorded capability wins", func(t *testing.T) { + t.Parallel() + ctx, rec := withCapabilityRefusalRecorder(context.Background()) + recordCapabilityRefusal(ctx, capabilityElicitation) + recordCapabilityRefusal(ctx, capabilitySampling) + assert.Equal(t, capabilityElicitation, rec.refused()) + }) +} + +// TestSDKAdapters_RecordCapabilityRefusal verifies both SDK requester adapters +// record the refusal into a context recorder exactly when the mcpcompat SDK +// refuses for lack of a session (ErrNoActiveSession) — and never for any other +// error, which would misclassify e.g. a client decline as a capability gap. +func TestSDKAdapters_RecordCapabilityRefusal(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + sdkErr error + wantRecorded bool // per-adapter capability name asserted below + }{ + { + name: "no active session records the capability", + sdkErr: server.ErrNoActiveSession, + wantRecorded: true, + }, + { + name: "wrapped no active session records too", + sdkErr: errors.Join(errors.New("outer"), server.ErrNoActiveSession), + wantRecorded: true, + }, + { + name: "other errors do not record", + sdkErr: errors.New("client declined"), + wantRecorded: false, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + elicit := &sdkElicitationAdapter{mcpServer: &fakeSDKElicitationRequester{err: tc.sdkErr}} + ctx, rec := withCapabilityRefusalRecorder(context.Background()) + _, err := elicit.RequestElicitation(ctx, vmcp.ElicitationRequest{Message: "m"}) + require.Error(t, err, "the SDK error must still propagate unchanged") + if tc.wantRecorded { + assert.Equal(t, capabilityElicitation, rec.refused()) + } else { + assert.Empty(t, rec.refused()) + } + + sample := &sdkSamplingAdapter{mcpServer: &fakeSDKSamplingRequester{err: tc.sdkErr}} + ctx, rec = withCapabilityRefusalRecorder(context.Background()) + _, err = sample.RequestSampling(ctx, vmcp.SamplingRequest{}) + require.Error(t, err, "the SDK error must still propagate unchanged") + if tc.wantRecorded { + assert.Equal(t, capabilitySampling, rec.refused()) + } else { + assert.Empty(t, rec.refused()) + } + }) + } +} diff --git a/pkg/vmcp/server/modern_dispatch.go b/pkg/vmcp/server/modern_dispatch.go new file mode 100644 index 0000000000..e7d79d4dc1 --- /dev/null +++ b/pkg/vmcp/server/modern_dispatch.go @@ -0,0 +1,688 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// Standard JSON-RPC 2.0 reserved error codes (spec-fixed, never change). Kept +// as a local block, used by both this file and modern_envelope.go's +// writeModernError status mapping, rather than imported from mcpcompat (which +// defines equivalents like mcp.METHOD_NOT_FOUND): mcpcompat is the SDK's +// wire-protocol vocabulary, while the Modern vMCP layer already sources its +// own codes from two other places -- mcpparser.JSONRPCCodeDenied (403, shared +// with the Legacy call gate) and the classifier's app-space -3202x codes. +// Pulling these four from mcpcompat too would split one small, unchanging set +// of constants across three packages for no benefit. +const ( + jsonRPCCodeInvalidRequest = -32600 + jsonRPCCodeMethodNotFound = -32601 + jsonRPCCodeInvalidParams = -32602 + jsonRPCCodeInternalError = -32603 +) + +// dispatchModern serves a single MCP 2026-07-28 ("Modern") stateless request +// by dispatching directly to the stateless vMCP core, bypassing the SDK +// Serve/session layer entirely. classifyingHandler routes here for every +// well-formed Modern request. +// +// Because this path bypasses the SDK server, it re-homes the SDK's +// pre-dispatch authorization gate itself (see the per-method blocks below) -- +// mirroring authzCallGate exactly, including its fail-open posture on +// non-authorization Check* errors. Do not add a case here without also +// deciding its gating: an ungated write would let a Cedar-denied call reach a +// backend. +func (s *Server) dispatchModern(w http.ResponseWriter, r *http.Request, parsed *mcpparser.ParsedMCPRequest) { + // A notification (no id) MUST get 202 with no body and no dispatch, per the + // Streamable HTTP spec's handling of a POST body containing only + // responses/notifications. parser.go's real parse path sets IsRequest true + // for every decoded jsonrpc2.Request -- calls AND notifications alike -- so + // IsRequest cannot distinguish them; absent id (nil) is the actual + // notification signal (parseMCPRequest leaves ParsedMCPRequest.ID nil when + // the JSON-RPC id is absent). + if parsed.ID == nil { + w.WriteHeader(http.StatusAccepted) + return + } + + // Defensive/unreachable today -- ParsingMiddleware rejects a + // JSON-RPC batch (leading '[') with HTTP 400 / -32600 before a + // ParsedMCPRequest is ever built (parser.go's IsBatchRequest check, ~line + // 119), so dispatchModern never sees one and IsBatch is hardcoded false + // (parser.go ~line 61-65). This also closes the batch blind spot + // call_gate.go used to document. Do not build batch parsing here; this + // guard is just a backstop if the parser ever stops rejecting batches + // upstream. + if parsed.IsBatch { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidRequest, "batch requests are not supported") + return + } + + ctx := r.Context() + // Sanctioned transport-boundary identity read (matches authzCallGate and + // every Serve-path handler). ctx itself is passed unmodified into every + // core.* call below: the stateless backend client reads forwarded headers + // off this exact context per call, so detaching or wrapping it would + // silently break forwarded-header backend auth. + identity, _ := auth.IdentityFromContext(ctx) + + switch parsed.Method { + case "tools/list": + s.dispatchModernToolsList(ctx, w, parsed, identity) + case "resources/list": + s.dispatchModernResourcesList(ctx, w, parsed, identity) + case "resources/templates/list": + s.dispatchModernResourceTemplatesList(ctx, w, parsed, identity) + case "prompts/list": + s.dispatchModernPromptsList(ctx, w, parsed, identity) + case methodServerDiscover: + s.dispatchModernDiscover(ctx, w, parsed, identity) + case "tools/call": + s.dispatchModernToolCall(ctx, w, parsed, identity) + case "resources/read": + s.dispatchModernResourceRead(ctx, w, parsed, identity) + case "prompts/get": + s.dispatchModernPromptGet(ctx, w, parsed, identity) + case "completion/complete": + s.dispatchModernComplete(ctx, w, parsed, identity) + case methodSubscriptionsListen: + // Ungated, same bucket as the list verbs and discover -- see + // dispatchModernSubscriptionsListen for why, and for the one future + // change that would require gating it. + s.dispatchModernSubscriptionsListen(ctx, w, parsed, identity) + case "ping": + // ping does NOT exist in 2026-07-28: `ping` appears nowhere in + // schema/draft/schema.ts, and go-sdk lists it among the methods removed + // for this revision (server.go:1880), answering -32601. So answering it + // at all is deliberate LENIENCY toward a client that pings anyway, not + // spec conformance -- an earlier version of this comment claimed the + // latter, which was wrong twice over: the method is gone, and the bare {} + // below omits `resultType`, which schema.ts's Result MUSTs on every + // result ("Servers implementing this protocol version MUST include this + // field"). + // + // Kept as-is rather than switched to -32601 because that is a behavior + // change to a pre-existing path, out of scope for the change that + // introduced this comment fix. It is inert in practice: a go-sdk Modern + // client never pings (startKeepalive sits past an early return), so + // nothing observable depends on either answer today. If you are here to + // tighten it, -32601 is the conformant answer. + // + // It is ungated (unauthenticated liveness, same bucket as initialize -- + // no Check*). Do not route it through the envelope builders above: they + // would stamp resultType and _meta.serverInfo, which the Legacy SDK path + // does not do for ping either (annotateServerInfo and + // setCompleteResultType both early-return for it, go-sdk + // server.go:1929-1945,1992), so the bare {} at least keeps the two paths + // answering alike. + writeModernResult(w, parsed.ID, struct{}{}) + default: + writeModernError(w, parsed.ID, jsonRPCCodeMethodNotFound, "method not found") + } +} + +// The four list-dispatch helpers below (tools/list, resources/list, +// resources/templates/list, prompts/list) each page the full +// admission-filtered set from the matching core.List* through paginateModern, +// emitting a nextCursor while items remain. The cursor is stateless keyset +// pagination over the aggregated ordering -- see modern_pagination.go for the +// wire contract and why a self-describing cursor is what a sessionless +// revision requires. This is unrelated to the aggregator's UPSTREAM +// cursor-following for internal discovery (#5851); that's a different layer. +// +// A List*/Discover failure logs the full error server-side and returns a +// generic -32603 message to the client (writeModernListError below): unlike +// the call/read/get verbs, these errors come from aggregation and routing +// plumbing (backend IDs, upstream addressing), and security.md forbids +// leaking that detail to callers. An invalid CURSOR is the exception: it is +// caller input rather than a server-side fault, so writeModernPageError maps +// it to -32602 instead. +func (s *Server) dispatchModernToolsList( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + tools, err := s.core.ListTools(ctx, identity) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + cursor, err := modernRequestCursor(parsed.Params) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + page, next, err := paginateModern(tools, func(t vmcp.Tool) string { return t.Name }, + cursorKindTools, cursor) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + result, err := newModernToolsList(page, s.config.Name, s.config.Version, next) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + writeModernResult(w, parsed.ID, result) +} + +func (s *Server) dispatchModernResourcesList( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + resources, err := s.core.ListResources(ctx, identity) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + cursor, err := modernRequestCursor(parsed.Params) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + page, next, err := paginateModern(resources, func(r vmcp.Resource) string { return r.URI }, + cursorKindResources, cursor) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + writeModernResult(w, parsed.ID, newModernResourcesList(page, s.config.Name, s.config.Version, next)) +} + +func (s *Server) dispatchModernResourceTemplatesList( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + templates, err := s.core.ListResourceTemplates(ctx, identity) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + cursor, err := modernRequestCursor(parsed.Params) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + page, next, err := paginateModern(templates, func(t vmcp.ResourceTemplate) string { return t.URITemplate }, + cursorKindResourceTemplates, cursor) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + writeModernResult(w, parsed.ID, newModernResourceTemplatesList(page, s.config.Name, s.config.Version, next)) +} + +func (s *Server) dispatchModernPromptsList( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + prompts, err := s.core.ListPrompts(ctx, identity) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + cursor, err := modernRequestCursor(parsed.Params) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + page, next, err := paginateModern(prompts, func(p vmcp.Prompt) string { return p.Name }, + cursorKindPrompts, cursor) + if err != nil { + writeModernPageError(ctx, w, parsed.ID, parsed.Method, err) + return + } + writeModernResult(w, parsed.ID, newModernPromptsList(page, s.config.Name, s.config.Version, next)) +} + +// dispatchModernDiscover serves server/discover, Modern's replacement for +// initialize+capability negotiation, as a post-admission capability-flags +// envelope: it calls core.Discover, which applies the same admission-filtered +// code paths the four list verbs use and collapses each to presence/absence, +// returning NO descriptor arrays -- so the response reflects only what this +// identity may reach (ListBackends's filterUnauthorized=true is the existing +// post-admission-presence precedent, core_vmcp.go:446). Like the list verbs, +// this is ungated: there is no separate Check* for discover, since it leaks +// no more than tools/list already does. +// +// This runs ONE backend fan-out per call via core.Discover (aggregatedView is +// uncached, so calling ListTools/ListResources/ListResourceTemplates/ +// ListPrompts independently here used to cost four -- and those four weren't +// even a consistent snapshot of the aggregated view). A single fan-out per +// request is fine for now, but a probe the spec expects to be cheap across +// requests too. There is no cross-request cache; add a short-TTL per-identity +// capability cache only if profiling shows the per-request fan-out cost +// matters (#5761, tracked separately, not blocking here). +func (s *Server) dispatchModernDiscover( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + caps, err := s.core.Discover(ctx, identity) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + result := newModernDiscover( + caps.HasTools, caps.HasResources, caps.HasResourceTemplates, caps.HasPrompts, + s.config.Name, s.config.Version, + ) + writeModernResult(w, parsed.ID, result) +} + +// dispatchModernToolCall re-homes authzCallGate's tools/call branch plus the +// post-dispatch TOCTOU reclassification (see writeModernDispatchError). +func (s *Server) dispatchModernToolCall( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + if hasNonObjectArguments(parsed.Params) { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidParams, "arguments must be an object") + return + } + if s.authzGateEnabled && gateDenied(ctx, parsed.Method, + s.core.CheckToolCall(ctx, identity, parsed.ResourceID, parsed.Arguments)) { + writeModernDenied(w, parsed.ID, vmcp.DenyMessageToolCall) + return + } + // The refusal recorder classifies a backend tool that demanded mid-call + // elicitation/sampling with no downstream session to forward it to — see + // modern_capability_refusal.go and writeModernCallFailure. + ctx, refusal := withCapabilityRefusalRecorder(ctx) + result, err := s.core.CallTool(ctx, identity, parsed.ResourceID, parsed.Arguments, parsed.Meta) + if err != nil { + // An unadvertised tool name is caller input, not a server fault, so it must + // not launder into writeModernCallFailure's generic -32603. core.CallTool + // holds tools/call to the advertised view (core_calls.go), so this is the + // answer for a name hidden from tools/list by excludeAllTools/excludeAll/ + // filter -- matching what the Legacy path already returns for a tool it + // never registered on the session (-32602, via toolhive-core mcpcompat's + // translateUnknownToolError, which rewrites go-sdk's "unknown tool" message + // to `tool "X" not found`), which writeModernError maps to HTTP 400. Same + // code either way; the two eras differ only in message text. + // + // Deliberately NOT folded into writeModernDispatchError: that helper is + // shared with resources/read, prompts/get and completion/complete, whose + // not-found classification is a separate decision. The message names no + // tool: an authorization denial is classified ahead of this (in + // core.CallTool, which authorizes before checking the advertised view). + // Omitting the name is a conservative choice rather than a mitigation -- + // a denial already answers 403 + JSONRPCCodeDenied against this 400 + + // -32602, so the two are distinguishable either way. + // + // ErrAuthorizationFailed is excluded explicitly so this cannot depend on + // branch order: authorizeToolCall wraps with a double %w (core_checks.go:84), + // so errors.Is matches through both, and an Admission implementation whose + // error happened to carry ErrNotFound would otherwise answer 400 instead of + // 403. Not reachable today (pkg/authz does not import pkg/vmcp), but + // ErrNotFound is exported and Admission is a public interface. + if errors.Is(err, vmcp.ErrNotFound) && !errors.Is(err, vmcp.ErrAuthorizationFailed) { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidParams, "unknown tool") + return + } + writeModernCallFailure(w, parsed, refusal, vmcp.DenyMessageToolCall, err) + return + } + // Label the audit backend on the success path only. The stateless dispatcher + // cannot pre-resolve the backend (routing is core-internal), so unlike the + // Legacy handlers -- which set the label before the call and thus keep it on + // backend-call failures -- a Modern backend-call failure audits without + // backend_name. Accepted: the event still records the tool and the outcome. + // (Same applies to resources/read and prompts/get below.) + if result.BackendID != "" { + if bi, ok := audit.BackendInfoFromContext(ctx); ok && bi != nil { + bi.BackendName = s.backendDisplayName(ctx, result.BackendID) + } + } + writeModernResult(w, parsed.ID, newModernCallToolResult(result, s.config.Name, s.config.Version)) +} + +// dispatchModernResourceRead re-homes authzCallGate's resources/read branch +// plus the post-dispatch TOCTOU reclassification (see writeModernDispatchError). +func (s *Server) dispatchModernResourceRead( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + if s.authzGateEnabled && gateDenied(ctx, parsed.Method, + s.core.CheckResourceRead(ctx, identity, parsed.ResourceID)) { + writeModernDenied(w, parsed.ID, vmcp.DenyMessageResourceRead) + return + } + ctx, refusal := withCapabilityRefusalRecorder(ctx) + result, err := s.core.ReadResource(ctx, identity, parsed.ResourceID) + if err != nil { + writeModernCallFailure(w, parsed, refusal, vmcp.DenyMessageResourceRead, err) + return + } + if result.BackendID != "" { + if bi, ok := audit.BackendInfoFromContext(ctx); ok && bi != nil { + bi.BackendName = s.backendDisplayName(ctx, result.BackendID) + } + } + writeModernResult(w, parsed.ID, newModernReadResourceResult(result, s.config.Name, s.config.Version)) +} + +// dispatchModernPromptGet re-homes authzCallGate's prompts/get branch plus the +// post-dispatch TOCTOU reclassification (see writeModernDispatchError). +func (s *Server) dispatchModernPromptGet( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + // hasNonObjectArguments only checks object-ness, not per-value typing: the + // SDK's GetPromptParams.Arguments is map[string]string, so it also rejects + // a non-string argument VALUE (e.g. {"x":123}) at decode. Modern accepts + // that shape -- narrower parity (object-shape only), not a behavior gap + // worth closing here. + if hasNonObjectArguments(parsed.Params) { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidParams, "arguments must be an object") + return + } + if s.authzGateEnabled && gateDenied(ctx, parsed.Method, + s.core.CheckPromptGet(ctx, identity, parsed.ResourceID)) { + writeModernDenied(w, parsed.ID, vmcp.DenyMessagePromptGet) + return + } + ctx, refusal := withCapabilityRefusalRecorder(ctx) + result, err := s.core.GetPrompt(ctx, identity, parsed.ResourceID, parsed.Arguments) + if err != nil { + writeModernCallFailure(w, parsed, refusal, vmcp.DenyMessagePromptGet, err) + return + } + if result.BackendID != "" { + if bi, ok := audit.BackendInfoFromContext(ctx); ok && bi != nil { + bi.BackendName = s.backendDisplayName(ctx, result.BackendID) + } + } + writeModernResult(w, parsed.ID, newModernGetPromptResult(result, s.config.Name, s.config.Version)) +} + +// modernCompleteWireParams is the completion/complete request params, decoded +// directly from parsed.Params. It mirrors go-sdk's CompleteParams/ +// CompleteReference/CompleteParamsArgument/CompleteContext field-for-field +// (protocol.go:577-648 in go-sdk@v1.7.0-pre.3) rather than reusing an SDK +// type: mcpcompat/mcp (ToolHive's import, now built on go-sdk v1.7.0-pre.3) +// carries an equivalent CompleteParams, but its Ref field is untyped (any), so +// there is no strongly-typed Modern reference shape to decode into directly. +// It stays local rather than adding JSON tags to vmcp.CompletionRef -- the +// domain type intentionally carries no wire coupling (anti-pattern #5, no +// mcp-go types crossing the core boundary). +type modernCompleteWireParams struct { + Ref *struct { + Type string `json:"type"` + Name string `json:"name,omitempty"` + URI string `json:"uri,omitempty"` + } `json:"ref"` + Argument struct { + Name string `json:"name"` + Value string `json:"value"` + } `json:"argument"` + Context *struct { + Arguments map[string]string `json:"arguments,omitempty"` + } `json:"context,omitempty"` +} + +// dispatchModernComplete serves completion/complete. Unlike tools/call, +// resources/read, and prompts/get, there is no pre-dispatch Check* gate here +// -- call_gate.go documents this as a conscious choice: core.Complete +// authorizes the underlying prompt/resource ref at dispatch (the same +// get/read decision GetPrompt/ReadResource enforce), so gating on the wire +// would just duplicate that check ahead of an admission decision that isn't +// argument-conditional the way the gate's fast path assumes. An admission +// denial from core.Complete still reclassifies to 403 via +// writeModernDispatchError, exactly like the three gated verbs. +// +// This handler does not label the audit BackendInfo the way the three gated +// verbs above do: the Legacy coreCompletionHandler never set backend_name +// either, so completion is a pre-existing gap on both paths, not something +// introduced here. +func (s *Server) dispatchModernComplete( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + var params modernCompleteWireParams + if err := json.Unmarshal(parsed.Params, ¶ms); err != nil || params.Ref == nil || params.Ref.Type == "" { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidParams, "invalid completion/complete params: missing ref") + return + } + if params.Argument.Name == "" { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidParams, "invalid completion/complete params: missing argument.name") + return + } + ref := vmcp.CompletionRef{Type: params.Ref.Type, Name: params.Ref.Name, URI: params.Ref.URI} + + var contextArgs map[string]string + if params.Context != nil { + contextArgs = params.Context.Arguments + } + + result, err := s.core.Complete(ctx, identity, ref, params.Argument.Name, params.Argument.Value, contextArgs) + if err != nil { + writeModernDispatchError(w, parsed.ID, completionDenyMessage(ref.Type), err) + return + } + writeModernResult(w, parsed.ID, newModernComplete(result, s.config.Name, s.config.Version)) +} + +// hasNonObjectArguments reports whether parsed.Params carries an "arguments" +// field that is present but NOT a JSON object (e.g. a string or array). +// +// The parser (handleNamedResourceMethod, parser.go:307) type-asserts +// paramsMap["arguments"].(map[string]interface{}) and silently drops the +// value to nil on a mismatch -- indistinguishable from "arguments absent" by +// the time ParsedMCPRequest.Arguments is built. The SDK path also rejects +// this shape before authz/the core is ever reached (coreToolHandler +// shape-checks req.Params.Arguments in serve_handlers.go; prompts/get gets +// the same pre-dispatch rejection for free because mcpcompat's +// GetPromptParams.Arguments is a concrete map[string]string, so a non-object +// value fails JSON decode before the handler runs) -- this function matches +// that TIMING, not the SDK's wire shape: the SDK's tools/call rejection +// surfaces as a 200 IsError tool result (conversion path), whereas this is a +// genuine JSON-RPC -32602, consistent with Modern's other protocol-level +// rejections (-32600/-32601). Modern must reject the same shape here, on the +// raw params, before that type information is lost -- otherwise a non-object +// arguments value silently authorizes and dispatches as a no-args call, +// diverging from the SDK path and potentially changing an +// argument-conditional authz decision. An absent or explicit-null "arguments" +// is a legitimate no-args call and is not rejected. +func hasNonObjectArguments(params json.RawMessage) bool { + var raw struct { + Arguments json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal(params, &raw); err != nil || raw.Arguments == nil { + return false + } + var obj map[string]any + return json.Unmarshal(raw.Arguments, &obj) != nil +} + +// gateDenied runs the PRE-dispatch admission classification for a gated +// method's Check* result, mirroring authzCallGate exactly: only an +// errors.Is(checkErr, vmcp.ErrAuthorizationFailed) denial returns true. Any +// other error falls through to the WARN+admit branch below, but only +// CheckToolCall can actually produce one: it re-aggregates +// (c.aggregatedView) and returns that error unwrapped on failure, so a +// tools/call gate can fail OPEN on an aggregation/backend-plumbing outage. +// CheckResourceRead and CheckPromptGet need no aggregated view and always +// wrap their error as vmcp.ErrAuthorizationFailed (core_checks.go), so their +// gates never take this fail-open path in practice. This WARN is the only +// operational signal of that fail-open outage admitting traffic; do not +// remove it. +func gateDenied(ctx context.Context, method string, checkErr error) bool { + if checkErr == nil { + return false + } + if errors.Is(checkErr, vmcp.ErrAuthorizationFailed) { + return true + } + slog.WarnContext(ctx, "vmcp authz gate: non-authorization error, admitting request", + "method", method, "error", checkErr) + return false +} + +// writeModernListError logs a List*/Discover failure server-side with the +// full error and writes a generic -32603 message to the client. Unlike +// writeModernDispatchError's call/read/get verbs, these errors surface +// aggregation and routing plumbing (backend IDs, upstream addressing), and +// security.md forbids exposing that detail to callers. +func writeModernListError(ctx context.Context, w http.ResponseWriter, id any, method string, err error) { + slog.ErrorContext(ctx, "vmcp modern dispatch: list/discover failed", "method", method, "error", err) + writeModernError(w, id, jsonRPCCodeInternalError, "internal error") +} + +// writeModernPageError classifies a paginateModern failure. An invalid cursor is +// bad caller input, so it gets -32602 -- matching the spec's "handle invalid +// cursors gracefully" and go-sdk, which returns ErrInvalidParams for a cursor it +// cannot decode. The message deliberately does not say WHY the cursor was +// rejected: clients must treat cursors as opaque, so describing the internal +// encoding would invite them to construct one. Anything else here is an encode +// failure, i.e. a server-side fault, and falls through to -32603. +func writeModernPageError(ctx context.Context, w http.ResponseWriter, id any, method string, err error) { + if errors.Is(err, errInvalidModernCursor) { + writeModernError(w, id, jsonRPCCodeInvalidParams, "invalid cursor") + return + } + writeModernListError(ctx, w, id, method, err) +} + +// writeModernCallFailure classifies a POST-dispatch error from the three call +// verbs (tools/call, resources/read, prompts/get), layering the mid-call +// capability-refusal case on top of writeModernDispatchError's authz/generic +// classification: +// +// - Authorization denials keep absolute priority (403 + denyMsg, so the +// audit middleware still logs "denied") — a refusal can never mask one. +// - A recorded refusal with the capability NOT declared in the request's +// _meta clientCapabilities is the draft schema's +// MissingRequiredClientCapabilityError: code -32021 with +// data.requiredCapabilities naming what the server needed (an actionable +// code a client can surface, where -32603 is opaque) — served at HTTP +// 200, a deliberate, documented deviation from the spec-mandated 400; +// see writeModernMissingCapability. +// - A recorded refusal with the capability DECLARED is deliberately NOT +// -32021: the client did declare it, so blaming the client would be +// wrong, and declaring the capability on a retry cannot help. The spec's answer for a declared +// capability is resultType "input_required" (multi-round retrieval, +// SEP-2322), which this dispatcher does not implement — and the +// 2026-07-28 error vocabulary has no "operation not supported" code, so +// there is no conformant code for this case at all (see the "unavailable +// to Modern clients" limitation in +// docs/arch/10-virtual-mcp-architecture.md). It stays -32603, with a +// vMCP-owned message that names the real cause instead of the backend's +// laundered error string. +// +// Do not fold this into writeModernDispatchError: completion/complete also +// uses that helper but never installs a refusal recorder (SEP-2322 does not +// extend to completion/complete, so an input_required-shaped failure there +// has no spec-defined classification). +func writeModernCallFailure( + w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, refusal *capabilityRefusalRecorder, + denyMsg string, err error, +) { + if capName := refusal.refused(); capName != "" && !errors.Is(err, vmcp.ErrAuthorizationFailed) { + if !modernClientDeclaredCapability(parsed.Meta, capName) { + writeModernMissingCapability(w, parsed.ID, capName) + return + } + writeModernError(w, parsed.ID, jsonRPCCodeInternalError, fmt.Sprintf( + "backend requires the %q client capability mid-call; serving it on protocol %s requires "+ + "multi-round retrieval (SEP-2322), which this server does not implement", + capName, mcpparser.MCPVersionModern)) + return + } + writeModernDispatchError(w, parsed.ID, denyMsg, err) +} + +// writeModernMissingCapability writes the MissingRequiredClientCapabilityError +// (-32021, data.requiredCapabilities typed as a ClientCapabilities object) at +// HTTP 200. +// +// STATUS DEVIATION, deliberate and load-bearing: SEP-2575 MUSTs HTTP 400 for +// this error, but go-sdk's streamable client (v1.7.0-pre.3) treats any +// non-transient 4xx as a CONNECTION failure, not a call failure — its +// transient set is only 500/502/503/504/429, so a 400 falls through +// checkResponse to fail(), which closes the session permanently. One refused +// elicitation would kill every subsequent request on the client. The JSON-RPC +// body below stays fully conformant (which is what a non-go-sdk client +// parses); only the transport status deviates. Tracked upstream as +// go-sdk#1117 — revisit the 200 when that is fixed. +// +// This is written directly rather than through mcpparser.WriteClassificationError +// on purpose: that helper hard-codes HTTP 400, which is correct for its other +// callers (pre-dispatch classification rejections a client never retries on a +// live session) and must not be changed for this one. +// +// The message names both the capability AND the gateway limitation, so a +// caller that reacts to -32021 by declaring the capability on a retry +// learns immediately that doing so will not help here (the declared case is served by +// writeModernCallFailure's -32603 branch, not by MRTR). +// +// id follows writeModernError (modern_envelope.go): absent (via +// transportsession.HasJSONRPCID) is encoded by omitting the "id" key, never +// as null. +func writeModernMissingCapability(w http.ResponseWriter, id any, capName string) { + envelope := map[string]any{ + "jsonrpc": "2.0", + "error": map[string]any{ + "code": mcpparser.CodeMissingClientCapability, + "message": fmt.Sprintf( + "this tool requires the %q client capability mid-call, which the request did not declare; "+ + "note that declaring it will not help on protocol %s — serving mid-call capability "+ + "requests needs multi-round retrieval (SEP-2322), which this server does not implement", + capName, mcpparser.MCPVersionModern), + "data": map[string]any{ + "requiredCapabilities": map[string]any{capName: map[string]any{}}, + }, + }, + } + if transportsession.HasJSONRPCID(id) { + envelope["id"] = id + } + writeModernEnvelope(w, http.StatusOK, envelope) +} + +// writeModernDispatchError classifies a POST-dispatch error from +// CallTool/ReadResource/GetPrompt. Check* and the real call each re-aggregate +// independently (documented "aggregates twice" on CheckToolCall), so a +// concurrent backend health flip, cache refresh, or annotation change +// (TOCTOU) can have Check* allow and the call itself deny. That denial MUST +// still surface as 403 + denyMsg -- the same as the pre-dispatch gate -- so +// the audit middleware logs it as "denied" rather than "failure"; it is +// therefore tested FIRST, before falling through to the generic internal +// error. +// +// A domain error that carries its own stable JSON-RPC code and data +// (mcpparser.CodedError — e.g. the rate limiter's 429 with +// data.retryAfterSeconds) is written with that code rather than laundered +// into -32603. This is the Modern counterpart of the SDK path's +// conversion.ErrorToToolResult, whose CodedError branch preserves the same +// code/data in an IsError tool result's structuredContent because the SDK +// tool-handler seam cannot emit a custom JSON-RPC error object. This +// dispatcher owns the envelope, so it emits the real thing (mirroring how +// #6061 classified mid-call capability refusals as -32021 instead of +// -32603). Authorization denials are still tested first: a coded error can +// never mask a denial's 403. +// +// The -32603 message reuses err.Error() verbatim. This matches the SDK path's +// existing posture rather than inventing a new one: conversion.ErrorToToolResult's +// generic branch, and the resources/read/prompts/get Serve handlers +// (serve_handlers.go), already surface the raw error text for a non-coded, +// non-authz error. Re-sanitizing here would just diverge from what the SDK +// path already exposes for the identical failure. +func writeModernDispatchError(w http.ResponseWriter, id any, denyMsg string, err error) { + // Denial first (matches conversion.ErrorToToolResult): an error that is + // both a CodedError and wraps ErrAuthorizationFailed must render as the + // denial, never as retry-shaped coded data (the sets are disjoint today; + // this ordering is the invariant). + if errors.Is(err, vmcp.ErrAuthorizationFailed) { + writeModernDenied(w, id, denyMsg) + return + } + var coded mcpparser.CodedError + if errors.As(err, &coded) { + writeModernCodedError(w, id, err, coded) + return + } + writeModernError(w, id, jsonRPCCodeInternalError, err.Error()) +} diff --git a/pkg/vmcp/server/modern_dispatch_test.go b/pkg/vmcp/server/modern_dispatch_test.go new file mode 100644 index 0000000000..0fd5cd4f67 --- /dev/null +++ b/pkg/vmcp/server/modern_dispatch_test.go @@ -0,0 +1,982 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/ratelimit" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/core" +) + +// modernFakeCore is a core.VMCP whose only exercised methods are the ones +// dispatchModern can reach. The embedded nil interface satisfies the rest and +// panics if dispatch ever calls a method it should not (e.g. a list call for a +// notification/batch/unknown-method request that must short-circuit before +// reaching the core at all). +type modernFakeCore struct { + core.VMCP + + tools []vmcp.Tool + resources []vmcp.Resource + templates []vmcp.ResourceTemplate + prompts []vmcp.Prompt + + discoverCaps core.DiscoverCapabilities + discoverErr error + // discoverCalls counts Discover invocations, so a test can assert the + // subscriptions/listen pre-check skipped the backend fan-out entirely (the + // response is identical either way, so only the count is observable). + discoverCalls int + // discoverIdentity records the identity Discover was called with, so + // per-identity filtering of a computed set can be pinned rather than assumed. + discoverIdentity *auth.Identity + + checkToolErr, checkResourceErr, checkPromptErr error + callToolErr, readResourceErr, getPromptErr error + completeErr error + listErr error // returned by every List* method + + checkCalled, callCalled bool + + // backendID is stamped onto CallTool/ReadResource/GetPrompt's returned + // result, mirroring the real core's post-routing BackendID assignment + // (core_calls.go). Empty by default, matching a fake that never resolves + // a backend. + backendID string + + // gotCtx captures the context handed to CallTool, so a test can assert a + // value set on the inbound request survives into the core call unmodified. + gotCtx context.Context +} + +func (f *modernFakeCore) ListTools(context.Context, *auth.Identity) ([]vmcp.Tool, error) { + return f.tools, f.listErr +} + +func (f *modernFakeCore) ListResources(context.Context, *auth.Identity) ([]vmcp.Resource, error) { + return f.resources, nil +} + +func (f *modernFakeCore) ListResourceTemplates(context.Context, *auth.Identity) ([]vmcp.ResourceTemplate, error) { + return f.templates, nil +} + +func (f *modernFakeCore) ListPrompts(context.Context, *auth.Identity) ([]vmcp.Prompt, error) { + return f.prompts, nil +} + +func (f *modernFakeCore) Discover(_ context.Context, identity *auth.Identity) (core.DiscoverCapabilities, error) { + f.discoverCalls++ + f.discoverIdentity = identity + return f.discoverCaps, f.discoverErr +} + +func (f *modernFakeCore) CheckToolCall(context.Context, *auth.Identity, string, map[string]any) error { + f.checkCalled = true + return f.checkToolErr +} + +func (f *modernFakeCore) CheckResourceRead(context.Context, *auth.Identity, string) error { + f.checkCalled = true + return f.checkResourceErr +} + +func (f *modernFakeCore) CheckPromptGet(context.Context, *auth.Identity, string) error { + f.checkCalled = true + return f.checkPromptErr +} + +func (f *modernFakeCore) CallTool( + ctx context.Context, _ *auth.Identity, _ string, _ map[string]any, _ map[string]any, +) (*vmcp.ToolCallResult, error) { + f.callCalled = true + f.gotCtx = ctx + if f.callToolErr != nil { + return nil, f.callToolErr + } + return &vmcp.ToolCallResult{Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "ok"}}, BackendID: f.backendID}, nil +} + +func (f *modernFakeCore) ReadResource( + ctx context.Context, _ *auth.Identity, uri string, +) (*vmcp.ResourceReadResult, error) { + f.callCalled = true + f.gotCtx = ctx + if f.readResourceErr != nil { + return nil, f.readResourceErr + } + return &vmcp.ResourceReadResult{Contents: []vmcp.ResourceContent{{URI: uri, Text: "body"}}, BackendID: f.backendID}, nil +} + +func (f *modernFakeCore) GetPrompt( + ctx context.Context, _ *auth.Identity, _ string, _ map[string]any, +) (*vmcp.PromptGetResult, error) { + f.callCalled = true + f.gotCtx = ctx + if f.getPromptErr != nil { + return nil, f.getPromptErr + } + return &vmcp.PromptGetResult{ + Messages: []vmcp.PromptMessage{{Role: "user", Content: vmcp.Content{Type: vmcp.ContentTypeText, Text: "hi"}}}, + BackendID: f.backendID, + }, nil +} + +func (f *modernFakeCore) Complete( + ctx context.Context, _ *auth.Identity, _ vmcp.CompletionRef, _, _ string, _ map[string]string, +) (*vmcp.CompletionResult, error) { + f.callCalled = true + f.gotCtx = ctx + if f.completeErr != nil { + return nil, f.completeErr + } + return &vmcp.CompletionResult{Values: []string{"opt1", "opt2"}, Total: 2}, nil +} + +// dispatchModernTest builds a Server over fakeCore and drives dispatchModern +// directly (Step 3 wires this into the real handler chain), returning the +// decoded JSON-RPC envelope and the recorder for status/header assertions. +func dispatchModernTest( + reqCtx context.Context, t *testing.T, fakeCore *modernFakeCore, authzEnabled bool, parsed *mcpparser.ParsedMCPRequest, +) (*httptest.ResponseRecorder, map[string]any) { + t.Helper() + + s := &Server{ + config: &Config{Name: testServerName, Version: testServerVersion}, + core: fakeCore, + authzGateEnabled: authzEnabled, + } + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(reqCtx) + rec := httptest.NewRecorder() + + s.dispatchModern(rec, req, parsed) + + var body map[string]any + if rec.Body.Len() > 0 { + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + } + return rec, body +} + +// TestDispatchModern_ControlFlow covers the guards that must short-circuit +// before any core call: a batch gets -32600, and an unrecognized method gets +// -32601. None of these may reach the core (modernFakeCore's embedded nil +// core.VMCP panics if they did). The notification guard is NOT in this table +// -- see TestDispatchModern_NotificationRealParser: a hand-built +// ParsedMCPRequest{IsRequest:false} would not have caught the regression +// where the real parser always sets IsRequest true. +func TestDispatchModern_ControlFlow(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + parsed *mcpparser.ParsedMCPRequest + wantStatus int + wantCode float64 + wantIDKey bool + }{ + { + name: "batch request returns -32600 invalid request", + parsed: &mcpparser.ParsedMCPRequest{Method: "tools/call", IsRequest: true, IsBatch: true, ID: "1"}, + wantStatus: http.StatusOK, + wantCode: -32600, + wantIDKey: true, + }, + { + // A batch with a raw-null id exercises writeModernError's + // absent-id path through the real dispatch route, not just the + // unit-level writer tests in modern_envelope_test.go: + // schema/2025-11-25 types the error id as optional string|number, + // so MCP omits the "id" key entirely here, never "id":null. + // + // The real parser can never actually produce this state: it sets ID + // from jsonrpc2.ID.Raw(), which yields only nil, int64, or string + // (parser.go), and a plain Go nil id can't reach this branch either -- + // dispatchModern's earlier `parsed.ID == nil` notification guard (an + // interface nil check, :54) would short-circuit it to 202 first. + // json.RawMessage("null") is constructed directly here solely to drive + // the writer's defensive absent-id backstop through this route. + name: "batch request with raw-null id omits the id key", + parsed: &mcpparser.ParsedMCPRequest{Method: "tools/call", IsRequest: true, IsBatch: true, ID: json.RawMessage("null")}, + wantStatus: http.StatusOK, + wantCode: -32600, + wantIDKey: false, + }, + { + name: "unknown method returns -32601 method not found", + parsed: &mcpparser.ParsedMCPRequest{Method: "roots/list", IsRequest: true, ID: "1"}, + wantStatus: http.StatusNotFound, + wantCode: -32601, + wantIDKey: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec, body := dispatchModernTest(t.Context(), t, &modernFakeCore{}, true, tt.parsed) + + assert.Equal(t, tt.wantStatus, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope") + assert.Equal(t, tt.wantCode, errObj["code"]) + + _, hasID := body["id"] + assert.Equal(t, tt.wantIDKey, hasID, `"id" key presence`) + }) + } +} + +// TestDispatchModern_NotificationRealParser drives the REAL pkg/mcp parser +// (ParsingMiddleware, the same middleware installed in front of dispatchModern +// in production) over a genuine no-id JSON-RPC notification body, rather than +// a hand-built ParsedMCPRequest. This matters because parser.go's real parse +// path (parseMCPRequest) always sets IsRequest:true for a decoded +// jsonrpc2.Request -- calls AND notifications alike -- so `!parsed.IsRequest` +// never fires for a real notification; only an absent id does. A fabricated +// ParsedMCPRequest{IsRequest:false} would pass either check and silently mask +// that bug, which is exactly what happened here. +func TestDispatchModern_NotificationRealParser(t *testing.T) { + t.Parallel() + + body := []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{}}`) + req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + + var parsed *mcpparser.ParsedMCPRequest + handler := mcpparser.ParsingMiddleware(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + parsed = mcpparser.GetParsedMCPRequest(r.Context()) + })) + handler.ServeHTTP(httptest.NewRecorder(), req) + + require.NotNil(t, parsed, "the real parser must produce a ParsedMCPRequest for a notification body") + require.True(t, parsed.IsRequest, "the real parser sets IsRequest true even for a notification") + require.Nil(t, parsed.ID, "a notification has no JSON-RPC id") + + fc := &modernFakeCore{} + rec, _ := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusAccepted, rec.Code) + assert.Empty(t, rec.Body.Bytes()) + assert.False(t, fc.callCalled, "a notification must not reach the core") +} + +// TestDispatchModern_PingRealParser drives ping through the REAL +// ParsingMiddleware + classifyingHandler chain with genuine Modern signaling +// (MCP-Protocol-Version header, Mcp-Method header, body _meta) rather than a +// hand-built ParsedMCPRequest -- guarding against the ping case regressing +// back to the unrecognized-method default (-32601/404), which a fabricated +// request would not catch. ping is deliberately absent from +// nameRequiredMethods, so no Mcp-Name header is sent. +func TestDispatchModern_PingRealParser(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "jsonrpc": "2.0", + "id": 1, + "method": "ping", + "params": { + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {} + } + } + }`) + req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcpparser.MCPVersionModern) + req.Header.Set("Mcp-Method", "ping") + + s := classifyingHandlerTestServer() + nextCalled := false + next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { nextCalled = true }) + + rec := httptest.NewRecorder() + mcpparser.ParsingMiddleware(s.classifyingHandler(next)).ServeHTTP(rec, req) + + require.False(t, nextCalled, "a well-formed Modern ping must dispatch, not fall through to the SDK") + require.Equal(t, http.StatusOK, rec.Code, "must not be the 404 a mis-routed -32601 would produce") + require.JSONEq(t, `{"jsonrpc":"2.0","id":1,"result":{}}`, rec.Body.String()) + + fc, ok := s.core.(*modernFakeCore) + require.True(t, ok) + assert.False(t, fc.callCalled, "ping must not reach the core") +} + +// TestDispatchModern_MethodRouting spot-checks that each method routes to the +// matching core call and produces an envelope with the expected resultType +// and top-level result key -- deep envelope shape (cacheable, serverInfo, +// empty-collection marshalling) is already covered by Step 1's +// modern_envelope_test.go. ping is the one exception: it wants a bare {} +// result (see the case's doc comment in modern_dispatch.go), so it opts out +// of the resultType/wantKey assertions via wantBareResult. +func TestDispatchModern_MethodRouting(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + parsed *mcpparser.ParsedMCPRequest + fakeCore *modernFakeCore + wantKey string // key expected inside result + wantBareResult bool // ping: result == {}, no resultType/_meta + }{ + { + name: "tools/list", + parsed: &mcpparser.ParsedMCPRequest{Method: "tools/list", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{tools: []vmcp.Tool{{Name: "echo", InputSchema: map[string]any{"type": "object"}}}}, + wantKey: "tools", + }, + { + name: "resources/list", + parsed: &mcpparser.ParsedMCPRequest{Method: "resources/list", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{resources: []vmcp.Resource{{Name: "info", URI: "embedded:info"}}}, + wantKey: "resources", + }, + { + name: "resources/templates/list", + parsed: &mcpparser.ParsedMCPRequest{Method: "resources/templates/list", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{ + templates: []vmcp.ResourceTemplate{{Name: "logs", URITemplate: "file:///{date}.txt"}}, + }, + wantKey: "resourceTemplates", + }, + { + name: "prompts/list", + parsed: &mcpparser.ParsedMCPRequest{Method: "prompts/list", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{prompts: []vmcp.Prompt{{Name: "review"}}}, + wantKey: "prompts", + }, + { + name: "server/discover", + parsed: &mcpparser.ParsedMCPRequest{Method: "server/discover", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{tools: []vmcp.Tool{{Name: "echo", InputSchema: map[string]any{"type": "object"}}}}, + wantKey: "capabilities", + }, + { + name: "tools/call", + parsed: &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "echo", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{}, + wantKey: "content", + }, + { + name: "resources/read", + parsed: &mcpparser.ParsedMCPRequest{Method: "resources/read", ResourceID: "embedded:info", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{}, + wantKey: "contents", + }, + { + name: "prompts/get", + parsed: &mcpparser.ParsedMCPRequest{Method: "prompts/get", ResourceID: "review", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{}, + wantKey: "messages", + }, + { + name: "completion/complete", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "completion/complete", IsRequest: true, ID: "1", + Params: json.RawMessage(`{"ref":{"type":"ref/prompt","name":"review"},"argument":{"name":"style","value":"terse"}}`), + }, + fakeCore: &modernFakeCore{}, + wantKey: "completion", + }, + { + name: "ping", + parsed: &mcpparser.ParsedMCPRequest{Method: "ping", IsRequest: true, ID: "1"}, + fakeCore: &modernFakeCore{}, + wantBareResult: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec, body := dispatchModernTest(t.Context(), t, tt.fakeCore, false, tt.parsed) + + require.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC result envelope, got %v", body) + if tt.wantBareResult { + assert.Empty(t, result, "ping must return a bare {} result with no resultType/_meta") + assert.False(t, tt.fakeCore.callCalled, "ping must not reach the core") + return + } + assert.Equal(t, "complete", result["resultType"]) + assert.Contains(t, result, tt.wantKey) + }) + } +} + +// TestDispatchModern_ListError asserts a core List* error maps to -32603 +// (HTTP 200), same posture as the other error paths -- guards the list +// branch against silently swallowing an aggregation failure. +func TestDispatchModern_ListError(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "tools/list", IsRequest: true, ID: "1"} + fc := &modernFakeCore{listErr: errors.New("aggregation exploded")} + + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope, got %v", body) + assert.Equal(t, float64(jsonRPCCodeInternalError), errObj["code"]) +} + +// TestDispatchModern_Discover asserts server/discover's flags reflect exactly +// what the fake core admits: an empty admitted set advertises no capability +// (no descriptor arrays either way), and a populated set flags only the +// features that came back non-empty. +func TestDispatchModern_Discover(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "server/discover", IsRequest: true, ID: "1"} + + t.Run("nothing admitted -- only the static completions capability", func(t *testing.T) { + t.Parallel() + + rec, body := dispatchModernTest(t.Context(), t, &modernFakeCore{}, false, parsed) + + require.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + caps, ok := result["capabilities"].(map[string]any) + require.True(t, ok) + assert.Equal(t, map[string]any{"completions": map[string]any{}}, caps, + "completions is unconditional; no admitted list-backed feature must advertise nothing else") + }) + + t.Run("admitted view flags only the populated features", func(t *testing.T) { + t.Parallel() + + fc := &modernFakeCore{ + discoverCaps: core.DiscoverCapabilities{HasTools: true, HasResourceTemplates: true}, + } + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + require.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + caps, ok := result["capabilities"].(map[string]any) + require.True(t, ok) + assert.Contains(t, caps, "tools") + assert.Contains(t, caps, "resources", "a resource template alone must still set the resources flag") + assert.Contains(t, caps, "completions", "completions is unconditional") + assert.NotContains(t, caps, "prompts") + }) + + t.Run("a core Discover error maps to -32603", func(t *testing.T) { + t.Parallel() + + fc := &modernFakeCore{discoverErr: errors.New("aggregation exploded")} + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope, got %v", body) + assert.Equal(t, float64(jsonRPCCodeInternalError), errObj["code"]) + }) +} + +// TestDispatchModern_Complete covers the completion/complete dispatch path: +// successful routing to core.Complete with no pre-dispatch Check* gate (see +// call_gate.go's comment on why completion/complete is intentionally not +// wire-gated), ref-type-scoped deny-message reclassification on +// ErrAuthorizationFailed (mirroring the SDK path's coreCompletionHandler), a +// non-authz core error mapping to -32603, and malformed params rejected +// before the core is ever reached. +func TestDispatchModern_Complete(t *testing.T) { + t.Parallel() + + validPromptRef := json.RawMessage( + `{"ref":{"type":"ref/prompt","name":"review"},"argument":{"name":"style","value":"terse"}}`) + validResourceRef := json.RawMessage( + `{"ref":{"type":"ref/resource","uri":"file:///{date}.txt"},"argument":{"name":"date","value":"2026"}}`) + + t.Run("success returns the completion envelope with no Check* call", func(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "completion/complete", IsRequest: true, ID: "1", Params: validPromptRef} + fc := &modernFakeCore{} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + require.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Equal(t, "complete", result["resultType"]) + completion, ok := result["completion"].(map[string]any) + require.True(t, ok) + assert.Equal(t, []any{"opt1", "opt2"}, completion["values"]) + assert.True(t, fc.callCalled) + assert.False(t, fc.checkCalled, "completion/complete has no pre-dispatch Check* gate") + }) + + t.Run("ErrAuthorizationFailed on a prompt ref reclassifies to 403 with the prompt-get deny message", func(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "completion/complete", IsRequest: true, ID: "1", Params: validPromptRef} + fc := &modernFakeCore{completeErr: fmt.Errorf("%w: denied", vmcp.ErrAuthorizationFailed)} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusForbidden, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok) + assert.Equal(t, float64(mcpparser.JSONRPCCodeDenied), errObj["code"]) + assert.Equal(t, vmcp.DenyMessagePromptGet, errObj["message"]) + }) + + t.Run("ErrAuthorizationFailed on a resource ref reclassifies to 403 with the resource-read deny message", func(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "completion/complete", IsRequest: true, ID: "1", Params: validResourceRef} + fc := &modernFakeCore{completeErr: fmt.Errorf("%w: denied", vmcp.ErrAuthorizationFailed)} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusForbidden, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok) + assert.Equal(t, vmcp.DenyMessageResourceRead, errObj["message"]) + }) + + t.Run("a non-authz core error maps to -32603", func(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "completion/complete", IsRequest: true, ID: "1", Params: validPromptRef} + fc := &modernFakeCore{completeErr: errors.New("backend exploded")} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok) + assert.Equal(t, float64(jsonRPCCodeInternalError), errObj["code"]) + }) + + t.Run("malformed params rejected before the core is reached", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + params json.RawMessage + }{ + {name: "missing ref", params: json.RawMessage(`{"argument":{"name":"x","value":"y"}}`)}, + { + name: "ref missing type", + params: json.RawMessage(`{"ref":{"name":"review"},"argument":{"name":"x","value":"y"}}`), + }, + { + name: "ref not an object", + params: json.RawMessage(`{"ref":"review","argument":{"name":"x","value":"y"}}`), + }, + { + name: "missing argument.name", + params: json.RawMessage(`{"ref":{"type":"ref/prompt","name":"review"},"argument":{"value":"y"}}`), + }, + { + name: "empty argument.name", + params: json.RawMessage(`{"ref":{"type":"ref/prompt","name":"review"},"argument":{"name":"","value":"y"}}`), + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{ + Method: "completion/complete", IsRequest: true, ID: "1", Params: tt.params, + } + fc := &modernFakeCore{} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusBadRequest, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Equal(t, float64(jsonRPCCodeInvalidParams), errObj["code"]) + assert.False(t, fc.callCalled, "malformed params must not reach the core") + }) + } + }) +} + +// TestDispatchModern_AuthzGate is the security-load-bearing table: it drives +// the re-homed gate for each of the three gated methods (tools/call, +// resources/read, prompts/get) through every branch authzCallGate itself +// supports, PLUS the dispatch-time (TOCTOU) reclassification that only exists +// here because dispatchModern -- unlike the gate -- also owns the real call. +func TestDispatchModern_AuthzGate(t *testing.T) { + t.Parallel() + + deny := fmt.Errorf("%w: policy said no", vmcp.ErrAuthorizationFailed) + infra := errors.New("aggregation exploded") + + type gatedCase struct { + method string + resourceID string + wantKey string + denyMsg string + } + cases := []gatedCase{ + {method: "tools/call", resourceID: "secret-tool", wantKey: "content", denyMsg: vmcp.DenyMessageToolCall}, + {method: "resources/read", resourceID: "file://secret", wantKey: "contents", denyMsg: vmcp.DenyMessageResourceRead}, + {method: "prompts/get", resourceID: "secret-prompt", wantKey: "messages", denyMsg: vmcp.DenyMessagePromptGet}, + } + + for _, c := range cases { + parsed := &mcpparser.ParsedMCPRequest{Method: c.method, ResourceID: c.resourceID, IsRequest: true, ID: "1"} + + t.Run(c.method+"/denied pre-dispatch returns 403", func(t *testing.T) { + t.Parallel() + fc := &modernFakeCore{checkToolErr: deny, checkResourceErr: deny, checkPromptErr: deny} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusForbidden, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok) + assert.Equal(t, float64(mcpparser.JSONRPCCodeDenied), errObj["code"]) + assert.Equal(t, c.denyMsg, errObj["message"]) + assert.True(t, fc.checkCalled, "Check* must have been invoked") + assert.False(t, fc.callCalled, "the real call must NOT run once Check* denies") + }) + + t.Run(c.method+"/[HIGH] denied at dispatch time (TOCTOU) still returns 403 not -32603", func(t *testing.T) { + t.Parallel() + // Check* allows (nil), but the real call denies -- the aggregation + // re-run at dispatch time disagrees with the pre-flight check. + fc := &modernFakeCore{callToolErr: deny, readResourceErr: deny, getPromptErr: deny} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusForbidden, rec.Code, "a dispatch-time authz denial must audit as 403/denied") + errObj, ok := body["error"].(map[string]any) + require.True(t, ok) + assert.Equal(t, float64(mcpparser.JSONRPCCodeDenied), errObj["code"], + "must be the denial code, never -32603 -- that would audit as failure instead of denied") + assert.Equal(t, c.denyMsg, errObj["message"]) + assert.True(t, fc.checkCalled) + assert.True(t, fc.callCalled, "the real call must still run when Check* allows") + }) + + t.Run(c.method+"/infra error at Check* fails open, dispatch still happens", func(t *testing.T) { + t.Parallel() + fc := &modernFakeCore{checkToolErr: infra, checkResourceErr: infra, checkPromptErr: infra} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusOK, rec.Code, "an infra error at Check* must not become a 403") + result, ok := body["result"].(map[string]any) + require.True(t, ok, "dispatch must still have proceeded and returned a result, got %v", body) + assert.Contains(t, result, c.wantKey) + assert.True(t, fc.callCalled) + }) + + t.Run(c.method+"/non-authz error at dispatch time maps to -32603 not 403", func(t *testing.T) { + t.Parallel() + fc := &modernFakeCore{callToolErr: infra, readResourceErr: infra, getPromptErr: infra} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok) + assert.Equal(t, float64(jsonRPCCodeInternalError), errObj["code"]) + }) + + t.Run(c.method+"/authz disabled skips Check* and dispatches", func(t *testing.T) { + t.Parallel() + // Even a would-be-denying Check* must never run when the gate is off. + fc := &modernFakeCore{checkToolErr: deny, checkResourceErr: deny, checkPromptErr: deny} + + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Contains(t, result, c.wantKey) + assert.False(t, fc.checkCalled, "Check* must not run when authzGateEnabled is false") + assert.True(t, fc.callCalled) + }) + } +} + +// TestDispatchModern_CodedErrorPreservesCodeAndData pins the CodedError branch +// of writeModernDispatchError: a domain error carrying its own stable JSON-RPC +// code and data (here the rate limiter's 429 with data.retryAfterSeconds) +// must reach the Modern client with that code and data intact, not be +// laundered into an opaque -32603 — parity with the SDK path, where +// conversion.ErrorToToolResult preserves the same code/data in +// structuredContent. HTTP status stays 200: go-sdk rejects a non-200 POST +// response before decoding its body, so a 429 would discard the error object +// -- and data.retryAfterSeconds -- unread (see writeModernCodedError). +func TestDispatchModern_CodedErrorPreservesCodeAndData(t *testing.T) { + t.Parallel() + + limited := &ratelimit.RateLimitedError{RetryAfter: 5 * time.Second} + + cases := []struct { + method string + resourceID string + }{ + {method: "tools/call", resourceID: "echo"}, + {method: "resources/read", resourceID: "file://doc"}, + {method: "prompts/get", resourceID: "review"}, + } + + for _, c := range cases { + t.Run(c.method+"/coded error keeps its code and data", func(t *testing.T) { + t.Parallel() + parsed := &mcpparser.ParsedMCPRequest{Method: c.method, ResourceID: c.resourceID, IsRequest: true, ID: "1"} + fc := &modernFakeCore{callToolErr: limited, readResourceErr: limited, getPromptErr: limited} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Equal(t, float64(ratelimit.CodeRateLimited), errObj["code"]) + assert.Equal(t, ratelimit.MessageRateLimited, errObj["message"]) + data, ok := errObj["data"].(map[string]any) + require.True(t, ok, "coded error data must be preserved, got %v", errObj) + assert.Equal(t, float64(5), data["retryAfterSeconds"]) + }) + } + + t.Run("wrapping a coded error keeps the code and the wrapped message", func(t *testing.T) { + t.Parallel() + parsed := &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "echo", IsRequest: true, ID: "1"} + fc := &modernFakeCore{callToolErr: fmt.Errorf("calling backend: %w", limited)} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Equal(t, float64(ratelimit.CodeRateLimited), errObj["code"], + "errors.As must unwrap to the coded error") + assert.Equal(t, "calling backend: "+ratelimit.MessageRateLimited, errObj["message"], + "message must keep the caller's wrap context") + }) + + t.Run("an authz denial wrapping never yields the coded branch", func(t *testing.T) { + t.Parallel() + // Belt-and-suspenders for the documented ordering: a denial keeps its + // 403 even if some future error value were both a denial and coded. + parsed := &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "echo", IsRequest: true, ID: "1"} + fc := &modernFakeCore{callToolErr: fmt.Errorf("%w: no", vmcp.ErrAuthorizationFailed)} + + rec, _ := dispatchModernTest(t.Context(), t, fc, true, parsed) + assert.Equal(t, http.StatusForbidden, rec.Code) + }) +} + +// TestDispatchModern_ArgumentsShape asserts the raw-params shape guard on +// tools/call and prompts/get: the parser (handleNamedResourceMethod) +// silently coerces a present-but-non-object "arguments" value to nil -- +// indistinguishable from "absent" -- so the guard must inspect parsed.Params +// directly, and must run BEFORE the authz gate and the real call. +// resources/read has no arguments and is intentionally not covered here. +func TestDispatchModern_ArgumentsShape(t *testing.T) { + t.Parallel() + + cases := []struct { + method string + resourceID string + wantKey string + }{ + {method: "tools/call", resourceID: "echo", wantKey: "content"}, + {method: "prompts/get", resourceID: "review", wantKey: "messages"}, + } + + for _, c := range cases { + t.Run(c.method+"/non-object arguments rejected before Check* or the call", func(t *testing.T) { + t.Parallel() + for _, raw := range []string{`"a string"`, `[1,2,3]`} { + parsed := &mcpparser.ParsedMCPRequest{ + Method: c.method, ResourceID: c.resourceID, IsRequest: true, ID: "1", + Params: json.RawMessage(`{"name":"` + c.resourceID + `","arguments":` + raw + `}`), + } + fc := &modernFakeCore{} + + rec, body := dispatchModernTest(t.Context(), t, fc, true, parsed) + + assert.Equal(t, http.StatusBadRequest, rec.Code, "-32602 invalid params maps to HTTP 400") + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope for arguments=%s, got %v", raw, body) + assert.Equal(t, float64(jsonRPCCodeInvalidParams), errObj["code"]) + assert.False(t, fc.checkCalled, "Check* must not run for invalid arguments shape") + assert.False(t, fc.callCalled, "the real call must not run for invalid arguments shape") + } + }) + + t.Run(c.method+"/absent arguments proceeds", func(t *testing.T) { + t.Parallel() + parsed := &mcpparser.ParsedMCPRequest{ + Method: c.method, ResourceID: c.resourceID, IsRequest: true, ID: "1", + Params: json.RawMessage(`{"name":"` + c.resourceID + `"}`), + } + fc := &modernFakeCore{} + + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Contains(t, result, c.wantKey) + assert.True(t, fc.callCalled) + }) + + t.Run(c.method+"/object arguments proceeds", func(t *testing.T) { + t.Parallel() + parsed := &mcpparser.ParsedMCPRequest{ + Method: c.method, ResourceID: c.resourceID, IsRequest: true, ID: "1", + Params: json.RawMessage(`{"name":"` + c.resourceID + `","arguments":{"x":1}}`), + } + fc := &modernFakeCore{} + + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Contains(t, result, c.wantKey) + assert.True(t, fc.callCalled) + }) + + t.Run(c.method+"/explicit null arguments proceeds", func(t *testing.T) { + t.Parallel() + // hasNonObjectArguments relies on json.Unmarshal("null", &raw.Arguments) + // leaving raw.Arguments nil with no error -- an explicit JSON null is a + // legitimate no-args call, same as an absent field, not a rejected shape. + parsed := &mcpparser.ParsedMCPRequest{ + Method: c.method, ResourceID: c.resourceID, IsRequest: true, ID: "1", + Params: json.RawMessage(`{"name":"` + c.resourceID + `","arguments":null}`), + } + fc := &modernFakeCore{} + + rec, body := dispatchModernTest(t.Context(), t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "got %v", body) + assert.Contains(t, result, c.wantKey) + assert.True(t, fc.callCalled) + }) + } +} + +// contextProbeKey is a private key used only to prove r.Context() reaches the +// core call unmodified (not detached, not re-wrapped): headerforward reads +// forwarded-header state off this exact context per call, so a detach would +// silently break backend auth on every Modern request. +type contextProbeKey struct{} + +func TestDispatchModern_ContextPassthrough(t *testing.T) { + t.Parallel() + + ctx := context.WithValue(t.Context(), contextProbeKey{}, "probe-value") + parsed := &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "echo", IsRequest: true, ID: "1"} + fc := &modernFakeCore{} + + _, _ = dispatchModernTest(ctx, t, fc, false, parsed) + + require.NotNil(t, fc.gotCtx) + assert.Equal(t, "probe-value", fc.gotCtx.Value(contextProbeKey{}), + "dispatchModern must pass r.Context() through to core.CallTool unmodified") +} + +// TestDispatchModern_LabelsAuditBackend mirrors TestServeHandlersLabelAuditBackend +// (serve_session_test.go) for the Modern path: after a successful tools/call, +// resources/read, or prompts/get, the dispatcher must write the registry-resolved +// backend name into the audit BackendInfo carried in the request context, using the +// same s.backendDisplayName resolution the Serve path's handlers use. completion/complete +// is deliberately excluded -- see dispatchModernComplete's doc comment on why that gap +// is pre-existing on both paths, not something this fix should introduce for Modern only. +func TestDispatchModern_LabelsAuditBackend(t *testing.T) { + t.Parallel() + + // BackendID "backend-x" != Name "github-mcp": a pass-through would record the ID. + reg := vmcp.NewImmutableRegistry([]vmcp.Backend{{ID: "backend-x", Name: "github-mcp"}}) + + tests := []struct { + name string + parsed *mcpparser.ParsedMCPRequest + }{ + {name: "tools/call", parsed: &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "echo", IsRequest: true, ID: "1"}}, + { + name: "resources/read", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "resources/read", ResourceID: "file://a", IsRequest: true, ID: "1", + }, + }, + { + name: "prompts/get", + parsed: &mcpparser.ParsedMCPRequest{ + Method: "prompts/get", ResourceID: "review", IsRequest: true, ID: "1", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + s := &Server{ + config: &Config{Name: testServerName, Version: testServerVersion}, + core: &modernFakeCore{backendID: "backend-x"}, + backendRegistry: reg, + } + + bi := &audit.BackendInfo{} + ctx := audit.WithBackendInfo(t.Context(), bi) + req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(ctx) + rec := httptest.NewRecorder() + + s.dispatchModern(rec, req, tt.parsed) + + require.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "github-mcp", bi.BackendName, + "the Modern dispatcher must label the audit event with the registry-resolved backend name") + }) + } +} + +// TestDispatchModern_NoBackendIDSkipsAuditLabel locks in the composite-tool case: a +// result with an empty BackendID (executeComposite never sets one, core_calls.go) must +// not touch the audit BackendInfo at all, matching the "no single serving backend" +// semantics documented on vmcp.ToolCallResult.BackendID. +func TestDispatchModern_NoBackendIDSkipsAuditLabel(t *testing.T) { + t.Parallel() + + parsed := &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "echo", IsRequest: true, ID: "1"} + fc := &modernFakeCore{} // backendID left empty + + bi := &audit.BackendInfo{} + ctx := audit.WithBackendInfo(t.Context(), bi) + rec, _ := dispatchModernTest(ctx, t, fc, false, parsed) + + assert.Equal(t, http.StatusOK, rec.Code) + assert.Empty(t, bi.BackendName, "an empty BackendID must leave the audit label untouched") +} diff --git a/pkg/vmcp/server/modern_envelope.go b/pkg/vmcp/server/modern_envelope.go new file mode 100644 index 0000000000..0da4e60260 --- /dev/null +++ b/pkg/vmcp/server/modern_envelope.go @@ -0,0 +1,676 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "encoding/json" + "fmt" + "log/slog" + "net/http" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/conversion" +) + +// This file hand-rolls the MCP 2026-07-28 ("Modern") stateless response +// envelope: the wire shape go-sdk@v1.7.0-pre.3 produces for a single-shot +// tools/list, resources/list, resources/templates/list, prompts/list, +// tools/call, resources/read, and prompts/get. Even in that version -- the one +// mcpcompat (ToolHive's go-sdk import path) now depends on -- these Modern +// types remain unexported, so this is a durable parallel serializer, not a +// stopgap: resultType and _meta.serverInfo are set by unexported SDK functions +// (setCompleteResultType, annotateServerInfo) that run inside the exact +// ServerSession dispatch this package bypasses for Modern stateless requests. +// A future go-sdk bump cannot just delete these structs and marshal SDK result +// types directly -- the Modern annotations would vanish with them. +// +// modernResultTypeComplete is the sole value dispatchModern ever needs: this +// package only performs single-shot dispatch, never the elicitation retry +// loop (resultType "input_required"), so every result built here is +// unconditionally "complete". Do not add Legacy-conditional branching to +// these structs -- this file is reached only for Modern requests. +const modernResultTypeComplete = "complete" + +// modernServerInfoKey is the go-sdk's MetaKeyServerInfo +// (protocol.go:2367 in go-sdk@v1.7.0-pre.3), reproduced by hand since +// mcpcompat/mcp (ToolHive's mcp import) does not re-export it, even though +// go-sdk itself does. +const modernServerInfoKey = "io.modelcontextprotocol/serverInfo" + +// modernServerInfo mirrors the go-sdk's Implementation type. +type modernServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// modernMeta carries the _meta.serverInfo entry the SDK attaches to every +// Modern result, unconditionally and together with resultType. +type modernMeta struct { + ServerInfo modernServerInfo `json:"io.modelcontextprotocol/serverInfo"` +} + +func newModernMeta(serverName, serverVersion string) modernMeta { + return modernMeta{ServerInfo: modernServerInfo{Name: serverName, Version: serverVersion}} +} + +// newModernResultMeta builds _meta for tools/call, resources/read, and +// prompts/get, the three builders whose domain result carries its own +// backend Meta (vmcp.ToolCallResult.Meta etc.) -- unlike the four list +// results and server/discover, which have no per-result backend meta and use +// the serverInfo-only newModernMeta above. The SDK path preserves backend +// meta via conversion.ToMCPMeta (serve_handlers.go); dropping it here would +// silently discard whatever the backend attached (progress tokens, trace +// ids, ...). +// +// Every reserved io.modelcontextprotocol/* key a backend set is stripped first, +// then vMCP's own serverInfo is stamped last so it always wins: the client's +// actual MCP peer is vMCP, not the backend, so no backend may speak for it. +// mcpparser.StripReservedMeta is the SAME helper the Legacy path reaches through +// conversion.ToMCPMeta -- keep it that way, or the two revisions drift (which is +// exactly what this function's previous forward-everything behavior caused). +// It clones, so the domain result's map is never mutated, and returns nil +// whenever nothing survives -- hence the fallback below, since this result +// always carries at least serverInfo. +func newModernResultMeta(backendMeta map[string]any, serverName, serverVersion string) map[string]any { + meta := mcpparser.StripReservedMeta(backendMeta) + if meta == nil { + meta = make(map[string]any, 1) + } + meta[modernServerInfoKey] = modernServerInfo{Name: serverName, Version: serverVersion} + return meta +} + +// modernCacheable mirrors the go-sdk's Cacheable struct (protocol.go:1168), +// with no omitempty: both fields are always present on the wire. +// +// CacheScope is deliberately hardcoded to "private", diverging from the SDK's +// "public" default. The four list results and resources/read vary by caller +// identity (admission.FilterTools/Resources/Prompts, AllowResourceRead), so +// advertising "public" would let a shared cache/intermediary serve one +// identity's admission-filtered view to another. TTLMs is 0 ("immediately +// stale") as an interim value for a re-aggregating gateway with no cache +// invalidation signal yet. +type modernCacheable struct { + TTLMs int `json:"ttlMs"` + CacheScope string `json:"cacheScope"` +} + +func newModernCacheable() modernCacheable { + return modernCacheable{TTLMs: 0, CacheScope: "private"} +} + +// modernToolsListResult is the tools/list wire result. +type modernToolsListResult struct { + ResultType string `json:"resultType"` + modernCacheable + Tools []mcp.Tool `json:"tools"` + NextCursor string `json:"nextCursor,omitempty"` + Meta modernMeta `json:"_meta"` +} + +// modernResourcesListResult is the resources/list wire result. +type modernResourcesListResult struct { + ResultType string `json:"resultType"` + modernCacheable + Resources []mcp.Resource `json:"resources"` + NextCursor string `json:"nextCursor,omitempty"` + Meta modernMeta `json:"_meta"` +} + +// modernResourceTemplatesListResult is the resources/templates/list wire result. +type modernResourceTemplatesListResult struct { + ResultType string `json:"resultType"` + modernCacheable + ResourceTemplates []mcp.ResourceTemplate `json:"resourceTemplates"` + NextCursor string `json:"nextCursor,omitempty"` + Meta modernMeta `json:"_meta"` +} + +// modernPromptsListResult is the prompts/list wire result. +type modernPromptsListResult struct { + ResultType string `json:"resultType"` + modernCacheable + Prompts []mcp.Prompt `json:"prompts"` + NextCursor string `json:"nextCursor,omitempty"` + Meta modernMeta `json:"_meta"` +} + +// modernCallToolResult is the tools/call wire result. Unlike the four list +// results and resources/read, it does NOT embed modernCacheable -- a tool +// call is an action, not a cacheable read, and the SDK never attaches +// Cacheable to CallToolResult. Meta is a map (not modernMeta) because it must +// carry both the backend's own result meta AND serverInfo -- see +// newModernResultMeta. +type modernCallToolResult struct { + ResultType string `json:"resultType"` + Content []mcp.Content `json:"content"` + StructuredContent any `json:"structuredContent,omitempty"` + IsError bool `json:"isError,omitempty"` + Meta map[string]any `json:"_meta"` +} + +// modernReadResourceResult is the resources/read wire result. Meta is a map +// for the same reason as modernCallToolResult -- see newModernResultMeta. +type modernReadResourceResult struct { + ResultType string `json:"resultType"` + modernCacheable + Contents []mcp.ResourceContents `json:"contents"` + Meta map[string]any `json:"_meta"` +} + +// modernGetPromptResult is the prompts/get wire result. Like tools/call, it +// does NOT embed modernCacheable -- the SDK never attaches Cacheable to +// GetPromptResult. Meta is a map for the same reason as modernCallToolResult +// -- see newModernResultMeta. +type modernGetPromptResult struct { + ResultType string `json:"resultType"` + Description string `json:"description,omitempty"` + Messages []mcp.PromptMessage `json:"messages"` + Meta map[string]any `json:"_meta"` +} + +// modernCompletionDetails mirrors the go-sdk's CompletionResultDetails +// (protocol.go:653 in go-sdk@v1.7.0-pre.3): Values has no omitempty tag +// there, so it doesn't here either -- an empty result still marshals +// "values":[] rather than omitting the field. +type modernCompletionDetails struct { + Values []string `json:"values"` + Total int `json:"total,omitempty"` + HasMore bool `json:"hasMore,omitempty"` +} + +// modernCompleteResult is the completion/complete wire result. Like +// tools/call and prompts/get, it does NOT embed modernCacheable -- the SDK +// never attaches Cacheable to CompleteResult. Meta is modernMeta (not a map) +// because vmcp.CompletionResult carries no backend meta to preserve -- unlike +// newModernResultMeta's callers, there is nothing to clone here. +type modernCompleteResult struct { + ResultType string `json:"resultType"` + Completion modernCompletionDetails `json:"completion"` + Meta modernMeta `json:"_meta"` +} + +// newModernComplete builds the completion/complete wire result from the +// core's CompletionResult. Values is initialized to a non-nil slice so an +// empty result marshals as [] not null, matching the four list builders. +// +// result is expected non-nil (core.Complete never returns a nil result +// without an error); the nil check below is defense-in-depth against a +// plausible aggregator bug, not an expected input. +func newModernComplete(result *vmcp.CompletionResult, serverName, serverVersion string) modernCompleteResult { + if result == nil { + result = &vmcp.CompletionResult{} + } + values := result.Values + if values == nil { + values = []string{} + } + return modernCompleteResult{ + ResultType: modernResultTypeComplete, + Completion: modernCompletionDetails{ + Values: values, + Total: result.Total, + HasMore: result.HasMore, + }, + Meta: newModernMeta(serverName, serverVersion), + } +} + +// modernDiscoverResult is the server/discover wire result -- Modern's +// replacement for initialize+capability negotiation. Unlike the four list +// results, it carries no descriptor arrays: Capabilities reuses mcpcompat's +// mcp.ServerCapabilities, whose per-feature fields are themselves the "flag" +// (a pointer field present iff the identity can reach that feature at all). +type modernDiscoverResult struct { + ResultType string `json:"resultType"` + modernCacheable + SupportedVersions []string `json:"supportedVersions"` + Capabilities mcp.ServerCapabilities `json:"capabilities"` + Meta modernMeta `json:"_meta"` +} + +// newModernDiscover builds the server/discover wire result. The capability +// flags themselves are shaped by newModernCapabilities below -- see its doc +// comment for what each flag means and why the push-related ones are false. +// This function only assembles the envelope around them. +// +// SupportedVersions enumerates every protocol version this vMCP endpoint +// actually serves, not just Modern: mcpparser.MCPVersionModern (this +// dispatcher) plus mcp.LATEST_PROTOCOL_VERSION, the one Legacy version the +// SDK path negotiates (mcpcompat's handleInitialize always responds with +// LATEST_PROTOCOL_VERSION regardless of what a client requests -- it does +// not support older Legacy revisions either). A discover-first client uses +// this list to decide whether it can skip the Legacy initialize handshake. +// +// CONSTRAINT: this function must only be reached with the Modern capability +// gate open. It hardcodes MCPVersionModern into SupportedVersions and never +// consults the gate; that is truthful only because classifyingHandler +// (classification.go) routes server/discover to the SDK when the gate is +// closed, and the SDK's stateful transport computes the truthful Legacy-only +// list there. Do NOT plumb the gate in here — that would create a second +// source of the version list. Anyone re-routing gated discover to +// dispatchModern must revisit this. +func newModernDiscover( + hasTools, hasResources, hasTemplates, hasPrompts bool, serverName, serverVersion string, +) modernDiscoverResult { + return modernDiscoverResult{ + ResultType: modernResultTypeComplete, + modernCacheable: newModernCacheable(), + SupportedVersions: []string{mcpparser.MCPVersionModern, mcp.LATEST_PROTOCOL_VERSION}, + Capabilities: newModernCapabilities(hasTools, hasResources, hasTemplates, hasPrompts), + Meta: newModernMeta(serverName, serverVersion), + } +} + +// newModernCapabilities shapes the four admission-filtered presence checks into +// the wire capability flags. +// +// hasTools/hasResources/hasTemplates/hasPrompts are the caller's +// len(core.List*(ctx, identity))>0 checks, so an advertisement reflects +// exclusively what this identity may reach, same as the four list results. +// hasTemplates folds into the "resources" capability: the wire protocol has no +// separate resource-templates flag. +// +// Completions is set unconditionally, unlike the identity-filtered flags: +// completion/complete has no per-identity admission of its own (core.Complete +// authorizes the underlying prompt/resource ref at dispatch, not the completions +// feature itself), and this dispatcher serves it for every caller -- matching how +// the Legacy/SDK path always wires server.WithCompletionHandler regardless of +// identity (serve.go). +// +// It is the SINGLE source of truth for what this dispatcher advertises: server/discover publishes it verbatim (above), and +// subscriptions/listen intersects a client's requested notification set against +// it (honoredSubscriptions, modern_subscriptions.go). Keeping one builder is +// what makes those two answers impossible to drift apart -- a client can never +// be told a push capability exists by one verb and denied it by the other. +// +// Every push-related flag below is deliberately false, and each `false` is a +// load-bearing statement about what vMCP can actually do, not an oversight: +// +// - Tools/Prompts/Resources ListChanged: dispatchModern creates no session, so +// the per-session list_changed machinery (buildListChangedSink, wired only +// from handleSessionRegistrationImpl in server.go) never runs for a Modern +// client, whatever revision its backends speak. +// - Resources Subscribe: this stateless single-shot dispatcher has no +// persistent connection to a client to push a server-initiated +// resources/updated notification over, so resources/subscribe is not +// advertised here and returns -32601 by design, not by oversight. vMCP does +// not forward backend resources/updated on the Legacy path either +// (ack-level only; see docs/arch/10-virtual-mcp-architecture.md). +// +// Flipping any of them to true is a promise to deliver that notification type, +// which today nothing in vMCP can keep. Implement delivery in the SAME change, +// or a Modern client will subscribe successfully and then wait forever. See +// #5743. +func newModernCapabilities(hasTools, hasResources, hasTemplates, hasPrompts bool) mcp.ServerCapabilities { + var caps mcp.ServerCapabilities + if hasTools { + caps.Tools = &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{} + } + if hasResources || hasTemplates { + caps.Resources = &struct { + Subscribe bool `json:"subscribe,omitempty"` + ListChanged bool `json:"listChanged,omitempty"` + }{} + } + if hasPrompts { + caps.Prompts = &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{} + } + caps.Completions = &struct{}{} + return caps +} + +// newModernToolsList builds the tools/list wire result from the core's +// admission-filtered domain tools. +// nextCursor is LAST on purpose: it sits among same-typed string parameters, and +// when it was in the middle a transposition with serverName compiled silently and +// put a server name in the cursor field. Keep new string params appended, not +// inserted. +func newModernToolsList( + tools []vmcp.Tool, serverName, serverVersion, nextCursor string, +) (modernToolsListResult, error) { + wireTools := make([]mcp.Tool, 0, len(tools)) + for _, t := range tools { + wireTool, err := modernToolFromDomain(t) + if err != nil { + return modernToolsListResult{}, err + } + wireTools = append(wireTools, wireTool) + } + return modernToolsListResult{ + ResultType: modernResultTypeComplete, + modernCacheable: newModernCacheable(), + Tools: wireTools, + NextCursor: nextCursor, + Meta: newModernMeta(serverName, serverVersion), + }, nil +} + +// newModernResourcesList builds the resources/list wire result from the +// core's admission-filtered domain resources. +func newModernResourcesList( + resources []vmcp.Resource, serverName, serverVersion, nextCursor string, +) modernResourcesListResult { + wireResources := make([]mcp.Resource, 0, len(resources)) + for _, r := range resources { + wireResources = append(wireResources, modernResourceFromDomain(r)) + } + return modernResourcesListResult{ + ResultType: modernResultTypeComplete, + modernCacheable: newModernCacheable(), + Resources: wireResources, + NextCursor: nextCursor, + Meta: newModernMeta(serverName, serverVersion), + } +} + +// newModernResourceTemplatesList builds the resources/templates/list wire +// result from the core's admission-filtered domain resource templates. +func newModernResourceTemplatesList( + templates []vmcp.ResourceTemplate, serverName, serverVersion, nextCursor string, +) modernResourceTemplatesListResult { + wireTemplates := make([]mcp.ResourceTemplate, 0, len(templates)) + for _, t := range templates { + wireTemplates = append(wireTemplates, modernResourceTemplateFromDomain(t)) + } + return modernResourceTemplatesListResult{ + ResultType: modernResultTypeComplete, + modernCacheable: newModernCacheable(), + ResourceTemplates: wireTemplates, + NextCursor: nextCursor, + Meta: newModernMeta(serverName, serverVersion), + } +} + +// newModernPromptsList builds the prompts/list wire result from the core's +// admission-filtered domain prompts. +func newModernPromptsList( + prompts []vmcp.Prompt, serverName, serverVersion, nextCursor string, +) modernPromptsListResult { + wirePrompts := make([]mcp.Prompt, 0, len(prompts)) + for _, p := range prompts { + wirePrompts = append(wirePrompts, modernPromptFromDomain(p)) + } + return modernPromptsListResult{ + ResultType: modernResultTypeComplete, + modernCacheable: newModernCacheable(), + Prompts: wirePrompts, + NextCursor: nextCursor, + Meta: newModernMeta(serverName, serverVersion), + } +} + +// newModernCallToolResult builds the tools/call wire result from the core's +// ToolCallResult. StructuredContent is omitted entirely (not merely +// omitempty-false) when the core did not set it, matching the SDK's +// omitempty behavior. +// +// result is non-nil on every call: dispatchModernToolCall (modern_dispatch.go) +// dereferences result.BackendID before calling this builder, so a nil result +// would already have panicked upstream. +func newModernCallToolResult(result *vmcp.ToolCallResult, serverName, serverVersion string) modernCallToolResult { + var structuredContent any + if len(result.StructuredContent) > 0 { + structuredContent = result.StructuredContent + } + return modernCallToolResult{ + ResultType: modernResultTypeComplete, + Content: conversion.ToMCPContents(result.Content), + StructuredContent: structuredContent, + IsError: result.IsError, + Meta: newModernResultMeta(result.Meta, serverName, serverVersion), + } +} + +// newModernReadResourceResult builds the resources/read wire result from the +// core's ResourceReadResult. +// +// result is non-nil on every call: dispatchModernResourceRead +// (modern_dispatch.go) dereferences result.BackendID before calling this +// builder, so a nil result would already have panicked upstream. +func newModernReadResourceResult( + result *vmcp.ResourceReadResult, serverName, serverVersion string, +) modernReadResourceResult { + return modernReadResourceResult{ + ResultType: modernResultTypeComplete, + modernCacheable: newModernCacheable(), + Contents: conversion.ToMCPResourceContents(result.Contents), + Meta: newModernResultMeta(result.Meta, serverName, serverVersion), + } +} + +// newModernGetPromptResult builds the prompts/get wire result from the +// core's PromptGetResult. +// +// result is non-nil on every call: dispatchModernPromptGet +// (modern_dispatch.go) dereferences result.BackendID before calling this +// builder, so a nil result would already have panicked upstream. +func newModernGetPromptResult(result *vmcp.PromptGetResult, serverName, serverVersion string) modernGetPromptResult { + return modernGetPromptResult{ + ResultType: modernResultTypeComplete, + Description: result.Description, + Messages: conversion.ToMCPPromptMessages(result.Messages), + Meta: newModernResultMeta(result.Meta, serverName, serverVersion), + } +} + +// modernToolFromDomain converts a vmcp.Tool to the wire mcp.Tool used in a +// tools/list result. Mirrors coreSessionTools' Legacy adaptation +// (serve_handlers.go) field-for-field, so Legacy and Modern report identical +// tool shapes -- do not reinvent this mapping if the domain type changes. +func modernToolFromDomain(t vmcp.Tool) (mcp.Tool, error) { + schemaJSON, err := json.Marshal(t.InputSchema) + if err != nil { + return mcp.Tool{}, fmt.Errorf("marshal input schema for tool %s: %w", t.Name, err) + } + wireTool := mcp.Tool{ + Name: t.Name, + Description: t.Description, + RawInputSchema: schemaJSON, + // A tool with no annotations still marshals "annotations":{} + // because mcpcompat's Tool.MarshalJSON (mcpcompat/mcp/tools.go:343) + // writes the field unconditionally. Shared with Legacy's + // coreSessionTools, so the real fix belongs in mcpcompat, not here -- + // fixing it only in this file would break Legacy/Modern parity. + Annotations: conversion.ToMCPToolAnnotations(t.Annotations), + } + // Unlike the required InputSchema above, OutputSchema is best-effort: on + // failure the tool is still advertised without it (matches + // coreSessionTools). + if t.OutputSchema != nil { + if outputSchemaJSON, marshalErr := json.Marshal(t.OutputSchema); marshalErr != nil { + slog.Warn("failed to marshal tool output schema", "tool", t.Name, "error", marshalErr) + } else { + wireTool.RawOutputSchema = outputSchemaJSON + } + } + return wireTool, nil +} + +// modernResourceFromDomain converts a vmcp.Resource to the wire mcp.Resource +// used in a resources/list result. +func modernResourceFromDomain(r vmcp.Resource) mcp.Resource { + return mcp.Resource{ + Name: r.Name, + URI: r.URI, + Description: r.Description, + MIMEType: r.MimeType, + } +} + +// modernResourceTemplateFromDomain converts a vmcp.ResourceTemplate to the +// wire mcp.ResourceTemplate used in a resources/templates/list result. +func modernResourceTemplateFromDomain(t vmcp.ResourceTemplate) mcp.ResourceTemplate { + return mcp.ResourceTemplate{ + Name: t.Name, + URITemplate: t.URITemplate, + Description: t.Description, + MIMEType: t.MimeType, + } +} + +// modernPromptFromDomain converts a vmcp.Prompt to the wire mcp.Prompt used +// in a prompts/list result. +func modernPromptFromDomain(p vmcp.Prompt) mcp.Prompt { + arguments := make([]mcp.PromptArgument, 0, len(p.Arguments)) + for _, arg := range p.Arguments { + arguments = append(arguments, mcp.PromptArgument{ + Name: arg.Name, + Description: arg.Description, + Required: arg.Required, + }) + } + return mcp.Prompt{ + Name: p.Name, + Description: p.Description, + Arguments: arguments, + } +} + +// writeModernResult writes the JSON-RPC success envelope +// {"jsonrpc":"2.0","id":,"result":} as a single HTTP 200 +// application/json response -- never SSE, no Mcp-Session-Id -- matching the +// Modern stateless wire contract. +// +// Unlike writeModernError/writeModernDenied below, id is ALWAYS emitted here, +// never conditionally omitted: MCP's JSONRPCResultResponse schema keeps id +// required (only the error response narrows it to optional), and dispatchModern +// never reaches this builder for a notification -- it returns 202 before +// dispatch when parsed.ID is nil (modern_dispatch.go). Making this conditional +// would be as wrong as emitting "id":null on an error response. +func writeModernResult(w http.ResponseWriter, id, result any) { + writeModernEnvelope(w, http.StatusOK, map[string]any{ + "jsonrpc": "2.0", + "id": id, + "result": result, + }) +} + +// writeModernError writes a JSON-RPC error envelope, deriving the HTTP +// status from the JSON-RPC code -- mirroring go-sdk's extractErrorStatus +// (streamable.go:1033-1061): +// - -32601 (method not found) -> 404: the 2026-07-28 spec MUSTs 404 for an +// unimplemented method. +// - -32602 (invalid params) -> 400: matches the SDK's own mapping. +// - everything else (-32600 batch, -32603 internal) -> 200: matches +// extractErrorStatus returning 0 for those codes -- the request was +// accepted and processed, so the failure is an application-level +// JSON-RPC error riding the transport, not a wire-level rejection. +// +// This is a protocol-level mapping, unrelated to authorization: only +// writeModernDenied changes status for a POLICY reason (403). +// +// MCP narrows base JSON-RPC 2.0 here: schema/2025-11-25 types the error +// response as `id?: RequestId` where RequestId = string | number, so an +// absent id is encoded by omitting the "id" key, never as null. See +// transportsession.HasJSONRPCID. +func writeModernError(w http.ResponseWriter, id any, code int, msg string) { + status := http.StatusOK + switch code { + case jsonRPCCodeMethodNotFound: + status = http.StatusNotFound + case jsonRPCCodeInvalidParams: + status = http.StatusBadRequest + } + envelope := map[string]any{ + "jsonrpc": "2.0", + "error": map[string]any{ + "code": code, + "message": msg, + }, + } + if transportsession.HasJSONRPCID(id) { + envelope["id"] = id + } + writeModernEnvelope(w, status, envelope) +} + +// writeModernCodedError writes a JSON-RPC error envelope for a domain error +// carrying its own stable code and data (mcpparser.CodedError), preserving +// both on the wire: {"code":coded.Code(),"message":err.Error(),"data":...}. +// message uses err.Error() rather than a canned string so callers can wrap a +// coded error with context via %w without losing it (mirrors +// conversion.CodedErrorResult, the Legacy seam's rendering of the same +// errors). +// +// Status is HTTP 200 deliberately: go-sdk's streamable client +// (v1.7.0-pre.3) rejects a non-200 POST response in checkResponse BEFORE +// decoding its body, so on any 4xx the JSON-RPC error object below — +// including data.retryAfterSeconds on the rate limiter's 429 — would be +// discarded unread. 200 is the only status on which the client surfaces the +// coded error to the caller. The request was accepted and processed; the +// failure is an application-level JSON-RPC error riding the transport. +// (writeModernMissingCapability in modern_dispatch.go documents a related +// but distinct 200-over-4xx deviation for -32021.) +// +// id follows writeModernError: absent (via transportsession.HasJSONRPCID) is +// encoded by omitting the "id" key, never as null. +func writeModernCodedError(w http.ResponseWriter, id any, err error, coded mcpparser.CodedError) { + errObj := map[string]any{ + "code": coded.Code(), + "message": err.Error(), + } + if data := coded.Data(); len(data) > 0 { + errObj["data"] = data + } + envelope := map[string]any{ + "jsonrpc": "2.0", + "error": errObj, + } + if transportsession.HasJSONRPCID(id) { + envelope["id"] = id + } + writeModernEnvelope(w, http.StatusOK, envelope) +} + +// writeModernDenied writes a JSON-RPC error envelope at HTTP 403 with +// mcpparser.JSONRPCCodeDenied, mirroring the Legacy call gate +// (call_gate.go) and pkg/authz.handleUnauthorized: the 403 status is what +// makes the audit middleware log the request as denied rather than failed. +// +// id follows writeModernError: absent (via transportsession.HasJSONRPCID) is +// encoded by omitting the "id" key, never as null. +func writeModernDenied(w http.ResponseWriter, id any, msg string) { + envelope := map[string]any{ + "jsonrpc": "2.0", + "error": map[string]any{ + "code": mcpparser.JSONRPCCodeDenied, + "message": msg, + }, + } + if transportsession.HasJSONRPCID(id) { + envelope["id"] = id + } + writeModernEnvelope(w, http.StatusForbidden, envelope) +} + +// writeModernEnvelope marshals envelope before writing headers/status, so a +// marshal failure never leaves a response half-written (mirrors +// classificationErrorBody's build-then-write ordering). +func writeModernEnvelope(w http.ResponseWriter, status int, envelope map[string]any) { + body, err := json.Marshal(envelope) + if err != nil { + // Unreachable in practice: every value passed to these writers is a + // wire struct defined in this file, a domain string, or an int code -- + // all JSON-marshalable. Fall back to a valid JSON-RPC error body + // rather than writing nothing. + slog.Error("failed to marshal Modern JSON-RPC envelope", "error", err) + body = []byte(`{"jsonrpc":"2.0","error":{"code":-32603,"message":"Internal error"}}`) + } + w.Header().Set("Content-Type", "application/json") + // Belt-and-suspenders behind the JSON-level cacheScope:"private" hint: + // an MCP-unaware intermediary (CDN, corporate proxy) wouldn't look at the + // body, so every Modern response also asserts it at the HTTP layer. + w.Header().Set("Cache-Control", "private, no-store") + w.WriteHeader(status) + //nolint:gosec // G104: writing a JSON-RPC response to an HTTP client + _, _ = w.Write(body) +} diff --git a/pkg/vmcp/server/modern_envelope_test.go b/pkg/vmcp/server/modern_envelope_test.go new file mode 100644 index 0000000000..c63c3abc6f --- /dev/null +++ b/pkg/vmcp/server/modern_envelope_test.go @@ -0,0 +1,906 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "encoding/json" + "maps" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/conversion" +) + +const ( + testServerName = "toolhive-vmcp" + testServerVersion = "0.1.0" +) + +// TestModernEnvelopeCommonFields asserts the invariants that apply to every +// Modern result: resultType:"complete", _meta.serverInfo, and Cacheable +// present on the four lists + resources/read but absent on tools/call and +// prompts/get. +func TestModernEnvelopeCommonFields(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + build func(t *testing.T) any + wantCacheable bool + }{ + { + name: "tools/list", + build: func(t *testing.T) any { + t.Helper() + result, err := newModernToolsList([]vmcp.Tool{{ + Name: "greet", + Description: "say hi", + InputSchema: map[string]any{"type": "object"}, + }}, testServerName, testServerVersion, "") + require.NoError(t, err) + return result + }, + wantCacheable: true, + }, + { + name: "resources/list", + build: func(*testing.T) any { + return newModernResourcesList([]vmcp.Resource{ + {Name: "info", URI: "embedded:info", MimeType: "text/plain"}, + }, testServerName, testServerVersion, "") + }, + wantCacheable: true, + }, + { + name: "resources/templates/list", + build: func(*testing.T) any { + return newModernResourceTemplatesList([]vmcp.ResourceTemplate{ + {Name: "logs", URITemplate: "file:///logs/{date}.txt"}, + }, testServerName, testServerVersion, "") + }, + wantCacheable: true, + }, + { + name: "prompts/list", + build: func(*testing.T) any { + return newModernPromptsList([]vmcp.Prompt{ + {Name: "code_review", Arguments: []vmcp.PromptArgument{{Name: "Code", Required: true}}}, + }, testServerName, testServerVersion, "") + }, + wantCacheable: true, + }, + { + name: "tools/call", + build: func(*testing.T) any { + return newModernCallToolResult(&vmcp.ToolCallResult{ + Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "hello"}}, + }, testServerName, testServerVersion) + }, + wantCacheable: false, + }, + { + name: "resources/read", + build: func(*testing.T) any { + return newModernReadResourceResult(&vmcp.ResourceReadResult{ + Contents: []vmcp.ResourceContent{{URI: "embedded:info", MimeType: "text/plain", Text: "hi"}}, + }, testServerName, testServerVersion) + }, + wantCacheable: true, + }, + { + name: "prompts/get", + build: func(*testing.T) any { + return newModernGetPromptResult(&vmcp.PromptGetResult{ + Messages: []vmcp.PromptMessage{{Role: "user", Content: vmcp.Content{Type: vmcp.ContentTypeText, Text: "hi"}}}, + }, testServerName, testServerVersion) + }, + wantCacheable: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(tt.build(t)) + require.NoError(t, err) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(raw, &decoded)) + + require.Equal(t, modernResultTypeComplete, decoded["resultType"], "resultType") + + meta, ok := decoded["_meta"].(map[string]any) + require.True(t, ok, "_meta must be present") + serverInfo, ok := meta[modernServerInfoKey].(map[string]any) + require.True(t, ok, "_meta.%s must be present", modernServerInfoKey) + require.Equal(t, testServerName, serverInfo["name"]) + require.Equal(t, testServerVersion, serverInfo["version"]) + + _, hasTTL := decoded["ttlMs"] + _, hasScope := decoded["cacheScope"] + require.Equal(t, tt.wantCacheable, hasTTL, "ttlMs presence") + require.Equal(t, tt.wantCacheable, hasScope, "cacheScope presence") + if tt.wantCacheable { + require.InDelta(t, 0, decoded["ttlMs"], 0) + require.Equal(t, "private", decoded["cacheScope"], + "vMCP results are admission-filtered per identity; \"public\" would leak across identities") + } + }) + } +} + +// TestModernResultMetaPreservesBackendMeta asserts _meta on tools/call, +// resources/read, and prompts/get carries BOTH the backend's own result.Meta +// keys AND serverInfo. Overwriting _meta with only serverInfo would silently +// discard whatever the backend attached (progress tokens, trace ids, ...), +// diverging from the SDK path's preservation via conversion.ToMCPMeta +// (serve_handlers.go). The no-backend-meta case (nil result.Meta) is already +// covered by TestModernEnvelopeCommonFields, which only asserts serverInfo. +func TestModernResultMetaPreservesBackendMeta(t *testing.T) { + t.Parallel() + + backendMeta := map[string]any{"progressToken": "tok-1", "traceId": "abc"} + + tests := []struct { + name string + build func() any + }{ + { + name: "tools/call", + build: func() any { + return newModernCallToolResult(&vmcp.ToolCallResult{Meta: backendMeta}, testServerName, testServerVersion) + }, + }, + { + name: "resources/read", + build: func() any { + return newModernReadResourceResult(&vmcp.ResourceReadResult{Meta: backendMeta}, testServerName, testServerVersion) + }, + }, + { + name: "prompts/get", + build: func() any { + return newModernGetPromptResult(&vmcp.PromptGetResult{Meta: backendMeta}, testServerName, testServerVersion) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(tt.build()) + require.NoError(t, err) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(raw, &decoded)) + + meta, ok := decoded["_meta"].(map[string]any) + require.True(t, ok, "_meta must be present") + require.Equal(t, "tok-1", meta["progressToken"], "backend meta key must survive") + require.Equal(t, "abc", meta["traceId"], "backend meta key must survive") + + serverInfo, ok := meta[modernServerInfoKey].(map[string]any) + require.True(t, ok, "_meta.%s must still be present alongside backend meta", modernServerInfoKey) + require.Equal(t, testServerName, serverInfo["name"]) + require.Equal(t, testServerVersion, serverInfo["version"]) + + // backendMeta itself must be untouched (copy before mutating caller input). + require.Len(t, backendMeta, 2, "the caller's Meta map must not be mutated") + }) + } +} + +// TestModernResultMetaOverwritesSpoofedServerInfo is a regression case for +// newModernResultMeta's clone-then-set order: if a backend result.Meta +// already contains the "io.modelcontextprotocol/serverInfo" key (e.g. an +// untrusted or misbehaving backend spoofing vMCP's own identity), the real +// vMCP serverInfo must still win on the wire. One builder is enough to pin +// the shared newModernResultMeta logic all three call/read/get builders use. +func TestModernResultMetaOverwritesSpoofedServerInfo(t *testing.T) { + t.Parallel() + + spoofed := map[string]any{ + modernServerInfoKey: map[string]any{"name": "attacker-server", "version": "666"}, + } + + raw, err := json.Marshal(newModernCallToolResult(&vmcp.ToolCallResult{Meta: spoofed}, testServerName, testServerVersion)) + require.NoError(t, err) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(raw, &decoded)) + + meta, ok := decoded["_meta"].(map[string]any) + require.True(t, ok, "_meta must be present") + serverInfo, ok := meta[modernServerInfoKey].(map[string]any) + require.True(t, ok) + require.Equal(t, testServerName, serverInfo["name"], "vMCP's real serverInfo must overwrite a backend-supplied one") + require.Equal(t, testServerVersion, serverInfo["version"]) +} + +// TestReservedResponseMetaStrippedByBothHelpers is the #5986 regression pin: +// the Legacy and Modern response builders must both route backend _meta through +// the SAME mcpparser.StripReservedMeta, so a backend cannot set a reserved +// io.modelcontextprotocol/* key on a result and have it reach the client as if +// vMCP had set it. +// +// It compares the two HELPERS, not the two end-to-end paths -- the outputs are +// deliberately not equal (Modern always mints its own serverInfo, Legacy never +// does, since the 2025-11-25 revision has no serverInfo _meta key). It also +// cannot speak for Legacy resources/read, which drops backend result _meta +// wholesale today (coreResourceHandler returns bare []mcp.ResourceContents) +// while Modern preserves it -- a pre-existing divergence, out of scope here. +func TestReservedResponseMetaStrippedByBothHelpers(t *testing.T) { + t.Parallel() + + // One fixture, fed to both builders: a hostile backend spoofing vMCP's + // identity and asserting protocol control, alongside metadata it is + // legitimately allowed to attach. + backendMeta := map[string]any{ + "io.modelcontextprotocol/serverInfo": map[string]any{"name": "attacker", "version": "666"}, + "io.modelcontextprotocol/protocolVersion": "1999-01-01", + "io.modelcontextprotocol/clientInfo": map[string]any{"name": "spoofed-client"}, + "io.modelcontextprotocol/clientCapabilities": map[string]any{}, + "io.modelcontextprotocol/logLevel": "debug", + "io.modelcontextprotocol/subscriptionId": "sub-1", + "io.modelcontextprotocol/futureThing": "whatever", + // Reserved but relay-semantic: MUST survive (see passthroughMetaKeys). + "io.modelcontextprotocol/related-task": map[string]any{"taskId": "t-1"}, + // Non-reserved backend metadata: MUST survive. + "progressToken": "tok-1", + "traceparent": "00-abc-def-01", + "custom": "keep-me", + } + snapshot := maps.Clone(backendMeta) + + assertClean := func(t *testing.T, meta map[string]any, wantServerInfo bool) { + t.Helper() + for k := range meta { + if k == mcpparser.ReservedMetaPrefix+"related-task" { + continue + } + if k == modernServerInfoKey && wantServerInfo { + continue + } + require.False(t, strings.HasPrefix(k, mcpparser.ReservedMetaPrefix), + "backend-set reserved key %q must not reach the client", k) + } + require.Equal(t, "tok-1", meta["progressToken"], "non-reserved backend meta must survive") + require.Equal(t, "00-abc-def-01", meta["traceparent"]) + require.Equal(t, "keep-me", meta["custom"]) + require.Equal(t, map[string]any{"taskId": "t-1"}, meta[mcpparser.ReservedMetaPrefix+"related-task"], + "end-to-end task correlation must be relayed, not terminated") + } + + t.Run("modern builder strips and re-stamps its own serverInfo", func(t *testing.T) { + t.Parallel() + + got := newModernResultMeta(backendMeta, testServerName, testServerVersion) + assertClean(t, got, true) + + serverInfo, ok := got[modernServerInfoKey].(modernServerInfo) + require.True(t, ok, "vMCP must stamp its own serverInfo") + require.Equal(t, testServerName, serverInfo.Name, "vMCP's serverInfo must beat the backend's") + require.Equal(t, testServerVersion, serverInfo.Version) + }) + + t.Run("legacy builder strips and stamps nothing", func(t *testing.T) { + t.Parallel() + + got := conversion.ToMCPMeta(backendMeta) + require.NotNil(t, got) + // ToMCPMeta hoists progressToken out of the map, so put it back before + // running the shared assertions. + flat := maps.Clone(got.AdditionalFields) + flat["progressToken"] = got.ProgressToken + assertClean(t, flat, false) + }) + + t.Run("the caller's map is never mutated", func(t *testing.T) { + t.Parallel() + + // Compare against a snapshot, not just the length: a swap-one-add-one + // mutation would keep the count identical. + require.Equal(t, snapshot, backendMeta, "neither builder may mutate the domain result's map") + }) +} + +// TestModernEnvelopeEmptyCollections asserts that an empty domain slice +// marshals to a JSON array ([]), never null. +func TestModernEnvelopeEmptyCollections(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + field string + build func(t *testing.T) any + }{ + { + name: "tools/list", + field: "tools", + build: func(t *testing.T) any { + t.Helper() + result, err := newModernToolsList(nil, testServerName, testServerVersion, "") + require.NoError(t, err) + return result + }, + }, + { + name: "resources/list", + field: "resources", + build: func(*testing.T) any { return newModernResourcesList(nil, testServerName, testServerVersion, "") }, + }, + { + name: "resources/templates/list", + field: "resourceTemplates", + build: func(*testing.T) any { + return newModernResourceTemplatesList(nil, testServerName, testServerVersion, "") + }, + }, + { + name: "prompts/list", + field: "prompts", + build: func(*testing.T) any { return newModernPromptsList(nil, testServerName, testServerVersion, "") }, + }, + { + name: "tools/call content", + field: "content", + build: func(*testing.T) any { + return newModernCallToolResult(&vmcp.ToolCallResult{}, testServerName, testServerVersion) + }, + }, + { + name: "resources/read contents", + field: "contents", + build: func(*testing.T) any { + return newModernReadResourceResult(&vmcp.ResourceReadResult{}, testServerName, testServerVersion) + }, + }, + { + name: "prompts/get messages", + field: "messages", + build: func(*testing.T) any { + return newModernGetPromptResult(&vmcp.PromptGetResult{}, testServerName, testServerVersion) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(tt.build(t)) + require.NoError(t, err) + + var decoded map[string]json.RawMessage + require.NoError(t, json.Unmarshal(raw, &decoded)) + require.JSONEq(t, "[]", string(decoded[tt.field]), "%s must marshal as [] not null", tt.field) + }) + } +} + +// TestModernCallToolResult covers tools/call-specific behavior not shared +// with the other builders: isError passthrough and conditional +// structuredContent. +func TestModernCallToolResult(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + domainResult *vmcp.ToolCallResult + wantIsError bool + wantStructured bool + }{ + { + name: "success", + domainResult: &vmcp.ToolCallResult{Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "ok"}}}, + }, + { + name: "isError true passthrough", + domainResult: &vmcp.ToolCallResult{ + Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "boom"}}, + IsError: true, + }, + wantIsError: true, + }, + { + name: "structuredContent present when set", + domainResult: &vmcp.ToolCallResult{StructuredContent: map[string]any{"count": float64(1)}}, + wantStructured: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(newModernCallToolResult(tt.domainResult, testServerName, testServerVersion)) + require.NoError(t, err) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(raw, &decoded)) + + gotIsError, _ := decoded["isError"].(bool) + require.Equal(t, tt.wantIsError, gotIsError) + + _, hasStructured := decoded["structuredContent"] + require.Equal(t, tt.wantStructured, hasStructured) + }) + } +} + +// TestModernCompleteTolerateNil asserts newModernComplete doesn't panic on a +// nil result -- defense-in-depth against a plausible aggregator bug, not an +// expected input (see its doc comment). newModernCallToolResult, +// newModernReadResourceResult, and newModernGetPromptResult carry no +// equivalent guard: their callers in modern_dispatch.go dereference +// result.BackendID before invoking the builder, so a nil result already +// panics upstream. +func TestModernCompleteTolerateNil(t *testing.T) { + t.Parallel() + + require.NotPanics(t, func() { + newModernComplete(nil, "other-server", "9.9.9") + }) +} + +// TestModernComplete pins the completion/complete wire shape from +// newModernComplete against the SDK's CompleteResult/CompletionResultDetails +// (protocol.go:653,660 in go-sdk@v1.7.0-pre.3): a "completion" object with +// values/total/hasMore, resultType, and _meta.serverInfo -- no Cacheable, +// matching the SDK (CompleteResult never embeds it). JSONEq's exact-match +// means an errant Cacheable field would fail these cases on its own, without +// a separate assertion. +func TestModernComplete(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + result *vmcp.CompletionResult + want string + }{ + { + name: "nil result marshals an empty completion", + result: nil, + want: `{ + "resultType": "complete", + "completion": {"values": []}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "nil Values slice marshals as [] not null", + result: &vmcp.CompletionResult{}, + want: `{ + "resultType": "complete", + "completion": {"values": []}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "values/total/hasMore pass through", + result: &vmcp.CompletionResult{Values: []string{"a", "b"}, Total: 5, HasMore: true}, + want: `{ + "resultType": "complete", + "completion": {"values": ["a", "b"], "total": 5, "hasMore": true}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(newModernComplete(tt.result, testServerName, testServerVersion)) + require.NoError(t, err) + require.JSONEq(t, tt.want, string(raw)) + }) + } +} + +// TestModernWriters covers the three HTTP write helpers. writeModernError's +// status derives from the JSON-RPC code (404 for method-not-found, 400 for +// invalid-params, 200 for everything else -- mirroring go-sdk's +// extractErrorStatus); writeModernDenied always changes the HTTP status to +// 403 for a POLICY reason -- that's the signal the audit middleware's +// determineOutcome keys off to log outcome:"denied", unrelated to the +// protocol-level mapping above. +func TestModernWriters(t *testing.T) { + t.Parallel() + + t.Run("writeModernResult", func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeModernResult(rec, "req-1", map[string]any{"ok": true}) + + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "application/json", rec.Header().Get("Content-Type")) + require.Equal(t, "private, no-store", rec.Header().Get("Cache-Control")) + + var decoded struct { + JSONRPC string `json:"jsonrpc"` + ID string `json:"id"` + Result map[string]any `json:"result"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &decoded)) + require.Equal(t, "2.0", decoded.JSONRPC) + require.Equal(t, "req-1", decoded.ID) + require.Equal(t, true, decoded.Result["ok"]) + }) + + // writeModernResult must ALWAYS emit "id", unlike writeModernError/ + // writeModernDenied below -- a result response's id is required by the + // MCP schema, not optional. This pins that a future refactor doesn't + // "helpfully" make it conditional to match its error-writing siblings. + t.Run("writeModernResult always emits id, even absent", func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeModernResult(rec, nil, map[string]any{"ok": true}) + + var asMap map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &asMap)) + _, hasID := asMap["id"] + require.True(t, hasID, `"id" key must be present on a result response even when id is nil`) + }) + + // writeModernError's "id" key follows transportsession.HasJSONRPCID: MCP + // narrows base JSON-RPC 2.0 so an absent id is encoded by omitting the + // key, never as "id":null (schema/2025-11-25 types the error response id + // as optional string|number). A map decode (not a tagged struct) is the + // only way to tell an absent key apart from a present null. + t.Run("writeModernError id presence", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + id any + wantIDKey bool + wantIDEcho string // substring the body must contain, verbatim, when wantIDKey + }{ + {name: "nil id is omitted", id: nil, wantIDKey: false}, + {name: "zero id is still present", id: int64(0), wantIDKey: true, wantIDEcho: `"id":0`}, + {name: "present id is present", id: "req-2", wantIDKey: true, wantIDEcho: `"id":"req-2"`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeModernError(rec, tt.id, jsonRPCCodeInternalError, "some message") + + var asMap map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &asMap)) + _, hasID := asMap["id"] + require.Equal(t, tt.wantIDKey, hasID, `"id" key presence`) + if tt.wantIDKey { + // A map decode would coerce a numeric id to float64, + // losing the verbatim-echo guarantee -- check the raw + // wire bytes instead. + require.Contains(t, rec.Body.String(), tt.wantIDEcho) + } + }) + } + }) + + t.Run("writeModernError", func(t *testing.T) { + t.Parallel() + + codeToStatus := []struct { + code int + wantStatus int + }{ + {jsonRPCCodeMethodNotFound, http.StatusNotFound}, + {jsonRPCCodeInvalidParams, http.StatusBadRequest}, + {jsonRPCCodeInvalidRequest, http.StatusOK}, + {jsonRPCCodeInternalError, http.StatusOK}, + } + + for _, tc := range codeToStatus { + rec := httptest.NewRecorder() + writeModernError(rec, "req-2", tc.code, "some message") + + require.Equal(t, tc.wantStatus, rec.Code, "code %d", tc.code) + require.Equal(t, "application/json", rec.Header().Get("Content-Type")) + require.Equal(t, "private, no-store", rec.Header().Get("Cache-Control")) + + var decoded struct { + Error struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &decoded)) + require.Equal(t, tc.code, decoded.Error.Code) + require.Equal(t, "some message", decoded.Error.Message) + } + }) + + t.Run("writeModernDenied", func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeModernDenied(rec, "req-3", "denied by policy") + + require.Equal(t, http.StatusForbidden, rec.Code, "a denial must change the HTTP status for audit outcome:\"denied\"") + require.Equal(t, "application/json", rec.Header().Get("Content-Type")) + require.Equal(t, "private, no-store", rec.Header().Get("Cache-Control")) + + var decoded struct { + Error struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &decoded)) + require.Equal(t, int(mcpparser.JSONRPCCodeDenied), decoded.Error.Code) + require.Equal(t, "denied by policy", decoded.Error.Message) + }) + + // writeModernDenied's "id" key follows the same HasJSONRPCID rule as + // writeModernError above. + t.Run("writeModernDenied id presence", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + id any + wantIDKey bool + wantIDEcho string + }{ + {name: "nil id is omitted", id: nil, wantIDKey: false}, + {name: "zero id is still present", id: int64(0), wantIDKey: true, wantIDEcho: `"id":0`}, + {name: "present id is present", id: "req-3", wantIDKey: true, wantIDEcho: `"id":"req-3"`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeModernDenied(rec, tt.id, "denied by policy") + + var asMap map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &asMap)) + _, hasID := asMap["id"] + require.Equal(t, tt.wantIDKey, hasID, `"id" key presence`) + if tt.wantIDKey { + require.Contains(t, rec.Body.String(), tt.wantIDEcho) + } + }) + } + }) +} + +// TestWriteModernEnvelopeMarshalFallback verifies the defensive fallback +// drops the "id" key (never "id":null) when the envelope fails to marshal. +// A channel is not JSON-marshalable, forcing json.Marshal to fail. +func TestWriteModernEnvelopeMarshalFallback(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + writeModernEnvelope(rec, http.StatusOK, map[string]any{"bad": make(chan int)}) + + const want = `{"jsonrpc":"2.0","error":{"code":-32603,"message":"Internal error"}}` + require.Equal(t, want, rec.Body.String()) + require.NotContains(t, rec.Body.String(), `"id"`) +} + +// TestModernDiscover asserts server/discover's capability-flags shape: a +// capability field is present iff the corresponding admitted list was +// non-empty, resources/templates fold into the single "resources" flag, and +// no descriptor arrays ever appear on the wire. +func TestModernDiscover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + hasTools, hasResources, hasTemplates, hasPrompts bool + want string + }{ + { + name: "nothing admitted -- only the static completions capability advertised", + want: `{ + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + "supportedVersions": ["2026-07-28", "2025-11-25"], + "capabilities": {"completions": {}}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "tools only", + hasTools: true, + want: `{ + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + "supportedVersions": ["2026-07-28", "2025-11-25"], + "capabilities": {"tools": {}, "completions": {}}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "resources only", + hasResources: true, + want: `{ + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + "supportedVersions": ["2026-07-28", "2025-11-25"], + "capabilities": {"resources": {}, "completions": {}}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "templates only still sets resources flag", + hasTemplates: true, + want: `{ + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + "supportedVersions": ["2026-07-28", "2025-11-25"], + "capabilities": {"resources": {}, "completions": {}}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "prompts only", + hasPrompts: true, + want: `{ + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + "supportedVersions": ["2026-07-28", "2025-11-25"], + "capabilities": {"prompts": {}, "completions": {}}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + { + name: "everything admitted", + hasTools: true, + hasResources: true, + hasTemplates: true, + hasPrompts: true, + want: `{ + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + "supportedVersions": ["2026-07-28", "2025-11-25"], + "capabilities": {"tools": {}, "resources": {}, "prompts": {}, "completions": {}}, + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "toolhive-vmcp", "version": "0.1.0"}} + }`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(newModernDiscover( + tt.hasTools, tt.hasResources, tt.hasTemplates, tt.hasPrompts, testServerName, testServerVersion, + )) + require.NoError(t, err) + require.JSONEq(t, tt.want, string(raw)) + }) + } +} + +// TestModernDescriptorFieldMapping pins the domain->wire descriptor mapping +// (the list-item shape, not the envelope wrapping it): tool name/description/ +// inputSchema (including the RawInputSchema round-trip)/outputSchema/ +// annotations, resource uri/name/mimeType, resource template uriTemplate, +// and prompt arguments[].required with omitempty. +func TestModernDescriptorFieldMapping(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + build func(t *testing.T) any + want string + }{ + { + name: "tool with annotations and output schema", + build: func(t *testing.T) any { + t.Helper() + readOnly := true + wireTool, err := modernToolFromDomain(vmcp.Tool{ + Name: "greet", + Description: "say hi", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"name": map[string]any{"type": "string"}}, + }, + OutputSchema: map[string]any{"type": "object"}, + Annotations: &vmcp.ToolAnnotations{ReadOnlyHint: &readOnly}, + }) + require.NoError(t, err) + return wireTool + }, + want: `{ + "name": "greet", + "description": "say hi", + "inputSchema": {"type": "object", "properties": {"name": {"type": "string"}}}, + "outputSchema": {"type": "object"}, + "annotations": {"readOnlyHint": true} + }`, + }, + { + // Pins the known documented behavior at the mapping site: a + // tool with no annotations still emits "annotations":{} because + // mcpcompat's Tool.MarshalJSON writes the field unconditionally. + name: "tool with no annotations still emits annotations:{}", + build: func(t *testing.T) any { + t.Helper() + wireTool, err := modernToolFromDomain(vmcp.Tool{ + Name: "noop", + InputSchema: map[string]any{"type": "object"}, + }) + require.NoError(t, err) + return wireTool + }, + want: `{"name": "noop", "inputSchema": {"type": "object"}, "annotations": {}}`, + }, + { + name: "resource", + build: func(*testing.T) any { + return modernResourceFromDomain(vmcp.Resource{ + Name: "info", URI: "embedded:info", Description: "info doc", MimeType: "text/plain", + }) + }, + want: `{"uri":"embedded:info","name":"info","description":"info doc","mimeType":"text/plain"}`, + }, + { + name: "resource template", + build: func(*testing.T) any { + return modernResourceTemplateFromDomain(vmcp.ResourceTemplate{ + Name: "logs", URITemplate: "file:///logs/{date}.txt", Description: "daily logs", MimeType: "text/plain", + }) + }, + want: `{"uriTemplate":"file:///logs/{date}.txt","name":"logs","description":"daily logs","mimeType":"text/plain"}`, + }, + { + name: "prompt arguments[].required with omitempty", + build: func(*testing.T) any { + return modernPromptFromDomain(vmcp.Prompt{ + Name: "code_review", + Description: "do a code review", + Arguments: []vmcp.PromptArgument{ + {Name: "Code", Required: true}, + {Name: "Language", Description: "optional hint"}, + }, + }) + }, + want: `{ + "name": "code_review", + "description": "do a code review", + "arguments": [ + {"name": "Code", "required": true}, + {"name": "Language", "description": "optional hint"} + ] + }`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + raw, err := json.Marshal(tt.build(t)) + require.NoError(t, err) + require.JSONEq(t, tt.want, string(raw)) + }) + } +} diff --git a/pkg/vmcp/server/modern_gate.go b/pkg/vmcp/server/modern_gate.go new file mode 100644 index 0000000000..21223ea1d2 --- /dev/null +++ b/pkg/vmcp/server/modern_gate.go @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +// modernDispatchBlockers enumerates the enabled features of THIS instance that +// the stateless Modern (2026-07-28) dispatch path cannot serve yet. It is the +// capability gate that decides whether vMCP advertises and serves the Modern +// revision: an empty result means every enabled feature is servable by +// dispatchModern and Modern requests are dispatched; a non-empty result means +// classifyingHandler keeps Modern-capable clients on Legacy (see the gate +// branch there for the wire mechanics). +// +// Contract for this list: +// +// - One entry per feature, each guarded by the narrowest signal that the +// feature is actually enabled on this instance, with a comment saying WHY +// the Modern path cannot serve it. "Cannot serve" means a Modern client +// would silently receive different behavior than the feature promises — +// not merely that the feature is session-flavored. A feature that Modern +// clients simply don't need (e.g. Redis-backed session sharing: Legacy +// clients keep their shared sessions, Modern clients are sessionless by +// design and store nothing) does NOT belong here; coexistence of that kind +// is asserted by test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_redis_test.go. +// +// - Every entry's comment must cite the issue tracking its Modern parity, so +// a stale entry (parity shipped, entry forgotten) surfaces in issue triage +// rather than silently keeping the instance off Modern. +// +// - When Modern parity lands for a feature, delete its entry. The gate's +// behavior is pinned by TestModernDispatchBlockers and +// TestClassifyingHandler_ModernCapabilityGate (classification_test.go) plus +// the full-handler pair in modern_gate_integration_test.go; deleting an +// entry must flip cases there, so parity work cannot silently ship without +// updating them. +// +// - Loud refusals are deliberately out of scope. A composite workflow with an +// elicitation step (config.WorkflowStepConfig type "elicitation") works over +// Legacy sessions and fails Modern clients with an explicit -32021 +// MissingRequiredClientCapabilityError (or -32603 when the client declared +// the capability) — see writeModernCallFailure. That is an honest error the +// client can act on, not silently different behavior, so it does not gate +// the whole instance off Modern for one workflow definition. Distinct from +// #6059, which covers a Modern BACKEND returning input_required; here vMCP +// itself is the elicitor. +// +// The result is derived from construction-time configuration only, so it is +// constant for the life of the Server; Serve logs it once at startup. +func (s *Server) modernDispatchBlockers() []string { + var blocked []string + + // Optimizer: find_tool/call_tool are Serve-layer, session-scoped meta-tools + // (serve_optimizer.go). Each session builds an FTS5 index over its advertised + // set and swaps the two meta-tools in place of the raw tools; the index is + // transport/session state and is deliberately NOT in the stateless core. + // dispatchModern serves tools/list and tools/call straight from + // core.ListTools/core.CallTool, so a Modern client of an optimizer-enabled + // instance would silently receive the full raw aggregated tool set and + // `tools/call find_tool` would fail -32603 "not found" — the optimizer + // feature would be invisibly disabled for exactly the newest clients. + // Modern parity needs an identity- or instance-scoped index to replace the + // session-scoped one; that work is tracked in #6089, and deleting this entry + // is its definition of done. Until it lands, an optimizer-enabled instance + // is Legacy-only. A non-nil s.optimizerFactory is a faithful "optimizer + // enabled" signal: sessionmanager.New's constructor guard rejects an + // optimizer without AdvertiseFromCore at startup, so the factory can never + // be enabled yet invisible here. + if s.optimizerFactory != nil { + blocked = append(blocked, blockerOptimizer) + } + + return blocked +} + +// blockerOptimizer names the optimizer's entry in the blocker list; +// TestModernDispatchBlockers asserts on it by name. +const blockerOptimizer = "optimizer" diff --git a/pkg/vmcp/server/modern_gate_integration_test.go b/pkg/vmcp/server/modern_gate_integration_test.go new file mode 100644 index 0000000000..9570c6aaa7 --- /dev/null +++ b/pkg/vmcp/server/modern_gate_integration_test.go @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "context" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpsdk "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" +) + +// supportedVersionsFromDiscover extracts result.supportedVersions from a decoded +// server/discover envelope, failing loudly when the field is absent or malformed +// so a "does not contain" assertion can never pass vacuously against the wrong +// response shape. +func supportedVersionsFromDiscover(t *testing.T, decoded map[string]any) []string { + t.Helper() + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC result envelope, got %v", decoded) + raw, ok := result["supportedVersions"].([]any) + require.True(t, ok, "expected result.supportedVersions, got %v", result) + require.NotEmpty(t, raw, "supportedVersions must never be empty: %v", result) + versions := make([]string, 0, len(raw)) + for _, v := range raw { + s, ok := v.(string) + require.True(t, ok, "supportedVersions entries must be strings: %v", raw) + versions = append(versions, s) + } + return versions +} + +// TestIntegration_ModernGate_OptimizerKeepsClientsOnLegacy pins the capability +// gate (modernDispatchBlockers, modern_gate.go) through the full handler chain +// for an optimizer-enabled instance — the configuration whose find_tool/ +// call_tool meta-tools are session-scoped and therefore unservable by the +// stateless Modern dispatch path: +// +// 1. server/discover must NOT advertise 2026-07-28: the gate lets it fall +// through to the SDK, whose stateful transport filters Modern out of the +// version list. This is the leg that steers a Modern-first client (go-sdk +// v1.7+ probes discover before initialize) onto the Legacy handshake, +// where sessions advertise find_tool/call_tool. +// 2. Any other well-formed Modern request is refused with a conformant +// 400 + -32022 UnsupportedProtocolVersionError listing the Legacy version +// — never go-sdk's plain-text stateful rejection, which carries no +// version list a client could negotiate down from. +// +// The contrast leg (an instance without the optimizer advertises and serves +// Modern) is TestIntegration_ModernGate_PlainInstanceAdvertisesModern; the two +// together stop the gate from being "simplified" away in either direction. +// Deleting the optimizer entry from modernDispatchBlockers (Modern optimizer +// parity) must flip this test. +func TestIntegration_ModernGate_OptimizerKeepsClientsOnLegacy(t *testing.T) { + t.Parallel() + + tool := vmcp.Tool{Name: "echo", Description: "echoes", InputSchema: map[string]any{"type": "object"}} + ts := buildTestServerWithOptions(t, newNoopMockFactory(t), serverOptions{ + tools: []vmcp.Tool{tool}, + optimizerFactory: func(_ context.Context, _ []mcpsdk.ServerTool) (optimizer.Optimizer, error) { + return &fakeOptimizer{}, nil + }, + }) + + discoverResp, discoverBody := postModern(t, ts.URL, "server/discover", nil, 1, "") + defer discoverResp.Body.Close() + require.Equal(t, http.StatusOK, discoverResp.StatusCode, + "gated discover must still be answered (fall-through to the SDK): %+v", discoverBody) + versions := supportedVersionsFromDiscover(t, discoverBody) + assert.NotContains(t, versions, "2026-07-28", + "an optimizer-enabled instance must not advertise Modern") + assert.Contains(t, versions, "2025-11-25", + "the Legacy version must stay advertised so the client can negotiate down") + + listResp, listBody := postModern(t, ts.URL, "tools/list", nil, 2, "") + defer listResp.Body.Close() + assert.Equal(t, http.StatusBadRequest, listResp.StatusCode, "decoded: %+v", listBody) + errObj, ok := listBody["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope, got %+v", listBody) + assert.Equal(t, float64(-32022), errObj["code"], + "refusal must be the conformant UnsupportedProtocolVersionError") + data, ok := errObj["data"].(map[string]any) + require.True(t, ok, "expected error.data with the version lists, got %+v", errObj) + assert.Equal(t, []any{"2025-11-25"}, data["supported"], + "the refusal must list the Legacy version so the client can negotiate down") +} + +// TestIntegration_ModernGate_PlainInstanceAdvertisesModern is the contrast leg +// of TestIntegration_ModernGate_OptimizerKeepsClientsOnLegacy: with no enabled +// feature the stateless path cannot serve, the gate is open, server/discover is +// answered by dispatchModernDiscover, and 2026-07-28 is advertised. +func TestIntegration_ModernGate_PlainInstanceAdvertisesModern(t *testing.T) { + t.Parallel() + + tool := vmcp.Tool{Name: "echo", Description: "echoes", InputSchema: map[string]any{"type": "object"}} + ts := buildTestServerWithOptions(t, newNoopMockFactory(t), serverOptions{tools: []vmcp.Tool{tool}}) + + resp, decoded := postModern(t, ts.URL, "server/discover", nil, 1, "") + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + assert.Contains(t, supportedVersionsFromDiscover(t, decoded), "2026-07-28", + "a plain instance must advertise Modern") +} diff --git a/pkg/vmcp/server/modern_pagination.go b/pkg/vmcp/server/modern_pagination.go new file mode 100644 index 0000000000..7bd6a0170c --- /dev/null +++ b/pkg/vmcp/server/modern_pagination.go @@ -0,0 +1,366 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "slices" + "strings" +) + +// Client-facing list pagination for the Modern (2026-07-28) dispatcher. +// +// WHY THIS IS NOT A PORT OF THE LEGACY BEHAVIOR. vMCP already follows +// pagination cursors in two places, but both are UPSTREAM -- vMCP acting as a +// client, walking a backend's pages to assemble the aggregated view +// (pkg/vmcp/client, pkg/vmcp/session/internal/backend, via +// pkg/vmcp/internal/pagination.ListAll; #5851). The giveaway that none of it is +// reusable here is the shape of the entry point: ListAll takes a `fetch` +// CALLBACK and drives it until the upstream stops yielding cursors, whereas this +// file takes a SLICE already fully in hand. ListAll holds no encode/decode logic +// and never mints a cursor -- it only consumes them. This is the opposite +// direction: vMCP as a SERVER, splitting its own already-aggregated list into +// pages for a downstream client. On the Legacy path the SDK's session-scoped +// feature store does that split; dispatchModern bypasses the SDK, so it must do +// it itself. +// +// Cursor TOKENS are consequently not interchangeable between the two revisions +// even though page BOUNDARIES are (see modernPageSize): go-sdk's Legacy server +// mints padded base64.URLEncoding over gob, this mints RawURLEncoding over JSON. +// Neither side ever needs to read the other's -- a client re-lists after +// negotiating, and the opacity rule forbids it from trying. +// +// THE STATELESS CONSTRAINT, AND WHY THE SPEC PERMITS THIS. Modern has no +// sessions, so a cursor may not denote server-held iteration state -- there is +// nowhere to keep it, and any two requests may land on different replicas. +// +// The governing text is the DRAFT pagination page +// (docs/specification/draft/server/utilities/pagination.mdx), not the 2025-11-25 +// one: they differ, and the draft is what 2026-07-28 ships. Its normative rules: +// +// - the cursor is "an opaque string token, representing a position in the +// result set" (Pagination Model); +// - clients "MUST treat cursors as opaque tokens" -- no assumptions about +// format, no parsing, no modifying (Implementation Guidelines 3); +// - servers "SHOULD provide stable cursors and handle invalid cursors +// gracefully" (Implementation Guidelines 1); +// - "Invalid cursors SHOULD result in an error with code -32602" (Error +// Handling). +// +// That opacity MUST, on its own, licenses this design: if a client may not +// interpret the token, the server is free to encode position INTO it rather than +// remembering it. No further justification is needed. +// +// ONE DRAFT-ONLY AMENDMENT, and it is load-bearing here. The 2025-11-25 bullet +// "Don't persist cursors across sessions" is GONE, replaced by: "Don't make any +// determination based on cursor value other than whether a non-null value was +// provided (e.g. an empty string is a valid cursor and thus MUST NOT be treated +// as the end of results)." The reason for the deletion is recorded in +// docs/seps/2567-sessionless-mcp.md under "Consequential spec edits" -- note that +// SEP is a list of directed edits, i.e. rationale, NOT the normative surface; +// the draft page above is. +// +// The consequence here: end-of-results MUST be signalled by OMITTING nextCursor, +// never by emitting an empty string, because a conformant client treats "" as a +// valid cursor and re-requests with it. Hence `omitempty` on every list result's +// NextCursor field (modern_envelope.go) -- without it a client loops forever on +// page one. +// +// A NOTE ON A TEMPTING ARGUMENT, flagged so nobody re-derives it as doctrine: +// one could argue a self-describing cursor is *required* now, since with sessions +// gone a server-held cursor has no protocol lifetime bounding it. The engineering +// half is true and worth knowing -- a stateful scheme would have to invent its own +// TTL and eviction. But it is NOT a spec consequence: the deleted bullet was +// advice to CLIENTS about persistence, and removing it withdraws a client-side +// caveat rather than creating a server-side obligation. SEP-2567 explicitly +// declines to speak here, putting "cursor stability and snapshot consistency" +// out of scope. Treat the TTL point as our reasoning; treat opacity as the +// authority. +// +// AGGREGATION ACROSS BACKENDS needs no per-backend positions. core.List* returns +// the complete, admission-filtered set in one call, so the cursor encodes a +// position in the AGGREGATED ordering; it never names a backend, and adding or +// removing one cannot invalidate it. Each page re-runs the aggregation -- the same +// per-request fan-out dispatchModernDiscover already documents, not a new one. +// Re-sorting per page is O(n log n) per request, measured at 0.76ms for n=1100, +// i.e. dwarfed by that fan-out. Do not "optimise" it with a cache. + +// modernPageSize is the maximum number of items in one Modern list page. +// +// It matches go-sdk's DefaultPageSize (server.go:36-37, value 1000), which is +// what the Legacy/SDK path uses for this same split: vMCP never calls mcpcompat's +// WithPageSize, so the SDK default is in force there. Keeping the two equal means +// a client sees the same page boundaries whichever revision it negotiates. +// +// This is a COPY, not a reference -- mcpcompat does not re-export the constant +// (raising that upstream is tracked separately). Provenance is not detection: if +// go-sdk changes its default, the Legacy path follows automatically and this one +// does not, silently falsifying the equality claimed above. +// +// KNOWN GAP, stated rather than papered over: nothing currently detects that +// drift. The Modern side IS pinned behaviourally in-package (a 1001-item corpus +// must yield a 1000-item first page), but not end-to-end: +// test/integration/vmcp's Over1000Tools regression exercises the Legacy split +// only -- its helpers are mcpcompat-backed, and mcpcompat cannot request +// 2026-07-28 (#5911) -- so no in-tree integration test pins the Modern split +// end-to-end. (A per-request Legacy/Modern pinning mechanism now exists in +// test/e2e's RawMCPClient, but the integration-tier helpers are still +// mcpcompat-backed.) +// +// So: treat the equality claimed above as an invariant maintained by REVIEW, not +// by CI, and re-check it on any go-sdk bump. This deliberately does not name a +// mechanism the missing Modern end-to-end coverage should use; whether that is a +// pinning harness, a loud failure at negotiation, or something else is +// undecided, and this comment must not go stale by presuming one. +const modernPageSize = 1000 + +// modernCursorMaxLen caps an inbound cursor before it is decoded. +// +// A legitimate cursor is tens of bytes (kind, one key, a small count). The +// request body cap is 8MB and list verbs are not rate-limited +// (ratelimit/decorator.go covers CallTool only), so without this a client could +// hand over a well-formed multi-megabyte cursor that decodes successfully -- +// measured at ~295ms of CPU for a 6MB key -- and repeat it unmetered. Rejecting +// on LENGTH before decoding keeps that O(1). 512 is generous: it accommodates a +// ~370-byte key, far beyond any real URI. +const modernCursorMaxLen = 512 + +// Cursor kinds, one per paginated list verb. The kind is carried inside the +// cursor and checked on decode so a cursor minted for one list cannot be +// replayed against another: without it, a tools cursor sent to prompts/list +// would be silently reinterpreted as a prompt name and return a plausible but +// meaningless page. +const ( + cursorKindTools = "tools" + cursorKindResources = "resources" + cursorKindResourceTemplates = "resourceTemplates" + cursorKindPrompts = "prompts" +) + +// errInvalidModernCursor marks a cursor that is over-long, malformed, not a JSON +// string, or well-formed but minted for a different list verb. Callers map it to +// -32602, which the draft pagination page's Error Handling section requires +// ("Invalid cursors SHOULD result in an error with code -32602") and which +// schema.ts classifies under invalid-params ("Pagination: Invalid or expired +// cursor values"). go-sdk does the same, returning ErrInvalidParams for a cursor +// it cannot decode (server.go:2091-2094). +var errInvalidModernCursor = errors.New("invalid pagination cursor") + +// modernCursor is the decoded cursor payload: a position in the sorted aggregated +// list, fully self-described so no server state is needed to interpret it. +// +// The JSON field names are single letters only to keep the token short; nothing +// outside this file may depend on them, since the token is opaque by spec. +type modernCursor struct { + Kind string `json:"k"` + LastKey string `json:"l"` + // Delivered is how many items sharing LastKey have already been sent, counted + // across every page so far. It is what makes the paginator safe against + // DUPLICATE keys -- see paginateModern for why a bare "strictly above + // LastKey" scan silently loses items. + Delivered int `json:"d,omitempty"` +} + +// encodeModernCursor builds the opaque token handed to the client as nextCursor. +// base64url keeps it safe in any JSON string; it is emphatically NOT encryption +// or authentication, and nothing secret may go in it. A client can trivially +// decode and forge one, which is harmless: a cursor only selects a position +// within a set the caller is already entitled to (admission filtering runs in +// core.List* before this file sees anything) and carries no identity, backend, or +// capability of its own. +func encodeModernCursor(kind, lastKey string, delivered int) (string, error) { + payload, err := json.Marshal(modernCursor{Kind: kind, LastKey: lastKey, Delivered: delivered}) + if err != nil { + return "", fmt.Errorf("encoding pagination cursor: %w", err) + } + return base64.RawURLEncoding.EncodeToString(payload), nil +} + +// decodeModernCursor recovers a position from a client-supplied cursor, rejecting +// anything that is not a well-formed cursor for wantKind. +// +// Every failure collapses to errInvalidModernCursor rather than surfacing the +// decode detail: the cursor is client-controlled input, and echoing back why it +// failed to parse tells a caller about the internal encoding it is required to +// treat as opaque. +func decodeModernCursor(wantKind, cursor string) (string, int, error) { + // Length first, so an oversized cursor costs nothing to reject. + if len(cursor) > modernCursorMaxLen { + return "", 0, errInvalidModernCursor + } + raw, err := base64.RawURLEncoding.DecodeString(cursor) + if err != nil { + return "", 0, errInvalidModernCursor + } + var decoded modernCursor + if err := json.Unmarshal(raw, &decoded); err != nil { + return "", 0, errInvalidModernCursor + } + // A syntactically valid decode is not enough: the kind must match the verb + // being served, an empty LastKey would make the scan below return the whole + // list (silently turning a page request into a full-list request), and a + // negative Delivered would rewind. + if decoded.Kind != wantKind || decoded.LastKey == "" || decoded.Delivered < 0 { + return "", 0, errInvalidModernCursor + } + return decoded.LastKey, decoded.Delivered, nil +} + +// paginateModern returns the page of items following cursor, plus the cursor for +// the page after it (empty when the returned page is the last one). +// +// KEY UNIQUENESS IS NOT ASSUMED, and this is the subtle part. An earlier version +// documented these keys as globally unique "since the aggregator's conflict +// resolver has already de-duplicated them across backends", and skipped forward +// with a bare `key > lastKey`. At the time that precondition held for exactly +// ONE of the four callers (Tool.Name; resource URIs, template strings and +// prompt names were plain-appended across backends with no de-duplication). +// The aggregator has since been fixed to resolve all four (#6060: +// aggregator/capability_conflicts.go), but this paginator is generic and MUST +// NOT depend on that caller invariant — its signature cannot enforce it, and a +// future caller may not uphold it. The collision-safety below stays. +// +// With a bare `>` scan, a collision at a page boundary permanently DROPS an item: +// the page ends on the first copy, the cursor names that key, and the scan then +// skips every copy of it, so the second appears on no page at all. Measured +// before this fix, a 1100-resource set with one duplicated URI delivered 1099. +// +// So the cursor carries (LastKey, Delivered) and the scan resumes *within* a run +// of equal keys rather than past it. Items sharing a key are interchangeable for +// this purpose -- Delivered is a position in the run, not an identity -- which +// keeps the walk correct without a secondary sort key this generic signature +// cannot see. A Delivered larger than the run (a copy vanished between pages) +// degrades to "resume at the next key" rather than failing. +// +// keyOf must still yield the item's natural identity; it just no longer has to be +// unique. +// +// items is never mutated: it is cloned before sorting, because it comes straight +// from core.List* and the caller's slice must not be reordered underneath it. +func paginateModern[T any](items []T, keyOf func(T) string, kind, cursor string) ([]T, string, error) { + // Sorting makes the ordering deterministic across requests, which keyset + // pagination requires: the aggregator's own order is not guaranteed stable + // between calls (backend fan-out completes concurrently). A STABLE sort keeps + // equal-keyed items in their incoming relative order, which costs nothing and + // makes the duplicate case reproducible within a snapshot. + sorted := slices.Clone(items) + slices.SortStableFunc(sorted, func(a, b T) int { + return strings.Compare(keyOf(a), keyOf(b)) + }) + + start := 0 + carried := 0 + cursorKey := "" + if cursor != "" { + lastKey, delivered, err := decodeModernCursor(kind, cursor) + if err != nil { + return nil, "", err + } + cursorKey = lastKey + // Binary search for the first key >= lastKey rather than scanning. The + // scan was O(n) per page, so a full walk was O(n^2); this makes it + // O(log n). A cursor whose key is no longer present lands at the next key + // above it rather than failing, which is the graceful degradation the + // "stable cursors" guidance is about -- BinarySearchFunc's insertion point + // gives that for free. + start, _ = slices.BinarySearchFunc(sorted, lastKey, func(item T, target string) int { + return strings.Compare(keyOf(item), target) + }) + // Then past the copies of lastKey already delivered, bounded by the run + // itself so a stale count cannot overshoot into the next key. + for range delivered { + if start >= len(sorted) || keyOf(sorted[start]) != lastKey { + break + } + start++ + carried++ + } + } + + end := min(start+modernPageSize, len(sorted)) + page := sorted[start:end] + + // A next cursor is emitted only when items actually remain. Emitting one on + // the final page would make a conforming client issue a guaranteed-empty extra + // round trip -- and emitting an empty string rather than omitting the field + // would make it loop on page one (see the file header). + if end == len(sorted) || len(page) == 0 { + return page, "", nil + } + + // Count the items sharing the last delivered key in this page, plus anything + // already counted for that same key on earlier pages, so a run spanning + // several pages accumulates instead of resetting. + lastKey := keyOf(page[len(page)-1]) + delivered := 0 + for i := len(page) - 1; i >= 0 && keyOf(page[i]) == lastKey; i-- { + delivered++ + } + if cursorKey == lastKey { + delivered += carried + } + + next, err := encodeModernCursor(kind, lastKey, delivered) + if err != nil { + return nil, "", err + } + return page, next, nil +} + +// modernRequestCursor extracts the optional cursor from a list request's params. +// +// An absent or explicitly null cursor reads as "first page" -- the field is +// optional, so neither is an error. A cursor that is PRESENT but not a string +// (e.g. `{"cursor": 123}`) is rejected with errInvalidModernCursor rather than +// silently read as "first page": the draft's Error Handling section requires +// invalid cursors to yield -32602, and schema.ts lists "Pagination: Invalid or +// expired cursor values" under the invalid-params class. +// +// Discarding that decode error was a real defect, not a style point: "a +// malformed params object was already rejected upstream" is false for this shape, +// because `{"cursor": 42}` is well-formed JSON and nothing upstream inspects the +// cursor's TYPE. The practical cost is a client with a serialization bug +// re-reading page one forever, each iteration paying a full unmetered backend +// fan-out. +// +// Decoding into json.RawMessage rather than straight into a string is deliberate, +// and it removes a whole hazard class rather than just relocating it. With a +// `Cursor string` field, encoding/json can populate the field AND return an error +// on the same call -- verified: `{"cursor":42,"cursor":"valid"}` yields +// ("valid", type error) -- so a discarded error there threw away a value already +// decoded, and an honoured one would reject a cursor it had successfully read. +// RawMessage captures bytes without type-checking, so there is no partial-decode +// state: duplicate keys resolve last-wins (Go's behaviour everywhere), and the +// single explicit string check below is the only place a type verdict is reached. +// +// Note the asymmetry with an empty-STRING cursor, which is deliberate and not the +// same case: `{"cursor": ""}` is "first page", matching go-sdk's server +// (`params.cursorPtr() == nil || *params.cursorPtr() == ""` selects the full +// sequence). vMCP never mints an empty cursor, so an empty string can only mean +// "no cursor supplied". +func modernRequestCursor(params json.RawMessage) (string, error) { + if len(params) == 0 { + return "", nil + } + var raw struct { + Cursor json.RawMessage `json:"cursor"` + } + if err := json.Unmarshal(params, &raw); err != nil || len(raw.Cursor) == 0 { + // A params envelope that does not decode at all, or carries no cursor + // field, is "first page"; a broken envelope is rejected upstream by the + // parser. + return "", nil + } + if string(raw.Cursor) == "null" { + return "", nil + } + var cursor string + if err := json.Unmarshal(raw.Cursor, &cursor); err != nil { + return "", errInvalidModernCursor + } + return cursor, nil +} diff --git a/pkg/vmcp/server/modern_pagination_test.go b/pkg/vmcp/server/modern_pagination_test.go new file mode 100644 index 0000000000..ca183c9676 --- /dev/null +++ b/pkg/vmcp/server/modern_pagination_test.go @@ -0,0 +1,725 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "math/rand/v2" + "net/http" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// identityKey keys a []string test corpus by its own element value. +func identityKey(s string) string { return s } + +// makeKeys builds n distinct keys that sort lexicographically in generation +// order, so a test can assert exact page contents without depending on sort +// subtleties. +// +// Corpora built ONLY from this helper are what hid the duplicate-key drop: it +// mints strictly unique keys, so no test built on it can exercise a collision. +// Anything asserting a completeness property must also be run against +// makeKeysWithDuplicates below. +func makeKeys(n int) []string { + keys := make([]string, 0, n) + for i := range n { + keys = append(keys, fmt.Sprintf("k%05d", i)) + } + return keys +} + +// makeKeysWithDuplicates builds an adversarial corpus of n items whose keys +// deliberately collide at positions that matter. The aggregator now resolves +// resource/template/prompt conflicts (#6060), but paginateModern is generic and +// deliberately does not depend on that caller invariant, so collisions remain +// the shape these completeness properties must be proven against. +// +// The collisions are placed by SORTED position, which is the subtle part and the +// reason a first attempt at this fixture was useless. paginateModern sorts before +// paging, so a distinctively-named duplicate key like "dup-boundary" does not stay +// where it was generated — it sorts among the "d" keys, nowhere near the page +// boundary, and the boundary case goes untested while looking tested. Collisions +// must therefore be introduced by overwriting a key with its own neighbour's +// value, which keeps the corpus sorted and pins the duplicate to the intended +// index. +// +// Two shapes are built: +// - a run spanning the page-1 boundary (indices pageSize-2 .. pageSize+2 share +// one key), so a page ends part-way through a run of equal keys. This is the +// shape that silently dropped an item. +// - a three-way tie well inside page 1. +func makeKeysWithDuplicates(n int) []string { + keys := make([]string, 0, n) + for i := range n { + keys = append(keys, fmt.Sprintf("k%05d", i)) + } + // Three-way tie inside page 1: indices 10, 11, 12 all take index 10's key. + for _, i := range []int{11, 12} { + if i < n { + keys[i] = keys[10] + } + } + // A run straddling the page-1 boundary: every index from pageSize-2 to + // pageSize+2 takes the key at pageSize-2, so the run begins on page 1 and + // continues onto page 2. + base := modernPageSize - 2 + if base >= 0 && base < n { + for i := base + 1; i <= modernPageSize+2 && i < n; i++ { + keys[i] = keys[base] + } + } + return keys +} + +// drainPages walks paginateModern to exhaustion, returning every key delivered +// and the number of pages it took. It fails if pagination does not terminate +// within a generous bound, so a cursor that fails to advance surfaces as a clear +// failure rather than an infinite loop. +func drainPages(t *testing.T, corpus []string) ([]string, int) { + t.Helper() + + // SHUFFLE before paging. This closes a blind spot that hid a second defect + // class: makeKeys returns pre-sorted keys, so deleting the sort from + // paginateModern entirely left this test green -- meaning the deterministic + // total order the whole keyset scheme rests on was unverified. The aggregator's + // real output is unordered (concurrent backend fan-out), so an unsorted input + // is also the realistic one. Deterministic seed keeps failures reproducible. + shuffled := slices.Clone(corpus) + rng := rand.New(rand.NewPCG(1, 2)) + rng.Shuffle(len(shuffled), func(i, j int) { shuffled[i], shuffled[j] = shuffled[j], shuffled[i] }) + corpus = shuffled + + // Non-nil so an empty corpus compares equal to makeKeys(0) rather than + // tripping on nil-vs-empty, which is not a distinction under test here. + seen := []string{} + cursor := "" + for pages := 1; pages <= 100; pages++ { + page, next, err := paginateModern(corpus, identityKey, cursorKindTools, cursor) + require.NoError(t, err) + seen = append(seen, page...) + if next == "" { + return seen, pages + } + cursor = next + } + t.Fatal("pagination did not terminate within 100 pages; cursor is not advancing") + return nil, 0 +} + +// TestPaginateModern_StaticSetDeliveredExactlyOnce is the central assertion of +// the pagination half: walking the cursors over an UNCHANGING set must yield the +// complete set, in order, with no gap and no duplicate. That is the property the +// whole scheme exists to provide, and the one a broken page boundary silently +// violates. +// +// The name says "static set" deliberately. The exactly-once guarantee holds only +// while the underlying set does not change mid-walk: with keyset pagination an +// item inserted BELOW the cursor after a page was served is never reported, since +// the scan only moves forward. That is inherent to the scheme (go-sdk's server +// behaves identically) and is explicitly out of scope per SEP-2567's "cursor +// stability and snapshot consistency" note -- conformant, not a bug. Do not widen +// this name back to an unqualified promise it cannot keep. +func TestPaginateModern_StaticSetDeliveredExactlyOnce(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + total int + wantPages int + }{ + {name: "empty set is a single empty page", total: 0, wantPages: 1}, + {name: "single item fits one page", total: 1, wantPages: 1}, + {name: "exactly one full page emits no cursor", total: modernPageSize, wantPages: 1}, + {name: "one item beyond a page spills to a second", total: modernPageSize + 1, wantPages: 2}, + {name: "the regression case: 1100 items", total: 1100, wantPages: 2}, + {name: "multiple full pages plus a tail", total: modernPageSize*2 + 7, wantPages: 3}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + corpus := makeKeys(tt.total) + seen, pages := drainPages(t, corpus) + + assert.Equal(t, tt.wantPages, pages, "unexpected page count") + // makeKeys is generated in sorted order, so the corpus doubles as the + // expected delivery order -- and because drainPages shuffles its input, + // matching it proves the paginator restored a deterministic total order + // rather than passing the input through. + assert.Equal(t, corpus, seen, "every item must be delivered exactly once, in sorted order") + + distinct := make(map[string]struct{}, len(seen)) + for _, k := range seen { + distinct[k] = struct{}{} + } + assert.Len(t, distinct, tt.total, "duplicates indicate overlapping pages") + }) + } +} + +// TestPaginateModern_DuplicateKeysAreNotDropped is the regression test for the +// blocking defect: with a bare "strictly above lastKey" scan, two items sharing a +// key at a page boundary caused BOTH to be skipped, so an item appeared on no +// page at all. A 1100-item resource corpus with one duplicated URI delivered 1099. +// +// When this defect shipped, key uniqueness held only for tools; resource URIs, +// resource-template URI templates and prompt names were plain-appended across +// backends with no conflict resolution. The aggregator now resolves those too +// (#6060), but the paginator is generic and must not depend on any caller's +// uniqueness invariant, and this test is what holds it to that. +// +// The corpus is adversarial by construction (see makeKeysWithDuplicates): a pair +// straddling the page-1 boundary, a run spanning it, and a three-way tie. Total +// count is asserted, because "delivered exactly n" is the property that fails +// when a collision eats an item. +func TestPaginateModern_DuplicateKeysAreNotDropped(t *testing.T) { + t.Parallel() + + for _, total := range []int{modernPageSize + 100, modernPageSize + 1, modernPageSize * 2} { + t.Run(fmt.Sprintf("run-spanning-boundary/total=%d", total), func(t *testing.T) { + t.Parallel() + + corpus := makeKeysWithDuplicates(total) + require.Len(t, corpus, total, "fixture must produce the requested item count") + + // Sanity-check the fixture actually collides; a non-colliding corpus + // would make this test silently vacuous, which is how the original + // defect survived. + counts := map[string]int{} + for _, k := range corpus { + counts[k]++ + } + boundaryKey := fmt.Sprintf("k%05d", modernPageSize-2) + assert.Greater(t, counts[boundaryKey], 1, + "fixture must duplicate a key spanning the page boundary") + assert.Equal(t, 3, counts[fmt.Sprintf("k%05d", 10)], + "fixture must include a three-way key tie") + + // The run must genuinely straddle the boundary: some copies on page 1, + // some on page 2. Otherwise this test cannot catch the drop. + sortedCorpus := slices.Clone(corpus) + slices.Sort(sortedCorpus) + onPage1 := 0 + for i := 0; i < modernPageSize && i < len(sortedCorpus); i++ { + if sortedCorpus[i] == boundaryKey { + onPage1++ + } + } + assert.Greater(t, onPage1, 0, "run must start on page 1") + assert.Less(t, onPage1, counts[boundaryKey], "run must continue onto page 2") + + seen, pages := drainPages(t, corpus) + + // The headline: nothing is lost. + assert.Len(t, seen, total, + "a duplicated key must not drop items; every item must be delivered") + assert.Greater(t, pages, 1, "corpus must actually paginate for this test to mean anything") + + // And the multiset matches exactly -- same keys, same multiplicities. + wantSorted := slices.Clone(corpus) + slices.Sort(wantSorted) + gotSorted := slices.Clone(seen) + slices.Sort(gotSorted) + assert.Equal(t, wantSorted, gotSorted, + "delivered multiset must equal the corpus, including duplicate multiplicities") + }) + } +} + +// TestPaginateModern_PageBoundaries pins the boundary behaviour the walk above +// exercises only indirectly: page size is capped, and a nextCursor appears if and +// only if items actually remain. +func TestPaginateModern_PageBoundaries(t *testing.T) { + t.Parallel() + + t.Run("first page is capped at the page size and carries a cursor", func(t *testing.T) { + t.Parallel() + page, next, err := paginateModern(makeKeys(1100), identityKey, cursorKindTools, "") + require.NoError(t, err) + assert.Len(t, page, modernPageSize) + assert.NotEmpty(t, next, "items remain, so a cursor is required") + }) + + t.Run("final page carries no cursor", func(t *testing.T) { + t.Parallel() + _, next, err := paginateModern(makeKeys(1100), identityKey, cursorKindTools, "") + require.NoError(t, err) + page, next2, err := paginateModern(makeKeys(1100), identityKey, cursorKindTools, next) + require.NoError(t, err) + assert.Len(t, page, 100) + assert.Empty(t, next2, "nothing remains, so emitting a cursor would force a wasted round trip") + }) + + t.Run("caller slice is not reordered", func(t *testing.T) { + t.Parallel() + // Deliberately unsorted, mirroring the aggregator's unordered fan-out. + corpus := []string{"zeta", "alpha", "mu"} + original := slices.Clone(corpus) + + page, _, err := paginateModern(corpus, identityKey, cursorKindTools, "") + require.NoError(t, err) + + assert.Equal(t, original, corpus, "paginateModern must not sort the caller's slice in place") + assert.Equal(t, []string{"alpha", "mu", "zeta"}, page, "the page itself must be sorted") + }) +} + +// TestPaginateModern_CursorValidation covers the cursor as untrusted input. A +// cursor is client-controlled, so every malformed or mismatched shape must be +// rejected rather than reinterpreted into a plausible-looking page. +func TestPaginateModern_CursorValidation(t *testing.T) { + t.Parallel() + + validOtherKind, err := encodeModernCursor(cursorKindPrompts, "k00001", 1) + require.NoError(t, err) + emptyKey, err := encodeModernCursor(cursorKindTools, "", 1) + require.NoError(t, err) + negativeDelivered, err := encodeModernCursor(cursorKindTools, "k00001", -1) + require.NoError(t, err) + + // A cursor over the length cap, rejected before any decode work happens. + oversized := base64.RawURLEncoding.EncodeToString( + []byte(`{"k":"tools","l":"` + strings.Repeat("A", modernCursorMaxLen*2) + `"}`)) + require.Greater(t, len(oversized), modernCursorMaxLen) + + tests := []struct { + name string + cursor string + }{ + {name: "not base64", cursor: "!!!not-base64!!!"}, + {name: "base64 of non-JSON", cursor: base64.RawURLEncoding.EncodeToString([]byte("plain"))}, + {name: "cursor minted for a different list verb", cursor: validOtherKind}, + {name: "cursor with an empty key", cursor: emptyKey}, + {name: "cursor with a negative delivered count", cursor: negativeDelivered}, + {name: "cursor longer than the length cap", cursor: oversized}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, _, err := paginateModern(makeKeys(10), identityKey, cursorKindTools, tt.cursor) + require.ErrorIs(t, err, errInvalidModernCursor) + }) + } + + t.Run("cursor naming a since-removed item resumes at the next key", func(t *testing.T) { + t.Parallel() + // "k00002" is gone from the corpus; the scan must continue from the next + // key above it rather than erroring or restarting. + cursor, err := encodeModernCursor(cursorKindTools, "k00002", 1) + require.NoError(t, err) + + page, _, err := paginateModern([]string{"k00001", "k00003", "k00004"}, identityKey, cursorKindTools, cursor) + require.NoError(t, err) + assert.Equal(t, []string{"k00003", "k00004"}, page) + }) + + t.Run("delivered count larger than the run resumes at the next key", func(t *testing.T) { + t.Parallel() + // Claims 9 copies of "dup" were delivered but only 2 exist: must not + // overshoot past "zed". + cursor, err := encodeModernCursor(cursorKindTools, "dup", 9) + require.NoError(t, err) + + page, _, err := paginateModern([]string{"dup", "dup", "zed"}, identityKey, cursorKindTools, cursor) + require.NoError(t, err) + assert.Equal(t, []string{"zed"}, page) + }) +} + +// TestModernCursor_RoundTrip checks the codec. It deliberately does NOT claim to +// test opacity: base64 provides encoding, not secrecy, and encodeModernCursor's +// own comment says it is "emphatically NOT encryption". Opacity is a rule imposed +// on CLIENTS by the spec, not a property the server can enforce, so asserting the +// key is absent from the token would be near-vacuous theatre. +func TestModernCursor_RoundTrip(t *testing.T) { + t.Parallel() + + encoded, err := encodeModernCursor(cursorKindResources, "file:///a/b.txt", 3) + require.NoError(t, err) + + gotKey, gotDelivered, err := decodeModernCursor(cursorKindResources, encoded) + require.NoError(t, err) + assert.Equal(t, "file:///a/b.txt", gotKey) + assert.Equal(t, 3, gotDelivered) +} + +// TestModernRequestCursor covers cursor extraction from request params, including +// the shapes that must read as "first page" and the shapes that must error. +func TestModernRequestCursor(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + params string + want string + wantErr bool + }{ + {name: "absent params", params: "", want: ""}, + {name: "empty object", params: `{}`, want: ""}, + {name: "null cursor", params: `{"cursor":null}`, want: ""}, + {name: "malformed params", params: `{`, want: ""}, + {name: "present cursor", params: `{"cursor":"abc"}`, want: "abc"}, + // An empty-string cursor is "first page", not an error: vMCP never mints + // one (nextCursor is omitted entirely at end-of-results), so it can only + // mean "no cursor supplied". Matches go-sdk's server. + {name: "empty string cursor is first page", params: `{"cursor":""}`, want: ""}, + // Present but not a string: must be an error, never a silent restart at + // page one. A client with a serialization bug needs a diagnosable -32602. + {name: "number cursor is rejected", params: `{"cursor":42}`, wantErr: true}, + {name: "object cursor is rejected", params: `{"cursor":{"k":"v"}}`, wantErr: true}, + {name: "array cursor is rejected", params: `{"cursor":["a"]}`, wantErr: true}, + {name: "bool cursor is rejected", params: `{"cursor":true}`, wantErr: true}, + // Duplicate JSON keys resolve last-wins, which is Go's behaviour + // everywhere and is why decoding into json.RawMessage matters: there is no + // partial-decode state to mishandle. Whichever value comes last is the one + // type-checked, so the same input is accepted or rejected purely on that + // value -- not on the presence of the duplicate. + {name: "duplicate cursor key takes the last value", params: `{"cursor":42,"cursor":"abc"}`, want: "abc"}, + {name: "duplicate cursor key rejects a bad last value", params: `{"cursor":"abc","cursor":42}`, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := modernRequestCursor(json.RawMessage(tt.params)) + if tt.wantErr { + require.ErrorIs(t, err, errInvalidModernCursor) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +// paginationFakeCore builds a fresh fake core holding `total` items of every +// kind. Each subtest gets its own instance rather than sharing one: modernFakeCore +// carries write-capable fields (checkCalled, callCalled, gotCtx), so a shared +// value would be a data race the moment a parallel subtest exercised a verb that +// writes them. +// +// Resource NAMES deliberately collide across all items while URIs stay distinct, +// and template names likewise. That makes the choice of pagination key +// observable: keying resources on Name instead of URI collapses the corpus and +// fails, where a fixture with unique names would have accepted either key. +func paginationFakeCore(total int) *modernFakeCore { + tools := make([]vmcp.Tool, 0, total) + resources := make([]vmcp.Resource, 0, total) + templates := make([]vmcp.ResourceTemplate, 0, total) + prompts := make([]vmcp.Prompt, 0, total) + for i := range total { + id := fmt.Sprintf("k%05d", i) + tools = append(tools, vmcp.Tool{Name: id, InputSchema: map[string]any{"type": "object"}}) + resources = append(resources, vmcp.Resource{URI: "file:///" + id, Name: "same-name"}) + templates = append(templates, vmcp.ResourceTemplate{URITemplate: "file:///{x}/" + id, Name: "same-name"}) + prompts = append(prompts, vmcp.Prompt{Name: id}) + } + return &modernFakeCore{tools: tools, resources: resources, templates: templates, prompts: prompts} +} + +// TestDispatchModernList_Pagination drives the four list verbs through +// dispatchModern to prove the wire result actually carries nextCursor and that a +// bad cursor is rejected as -32602 rather than -32603. +func TestDispatchModernList_Pagination(t *testing.T) { + t.Parallel() + + const total = modernPageSize + 5 + + tests := []struct { + method string + itemsField string + }{ + {method: "tools/list", itemsField: "tools"}, + {method: "resources/list", itemsField: "resources"}, + {method: "resources/templates/list", itemsField: "resourceTemplates"}, + {method: "prompts/list", itemsField: "prompts"}, + } + + for _, tt := range tests { + t.Run(tt.method, func(t *testing.T) { + t.Parallel() + fakeCore := paginationFakeCore(total) + + // Page 1: capped, and carries a cursor. + _, body := dispatchModernTest(t.Context(), t, fakeCore, false, &mcpparser.ParsedMCPRequest{ + ID: 1, Method: tt.method, Params: json.RawMessage(`{}`), + }) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "expected a result envelope, got %v", body) + items, ok := result[tt.itemsField].([]any) + require.True(t, ok, "expected %q in the result", tt.itemsField) + assert.Len(t, items, modernPageSize) + + cursor, ok := result["nextCursor"].(string) + require.True(t, ok, "a nextCursor must be present while items remain") + require.NotEmpty(t, cursor) + + // Page 2: the tail, and no further cursor. + params, err := json.Marshal(map[string]any{"cursor": cursor}) + require.NoError(t, err) + _, body2 := dispatchModernTest(t.Context(), t, fakeCore, false, &mcpparser.ParsedMCPRequest{ + ID: 2, Method: tt.method, Params: params, + }) + result2, ok := body2["result"].(map[string]any) + require.True(t, ok, "expected a result envelope, got %v", body2) + items2, ok := result2[tt.itemsField].([]any) + require.True(t, ok) + assert.Len(t, items2, 5, "the tail page must hold the remaining items") + assert.NotContains(t, result2, "nextCursor", "the final page must omit nextCursor entirely") + }) + } + + // Both an undecodable string cursor and a cursor of the wrong JSON type must + // surface as -32602. The non-string case is the one that used to fall through + // to "first page", silently serving page 1 to a client with a serialization + // bug instead of giving it a diagnosable error. + for _, tt := range []struct { + name string + params string + }{ + {name: "undecodable string cursor", params: `{"cursor":"!!!bogus!!!"}`}, + {name: "non-string cursor", params: `{"cursor":42}`}, + } { + t.Run("invalid cursor is invalid params, not internal error: "+tt.name, func(t *testing.T) { + t.Parallel() + + rec, body := dispatchModernTest(t.Context(), t, paginationFakeCore(10), false, &mcpparser.ParsedMCPRequest{ + ID: 3, Method: "tools/list", Params: json.RawMessage(tt.params), + }) + assert.Equal(t, http.StatusBadRequest, rec.Code) + + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope, got %v", body) + assert.Equal(t, float64(jsonRPCCodeInvalidParams), errObj["code"]) + assert.Equal(t, "invalid cursor", errObj["message"], + "the message must not describe the cursor encoding clients are told to treat as opaque") + + _, isResult := body["result"] + assert.False(t, isResult, "a rejected cursor must not also serve a page") + }) + } +} + +// TestDispatchModernList_CursorNotReplayableAcrossVerbs is the end-to-end twin of +// the kind check. Without it the kind field is self-consistently verified -- +// encode and decode both read the same site-local constant, so mutating one verb +// to mint another verb's kind passes every unit test while, in production, a +// prompts/list cursor would be accepted by tools/list. +// +// This replays a real cursor from one verb against every other verb through +// dispatchModern, which is the only way to catch that class of mistake. +func TestDispatchModernList_CursorNotReplayableAcrossVerbs(t *testing.T) { + t.Parallel() + + const total = modernPageSize + 5 + verbs := []string{"tools/list", "resources/list", "resources/templates/list", "prompts/list"} + + // Mint one real cursor per verb, through the dispatcher. + cursors := make(map[string]string, len(verbs)) + for _, verb := range verbs { + _, body := dispatchModernTest(t.Context(), t, paginationFakeCore(total), false, &mcpparser.ParsedMCPRequest{ + ID: 1, Method: verb, Params: json.RawMessage(`{}`), + }) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "%s: expected a result, got %v", verb, body) + cursor, ok := result["nextCursor"].(string) + require.True(t, ok, "%s: expected a nextCursor", verb) + cursors[verb] = cursor + } + + for _, minted := range verbs { + for _, replayed := range verbs { + if minted == replayed { + continue + } + t.Run(minted+" cursor rejected by "+replayed, func(t *testing.T) { + t.Parallel() + + params, err := json.Marshal(map[string]any{"cursor": cursors[minted]}) + require.NoError(t, err) + + rec, body := dispatchModernTest(t.Context(), t, paginationFakeCore(total), false, + &mcpparser.ParsedMCPRequest{ID: 9, Method: replayed, Params: params}) + + assert.Equal(t, http.StatusBadRequest, rec.Code) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected an error envelope, got %v", body) + assert.Equal(t, float64(jsonRPCCodeInvalidParams), errObj["code"]) + + _, isResult := body["result"] + assert.False(t, isResult, "a cross-verb cursor must not serve a page") + }) + } + } +} + +// TestModernList_EndOfResultsOmitsCursor pins the draft-only pagination +// amendment: end-of-results MUST be signalled by OMITTING nextCursor, never by +// emitting an empty string. Per the draft's Implementation Guidelines, "an empty +// string is a valid cursor and thus MUST NOT be treated as the end of results" -- +// so a client receiving `nextCursor: ""` would re-request with `cursor: ""`, +// which reads as "first page", and loop forever being served page one. +// +// This is a wire-serialization guarantee (the `omitempty` struct tags in +// modern_envelope.go), so it is asserted on the marshalled JSON rather than on +// the Go struct, where the zero value is indistinguishable from a dropped field. +func TestModernList_EndOfResultsOmitsCursor(t *testing.T) { + t.Parallel() + + for _, method := range []string{"tools/list", "resources/list", "resources/templates/list", "prompts/list"} { + t.Run(method, func(t *testing.T) { + t.Parallel() + + // A single short page: there is no next page, so no cursor may appear. + rec, body := dispatchModernTest(t.Context(), t, paginationFakeCore(1), false, + &mcpparser.ParsedMCPRequest{ID: 1, Method: method, Params: json.RawMessage(`{}`)}) + + result, ok := body["result"].(map[string]any) + require.True(t, ok, "expected a result envelope, got %v", body) + + assert.NotContains(t, result, "nextCursor", + "end-of-results must omit nextCursor entirely, not emit an empty string") + assert.NotContains(t, rec.Body.String(), `"nextCursor"`, + "the key must be absent from the wire bytes, not merely empty") + }) + } +} + +// TestModernPageSize_MatchesSDKDefaultBoundary pins the page-size value +// behaviourally rather than by reading the constant, so it fails on a real +// behaviour change instead of tautologically agreeing with whatever the constant +// says. +// +// modernPageSize is a hand-copied duplicate of go-sdk's DefaultPageSize (1000), +// which is what the Legacy/SDK path applies to the same split. mcpcompat does not +// re-export the constant, so nothing links the two: if go-sdk changes its +// default, Legacy follows and Modern does not. +// +// This pins the Modern half — a 1001-item corpus must produce exactly a 1000-item +// first page and a 1-item second. The Legacy half is NOT pinned here; doing so +// needs a harness that can hold a client on Legacy, which this package lacks +// (see the constant's doc comment). Asserting only the reachable half, and saying +// so, beats asserting a pairing this test cannot observe. +func TestModernPageSize_MatchesSDKDefaultBoundary(t *testing.T) { + t.Parallel() + + const sdkDefaultPageSize = 1000 + require.Equal(t, sdkDefaultPageSize, modernPageSize, + "modernPageSize must track go-sdk's DefaultPageSize; if go-sdk changed, update both and re-check Legacy") + + fakeCore := paginationFakeCore(sdkDefaultPageSize + 1) + + _, body := dispatchModernTest(t.Context(), t, fakeCore, false, &mcpparser.ParsedMCPRequest{ + ID: 1, Method: "tools/list", Params: json.RawMessage(`{}`), + }) + result, ok := body["result"].(map[string]any) + require.True(t, ok, "expected a result envelope, got %v", body) + items, ok := result["tools"].([]any) + require.True(t, ok) + assert.Len(t, items, sdkDefaultPageSize, "the first page must hold exactly the SDK default page size") + + cursor, ok := result["nextCursor"].(string) + require.True(t, ok, "one item remains, so a cursor is required") + + params, err := json.Marshal(map[string]any{"cursor": cursor}) + require.NoError(t, err) + _, body2 := dispatchModernTest(t.Context(), t, fakeCore, false, &mcpparser.ParsedMCPRequest{ + ID: 2, Method: "tools/list", Params: params, + }) + result2, ok := body2["result"].(map[string]any) + require.True(t, ok) + items2, ok := result2["tools"].([]any) + require.True(t, ok) + assert.Len(t, items2, 1, "the tail page must hold the single remaining item") +} + +// TestPaginateModern_DuplicateShapes covers the remaining two shapes of the +// duplicate-key defect, whose loss profile differs sharply and whose rule is +// easy to state wrongly. +// +// The general rule: every item whose key equals the LAST key delivered on a page +// was permanently dropped. So loss scaled with the size of the duplicate run +// straddling a page boundary, and a duplicate NOT at a boundary was harmless -- +// which is exactly why the bug was easy to miss. +// +// - single duplicate exactly at the boundary: 1 item lost +// - duplicate mid-page: 0 lost (harmless, but must stay correct) +// - run of 50 at the boundary: 50 lost +// - every key identical: an entire page lost, and page 2 came back empty +func TestPaginateModern_DuplicateShapes(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + corpus func() []string + }{ + { + name: "single duplicate exactly at the page boundary", + corpus: func() []string { + keys := makeKeys(modernPageSize + 100) + keys[modernPageSize] = keys[modernPageSize-1] + return keys + }, + }, + { + name: "duplicate mid-page is harmless but must stay correct", + corpus: func() []string { + keys := makeKeys(modernPageSize + 100) + keys[500] = keys[499] + return keys + }, + }, + { + name: "run of 50 copies straddling the boundary", + corpus: func() []string { + keys := makeKeys(modernPageSize + 100) + for i := modernPageSize - 25; i < modernPageSize+25; i++ { + keys[i] = keys[modernPageSize-25] + } + return keys + }, + }, + { + name: "every key identical", + corpus: func() []string { + keys := make([]string, modernPageSize+100) + for i := range keys { + keys[i] = "same" + } + return keys + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + corpus := tt.corpus() + seen, _ := drainPages(t, corpus) + + assert.Len(t, seen, len(corpus), + "every item must be delivered; a duplicated key must never drop one") + + wantSorted := slices.Clone(corpus) + slices.Sort(wantSorted) + gotSorted := slices.Clone(seen) + slices.Sort(gotSorted) + assert.Equal(t, wantSorted, gotSorted, + "delivered multiset must match the corpus, duplicate multiplicities included") + }) + } +} diff --git a/pkg/vmcp/server/modern_realbackend_integration_test.go b/pkg/vmcp/server/modern_realbackend_integration_test.go new file mode 100644 index 0000000000..6c6ef2e5ae --- /dev/null +++ b/pkg/vmcp/server/modern_realbackend_integration_test.go @@ -0,0 +1,576 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "maps" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --------------------------------------------------------------------------- +// Modern (2026-07-28) raw HTTP client helper +// --------------------------------------------------------------------------- + +// postModern sends a single Modern (2026-07-28) stateless JSON-RPC request to +// baseURL+"/mcp", hand-rolled per pkg/mcp/revision.go's wire contract rather +// than through go-sdk (importing it here would force an MVS bump): the +// MCP-Protocol-Version and Mcp-Method headers are mandatory on every Modern +// POST, Mcp-Name is set when non-empty (required only for tools/call, +// resources/read, prompts/get -- see nameRequiredMethods in revision.go), and +// _meta carries the reserved io.modelcontextprotocol/protocolVersion and +// clientCapabilities keys ClassifyRevision requires to admit the request as +// Modern in the first place. +// +// id == nil sends a notification: the "id" key is omitted from the body +// entirely (not set to JSON null), matching how parseMCPRequest distinguishes +// a call from a notification (dispatchModern, and jsonrpc2 before it, key off +// id being ABSENT, not nil). +// +// Returns the raw *http.Response (body re-readable: it is buffered and +// restored) alongside the JSON-RPC envelope decoded into a generic map, or a +// nil map for a body-less response (e.g. a notification's 202). +func postModern( + t *testing.T, baseURL, method string, params map[string]any, id any, mcpName string, +) (*http.Response, map[string]any) { + t.Helper() + + // Copy before mutating caller input (go-style rule): we inject _meta below, + // so clone the caller's params and any nested _meta rather than writing through. + params = maps.Clone(params) + if params == nil { + params = map[string]any{} + } + meta, _ := params["_meta"].(map[string]any) + meta = maps.Clone(meta) + if meta == nil { + meta = map[string]any{} + } + meta["io.modelcontextprotocol/protocolVersion"] = "2026-07-28" + // clientCapabilities is required on every Modern request; default to the + // empty declaration, but preserve a caller-supplied value so tests can + // exercise declared-capability behavior (see + // TestIntegration_Modern_RealBackend_MidCallCapabilityContract). + if _, ok := meta["io.modelcontextprotocol/clientCapabilities"]; !ok { + meta["io.modelcontextprotocol/clientCapabilities"] = map[string]any{} + } + params["_meta"] = meta + + body := map[string]any{ + "jsonrpc": "2.0", + "method": method, + "params": params, + } + if id != nil { + body["id"] = id + } + payload, err := json.Marshal(body) + require.NoError(t, err) + + // Bounded: several callers assert "resolves promptly, never hangs", and + // without a deadline a hang would only be caught by go test's package + // timeout. 30s is generous for an in-process round-trip even under CI + // -race load. + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + t.Cleanup(cancel) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/mcp", bytes.NewReader(payload)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", "2026-07-28") + req.Header.Set("Mcp-Method", method) + if mcpName != "" { + req.Header.Set("Mcp-Name", mcpName) + } + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + + respBody, err := io.ReadAll(resp.Body) + require.NoError(t, err) + resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + // Best-effort decode: a request that never reaches the Modern dispatcher + // (e.g. a malformed request rejected before dispatch) can come back as a + // plain-text error rather than JSON. Callers that expect a JSON-RPC + // envelope assert on decoded's fields directly, which fails informatively + // (nil map) if decoding didn't happen. + var decoded map[string]any + if len(respBody) > 0 { + _ = json.Unmarshal(respBody, &decoded) + } + return resp, decoded +} + +// --------------------------------------------------------------------------- +// Integration tests -- Modern (2026-07-28) stateless dispatch, real backend +// --------------------------------------------------------------------------- + +// TestIntegration_Modern_RealBackend_ToolCall verifies the load-bearing +// end-to-end round-trip: a Modern tools/call request travels through the real +// middleware chain (parsing/classification/dispatch), a real core, and a real +// backend, and comes back as a Modern envelope with no session established. +func TestIntegration_Modern_RealBackend_ToolCall(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "hello modern"}, + }, 1, "echo") + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + assert.Empty(t, resp.Header.Get("Mcp-Session-Id"), "Modern responses must never carry a session ID") + + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.Equal(t, "complete", result["resultType"]) + content, ok := result["content"].([]any) + require.True(t, ok && len(content) == 1) + first := content[0].(map[string]any) + assert.Equal(t, "text", first["type"]) + assert.Equal(t, "hello modern", first["text"]) + // IsError has an omitempty JSON tag, so a successful call omits the key + // entirely rather than marshaling it as false. + assert.NotEqual(t, true, result["isError"], "tool call must not be marked as an error") +} + +// TestIntegration_Modern_RealBackend_ToolsList verifies tools/list against the +// real backend's discovered tool set, with the Modern cacheability envelope. +func TestIntegration_Modern_RealBackend_ToolsList(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "tools/list", nil, 1, "") + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.Equal(t, "complete", result["resultType"]) + assert.Equal(t, "private", result["cacheScope"]) + _, hasTTL := result["ttlMs"] + assert.True(t, hasTTL, "ttlMs must be present even when zero") + + tools, ok := result["tools"].([]any) + require.True(t, ok && len(tools) == 1, "expected exactly the echo tool: %+v", result) + assert.Equal(t, "echo", tools[0].(map[string]any)["name"]) +} + +// TestIntegration_Modern_RealBackend_Discover verifies server/discover reports +// capability presence derived from the real backend's actual tool set: tools +// and completions present (the echo backend has a tool; completions is +// unconditional), resources and prompts absent (the echo backend exposes +// neither). +func TestIntegration_Modern_RealBackend_Discover(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "server/discover", nil, 1, "") + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.Equal(t, "private", result["cacheScope"]) + + caps, ok := result["capabilities"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + _, hasTools := caps["tools"] + _, hasCompletions := caps["completions"] + _, hasResources := caps["resources"] + _, hasPrompts := caps["prompts"] + assert.True(t, hasTools, "echo backend has a tool") + assert.True(t, hasCompletions, "completions is advertised unconditionally") + assert.False(t, hasResources, "echo backend exposes no resources") + assert.False(t, hasPrompts, "echo backend exposes no prompts") +} + +// TestIntegration_Modern_RealBackend_Complete verifies completion/complete +// routes to the core rather than 404ing as an unknown method. The echo +// backend has no prompts, so the referenced name is unroutable; core.Complete +// treats that leniently (empty candidates, not an error -- see +// coreVMCP.Complete), so this asserts a clean 200 completion object rather +// than a protocol-level rejection. +func TestIntegration_Modern_RealBackend_Complete(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "completion/complete", map[string]any{ + "ref": map[string]any{"type": "ref/prompt", "name": "nonexistent"}, + "argument": map[string]any{"name": "a", "value": ""}, + }, 1, "") + defer resp.Body.Close() + + require.NotEqual(t, http.StatusNotFound, resp.StatusCode, + "completion/complete must not be treated as an unknown method: decoded: %+v", decoded) + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + completion, ok := result["completion"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.Equal(t, []any{}, completion["values"], "unroutable ref yields empty candidates, not an error") +} + +// TestIntegration_Modern_RealBackend_Ping verifies ping returns a bare +// {"jsonrpc":"2.0","id":..,"result":{}} -- no resultType, no _meta -- per +// dispatchModern's documented deliberate bypass of the envelope builders for +// this method. +func TestIntegration_Modern_RealBackend_Ping(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "ping", nil, 7, "") + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + assert.Equal(t, map[string]any{ + "jsonrpc": "2.0", + "id": float64(7), + "result": map[string]any{}, + }, decoded) +} + +// TestIntegration_Modern_RealBackend_Notification verifies a Modern request +// with no "id" (a notification) is acknowledged with 202 and no body, per +// dispatchModern's ID-nil check -- which runs before any method dispatch. +func TestIntegration_Modern_RealBackend_Notification(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "tools/list", nil, nil, "") + defer resp.Body.Close() + + assert.Equal(t, http.StatusAccepted, resp.StatusCode) + assert.Nil(t, decoded) + leftover, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Empty(t, leftover, "a notification response must carry no body") +} + +// TestIntegration_Modern_RealBackend_UnknownMethod verifies a syntactically +// well-formed Modern request naming a method dispatchModern does not +// recognize is rejected with 404 + JSON-RPC -32601, per the draft spec's +// MUST-404-unimplemented-method rule (writeModernError). +func TestIntegration_Modern_RealBackend_UnknownMethod(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "resources/subscribe", map[string]any{"uri": "file:///x"}, 1, "") + defer resp.Body.Close() + + require.Equal(t, http.StatusNotFound, resp.StatusCode, "decoded: %+v", decoded) + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.EqualValues(t, -32601, errObj["code"]) +} + +// TestIntegration_Modern_RealBackend_ElicitingToolFailsCleanly pins the +// Modern-client CONTRACT for a backend tool that issues a mid-call +// server-initiated request (elicitation/create, sampling/createMessage): the +// call MUST resolve promptly (postModern is deadline-bounded) to an explicit +// JSON-RPC error — never a hang, never a fabricated success, and never a +// resultType "input_required" envelope, which this dispatcher does not emit +// (client-polled multi-round retrieval, SEP-2322, is unimplemented; +// modernResultTypeComplete is the only resultType modern_envelope.go builds). +// +// Deliberately NOT pinned: the specific error code. Today it is -32603, but +// the spec MUSTs -32021 MissingRequiredClientCapability for the undeclared +// case (at HTTP 400, which go-sdk's client escalates to permanent session +// death — the follow-up therefore plans -32021 at HTTP 200 as a documented +// deviation; see the "unavailable to Modern clients" section of +// docs/arch/10-virtual-mcp-architecture.md). Pinning -32603 here would make +// that spec-correcting follow-up look like a regression. TODO(follow-up): +// once the two-path capability-error contract lands, assert -32021 + +// data.requiredCapabilities for the undeclared case. +// +// This is the honest-unsupported contract for the surface the 2026-07-28 +// revision removed: server-initiated requests need a live session, Modern +// requests have none by design, and SEP-2577 additionally deprecates sampling +// outright. The Legacy-session behavior for the same backend tools is covered +// by the TestForwarding_* fixtures, whose downstream clients pin Legacy +// explicitly (see legacyPinningRoundTripper). +func TestIntegration_Modern_RealBackend_ElicitingToolFailsCleanly(t *testing.T) { + t.Parallel() + + backendURL := startForwardingBackend(t) + ts := newRealTestServer(t, backendURL) + + tests := []struct { + name string + tool string + }{ + {name: "elicitation", tool: fwdElicitTool}, + {name: "sampling", tool: fwdSampleTool}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + resp, decoded := postModern(t, ts.URL, "tools/call", map[string]any{"name": tc.tool}, 1, tc.tool) + defer resp.Body.Close() + + // Both today's -32603 and the follow-up's -32021 (documented HTTP + // 200 deviation) ride HTTP 200, so this holds across the fix. + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + require.NotContains(t, decoded, "result", + "must not fabricate a success or an input_required envelope: %+v", decoded) + _, ok := decoded["error"].(map[string]any) + require.True(t, ok, "the failure must be an explicit JSON-RPC error: %+v", decoded) + + // Nothing about the MESSAGE is asserted, deliberately. The refused + // request's name (elicitation/create etc.) reaches it only via + // go-sdk's "calling %q" error wrapping — a dependency's string + // (testing.md, test scope) — and the message also carries the + // pre-existing backend-workload-ID leak, which the follow-up + // capability-error contract fixes with vMCP-owned messages. Both + // message properties (names our string, no backend ID) are asserted + // THERE, where vMCP owns the text. + }) + } +} + +// TestIntegration_Modern_RealBackend_ProgressDropped pins today's Modern +// behavior for a backend tool that emits notifications/progress mid-call: the +// notification is silently dropped and the call still completes promptly with +// its result. There is nowhere for it to go — a Modern response is a single +// JSON body (writeModernResult), and request-scoped notifications would ride +// a per-request SSE response stream (SEP-2260) that the single-shot +// dispatcher does not produce. That is a vMCP streaming-dispatch gap, NOT a +// spec absence — this test converts that prose claim (see +// TestForwarding_Progress_RealBackend's pin comment and the arch doc) into an +// executable one, and MUST be revisited when streaming Modern dispatch is +// implemented: at that point the progress notification becomes deliverable +// and this test's premise changes. +// +// Promptness matters here: on a Legacy session the same fixture relays the +// notification; on Modern, a client waiting for it would hang to its own +// deadline. postModern's bounded context is the hang guard. +func TestIntegration_Modern_RealBackend_ProgressDropped(t *testing.T) { + t.Parallel() + + backendURL := startForwardingBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "tools/call", map[string]any{ + "name": fwdProgressTool, + // A progressToken is exactly what would solicit progress delivery if a + // channel existed; its presence makes "dropped" the strongest claim. + "_meta": map[string]any{"progressToken": fwdProgressToken}, + }, 1, fwdProgressTool) + defer resp.Body.Close() + + // The call completes with the tool's result; the mid-call progress + // notification had no channel and was dropped. A single JSON body can + // carry nothing else, so completing at all IS the drop assertion. + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + result, ok := decoded["result"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.Equal(t, "complete", result["resultType"]) + content, ok := result["content"].([]any) + require.True(t, ok && len(content) == 1, "decoded: %+v", decoded) + assert.Equal(t, "done", content[0].(map[string]any)["text"]) +} + +// TestIntegration_Modern_RealBackend_LoggingContract pins today's Modern +// contract for logging/setLevel, which the 2026-07-28 revision removed (the +// per-request logLevel _meta key replaces it): +// +// - A WELL-FORMED Modern logging/setLevel is an unknown method to +// dispatchModern: 404 + -32601, matching go-sdk's own server ("method +// removed in the new protocol"). +// - The request go-sdk v1.7.x ACTUALLY sends is malformed — its +// SetLoggingLevel omits the per-request _meta injection, producing a Modern +// header with no _meta protocolVersion — and vMCP's classifier CORRECTLY +// rejects that shape with 400 + -32020 (HeaderMismatch). Loosening this to +// accept the malformed request would be worse than the failing upstream +// call. This malformed shape is a permanent go-sdk v1.7.x fixture: +// modelcontextprotocol/go-sdk#1116 was closed wont-fix-by-design (the RPC +// is removed on Modern), so -32020 is the standing contract for any caller +// still using the removed RPC on a Modern session, not a temporary bug +// guard. +// +// See TestForwarding_Logging_RealBackend's pin comment for the full +// three-cause disposition (go-sdk's wont-fix SetLoggingLevel, method +// removal/streaming gap, SEP-2577 deprecation). +func TestIntegration_Modern_RealBackend_LoggingContract(t *testing.T) { + t.Parallel() + + backendURL := startForwardingBackend(t) + ts := newRealTestServer(t, backendURL) + + t.Run("well-formed setLevel is method-not-found", func(t *testing.T) { + t.Parallel() + resp, decoded := postModern(t, ts.URL, "logging/setLevel", map[string]any{"level": "debug"}, 1, "") + defer resp.Body.Close() + + require.Equal(t, http.StatusNotFound, resp.StatusCode, "decoded: %+v", decoded) + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.EqualValues(t, -32601, errObj["code"]) + }) + + t.Run("go-sdk's malformed setLevel is rejected -32020", func(t *testing.T) { + t.Parallel() + // Replicate the upstream bug's wire shape by hand: Modern header, but + // params carrying NO _meta protocolVersion (postModern would inject it, + // so build the request raw). + payload, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "logging/setLevel", + "params": map[string]any{"level": "debug"}, + }) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ts.URL+"/mcp", bytes.NewReader(payload)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", "2026-07-28") + req.Header.Set("Mcp-Method", "logging/setLevel") + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + require.Equal(t, http.StatusBadRequest, resp.StatusCode, "body: %s", body) + var decoded map[string]any + require.NoError(t, json.Unmarshal(body, &decoded), "body: %s", body) + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "body: %s", body) + assert.EqualValues(t, -32020, errObj["code"], + "the malformed request must be rejected as a header/_meta mismatch") + }) +} + +// TestIntegration_Modern_RealBackend_MalformedArguments verifies a +// syntactically valid tools/call whose "arguments" is present but not a JSON +// object is rejected with 400 + JSON-RPC -32602, per hasNonObjectArguments' +// pre-dispatch shape check. +func TestIntegration_Modern_RealBackend_MalformedArguments(t *testing.T) { + t.Parallel() + + backendURL := startRealMCPBackend(t) + ts := newRealTestServer(t, backendURL) + + resp, decoded := postModern(t, ts.URL, "tools/call", map[string]any{ + "name": "echo", + "arguments": "not-an-object", + }, 1, "echo") + defer resp.Body.Close() + + require.Equal(t, http.StatusBadRequest, resp.StatusCode, "decoded: %+v", decoded) + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.EqualValues(t, -32602, errObj["code"]) +} + +// TestIntegration_Modern_RealBackend_MidCallCapabilityContract pins the +// two-path error contract for a backend tool that demands mid-call +// elicitation/sampling during a Modern client's tools/call +// (writeModernCallFailure): +// +// - capability NOT declared in _meta clientCapabilities: the draft schema's +// MissingRequiredClientCapabilityError — code -32021 with +// data.requiredCapabilities typed as a ClientCapabilities object — served +// at HTTP 200, a deliberate, documented deviation from the spec-mandated +// 400 (go-sdk treats a non-transient 4xx as permanent connection death; +// see writeModernMissingCapability and go-sdk#1117). +// - capability DECLARED: -32021 would wrongly blame the client, and the +// 2026-07-28 vocabulary has no "operation not supported" code, so it is +// an explicit -32603 whose vMCP-owned message names the real cause: +// honouring a declared capability mid-call is multi-round retrieval +// (SEP-2322), which this server does not implement. +// +// Both paths emit vMCP-crafted messages, so — unlike the raw backend error +// chain they replace — they are asserted on ToolHive-owned strings, and the +// backend workload ID must NOT leak into them. +func TestIntegration_Modern_RealBackend_MidCallCapabilityContract(t *testing.T) { + t.Parallel() + + backendURL := startForwardingBackend(t) + ts := newRealTestServer(t, backendURL) + + tests := []struct { + name string + tool string + capability string // the capability the backend's tool demands mid-call + declared bool // whether the request's _meta declares it + }{ + {name: "elicitation undeclared", tool: fwdElicitTool, capability: "elicitation", declared: false}, + {name: "sampling undeclared", tool: fwdSampleTool, capability: "sampling", declared: false}, + {name: "elicitation declared", tool: fwdElicitTool, capability: "elicitation", declared: true}, + {name: "sampling declared", tool: fwdSampleTool, capability: "sampling", declared: true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + params := map[string]any{"name": tc.tool} + if tc.declared { + params["_meta"] = map[string]any{ + "io.modelcontextprotocol/clientCapabilities": map[string]any{ + tc.capability: map[string]any{}, + }, + } + } + resp, decoded := postModern(t, ts.URL, "tools/call", params, 1, tc.tool) + defer resp.Body.Close() + + // Both paths ride HTTP 200: -32603 per writeModernError's mapping, + // -32021 per writeModernMissingCapability's documented deviation. + require.Equal(t, http.StatusOK, resp.StatusCode, "decoded: %+v", decoded) + require.NotContains(t, decoded, "result", + "must not fabricate a success or an input_required envelope: %+v", decoded) + errObj, ok := decoded["error"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + msg, _ := errObj["message"].(string) + + // vMCP owns both messages: they must name the capability and the + // gateway limitation (SEP-2322), and must not leak the backend + // workload ID the raw error chain used to carry. + assert.Contains(t, msg, tc.capability) + assert.Contains(t, msg, "SEP-2322", + "the message must name multi-round retrieval as the gateway limitation") + assert.NotContains(t, msg, "real-backend", + "the crafted message must not leak the backend workload ID") + + if !tc.declared { + assert.EqualValues(t, -32021, errObj["code"]) + data, ok := errObj["data"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + required, ok := data["requiredCapabilities"].(map[string]any) + require.True(t, ok, "decoded: %+v", decoded) + assert.Contains(t, required, tc.capability, + "requiredCapabilities must name the capability the server needed") + return + } + assert.EqualValues(t, -32603, errObj["code"]) + }) + } +} diff --git a/pkg/vmcp/server/modern_subscriptions.go b/pkg/vmcp/server/modern_subscriptions.go new file mode 100644 index 0000000000..cab8d61117 --- /dev/null +++ b/pkg/vmcp/server/modern_subscriptions.go @@ -0,0 +1,378 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" +) + +// This file implements MCP 2026-07-28 ("Modern") subscriptions/listen, the +// revision's ONLY server->client push channel. Modern removed the standalone +// HTTP GET stream, so a client that wants unsolicited notifications opens one +// of these instead. +// +// WHY IT EXISTS EVEN THOUGH vMCP HAS NOTHING TO PUSH. Without a handler for +// this method a go-sdk v1.7 client cannot talk to vMCP at all: its Connect is +// Modern-first, and once server/discover succeeds it opens a listen stream +// whenever any list-changed handler is registered -- which mcpcompat's +// Initialize does unconditionally (installNotificationHandlers). A -32601 there +// fails Connect outright and tears the session down. So this handler is what +// stands between a Modern client and a usable session; it is not a facility +// nobody reaches. +// +// WHAT IT HONORS: nothing, today, and it says so on the wire. The honored set is +// computed by intersecting the client's requested notification types against +// newModernCapabilities -- the same builder server/discover publishes -- every +// push-related flag of which is deliberately false (see its doc comment for why, +// per flag). An empty honored set is the spec's own way to say "I support none +// of these", not a stub: SEP-2575 specifies the acknowledgement as carrying +// "the subset of the requested notifications the server has agreed to honor", +// and requires unsupported types to be omitted from it. go-sdk's reference +// server does exactly this, filtering the requested set through its own +// capabilities (allowedSubscriptions, server.go:1254+ in go-sdk@v1.7.0-pre.3). +// +// The distinction that matters: this is a truthful NEGATIVE declaration, not an +// empty success. It never reports a subscription as established and then +// silently drops notifications -- it enumerates, in the acknowledgement, that +// nothing is honored. Real delivery is tracked separately (#5743) and requires +// vMCP to first start ADVERTISING a push capability in newModernCapabilities; +// until then there is deliberately nothing for this stream to carry. +// +// Known limitation worth stating plainly: go-sdk's client ignores the +// acknowledgement (callSubscriptionsAckHandler returns (nil, nil), +// client.go:1457-1459), so against that client this honest declaration is not +// observable in behavior. It is correct on the wire regardless, and the +// alternative -- holding a stream open forever delivering nothing -- is strictly +// worse. + +// Modern subscription method and notification names, plus the reserved _meta +// key deliveries are tagged with. +// +// All three are NORMATIVE in schema/draft/schema.ts -- do not weaken them to +// "SDK convention". The relevant declarations: +// +// - SubscriptionsListenResultMeta requires (non-optional) +// "io.modelcontextprotocol/subscriptionId", whose value is "the JSON-RPC ID +// of the subscriptions/listen request that opened the stream (and equals +// this response's id)" -- so request-id-as-subscription-identity is +// specified, not invented. +// - SubscriptionsAcknowledgedNotification: "This notification MUST be the +// first message the server sends carrying the subscription's ID ... The +// server MUST NOT send any notification on the subscription before +// acknowledging it." +// - NotificationMetaObject: "The server MUST include this key on every +// notification delivered via a subscriptions/listen stream." That third one +// binds a DELIVERY implementation rather than this handler (which delivers +// nothing); it is restated at the #5743 hand-off below so it is not lost. +// +// go-sdk agrees and can be read as a cross-check (MetaKeySubscriptionID at +// protocol.go:2377-2379; server.go:1187-1250), but the schema is the authority. +// An earlier version of this comment attributed these to the SDK because the +// SEP-2575 prose does not restate them -- the schema does, and it governs. +const ( + methodSubscriptionsListen = "subscriptions/listen" + notificationSubscriptionsAcked = "notifications/subscriptions/acknowledged" + modernSubscriptionIDKey = "io.modelcontextprotocol/subscriptionId" +) + +// notificationSubscriptions is the per-type, per-URI opt-in set carried by both +// subscriptions/listen params and the acknowledgement. It mirrors go-sdk's +// NotificationSubscriptions field-for-field (protocol.go:2070-2082); mcpcompat +// does not re-export the type, so like the rest of the Modern envelope this is a +// hand-rolled parallel serializer. +// +// These four are the ENTIRE subscribable universe under SEP-2575. Progress, +// logging, elicitation, and sampling are structurally absent -- they are not +// notification types a client can opt into here, which is why this channel +// cannot carry them however it is implemented. +type notificationSubscriptions struct { + ToolsListChanged bool `json:"toolsListChanged,omitempty"` + PromptsListChanged bool `json:"promptsListChanged,omitempty"` + ResourcesListChanged bool `json:"resourcesListChanged,omitempty"` + ResourceSubscriptions []string `json:"resourceSubscriptions,omitempty"` +} + +// isEmpty reports whether nothing at all is subscribed. +func (n notificationSubscriptions) isEmpty() bool { + return !n.ToolsListChanged && !n.PromptsListChanged && + !n.ResourcesListChanged && len(n.ResourceSubscriptions) == 0 +} + +// subscriptionsListenParams is the decoded subscriptions/listen request params. +// Notifications is a pointer so an absent field is distinguishable from an +// explicit empty object: the spec makes it REQUIRED, and go-sdk rejects a nil +// one with invalid-params (server.go:1193-1195), so vMCP must too. +type subscriptionsListenParams struct { + Notifications *notificationSubscriptions `json:"notifications"` +} + +// subscriptionsAcknowledgedParams is the acknowledgement notification's params: +// the honored subset, plus the subscription id this stream is keyed by. +type subscriptionsAcknowledgedParams struct { + Notifications notificationSubscriptions `json:"notifications"` + Meta map[string]any `json:"_meta"` +} + +// subscriptionsListenResult is the response to subscriptions/listen, signalling +// that the subscription ended gracefully. Mirrors go-sdk's +// SubscriptionsListenResult (protocol.go:2110-2118): resultType "complete" plus +// a subscriptionId-tagged _meta. +type subscriptionsListenResult struct { + ResultType string `json:"resultType"` + Meta map[string]any `json:"_meta"` +} + +// dispatchModernSubscriptionsListen serves subscriptions/listen. +// +// It is UNGATED, in the same bucket as the four list verbs and server/discover: +// there is no Check* for it because it performs no write, reaches no backend, +// and discloses nothing beyond what server/discover already does -- the honored +// set it returns is derived from core.Discover, which is itself +// admission-filtered per identity. Should this ever begin honoring a +// subscription (i.e. actually delivering resource updates), revisit that: a +// per-URI resource subscription that delivers content WOULD need +// CheckResourceRead-equivalent gating per URI. +// +// The response is an SSE stream, not a single JSON body, because the protocol +// requires two messages on it: the acknowledgement notification first, then the +// result. go-sdk forces SSE for this method for the same reason +// (streamable.go:1645-1650). The stream is closed as soon as both are written +// because the honored set is empty and there is consequently nothing to keep it +// open for -- matching go-sdk's reference server, which blocks on the request +// context only when it has agreed to honor at least one subscription +// (server.go:1246-1250). +func (s *Server) dispatchModernSubscriptionsListen( + ctx context.Context, w http.ResponseWriter, parsed *mcpparser.ParsedMCPRequest, identity *auth.Identity, +) { + var params subscriptionsListenParams + if err := json.Unmarshal(parsed.Params, ¶ms); err != nil || params.Notifications == nil { + writeModernError(w, parsed.ID, jsonRPCCodeInvalidParams, + "invalid subscriptions/listen params: missing required 'notifications' field") + return + } + + // Resolve what this identity may reach, then shape it through the same + // capability builder server/discover publishes. + // + // The pre-check first: core.Discover is an un-rate-limited backend fan-out + // (ratelimit/decorator.go meters CallTool only), and while no push capability + // is advertised its result cannot change the answer -- every possible + // DiscoverCapabilities combination yields the same empty honored set, because + // the has* booleans only decide whether a capability POINTER is non-nil, never + // the listChanged/subscribe flags inside it. Calling it anyway meant N no-op + // requests cost N fan-outs, and that got worse once list verbs began paging. + // So probe the ceiling -- everything present -- and skip the fan-out when even + // that honors nothing. This is a pure optimisation: it cannot change the + // response, only what it costs. + ceiling := honoredSubscriptions(*params.Notifications, newModernCapabilities(true, true, true, true)) + honored := ceiling + if !ceiling.isEmpty() { + caps, err := s.core.Discover(ctx, identity) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + advertised := newModernCapabilities( + caps.HasTools, caps.HasResources, caps.HasResourceTemplates, caps.HasPrompts) + honored = honoredSubscriptions(*params.Notifications, advertised) + } + + // The subscription id is the listen request's own JSON-RPC id. Modern has no + // sessions, so this -- not an Mcp-Session-Id -- is what a delivery + // implementation would key streams by. + subscriptionMeta := map[string]any{modernSubscriptionIDKey: parsed.ID} + + // Build first: a marshal failure here is still reportable as -32603, because + // nothing has been written yet. + frames, err := buildModernListenFrames(parsed.ID, honored, subscriptionMeta) + if err != nil { + writeModernListError(ctx, w, parsed.ID, parsed.Method, err) + return + } + + if err := writeModernListenStream(w, frames); err != nil { + // Only reachable mid-stream: the frames were built successfully and the + // status line is already sent, so there is no way to turn this into a + // JSON-RPC error. Log it and let the connection drop. + slog.DebugContext(ctx, "vmcp modern dispatch: subscriptions/listen stream ended early", + "error", err) + return + } + + if !honored.isEmpty() { + // Unreachable while newModernCapabilities advertises no push capability, + // and deliberately not built out here: keeping a stream open is only half + // of honoring a subscription, and the delivery half does not exist yet. + // Whoever flips a capability flag owns adding both (#5743) -- this WARN + // exists so that, if a flag is flipped without it, the gap is loud in the + // logs instead of presenting as a silently idle client subscription. + // + // Hand-off note for that work, since it is a MUST and easy to miss: per + // schema.ts's NotificationMetaObject, the server MUST include + // io.modelcontextprotocol/subscriptionId on EVERY notification delivered + // via a subscriptions/listen stream -- not just on the acknowledgement + // this handler already tags. A delivery implementation that pushes a bare + // notifications/tools/list_changed onto the stream is non-conformant even + // though the subscription itself was acknowledged correctly. + // + // Logged by COUNT, never by URI: resourceSubscriptions is client-supplied + // and only length-capped, so echoing the array into the log sink on an + // unmetered verb would let a caller choose how much it writes per request. + slog.WarnContext(ctx, "vmcp modern dispatch: subscriptions/listen honored a subscription "+ + "but vMCP has no delivery mechanism; notifications will not arrive", + "subscription_id", parsed.ID, + "tools_list_changed", honored.ToolsListChanged, + "prompts_list_changed", honored.PromptsListChanged, + "resources_list_changed", honored.ResourcesListChanged, + "resource_subscription_count", len(honored.ResourceSubscriptions)) + } +} + +// honoredSubscriptions intersects a client's requested notification set with +// what vMCP advertises, returning only the types it can actually honor. +// +// This is the whole of the per-type AND per-URI opt-in filtering SEP-2575 +// mandates ("the server MUST NOT send notification types the client has not +// explicitly requested"): a type survives only if the client asked for it AND +// the matching capability is advertised. Because it takes capabilities as a +// parameter rather than reading a package-level constant, it needs no edit when +// a capability flips -- and it is independently testable across combinations +// the live advertisement cannot currently produce. +// +// resourceSubscriptions is filtered as a unit rather than per-URI: the +// Subscribe capability is server-wide, so it either gates all requested URIs or +// none. Per-URI admission (dropping URIs this identity may not read) belongs +// with delivery, since it is only observable once updates actually flow. +func honoredSubscriptions( + want notificationSubscriptions, advertised mcp.ServerCapabilities, +) notificationSubscriptions { + var honored notificationSubscriptions + if want.ToolsListChanged && advertised.Tools != nil && advertised.Tools.ListChanged { + honored.ToolsListChanged = true + } + if want.PromptsListChanged && advertised.Prompts != nil && advertised.Prompts.ListChanged { + honored.PromptsListChanged = true + } + if want.ResourcesListChanged && advertised.Resources != nil && advertised.Resources.ListChanged { + honored.ResourcesListChanged = true + } + if len(want.ResourceSubscriptions) > 0 && advertised.Resources != nil && advertised.Resources.Subscribe { + honored.ResourceSubscriptions = want.ResourceSubscriptions + } + return honored +} + +// buildModernListenFrames marshals the two frames the listen response carries: +// the mandatory initial acknowledgement notification, then the terminating +// result. +// +// It is separate from writeModernListenStream so that everything which can fail +// BEFORE any byte is written stays on a path that can still produce a proper +// JSON-RPC error. Folding these into the writer meant a marshal failure returned +// after headers were nominally "committed" when in fact nothing had been written +// at all, so the client received a bare HTTP 200 with an empty body for a request +// carrying an id. Now the caller maps a build failure to -32603, matching +// writeModernEnvelope's own build-then-write ordering, and only a genuine +// mid-stream write failure is unreportable. +func buildModernListenFrames( + id any, honored notificationSubscriptions, subscriptionMeta map[string]any, +) ([][]byte, error) { + ack, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "method": notificationSubscriptionsAcked, + "params": subscriptionsAcknowledgedParams{ + Notifications: honored, + Meta: subscriptionMeta, + }, + }) + if err != nil { + return nil, fmt.Errorf("marshalling subscriptions acknowledgement: %w", err) + } + result, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "result": subscriptionsListenResult{ + ResultType: modernResultTypeComplete, + Meta: subscriptionMeta, + }, + }) + if err != nil { + return nil, fmt.Errorf("marshalling subscriptions listen result: %w", err) + } + return [][]byte{ack, result}, nil +} + +// writeModernListenStream writes pre-marshalled frames as an SSE response and +// closes it. Everything it can return is a mid-stream failure -- the frames are +// already built and the status line is already gone -- so its error is only ever +// loggable, never reportable to the client. +func writeModernListenStream(w http.ResponseWriter, frames [][]byte) error { + // http.ResponseController rather than a w.(http.Flusher) type assertion. + // + // The assertion would not establish what it appears to. Every ResponseWriter + // wrapper in vMCP's middleware chain -- audit/auditor.go, + // telemetry/middleware.go, bodylimit/middleware.go, mcp/tool_filter.go -- + // defines Flush() UNCONDITIONALLY and forwards only if the writer it wraps + // supports it. So the assertion succeeds whenever a wrapper is present, + // whatever the writer underneath can do, and says nothing about whether a + // flush will actually reach the socket. ResponseController unwraps the chain + // (via Unwrap) and reports a real answer. + // + // The tradeoff, stated because it costs something: there is no longer a + // PRE-write capability check. ResponseController exposes no "can you flush" + // predicate -- the only way to find out is to call Flush, which implicitly + // commits the header -- so an unflushable writer now surfaces as a mid-stream + // error rather than a -32603 before anything is written. That is deliberate: + // the old assertion answered early but WRONGLY, and an early wrong answer on a + // path no deployment can reach is worse than a correct late one. If Go ever + // grows a flushability predicate, move this back ahead of WriteHeader. + rc := http.NewResponseController(w) + + w.Header().Set("Content-Type", "text/event-stream") + // no-store as well as no-cache: no-cache still permits a shared cache to + // STORE the response and revalidate, which for a per-identity stream is wrong. + w.Header().Set("Cache-Control", "no-cache, no-store, no-transform, private") + w.Header().Set("Connection", "keep-alive") + // Draft Streamable HTTP, "When initiating an SSE stream, servers SHOULD + // include the X-Accel-Buffering: no header", so a reverse proxy (nginx and + // friends) does not accumulate events before forwarding them. This is the + // only SSE-emitting path in the codebase, so there is nowhere else to copy + // the header from. + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + + for _, frame := range frames { + // Frame format is load-bearing, not cosmetic: an SSE event is terminated + // by a BLANK line, so the trailing "\n\n" is what makes a parser dispatch + // the event. Emitting a single "\n" leaves both frames inside one + // never-dispatched event -- the client connects and then silently receives + // nothing, which is exactly the mutation that escaped both this package's + // tests and the real-client integration suite. + // + // The "event: message" line is NOT mandated: the draft Streamable HTTP + // page specifies the content type and the stream lifecycle but never an + // SSE event name, and it explicitly drops resumability ("Resumable SSE + // streams via Last-Event-ID are not supported"), which is why no "id:" + // field is emitted either. The name is kept for symmetry with go-sdk's + // framing; do not add an "id:" to match it. + if _, err := fmt.Fprintf(w, "event: message\ndata: %s\n\n", frame); err != nil { + return fmt.Errorf("writing subscriptions/listen frame: %w", err) + } + // A writer that genuinely cannot flush would buffer both frames until the + // handler returned, which for a stream read incrementally is + // indistinguishable from a hang. Report it rather than hanging. + if err := rc.Flush(); err != nil { + return fmt.Errorf("flushing subscriptions/listen frame: %w", err) + } + } + return nil +} diff --git a/pkg/vmcp/server/modern_subscriptions_test.go b/pkg/vmcp/server/modern_subscriptions_test.go new file mode 100644 index 0000000000..36513551d8 --- /dev/null +++ b/pkg/vmcp/server/modern_subscriptions_test.go @@ -0,0 +1,464 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp/core" +) + +// pushCaps builds an mcp.ServerCapabilities advertising the given push flags. +// The anonymous struct types must be spelled out to match mcp.ServerCapabilities +// field-for-field; this helper keeps that noise out of the table below and lets +// tests construct advertisements the live newModernCapabilities cannot currently +// produce (every push flag it emits is false). +func pushCaps(toolsLC, promptsLC, resourcesLC, subscribe bool) mcp.ServerCapabilities { + var caps mcp.ServerCapabilities + caps.Tools = &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{ListChanged: toolsLC} + caps.Prompts = &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{ListChanged: promptsLC} + caps.Resources = &struct { + Subscribe bool `json:"subscribe,omitempty"` + ListChanged bool `json:"listChanged,omitempty"` + }{Subscribe: subscribe, ListChanged: resourcesLC} + return caps +} + +// allWanted is a client asking for every subscribable type there is. +func allWanted() notificationSubscriptions { + return notificationSubscriptions{ + ToolsListChanged: true, + PromptsListChanged: true, + ResourcesListChanged: true, + ResourceSubscriptions: []string{"file:///a", "file:///b"}, + } +} + +// TestHonoredSubscriptions is the central assertion of this change: a +// notification type is honored if and only if the client asked for it AND the +// matching capability is advertised. That intersection is the whole of SEP-2575's +// per-type/per-URI opt-in filtering, and it is what guarantees the +// acknowledgement can never promise more than server/discover advertises. +// +// The advertised side is a parameter rather than read from a constant, so this +// covers the honoring combinations that the live advertisement cannot currently +// reach -- which is exactly what makes the handler correct-by-construction the +// day a capability flag flips. +func TestHonoredSubscriptions(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + want notificationSubscriptions + advertised mcp.ServerCapabilities + expect notificationSubscriptions + }{ + { + name: "live advertisement honors nothing even when everything is requested", + want: allWanted(), + advertised: newModernCapabilities(true, true, true, true), + expect: notificationSubscriptions{}, + }, + { + name: "all advertised and all wanted honors everything", + want: allWanted(), + advertised: pushCaps(true, true, true, true), + expect: allWanted(), + }, + { + name: "nothing wanted honors nothing even when all advertised", + want: notificationSubscriptions{}, + advertised: pushCaps(true, true, true, true), + expect: notificationSubscriptions{}, + }, + { + // Named for what it asserts: everything is advertised, only tools is + // wanted, so ONLY tools comes back -- the three advertised-but-unwanted + // types must be absent. Exact equality is what pins that. + name: "only the wanted subset is honored when all are advertised", + want: notificationSubscriptions{ToolsListChanged: true}, + advertised: pushCaps(true, true, true, true), + expect: notificationSubscriptions{ToolsListChanged: true}, + }, + { + name: "wanted but not advertised is dropped per type", + want: allWanted(), + advertised: pushCaps(true, false, false, false), + expect: notificationSubscriptions{ToolsListChanged: true}, + }, + { + name: "resource subscriptions are gated by subscribe, not by resources listChanged", + want: notificationSubscriptions{ResourceSubscriptions: []string{"file:///a"}}, + advertised: pushCaps(false, false, true, false), + expect: notificationSubscriptions{}, + }, + { + name: "resource subscriptions survive when subscribe is advertised", + want: notificationSubscriptions{ResourceSubscriptions: []string{"file:///a"}}, + advertised: pushCaps(false, false, false, true), + expect: notificationSubscriptions{ResourceSubscriptions: []string{"file:///a"}}, + }, + { + name: "absent capability pointers honor nothing", + want: allWanted(), + advertised: mcp.ServerCapabilities{}, + expect: notificationSubscriptions{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.expect, honoredSubscriptions(tt.want, tt.advertised)) + }) + } +} + +// listenSSE drives dispatchModern for a subscriptions/listen request and splits +// the SSE body into its decoded JSON-RPC frames. It goes through dispatchModern +// rather than calling the handler directly so the method-switch wiring is +// covered too. +func listenSSE(t *testing.T, fakeCore *modernFakeCore, params any, id any) (*httptest.ResponseRecorder, []map[string]any) { + t.Helper() + + raw, err := json.Marshal(params) + require.NoError(t, err) + + s := &Server{ + config: &Config{Name: testServerName, Version: testServerVersion}, + core: fakeCore, + } + req := httptest.NewRequest(http.MethodPost, "/mcp", nil).WithContext(context.Background()) + rec := httptest.NewRecorder() + + s.dispatchModern(rec, req, &mcpparser.ParsedMCPRequest{ + ID: id, + Method: methodSubscriptionsListen, + Params: raw, + }) + + // Only strict-parse an actual stream; a rejected listen answers with a plain + // JSON error envelope and has no frames. + if rec.Header().Get("Content-Type") != "text/event-stream" { + return rec, nil + } + return rec, parseSSEStrict(t, rec.Body.String()) +} + +// TestDispatchModernSubscriptionsListen_AcknowledgesEmptyHonoredSet covers the +// full wire contract for the only path the live advertisement can produce: an +// SSE response carrying the mandatory acknowledgement first and the terminating +// result second, both tagged with the subscription id, and an honored set that +// is explicitly empty rather than absent. +// +// discoverCaps is all-true so the test proves the empty honored set comes from +// the capability advertisement (which advertises no push flags) and not merely +// from the identity having nothing to reach. +func TestDispatchModernSubscriptionsListen_AcknowledgesEmptyHonoredSet(t *testing.T) { + t.Parallel() + + fakeCore := &modernFakeCore{discoverCaps: core.DiscoverCapabilities{ + HasTools: true, HasResources: true, HasResourceTemplates: true, HasPrompts: true, + }} + rec, frames := listenSSE(t, fakeCore, subscriptionsListenParams{ + Notifications: ¬ificationSubscriptions{ + ToolsListChanged: true, + PromptsListChanged: true, + ResourcesListChanged: true, + ResourceSubscriptions: []string{"file:///a"}, + }, + }, "sub-1") + + assert.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "text/event-stream", rec.Header().Get("Content-Type")) + require.Len(t, frames, 2, "expected an acknowledgement frame followed by a result frame") + + // Frame 1: the acknowledgement MUST come first, and MUST report an empty + // honored set. An absent "notifications" object would be a different (and + // wrong) statement -- "I did not answer" rather than "I honor none". + ack := frames[0] + assert.Equal(t, notificationSubscriptionsAcked, ack["method"]) + ackParams, ok := ack["params"].(map[string]any) + require.True(t, ok, "acknowledgement must carry params") + honored, ok := ackParams["notifications"].(map[string]any) + require.True(t, ok, "acknowledgement must carry an explicit notifications object") + assert.Empty(t, honored, "no push capability is advertised, so nothing may be honored") + + ackMeta, ok := ackParams["_meta"].(map[string]any) + require.True(t, ok, "acknowledgement must be tagged with the subscription id") + assert.Equal(t, "sub-1", ackMeta[modernSubscriptionIDKey]) + + // Frame 2: the terminating result, same subscription id. Its presence is what + // makes the stream honestly finite: there is nothing to deliver, so the + // subscription ends immediately instead of idling open forever. + result := frames[1] + assert.Equal(t, "sub-1", result["id"]) + resultBody, ok := result["result"].(map[string]any) + require.True(t, ok, "result frame must carry a result") + assert.Equal(t, modernResultTypeComplete, resultBody["resultType"]) + resultMeta, ok := resultBody["_meta"].(map[string]any) + require.True(t, ok, "result must be tagged with the subscription id") + assert.Equal(t, "sub-1", resultMeta[modernSubscriptionIDKey]) +} + +// TestDispatchModernSubscriptionsListen_Errors covers the two rejections that +// must NOT open a stream: a missing required "notifications" field, and a failed +// capability resolution. +func TestDispatchModernSubscriptionsListen_Errors(t *testing.T) { + t.Parallel() + + // NOTE: there is deliberately no "capability resolution failure" case here. + // The M2 pre-check skips core.Discover entirely while no push capability is + // advertised, so the fan-out -- and its error path -- is unreachable from this + // handler today. Testing it would require a capability flag that + // newModernCapabilities never sets, and asserting a branch the code cannot + // reach would be theatre. TestNewModernCapabilities_AdvertisesNoPushCapability + // is what pins the precondition that makes it unreachable; if a flag is ever + // flipped, the error path becomes live and wants a case here. + tests := []struct { + name string + params any + wantCode float64 + wantStatus int + }{ + { + name: "missing notifications field", + params: map[string]any{}, + wantCode: jsonRPCCodeInvalidParams, + wantStatus: http.StatusBadRequest, + }, + { + name: "explicitly null notifications field", + params: map[string]any{"notifications": nil}, + wantCode: jsonRPCCodeInvalidParams, + wantStatus: http.StatusBadRequest, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec, _ := listenSSE(t, &modernFakeCore{}, tt.params, 7) + + assert.Equal(t, tt.wantStatus, rec.Code) + assert.NotEqual(t, "text/event-stream", rec.Header().Get("Content-Type"), + "a rejected listen must not open a stream") + + var body map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "expected a JSON-RPC error envelope") + assert.Equal(t, tt.wantCode, errObj["code"]) + }) + } +} + +// parseSSEStrict parses an SSE body the way a compliant client does, and is +// deliberately STRICT where the previous helper was lenient. +// +// That leniency hid a real gap: it split on "\n\n" and then fell back to matching +// any line beginning "data: ", so replacing the blank-line event terminator with +// a single "\n" left the whole suite green -- including the real-client +// integration suite -- even though a conformant parser would see one +// never-terminated event and dispatch neither frame. The ack-then-result contract +// was therefore unverifiable. +// +// Rules enforced here, per the SSE grammar: events are separated by a blank +// line, every event must be terminated (the body ends with a blank line), and +// each event carries exactly one "data:" field. +func parseSSEStrict(t *testing.T, body string) []map[string]any { + t.Helper() + + require.NotEmpty(t, body, "expected an SSE body") + require.True(t, strings.HasSuffix(body, "\n\n"), + "SSE body must end with a blank line; an unterminated final event is never dispatched by a client") + + var frames []map[string]any + for _, block := range strings.Split(strings.TrimSuffix(body, "\n\n"), "\n\n") { + require.NotEmpty(t, block, "empty SSE event block: two consecutive blank lines") + + var data string + seenData := false + for _, line := range strings.Split(block, "\n") { + field, value, ok := strings.Cut(line, ":") + require.True(t, ok, "malformed SSE line %q: no field separator", line) + value = strings.TrimPrefix(value, " ") + switch field { + case "event": + assert.Equal(t, "message", value, "MCP SSE frames must use the 'message' event type") + case "data": + require.False(t, seenData, "more than one data field in a single SSE event") + data, seenData = value, true + default: + t.Fatalf("unexpected SSE field %q", field) + } + } + require.True(t, seenData, "SSE event carried no data field") + + var frame map[string]any + require.NoError(t, json.Unmarshal([]byte(data), &frame), "SSE data must be valid JSON") + frames = append(frames, frame) + } + return frames +} + +// TestDispatchModernSubscriptionsListen_WireContract pins the literal strings and +// envelope fields the protocol requires. Every one of these was previously free: +// mutating the method name, the acknowledgement name, or the subscriptionId key +// left the whole package green, and those last two are precisely the constants +// whose provenance was least certain. +func TestDispatchModernSubscriptionsListen_WireContract(t *testing.T) { + t.Parallel() + + // Literal, not the constant: asserting against the constant would pass even if + // the constant itself were mistyped, which is the mutation that escaped. + assert.Equal(t, "subscriptions/listen", methodSubscriptionsListen) + assert.Equal(t, "notifications/subscriptions/acknowledged", notificationSubscriptionsAcked) + assert.Equal(t, "io.modelcontextprotocol/subscriptionId", modernSubscriptionIDKey) + + rec, frames := listenSSE(t, &modernFakeCore{}, subscriptionsListenParams{ + Notifications: ¬ificationSubscriptions{ToolsListChanged: true}, + }, "sub-wire") + + assert.Equal(t, http.StatusOK, rec.Code) + require.Len(t, frames, 2) + + // Both frames must declare JSON-RPC 2.0 -- unasserted before, and a protocol + // violation if dropped. + for i, frame := range frames { + assert.Equal(t, "2.0", frame["jsonrpc"], "frame %d must carry jsonrpc 2.0", i) + } + + // The acknowledgement must be FIRST and must be the ack method: schema.ts + // makes both a MUST ("This notification MUST be the first message the server + // sends carrying the subscription's ID"). + assert.Equal(t, "notifications/subscriptions/acknowledged", frames[0]["method"], + "the first frame must be the acknowledgement, by literal name") + assert.NotContains(t, frames[0], "id", "a notification must not carry an id") + + // The result frame closes the subscription and echoes the request id. + assert.Equal(t, "sub-wire", frames[1]["id"]) + assert.NotContains(t, frames[1], "method", "a response must not carry a method") + + // The subscriptionId key, by literal string, on both frames. + ackParams, ok := frames[0]["params"].(map[string]any) + require.True(t, ok) + ackMeta, ok := ackParams["_meta"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "sub-wire", ackMeta["io.modelcontextprotocol/subscriptionId"]) + + resultBody, ok := frames[1]["result"].(map[string]any) + require.True(t, ok) + resultMeta, ok := resultBody["_meta"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "sub-wire", resultMeta["io.modelcontextprotocol/subscriptionId"]) +} + +// TestNewModernCapabilities_AdvertisesNoPushCapability is the guard the whole +// empty-honored-set design rests on. +// +// The argument for this handler being conformant rather than a silent no-op has +// two halves: (1) the advertisement and the honored set are computed from one +// source, which the code enforces, and (2) that source advertises nothing +// pushable -- which, before this test, was enforced by nothing at all. The safety +// property was "no line of source sets these true", and a runtime WARN only fires +// after such a line ships, in production, per request, naming a developer action +// an operator cannot remedy. +// +// So it is asserted here, at build time, for every combination of the presence +// flags. +// +// IF YOU FLIP ONE OF THESE: you are promising to deliver that notification type. +// Nothing in vMCP can keep that promise yet -- dispatchModern creates no session, +// so the list_changed machinery never runs for a Modern client, and backend +// resources/updated is not forwarded even on Legacy. Implement delivery in the +// same change (#5743), including the schema MUST that every notification on a +// listen stream carries io.modelcontextprotocol/subscriptionId, and update this +// test deliberately rather than to make it pass. +// +// FLIPPING A FLAG ALSO RE-ENABLES CODE THAT IS CURRENTLY UNREACHABLE, and it +// carries no test coverage today precisely because it cannot be entered. Restore +// coverage for all three when you flip: +// +// - the core.Discover call in dispatchModernSubscriptionsListen (the ceiling +// pre-check skips it entirely while nothing can be honored), AND +// - that call's error path, which maps a fan-out failure to -32603. Its test was +// removed rather than left asserting an unenterable branch. +// - the handler's non-empty-honored branch, including its WARN. +func TestNewModernCapabilities_AdvertisesNoPushCapability(t *testing.T) { + t.Parallel() + + for _, hasTools := range []bool{false, true} { + for _, hasResources := range []bool{false, true} { + for _, hasTemplates := range []bool{false, true} { + for _, hasPrompts := range []bool{false, true} { + caps := newModernCapabilities(hasTools, hasResources, hasTemplates, hasPrompts) + + if caps.Tools != nil { + assert.False(t, caps.Tools.ListChanged, + "tools.listChanged must stay false until delivery exists (#5743)") + } + if caps.Prompts != nil { + assert.False(t, caps.Prompts.ListChanged, + "prompts.listChanged must stay false until delivery exists (#5743)") + } + if caps.Resources != nil { + assert.False(t, caps.Resources.ListChanged, + "resources.listChanged must stay false until delivery exists (#5743)") + assert.False(t, caps.Resources.Subscribe, + "resources.subscribe must stay false until delivery exists (#5743)") + } + + // And the consequence: a client asking for everything is + // honored nothing, whatever it can reach. + assert.True(t, honoredSubscriptions(allWanted(), caps).isEmpty(), + "no capability combination may honor a subscription") + } + } + } + } +} + +// TestDispatchModernSubscriptionsListen_SkipsFanOutWhenNothingHonorable pins the +// pre-check: core.Discover is an un-rate-limited backend fan-out whose result +// cannot change the answer while no push capability is advertised, so it must not +// be called at all. Without this, N no-op listens cost N fan-outs. +// +// discoverCalls counting is the point -- the response is identical either way, so +// only the call count distinguishes the optimisation being present from absent. +func TestDispatchModernSubscriptionsListen_SkipsFanOutWhenNothingHonorable(t *testing.T) { + t.Parallel() + + fakeCore := &modernFakeCore{discoverCaps: core.DiscoverCapabilities{ + HasTools: true, HasResources: true, HasResourceTemplates: true, HasPrompts: true, + }} + + _, frames := listenSSE(t, fakeCore, subscriptionsListenParams{ + Notifications: ¬ificationSubscriptions{ + ToolsListChanged: true, + PromptsListChanged: true, + ResourcesListChanged: true, + ResourceSubscriptions: []string{"file:///a", "file:///b"}, + }, + }, "sub-nofanout") + + require.Len(t, frames, 2, "the response must be unaffected by the pre-check") + assert.Zero(t, fakeCore.discoverCalls, + "core.Discover must not be called when no push capability is advertised") +} diff --git a/pkg/vmcp/server/schema_fidelity_regression_test.go b/pkg/vmcp/server/schema_fidelity_regression_test.go new file mode 100644 index 0000000000..97c42723b8 --- /dev/null +++ b/pkg/vmcp/server/schema_fidelity_regression_test.go @@ -0,0 +1,146 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_ToolSchemaWithTypeObject_ProjectedIntact guards schema +// fidelity for the SUPPORTED case: a tool whose InputSchema includes a +// top-level "type":"object" is projected to the client UNCHANGED in tools/list, +// with its properties and required array preserved. The Serve path +// (coreSessionTools in serve_handlers.go) marshals the core's InputSchema into +// RawInputSchema; the go-sdk bridge's normalizeObjectSchema passes object-typed +// schemas through verbatim, so properties/required survive. +// +// All assertions use encoding/json and generic maps — not SDK types — so the +// check is against the wire representation the client sees. +func TestRegression_ToolSchemaWithTypeObject_ProjectedIntact(t *testing.T) { + t.Parallel() + + originalSchema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "x": map[string]any{"type": "string"}, + }, + "required": []any{"x"}, + } + fc := &fakeCore{tools: []vmcp.Tool{{ + Name: "schema-tool", + Description: "a schema-fidelity test tool", + InputSchema: originalSchema, + }}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode, "tools/list should succeed") + + env, body := readServeJSONRPC(t, resp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list must have a result; body: %s", string(body)) + + tools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + + // Find the projected tool by name. + var found map[string]any + for _, tRaw := range tools { + tm, ok := tRaw.(map[string]any) + if !ok { + continue + } + if tm["name"] == "schema-tool" { + found = tm + break + } + } + require.NotNil(t, found, "schema-tool must be present in tools/list; tools: %v", tools) + + inputSchema, ok := found["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object; tool: %v", found) + + // The projected schema must equal the original map exactly: properties and + // required preserved, type unchanged. + assert.Equal(t, originalSchema, inputSchema, + "an object-typed inputSchema must be projected intact; got %v", inputSchema) + assert.Equal(t, "object", inputSchema["type"], + "the top-level type must remain \"object\"; got %v", inputSchema) +} + +// TestRegression_ToolSchemaWithoutTypeObject_ProjectedIntact guards the go-sdk +// bridge's handling of a schema that omits the top-level "type":"object". +// Earlier mcpcompat releases dropped properties/required for such schemas (a +// fidelity gap versus mcp-go); as of toolhive-core v0.0.28 normalizeObjectSchema +// preserves them and supplies the missing "type":"object". This test now pins +// that intact projection so any future regression is a deliberate, visible flip. +func TestRegression_ToolSchemaWithoutTypeObject_ProjectedIntact(t *testing.T) { + t.Parallel() + + originalSchema := map[string]any{ + "properties": map[string]any{ + "x": map[string]any{"type": "string"}, + }, + "required": []any{"x"}, + // NOTE: deliberately NO top-level "type": "object". + } + fc := &fakeCore{tools: []vmcp.Tool{{ + Name: "schema-tool", + Description: "a schema-fidelity test tool", + InputSchema: originalSchema, + }}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode, "tools/list should succeed") + + env, body := readServeJSONRPC(t, resp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list must have a result; body: %s", string(body)) + + tools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + + var found map[string]any + for _, tRaw := range tools { + tm, ok := tRaw.(map[string]any) + if !ok { + continue + } + if tm["name"] == "schema-tool" { + found = tm + break + } + } + require.NotNil(t, found, "schema-tool must be present in tools/list; tools: %v", tools) + + inputSchema, ok := found["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object; tool: %v", found) + + // The schema is projected intact: properties and required are preserved and + // the missing top-level "type":"object" is supplied (toolhive-core v0.0.28). + assert.Equal(t, "object", inputSchema["type"], + "a missing top-level type must be supplied as \"object\"; got %v", inputSchema) + assert.Equal(t, originalSchema["properties"], inputSchema["properties"], + "properties must be preserved; got %v", inputSchema) + assert.Equal(t, originalSchema["required"], inputSchema["required"], + "required must be preserved; got %v", inputSchema) +} diff --git a/pkg/vmcp/server/sdk_elicitation_adapter.go b/pkg/vmcp/server/sdk_elicitation_adapter.go index 3cf0eb7ebd..9ed38864e3 100644 --- a/pkg/vmcp/server/sdk_elicitation_adapter.go +++ b/pkg/vmcp/server/sdk_elicitation_adapter.go @@ -7,15 +7,15 @@ package server import ( "context" - "maps" - - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" + "errors" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/conversion" ) -// sdkElicitationAdapter wraps mark3labs MCPServer to implement vmcp.ElicitationRequester. +// sdkElicitationAdapter wraps mcpcompat MCPServer to implement vmcp.ElicitationRequester. // // It is the sole point where mcp-go elicitation types appear: it translates the // domain ElicitationRequest/ElicitationResult to/from the SDK types, keeping the @@ -24,22 +24,22 @@ import ( // Per MCP 2025-06-18 spec: The SDK handles JSON-RPC ID correlation internally, // so the adapter does not manage IDs. // -// Thread-safety: Safe for concurrent calls. The mark3labs MCPServer is thread-safe. +// Thread-safety: Safe for concurrent calls. The mcpcompat MCPServer is thread-safe. type sdkElicitationAdapter struct { - // mcpServer is the mark3labs SDK server instance that handles elicitation protocol. + // mcpServer is the mcpcompat SDK server instance that handles elicitation protocol. // Typed as the minimal mcpElicitationRequester seam so the translation logic // can be unit-tested against a fake SDK; *server.MCPServer satisfies it. mcpServer mcpElicitationRequester } -// mcpElicitationRequester is the minimal slice of the mark3labs SDK that the +// mcpElicitationRequester is the minimal slice of the mcpcompat SDK that the // adapter depends on. *server.MCPServer satisfies it in production; tests // substitute a fake to verify domain ⇄ mcp-go translation without a live session. type mcpElicitationRequester interface { RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) } -// NewSDKElicitationAdapter creates a new elicitation adapter that wraps the mark3labs SDK server. +// NewSDKElicitationAdapter creates a new elicitation adapter that wraps the mcpcompat SDK server. // // The returned adapter implements vmcp.ElicitationRequester by translating the domain // request to/from mcp-go types and delegating to the SDK's RequestElicitation method. @@ -59,11 +59,11 @@ func NewSDKElicitationAdapter(mcpServer *server.MCPServer) vmcp.ElicitationReque } // RequestElicitation translates the domain request to mcp-go, delegates to the -// mark3labs SDK's RequestElicitation method, and translates the response back. +// mcpcompat SDK's RequestElicitation method, and translates the response back. // // This is a synchronous blocking call that: // 1. Maps the domain ElicitationRequest to an mcp.ElicitationRequest -// 2. Forwards the request to the mark3labs SDK +// 2. Forwards the request to the mcpcompat SDK // 3. Blocks until the client responds or timeout occurs // 4. Maps the SDK's mcp.ElicitationResult back to the domain ElicitationResult // @@ -88,14 +88,21 @@ func (a *sdkElicitationAdapter) RequestElicitation( RequestedSchema: req.RequestedSchema, }, } - // Only attach _meta when the caller actually set it. NewMetaFromMap mutates - // its argument (it deletes progressToken), so copy first to avoid mutating - // the caller's map. - if req.Meta != nil { - mcpReq.Params.Meta = mcp.NewMetaFromMap(maps.Clone(req.Meta)) - } + // req.Meta came from the BACKEND's elicitation/create (forwarding.go's + // newElicitationForwarder), so it crosses the same trust boundary as a backend + // result: route it through conversion.ToMCPMeta, which strips the reserved + // io.modelcontextprotocol/* keys before they reach the downstream client. A + // leaked protocolVersion here is worse than on a result -- this is a + // server->client REQUEST, where a go-sdk client may validate it. + // + // ToMCPMeta also hoists progressToken and copies (never mutating the caller's + // map), and returns nil for empty input -- so _meta stays absent when the + // backend set none, or set only reserved keys. Do not swap in + // mcp.NewMetaFromMap: it returns a non-nil *Meta for nil input, which would + // start emitting an empty _meta. + mcpReq.Params.Meta = conversion.ToMCPMeta(req.Meta) - // Delegate to the mark3labs SDK's RequestElicitation method. + // Delegate to the mcpcompat SDK's RequestElicitation method. // The SDK will: // 1. Extract session ID from context (set by SDK middleware) // 2. Generate JSON-RPC ID for the request @@ -107,6 +114,18 @@ func (a *sdkElicitationAdapter) RequestElicitation( // We don't need to manage any of this - it's all handled by the SDK. resp, err := a.mcpServer.RequestElicitation(ctx, mcpReq) if err != nil { + // A refusal for lack of a downstream session (Modern ingress: the + // stateless dispatch installed no ClientSession) is recorded so + // dispatchModern can classify the resulting call failure as a + // -32021 MissingRequiredClientCapabilityError instead of an opaque + // internal error. The error itself still propagates unchanged: it is + // the answer to the BACKEND's elicitation request, and the typed + // sentinel does not survive that wire round-trip — the recorder is + // the only channel back to the dispatcher (see + // modern_capability_refusal.go). + if errors.Is(err, server.ErrNoActiveSession) { + recordCapabilityRefusal(ctx, capabilityElicitation) + } return nil, err } diff --git a/pkg/vmcp/server/sdk_elicitation_adapter_test.go b/pkg/vmcp/server/sdk_elicitation_adapter_test.go index 8c28877593..e5d73aa768 100644 --- a/pkg/vmcp/server/sdk_elicitation_adapter_test.go +++ b/pkg/vmcp/server/sdk_elicitation_adapter_test.go @@ -8,11 +8,11 @@ import ( "errors" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -142,9 +142,71 @@ func TestSDKElicitationAdapter_NilMetaProducesNoMeta(t *testing.T) { assert.Nil(t, fake.gotRequest.Params.Meta) } +// TestSDKElicitationAdapter_StripsReservedMeta is the #5986 regression pin for +// the server->client REQUEST direction. req.Meta here originates from the +// BACKEND's elicitation/create (newElicitationForwarder in pkg/vmcp/client), so +// it crosses the same trust boundary as a backend result: a reserved +// io.modelcontextprotocol/* key must not reach the downstream client, which may +// validate it (a leaked protocolVersion is worse on a request than on a result). +func TestSDKElicitationAdapter_StripsReservedMeta(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + meta map[string]any + wantNilMeta bool + wantAdditions map[string]any + }{ + { + name: "reserved keys stripped, backend metadata preserved", + meta: map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/serverInfo": map[string]any{"name": "attacker"}, + "trace": "abc", + }, + wantAdditions: map[string]any{"trace": "abc"}, + }, + { + // Collapses to no _meta at all rather than an empty object, matching + // the nil-Meta behavior above. + name: "only reserved keys produces no _meta", + meta: map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + }, + wantNilMeta: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + fake := &fakeSDKElicitationRequester{ + result: &mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{Action: mcp.ElicitationResponseActionAccept}, + }, + } + adapter := &sdkElicitationAdapter{mcpServer: fake} + + _, err := adapter.RequestElicitation(context.Background(), vmcp.ElicitationRequest{ + Message: "Confirm?", + Meta: tt.meta, + }) + require.NoError(t, err) + + if tt.wantNilMeta { + assert.Nil(t, fake.gotRequest.Params.Meta) + return + } + require.NotNil(t, fake.gotRequest.Params.Meta) + assert.Equal(t, tt.wantAdditions, fake.gotRequest.Params.Meta.AdditionalFields) + }) + } +} + // TestSDKElicitationAdapter_MetaIsCopied verifies that translating a request -// with Meta does not mutate the caller's map (NewMetaFromMap deletes -// progressToken from its argument). +// with Meta does not mutate the caller's map (the SDK's NewMetaFromMap deletes +// progressToken from its argument; conversion.ToMCPMeta copies instead). func TestSDKElicitationAdapter_MetaIsCopied(t *testing.T) { t.Parallel() @@ -179,7 +241,7 @@ func TestSDKElicitationAdapter_Integration(t *testing.T) { } // TestServer_MCPServer_ReturnsSameInstance verifies that (*Server).MCPServer -// returns the exact mark3labs server pointer stored at construction time. +// returns the exact mcpcompat server pointer stored at construction time. // Identity matters because ClientSession correlation is keyed to the server // that received the initialize request; embedders building their own // elicitation requester must receive the authoritative instance. diff --git a/pkg/vmcp/server/sdk_notifier_adapter.go b/pkg/vmcp/server/sdk_notifier_adapter.go new file mode 100644 index 0000000000..deb4fda2c7 --- /dev/null +++ b/pkg/vmcp/server/sdk_notifier_adapter.go @@ -0,0 +1,91 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// sdkNotifierAdapter wraps mcpcompat MCPServer to implement vmcp.ClientNotifier. +// +// It is the sole point where the SDK's SendNotificationToClient appears: it +// translates the domain ProgressNotification/LogMessage into the SDK's +// (method, params) form and delegates to the server, keeping the backend client +// free of SDK coupling. It mirrors sdkElicitationAdapter. +// +// Forwarding is best-effort: when the context carries no downstream session +// (health probes, capability listing) the SDK returns ErrNoActiveSession, which +// the adapter swallows so the backend's notification is silently dropped rather +// than surfaced as an error. +// +// Thread-safety: Safe for concurrent calls. The mcpcompat MCPServer is thread-safe. +type sdkNotifierAdapter struct { + // mcpServer is the mcpcompat SDK server. Typed as the minimal mcpClientNotifier + // seam so the translation logic can be unit-tested against a fake SDK; + // *server.MCPServer satisfies it. + mcpServer mcpClientNotifier +} + +// mcpClientNotifier is the minimal slice of the mcpcompat SDK that the adapter +// depends on. *server.MCPServer satisfies it in production; tests substitute a +// fake to verify domain -> mcp-go translation without a live session. +type mcpClientNotifier interface { + SendNotificationToClient(ctx context.Context, method string, params map[string]any) error +} + +// NewSDKNotifierAdapter creates a notifier adapter that wraps the mcpcompat SDK +// server. The returned adapter implements vmcp.ClientNotifier by translating the +// domain notification into the SDK's (method, params) form and delegating to +// SendNotificationToClient. Pass the *server.MCPServer obtained from +// (*Server).MCPServer() so notifications correlate with the server handling +// incoming client sessions. +func NewSDKNotifierAdapter(mcpServer *server.MCPServer) vmcp.ClientNotifier { + return &sdkNotifierAdapter{mcpServer: mcpServer} +} + +// NotifyProgress forwards a notifications/progress message to the downstream +// client. Total and Message are omitted from the wire params when unset (zero / +// empty) so a backend that sends only a progress value does not emit a spurious +// total:0. Best-effort: a missing downstream session is not an error. +func (a *sdkNotifierAdapter) NotifyProgress(ctx context.Context, n vmcp.ProgressNotification) error { + params := map[string]any{ + "progressToken": n.ProgressToken, + "progress": n.Progress, + } + if n.Total != 0 { + params["total"] = n.Total + } + if n.Message != "" { + params["message"] = n.Message + } + return a.send(ctx, vmcp.MethodProgressNotification, params) +} + +// NotifyLog forwards a notifications/message (logging) message to the downstream +// client. Logger is omitted when empty. Best-effort: a missing downstream +// session is not an error. +func (a *sdkNotifierAdapter) NotifyLog(ctx context.Context, n vmcp.LogMessage) error { + params := map[string]any{ + "level": n.Level, + "data": n.Data, + } + if n.Logger != "" { + params["logger"] = n.Logger + } + return a.send(ctx, vmcp.MethodLogNotification, params) +} + +// send delegates to the SDK and swallows the no-active-session error so +// forwarding stays best-effort per the ClientNotifier contract. +func (a *sdkNotifierAdapter) send(ctx context.Context, method string, params map[string]any) error { + err := a.mcpServer.SendNotificationToClient(ctx, method, params) + if errors.Is(err, server.ErrNoActiveSession) { + return nil + } + return err +} diff --git a/pkg/vmcp/server/sdk_notifier_adapter_test.go b/pkg/vmcp/server/sdk_notifier_adapter_test.go new file mode 100644 index 0000000000..264c137737 --- /dev/null +++ b/pkg/vmcp/server/sdk_notifier_adapter_test.go @@ -0,0 +1,122 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// fakeSDKNotifier captures the (method, params) the adapter forwards and returns +// a canned error, letting us assert the adapter's domain -> mcp-go mapping and +// its best-effort no-session handling. +type fakeSDKNotifier struct { + gotMethod string + gotParams map[string]any + err error +} + +func (f *fakeSDKNotifier) SendNotificationToClient(_ context.Context, method string, params map[string]any) error { + f.gotMethod = method + f.gotParams = params + return f.err +} + +func TestSDKNotifierAdapter_NotifyProgress(t *testing.T) { + t.Parallel() + + t.Run("includes total and message when set", func(t *testing.T) { + t.Parallel() + + fake := &fakeSDKNotifier{} + adapter := &sdkNotifierAdapter{mcpServer: fake} + + err := adapter.NotifyProgress(context.Background(), vmcp.ProgressNotification{ + ProgressToken: "tok-1", + Progress: 0.5, + Total: 1.0, + Message: "halfway", + }) + require.NoError(t, err) + assert.Equal(t, vmcp.MethodProgressNotification, fake.gotMethod) + assert.Equal(t, "tok-1", fake.gotParams["progressToken"]) + assert.Equal(t, 0.5, fake.gotParams["progress"]) + assert.Equal(t, 1.0, fake.gotParams["total"]) + assert.Equal(t, "halfway", fake.gotParams["message"]) + }) + + t.Run("omits total and message when unset", func(t *testing.T) { + t.Parallel() + + fake := &fakeSDKNotifier{} + adapter := &sdkNotifierAdapter{mcpServer: fake} + + err := adapter.NotifyProgress(context.Background(), vmcp.ProgressNotification{ + ProgressToken: "tok-1", + Progress: 0.5, + }) + require.NoError(t, err) + _, hasTotal := fake.gotParams["total"] + _, hasMessage := fake.gotParams["message"] + assert.False(t, hasTotal) + assert.False(t, hasMessage) + }) +} + +func TestSDKNotifierAdapter_NotifyLog(t *testing.T) { + t.Parallel() + + fake := &fakeSDKNotifier{} + adapter := &sdkNotifierAdapter{mcpServer: fake} + + err := adapter.NotifyLog(context.Background(), vmcp.LogMessage{ + Level: "info", + Logger: "backend", + Data: "something happened", + }) + require.NoError(t, err) + assert.Equal(t, vmcp.MethodLogNotification, fake.gotMethod) + assert.Equal(t, "info", fake.gotParams["level"]) + assert.Equal(t, "backend", fake.gotParams["logger"]) + assert.Equal(t, "something happened", fake.gotParams["data"]) +} + +// TestSDKNotifierAdapter_NoSessionIsNoOp verifies that a missing downstream +// session (ErrNoActiveSession) is swallowed so forwarding stays best-effort. +func TestSDKNotifierAdapter_NoSessionIsNoOp(t *testing.T) { + t.Parallel() + + fake := &fakeSDKNotifier{err: server.ErrNoActiveSession} + adapter := &sdkNotifierAdapter{mcpServer: fake} + + assert.NoError(t, adapter.NotifyProgress(context.Background(), vmcp.ProgressNotification{Progress: 1})) + assert.NoError(t, adapter.NotifyLog(context.Background(), vmcp.LogMessage{Level: "info"})) +} + +// TestSDKNotifierAdapter_OtherErrorIsPropagated verifies that a non-session +// error is surfaced (not swallowed). +func TestSDKNotifierAdapter_OtherErrorIsPropagated(t *testing.T) { + t.Parallel() + + fake := &fakeSDKNotifier{err: errors.New("transport closed")} + adapter := &sdkNotifierAdapter{mcpServer: fake} + + require.Error(t, adapter.NotifyProgress(context.Background(), vmcp.ProgressNotification{Progress: 1})) + require.Error(t, adapter.NotifyLog(context.Background(), vmcp.LogMessage{Level: "info"})) +} + +func TestSDKNotifierAdapter_Integration(t *testing.T) { + t.Parallel() + + mcpServer := server.NewMCPServer("test", "1.0.0") + adapter := NewSDKNotifierAdapter(mcpServer) + assert.NotNil(t, adapter) +} diff --git a/pkg/vmcp/server/sdk_sampling_adapter.go b/pkg/vmcp/server/sdk_sampling_adapter.go new file mode 100644 index 0000000000..c113a3bd8c --- /dev/null +++ b/pkg/vmcp/server/sdk_sampling_adapter.go @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// sdkSamplingAdapter wraps mcpcompat MCPServer to implement vmcp.SamplingRequester. +// +// It is the sole point where mcp-go sampling types appear: it translates the +// domain SamplingRequest/SamplingResult to/from the SDK types, keeping the +// backend client and the rest of vmcp free of SDK coupling. It mirrors +// sdkElicitationAdapter. +// +// Thread-safety: Safe for concurrent calls. The mcpcompat MCPServer is thread-safe. +type sdkSamplingAdapter struct { + // mcpServer is the mcpcompat SDK server instance that handles the sampling + // protocol. Typed as the minimal mcpSamplingRequester seam so the translation + // logic can be unit-tested against a fake SDK; *server.MCPServer satisfies it. + mcpServer mcpSamplingRequester +} + +// mcpSamplingRequester is the minimal slice of the mcpcompat SDK that the +// adapter depends on. *server.MCPServer satisfies it in production; tests +// substitute a fake to verify domain <-> mcp-go translation without a live session. +type mcpSamplingRequester interface { + RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) +} + +// NewSDKSamplingAdapter creates a new sampling adapter that wraps the mcpcompat +// SDK server. The returned adapter implements vmcp.SamplingRequester by +// translating the domain request to/from mcp-go types and delegating to the +// SDK's RequestSampling method. Session management and JSON-RPC ID correlation +// are handled entirely by the SDK. Pass the *server.MCPServer obtained from +// (*Server).MCPServer() so the returned requester correlates with the server +// handling incoming client sessions. +func NewSDKSamplingAdapter(mcpServer *server.MCPServer) vmcp.SamplingRequester { + return &sdkSamplingAdapter{mcpServer: mcpServer} +} + +// RequestSampling translates the domain request to mcp-go, delegates to the +// mcpcompat SDK's RequestSampling method, and translates the response back. +// +// This is a synchronous blocking call: it maps the domain SamplingRequest to an +// mcp.CreateMessageRequest, forwards it to the SDK (which routes it to the +// correct client session and blocks until the client responds), and maps the +// SDK's mcp.CreateMessageResult back to the domain SamplingResult. +// +// Returns an error if the request fails, times out, the context carries no +// downstream session, or the client did not advertise the sampling capability. +func (a *sdkSamplingAdapter) RequestSampling( + ctx context.Context, + req vmcp.SamplingRequest, +) (*vmcp.SamplingResult, error) { + mcpReq := mcp.CreateMessageRequest{ + CreateMessageParams: mcp.CreateMessageParams{ + Messages: toMCPSamplingMessages(req.Messages), + ModelPreferences: toMCPModelPreferences(req.ModelPreferences), + SystemPrompt: req.SystemPrompt, + IncludeContext: req.IncludeContext, + Temperature: req.Temperature, + MaxTokens: req.MaxTokens, + StopSequences: req.StopSequences, + Metadata: req.Metadata, + }, + } + + resp, err := a.mcpServer.RequestSampling(ctx, mcpReq) + if err != nil { + // Mirror of sdkElicitationAdapter's refusal recording — see the + // comment there and modern_capability_refusal.go. + if errors.Is(err, server.ErrNoActiveSession) { + recordCapabilityRefusal(ctx, capabilitySampling) + } + return nil, err + } + + return &vmcp.SamplingResult{ + Role: string(resp.Role), + Content: resp.Content, + Model: resp.Model, + StopReason: resp.StopReason, + }, nil +} + +// toMCPSamplingMessages maps domain sampling messages to the SDK type. +func toMCPSamplingMessages(msgs []vmcp.SamplingMessage) []mcp.SamplingMessage { + if msgs == nil { + return nil + } + out := make([]mcp.SamplingMessage, len(msgs)) + for i, m := range msgs { + out[i] = mcp.SamplingMessage{ + Role: mcp.Role(m.Role), + Content: m.Content, + } + } + return out +} + +// toMCPModelPreferences maps the domain model preferences to the SDK type, +// preserving a nil (unset) value. +func toMCPModelPreferences(p *vmcp.ModelPreferences) *mcp.ModelPreferences { + if p == nil { + return nil + } + hints := make([]mcp.ModelHint, len(p.Hints)) + for i, h := range p.Hints { + hints[i] = mcp.ModelHint{Name: h.Name} + } + return &mcp.ModelPreferences{ + Hints: hints, + CostPriority: p.CostPriority, + SpeedPriority: p.SpeedPriority, + IntelligencePriority: p.IntelligencePriority, + } +} diff --git a/pkg/vmcp/server/sdk_sampling_adapter_test.go b/pkg/vmcp/server/sdk_sampling_adapter_test.go new file mode 100644 index 0000000000..2ec8a44b1c --- /dev/null +++ b/pkg/vmcp/server/sdk_sampling_adapter_test.go @@ -0,0 +1,118 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// fakeSDKSamplingRequester captures the translated mcp-go request and returns a +// canned response/error, letting us assert the adapter's domain <-> mcp-go mapping. +type fakeSDKSamplingRequester struct { + gotRequest mcp.CreateMessageRequest + result *mcp.CreateMessageResult + err error +} + +func (f *fakeSDKSamplingRequester) RequestSampling( + _ context.Context, request mcp.CreateMessageRequest, +) (*mcp.CreateMessageResult, error) { + f.gotRequest = request + return f.result, f.err +} + +func TestSDKSamplingAdapter_RequestSampling(t *testing.T) { + t.Parallel() + + t.Run("round-trips request and result", func(t *testing.T) { + t.Parallel() + + fake := &fakeSDKSamplingRequester{ + result: &mcp.CreateMessageResult{ + SamplingMessage: mcp.SamplingMessage{ + Role: mcp.Role("assistant"), + Content: map[string]any{"type": "text", "text": "hi"}, + }, + Model: "test-model", + StopReason: "endTurn", + }, + } + adapter := &sdkSamplingAdapter{mcpServer: fake} + + req := vmcp.SamplingRequest{ + Messages: []vmcp.SamplingMessage{ + {Role: "user", Content: map[string]any{"type": "text", "text": "hello"}}, + }, + ModelPreferences: &vmcp.ModelPreferences{ + Hints: []vmcp.ModelHint{{Name: "claude"}}, + CostPriority: 0.5, + }, + SystemPrompt: "be brief", + MaxTokens: 100, + Temperature: 0.7, + StopSequences: []string{"STOP"}, + } + + res, err := adapter.RequestSampling(context.Background(), req) + require.NoError(t, err) + require.NotNil(t, res) + + // Request translation: domain fields map onto mcp.CreateMessageParams. + require.Len(t, fake.gotRequest.Messages, 1) + assert.Equal(t, mcp.Role("user"), fake.gotRequest.Messages[0].Role) + assert.Equal(t, "be brief", fake.gotRequest.SystemPrompt) + assert.Equal(t, 100, fake.gotRequest.MaxTokens) + assert.Equal(t, []string{"STOP"}, fake.gotRequest.StopSequences) + require.NotNil(t, fake.gotRequest.ModelPreferences) + require.Len(t, fake.gotRequest.ModelPreferences.Hints, 1) + assert.Equal(t, "claude", fake.gotRequest.ModelPreferences.Hints[0].Name) + + // Result translation: mcp fields map back onto domain SamplingResult. + assert.Equal(t, "assistant", res.Role) + assert.Equal(t, "test-model", res.Model) + assert.Equal(t, "endTurn", res.StopReason) + assert.Equal(t, map[string]any{"type": "text", "text": "hi"}, res.Content) + }) + + t.Run("nil model preferences stays nil", func(t *testing.T) { + t.Parallel() + + fake := &fakeSDKSamplingRequester{ + result: &mcp.CreateMessageResult{SamplingMessage: mcp.SamplingMessage{Role: mcp.Role("assistant")}}, + } + adapter := &sdkSamplingAdapter{mcpServer: fake} + + _, err := adapter.RequestSampling(context.Background(), vmcp.SamplingRequest{MaxTokens: 10}) + require.NoError(t, err) + assert.Nil(t, fake.gotRequest.ModelPreferences) + }) + + t.Run("SDK error is propagated", func(t *testing.T) { + t.Parallel() + + fake := &fakeSDKSamplingRequester{err: errors.New("no active session")} + adapter := &sdkSamplingAdapter{mcpServer: fake} + + res, err := adapter.RequestSampling(context.Background(), vmcp.SamplingRequest{MaxTokens: 10}) + require.Error(t, err) + assert.Nil(t, res) + }) +} + +func TestSDKSamplingAdapter_Integration(t *testing.T) { + t.Parallel() + + mcpServer := server.NewMCPServer("test", "1.0.0") + adapter := NewSDKSamplingAdapter(mcpServer) + assert.NotNil(t, adapter) +} diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 2a7ac4304b..b80a00dc7f 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -10,9 +10,8 @@ import ( "net/http" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/audit" asrunner "github.com/stacklok/toolhive/pkg/authserver/runner" "github.com/stacklok/toolhive/pkg/telemetry" @@ -29,14 +28,12 @@ import ( // It carries the subset of the legacy server.Config that configures the HTTP/SDK // runtime, plus the cross-cutting TelemetryProvider/AuditConfig that are consumed // by both the core (core.New) and the transport (Serve) — not a clean partition. -// The fields the core exclusively owns (Aggregator, Router, BackendClient, Authz) -// live on core.Config instead and are absent here. Two collaborators are shared -// rather than core-owned, so they appear on both configs: BackendRegistry (the -// Serve session layer also enumerates backends when creating a session) and the -// composite-tool WorkflowDefs (carried indirectly via SessionManagerConfig, whose -// ComposerFactory closes over the core-owned router/backend client). The -// composition root passes the same BackendRegistry instance to both core.New and -// Serve and assembles SessionManagerConfig. +// The fields the core exclusively owns (Aggregator, Router, BackendClient, Authz, +// composite-tool WorkflowDefs) live on core.Config instead and are absent here. +// One collaborator is shared rather than core-owned, so it appears on both +// configs: BackendRegistry (the Serve session layer also enumerates backends when +// creating a session). The composition root passes the same BackendRegistry +// instance to both core.New and Serve and assembles SessionManagerConfig. // // The name intentionally stutters as server.ServerConfig: it is mandated by the // New/Serve split, and the package's existing Config is the legacy transport @@ -67,6 +64,9 @@ type ServerConfig struct { // WriteTimeout is the HTTP server write deadline. Values below the default // (30 s) are raised to it. WriteTimeout time.Duration + // HeartbeatInterval configures the SSE keep-alive ping interval on GET + // connections (default: 30s when zero). + HeartbeatInterval time.Duration // AuthMiddleware is the optional authentication middleware applied to MCP routes. // If nil, no authentication is required. @@ -111,17 +111,11 @@ type ServerConfig struct { // SessionManagerConfig is the pre-built construction config for the vMCP session // manager (sessionmanager.New). FactoryConfig.Base is the required underlying - // MultiSessionFactory; the config also carries the composite-tool WorkflowDefs and - // ComposerFactory, the optimizer factory/config, and the telemetry provider. - // Required; Serve returns an error when it is nil. The composition root assembles - // it because the ComposerFactory closes over the core-owned router and backend - // client. - // - // Caller responsibility: unlike server.New, Serve does NOT run validateWorkflows on - // FactoryConfig.WorkflowDefs — the composition root must validate composite-tool - // definitions before assembling this config (sessionmanager.New only checks the - // WorkflowDefs/ComposerFactory pairing). This responsibility moves here with the - // relocation and matters when server.New is routed through Serve in Phase 3. + // MultiSessionFactory; the config also carries the optimizer factory/config, the + // telemetry provider, the session-cache capacity, and the AdvertiseFromCore flag. + // Required; Serve returns an error when it is nil. Composite-tool workflows are + // NOT carried here — they are core-owned (core.Config.WorkflowDefs, validated in + // core.New). // // Caller responsibility (AC2, the single-aggregation contract): FactoryConfig.Base // MUST be constructed WITHOUT a session.WithAggregator option on the Serve path. On @@ -136,9 +130,10 @@ type ServerConfig struct { // // Caller responsibility (optimizer): to enable the optimizer on the Serve path, set // FactoryConfig.OptimizerConfig/OptimizerFactory AND FactoryConfig.AdvertiseFromCore. - // Serve then builds a per-session optimizer over the core's tools (serve_optimizer.go); - // AdvertiseFromCore suppresses the factory's own optimizer decorator so the shared FTS5 - // store is not double-indexed (see FactoryConfig.AdvertiseFromCore). + // Serve then builds a per-session optimizer over the core's tools (serve_optimizer.go). + // sessionmanager.New rejects an optimizer without AdvertiseFromCore, so a composition + // root that forgets the flag fails at startup instead of double-indexing the shared + // FTS5 store (see FactoryConfig.AdvertiseFromCore). SessionManagerConfig *sessionmanager.FactoryConfig // TelemetryProvider is the cross-cutting telemetry provider (also consumed by @@ -214,16 +209,85 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) // the mcp-go server for the same reason). hooks := &server.Hooks{} - // Build the mcp-go server (mirrors server.New): tools and resources are - // registered per session. listChanged is advertised because a backend change - // sends list_changed to the sessions (serve_refresh.go). + // The completion handler is a single global handler (unlike the per-session + // tool/resource/prompt handlers), so it is wired here at construction. It closes + // over the late-assigned srv — safe because WithCompletionHandler only stores the + // handler; it is invoked at request time, always after srv is assigned below. This + // mirrors the hooks' close-over-srv pattern. Setting the handler makes the shim + // auto-advertise the completions capability at initialize. + var srv *Server + completionHandler := func(ctx context.Context, req mcp.CompleteRequest) (*mcp.CompleteResult, error) { + return srv.coreCompletionHandler(ctx, req) + } + + // Resource subscribe/unsubscribe handlers, wired here alongside the completion + // handler for the same close-over-late-srv reason. They accept the subscription + // at ack level (enforce session binding, validate the URI is an advertised, + // admitted resource) and return nil so go-sdk records it; vMCP does NOT forward + // backend resources/updated notifications yet (see coreSubscribeHandler and the + // architecture doc's subscription note). + subscribeHandler := func(ctx context.Context, uri string) error { + return srv.coreSubscribeHandler(ctx, "resources/subscribe", uri) + } + unsubscribeHandler := func(ctx context.Context, uri string) error { + return srv.coreSubscribeHandler(ctx, "resources/unsubscribe", uri) + } + + // Build the mcp-go server (mirrors server.New): tools, resources, and + // prompts are registered dynamically, so capabilities start disabled — + // except resources, which additionally advertise subscribe support so + // spec-compliant clients issue resources/subscribe (the handlers above + // answer it at ack level). + // + // tools.listChanged (#5748) and resources.listChanged/prompts.listChanged + // (#5969) are now advertised: a persistent backend connection's + // notifications/tools/list_changed, notifications/resources/list_changed, + // or notifications/prompts/list_changed drives a per-kind session resync + // (serve_list_changed.go) that replaces the SDK session's tool, resource + // (+ resource-template), or prompt store via SetSessionTools/ + // SetSessionResources/SetSessionResourceTemplates/SetSessionPrompts, and + // go-sdk auto-emits the corresponding list_changed notification to the + // downstream client whenever that capability is on. WithPromptCapabilities + // also fixes a spec-conformance defect: vMCP already served prompts per + // session (SessionWithPrompts) but never declared the prompts capability at + // initialize, violating MCP 2025-11-25's "Servers that support prompts MUST + // declare the prompts capability during initialization". + // + // R1 (verified against toolhive-core v0.0.34 / go-sdk v1.7.0-pre.3): go-sdk's own + // AddTool/RemoveTools (and the resource/prompt equivalents) call a shared + // changeAndNotify that debounces (10ms) and then broadcasts the + // notification to EVERY session currently connected to this *gosdk.Server + // (mcpcompat/server/server.go's srv is one instance shared across all + // sessions) — not just the session whose overlay changed. go-sdk's own + // bind() adds a session to that broadcast list as soon as the transport + // connects (Connect), before initialize/notifications initialized even + // round-trip; go-sdk's shared.go itself documents this as a known upstream + // gap ("potential spec violation... when the feature list changes before + // the session ... is initialized"). Concretely: every session's + // registration-time per-session AddTool/resource/prompt calls + // (setSessionToolsDirect/setSessionResourcesDirect/ + // setSessionResourceTemplatesDirect/setSessionPromptsDirect) now cause a + // list_changed broadcast to every OTHER already-connected session too, not + // only the one just registered — this is inherent go-sdk behavior, not + // something introduced by the #5748/#5969 resync sinks, and there is no + // supported hook to scope or suppress it short of patching the vendored SDK + // (out of scope here). It is accepted as a benign nuisance for all three + // capability kinds: MCP clients treat list_changed as "you may want to + // refetch," so an extra notification only costs an idempotent list + // round-trip — it never changes what a given session's own list actually + // returns (each session's advertised set still comes solely from its own + // overlay). See docs/arch/10-virtual-mcp-architecture.md for the + // propagation writeup. mcpServer := server.NewMCPServer( serveCfg.Name, serveCfg.Version, server.WithToolCapabilities(true), - server.WithResourceCapabilities(false, true), + server.WithResourceCapabilities(true, true), + server.WithPromptCapabilities(true), server.WithLogging(), server.WithHooks(hooks), + server.WithCompletionHandler(completionHandler), + server.WithSubscribeHandlers(subscribeHandler, unsubscribeHandler), ) // Two-phase session-creation wiring (relocated from server.New). The pluggable @@ -255,7 +319,7 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) // leak it (mirroring the closeStorageOnErr guard on the sibling sessionDataStorage). sessionManager := transportsession.NewManager(serveCfg.SessionTTL, transportsession.NewStreamableSession) - srv := &Server{ + srv = &Server{ config: serveCfg, core: v, mcpServer: mcpServer, @@ -272,13 +336,18 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) statusReporter: cfg.StatusReporter, } - if cfg.TelemetryProvider != nil { - // Gauges are diagnostics only: a registration failure must not stop the - // server (and an error return here would leak the components built above). - if err := monitorSessionCardinality(cfg.TelemetryProvider.MeterProvider(), srv); err != nil { - slog.Warn("failed to register session cardinality gauges", "error", err) - } - } + // Server-lifetime parent context for asynchronous tools/resources/prompts + // list_changed resync work (#5748, #5969). Cancelling it on Stop aborts any + // in-flight backend re-aggregation a late notification kicked off, so + // nothing outlives the server. context.Background() is the parent (not the + // Serve ctx, which may be request-scoped) so the resync lifetime is bounded + // solely by Stop. + resyncCtx, resyncCancel := context.WithCancel(context.Background()) + srv.resyncBaseCtx = resyncCtx + srv.shutdownFuncs = append(srv.shutdownFuncs, func(context.Context) error { + resyncCancel() + return nil + }) if optimizerCleanup != nil { srv.shutdownFuncs = append(srv.shutdownFuncs, optimizerCleanup) @@ -291,28 +360,29 @@ func Serve(ctx context.Context, v core.VMCP, cfg *ServerConfig) (*Server, error) return v.Close() }) - // Register the SDK hooks against the assembled *Server (relocated from + // Register the three SDK hooks against the assembled *Server (relocated from // server.New). They drive phase-2 of two-phase creation (OnRegisterSession) and // the cross-pod Redis re-hydration of per-session tools (OnBeforeListTools / - // OnBeforeCallTool) and resources (OnBeforeListResources / OnBeforeReadResource). + // OnBeforeCallTool); the *Server receiver methods are unchanged. The call hook + // passes the requested tool so a partial store can be topped up. hooks.AddOnRegisterSession(func(hookCtx context.Context, session server.ClientSession) { srv.handleSessionRegistration(hookCtx, session) }) - hooks.AddOnUnregisterSession(func(_ context.Context, session server.ClientSession) { - srv.activeClientSessions.remove(session.SessionID()) - }) hooks.AddBeforeListTools(func(hookCtx context.Context, _ any, _ *mcp.ListToolsRequest) { srv.lazyInjectSessionTools(hookCtx, "") }) hooks.AddBeforeCallTool(func(hookCtx context.Context, _ any, req *mcp.CallToolRequest) { srv.lazyInjectSessionTools(hookCtx, req.Params.Name) }) - hooks.AddBeforeListResources(func(hookCtx context.Context, _ any, _ *mcp.ListResourcesRequest) { - srv.lazyInjectSessionResources(hookCtx) - }) - hooks.AddBeforeReadResource(func(hookCtx context.Context, _ any, _ *mcp.ReadResourceRequest) { - srv.lazyInjectSessionResources(hookCtx) - }) + + // Surface the capability gate's verdict once at startup: the blocker list is + // derived from construction-time configuration and cannot change afterwards, + // so this single line is the operator-visible record of why Modern-capable + // clients of this instance negotiate down to Legacy (see modern_gate.go). + if blocked := srv.modernDispatchBlockers(); len(blocked) > 0 { + slog.Warn("MCP 2026-07-28 (Modern) dispatch disabled: enabled features require the session (Legacy) path", + "features", blocked) + } // Disarm the close-on-error guard: the Server is fully constructed. closeStorageOnErr = false @@ -351,6 +421,7 @@ func buildServeConfig(cfg *ServerConfig) *Config { EndpointPath: cfg.EndpointPath, SessionTTL: cfg.SessionTTL, WriteTimeout: cfg.WriteTimeout, + HeartbeatInterval: cfg.HeartbeatInterval, AuthMiddleware: cfg.AuthMiddleware, AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, diff --git a/pkg/vmcp/server/serve_crosspod_rebind_test.go b/pkg/vmcp/server/serve_crosspod_rebind_test.go index 8f08ead134..85cf7338f0 100644 --- a/pkg/vmcp/server/serve_crosspod_rebind_test.go +++ b/pkg/vmcp/server/serve_crosspod_rebind_test.go @@ -48,24 +48,19 @@ const ( ) // countingFactory wraps a MultiSessionFactory and counts phase-2 creations -// (MakeSessionWithID) and cross-pod restores (RestoreSession). When gate is set, -// MakeSessionWithID blocks until it is closed, which holds the session record as an -// unbound placeholder. +// (MakeSessionWithID) and cross-pod restores (RestoreSession). type countingFactory struct { base vmcpsession.MultiSessionFactory - gate chan struct{} makes atomic.Int32 restores atomic.Int32 } func (f *countingFactory) MakeSessionWithID( ctx context.Context, id string, identity *auth.Identity, backends []*vmcp.Backend, + sink vmcpsession.ListChangedSink, ) (vmcpsession.MultiSession, error) { f.makes.Add(1) - if f.gate != nil { - <-f.gate - } - return f.base.MakeSessionWithID(ctx, id, identity, backends) + return f.base.MakeSessionWithID(ctx, id, identity, backends, sink) } func (f *countingFactory) RestoreSession( @@ -99,18 +94,17 @@ type crossPod struct { // startCrossPod builds one Serve-path vMCP replica whose vMCP session manager uses // the shared store, and returns the full shared Handler (auth middleware included). -// A non-nil gate blocks the replica's phase-2 creation until it is closed. -func startCrossPod(t *testing.T, shared transportsession.DataStorage, gate chan struct{}) *crossPod { +func startCrossPod(t *testing.T, shared transportsession.DataStorage) *crossPod { t.Helper() - return startCrossPodWithTTL(t, shared, gate, time.Hour) + return startCrossPodWithTTL(t, shared, time.Hour) } // startCrossPodWithTTL is startCrossPod with the given session TTL. -func startCrossPodWithTTL(t *testing.T, shared transportsession.DataStorage, gate chan struct{}, ttl time.Duration) *crossPod { +func startCrossPodWithTTL(t *testing.T, shared transportsession.DataStorage, ttl time.Duration) *crossPod { t.Helper() authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() require.NoError(t, authReg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, strategies.NewUnauthenticatedStrategy())) - factory := &countingFactory{base: vmcpsession.NewSessionFactory(authReg), gate: gate} + factory := &countingFactory{base: vmcpsession.NewSessionFactory(authReg)} registry := vmcp.NewImmutableRegistry([]vmcp.Backend{}) fc := &fakeCore{ tools: []vmcp.Tool{{Name: "echo", Description: "echo"}}, @@ -287,8 +281,8 @@ func TestCrossPodSessionRebind(t *testing.T) { shared, err := transportsession.NewLocalSessionDataStorage(time.Hour) require.NoError(t, err) t.Cleanup(func() { _ = shared.Close() }) - podA := startCrossPod(t, shared, nil) - podB := startCrossPod(t, shared, nil) + podA := startCrossPod(t, shared) + podB := startCrossPod(t, shared) sessionID := initSession(t, shared, podA, "alice") aliceBinding := storedBinding(t, shared, sessionID) @@ -345,88 +339,12 @@ func TestCrossPodSessionRebind(t *testing.T) { } } -// TestCrossPodGETDuringRegistration covers the window before the session has a -// binding: alice's initialize on pod A is still registering (phase-2 creation is -// blocked) when eve sends a GET with the session ID to pod B. Only the session's own -// initialize may create the session, so pod B must not run phase-2 creation for -// eve's GET; it waits for alice's binding and then rejects eve. -func TestCrossPodGETDuringRegistration(t *testing.T) { - t.Parallel() - - shared, err := transportsession.NewLocalSessionDataStorage(time.Hour) - require.NoError(t, err) - t.Cleanup(func() { _ = shared.Close() }) - gate := make(chan struct{}) - podA := startCrossPod(t, shared, gate) - podB := startCrossPod(t, shared, nil) - - // alice's initialize on pod A. mcp-go writes the response headers (with the - // session ID) before it runs the registration hook, which then blocks in the gate. - initW := newSSEWriter() - initDone := make(chan struct{}) - go func() { - defer close(initDone) - podA.handler.ServeHTTP(initW, newCrossPodRequest(context.Background(), http.MethodPost, "alice", "", - `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18",`+ - `"capabilities":{},"clientInfo":{"name":"t","version":"1"}}}`)) - }() - select { - case <-initW.headerWritten: - case <-time.After(5 * time.Second): - t.Fatal("initialize did not write response headers") - } - sessionID := initW.Header().Get("Mcp-Session-Id") - require.NotEmpty(t, sessionID) - require.Eventually(t, func() bool { return podA.factory.makes.Load() == 1 }, - 2*time.Second, 5*time.Millisecond, "alice's registration must be blocked in phase-2 creation") - require.Empty(t, storedBinding(t, shared, sessionID), "the session must not be bound yet") - - // eve's GET on pod B while alice's registration is pending. - getCtx, cancelGET := context.WithCancel(context.Background()) - getW := newSSEWriter() - getDone := make(chan struct{}) - go func() { - defer close(getDone) - podB.handler.ServeHTTP(getW, newCrossPodRequest(getCtx, http.MethodGet, "eve", sessionID, "")) - }() - - // Give eve's GET time to reach the registration hook, then let alice's - // registration finish. - time.Sleep(200 * time.Millisecond) - close(gate) - select { - case <-initDone: - case <-time.After(5 * time.Second): - t.Fatal("alice's registration did not finish") - } - select { - case <-getW.headerWritten: - case <-getDone: - case <-time.After(5 * time.Second): - t.Fatal("GET did not write response headers") - } - cancelGET() - <-getDone - - assert.Equal(t, http.StatusForbidden, getW.status(), "eve's GET must be refused once alice's binding is written") - assert.Equal(t, int32(0), podB.factory.makes.Load(), - "a GET on pod B must not call CreateSession/MakeSessionWithID") - binding := storedBinding(t, shared, sessionID) - alice := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ - Subject: "alice", Claims: map[string]any{"iss": crossPodIssuer, "sub": "alice"}, - }} - assert.NoError(t, vmcpsession.ValidateCaller(binding, alice), "the session must be bound to alice") - - body := callEcho(t, podB, "alice", sessionID) - assert.Contains(t, body, `"ok"`, "alice must own her session after eve's GET") - body = callEcho(t, podB, "eve", sessionID) - assert.Contains(t, body, "Unauthorized", "eve must not be able to use alice's session") -} - // TestCrossPodRehydration covers a session initialized on pod A and used on pod B, // which never ran the registration for it: pod B restores the session from the -// shared store and fills the per-session tool and resource stores from the core, -// without creating the session again. +// shared store and fills the per-session tool store from the core, without +// creating the session again. Resources are not covered: mcpcompat runs no hook +// before resources/list or resources/read, so pod B serves no resources for the +// session (upstream gap). func TestCrossPodRehydration(t *testing.T) { t.Parallel() @@ -438,11 +356,6 @@ func TestCrossPodRehydration(t *testing.T) { }{ {name: "tools/list", method: "tools/list", params: `{}`, want: `"name":"echo"`}, {name: "tools/call", method: "tools/call", params: `{"name":"echo","arguments":{}}`, want: `"text":"ok"`}, - {name: "resources/list", method: "resources/list", params: `{}`, want: `"uri":"` + crossPodResourceURI + `"`}, - { - name: "resources/read", method: "resources/read", - params: `{"uri":"` + crossPodResourceURI + `"}`, want: `"text":"resource-body"`, - }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { @@ -450,8 +363,8 @@ func TestCrossPodRehydration(t *testing.T) { shared, err := transportsession.NewLocalSessionDataStorage(time.Hour) require.NoError(t, err) t.Cleanup(func() { _ = shared.Close() }) - podA := startCrossPod(t, shared, nil) - podB := startCrossPod(t, shared, nil) + podA := startCrossPod(t, shared) + podB := startCrossPod(t, shared) sessionID := initSession(t, shared, podA, "alice") binding := storedBinding(t, shared, sessionID) @@ -478,7 +391,7 @@ func TestListenStreamGuard(t *testing.T) { shared, err := transportsession.NewLocalSessionDataStorage(time.Hour) require.NoError(t, err) t.Cleanup(func() { _ = shared.Close() }) - return shared, startCrossPod(t, shared, nil), startCrossPod(t, shared, nil) + return shared, startCrossPod(t, shared), startCrossPod(t, shared) } t.Run("GET on the owner pod", func(t *testing.T) { @@ -549,8 +462,8 @@ func TestCrossPodIdleSweepKeepsSessionInUse(t *testing.T) { shared, err := transportsession.NewLocalSessionDataStorage(time.Hour) require.NoError(t, err) t.Cleanup(func() { _ = shared.Close() }) - podA := startCrossPodWithTTL(t, shared, nil, ttl) - podB := startCrossPodWithTTL(t, shared, nil, ttl) + podA := startCrossPodWithTTL(t, shared, ttl) + podB := startCrossPodWithTTL(t, shared, ttl) sessionID := initSession(t, shared, podA, "alice") require.Contains(t, callEcho(t, podB, "alice", sessionID), `"ok"`) diff --git a/pkg/vmcp/server/serve_handlers.go b/pkg/vmcp/server/serve_handlers.go index 401ad75d39..c39b25c75e 100644 --- a/pkg/vmcp/server/serve_handlers.go +++ b/pkg/vmcp/server/serve_handlers.go @@ -9,38 +9,18 @@ import ( "errors" "fmt" "log/slog" - "time" - - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" - "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" ) -const ( - // sessionRegistrationWaitTimeout bounds how long a call waits for the - // session's registration (OnRegisterSession) to write the identity binding. - sessionRegistrationWaitTimeout = 10 * time.Second - - // sessionRegistrationPollInterval is the first step between two reads of the - // session record while waiting for the registration. Each step doubles, up - // to sessionRegistrationMaxPollInterval. - sessionRegistrationPollInterval = 25 * time.Millisecond - - // sessionRegistrationMaxPollInterval caps the step between two reads. - sessionRegistrationMaxPollInterval = 500 * time.Millisecond -) - -// errSessionNotRegistered is returned when a call gives up waiting for the -// session's registration. The session is not terminated for it. -var errSessionNotRegistered = errors.New("session registration not finished") - // This file holds the Serve-path (s.core != nil) capability wiring. On the Serve // path the core VMCP is the single authoritative aggregation: session // registration sources the advertised tool/resource set from core.ListTools / @@ -68,8 +48,10 @@ var errSessionNotRegistered = errors.New("session registration not finished") // ServerConfig.SessionManagerConfig); otherwise the factory would aggregate a second, // divergent set whose routing table this path discards. // -// Prompts are intentionally not injected: per-session prompt injection is not yet -// supported by the SDK (parity with the legacy path, which also omits them). +// Prompts are injected here too, alongside tools and resources: the SDK's +// SessionWithPrompts interface (toolhive-core mcpcompat) supports per-session prompt +// injection, so core.ListPrompts sources the advertised set and corePromptHandler +// routes prompts/get through core.GetPrompt (an exact mirror of the resource path). func (s *Server) injectCoreSessionCapabilities(ctx context.Context, session server.ClientSession) error { sessionID := session.SessionID() @@ -89,6 +71,16 @@ func (s *Server) injectCoreSessionCapabilities(ctx context.Context, session serv slog.Error("failed to list core resources for session", "session_id", sessionID, "error", err) return err } + resourceTemplates, err := s.coreSessionResourceTemplates(ctx, sessionID, identity) + if err != nil { + slog.Error("failed to list core resource templates for session", "session_id", sessionID, "error", err) + return err + } + prompts, err := s.coreSessionPrompts(ctx, sessionID, identity) + if err != nil { + slog.Error("failed to list core prompts for session", "session_id", sessionID, "error", err) + return err + } if len(resources) > 0 { if err := setSessionResourcesDirect(session, resources); err != nil { @@ -96,6 +88,18 @@ func (s *Server) injectCoreSessionCapabilities(ctx context.Context, session serv return err } } + if len(resourceTemplates) > 0 { + if err := setSessionResourceTemplatesDirect(session, resourceTemplates); err != nil { + slog.Error("failed to add session resource templates", "session_id", sessionID, "error", err) + return err + } + } + if len(prompts) > 0 { + if err := setSessionPromptsDirect(session, prompts); err != nil { + slog.Error("failed to add session prompts", "session_id", sessionID, "error", err) + return err + } + } if len(tools) > 0 { if err := setSessionToolsDirect(session, tools); err != nil { slog.Error("failed to add session tools", "session_id", sessionID, "error", err) @@ -106,7 +110,9 @@ func (s *Server) injectCoreSessionCapabilities(ctx context.Context, session serv slog.Info("session capabilities injected from core", "session_id", sessionID, "tool_count", len(tools), - "resource_count", len(resources)) + "resource_count", len(resources), + "resource_template_count", len(resourceTemplates), + "prompt_count", len(prompts)) return nil } @@ -172,19 +178,87 @@ func (s *Server) coreSessionResources( sdkResources := make([]server.ServerResource, 0, len(domainResources)) for _, domainResource := range domainResources { - sdkResources = append(sdkResources, server.ServerResource{ - Resource: mcp.Resource{ + sdkResources = append(sdkResources, server.NewServerResourceWithResult( + mcp.Resource{ Name: domainResource.Name, URI: domainResource.URI, Description: domainResource.Description, MIMEType: domainResource.MimeType, }, - Handler: s.coreResourceHandler(sessionID, domainResource.URI, s.backendDisplayName(ctx, domainResource.BackendID)), - }) + s.coreResourceHandler(sessionID, domainResource.URI, s.backendDisplayName(ctx, domainResource.BackendID)), + )) } return sdkResources, nil } +// coreSessionResourceTemplates queries the core for the resource templates +// advertised to identity and adapts them to SDK ServerResourceTemplates whose +// handlers route through core.ReadResource. Mirrors coreSessionResources. +// +// The SDK serves resources/read for a URI matching a registered template through +// the template's handler, so the handler receives the concrete (expanded) URI on +// the request. It routes that URI through core.ReadResource, which resolves it via +// the router's template-match fallback (session_router.go) — there is no dedicated +// template read method. The backend display name is resolved here per template for +// audit labelling (see coreResourceTemplateHandler). +func (s *Server) coreSessionResourceTemplates( + ctx context.Context, sessionID string, identity *auth.Identity, +) ([]server.ServerResourceTemplate, error) { + domainTemplates, err := s.core.ListResourceTemplates(ctx, identity) + if err != nil { + return nil, fmt.Errorf("core ListResourceTemplates: %w", err) + } + + sdkTemplates := make([]server.ServerResourceTemplate, 0, len(domainTemplates)) + for _, domainTemplate := range domainTemplates { + sdkTemplates = append(sdkTemplates, server.NewServerResourceTemplateWithResult( + mcp.ResourceTemplate{ + Name: domainTemplate.Name, + URITemplate: domainTemplate.URITemplate, + Description: domainTemplate.Description, + MIMEType: domainTemplate.MimeType, + }, + s.coreResourceTemplateHandler( + sessionID, s.backendDisplayName(ctx, domainTemplate.BackendID)), + )) + } + return sdkTemplates, nil +} + +// coreSessionPrompts queries the core for the prompts advertised to identity and +// adapts them to SDK ServerPrompts whose handlers route through core.GetPrompt. The +// backend display name is resolved here and captured in the handler closure for audit +// labelling (see corePromptHandler). Mirrors coreSessionResources. +func (s *Server) coreSessionPrompts( + ctx context.Context, sessionID string, identity *auth.Identity, +) ([]server.ServerPrompt, error) { + domainPrompts, err := s.core.ListPrompts(ctx, identity) + if err != nil { + return nil, fmt.Errorf("core ListPrompts: %w", err) + } + + sdkPrompts := make([]server.ServerPrompt, 0, len(domainPrompts)) + for _, domainPrompt := range domainPrompts { + arguments := make([]mcp.PromptArgument, 0, len(domainPrompt.Arguments)) + for _, arg := range domainPrompt.Arguments { + arguments = append(arguments, mcp.PromptArgument{ + Name: arg.Name, + Description: arg.Description, + Required: arg.Required, + }) + } + sdkPrompts = append(sdkPrompts, server.ServerPrompt{ + Prompt: mcp.Prompt{ + Name: domainPrompt.Name, + Description: domainPrompt.Description, + Arguments: arguments, + }, + Handler: s.corePromptHandler(sessionID, domainPrompt.Name, s.backendDisplayName(ctx, domainPrompt.BackendID)), + }) + } + return sdkPrompts, nil +} + // coreToolHandler builds the SDK handler for a Serve-path tool. It labels the audit // event with backendName (the serving backend, pre-resolved at registration), // enforces the session's identity binding, then delegates to core.CallTool with the @@ -200,6 +274,10 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server bi.BackendName = backendName } + // Shape-check on the SDK decode first: a non-object payload is rejected before the + // core is reached. This validates request shape for the parse we forward too — both + // decode the same request bytes, so a valid transport parse implies a valid SDK + // shape (the substitution below only changes which equal-shaped map is used). args, ok := req.Params.Arguments.(map[string]any) if !ok { wrappedErr := fmt.Errorf("%w: arguments must be object, got %T", vmcp.ErrInvalidInput, req.Params.Arguments) @@ -207,6 +285,16 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server return mcp.NewToolResultError(wrappedErr.Error()), nil } + // Prefer the transport-parsed argument map so the pre-dispatch authz gate's + // decision, this handler's enforced decision, and the forwarded backend call + // all derive from one decode (#5845). A nil result means no matching parse is + // available (batch / embedders bypassing transport middleware / method or tool + // mismatch); those paths keep the SDK decode above and still make a single + // decision on a single map, so gate and dispatch cannot diverge. + if pa := gateParsedArgs(ctx, toolName); pa != nil { + args = pa + } + caller, _ := auth.IdentityFromContext(ctx) if err := s.enforceSessionBinding(ctx, sessionID, caller); err != nil { s.terminateOnBindingFailure(sessionID, toolName, err) @@ -227,12 +315,33 @@ func (s *Server) coreToolHandler(sessionID, toolName, backendName string) server } } +// gateParsedArgs returns the argument map the pre-dispatch authz gate decided on +// (pkg/mcp's transport parse), so the gated decision, this handler's enforced +// decision, and the forwarded backend call all share one decode (#5845). It returns +// parsed.Arguments only when the context carries a ParsedMCPRequest for the SAME +// tools/call on the SAME tool with non-nil arguments. +// +// A nil result tells the caller to keep the SDK decode: batch requests, embedders +// that bypass the transport middleware, or a method/tool mismatch leave no matching +// parse. Those paths make a single decision on a single map (the SDK decode), so +// there is still no allow-then-deny divergence between gate and dispatch. +func gateParsedArgs(ctx context.Context, toolName string) map[string]any { + parsed := mcpparser.GetParsedMCPRequest(ctx) + if parsed == nil || parsed.Method != "tools/call" || parsed.ResourceID != toolName || parsed.Arguments == nil { + return nil + } + return parsed.Arguments +} + // coreResourceHandler builds the SDK handler for a Serve-path resource. It mirrors // coreToolHandler: audit label, binding check, then core.ReadResource with explicit identity. +// The result-returning shape forwards the backend's non-reserved _meta (trace ids, +// custom fields) onto the wire result — reserved io.modelcontextprotocol/* keys are +// stripped by conversion.ToMCPMeta, matching every other Legacy egress (#5986). func (s *Server) coreResourceHandler( sessionID, uri, backendName string, -) func(ctx context.Context, req mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { - return func(ctx context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { +) server.ResourceResultHandlerFunc { + return func(ctx context.Context, _ mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { if bi, ok := audit.BackendInfoFromContext(ctx); ok && bi != nil { bi.BackendName = backendName } @@ -246,12 +355,189 @@ func (s *Server) coreResourceHandler( result, err := s.core.ReadResource(ctx, caller, uri) if err != nil { if errors.Is(err, vmcp.ErrAuthorizationFailed) { - return nil, errors.New("read denied by authorization policy") + return nil, errors.New(vmcp.DenyMessageResourceRead) + } + return nil, err + } + return &mcp.ReadResourceResult{ + Result: mcp.Result{Meta: conversion.ToMCPMeta(result.Meta)}, + Contents: conversion.ToMCPResourceContents(result.Contents), + }, nil + } +} + +// coreResourceTemplateHandler builds the SDK handler for a Serve-path resource +// template. It delegates to coreResourceHandler with the concrete URI read from +// the request (req.Params.URI) rather than a fixed URI captured at registration, +// because one template serves a whole family of URIs. coreResourceHandler routes +// that URI through core.ReadResource, which resolves it via the router's +// template-match fallback. +func (s *Server) coreResourceTemplateHandler( + sessionID, backendName string, +) server.ResourceResultHandlerFunc { + return func(ctx context.Context, req mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + return s.coreResourceHandler(sessionID, req.Params.URI, backendName)(ctx, req) + } +} + +// corePromptHandler builds the SDK handler for a Serve-path prompt. It mirrors +// coreResourceHandler: audit label, binding check, then core.GetPrompt with explicit +// identity. The request's string-typed prompt arguments are widened to map[string]any +// for the domain call, and the domain result is converted to *mcp.GetPromptResult. +func (s *Server) corePromptHandler( + sessionID, promptName, backendName string, +) server.PromptHandlerFunc { + return func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + if bi, ok := audit.BackendInfoFromContext(ctx); ok && bi != nil { + bi.BackendName = backendName + } + + caller, _ := auth.IdentityFromContext(ctx) + if err := s.enforceSessionBinding(ctx, sessionID, caller); err != nil { + s.terminateOnBindingFailure(sessionID, promptName, err) + return nil, fmt.Errorf("unauthorized: %w", err) + } + + args := make(map[string]any, len(req.Params.Arguments)) + for k, v := range req.Params.Arguments { + args[k] = v + } + + result, err := s.core.GetPrompt(ctx, caller, promptName, args) + if err != nil { + if errors.Is(err, vmcp.ErrAuthorizationFailed) { + return nil, errors.New(vmcp.DenyMessagePromptGet) } return nil, err } - return conversion.ToMCPResourceContents(result.Contents), nil + return &mcp.GetPromptResult{ + Result: mcp.Result{Meta: conversion.ToMCPMeta(result.Meta)}, + Description: result.Description, + Messages: conversion.ToMCPPromptMessages(result.Messages), + }, nil + } +} + +// coreCompletionHandler is the SDK completion/complete handler on the Serve path. +// Unlike the per-session tool/resource/prompt handlers (whose closures capture a +// sessionID at registration), completion is a single global handler installed on the +// mcp-go server via server.WithCompletionHandler, so it resolves the session ID from +// the SDK request context. It mirrors the other core handlers: audit label, +// enforceSessionBinding, then routes through core.Complete with the caller's explicit +// identity resolved at the transport boundary. +// +// It converts the mcp-go CompleteRequest ref (a PromptReference or ResourceReference, +// typed as any) into the domain vmcp.CompletionRef, and the domain +// *vmcp.CompletionResult back into *mcp.CompleteResult. An admission denial surfaces +// as the standard deny message for the referenced capability kind. +func (s *Server) coreCompletionHandler( + ctx context.Context, req mcp.CompleteRequest, +) (*mcp.CompleteResult, error) { + ref, err := completionRefFromMCP(req.Params.Ref) + if err != nil { + return nil, err + } + + sessionID := "" + if sess := server.ClientSessionFromContext(ctx); sess != nil { + sessionID = sess.SessionID() + } + + caller, _ := auth.IdentityFromContext(ctx) + if err := s.enforceSessionBinding(ctx, sessionID, caller); err != nil { + s.terminateOnBindingFailure(sessionID, "completion/complete", err) + return nil, fmt.Errorf("unauthorized: %w", err) + } + + var contextArgs map[string]string + if req.Params.Context != nil { + contextArgs = req.Params.Context.Arguments + } + + result, err := s.core.Complete( + ctx, caller, ref, req.Params.Argument.Name, req.Params.Argument.Value, contextArgs, + ) + if err != nil { + if errors.Is(err, vmcp.ErrAuthorizationFailed) { + return nil, errors.New(completionDenyMessage(ref.Type)) + } + return nil, err } + + return &mcp.CompleteResult{ + Completion: mcp.CompletionResultDetails{ + Values: result.Values, + Total: result.Total, + HasMore: result.HasMore, + }, + }, nil +} + +// coreSubscribeHandler answers resources/subscribe and resources/unsubscribe at +// ack level. Like coreCompletionHandler it is a single global handler installed on +// the mcp-go server (not a per-session closure), so it resolves the session ID from +// the SDK request context. It mirrors the other core handlers: resolve the caller's +// explicit identity at the transport boundary, enforce the session binding, then +// validate that uri is an advertised resource the caller may read (core.LookupResource +// applies the same admission decision ListResources/ReadResource enforce). An unknown +// or admission-denied URI is rejected; on success it returns nil so go-sdk records the +// subscription and answers the client with a success ack. +// +// Scope limitation (intentional): vMCP accepts the subscription but does NOT yet +// forward backend resources/updated notifications to the client — doing so needs +// persistent per-session backend connections, which is out of scope here. capability +// is the method name used for binding-failure diagnostics. +func (s *Server) coreSubscribeHandler(ctx context.Context, capability, uri string) error { + sessionID := "" + if sess := server.ClientSessionFromContext(ctx); sess != nil { + sessionID = sess.SessionID() + } + + caller, _ := auth.IdentityFromContext(ctx) + if err := s.enforceSessionBinding(ctx, sessionID, caller); err != nil { + s.terminateOnBindingFailure(sessionID, capability, err) + return fmt.Errorf("unauthorized: %w", err) + } + + // Validate the URI is an advertised resource the caller may read; an unknown or + // admission-denied URI is rejected rather than silently accepted. + if _, err := s.core.LookupResource(ctx, caller, uri); err != nil { + return fmt.Errorf("%s: %w", capability, err) + } + return nil +} + +// completionRefFromMCP converts the mcp-go CompleteParams.Ref (typed as any, holding +// a PromptReference or ResourceReference, possibly decoded as a map from JSON) into a +// domain vmcp.CompletionRef. It accepts both the typed shim structs and the generic +// map[string]any the SDK may hand back after a JSON round-trip. +func completionRefFromMCP(raw any) (vmcp.CompletionRef, error) { + switch r := raw.(type) { + case mcp.PromptReference: + return vmcp.CompletionRef{Type: r.Type, Name: r.Name}, nil + case mcp.ResourceReference: + return vmcp.CompletionRef{Type: r.Type, URI: r.URI}, nil + case map[string]any: + refType, _ := r["type"].(string) + name, _ := r["name"].(string) + uri, _ := r["uri"].(string) + if refType == "" { + return vmcp.CompletionRef{}, fmt.Errorf("%w: completion ref missing type", vmcp.ErrInvalidInput) + } + return vmcp.CompletionRef{Type: refType, Name: name, URI: uri}, nil + default: + return vmcp.CompletionRef{}, fmt.Errorf("%w: unsupported completion ref shape %T", vmcp.ErrInvalidInput, raw) + } +} + +// completionDenyMessage maps a completion ref type to the deny message of the +// underlying capability decision (prompt-get or resource-read), so an admission +// denial on completion reuses the same client-facing wording as a direct get/read. +func completionDenyMessage(refType string) string { + if refType == vmcp.CompletionRefTypeResource { + return vmcp.DenyMessageResourceRead + } + return vmcp.DenyMessagePromptGet } // enforceSessionBinding validates caller against the session's stored identity @@ -266,87 +552,19 @@ func (s *Server) coreResourceHandler( // ErrSessionOwnerUnknown. Unlike the legacy GetAdaptedTools handler, which terminates // only on ErrUnauthorizedCaller/ErrNilCaller, terminateOnBindingFailure terminates on // any rejection here — intentional fail-closed behavior, not a bug. -// -// A call that arrives while the session is still being registered (the client -// got the session ID from initialize, but OnRegisterSession has not yet written -// the binding) waits for the registration instead of failing; see -// waitForRegisteredSession. If the wait ends first, errSessionNotRegistered is -// returned and the session is kept. func (s *Server) enforceSessionBinding(ctx context.Context, sessionID string, caller *auth.Identity) error { - sess, err := s.waitForRegisteredSession(ctx, sessionID) - if err != nil { - return err + sess, ok := s.vmcpSessionMgr.GetMultiSession(ctx, sessionID) + if !ok { + return sessiontypes.ErrUnauthorizedCaller } // Single-key read: avoid GetMetadata()'s per-call full-map copy on this hot path. storedBinding, _ := sess.GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding) return vmcpsession.ValidateCaller(storedBinding, caller) } -// waitForRegisteredSession returns the session's MultiSession. When the session -// is not available and its stored record is a live placeholder (Generate wrote -// it, CreateSession has not finished on any pod), it waits for the registration -// (see waitWhileRegistrationPending) and then looks the session up once more. -func (s *Server) waitForRegisteredSession(ctx context.Context, sessionID string) (vmcpsession.MultiSession, error) { - if sess, ok := s.vmcpSessionMgr.GetMultiSession(ctx, sessionID); ok { - return sess, nil - } - if !s.registrationPending(ctx, sessionID) { - return nil, sessiontypes.ErrUnauthorizedCaller - } - if s.waitWhileRegistrationPending(ctx, sessionID) { - return nil, errSessionNotRegistered - } - if sess, ok := s.vmcpSessionMgr.GetMultiSession(ctx, sessionID); ok { - return sess, nil - } - return nil, sessiontypes.ErrUnauthorizedCaller -} - -// waitWhileRegistrationPending waits, for up to sessionRegistrationWaitTimeout, -// while the stored session record is a live placeholder whose registration has -// not written the identity binding yet. Each step reads the record once; the -// step starts at sessionRegistrationPollInterval and doubles up to -// sessionRegistrationMaxPollInterval. It reports whether the registration was -// still pending when it gave up (timeout or ctx done). -func (s *Server) waitWhileRegistrationPending(ctx context.Context, sessionID string) (stillPending bool) { - deadline := time.Now().Add(sessionRegistrationWaitTimeout) - step := sessionRegistrationPollInterval - for s.registrationPending(ctx, sessionID) { - if !time.Now().Before(deadline) { - return true - } - timer := time.NewTimer(min(step, time.Until(deadline))) - select { - case <-ctx.Done(): - timer.Stop() - return true - case <-timer.C: - } - step = min(2*step, sessionRegistrationMaxPollInterval) - } - return false -} - -// registrationPending reports whether the stored session record is a live -// placeholder: present, not terminated, and without an identity binding. -func (s *Server) registrationPending(ctx context.Context, sessionID string) bool { - metadata, err := s.sessionDataStorage.Load(ctx, sessionID) - if err != nil || metadata[sessionmanager.MetadataKeyTerminated] == sessionmanager.MetadataValTrue { - return false - } - _, bound := metadata[vmcpsession.MetadataKeyIdentityBinding] - return !bound -} - // terminateOnBindingFailure logs the hijack-prevention event and terminates the // session, mirroring GetAdaptedTools' handling of ErrUnauthorizedCaller/ErrNilCaller. -// A session whose registration has not finished yet is not terminated. func (s *Server) terminateOnBindingFailure(sessionID, capability string, err error) { - if errors.Is(err, errSessionNotRegistered) { - slog.Warn("session registration did not finish in time", - "session_id", sessionID, "capability", capability) - return - } slog.Warn("caller authorization failed, terminating session", "session_id", sessionID, "capability", capability, "error", err) if _, termErr := s.vmcpSessionMgr.Terminate(sessionID); termErr != nil { diff --git a/pkg/vmcp/server/serve_handlers_test.go b/pkg/vmcp/server/serve_handlers_test.go new file mode 100644 index 0000000000..b2e3d2c413 --- /dev/null +++ b/pkg/vmcp/server/serve_handlers_test.go @@ -0,0 +1,142 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// withParsedRequest installs a ParsedMCPRequest on ctx under the same context key +// GetParsedMCPRequest reads, mirroring what the transport parsing middleware does. +func withParsedRequest(ctx context.Context, parsed *mcpparser.ParsedMCPRequest) context.Context { + return context.WithValue(ctx, mcpparser.MCPRequestContextKey, parsed) +} + +// TestCoreToolHandler_UsesGateParsedArguments verifies that when the context carries the +// transport parse for this tools/call, the handler forwards the parser's argument map to +// core.CallTool — the same map the pre-dispatch authz gate decided on (#5845) — rather +// than the SDK's decode of the request. +func TestCoreToolHandler_UsesGateParsedArguments(t *testing.T) { + t.Parallel() + + const toolName = "t" + fc := &fakeCore{tools: []vmcp.Tool{{Name: toolName}}} + srv, sessionID, _ := registerServeSession(t, fc) + + // Disjoint keys so the assertion distinguishes a full replacement (correct) from a + // merge: a merge would leak "sdkonly" — a key the gate never authorized on — into the + // map forwarded to the backend, which is exactly the allow-then-execute-different-args + // hole #5845 closes. + ctx := withParsedRequest(t.Context(), &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + ResourceID: toolName, + Arguments: map[string]any{"src": "parser", "parseronly": "y"}, + }) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{ + Name: toolName, + Arguments: map[string]any{"src": "sdk", "sdkonly": "x"}, + }} + + res, err := srv.coreToolHandler(sessionID, toolName, "")(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + assert.False(t, res.IsError) + + gotArgs, _ := fc.lastCallToolArgs.Load().(map[string]any) + assert.Equal(t, map[string]any{"src": "parser", "parseronly": "y"}, gotArgs, + "the handler must forward exactly the gate's transport-parsed arguments, not the SDK decode") + assert.NotContains(t, gotArgs, "sdkonly", + "SDK-only keys the gate never authorized on must not reach the backend") +} + +// TestCoreToolHandler_FallsBackToSDKArguments verifies the handler keeps the SDK decode +// whenever no matching transport parse is available: a nil parse, a non-tools/call +// method, a resource-ID mismatch, or nil parsed arguments. Each of these must derive the +// call from the SDK's single decode so gate and dispatch cannot diverge. +func TestCoreToolHandler_FallsBackToSDKArguments(t *testing.T) { + t.Parallel() + + const toolName = "t" + + tests := []struct { + name string + parsed *mcpparser.ParsedMCPRequest + }{ + {"nil parse", nil}, + { + "method mismatch", + &mcpparser.ParsedMCPRequest{Method: "resources/read", ResourceID: toolName, Arguments: map[string]any{"src": "parser"}}, + }, + { + "resource-id mismatch", + &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: "other", Arguments: map[string]any{"src": "parser"}}, + }, + { + "nil arguments", + &mcpparser.ParsedMCPRequest{Method: "tools/call", ResourceID: toolName, Arguments: nil}, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + fc := &fakeCore{tools: []vmcp.Tool{{Name: toolName}}} + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := t.Context() + if tc.parsed != nil { + ctx = withParsedRequest(ctx, tc.parsed) + } + req := mcp.CallToolRequest{Params: mcp.CallToolParams{ + Name: toolName, + Arguments: map[string]any{"src": "sdk"}, + }} + + res, err := srv.coreToolHandler(sessionID, toolName, "")(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + assert.False(t, res.IsError) + + gotArgs, _ := fc.lastCallToolArgs.Load().(map[string]any) + assert.Equal(t, "sdk", gotArgs["src"], + "without a matching transport parse the handler must use the SDK decode") + }) + } +} + +// TestCoreToolHandler_NonObjectSDKArgsRejectedDespiteParse verifies the SDK type check +// runs first: a non-object SDK arguments payload is rejected with ErrInvalidInput even +// when a matching transport parse exists, so the gate's parse never masks a malformed +// request shape. +func TestCoreToolHandler_NonObjectSDKArgsRejectedDespiteParse(t *testing.T) { + t.Parallel() + + const toolName = "t" + fc := &fakeCore{tools: []vmcp.Tool{{Name: toolName}}} + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := withParsedRequest(t.Context(), &mcpparser.ParsedMCPRequest{ + Method: "tools/call", + ResourceID: toolName, + Arguments: map[string]any{"src": "parser"}, + }) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{ + Name: toolName, + Arguments: "not-an-object", + }} + + res, err := srv.coreToolHandler(sessionID, toolName, "")(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + assert.True(t, res.IsError, "a non-object SDK arguments payload must be rejected") + assert.Equal(t, int32(0), fc.callToolCalls.Load(), + "the core must not be reached when the SDK arguments are not an object") +} diff --git a/pkg/vmcp/server/serve_health_probe_test.go b/pkg/vmcp/server/serve_health_probe_test.go index 2950f6a4b3..3b39a80e90 100644 --- a/pkg/vmcp/server/serve_health_probe_test.go +++ b/pkg/vmcp/server/serve_health_probe_test.go @@ -11,11 +11,11 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/server" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/health" @@ -102,7 +102,8 @@ func TestServeSkipsHealthProbeSessions(t *testing.T) { assert.False(t, state.makeWithIDCalled.Load(), "a probe must not open backend connections") assert.Zero(t, fc.listToolsCalls.Load(), "a probe must not list capabilities") assert.Empty(t, sess.tools, "a probe must not get tools injected") - assert.Zero(t, srv.activeClientSessionCount(), "a probe must not be counted as an active session") + _, registered := srv.refreshSinks.Load(sess.id) + assert.False(t, registered, "a probe must not be registered for capability refreshes") } // TestPeerProbeHeaderCannotCall verifies that the client-controlled probe header @@ -127,7 +128,7 @@ func TestPeerProbeHeaderCannotCall(t *testing.T) { shared, err := transportsession.NewLocalSessionDataStorage(time.Hour) require.NoError(t, err) t.Cleanup(func() { _ = shared.Close() }) - pod := startCrossPod(t, shared, nil) + pod := startCrossPod(t, shared) sessionID := initSession(t, shared, pod, "alice") rec := httptest.NewRecorder() diff --git a/pkg/vmcp/server/serve_lifecycle_test.go b/pkg/vmcp/server/serve_lifecycle_test.go index 609be45eee..89ea57cf5a 100644 --- a/pkg/vmcp/server/serve_lifecycle_test.go +++ b/pkg/vmcp/server/serve_lifecycle_test.go @@ -149,8 +149,9 @@ func TestServeWithOptimizerStartsAndStopsCleanly(t *testing.T) { cfg := testMinimalServeConfig() cfg.SessionManagerConfig = &sessionmanager.FactoryConfig{ - Base: testMinimalFactory(), - OptimizerConfig: &optimizer.Config{}, + Base: testMinimalFactory(), + OptimizerConfig: &optimizer.Config{}, + AdvertiseFromCore: true, } srv, err := Serve(context.Background(), &stubVMCP{}, cfg) diff --git a/pkg/vmcp/server/serve_list_changed.go b/pkg/vmcp/server/serve_list_changed.go new file mode 100644 index 0000000000..afadf81319 --- /dev/null +++ b/pkg/vmcp/server/serve_list_changed.go @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "fmt" + "log/slog" + "sync" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp/headerforward" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// This file holds the tools/resources/prompts list_changed propagation added +// by #5748 (tools) and extended by #5969 (resources/resource templates and +// prompts): a persistent backend connection observing +// notifications/tools/list_changed, notifications/resources/list_changed, or +// notifications/prompts/list_changed invalidates the shared capability cache +// and resyncs the affected session's advertised set for that capability kind, +// so a client that already initialized eventually sees the backend's change +// reflected in tools/list, resources/list, resources/templates/list, or +// prompts/list without reconnecting. +// +// Scope: additions AND removals propagate for tools; for resources/resource +// templates/prompts, only ADDITIONS propagate today — see the add-only caveat +// on resyncSessionResources/resyncSessionPrompts and +// docs/arch/10-virtual-mcp-architecture.md. Cross-pod session restore +// (RestoreSession) does not thread a sink at all (no live ClientSession to +// resync there). + +// listChangedResyncWorker coalesces a session's list_changed resyncs for ONE +// capability kind (tools, resources, or prompts). buildListChangedSink builds +// one worker per (session, kind): a tools notification resyncing tools alone, +// without touching resources/prompts, matters because each capability's +// resync REPLACES that capability's session store, and re-deriving one kind +// must not force a redundant re-apply (and possible spurious downstream +// list_changed) of the others. +// +// The sink runs on the backend's receive-loop goroutine and must not block +// (vmcpsession.ListChangedSink contract), so trigger() only flips a dirty flag +// and, at most, starts ONE worker goroutine — it never does the resync inline +// and never spawns a goroutine per notification. At most one resync is ever in +// flight per (session, kind); notifications that arrive while one runs +// collapse into a single follow-up run. This bounds a misbehaving backend to +// O(1) concurrent resync work per (session, kind) regardless of how fast it +// emits notifications. +type listChangedResyncWorker struct { + // run performs exactly one resync using the given (server-lifetime) context. + // It is self-contained: liveness check, context reconstruction, and the + // re-derive-and-replace of the session's advertised set for this worker's + // capability kind (tools, resources, or prompts). + run func(ctx context.Context) + // baseCtx bounds every run to the server lifetime (cancelled on Stop) so a + // late notification cannot start work that outlives the server. + baseCtx context.Context + + mu sync.Mutex + running bool + dirty bool +} + +// trigger requests a resync. It is non-blocking and safe to call from the +// backend receive-loop goroutine: if a resync is already running it just marks +// the state dirty (coalescing), otherwise it starts the single worker goroutine. +func (w *listChangedResyncWorker) trigger() { + w.mu.Lock() + defer w.mu.Unlock() + if w.running { + w.dirty = true + return + } + w.running = true + go w.loop() +} + +// loop runs resyncs until no further notification arrived during the last run. +// It is the single worker goroutine started by trigger; it exits when idle so +// an idle session holds no goroutine. +func (w *listChangedResyncWorker) loop() { + for { + w.run(w.baseCtx) + + w.mu.Lock() + if !w.dirty { + w.running = false + w.mu.Unlock() + return + } + w.dirty = false + w.mu.Unlock() + } +} + +// buildListChangedSink returns the sink passed to +// SessionManager.CreateSession for a newly-registered session. The returned +// sink runs on the backend receive-loop goroutine and only hands work off to a +// per-(session, kind) coalescing worker (see listChangedResyncWorker) — it +// never does the cache purge or backend re-aggregation inline. +// +// It builds one listChangedResyncWorker per ChangeKind (tools, resources, +// prompts) rather than a single shared worker: each worker's run closure +// resyncs only its own capability kind, so a tools notification cannot force a +// spurious re-apply (and possible downstream list_changed) of the session's +// resources/prompts overlay, and vice versa. +// +// Each worker closes over the SDK ClientSession, the session ID, and the +// caller's identity AND per-request forwarded headers captured AT +// REGISTRATION TIME (issue #5748, decision B1/Option 2). The async resync +// reconstructs a context carrying that identity + those headers before calling +// into the core, because the capability cache key and the outbound backend +// authentication are derived from the CONTEXT (auth.IdentityFromContext / +// headerforward.ForwardedHeadersFromContext — see +// pkg/vmcp/aggregator/caching_aggregator.go and pkg/vmcp/client), NOT from an +// explicit identity argument. Without this the resync would enumerate +// backends unauthenticated and could advertise metadata the principal's own +// credentials would not surface (or wrongly drop credential-gated +// tools/resources/prompts), while replacing the correctly-scoped registration +// set. +// +// Token-staleness risk: the captured identity is a snapshot. If its upstream +// tokens are refreshed after registration (see #5323), a resync's +// core.ListTools/ListResources/ListPrompts call authorizes/aggregates using +// the (possibly stale) captured tokens, not the live per-request ones. This +// only affects the accuracy of the asynchronous resync's own view — every +// live call still authenticates via the fresh per-request identity through +// enforceSessionBinding, so staleness here cannot grant a call that would +// otherwise be denied. +func (s *Server) buildListChangedSink( + sessionID string, session server.ClientSession, identity *auth.Identity, forwardedHeaders map[string]string, +) vmcpsession.ListChangedSink { + newWorker := func(kind vmcpsession.ChangeKind) *listChangedResyncWorker { + return &listChangedResyncWorker{ + baseCtx: s.resyncBaseCtx, + run: func(ctx context.Context) { + s.runListChangedResync(ctx, sessionID, session, identity, forwardedHeaders, kind) + }, + } + } + workers := map[vmcpsession.ChangeKind]*listChangedResyncWorker{ + vmcpsession.KindTools: newWorker(vmcpsession.KindTools), + vmcpsession.KindResources: newWorker(vmcpsession.KindResources), + vmcpsession.KindPrompts: newWorker(vmcpsession.KindPrompts), + } + return func(_ context.Context, backendWorkloadID string, kind vmcpsession.ChangeKind) { + worker, ok := workers[kind] + if !ok { + slog.Debug("ignoring list_changed notification of unknown kind", + "session_id", sessionID, "backend_id", backendWorkloadID, "kind", kind) + return + } + slog.Debug("backend reported list_changed; scheduling session resync", + "session_id", sessionID, "backend_id", backendWorkloadID, "kind", kind) + worker.trigger() + } +} + +// runListChangedResync performs one coalesced resync for a session, for the +// given capability kind. It runs on the worker goroutine (off the receive +// loop). Order matters: +// 1. Liveness guard — skip (and do no work) if the session was terminated, so +// a storm of notifications for a dead session cannot drive re-aggregation. +// 2. Reconstruct the registration-time request context (identity + forwarded +// headers) atop the server-lifetime base context, so the cache key and +// backend auth match a real request from this principal. +// 3. Invalidate the capability cache so the re-derivation below re-sweeps the +// backend rather than reading the entry it just purged. +// 4. Re-derive and REPLACE the session's advertised set for kind: KindTools +// resyncs tools, KindResources resyncs both resources and resource +// templates (MCP 2025-11-25 has no separate wire method for template +// changes), and KindPrompts resyncs prompts. +func (s *Server) runListChangedResync( + baseCtx context.Context, + sessionID string, + session server.ClientSession, + identity *auth.Identity, + forwardedHeaders map[string]string, + kind vmcpsession.ChangeKind, +) { + // Liveness guard: if the session is gone (terminated/expired) there is + // nothing to resync, and doing the work would waste a full backend sweep. + if _, ok := s.vmcpSessionMgr.GetMultiSession(baseCtx, sessionID); !ok { + slog.Debug("skipping list_changed resync for terminated session", "session_id", sessionID, "kind", kind) + return + } + + ctx := auth.WithIdentity(baseCtx, identity) + if len(forwardedHeaders) > 0 { + ctx = headerforward.WithForwardedHeaders(ctx, forwardedHeaders) + } + + // Invalidate FIRST so the re-derivation below re-sweeps the backend instead + // of re-reading the stale cached aggregation. The purge is global (all + // identities, all kinds) — see InvalidateCapabilityCache — but coalescing + // bounds it to at most once per resync burst per kind. + s.core.InvalidateCapabilityCache() + + var err error + switch kind { + case vmcpsession.KindTools: + err = s.resyncSessionTools(ctx, session, sessionID, identity) + case vmcpsession.KindResources: + err = s.resyncSessionResources(ctx, session, sessionID, identity) + case vmcpsession.KindPrompts: + err = s.resyncSessionPrompts(ctx, session, sessionID, identity) + default: + slog.Debug("skipping resync for unknown list_changed kind", "session_id", sessionID, "kind", kind) + return + } + if err != nil { + slog.Warn("failed to resync session after backend list_changed", + "session_id", sessionID, "kind", kind, "error", err) + } +} + +// resyncSessionTools re-derives sessionID's advertised tool set (via +// serveSessionTools — the core, cache now cold, plus optimizer meta-tools when +// enabled) and REPLACES the SDK session's tool store with it, so a tool the +// backend removed (and the core therefore no longer advertises) disappears +// rather than lingering (unlike setSessionToolsDirect's registration-time +// MERGE, which only adds). The go-sdk server then auto-emits +// notifications/tools/list_changed to the downstream client, since serve.go +// enables WithToolCapabilities(true). +// +// ctx must already carry the resyncing principal's identity and forwarded +// headers (runListChangedResync builds it) so serveSessionTools -> +// core.ListTools enumerates backends with the correct credentials and cache key. +func (s *Server) resyncSessionTools( + ctx context.Context, session server.ClientSession, sessionID string, identity *auth.Identity, +) error { + tools, err := s.serveSessionTools(ctx, sessionID, identity) + if err != nil { + return fmt.Errorf("resync session tools: core ListTools for session %s: %w", sessionID, err) + } + + sessionWithTools, ok := session.(server.SessionWithTools) + if !ok { + return fmt.Errorf("resync session tools: session %s does not support per-session tools", sessionID) + } + + toolMap := make(map[string]server.ServerTool, len(tools)) + for _, tool := range tools { + toolMap[tool.Tool.Name] = tool + } + sessionWithTools.SetSessionTools(toolMap) + + slog.Debug("resynced session tools after backend list_changed", + "session_id", sessionID, "tool_count", len(tools)) + return nil +} + +// resyncSessionResources re-derives sessionID's advertised resources AND +// resource templates (via coreSessionResources/coreSessionResourceTemplates — +// the core, cache now cold) and REPLACES the SDK session's resource and +// resource-template stores with them, mirroring resyncSessionTools. Both are +// re-derived together because MCP 2025-11-25 has no separate wire method for +// resource-template changes: notifications/resources/list_changed covers both. +// +// ctx must already carry the resyncing principal's identity and forwarded +// headers (runListChangedResync builds it) so the coreSession* helpers -> +// core.ListResources/ListResourceTemplates enumerate backends with the +// correct credentials and cache key. +// +// Add-only removal limitation: unlike resyncSessionTools, replacing the +// overlay here propagates ADDITIONS but not REMOVALS. toolhive-core's +// mcpcompat per-session sync for resources/resource templates +// (syncSessionResources/syncSessionResourceTemplates) is add-only — unlike +// syncSessionTools, which also calls RemoveTools — so a backend-removed +// resource/template stays registered on the session's go-sdk server (listed +// until re-initialize) even though this overlay no longer contains it. The +// overlay itself IS replaced (not merged) here, so when toolhive-core gains +// removal reconciliation (tracked follow-up: stacklok/toolhive-core#184), +// removals start propagating with no change needed in this function. +// Advertising listChanged:true stays honest — notifications ARE emitted on +// change (except that a pure removal which EMPTIES the overlay issues no Add* +// and therefore triggers no downstream list_changed, since go-sdk notifies +// only via Add*/RemoveTools); listChanged promises notification, not list +// minimization — but until that follow-up lands, a refetch after a pure +// removal returns a stale superset. +func (s *Server) resyncSessionResources( + ctx context.Context, session server.ClientSession, sessionID string, identity *auth.Identity, +) error { + resources, err := s.coreSessionResources(ctx, sessionID, identity) + if err != nil { + return fmt.Errorf("resync session resources: core ListResources for session %s: %w", sessionID, err) + } + sessionWithResources, ok := session.(server.SessionWithResources) + if !ok { + return fmt.Errorf("resync session resources: session %s does not support per-session resources", sessionID) + } + resourceMap := make(map[string]server.ServerResource, len(resources)) + for _, res := range resources { + resourceMap[res.Resource.URI] = res + } + sessionWithResources.SetSessionResources(resourceMap) + + templates, err := s.coreSessionResourceTemplates(ctx, sessionID, identity) + if err != nil { + return fmt.Errorf("resync session resource templates: core ListResourceTemplates for session %s: %w", sessionID, err) + } + sessionWithTemplates, ok := session.(server.SessionWithResourceTemplates) + if !ok { + return fmt.Errorf( + "resync session resource templates: session %s does not support per-session resource templates", sessionID) + } + templateMap := make(map[string]server.ServerResourceTemplate, len(templates)) + for _, tmpl := range templates { + templateMap[tmpl.Template.URITemplate] = tmpl + } + sessionWithTemplates.SetSessionResourceTemplates(templateMap) + + slog.Debug("resynced session resources after backend list_changed", + "session_id", sessionID, "resource_count", len(resources), "resource_template_count", len(templates)) + return nil +} + +// resyncSessionPrompts re-derives sessionID's advertised prompt set (via +// coreSessionPrompts — the core, cache now cold) and REPLACES the SDK +// session's prompt store with it, mirroring resyncSessionTools. +// +// ctx must already carry the resyncing principal's identity and forwarded +// headers (runListChangedResync builds it) so coreSessionPrompts -> +// core.ListPrompts enumerates backends with the correct credentials and cache +// key. +// +// Add-only removal limitation: as with resyncSessionResources, replacing the +// overlay here propagates ADDITIONS but not REMOVALS. toolhive-core's +// mcpcompat per-session sync for prompts (syncSessionPrompts) is add-only, so +// a backend-removed prompt stays registered on the session's go-sdk server +// (listed until re-initialize) even though this overlay no longer contains +// it. The overlay itself IS replaced (not merged) here, so when toolhive-core +// gains removal reconciliation (tracked follow-up: stacklok/toolhive-core#184), +// removals start propagating with no change needed in this function. +// Advertising listChanged:true stays honest — notifications ARE emitted on +// change (except that a pure removal which EMPTIES the overlay issues no Add* +// and therefore triggers no downstream list_changed, since go-sdk notifies +// only via Add*/RemoveTools); listChanged promises notification, not list +// minimization — but until that follow-up lands, a refetch after a pure +// removal returns a stale superset. +func (s *Server) resyncSessionPrompts( + ctx context.Context, session server.ClientSession, sessionID string, identity *auth.Identity, +) error { + prompts, err := s.coreSessionPrompts(ctx, sessionID, identity) + if err != nil { + return fmt.Errorf("resync session prompts: core ListPrompts for session %s: %w", sessionID, err) + } + + sessionWithPrompts, ok := session.(server.SessionWithPrompts) + if !ok { + return fmt.Errorf("resync session prompts: session %s does not support per-session prompts", sessionID) + } + + promptMap := make(map[string]server.ServerPrompt, len(prompts)) + for _, p := range prompts { + promptMap[p.Prompt.Name] = p + } + sessionWithPrompts.SetSessionPrompts(promptMap) + + slog.Debug("resynced session prompts after backend list_changed", + "session_id", sessionID, "prompt_count", len(prompts)) + return nil +} diff --git a/pkg/vmcp/server/serve_list_changed_test.go b/pkg/vmcp/server/serve_list_changed_test.go new file mode 100644 index 0000000000..654c10a901 --- /dev/null +++ b/pkg/vmcp/server/serve_list_changed_test.go @@ -0,0 +1,641 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/headerforward" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// fakeToolsSession is a minimal server.ClientSession + server.SessionWithTools +// fake used to test resyncSessionTools/runListChangedResync without a live SDK +// session. It records every SetSessionTools call so tests can assert +// replacement semantics. +type fakeToolsSession struct { + id string + + mu sync.Mutex + tools map[string]server.ServerTool + setSessionToolsCalls int +} + +func (*fakeToolsSession) Initialize() {} +func (*fakeToolsSession) Initialized() bool { return true } +func (*fakeToolsSession) NotificationChannel() chan<- mcp.JSONRPCNotification { return nil } +func (f *fakeToolsSession) SessionID() string { return f.id } + +func (f *fakeToolsSession) GetSessionTools() map[string]server.ServerTool { + f.mu.Lock() + defer f.mu.Unlock() + out := make(map[string]server.ServerTool, len(f.tools)) + for k, v := range f.tools { + out[k] = v + } + return out +} + +func (f *fakeToolsSession) SetSessionTools(tools map[string]server.ServerTool) { + f.mu.Lock() + defer f.mu.Unlock() + f.tools = tools + f.setSessionToolsCalls++ +} + +var _ server.ClientSession = (*fakeToolsSession)(nil) +var _ server.SessionWithTools = (*fakeToolsSession)(nil) + +// fakeCapsSession is a minimal server.ClientSession + SessionWithResources + +// SessionWithResourceTemplates + SessionWithPrompts fake used to test +// resyncSessionResources/resyncSessionPrompts/runListChangedResync without a +// live SDK session. It records every SetSession* call so tests can assert +// replacement semantics (#5969, mirroring fakeToolsSession for #5748). +type fakeCapsSession struct { + id string + + mu sync.Mutex + resources map[string]server.ServerResource + resourceTemplates map[string]server.ServerResourceTemplate + prompts map[string]server.ServerPrompt + setSessionResourcesCalls int + setSessionResourceTemplatesCalls int + setSessionPromptsCalls int +} + +func (*fakeCapsSession) Initialize() {} +func (*fakeCapsSession) Initialized() bool { return true } +func (*fakeCapsSession) NotificationChannel() chan<- mcp.JSONRPCNotification { return nil } +func (f *fakeCapsSession) SessionID() string { return f.id } + +func (f *fakeCapsSession) GetSessionResources() map[string]server.ServerResource { + f.mu.Lock() + defer f.mu.Unlock() + out := make(map[string]server.ServerResource, len(f.resources)) + for k, v := range f.resources { + out[k] = v + } + return out +} + +func (f *fakeCapsSession) SetSessionResources(resources map[string]server.ServerResource) { + f.mu.Lock() + defer f.mu.Unlock() + f.resources = resources + f.setSessionResourcesCalls++ +} + +func (f *fakeCapsSession) GetSessionResourceTemplates() map[string]server.ServerResourceTemplate { + f.mu.Lock() + defer f.mu.Unlock() + out := make(map[string]server.ServerResourceTemplate, len(f.resourceTemplates)) + for k, v := range f.resourceTemplates { + out[k] = v + } + return out +} + +func (f *fakeCapsSession) SetSessionResourceTemplates(templates map[string]server.ServerResourceTemplate) { + f.mu.Lock() + defer f.mu.Unlock() + f.resourceTemplates = templates + f.setSessionResourceTemplatesCalls++ +} + +func (f *fakeCapsSession) GetSessionPrompts() map[string]server.ServerPrompt { + f.mu.Lock() + defer f.mu.Unlock() + out := make(map[string]server.ServerPrompt, len(f.prompts)) + for k, v := range f.prompts { + out[k] = v + } + return out +} + +func (f *fakeCapsSession) SetSessionPrompts(prompts map[string]server.ServerPrompt) { + f.mu.Lock() + defer f.mu.Unlock() + f.prompts = prompts + f.setSessionPromptsCalls++ +} + +var _ server.ClientSession = (*fakeCapsSession)(nil) +var _ server.SessionWithResources = (*fakeCapsSession)(nil) +var _ server.SessionWithResourceTemplates = (*fakeCapsSession)(nil) +var _ server.SessionWithPrompts = (*fakeCapsSession)(nil) + +// resourcesOnlySession implements server.ClientSession + SessionWithResources +// but NOT SessionWithResourceTemplates, so a resyncSessionResources test can +// prove the templates-half type assertion fails loudly. It provides its own +// resource methods (rather than embedding fakeCapsSession, which would also +// promote the template/prompt methods and satisfy those interfaces). +type resourcesOnlySession struct { + fake fakeCapsSession +} + +func (*resourcesOnlySession) Initialize() {} +func (*resourcesOnlySession) Initialized() bool { return true } +func (*resourcesOnlySession) NotificationChannel() chan<- mcp.JSONRPCNotification { return nil } +func (s *resourcesOnlySession) SessionID() string { return s.fake.id } +func (s *resourcesOnlySession) GetSessionResources() map[string]server.ServerResource { + return s.fake.GetSessionResources() +} +func (s *resourcesOnlySession) SetSessionResources(resources map[string]server.ServerResource) { + s.fake.SetSessionResources(resources) +} + +var _ server.ClientSession = (*resourcesOnlySession)(nil) +var _ server.SessionWithResources = (*resourcesOnlySession)(nil) + +// stubSessionManager is a minimal SessionManager for the #5748 resync tests. +// Only GetMultiSession is meaningful (its returned liveness bool drives the +// resync liveness guard); the rest satisfy the interface and fail loudly if a +// test reaches them unexpectedly. +type stubSessionManager struct { + alive bool +} + +func (m *stubSessionManager) GetMultiSession(context.Context, string) (vmcpsession.MultiSession, bool) { + return nil, m.alive +} +func (*stubSessionManager) Generate() string { panic("stubSessionManager: Generate unexpected") } +func (*stubSessionManager) Validate(string) (bool, error) { + panic("stubSessionManager: Validate unexpected") +} +func (*stubSessionManager) Terminate(string) (bool, error) { return false, nil } +func (*stubSessionManager) CreateSession( + context.Context, string, vmcpsession.ListChangedSink, +) (vmcpsession.MultiSession, error) { + panic("stubSessionManager: CreateSession unexpected") +} +func (*stubSessionManager) DecorateSession(string, func(vmcpsession.MultiSession) vmcpsession.MultiSession) error { + panic("stubSessionManager: DecorateSession unexpected") +} +func (*stubSessionManager) NotifyBackendExpired(string, string, map[string]string) {} + +var _ SessionManager = (*stubSessionManager)(nil) + +// TestResyncSessionTools_ReplacesRatherThanMerges verifies (#5748) that +// resyncSessionTools REPLACES the session's tool store with the freshly +// core-derived set — unlike setSessionToolsDirect's registration-time MERGE — +// so a tool the backend removed (and the core therefore no longer advertises) +// disappears rather than lingering. +func TestResyncSessionTools_ReplacesRatherThanMerges(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "kept"}, {Name: "added"}}} + srv := &Server{core: fc} + + sess := &fakeToolsSession{id: "sess-1", tools: map[string]server.ServerTool{ + "kept": {Tool: mcp.Tool{Name: "kept"}}, + "removed": {Tool: mcp.Tool{Name: "removed"}}, // must disappear after resync + }} + + err := srv.resyncSessionTools(context.Background(), sess, "sess-1", nil) + require.NoError(t, err) + + got := sess.GetSessionTools() + assert.Len(t, got, 2) + assert.Contains(t, got, "kept") + assert.Contains(t, got, "added") + assert.NotContains(t, got, "removed", "resync must REPLACE the tool store, dropping a no-longer-advertised tool") + assert.Equal(t, 1, sess.setSessionToolsCalls) +} + +// TestResyncSessionTools_SessionWithoutToolSupport verifies resyncSessionTools +// fails loudly (rather than silently no-opping) when the session does not +// implement server.SessionWithTools. +func TestResyncSessionTools_SessionWithoutToolSupport(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv := &Server{core: fc} + + // plainSession implements only server.ClientSession, not SessionWithTools. + type plainSession struct{ server.ClientSession } + sess := plainSession{&fakeToolsSession{id: "sess-1"}} + + err := srv.resyncSessionTools(context.Background(), sess, "sess-1", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not support per-session tools") +} + +// TestResyncSessionResources_SessionWithoutResourceSupport verifies (#5969) +// resyncSessionResources fails loudly (rather than silently no-opping) when +// the session does not implement server.SessionWithResources. +func TestResyncSessionResources_SessionWithoutResourceSupport(t *testing.T) { + t.Parallel() + + fc := &fakeCore{resources: []vmcp.Resource{{Name: "r", URI: "r"}}} + srv := &Server{core: fc} + + // plainSession implements only server.ClientSession, not SessionWithResources. + type plainSession struct{ server.ClientSession } + sess := plainSession{&fakeCapsSession{id: "sess-1"}} + + err := srv.resyncSessionResources(context.Background(), sess, "sess-1", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not support per-session resources") +} + +// TestResyncSessionResources_SessionWithoutTemplateSupport verifies (#5969) +// resyncSessionResources fails loudly when the session supports per-session +// resources but not resource templates (SessionWithResourceTemplates). The +// resource store is set before the template assertion fails, so this pins the +// exact error surfaced for the templates half. +func TestResyncSessionResources_SessionWithoutTemplateSupport(t *testing.T) { + t.Parallel() + + fc := &fakeCore{resources: []vmcp.Resource{{Name: "r", URI: "r"}}} + srv := &Server{core: fc} + + // resourcesOnlySession implements SessionWithResources but NOT + // SessionWithResourceTemplates. + sess := &resourcesOnlySession{fake: fakeCapsSession{id: "sess-1"}} + + err := srv.resyncSessionResources(context.Background(), sess, "sess-1", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not support per-session resource templates") +} + +// TestResyncSessionPrompts_SessionWithoutPromptSupport verifies (#5969) +// resyncSessionPrompts fails loudly (rather than silently no-opping) when the +// session does not implement server.SessionWithPrompts. +func TestResyncSessionPrompts_SessionWithoutPromptSupport(t *testing.T) { + t.Parallel() + + fc := &fakeCore{prompts: []vmcp.Prompt{{Name: "p"}}} + srv := &Server{core: fc} + + // plainSession implements only server.ClientSession, not SessionWithPrompts. + type plainSession struct{ server.ClientSession } + sess := plainSession{&fakeCapsSession{id: "sess-1"}} + + err := srv.resyncSessionPrompts(context.Background(), sess, "sess-1", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not support per-session prompts") +} + +// TestResyncSessionResources_ReplacesRatherThanMerges verifies (#5969) that +// resyncSessionResources REPLACES both the session's resource store AND its +// resource-template store with the freshly core-derived sets — unlike +// setSessionResourcesDirect/setSessionResourceTemplatesDirect's +// registration-time MERGE — so a resource/template the backend removed (and +// the core therefore no longer advertises) disappears rather than lingering. +func TestResyncSessionResources_ReplacesRatherThanMerges(t *testing.T) { + t.Parallel() + + fc := &fakeCore{ + resources: []vmcp.Resource{{Name: "kept", URI: "kept"}, {Name: "added", URI: "added"}}, + resourceTemplates: []vmcp.ResourceTemplate{{Name: "kept-tmpl", URITemplate: "kept-tmpl"}}, + } + srv := &Server{core: fc} + + sess := &fakeCapsSession{ + id: "sess-1", + resources: map[string]server.ServerResource{ + "kept": {Resource: mcp.Resource{Name: "kept", URI: "kept"}}, + "removed": {Resource: mcp.Resource{Name: "removed", URI: "removed"}}, // must disappear + }, + resourceTemplates: map[string]server.ServerResourceTemplate{ + "kept-tmpl": {Template: mcp.ResourceTemplate{Name: "kept-tmpl", URITemplate: "kept-tmpl"}}, + "removed-tmpl": {Template: mcp.ResourceTemplate{Name: "removed-tmpl", URITemplate: "removed-tmpl"}}, // must disappear + }, + } + + err := srv.resyncSessionResources(context.Background(), sess, "sess-1", nil) + require.NoError(t, err) + + gotResources := sess.GetSessionResources() + assert.Len(t, gotResources, 2) + assert.Contains(t, gotResources, "kept") + assert.Contains(t, gotResources, "added") + assert.NotContains(t, gotResources, "removed", + "resync must REPLACE the resource store, dropping a no-longer-advertised resource") + assert.Equal(t, 1, sess.setSessionResourcesCalls) + + gotTemplates := sess.GetSessionResourceTemplates() + assert.Len(t, gotTemplates, 1) + assert.Contains(t, gotTemplates, "kept-tmpl") + assert.NotContains(t, gotTemplates, "removed-tmpl", + "resync must REPLACE the resource-template store, dropping a no-longer-advertised template") + assert.Equal(t, 1, sess.setSessionResourceTemplatesCalls) +} + +// TestResyncSessionPrompts_ReplacesRatherThanMerges verifies (#5969) that +// resyncSessionPrompts REPLACES the session's prompt store with the freshly +// core-derived set — unlike setSessionPromptsDirect's registration-time +// MERGE — so a prompt the backend removed (and the core therefore no longer +// advertises) disappears rather than lingering. +func TestResyncSessionPrompts_ReplacesRatherThanMerges(t *testing.T) { + t.Parallel() + + fc := &fakeCore{prompts: []vmcp.Prompt{{Name: "kept"}, {Name: "added"}}} + srv := &Server{core: fc} + + sess := &fakeCapsSession{id: "sess-1", prompts: map[string]server.ServerPrompt{ + "kept": {Prompt: mcp.Prompt{Name: "kept"}}, + "removed": {Prompt: mcp.Prompt{Name: "removed"}}, // must disappear after resync + }} + + err := srv.resyncSessionPrompts(context.Background(), sess, "sess-1", nil) + require.NoError(t, err) + + got := sess.GetSessionPrompts() + assert.Len(t, got, 2) + assert.Contains(t, got, "kept") + assert.Contains(t, got, "added") + assert.NotContains(t, got, "removed", "resync must REPLACE the prompt store, dropping a no-longer-advertised prompt") + assert.Equal(t, 1, sess.setSessionPromptsCalls) +} + +// TestListChangedResyncWorker_CoalescesAndNonBlocking verifies the B-HIGH-2 +// hand-off: trigger() returns immediately, never spawns a goroutine per call, +// runs at most one resync at a time, and coalesces concurrent triggers into a +// single follow-up run (dirty flag). +func TestListChangedResyncWorker_CoalescesAndNonBlocking(t *testing.T) { + t.Parallel() + + var ( + runs atomic.Int32 + inFlight atomic.Int32 + maxSeen atomic.Int32 + release = make(chan struct{}) + ) + firstStarted := make(chan struct{}, 1) + + w := &listChangedResyncWorker{ + baseCtx: context.Background(), + run: func(context.Context) { + n := inFlight.Add(1) + for { + m := maxSeen.Load() + if n <= m || maxSeen.CompareAndSwap(m, n) { + break + } + } + if runs.Add(1) == 1 { + firstStarted <- struct{}{} + <-release // hold the first run so subsequent triggers coalesce + } + inFlight.Add(-1) + }, + } + + // First trigger starts the (blocked) worker goroutine. + w.trigger() + select { + case <-firstStarted: + case <-time.After(5 * time.Second): + t.Fatal("worker did not start") + } + + // While the first run is blocked, fire many more triggers. Each must return + // immediately (non-blocking) and collapse into a single dirty re-run. + for range 50 { + w.trigger() + } + + close(release) // let the first run finish; one coalesced re-run should follow + + require.Eventually(t, func() bool { + w.mu.Lock() + defer w.mu.Unlock() + return !w.running + }, 5*time.Second, 5*time.Millisecond, "worker should become idle") + + assert.Equal(t, int32(1), maxSeen.Load(), "at most one resync may run at a time") + assert.Equal(t, int32(2), runs.Load(), + "50 triggers during one in-flight run must coalesce into exactly one follow-up run") +} + +// TestRunListChangedResync verifies (#5748, extended by #5969) the resync +// body run off the receive loop, per capability kind: it skips terminated +// sessions (liveness guard), and for a live session it invalidates the cache +// and reconstructs a context carrying the captured identity + forwarded +// headers (B-HIGH-1) before re-deriving/replacing ONLY the store belonging to +// the notification's kind. +func TestRunListChangedResync(t *testing.T) { + t.Parallel() + + const fwdKey = "X-Tenant" + identity := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: "alice"}} + fwd := map[string]string{fwdKey: "acme"} + + tests := []struct { + name string + kind vmcpsession.ChangeKind + }{ + {"KindTools", vmcpsession.KindTools}, + {"KindResources", vmcpsession.KindResources}, + {"KindPrompts", vmcpsession.KindPrompts}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + t.Run("terminated session: skipped entirely", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + tools: []vmcp.Tool{{Name: "fresh"}}, + resources: []vmcp.Resource{{Name: "fresh", URI: "fresh"}}, + resourceTemplates: []vmcp.ResourceTemplate{{Name: "fresh", URITemplate: "fresh"}}, + prompts: []vmcp.Prompt{{Name: "fresh"}}, + } + srv := &Server{core: fc, vmcpSessionMgr: &stubSessionManager{alive: false}} + toolsSess := &fakeToolsSession{id: "sess-1"} + capsSess := &fakeCapsSession{id: "sess-1"} + + var sess server.ClientSession = toolsSess + if tc.kind != vmcpsession.KindTools { + sess = capsSess + } + + srv.runListChangedResync(context.Background(), "sess-1", sess, identity, fwd, tc.kind) + + assert.Equal(t, int32(0), fc.invalidateCacheCalls.Load(), "terminated session must not invalidate the cache") + assert.Equal(t, int32(0), fc.listToolsCalls.Load(), "terminated session must not re-aggregate") + assert.Equal(t, int32(0), fc.listResourcesCalls.Load(), "terminated session must not re-aggregate") + assert.Equal(t, int32(0), fc.listResourceTemplatesCalls.Load(), "terminated session must not re-aggregate") + assert.Equal(t, int32(0), fc.listPromptsCalls.Load(), "terminated session must not re-aggregate") + }) + + t.Run("live session: invalidates, reconstructs ctx, replaces only this kind", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + tools: []vmcp.Tool{{Name: "fresh"}}, + resources: []vmcp.Resource{{Name: "fresh", URI: "fresh"}}, + resourceTemplates: []vmcp.ResourceTemplate{{Name: "fresh", URITemplate: "fresh"}}, + prompts: []vmcp.Prompt{{Name: "fresh"}}, + } + srv := &Server{core: fc, vmcpSessionMgr: &stubSessionManager{alive: true}} + toolsSess := &fakeToolsSession{id: "sess-1", tools: map[string]server.ServerTool{ + "stale": {Tool: mcp.Tool{Name: "stale"}}, + }} + capsSess := &fakeCapsSession{ + id: "sess-1", + resources: map[string]server.ServerResource{ + "stale": {Resource: mcp.Resource{Name: "stale", URI: "stale"}}, + }, + resourceTemplates: map[string]server.ServerResourceTemplate{ + "stale": {Template: mcp.ResourceTemplate{Name: "stale", URITemplate: "stale"}}, + }, + prompts: map[string]server.ServerPrompt{ + "stale": {Prompt: mcp.Prompt{Name: "stale"}}, + }, + } + + var sess server.ClientSession = toolsSess + if tc.kind != vmcpsession.KindTools { + sess = capsSess + } + + srv.runListChangedResync(context.Background(), "sess-1", sess, identity, fwd, tc.kind) + + assert.Equal(t, int32(1), fc.invalidateCacheCalls.Load(), "live session must invalidate the cache once") + + var resyncCtx context.Context + switch tc.kind { + case vmcpsession.KindTools: + assert.Equal(t, int32(1), fc.listToolsCalls.Load()) + assert.Equal(t, int32(0), fc.listResourcesCalls.Load()) + assert.Equal(t, int32(0), fc.listResourceTemplatesCalls.Load()) + assert.Equal(t, int32(0), fc.listPromptsCalls.Load()) + assert.Equal(t, 1, toolsSess.setSessionToolsCalls) + got := toolsSess.GetSessionTools() + require.Len(t, got, 1) + assert.Contains(t, got, "fresh") + assert.NotContains(t, got, "stale", "resync must replace the stale set") + box := fc.lastListToolsCtx.Load() + require.NotNil(t, box, "ListTools must have been called") + resyncCtx = box.ctx + case vmcpsession.KindResources: + assert.Equal(t, int32(0), fc.listToolsCalls.Load()) + assert.Equal(t, int32(1), fc.listResourcesCalls.Load()) + assert.Equal(t, int32(1), fc.listResourceTemplatesCalls.Load()) + assert.Equal(t, int32(0), fc.listPromptsCalls.Load()) + assert.Equal(t, 1, capsSess.setSessionResourcesCalls) + assert.Equal(t, 1, capsSess.setSessionResourceTemplatesCalls) + gotResources := capsSess.GetSessionResources() + require.Len(t, gotResources, 1) + assert.Contains(t, gotResources, "fresh") + assert.NotContains(t, gotResources, "stale", "resync must replace the stale set") + gotTemplates := capsSess.GetSessionResourceTemplates() + require.Len(t, gotTemplates, 1) + assert.Contains(t, gotTemplates, "fresh") + assert.NotContains(t, gotTemplates, "stale", "resync must replace the stale set") + box := fc.lastListResourcesCtx.Load() + require.NotNil(t, box, "ListResources must have been called") + resyncCtx = box.ctx + case vmcpsession.KindPrompts: + assert.Equal(t, int32(0), fc.listToolsCalls.Load()) + assert.Equal(t, int32(0), fc.listResourcesCalls.Load()) + assert.Equal(t, int32(0), fc.listResourceTemplatesCalls.Load()) + assert.Equal(t, int32(1), fc.listPromptsCalls.Load()) + assert.Equal(t, 1, capsSess.setSessionPromptsCalls) + got := capsSess.GetSessionPrompts() + require.Len(t, got, 1) + assert.Contains(t, got, "fresh") + assert.NotContains(t, got, "stale", "resync must replace the stale set") + box := fc.lastListPromptsCtx.Load() + require.NotNil(t, box, "ListPrompts must have been called") + resyncCtx = box.ctx + } + + // B-HIGH-1 (all kinds): the resync must call the core with a context + // carrying the captured identity AND forwarded headers, so the cache + // key and outbound backend auth match a real request from this + // principal. Asserting this for every kind guards a future refactor + // from silently dropping identity/header scoping on the + // resources/prompts paths. + require.NotNil(t, resyncCtx) + gotID, ok := auth.IdentityFromContext(resyncCtx) + require.True(t, ok, "resync context must carry the captured identity") + assert.Equal(t, "alice", gotID.Subject) + assert.Equal(t, "acme", headerforward.ForwardedHeadersFromContext(resyncCtx)[fwdKey], + "resync context must carry the captured forwarded headers") + }) + }) + } +} + +// TestBuildListChangedSink_DispatchesByKind verifies (#5969) the sink's +// per-kind worker dispatch: a KindResources or KindPrompts notification +// triggers its own worker (invalidating the cache and resyncing only that +// kind's store), and an unknown ChangeKind is dropped synchronously (no +// worker goroutine, no cache invalidation, no resync). +func TestBuildListChangedSink_DispatchesByKind(t *testing.T) { + t.Parallel() + + t.Run("KindResources triggers a resources-only resync", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "fresh", URI: "fresh"}}} + srv := &Server{ + core: fc, + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + } + sess := &fakeCapsSession{id: "sess-1"} + + sink := srv.buildListChangedSink("sess-1", sess, nil, nil) + sink(context.Background(), "backend-1", vmcpsession.KindResources) + + require.Eventually(t, func() bool { return fc.invalidateCacheCalls.Load() >= 1 }, + 2*time.Second, 10*time.Millisecond, "KindResources must invalidate the cache") + require.Eventually(t, func() bool { return len(sess.GetSessionResources()) == 1 }, + 2*time.Second, 10*time.Millisecond, "KindResources must resync the session's resources") + assert.Equal(t, int32(0), fc.listToolsCalls.Load(), "a resources notification must not resync tools") + assert.Equal(t, int32(0), fc.listPromptsCalls.Load(), "a resources notification must not resync prompts") + }) + + t.Run("KindPrompts triggers a prompts-only resync", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{prompts: []vmcp.Prompt{{Name: "fresh"}}} + srv := &Server{ + core: fc, + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + } + sess := &fakeCapsSession{id: "sess-1"} + + sink := srv.buildListChangedSink("sess-1", sess, nil, nil) + sink(context.Background(), "backend-1", vmcpsession.KindPrompts) + + require.Eventually(t, func() bool { return fc.invalidateCacheCalls.Load() >= 1 }, + 2*time.Second, 10*time.Millisecond, "KindPrompts must invalidate the cache") + require.Eventually(t, func() bool { return len(sess.GetSessionPrompts()) == 1 }, + 2*time.Second, 10*time.Millisecond, "KindPrompts must resync the session's prompts") + assert.Equal(t, int32(0), fc.listToolsCalls.Load(), "a prompts notification must not resync tools") + assert.Equal(t, int32(0), fc.listResourcesCalls.Load(), "a prompts notification must not resync resources") + }) + + t.Run("unknown kind is dropped synchronously", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{tools: []vmcp.Tool{{Name: "fresh"}}} + srv := &Server{ + core: fc, + vmcpSessionMgr: &stubSessionManager{alive: true}, + resyncBaseCtx: context.Background(), + } + sess := &fakeToolsSession{id: "sess-1"} + + sink := srv.buildListChangedSink("sess-1", sess, nil, nil) + sink(context.Background(), "backend-1", vmcpsession.ChangeKind("unknown")) + + assert.Equal(t, int32(0), fc.invalidateCacheCalls.Load(), "unknown kind must not invalidate the cache") + assert.Equal(t, 0, sess.setSessionToolsCalls, "unknown kind must not resync anything") + }) +} diff --git a/pkg/vmcp/server/serve_optimizer.go b/pkg/vmcp/server/serve_optimizer.go index f00ad34f5d..0f5a6cae00 100644 --- a/pkg/vmcp/server/serve_optimizer.go +++ b/pkg/vmcp/server/serve_optimizer.go @@ -9,9 +9,8 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" diff --git a/pkg/vmcp/server/serve_optimizer_test.go b/pkg/vmcp/server/serve_optimizer_test.go index e15d35ddcc..0b1a5948ef 100644 --- a/pkg/vmcp/server/serve_optimizer_test.go +++ b/pkg/vmcp/server/serve_optimizer_test.go @@ -13,12 +13,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/ratelimit" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/pkg/vmcp/server/serve_refresh.go b/pkg/vmcp/server/serve_refresh.go index 8983cf415c..2a05ef6735 100644 --- a/pkg/vmcp/server/serve_refresh.go +++ b/pkg/vmcp/server/serve_refresh.go @@ -8,24 +8,22 @@ import ( "log/slog" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/health" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) -// This file refreshes the tools and resources of the sessions registered on this -// pod when the set of exposed backends changes. Each session's tool and resource -// store is filled once (at registration or by lazy injection), so without a -// refresh a session keeps the tool set it saw first. A refresh clears the stores -// and sends list_changed; the next tools/list or resources/list fills them again -// from the core with the live request identity (lazyInjectSessionTools / -// lazyInjectSessionResources). +// This file refreshes the capabilities of the sessions registered on this pod +// when the set of exposed backends changes (a backend's health crosses the +// advertise boundary, or a dynamic registry adds or removes a backend). Upstream +// resyncs a session only when a connected backend sends list_changed +// (serve_list_changed.go); a backend that comes back or joins the group sends +// nothing, so without this a session keeps the capabilities it saw first. // -// Scope: only sessions registered on this pod (activeClientSessions) are -// refreshed, as they are the only ones that can hold a listen stream here. Every -// pod runs its own health monitor and refreshes its own sessions. +// A refresh triggers each session's list_changed sink for tools, resources and +// prompts, the same coalesced resync a backend notification starts. Scope: only +// sessions registered on this pod (refreshSinks). Every pod runs its own health +// monitor and refreshes its own sessions. const ( // registryVersionCheckInterval is how often a dynamic backend registry is @@ -59,10 +57,10 @@ func (s *Server) startCapabilityRefresh() func(context.Context) error { } // scheduleOnExposureChange schedules a refresh when a backend crosses the -// exposure boundary (see health.IsStatusEligibleForExposure). +// advertise boundary (see health.ShouldAdvertise). func scheduleOnExposureChange(sched *refreshScheduler, event health.StatusChangeEvent) { - wasEligible := health.IsStatusEligibleForExposure(event.PreviousStatus) - isEligible := health.IsStatusEligibleForExposure(event.Status) + wasEligible := health.ShouldAdvertise(event.PreviousStatus) + isEligible := health.ShouldAdvertise(event.Status) if wasEligible == isEligible { return } @@ -96,31 +94,28 @@ func watchRegistryVersion( } } -// refreshSessionCapabilities clears the tool and resource store of every session -// registered on this pod and tells its client to list them again. It is the -// refreshScheduler run function; the mode is only logged, because each store is -// refilled from the core, which already reflects the change. +// refreshSessionCapabilities triggers the list_changed sink of every session +// registered on this pod, and drops the sinks of sessions that ended. It is the +// refreshScheduler run function; the mode is only logged, because each resync +// re-derives the set from the core, which already reflects the change. func (s *Server) refreshSessionCapabilities(trigger string, mode refreshTargetMode) { + ctx := context.Background() refreshed := 0 - for _, sess := range s.activeClientSessions.snapshot() { - sessionID := sess.SessionID() - if withTools, ok := sess.(server.SessionWithTools); ok { - withTools.SetSessionTools(nil) - s.notifySession(sessionID, mcp.MethodNotificationToolsListChanged) + s.refreshSinks.Range(func(key, value any) bool { + sessionID, _ := key.(string) + sink, _ := value.(vmcpsession.ListChangedSink) + if terminated, err := s.vmcpSessionMgr.Validate(sessionID); terminated || err != nil || sink == nil { + s.refreshSinks.Delete(sessionID) + return true } - if withResources, ok := sess.(server.SessionWithResources); ok { - withResources.SetSessionResources(nil) - s.notifySession(sessionID, mcp.MethodNotificationResourcesListChanged) + for _, kind := range []vmcpsession.ChangeKind{ + vmcpsession.KindTools, vmcpsession.KindResources, vmcpsession.KindPrompts, + } { + sink(ctx, trigger, kind) } refreshed++ - } + return true + }) slog.Info("refreshed session capabilities after backend change", "trigger", trigger, "mode", mode.String(), "session_count", refreshed) } - -func (s *Server) notifySession(sessionID, method string) { - if err := s.mcpServer.SendNotificationToSpecificClient(sessionID, method, nil); err != nil { - slog.Debug("failed to send list_changed notification", - "session_id", sessionID, "method", method, "error", err) - } -} diff --git a/pkg/vmcp/server/serve_refresh_test.go b/pkg/vmcp/server/serve_refresh_test.go index ed2eddaefd..f506ac56da 100644 --- a/pkg/vmcp/server/serve_refresh_test.go +++ b/pkg/vmcp/server/serve_refresh_test.go @@ -8,8 +8,6 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" @@ -17,6 +15,7 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/health" "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) func TestScheduleOnExposureChange(t *testing.T) { @@ -72,29 +71,15 @@ func TestWatchRegistryVersionSchedulesOnChange(t *testing.T) { 2*time.Second, 10*time.Millisecond, "a registry change must schedule a refresh") } -// notifyingSession is a fake SDK session with tool and resource stores and a -// notification channel, so a test can observe list_changed notifications. -type notifyingSession struct { - fakeSDKResourceSession - notifications chan mcp.JSONRPCNotification -} - -func (n *notifyingSession) NotificationChannel() chan<- mcp.JSONRPCNotification { - return n.notifications -} - -// TestRefreshSessionCapabilitiesClearsAndNotifies verifies that a refresh clears -// the tool and resource store of a session registered on this pod, sends -// list_changed for both, and that the next request fills the stores again. -func TestRefreshSessionCapabilitiesClearsAndNotifies(t *testing.T) { +// TestRefreshSessionCapabilitiesTriggersSinks verifies that a refresh triggers the +// list_changed sink of each live session registered on this pod for tools, +// resources and prompts, and drops the sink of a session that ended. +func TestRefreshSessionCapabilitiesTriggersSinks(t *testing.T) { t.Parallel() - const uri = "file:///doc.txt" ctrl := gomock.NewController(t) - tools := []vmcp.Tool{{Name: "t"}} - factory, _ := newToolSessionFactory(t, ctrl, tools) - fc := &fakeCore{tools: tools, resources: []vmcp.Resource{{URI: uri, Name: "doc"}}} - srv, err := Serve(context.Background(), fc, &ServerConfig{ + factory, _ := newToolSessionFactory(t, ctrl, nil) + srv, err := Serve(context.Background(), &fakeCore{}, &ServerConfig{ SessionTTL: time.Minute, SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), @@ -102,32 +87,22 @@ func TestRefreshSessionCapabilitiesClearsAndNotifies(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = srv.Stop(context.Background()) }) - sess := ¬ifyingSession{ - fakeSDKResourceSession: fakeSDKResourceSession{ - fakeSDKSession: fakeSDKSession{id: srv.vmcpSessionMgr.Generate(), tools: map[string]server.ServerTool{}}, - resources: map[string]server.ServerResource{}, - }, - notifications: make(chan mcp.JSONRPCNotification, 10), - } - require.NoError(t, srv.mcpServer.RegisterSession(context.Background(), sess)) - require.Contains(t, sess.tools, "t", "registration must inject the tools") - require.Contains(t, sess.resources, uri, "registration must inject the resources") + live := srv.vmcpSessionMgr.Generate() + var kinds []vmcpsession.ChangeKind + srv.refreshSinks.Store(live, vmcpsession.ListChangedSink( + func(_ context.Context, _ string, kind vmcpsession.ChangeKind) { kinds = append(kinds, kind) })) + srv.refreshSinks.Store("ended-session", vmcpsession.ListChangedSink( + func(context.Context, string, vmcpsession.ChangeKind) { + t.Error("the sink of an ended session must not run") + })) srv.refreshSessionCapabilities("b", refreshSessionsContainingBackend) - assert.Empty(t, sess.tools, "refresh must clear the tool store") - assert.Empty(t, sess.resources, "refresh must clear the resource store") - var methods []string - for len(sess.notifications) > 0 { - methods = append(methods, (<-sess.notifications).Method) - } - assert.ElementsMatch(t, []string{ - mcp.MethodNotificationToolsListChanged, mcp.MethodNotificationResourcesListChanged, - }, methods) - - hookCtx := srv.mcpServer.WithContext(context.Background(), sess) - srv.lazyInjectSessionTools(hookCtx, "") - srv.lazyInjectSessionResources(hookCtx) - assert.Contains(t, sess.tools, "t", "the next request must fill the tool store again") - assert.Contains(t, sess.resources, uri, "the next request must fill the resource store again") + assert.ElementsMatch(t, []vmcpsession.ChangeKind{ + vmcpsession.KindTools, vmcpsession.KindResources, vmcpsession.KindPrompts, + }, kinds) + _, kept := srv.refreshSinks.Load(live) + assert.True(t, kept, "a live session keeps its sink") + _, stillThere := srv.refreshSinks.Load("ended-session") + assert.False(t, stillThere, "the sink of an ended session is dropped") } diff --git a/pkg/vmcp/server/serve_registration_wait_test.go b/pkg/vmcp/server/serve_registration_wait_test.go deleted file mode 100644 index bdb447a6f6..0000000000 --- a/pkg/vmcp/server/serve_registration_wait_test.go +++ /dev/null @@ -1,223 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "sync/atomic" - "testing" - "time" - - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "go.uber.org/mock/gomock" - - "github.com/stacklok/toolhive/pkg/auth" - transportsession "github.com/stacklok/toolhive/pkg/transport/session" - "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" - vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" - sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" -) - -// gatedSessionFactory holds MakeSessionWithID until release is closed, so a test -// can send a call while the session registration is still running. -type gatedSessionFactory struct { - vmcpsession.MultiSessionFactory - release chan struct{} -} - -func (f *gatedSessionFactory) MakeSessionWithID( - ctx context.Context, id string, identity *auth.Identity, backends []*vmcp.Backend, -) (vmcpsession.MultiSession, error) { - <-f.release - return f.MultiSessionFactory.MakeSessionWithID(ctx, id, identity, backends) -} - -// startGatedRegistration builds a Serve server, generates a session placeholder and -// starts its registration, which blocks in the session factory until the returned -// release func is called. -func startGatedRegistration(t *testing.T) (*Server, string, func()) { - t.Helper() - ctrl := gomock.NewController(t) - tools := []vmcp.Tool{{Name: "t"}} - base, _ := newToolSessionFactory(t, ctrl, tools) - factory := &gatedSessionFactory{MultiSessionFactory: base, release: make(chan struct{})} - - srv, err := Serve(context.Background(), &fakeCore{tools: tools}, &ServerConfig{ - SessionTTL: time.Minute, - SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, - BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), - }) - require.NoError(t, err) - - sessionID := srv.vmcpSessionMgr.Generate() - require.NotEmpty(t, sessionID) - registered := make(chan struct{}) - go func() { - defer close(registered) - srv.handleSessionRegistration(context.Background(), - &fakeSDKSession{id: sessionID, tools: map[string]server.ServerTool{}}) - }() - - released := false - release := func() { - if !released { - released = true - close(factory.release) - } - } - t.Cleanup(func() { - release() - <-registered - _ = srv.Stop(context.Background()) - }) - return srv, sessionID, release -} - -// TestEnforceSessionBindingWaitsForRegistration covers a call that arrives after -// initialize returned the session ID but before the registration wrote the -// identity binding. -func TestEnforceSessionBindingWaitsForRegistration(t *testing.T) { - t.Parallel() - - t.Run("call waits and succeeds once the registration finishes", func(t *testing.T) { - t.Parallel() - srv, sessionID, release := startGatedRegistration(t) - - result := make(chan error, 1) - go func() { result <- srv.enforceSessionBinding(context.Background(), sessionID, nil) }() - - select { - case err := <-result: - t.Fatalf("enforceSessionBinding returned before the registration finished: %v", err) - case <-time.After(150 * time.Millisecond): - } - - release() - select { - case err := <-result: - require.NoError(t, err) - case <-time.After(5 * time.Second): - t.Fatal("enforceSessionBinding did not return after the registration finished") - } - }) - - t.Run("call that gives up waiting does not terminate the session", func(t *testing.T) { - t.Parallel() - srv, sessionID, release := startGatedRegistration(t) - - ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer cancel() - req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "t", Arguments: map[string]any{}}} - res, err := srv.coreToolHandler(sessionID, "t", "")(ctx, req) - require.NoError(t, err) - require.True(t, res.IsError) - - isTerminated, err := srv.vmcpSessionMgr.Validate(sessionID) - require.NoError(t, err, "the placeholder must still exist") - assert.False(t, isTerminated, "the session must not be terminated on 'not registered yet'") - - release() - require.Eventually(t, func() bool { - return srv.enforceSessionBinding(context.Background(), sessionID, nil) == nil - }, 5*time.Second, 25*time.Millisecond, "the session must be usable after the registration") - }) - - t.Run("unknown session is rejected without waiting", func(t *testing.T) { - t.Parallel() - srv, _, _ := startGatedRegistration(t) - - start := time.Now() - err := srv.enforceSessionBinding(context.Background(), "no-such-session", nil) - require.ErrorIs(t, err, sessiontypes.ErrUnauthorizedCaller) - assert.Less(t, time.Since(start), time.Second) - }) -} - -// countingStorage counts Load calls on a DataStorage. -type countingStorage struct { - transportsession.DataStorage - loads atomic.Int32 -} - -func (c *countingStorage) Load(ctx context.Context, id string) (map[string]string, error) { - c.loads.Add(1) - return c.DataStorage.Load(ctx, id) -} - -// TestRegistrationWaitBacksOff verifies that a call waiting for a registration -// reads the session store once per step, with a growing step, instead of a full -// session lookup plus a placeholder check every 25 ms. -func TestRegistrationWaitBacksOff(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - tools := []vmcp.Tool{{Name: "t"}} - base, _ := newToolSessionFactory(t, ctrl, tools) - srv, err := Serve(context.Background(), &fakeCore{tools: tools}, &ServerConfig{ - SessionTTL: time.Minute, - SessionManagerConfig: &sessionmanager.FactoryConfig{Base: base}, - BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), - }) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Stop(context.Background()) }) - - local, err := transportsession.NewLocalSessionDataStorage(time.Minute) - require.NoError(t, err) - storage := &countingStorage{DataStorage: local} - mgr, cleanup, err := sessionmanager.New(storage, &sessionmanager.FactoryConfig{Base: base}, - vmcp.NewImmutableRegistry([]vmcp.Backend{})) - require.NoError(t, err) - t.Cleanup(func() { _ = cleanup(context.Background()) }) - srv.vmcpSessionMgr = mgr - _ = srv.sessionDataStorage.Close() - srv.sessionDataStorage = storage - - // A placeholder whose registration never runs stays pending. - sessionID := mgr.Generate() - storage.loads.Store(0) - - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - err = srv.enforceSessionBinding(ctx, sessionID, nil) - require.ErrorIs(t, err, errSessionNotRegistered) - assert.Less(t, storage.loads.Load(), int32(15), "a one-second wait must read the store only a few times") -} - -// TestUnregisterDuringRegistration verifies that a session mcp-go unregisters -// while its registration is still running is not left in the set of active -// sessions (which feeds the gauges and the capability refresh). -func TestUnregisterDuringRegistration(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - tools := []vmcp.Tool{{Name: "t"}} - base, _ := newToolSessionFactory(t, ctrl, tools) - factory := &gatedSessionFactory{MultiSessionFactory: base, release: make(chan struct{})} - srv, err := Serve(context.Background(), &fakeCore{tools: tools}, &ServerConfig{ - SessionTTL: time.Minute, - SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, - BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), - }) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Stop(context.Background()) }) - - sessionID := srv.vmcpSessionMgr.Generate() - sess := &fakeSDKSession{id: sessionID, tools: map[string]server.ServerTool{}} - registered := make(chan error, 1) - go func() { registered <- srv.mcpServer.RegisterSession(context.Background(), sess) }() - - // Wait until the registration is blocked in the session factory, then - // unregister the session, as a DELETE or the end of a GET would. - require.Eventually(t, func() bool { return srv.inFlightRegistrations.Load() == 1 }, - 2*time.Second, 5*time.Millisecond) - srv.mcpServer.UnregisterSession(context.Background(), sessionID) - close(factory.release) - require.NoError(t, <-registered) - - assert.Zero(t, srv.activeClientSessionCount(), "an unregistered session must not stay active") -} diff --git a/pkg/vmcp/server/serve_resource_inject_test.go b/pkg/vmcp/server/serve_resource_inject_test.go deleted file mode 100644 index 51a59fe72e..0000000000 --- a/pkg/vmcp/server/serve_resource_inject_test.go +++ /dev/null @@ -1,75 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "testing" - "time" - - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "go.uber.org/mock/gomock" - - "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" -) - -// fakeSDKResourceSession adds a per-session resource store to fakeSDKSession. -type fakeSDKResourceSession struct { - fakeSDKSession - resources map[string]server.ServerResource -} - -var _ server.SessionWithResources = (*fakeSDKResourceSession)(nil) - -func (f *fakeSDKResourceSession) GetSessionResources() map[string]server.ServerResource { - return f.resources -} - -func (f *fakeSDKResourceSession) SetSessionResources(r map[string]server.ServerResource) { - f.resources = r -} - -// TestServeLazyInjectsResourcesForRehydratedSession covers a session registered on -// another pod: its SDK session on this pod has an empty resource store, which the -// before-list/before-read resource hooks fill from the core. -func TestServeLazyInjectsResourcesForRehydratedSession(t *testing.T) { - t.Parallel() - - const uri = "file:///doc.txt" - ctrl := gomock.NewController(t) - factory, _ := newToolSessionFactory(t, ctrl, nil) - fc := &fakeCore{resources: []vmcp.Resource{{URI: uri, Name: "doc"}}} - srv, err := Serve(context.Background(), fc, &ServerConfig{ - SessionTTL: time.Minute, - SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, - BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), - }) - require.NoError(t, err) - t.Cleanup(func() { _ = srv.Stop(context.Background()) }) - - preexisting := map[string]server.ServerResource{"file:///other": {Resource: mcp.Resource{URI: "file:///other"}}} - tests := []struct { - name string - ctx context.Context - resources map[string]server.ServerResource - wantURI bool - }{ - {"empty store is filled from the core", context.Background(), map[string]server.ServerResource{}, true}, - {"populated store is left unchanged", context.Background(), preexisting, false}, - {"health probe gets no resources", context.WithValue(context.Background(), peerProbeKey{}, true), map[string]server.ServerResource{}, false}, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - sess := &fakeSDKResourceSession{fakeSDKSession: fakeSDKSession{id: "s"}, resources: tc.resources} - srv.lazyInjectSessionResources(srv.mcpServer.WithContext(tc.ctx, sess)) - _, got := sess.resources[uri] - assert.Equal(t, tc.wantURI, got) - }) - } -} diff --git a/pkg/vmcp/server/serve_session_termination_test.go b/pkg/vmcp/server/serve_session_termination_test.go new file mode 100644 index 0000000000..1819ab1361 --- /dev/null +++ b/pkg/vmcp/server/serve_session_termination_test.go @@ -0,0 +1,144 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + tcredis "github.com/stacklok/toolhive-core/redis" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + sessionfactorymocks "github.com/stacklok/toolhive/pkg/vmcp/session/mocks" +) + +// sharedRedisSessionKeyPrefix is the Redis key prefix shared by the origin +// server's session-data storage and the out-of-band terminating Manager in +// TestRegression_OriginPod_RejectsRequestAfterCrossPodTermination, so both see +// the same underlying keys in the shared miniredis instance. +const sharedRedisSessionKeyPrefix = "test:vmcp:session:" + +// TestRegression_OriginPod_RejectsRequestAfterCrossPodTermination pins the +// real U5 fix (toolhive-core v0.0.32, #5742): a session present in the origin +// pod's LOCAL bookkeeping but terminated in the SHARED (Redis) session store +// by a second pod must be rejected on the very next request to the origin +// pod, not served as if still live. +// +// This is distinct from the existing DELETE-path termination tests +// (TestRegression_TerminatedSessionRejected et al., session_lifecycle_regression_test.go): +// those terminate through the SAME server instance, so the SDK's local +// session bookkeeping (forgetSession) is dropped in lockstep with the storage +// delete. Here, termination happens on a SECOND Manager that only touches the +// shared store -- the origin pod's local go-sdk session and node-local cache +// entry are never told directly. The origin pod must instead detect the +// termination by re-validating the shared store on the next request (see +// mcpcompat/server/transports.go's "Local-session validation, issue #156, +// item U5" comment), and drop its local bookkeeping only then. +func TestRegression_OriginPod_RejectsRequestAfterCrossPodTermination(t *testing.T) { + t.Parallel() + + mr := miniredis.RunT(t) + + // Origin server ("pod A"): a Serve-built server backed by Redis-backed + // session-data storage, so its vMCP session manager's Validate() answers + // from the shared store rather than an in-process-only map. + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + testTool := vmcp.Tool{Name: "t"} + factory, _ := newToolSessionFactory(t, ctrl, []vmcp.Tool{testTool}) + fc := &fakeCore{tools: []vmcp.Tool{testTool}} + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionStorage: &vmcpconfig.SessionStorageConfig{ + Provider: "redis", + Address: mr.Addr(), + KeyPrefix: sharedRedisSessionKeyPrefix, + }, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-11-25", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return ok + }, 2*time.Second, 10*time.Millisecond, "session should be registered locally on the origin pod") + + // Sanity: the origin pod must currently regard the session as live, or the + // post-termination 404 assertion below would pass vacuously. + liveResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": map[string]any{}, + }, sessionID) + require.Equal(t, http.StatusOK, liveResp.StatusCode, "session must be live before cross-pod termination") + liveResp.Body.Close() + + // A second Manager ("pod B") over the SAME shared Redis store terminates + // the session out-of-band. This does NOT go through the origin pod's + // in-process forgetSession bookkeeping -- it only mutates the shared + // store, exactly like a DELETE served by a different replica. + terminatorFactory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + terminatorStorage, err := transportsession.NewRedisSessionDataStorage( + context.Background(), + tcredis.Config{Addr: mr.Addr()}, + sharedRedisSessionKeyPrefix, + time.Hour, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = terminatorStorage.Close() }) + + smB, cleanupB, err := sessionmanager.New( + terminatorStorage, + &sessionmanager.FactoryConfig{Base: terminatorFactory}, + vmcp.NewImmutableRegistry([]vmcp.Backend{}), + ) + require.NoError(t, err) + t.Cleanup(func() { _ = cleanupB(context.Background()) }) + + isNotAllowed, err := smB.Terminate(sessionID) + require.NoError(t, err) + require.False(t, isNotAllowed) + + // The origin pod must reject the NEXT request for the cross-pod-terminated + // session with HTTP 404 -- the streamable transport's validate-on-every- + // request path must observe the termination via the shared store and + // refuse to serve it from local state. + rejectResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", "id": 3, "method": "tools/list", "params": map[string]any{}, + }, sessionID) + defer rejectResp.Body.Close() + assert.Equal(t, http.StatusNotFound, rejectResp.StatusCode, + "a session terminated on another pod's shared store must be rejected as 404 "+ + "on the very next request to the origin pod") +} diff --git a/pkg/vmcp/server/serve_session_test.go b/pkg/vmcp/server/serve_session_test.go index ecdd783dab..8451d3b359 100644 --- a/pkg/vmcp/server/serve_session_test.go +++ b/pkg/vmcp/server/serve_session_test.go @@ -12,17 +12,20 @@ import ( "io" "net/http" "net/http/httptest" + "sort" + "strings" + "sync" "sync/atomic" "testing" "time" "github.com/alicebob/miniredis/v2" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/ratelimit" @@ -49,6 +52,20 @@ import ( // toolSessionState tracks observable behaviour of the mock session factory below. type toolSessionState struct { makeWithIDCalled atomic.Bool + + // sinkMu guards capturedSink: the sink passed to MakeSessionWithID (#5748), + // captured so tests can invoke it directly to simulate an asynchronous + // backend notification firing after registration. + sinkMu sync.Mutex + capturedSink vmcpsession.ListChangedSink +} + +// sink returns the ListChangedSink MakeSessionWithID most recently received, or +// nil if none was supplied. +func (s *toolSessionState) sink() vmcpsession.ListChangedSink { + s.sinkMu.Lock() + defer s.sinkMu.Unlock() + return s.capturedSink } // newToolSessionFactory creates a MockMultiSessionFactory whose MakeSessionWithID @@ -60,9 +77,16 @@ func newToolSessionFactory( t.Helper() state := &toolSessionState{} factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func( + _ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, sink vmcpsession.ListChangedSink, + ) (vmcpsession.MultiSession, error) { state.makeWithIDCalled.Store(true) + if sink != nil { + state.sinkMu.Lock() + state.capturedSink = sink + state.sinkMu.Unlock() + } mock := sessionmocks.NewMockMultiSession(ctrl) mock.EXPECT().ID().Return(id).AnyTimes() mock.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes() @@ -108,55 +132,152 @@ func newToolSessionFactory( // session factory: the factory establishes the session record, the core supplies the // advertised set and call routing. type fakeCore struct { - tools []vmcp.Tool - resources []vmcp.Resource - prompts []vmcp.Prompt - - listToolsCalls atomic.Int32 - listResourcesCalls atomic.Int32 - callToolCalls atomic.Int32 - readResourceCalls atomic.Int32 - lastCallToolName atomic.Value // string - lastReadURI atomic.Value // string + tools []vmcp.Tool + resources []vmcp.Resource + resourceTemplates []vmcp.ResourceTemplate + prompts []vmcp.Prompt + + listToolsCalls atomic.Int32 + listResourcesCalls atomic.Int32 + listResourceTemplatesCalls atomic.Int32 + listPromptsCalls atomic.Int32 + callToolCalls atomic.Int32 + readResourceCalls atomic.Int32 + getPromptCalls atomic.Int32 + lastCallToolName atomic.Value // string + lastCallToolArgs atomic.Value // map[string]any + lastReadURI atomic.Value // string + lastGetPromptName atomic.Value // string + lastGetPromptArgs atomic.Value // map[string]any // healthMarkedCalls counts CallTool/ReadResource calls whose context carries // the health-check marker (which switches outgoing auth to service credentials). healthMarkedCalls atomic.Int32 lastCallSubject atomic.Value // string: the identity CallTool ran with - callErr error // when set, CallTool returns it (e.g. vmcp.ErrAuthorizationFailed) - readErr error // when set, ReadResource returns it + completeCalls atomic.Int32 + lastCompleteRef atomic.Value // vmcp.CompletionRef + completeValues []string // returned by Complete when completeErr is nil + + callErr error // when set, CallTool returns it (e.g. vmcp.ErrAuthorizationFailed) + readErr error // when set, ReadResource returns it + readMeta map[string]any // when set, ReadResource sets it as result.Meta on success + promptErr error // when set, GetPrompt returns it (e.g. vmcp.ErrAuthorizationFailed) + completeErr error // when set, Complete returns it (e.g. vmcp.ErrAuthorizationFailed) + lookupResourceErr error // when set, LookupResource returns it for an ADVERTISED URI (admission denial) + + // invalidateCacheCalls counts InvalidateCapabilityCache invocations, so tests + // covering the list_changed sink can assert the cache was re-swept (#5748). + invalidateCacheCalls atomic.Int32 + + // lastListToolsCtx records the context ListTools last received, so #5748 + // tests can assert the async resync reconstructed a context carrying the + // registration-time identity + forwarded headers (B-HIGH-1). Wrapped in a + // fixed struct type because atomic.Value panics when Store sees differing + // concrete context types across calls (HTTP request ctx vs Background). + lastListToolsCtx atomic.Pointer[ctxBox] + // lastListResourcesCtx and lastListPromptsCtx are the #5969 analogues, so + // the resources/prompts resyncs are guarded against silently dropping + // identity/header scoping too (auth-scoping regression guard). + lastListResourcesCtx atomic.Pointer[ctxBox] + lastListPromptsCtx atomic.Pointer[ctxBox] + + // callIdentitiesMu guards callIdentities. Additive, nil-safe recording of + // the identity CallTool observed per invocation (issue #156 item U3 / + // #5742 regression guard): nil until a test opts in via + // enableCallIdentityRecording, so existing fakeCore users that never call + // it are unaffected. Keyed by the request-supplied "caller" argument (which + // travels in the request BODY, independent of the context-bridged identity), + // with the observed identity as the value. A test asserts value.Subject == + // key: since key and value come from independent sources, a concurrency + // regression that SWAPPED identities between two racing requests (request + // caller=p1 handled with p2's identity) records a mismatched pair and fails, + // not just a collapse where one key goes missing. + callIdentitiesMu sync.Mutex + callIdentities map[string]*auth.Identity } +// ctxBox wraps a context.Context for atomic.Pointer storage on fakeCore. +type ctxBox struct{ ctx context.Context } + var _ core.VMCP = (*fakeCore)(nil) -func (f *fakeCore) ListTools(context.Context, *auth.Identity) ([]vmcp.Tool, error) { +func (f *fakeCore) ListTools(ctx context.Context, _ *auth.Identity) ([]vmcp.Tool, error) { f.listToolsCalls.Add(1) + f.lastListToolsCtx.Store(&ctxBox{ctx: ctx}) return f.tools, nil } func (f *fakeCore) CallTool( - ctx context.Context, caller *auth.Identity, name string, _ map[string]any, _ map[string]any, + ctx context.Context, identity *auth.Identity, name string, args map[string]any, _ map[string]any, ) (*vmcp.ToolCallResult, error) { f.callToolCalls.Add(1) - if caller != nil { - f.lastCallSubject.Store(caller.Subject) + if identity != nil { + f.lastCallSubject.Store(identity.Subject) } if health.IsHealthCheck(ctx) { f.healthMarkedCalls.Add(1) } f.lastCallToolName.Store(name) + if args != nil { + f.lastCallToolArgs.Store(args) + } + f.recordCallIdentity(args, identity) if f.callErr != nil { return nil, f.callErr } return &vmcp.ToolCallResult{Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "ok"}}}, nil } -func (f *fakeCore) ListResources(context.Context, *auth.Identity) ([]vmcp.Resource, error) { +// enableCallIdentityRecording turns on per-call identity recording (see +// callIdentities). Call before Serve; a fakeCore that never calls this keeps +// the zero-value nil map, so recordCallIdentity/CallTool remain a no-op and +// existing callers are unaffected. +func (f *fakeCore) enableCallIdentityRecording() { + f.callIdentitiesMu.Lock() + defer f.callIdentitiesMu.Unlock() + f.callIdentities = make(map[string]*auth.Identity) +} + +// recordCallIdentity stores the observed identity keyed by the request-supplied +// "caller" argument (from args, which travels in the request body independent of +// the context-bridged identity) when recording is enabled; nil-safe no-op +// otherwise (callIdentities is nil for every fakeCore that never calls +// enableCallIdentityRecording). Keying on the request-body caller rather than +// the identity's own Subject is what makes callIdentityFor(caller).Subject == +// caller a non-tautological, swap-detecting assertion. +func (f *fakeCore) recordCallIdentity(args map[string]any, identity *auth.Identity) { + f.callIdentitiesMu.Lock() + defer f.callIdentitiesMu.Unlock() + if f.callIdentities == nil { + return + } + key := "" + if caller, ok := args["caller"].(string); ok { + key = caller + } + f.callIdentities[key] = identity +} + +// callIdentityFor returns the identity CallTool observed for the request whose +// "caller" argument equals key, or nil if no such call was recorded. +func (f *fakeCore) callIdentityFor(key string) *auth.Identity { + f.callIdentitiesMu.Lock() + defer f.callIdentitiesMu.Unlock() + return f.callIdentities[key] +} + +func (f *fakeCore) ListResources(ctx context.Context, _ *auth.Identity) ([]vmcp.Resource, error) { f.listResourcesCalls.Add(1) + f.lastListResourcesCtx.Store(&ctxBox{ctx: ctx}) return f.resources, nil } +func (f *fakeCore) ListResourceTemplates(context.Context, *auth.Identity) ([]vmcp.ResourceTemplate, error) { + f.listResourceTemplatesCalls.Add(1) + return f.resourceTemplates, nil +} + func (f *fakeCore) ReadResource(ctx context.Context, _ *auth.Identity, uri string) (*vmcp.ResourceReadResult, error) { f.readResourceCalls.Add(1) if health.IsHealthCheck(ctx) { @@ -166,17 +287,46 @@ func (f *fakeCore) ReadResource(ctx context.Context, _ *auth.Identity, uri strin if f.readErr != nil { return nil, f.readErr } - return &vmcp.ResourceReadResult{Contents: []vmcp.ResourceContent{{URI: uri, Text: "resource-body"}}}, nil + return &vmcp.ResourceReadResult{ + Contents: []vmcp.ResourceContent{{URI: uri, Text: "resource-body"}}, + Meta: f.readMeta, + }, nil } -func (f *fakeCore) ListPrompts(context.Context, *auth.Identity) ([]vmcp.Prompt, error) { +func (f *fakeCore) ListPrompts(ctx context.Context, _ *auth.Identity) ([]vmcp.Prompt, error) { + f.listPromptsCalls.Add(1) + f.lastListPromptsCtx.Store(&ctxBox{ctx: ctx}) return f.prompts, nil } -func (*fakeCore) GetPrompt( - context.Context, *auth.Identity, string, map[string]any, +func (f *fakeCore) GetPrompt( + _ context.Context, _ *auth.Identity, name string, args map[string]any, ) (*vmcp.PromptGetResult, error) { - return &vmcp.PromptGetResult{}, nil + f.getPromptCalls.Add(1) + f.lastGetPromptName.Store(name) + if args != nil { + f.lastGetPromptArgs.Store(args) + } + if f.promptErr != nil { + return nil, f.promptErr + } + return &vmcp.PromptGetResult{ + Description: "prompt-desc", + Messages: []vmcp.PromptMessage{ + {Role: "user", Content: vmcp.Content{Type: vmcp.ContentTypeText, Text: "prompt-body"}}, + }, + }, nil +} + +func (f *fakeCore) Complete( + _ context.Context, _ *auth.Identity, ref vmcp.CompletionRef, _, _ string, _ map[string]string, +) (*vmcp.CompletionResult, error) { + f.completeCalls.Add(1) + f.lastCompleteRef.Store(ref) + if f.completeErr != nil { + return nil, f.completeErr + } + return &vmcp.CompletionResult{Values: f.completeValues, Total: len(f.completeValues)}, nil } func (f *fakeCore) LookupTool(_ context.Context, _ *auth.Identity, name string) (*vmcp.Tool, error) { @@ -192,6 +342,12 @@ func (f *fakeCore) LookupTool(_ context.Context, _ *auth.Identity, name string) func (f *fakeCore) LookupResource(_ context.Context, _ *auth.Identity, uri string) (*vmcp.Resource, error) { for i := range f.resources { if f.resources[i].URI == uri { + // An advertised resource whose admission is denied: exercises the + // subscribe/completion admission-denial path (the URI is known, but the + // caller may not read it). + if f.lookupResourceErr != nil { + return nil, f.lookupResourceErr + } res := f.resources[i] return &res, nil } @@ -203,6 +359,18 @@ func (*fakeCore) LookupPrompt(context.Context, *auth.Identity, string) (*vmcp.Pr return nil, vmcp.ErrNotFound } +// Check* mirror the Call*/Read* error injection so a pre-flight gate over this fake +// sees the same allow/deny the call path would. +func (f *fakeCore) CheckToolCall(context.Context, *auth.Identity, string, map[string]any) error { + return f.callErr +} + +func (f *fakeCore) CheckResourceRead(context.Context, *auth.Identity, string) error { + return f.readErr +} + +func (*fakeCore) CheckPromptGet(context.Context, *auth.Identity, string) error { return nil } + func (*fakeCore) ListBackends(context.Context, *auth.Identity, bool) ([]vmcp.Backend, error) { return nil, nil } @@ -211,10 +379,16 @@ func (*fakeCore) LookupBackend(context.Context, *auth.Identity, string) (*vmcp.B return nil, vmcp.ErrNotFound } +func (*fakeCore) Discover(context.Context, *auth.Identity) (core.DiscoverCapabilities, error) { + return core.DiscoverCapabilities{}, nil +} + func (*fakeCore) Close() error { return nil } func (*fakeCore) BackendHealth() health.Reporter { return nil } +func (f *fakeCore) InvalidateCapabilityCache() { f.invalidateCacheCalls.Add(1) } + // TestServeRegistersSessionHooks verifies that Serve registers the OnRegisterSession // hook and wires two-phase session creation: an MCP initialize triggers the hook, // which creates the session record (MakeSessionWithID) and injects the core's advertised @@ -296,6 +470,126 @@ func TestServeRegistersSessionHooks(t *testing.T) { "tools/list must reuse the fixed-at-initialize set, not re-aggregate via the core") } +// TestServeToolsListPagination_CompleteSet pins #5742 §5: when vMCP's aggregated +// tool set exceeds a single page, a downstream client that follows the MCP +// pagination cursor receives the COMPLETE set — every advertised tool reachable, +// no duplicates, no omissions. That completeness is the real contract; how the +// server chunks it is not. +// +// Background: go-sdk paginates server-side list responses at DefaultPageSize +// (1000), emitting nextCursor; mcp-go (pre-migration) returned the whole set in +// one page, so a naive reader issuing a single tools/list and ignoring +// nextCursor silently drops the tail — the regression the migration risked. +// §5 sanctions either branch ("either the server page size is raised OR the +// test exercises cursor-following"), and mcpcompat's own WithPageSize doc tells +// aggregators with >1000 tools to raise it. So this test deliberately does NOT +// assert that pagination happened — asserting nextCursor would go red on +// exactly that recommended fix even though downstream clients are strictly +// better off and completeness still holds. It drives the cursor loop and +// asserts only the recovered set, which is invariant under both branches. +// +// The set is generated deterministically (tool-0000 … tool-1499) so the +// assertion checks exact membership, not just the count. totalTools is a local +// constant rather than derived from go-sdk's DefaultPageSize: the mcpcompat +// import doesn't re-export that constant and go-sdk is an indirect dependency, +// so pinning 1500 > 1000 by hand is cheaper than promoting it for one test. +func TestServeToolsListPagination_CompleteSet(t *testing.T) { + t.Parallel() + + const totalTools = 1500 // > go-sdk DefaultPageSize (1000) today, so the server paginates + + ctrl := gomock.NewController(t) + tools := make([]vmcp.Tool, totalTools) + for i := range tools { + tools[i] = vmcp.Tool{Name: fmt.Sprintf("tool-%04d", i), Description: "pagination test tool"} + } + factory, state := newToolSessionFactory(t, ctrl, tools) + fc := &fakeCore{tools: tools} + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, state.makeWithIDCalled.Load, 2*time.Second, 10*time.Millisecond) + + // Drive the cursor loop exactly as a real client does: the first request + // carries NO cursor (an empty string is only accepted by go-sdk leniency — + // mcp/server.go short-circuits "" before decodeCursor — not by any protocol + // guarantee, and the spec shows the initial call cursorless), then each + // non-empty nextCursor is echoed back until the server stops issuing one. + const maxPages = 10 // generous bound; 1500 tools at page size 1000 needs 2 + var ( + gotNames []string + cursor string + ) + for pages := 1; ; pages++ { + require.LessOrEqual(t, pages, maxPages, "pagination must terminate (no infinite cursor)") + + params := map[string]any{} + if cursor != "" { + params["cursor"] = cursor + } + listResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": pages + 100, // stay clear of the initialize id + "method": "tools/list", + "params": params, + }, sessionID) + require.Equal(t, http.StatusOK, listResp.StatusCode, "tools/list page %d should succeed", pages) + + env, _ := readServeJSONRPC(t, listResp) + listResp.Body.Close() + require.Nil(t, env["error"], "tools/list must not return a JSON-RPC error: %v", env) + + gotNames = append(gotNames, toolNamesFromListResult(t, env)...) + result, ok := env["result"].(map[string]any) + require.True(t, ok) + next, _ := result["nextCursor"].(string) + if next == "" { + break + } + cursor = next + } + + // The complete set must be reachable across the pages, with no duplicates + // and no omissions. require.Equal on the sorted slices dumps a single diff + // on mismatch — far more readable in CI than 1500 per-index assert failures + // when a dropped-and-duplicated pair shifts every later index. + want := make([]string, totalTools) + for i := range want { + want[i] = fmt.Sprintf("tool-%04d", i) + } + sort.Strings(gotNames) + require.Equal(t, want, gotNames, + "cursor-following must recover exactly the advertised set — no duplicates, no omissions") +} + // fakeSDKSession is a minimal server.ClientSession + server.SessionWithTools used to // drive lazyInjectSessionTools directly (the SDK's session-context plumbing is otherwise // only reachable over HTTP, via MCPServer.WithContext). Its tool store is the observable @@ -713,6 +1007,251 @@ func registerServeSessionWithRegistry( return srv, sessionID, ts.URL } +// registerServeSessionCapturingSink is registerServeSession plus the +// toolSessionState so a test can retrieve the ListChangedSink (#5748) +// MakeSessionWithID received for the registered session and invoke it +// directly, simulating an asynchronous backend notification firing after +// registration completes. +func registerServeSessionCapturingSink( + t *testing.T, fc *fakeCore, +) (sessionID, baseURL string, state *toolSessionState) { + t.Helper() + ctrl := gomock.NewController(t) + factory, state := newToolSessionFactory(t, ctrl, fc.tools) + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID = initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, func() bool { _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID); return ok }, + 2*time.Second, 10*time.Millisecond, "session should be registered") + require.Eventually(t, func() bool { return state.sink() != nil }, + 2*time.Second, 10*time.Millisecond, "MakeSessionWithID should have received a non-nil sink") + return sessionID, ts.URL, state +} + +// TestListChangedSink_EndToEnd_ResyncsRegisteredSession drives a real session +// registration over HTTP, then invokes the captured sink directly (simulating +// a backend's asynchronous notifications/tools/list_changed) and verifies, +// via a real tools/list request, that the session's advertised set reflects +// the core's now-changed tool set — including a REMOVED tool disappearing — +// and that the cache was invalidated (#5748). +func TestListChangedSink_EndToEnd_ResyncsRegisteredSession(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "kept"}, {Name: "removed"}}} + sessionID, baseURL, state := registerServeSessionCapturingSink(t, fc) + + listResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": map[string]any{}, + }, sessionID) + env, _ := readServeJSONRPC(t, listResp) + listResp.Body.Close() + before := toolNamesFromListResult(t, env) + assert.ElementsMatch(t, []string{"kept", "removed"}, before) + + // The core's advertised set changes (as if the backend's own tool set + // changed): "removed" is gone, "added" is new. + fc.tools = []vmcp.Tool{{Name: "kept"}, {Name: "added"}} + + sink := state.sink() + require.NotNil(t, sink) + sink(context.Background(), "some-backend", "tools") + + require.Eventually(t, func() bool { + return fc.invalidateCacheCalls.Load() >= 1 + }, 2*time.Second, 10*time.Millisecond, "sink must invalidate the capability cache") + + require.Eventually(t, func() bool { + listResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", "id": 3, "method": "tools/list", "params": map[string]any{}, + }, sessionID) + defer listResp.Body.Close() + env, _ := readServeJSONRPC(t, listResp) + got := toolNamesFromListResult(t, env) + return assert.ObjectsAreEqualValues([]string{"added", "kept"}, sortedCopy(got)) + }, 2*time.Second, 10*time.Millisecond, "tools/list must reflect the resynced (replaced) tool set") +} + +// TestListChangedSink_EndToEnd_ResyncsRegisteredSession_Resources mirrors +// TestListChangedSink_EndToEnd_ResyncsRegisteredSession for resources (#5969): +// a resources/list_changed sink invocation invalidates the cache and resyncs +// the session's resource AND resource-template stores. Unlike tools, this +// pins the documented ADD-ONLY limitation (resyncSessionResources doc +// comment; toolhive-core mcpcompat per-session resource/template sync has no +// removal reconciliation): an addition ("added"/"added-tmpl") becomes visible, +// but a removed entry ("res-removed"/"tmpl-removed") STILL appears — it is +// not expected to disappear here, unlike the tools contrast test above. +func TestListChangedSink_EndToEnd_ResyncsRegisteredSession_Resources(t *testing.T) { + t.Parallel() + + fc := &fakeCore{ + resources: []vmcp.Resource{ + {Name: "res-kept", URI: "res-kept"}, + {Name: "res-removed", URI: "res-removed"}, + }, + resourceTemplates: []vmcp.ResourceTemplate{ + {Name: "tmpl-kept", URITemplate: "tmpl-kept"}, + {Name: "tmpl-removed", URITemplate: "tmpl-removed"}, + }, + } + sessionID, baseURL, state := registerServeSessionCapturingSink(t, fc) + + // Before snapshot: the added entries must be absent and the to-be-removed + // entries present, so the post-resync assertions cannot pass vacuously. + beforeResources := postServeMCPAndReadBody(t, baseURL, "resources/list", sessionID) + beforeTemplates := postServeMCPAndReadBody(t, baseURL, "resources/templates/list", sessionID) + assert.NotContains(t, beforeResources, "res-added") + assert.Contains(t, beforeResources, "res-removed") + assert.NotContains(t, beforeTemplates, "tmpl-added") + assert.Contains(t, beforeTemplates, "tmpl-removed") + + // The core's advertised set changes: "res-removed"/"tmpl-removed" are gone + // (as if the backend removed them), "res-added"/"tmpl-added" are new. + fc.resources = []vmcp.Resource{ + {Name: "res-kept", URI: "res-kept"}, + {Name: "res-added", URI: "res-added"}, + } + fc.resourceTemplates = []vmcp.ResourceTemplate{ + {Name: "tmpl-kept", URITemplate: "tmpl-kept"}, + {Name: "tmpl-added", URITemplate: "tmpl-added"}, + } + + sink := state.sink() + require.NotNil(t, sink) + sink(context.Background(), "some-backend", "resources") + + require.Eventually(t, func() bool { + return fc.invalidateCacheCalls.Load() >= 1 + }, 2*time.Second, 10*time.Millisecond, "sink must invalidate the capability cache") + + require.Eventually(t, func() bool { + resourcesBody := postServeMCPAndReadBody(t, baseURL, "resources/list", sessionID) + templatesBody := postServeMCPAndReadBody(t, baseURL, "resources/templates/list", sessionID) + return strings.Contains(resourcesBody, "res-kept") && + strings.Contains(resourcesBody, "res-added") && + strings.Contains(templatesBody, "tmpl-kept") && + strings.Contains(templatesBody, "tmpl-added") + }, 2*time.Second, 10*time.Millisecond, "resources/list and resources/templates/list must reflect the resynced additions") + + // ADD-ONLY GUARD (tracked toolhive-core follow-up, see resyncSessionResources + // doc comment): a removed resource/template is NOT expected to disappear — + // toolhive-core's mcpcompat per-session sync only adds, never removes, so + // the overlay being replaced here does not (yet) reach the live SDK server's + // removal path the way tools does. + resourcesBody := postServeMCPAndReadBody(t, baseURL, "resources/list", sessionID) + templatesBody := postServeMCPAndReadBody(t, baseURL, "resources/templates/list", sessionID) + assert.Contains(t, resourcesBody, "res-removed", + "add-only limitation: a removed resource still lingers until the toolhive-core follow-up lands") + assert.Contains(t, templatesBody, "tmpl-removed", + "add-only limitation: a removed resource template still lingers until the toolhive-core follow-up lands") +} + +// TestListChangedSink_EndToEnd_ResyncsRegisteredSession_Prompts mirrors +// TestListChangedSink_EndToEnd_ResyncsRegisteredSession for prompts (#5969), +// including the same ADD-ONLY guard as the resources test above. +func TestListChangedSink_EndToEnd_ResyncsRegisteredSession_Prompts(t *testing.T) { + t.Parallel() + + fc := &fakeCore{prompts: []vmcp.Prompt{{Name: "prompt-kept"}, {Name: "prompt-removed"}}} + sessionID, baseURL, state := registerServeSessionCapturingSink(t, fc) + + // Before snapshot: the added prompt must be absent and the to-be-removed + // prompt present, so the post-resync assertions cannot pass vacuously. + beforePrompts := postServeMCPAndReadBody(t, baseURL, "prompts/list", sessionID) + assert.NotContains(t, beforePrompts, "prompt-added") + assert.Contains(t, beforePrompts, "prompt-removed") + + fc.prompts = []vmcp.Prompt{{Name: "prompt-kept"}, {Name: "prompt-added"}} + + sink := state.sink() + require.NotNil(t, sink) + sink(context.Background(), "some-backend", "prompts") + + require.Eventually(t, func() bool { + return fc.invalidateCacheCalls.Load() >= 1 + }, 2*time.Second, 10*time.Millisecond, "sink must invalidate the capability cache") + + require.Eventually(t, func() bool { + body := postServeMCPAndReadBody(t, baseURL, "prompts/list", sessionID) + return strings.Contains(body, "prompt-kept") && strings.Contains(body, "prompt-added") + }, 2*time.Second, 10*time.Millisecond, "prompts/list must reflect the resynced (added) prompt") + + // ADD-ONLY GUARD (tracked toolhive-core follow-up, see resyncSessionPrompts + // doc comment): a removed prompt is NOT expected to disappear — + // toolhive-core's mcpcompat per-session prompt sync only adds, never + // removes. + body := postServeMCPAndReadBody(t, baseURL, "prompts/list", sessionID) + assert.Contains(t, body, "prompt-removed", + "add-only limitation: a removed prompt still lingers until the toolhive-core follow-up lands") +} + +// postServeMCPAndReadBody sends a JSON-RPC list request over the Streamable +// HTTP protocol and returns the raw response body as a string, for substring +// assertions against the (un-parsed) list result. +func postServeMCPAndReadBody(t *testing.T, baseURL, method, sessionID string) string { + t.Helper() + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", "id": 99, "method": method, "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + return string(body) +} + +// toolNamesFromListResult extracts the tool names from a tools/list JSON-RPC +// response envelope. +func toolNamesFromListResult(t *testing.T, env map[string]any) []string { + t.Helper() + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list response must have a result object; env: %v", env) + rawTools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + names := make([]string, 0, len(rawTools)) + for _, rt := range rawTools { + tm, ok := rt.(map[string]any) + require.True(t, ok) + name, _ := tm["name"].(string) + names = append(names, name) + } + return names +} + +// sortedCopy returns a sorted copy of ss for order-independent comparison. +func sortedCopy(ss []string) []string { + out := make([]string, len(ss)) + copy(out, ss) + sort.Strings(out) + return out +} + // TestServeCoreToolHandler exercises the Serve tool handler's error/denial branches by // invoking the handler closure directly against a registered (anonymous) session: // the happy path returns the core result; an ErrAuthorizationFailed is genericized so the @@ -856,6 +1395,508 @@ func TestServeCoreResourceHandler(t *testing.T) { assert.Contains(t, err.Error(), "read denied by authorization policy") assert.NotContains(t, err.Error(), "cedar said no", "the underlying authorizer error must not leak") }) + + t.Run("backend _meta is forwarded and reserved keys stripped", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + resources: []vmcp.Resource{{Name: "doc", URI: uri}}, + readMeta: map[string]any{ + "traceId": "trace-abc", + // Reserved protocol key: must be stripped on egress (#5986). + "io.modelcontextprotocol/protocolVersion": mcp.LATEST_PROTOCOL_VERSION, + }, + } + srv, sessionID, _ := registerServeSession(t, fc) + + res, err := srv.coreResourceHandler(sessionID, uri, "")(context.Background(), mcp.ReadResourceRequest{}) + require.NoError(t, err) + require.NotNil(t, res.Meta, "backend resource _meta must reach the wire result") + assert.Equal(t, "trace-abc", res.Meta.AdditionalFields["traceId"]) + _, reservedPresent := res.Meta.AdditionalFields["io.modelcontextprotocol/protocolVersion"] + assert.False(t, reservedPresent, "reserved io.modelcontextprotocol/* keys must be stripped") + }) +} + +// TestServeCoreResourceTemplateHandler covers the Serve resource-template path: +// the core's resource template is advertised by the registration builder, the +// handler routes a read of an EXPANDED URI (taken from the request, not a fixed +// URI) through core.ReadResource, and an ErrAuthorizationFailed is genericized. +func TestServeCoreResourceTemplateHandler(t *testing.T) { + t.Parallel() + + const uriTemplate = "file:///logs/{date}.txt" + const expandedURI = "file:///logs/2025-01-01.txt" + + t.Run("advertises the template and routes the expanded read through core.ReadResource", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resourceTemplates: []vmcp.ResourceTemplate{ + {Name: "Daily log", URITemplate: uriTemplate, MimeType: "text/plain"}, + }} + srv, sessionID, _ := registerServeSession(t, fc) + + templates, err := srv.coreSessionResourceTemplates(t.Context(), sessionID, nil) + require.NoError(t, err) + require.Len(t, templates, 1) + assert.Equal(t, uriTemplate, templates[0].Template.URITemplate) + assert.Equal(t, "text/plain", templates[0].Template.MIMEType) + + // The template handler reads the concrete URI from the request, not a fixed URI. + req := mcp.ReadResourceRequest{Params: mcp.ReadResourceParams{URI: expandedURI}} + contents, err := srv.coreResourceTemplateHandler(sessionID, "")(t.Context(), req) + require.NoError(t, err) + require.NotEmpty(t, contents) + + gotURI, _ := fc.lastReadURI.Load().(string) + assert.Equal(t, expandedURI, gotURI, + "the template handler must route the read through core.ReadResource with the expanded URI") + }) + + t.Run("admission filtering withholds a denied template", func(t *testing.T) { + t.Parallel() + // The core (fakeCore stands in for it) returns the admission-filtered set; + // an empty set means no template is advertised for this identity. + fc := &fakeCore{resourceTemplates: nil} + srv, sessionID, _ := registerServeSession(t, fc) + + templates, err := srv.coreSessionResourceTemplates(t.Context(), sessionID, nil) + require.NoError(t, err) + assert.Empty(t, templates, "a denied/withheld template must not be advertised") + }) + + t.Run("authorization denial yields a generic message", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + resourceTemplates: []vmcp.ResourceTemplate{{Name: "Daily log", URITemplate: uriTemplate}}, + readErr: fmt.Errorf("%w: cedar said no", vmcp.ErrAuthorizationFailed), + } + srv, sessionID, _ := registerServeSession(t, fc) + + req := mcp.ReadResourceRequest{Params: mcp.ReadResourceParams{URI: expandedURI}} + _, err := srv.coreResourceTemplateHandler(sessionID, "")(t.Context(), req) + require.Error(t, err) + assert.Contains(t, err.Error(), "read denied by authorization policy") + assert.NotContains(t, err.Error(), "cedar said no", "the underlying authorizer error must not leak") + }) +} + +// TestServeResourceTemplatesEndToEnd drives the resource-templates lane over the +// full Streamable HTTP protocol (the conformance gap this lane closes): +// +// - resources/templates/list advertises the backend's template, and +// - resources/read of an EXPANDED URI matching the template is served (routed +// through core.ReadResource) and returns contents, and +// - a denied read returns an authorization error rather than -32002 +// "Resource not found" — the pre-fix symptom was that vMCP never registered +// the template, so every templated read fell through to not-found. +func TestServeResourceTemplatesEndToEnd(t *testing.T) { + t.Parallel() + + const uriTemplate = "file:///logs/{date}.txt" + const expandedURI = "file:///logs/2025-01-01.txt" + + t.Run("lists the template and serves an expanded read", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resourceTemplates: []vmcp.ResourceTemplate{ + {Name: "Daily log", URITemplate: uriTemplate, MimeType: "text/plain"}, + }} + _, sessionID, baseURL := registerServeSession(t, fc) + + // resources/templates/list must advertise the backend's template. + listResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "resources/templates/list", + "params": map[string]any{}, + }, sessionID) + defer listResp.Body.Close() + require.Equal(t, http.StatusOK, listResp.StatusCode) + + env, raw := readServeJSONRPC(t, listResp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "list must have a result; body: %s", string(raw)) + templates, ok := result["resourceTemplates"].([]any) + require.True(t, ok, "result.resourceTemplates must be present; body: %s", string(raw)) + require.Len(t, templates, 1) + tmpl := templates[0].(map[string]any) + assert.Equal(t, uriTemplate, tmpl["uriTemplate"]) + + // resources/read of a URI expanded from the template must be served through + // core.ReadResource and return contents (not -32002). + readResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 3, + "method": "resources/read", + "params": map[string]any{"uri": expandedURI}, + }, sessionID) + defer readResp.Body.Close() + require.Equal(t, http.StatusOK, readResp.StatusCode) + + readEnv, readRaw := readServeJSONRPC(t, readResp) + require.NotContains(t, readEnv, "error", "a templated read must not error; body: %s", string(readRaw)) + readResult, ok := readEnv["result"].(map[string]any) + require.True(t, ok, "read must have a result; body: %s", string(readRaw)) + contents, ok := readResult["contents"].([]any) + require.True(t, ok, "result.contents must be present; body: %s", string(readRaw)) + require.NotEmpty(t, contents) + + gotURI, _ := fc.lastReadURI.Load().(string) + assert.Equal(t, expandedURI, gotURI, "the expanded URI must reach core.ReadResource") + }) + + t.Run("denied templated read returns an authorization error, not -32002", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + resourceTemplates: []vmcp.ResourceTemplate{{Name: "Daily log", URITemplate: uriTemplate}}, + readErr: fmt.Errorf("%w: cedar said no", vmcp.ErrAuthorizationFailed), + } + _, sessionID, baseURL := registerServeSession(t, fc) + + readResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "resources/read", + "params": map[string]any{"uri": expandedURI}, + }, sessionID) + defer readResp.Body.Close() + + env, raw := readServeJSONRPC(t, readResp) + readErr, ok := env["error"].(map[string]any) + require.True(t, ok, "a denied read must return a JSON-RPC error; body: %s", string(raw)) + + // The whole point: the template is registered, so the read reaches the + // handler and is DENIED — it is not the SDK's -32002 "Resource not found" + // that the pre-fix (never-registered) path returned. + if code, ok := readErr["code"].(float64); ok { + assert.NotEqual(t, float64(-32002), code, "denial must not surface as -32002 Resource not found") + } + msg, _ := readErr["message"].(string) + assert.Contains(t, msg, "read denied by authorization policy") + assert.NotContains(t, msg, "cedar said no", "the underlying authorizer detail must not leak") + }) +} + +// TestServeCorePromptHandler covers the Serve prompt path: the core's prompt is +// advertised by the registration builder, the handler routes a get through +// core.GetPrompt (widening the string args to map[string]any and converting the +// result), and an ErrAuthorizationFailed is genericized to the prompt denial message. +func TestServeCorePromptHandler(t *testing.T) { + t.Parallel() + + const promptName = "greeting" + + t.Run("advertises and routes the prompt through core.GetPrompt", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{prompts: []vmcp.Prompt{{ + Name: promptName, + Description: "a greeting prompt", + Arguments: []vmcp.PromptArgument{{Name: "name", Description: "who to greet", Required: true}}, + }}} + srv, sessionID, _ := registerServeSession(t, fc) + + prompts, err := srv.coreSessionPrompts(t.Context(), sessionID, nil) + require.NoError(t, err) + require.Len(t, prompts, 1) + assert.Equal(t, promptName, prompts[0].Prompt.Name) + require.Len(t, prompts[0].Prompt.Arguments, 1) + assert.Equal(t, "name", prompts[0].Prompt.Arguments[0].Name) + assert.True(t, prompts[0].Prompt.Arguments[0].Required) + + req := mcp.GetPromptRequest{Params: mcp.GetPromptParams{ + Name: promptName, + Arguments: map[string]string{"name": "world"}, + }} + result, err := srv.corePromptHandler(sessionID, promptName, "")(t.Context(), req) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "prompt-desc", result.Description) + require.Len(t, result.Messages, 1) + + gotName, _ := fc.lastGetPromptName.Load().(string) + assert.Equal(t, promptName, gotName, "the handler must route the get through core.GetPrompt with the name") + gotArgs, _ := fc.lastGetPromptArgs.Load().(map[string]any) + assert.Equal(t, "world", gotArgs["name"], "the handler must widen the string prompt args to map[string]any") + }) + + t.Run("authorization denial yields the prompt denial message", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + prompts: []vmcp.Prompt{{Name: promptName}}, + promptErr: fmt.Errorf("%w: cedar said no", vmcp.ErrAuthorizationFailed), + } + srv, sessionID, _ := registerServeSession(t, fc) + + _, err := srv.corePromptHandler(sessionID, promptName, "")(t.Context(), mcp.GetPromptRequest{}) + require.Error(t, err) + assert.Contains(t, err.Error(), vmcp.DenyMessagePromptGet) + assert.NotContains(t, err.Error(), "cedar said no", "the underlying authorizer error must not leak") + }) +} + +// TestServeCoreCompletionHandler covers coreCompletionHandler's admission-denial +// mapping directly (mirroring TestServeCorePromptHandler): completion reuses the +// underlying capability's deny wording, so a Cedar-denied prompt ref surfaces the +// prompt deny message and a denied resource ref surfaces the resource deny +// message — and neither leaks the raw authorizer detail. +func TestServeCoreCompletionHandler(t *testing.T) { + t.Parallel() + + const promptName = "greeting" + const uri = "file:///doc.txt" + + t.Run("prompt-ref authorization denial yields the prompt deny message", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + prompts: []vmcp.Prompt{{Name: promptName}}, + completeErr: fmt.Errorf("%w: cedar said no", vmcp.ErrAuthorizationFailed), + } + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := sdkSessionContext(t.Context(), srv, sessionID) + _, err := srv.coreCompletionHandler(ctx, mcp.CompleteRequest{Params: mcp.CompleteParams{ + Ref: mcp.PromptReference{Type: vmcp.CompletionRefTypePrompt, Name: promptName}, + Argument: mcp.CompleteArgument{Name: "name", Value: "wor"}, + }}) + require.Error(t, err) + assert.Contains(t, err.Error(), vmcp.DenyMessagePromptGet) + assert.NotContains(t, err.Error(), "cedar said no", "the underlying authorizer error must not leak") + }) + + t.Run("resource-ref authorization denial yields the resource deny message", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{ + resources: []vmcp.Resource{{Name: "doc", URI: uri}}, + completeErr: fmt.Errorf("%w: cedar said no", vmcp.ErrAuthorizationFailed), + } + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := sdkSessionContext(t.Context(), srv, sessionID) + _, err := srv.coreCompletionHandler(ctx, mcp.CompleteRequest{Params: mcp.CompleteParams{ + Ref: mcp.ResourceReference{Type: vmcp.CompletionRefTypeResource, URI: uri}, + Argument: mcp.CompleteArgument{Name: "path", Value: "d"}, + }}) + require.Error(t, err) + assert.Contains(t, err.Error(), vmcp.DenyMessageResourceRead) + assert.NotContains(t, err.Error(), "cedar said no", "the underlying authorizer error must not leak") + }) +} + +// TestServeCompletionEndToEnd drives the completion/complete lane over the full +// Streamable HTTP protocol (the conformance gap this lane closes): the SDK +// auto-advertises the completions capability (WithCompletionHandler is wired at +// construction), and a completion/complete request for a prompt ref routes through +// core.Complete and returns the backend's candidate values. +func TestServeCompletionEndToEnd(t *testing.T) { + t.Parallel() + + const promptName = "greeting" + + fc := &fakeCore{ + prompts: []vmcp.Prompt{{Name: promptName, Arguments: []vmcp.PromptArgument{{Name: "name"}}}}, + completeValues: []string{"world", "worf", "wormhole"}, + } + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "completion/complete", + "params": map[string]any{ + "ref": map[string]any{"type": "ref/prompt", "name": promptName}, + "argument": map[string]any{"name": "name", "value": "wor"}, + }, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + env, raw := readServeJSONRPC(t, resp) + require.NotContains(t, env, "error", "completion must not error; body: %s", string(raw)) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "completion must have a result; body: %s", string(raw)) + completion, ok := result["completion"].(map[string]any) + require.True(t, ok, "result.completion must be present; body: %s", string(raw)) + values, ok := completion["values"].([]any) + require.True(t, ok, "result.completion.values must be present; body: %s", string(raw)) + require.Len(t, values, 3) + assert.Equal(t, "world", values[0]) + + require.Equal(t, int32(1), fc.completeCalls.Load(), "the request must route through core.Complete exactly once") + gotRef, _ := fc.lastCompleteRef.Load().(vmcp.CompletionRef) + assert.Equal(t, vmcp.CompletionRefTypePrompt, gotRef.Type) + assert.Equal(t, promptName, gotRef.Name, "the prompt ref must reach core.Complete") +} + +// TestServeSubscriptionsEndToEnd drives resources/subscribe and resources/unsubscribe +// over the full Streamable HTTP protocol (the conformance gap these ack-level lanes +// close): initialize advertises resources.subscribe=true, subscribing to an advertised +// resource URI succeeds, and an unknown URI is rejected. vMCP accepts the subscription +// at ack level; it does NOT forward backend resources/updated notifications (out of +// scope). The pre-fix symptom was -32601 "server does not support resource subscriptions". +func TestServeSubscriptionsEndToEnd(t *testing.T) { + t.Parallel() + + const uri = "file:///doc.txt" + + t.Run("initialize advertises subscribe support", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "doc", URI: uri}}} + ctrl := gomock.NewController(t) + factory, _ := newToolSessionFactory(t, ctrl, fc.tools) + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + + env, raw := readServeJSONRPC(t, initResp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "initialize must have a result; body: %s", string(raw)) + caps, ok := result["capabilities"].(map[string]any) + require.True(t, ok, "result.capabilities must be present; body: %s", string(raw)) + resources, ok := caps["resources"].(map[string]any) + require.True(t, ok, "capabilities.resources must be advertised; body: %s", string(raw)) + assert.Equal(t, true, resources["subscribe"], "resources.subscribe must be advertised; body: %s", string(raw)) + }) + + t.Run("subscribe and unsubscribe an advertised resource succeed", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "doc", URI: uri}}} + _, sessionID, baseURL := registerServeSession(t, fc) + + subResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "resources/subscribe", + "params": map[string]any{"uri": uri}, + }, sessionID) + defer subResp.Body.Close() + require.Equal(t, http.StatusOK, subResp.StatusCode) + + subEnv, subRaw := readServeJSONRPC(t, subResp) + require.NotContains(t, subEnv, "error", "subscribe must succeed (regression: -32601); body: %s", string(subRaw)) + _, ok := subEnv["result"] + require.True(t, ok, "subscribe must return a result; body: %s", string(subRaw)) + + unsubResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 3, + "method": "resources/unsubscribe", + "params": map[string]any{"uri": uri}, + }, sessionID) + defer unsubResp.Body.Close() + require.Equal(t, http.StatusOK, unsubResp.StatusCode) + + unsubEnv, unsubRaw := readServeJSONRPC(t, unsubResp) + require.NotContains(t, unsubEnv, "error", "unsubscribe must succeed; body: %s", string(unsubRaw)) + }) + + t.Run("subscribe to an unknown resource is rejected", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "doc", URI: uri}}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "resources/subscribe", + "params": map[string]any{"uri": "file:///unknown.txt"}, + }, sessionID) + defer resp.Body.Close() + + env, raw := readServeJSONRPC(t, resp) + _, ok := env["error"].(map[string]any) + require.True(t, ok, "subscribing to an unadvertised resource must return an error; body: %s", string(raw)) + }) +} + +// TestServeCoreSubscribeHandler covers coreSubscribeHandler's branches directly: an +// advertised resource is accepted (nil), an unknown/admission-denied URI is rejected, +// and a session-binding failure is enforced before the resource lookup. +func TestServeCoreSubscribeHandler(t *testing.T) { + t.Parallel() + + const uri = "file:///doc.txt" + + t.Run("advertised resource is accepted", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "doc", URI: uri}}} + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := sdkSessionContext(t.Context(), srv, sessionID) + err := srv.coreSubscribeHandler(ctx, "resources/subscribe", uri) + require.NoError(t, err) + }) + + t.Run("unknown resource is rejected", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "doc", URI: uri}}} + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := sdkSessionContext(t.Context(), srv, sessionID) + err := srv.coreSubscribeHandler(ctx, "resources/subscribe", "file:///nope.txt") + require.Error(t, err) + assert.ErrorIs(t, err, vmcp.ErrNotFound) + }) + + t.Run("admission-denied advertised resource is rejected, not acked", func(t *testing.T) { + t.Parallel() + // The URI is advertised, but LookupResource applies the same admission + // decision resources/read enforces and denies it: subscribe must reject + // rather than silently ack. + fc := &fakeCore{ + resources: []vmcp.Resource{{Name: "doc", URI: uri}}, + lookupResourceErr: fmt.Errorf("%w: cedar said no", vmcp.ErrAuthorizationFailed), + } + srv, sessionID, _ := registerServeSession(t, fc) + + ctx := sdkSessionContext(t.Context(), srv, sessionID) + err := srv.coreSubscribeHandler(ctx, "resources/subscribe", uri) + require.Error(t, err) + assert.ErrorIs(t, err, vmcp.ErrAuthorizationFailed) + }) + + t.Run("session binding is enforced", func(t *testing.T) { + t.Parallel() + fc := &fakeCore{resources: []vmcp.Resource{{Name: "doc", URI: uri}}} + srv, _, _ := registerServeSession(t, fc) + + // An unknown session ID has no binding record, so enforceSessionBinding fails + // closed before the resource lookup. + ctx := sdkSessionContext(t.Context(), srv, "not-a-real-session") + err := srv.coreSubscribeHandler(ctx, "resources/subscribe", uri) + require.Error(t, err) + assert.Contains(t, err.Error(), "unauthorized") + }) +} + +// sdkSessionContext returns ctx carrying an SDK ClientSession for sessionID, the +// mechanism the global completion/subscribe handlers use to recover the session ID +// (server.ClientSessionFromContext). It mirrors the SDK's per-request context +// plumbing without going over HTTP. +func sdkSessionContext(ctx context.Context, srv *Server, sessionID string) context.Context { + return srv.mcpServer.WithContext(ctx, &fakeSDKSession{id: sessionID, tools: map[string]server.ServerTool{}}) } // TestServeHandlersLabelAuditBackend verifies the Serve-path audit labelling (#5512): @@ -903,6 +1944,22 @@ func TestServeHandlersLabelAuditBackend(t *testing.T) { "labelling must not re-aggregate — core.ListResources is not called again during the resource read") }) + t.Run("prompt handler labels the backend and does not re-aggregate", func(t *testing.T) { + t.Parallel() + const promptName = "greeting" + fc := &fakeCore{prompts: []vmcp.Prompt{{Name: promptName}}} + srv, sessionID, _ := registerServeSession(t, fc) + listPromptsAtRegistration := fc.listPromptsCalls.Load() + + bi := &audit.BackendInfo{} + ctx := audit.WithBackendInfo(context.Background(), bi) + _, err := srv.corePromptHandler(sessionID, promptName, "prompts-backend")(ctx, mcp.GetPromptRequest{}) + require.NoError(t, err) + assert.Equal(t, "prompts-backend", bi.BackendName) + assert.Equal(t, listPromptsAtRegistration, fc.listPromptsCalls.Load(), + "labelling must not re-aggregate — core.ListPrompts is not called again during the prompt get") + }) + t.Run("no BackendInfo in context is a safe no-op", func(t *testing.T) { t.Parallel() fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} @@ -975,77 +2032,88 @@ func TestServeResolvesBackendNameEndToEnd(t *testing.T) { "the handler must label the audit event with the registry-resolved name, not the raw BackendID") } -// TestServeOmitsPrompts locks in the intentional prompt omission: even when the core -// advertises a prompt, the Serve path injects only tools/resources (the SDK has no -// per-session prompt support), so a prompts/list does not surface it. -func TestServeOmitsPrompts(t *testing.T) { +// TestServeServesPrompts locks in end-to-end per-session prompt serving: the core +// advertises a prompt, the Serve path injects it via SessionWithPrompts, and a +// prompts/list surfaces it while a prompts/get routes through core.GetPrompt and +// returns the core result (the fix for the "unknown prompt" -32602 regression). +func TestServeServesPrompts(t *testing.T) { t.Parallel() - const promptName = "serve-only-prompt" - fc := &fakeCore{prompts: []vmcp.Prompt{{Name: promptName}}} + const promptName = "serve-prompt" + fc := &fakeCore{prompts: []vmcp.Prompt{{Name: promptName, Description: "a served prompt"}}} _, sessionID, baseURL := registerServeSession(t, fc) - resp := postServeMCP(t, baseURL, map[string]any{ + listResp := postServeMCP(t, baseURL, map[string]any{ "jsonrpc": "2.0", "id": 2, "method": "prompts/list", "params": map[string]any{}, }, sessionID) - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) + defer listResp.Body.Close() + listBody, err := io.ReadAll(listResp.Body) require.NoError(t, err) - assert.NotContains(t, string(body), promptName, - "the Serve path must not advertise prompts even when the core supplies them") + assert.Contains(t, string(listBody), promptName, + "the Serve path must advertise prompts injected from the core") + + getResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 3, + "method": "prompts/get", + "params": map[string]any{"name": promptName}, + }, sessionID) + defer getResp.Body.Close() + getBody, err := io.ReadAll(getResp.Body) + require.NoError(t, err) + assert.Equal(t, http.StatusOK, getResp.StatusCode, "prompts/get should succeed; body: %s", string(getBody)) + assert.Contains(t, string(getBody), "prompt-body", + "prompts/get must return the core result, not a -32602 unknown-prompt error") + assert.NotContains(t, string(getBody), "-32602", + "prompts/get must not report the prompt as unknown") + assert.Equal(t, int32(1), fc.getPromptCalls.Load(), + "prompts/get must route through core.GetPrompt exactly once") } // postServeMCP sends a JSON-RPC POST to the given Streamable HTTP base URL. It is the // package-server-internal analogue of the postMCP helper in the external suite. func postServeMCP(t *testing.T, baseURL string, body map[string]any, sessionID string) *http.Response { t.Helper() - rawBody, err := json.Marshal(body) + resp, err := doServeMCP(baseURL, body, sessionID) require.NoError(t, err) + return resp +} + +// doServeMCP performs a tools POST and returns the response or an error. It never +// calls require/assert, so it is safe to invoke from worker goroutines where +// FailNow/Goexit would run off the test goroutine and misreport. postServeMCP +// wraps it for the common test-goroutine case. +func doServeMCP(baseURL string, body map[string]any, sessionID string) (*http.Response, error) { + return doServeMCPWithHeaders(baseURL, body, sessionID, nil) +} + +// doServeMCPWithHeaders is doServeMCP plus caller-supplied extra headers (e.g. +// X-Test-Principal), for tests that need per-request header control. doServeMCP +// delegates here with a nil header map (the range is then a no-op), so both +// share one request-building path. +func doServeMCPWithHeaders( + baseURL string, body map[string]any, sessionID string, headers map[string]string, +) (*http.Response, error) { + rawBody, err := json.Marshal(body) + if err != nil { + return nil, err + } req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, baseURL+"/mcp", bytes.NewReader(rawBody)) - require.NoError(t, err) + if err != nil { + return nil, err + } req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") if sessionID != "" { req.Header.Set("Mcp-Session-Id", sessionID) } + for k, v := range headers { + req.Header.Set(k, v) + } - resp, err := http.DefaultClient.Do(req) - require.NoError(t, err) - return resp -} - -// recordingSessionIDManager records Terminate calls. -type recordingSessionIDManager struct { - terminated []string -} - -func (*recordingSessionIDManager) Generate() string { return "generated" } -func (*recordingSessionIDManager) Validate(string) (bool, error) { return false, nil } -func (m *recordingSessionIDManager) Terminate(id string) (bool, error) { - m.terminated = append(m.terminated, id) - return false, nil -} - -// TestSessionIDManagerResolver verifies that a request (DELETE included) reaches -// the vMCP session manager, while mcp-go's idle sweeper (nil request) gets a -// manager whose Terminate leaves the shared session record alone. -func TestSessionIDManagerResolver(t *testing.T) { - t.Parallel() - - mgr := &recordingSessionIDManager{} - resolver := sessionIDManagerResolver{manager: mgr} - - _, err := resolver.ResolveSessionIdManager(nil).Terminate("swept") - require.NoError(t, err) - assert.Empty(t, mgr.terminated, "the idle sweeper must not terminate the shared session") - assert.Equal(t, "generated", resolver.ResolveSessionIdManager(nil).Generate()) - - req := httptest.NewRequest(http.MethodDelete, "/mcp", nil) - _, err = resolver.ResolveSessionIdManager(req).Terminate("deleted") - require.NoError(t, err) - assert.Equal(t, []string{"deleted"}, mgr.terminated, "a client DELETE must terminate the session") + return http.DefaultClient.Do(req) } diff --git a/pkg/vmcp/server/serve_test.go b/pkg/vmcp/server/serve_test.go index 125d435747..067d3e3d8d 100644 --- a/pkg/vmcp/server/serve_test.go +++ b/pkg/vmcp/server/serve_test.go @@ -45,6 +45,9 @@ func (*stubVMCP) CallTool( func (*stubVMCP) ListResources(context.Context, *auth.Identity) ([]vmcp.Resource, error) { return nil, nil } +func (*stubVMCP) ListResourceTemplates(context.Context, *auth.Identity) ([]vmcp.ResourceTemplate, error) { + return nil, nil +} func (*stubVMCP) ReadResource(context.Context, *auth.Identity, string) (*vmcp.ResourceReadResult, error) { return nil, nil } @@ -54,6 +57,11 @@ func (*stubVMCP) GetPrompt( ) (*vmcp.PromptGetResult, error) { return nil, nil } +func (*stubVMCP) Complete( + context.Context, *auth.Identity, vmcp.CompletionRef, string, string, map[string]string, +) (*vmcp.CompletionResult, error) { + return nil, nil +} func (*stubVMCP) LookupTool(context.Context, *auth.Identity, string) (*vmcp.Tool, error) { return nil, nil } @@ -63,14 +71,23 @@ func (*stubVMCP) LookupResource(context.Context, *auth.Identity, string) (*vmcp. func (*stubVMCP) LookupPrompt(context.Context, *auth.Identity, string) (*vmcp.Prompt, error) { return nil, nil } +func (*stubVMCP) CheckToolCall(context.Context, *auth.Identity, string, map[string]any) error { + return nil +} +func (*stubVMCP) CheckResourceRead(context.Context, *auth.Identity, string) error { return nil } +func (*stubVMCP) CheckPromptGet(context.Context, *auth.Identity, string) error { return nil } func (*stubVMCP) ListBackends(context.Context, *auth.Identity, bool) ([]vmcp.Backend, error) { return nil, nil } func (*stubVMCP) LookupBackend(context.Context, *auth.Identity, string) (*vmcp.Backend, error) { return nil, nil } +func (*stubVMCP) Discover(context.Context, *auth.Identity) (core.DiscoverCapabilities, error) { + return core.DiscoverCapabilities{}, nil +} func (s *stubVMCP) Close() error { s.closed = true; return nil } func (*stubVMCP) BackendHealth() health.Reporter { return nil } +func (*stubVMCP) InvalidateCapabilityCache() {} // stubWatcher is a non-nil Watcher for the drift-guard test; its behavior is not // exercised there. @@ -346,6 +363,7 @@ func TestBuildServeConfigMapsSharedFields(t *testing.T) { EndpointPath: "/e", SessionTTL: time.Second, WriteTimeout: time.Second, + HeartbeatInterval: time.Second, AuthMiddleware: func(h http.Handler) http.Handler { return h }, AuthInfoHandler: http.NewServeMux(), PassthroughHeaders: []string{"x-test"}, diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 46f493b09d..66ed3bdcc8 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -19,11 +19,9 @@ import ( "os" "strings" "sync" - "sync/atomic" "time" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/server" tcredis "github.com/stacklok/toolhive-core/redis" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" @@ -95,13 +93,17 @@ const ( // per tool call; a short TTL collapses bursts of calls into ~one sweep while keeping // staleness tighter than the legacy once-per-session refresh. capabilityCacheTTL = 30 * time.Second - - // defaultSessionRegistrationTimeout bounds the post-initialize session - // registration (opening backend connections and injecting capabilities), - // which runs detached from the initialize request's cancellation. - defaultSessionRegistrationTimeout = 45 * time.Second ) +// heartbeatInterval returns the configured heartbeat interval, or the default +// when the configured value is zero or negative (unset or invalid). +func heartbeatInterval(d time.Duration) time.Duration { + if d <= 0 { + return defaultHeartbeatInterval + } + return d +} + //go:generate mockgen -destination=mocks/mock_watcher.go -package=mocks -source=server.go Watcher // Watcher is the interface for Kubernetes backend watcher integration. @@ -142,6 +144,16 @@ type Config struct { // defaultWriteTimeout are raised to it; larger values let configured backend // request timeouts finish and leave time to write the MCP response. WriteTimeout time.Duration + // HeartbeatInterval configures the SSE keep-alive ping interval on GET + // connections. When zero, the Handler defaults to defaultHeartbeatInterval (30s). + // Prevents proxies/load balancers from closing idle SSE connections. + // + // This is intentionally plumbed end-to-end (ServerConfig → Config → Handler → + // WithHeartbeatInterval) ahead of any CLI-flag or CRD wiring. No external entry + // point sets it yet, so in practice it is always the default unless an embedder + // assigns it programmatically; the plumbing exists so surfacing a flag/field + // later is a one-line change rather than a re-thread through the server. + HeartbeatInterval time.Duration // AuthMiddleware is the optional authentication middleware to apply to MCP routes. // If nil, no authentication is required. @@ -268,6 +280,14 @@ type Server struct { // server (BackendHealth keeps a defensive nil guard regardless). core core.VMCP + // authzGateEnabled installs the pre-dispatch authorization CallGate in Handler. + // Authz is core-owned and stripped from the transport ServerConfig (see + // deriveServerConfig / buildServeConfig), so it never reaches s.config; New — + // the only composition path that knows the authz config — sets this flag from + // cfg.Authz != nil after Serve builds the Server. It stays false for direct-Serve + // callers, which have no way to configure authorization and run allow-all cores. + authzGateEnabled bool + // optimizerFactory builds a per-session optimizer over the core's advertised // tools. Set only when the optimizer is enabled (nil otherwise). When non-nil, // session registration advertises find_tool/ @@ -276,7 +296,7 @@ type Server struct { // session manager; this is the resolved factory surfaced via Manager.OptimizerFactory. optimizerFactory func(context.Context, []server.ServerTool) (optimizer.Optimizer, error) - // MCP protocol server (mark3labs/mcp-go) + // MCP protocol server (stacklok/toolhive-core/mcpcompat) mcpServer *server.MCPServer // HTTP server for Streamable HTTP transport @@ -322,14 +342,19 @@ type Server struct { // returns (or, for an embedded Handler, after the caller stops serving it). shutdownFuncs []func(context.Context) error - // activeClientSessions holds the SDK sessions registered on this replica; it - // feeds the session cardinality gauges and selects the sessions a capability - // refresh clears (serve_refresh.go). - activeClientSessions clientSessionSet - - // inFlightRegistrations counts session registrations currently running the - // post-initialize hook; it feeds the session cardinality gauges. - inFlightRegistrations atomic.Int64 + // refreshSinks maps each session registered on this replica to its + // list_changed sink, so a backend health or registry change can resync the + // session's capabilities (serve_refresh.go). The next refresh drops the + // entries of sessions that ended. + refreshSinks sync.Map + + // resyncBaseCtx is the server-lifetime parent context for asynchronous + // tools/resources/prompts list_changed resync work (#5748, #5969). It is + // cancelled by a shutdownFunc on Stop, so an in-flight backend + // re-aggregation triggered by a late notification cannot outlive the + // server. Set by Serve; nil for direct-Serve callers that never register a + // list_changed sink. + resyncBaseCtx context.Context } // buildSessionDataStorage constructs the DataStorage backend from cfg. @@ -496,9 +521,21 @@ func New( OptimizerFactory: cfg.OptimizerFactory, TelemetryProvider: cfg.TelemetryProvider, AdvertiseFromCore: true, - // New sessions connect only to the backends the core's health monitor - // reports as eligible (nil when health monitoring is disabled). - HealthStatusProvider: coreVMCP.BackendHealth(), + // Gate new-session backend connects on health status (#5861). Without this a + // backend the monitor already knows is bad is still re-attempted by every new + // session, and session creation blocks on it — so one bad backend sets the + // floor for the whole tenant's initialize latency. BackendHealth() returns a + // true nil interface when monitoring is disabled, which the session manager + // reads as "attempt every backend" (the prior behaviour). + // + // A skipped backend's tools stay advertised and callable, but only because + // AdvertiseFromCore is set just above: tools/call then routes through + // core.CallTool over the core's aggregated view rather than the session's own + // routing table, which is built solely from backends that connected. The + // session does lose list_changed propagation for a skipped backend, partly + // offset by InvalidateCapabilityCache being global (serve_list_changed.go) — + // so any other backend's notification opportunistically sweeps in its changes. + BackendHealth: coreVMCP.BackendHealth(), // Close cached sessions idle on this replica for a full session TTL. IdleTTL: resolved.SessionTTL, } @@ -508,9 +545,29 @@ func New( return nil, err } + // Enable the pre-dispatch authorization gate when Cedar policies are configured. + // Authz is core-owned and not carried on the transport ServerConfig Serve builds + // from, so New — which holds cfg.Authz — sets the flag here (see Server.authzGateEnabled). + srv.authzGateEnabled = cfg.Authz != nil + // Bind the elicitation adapter to the SDK server Serve built so composite-workflow // elicitation reaches the same mcp-go server that serves client traffic. - elicitation.bind(NewSDKElicitationAdapter(srv.MCPServer())) + elicitAdapter := NewSDKElicitationAdapter(srv.MCPServer()) + elicitation.bind(elicitAdapter) + + // Bind the server->client forwarders onto the concrete backend client so a + // backend's mid-call elicitation, sampling, and progress/logging traffic is + // relayed to the downstream client on the same session. This mirrors the + // elicitation late-binding above (the SDK adapters wrap the mcp-go server Serve + // just built, so they cannot exist before this point); a backend client that + // does not implement the binder simply does not forward server->client traffic. + if binder, ok := backendClient.(vmcp.ClientForwarderBinder); ok { + binder.BindForwarders( + elicitAdapter, + NewSDKSamplingAdapter(srv.MCPServer()), + NewSDKNotifierAdapter(srv.MCPServer()), + ) + } closeCoreOnErr = false // Serve succeeded; srv.Stop now owns the core's lifecycle. return srv, nil @@ -520,28 +577,28 @@ func New( // This enables embedding the vmcp server inside another HTTP server or framework. // // The returned handler includes all routes (health, metrics, well-known, MCP) -// and the full middleware chain (recovery, body limit, header validation, auth, -// rate limit, audit, MCP parsing, telemetry). +// and the full middleware chain (recovery, body limit, header validation, +// audit, auth, MCP parsing, telemetry). // // Each call builds a fresh handler. The method is safe to call multiple times. // All returned handlers share the same underlying MCPServer and SessionManager, // so callers should not serve concurrent traffic through multiple handlers. func (s *Server) Handler(_ context.Context) (http.Handler, error) { - // Create Streamable HTTP server with ToolHive session management - streamableServer := server.NewStreamableHTTPServer( - s.mcpServer, + // Create Streamable HTTP server with ToolHive session management. + streamableOpts := []server.StreamableHTTPOption{ server.WithEndpointPath(s.config.EndpointPath), - server.WithSessionIdManagerResolver(sessionIDManagerResolver{manager: s.vmcpSessionMgr}), - server.WithHeartbeatInterval(defaultHeartbeatInterval), - // Sweep the SDK's per-session tool/resource stores for sessions that go idle - // on this replica without a DELETE, using the same TTL as the session - // storage. The sweep clears only this replica's state (see - // sessionIDManagerResolver); the storage TTL expires the shared record. - server.WithSessionIdleTTL(s.config.SessionTTL), - ) - // Stop the idle sweeper's goroutine on Stop. Shutdown only cancels the sweeper - // here: the StreamableHTTPServer is used as a handler and owns no listener. - s.shutdownFuncs = append(s.shutdownFuncs, streamableServer.Shutdown) + server.WithSessionIdManager(s.vmcpSessionMgr), + server.WithHeartbeatInterval(heartbeatInterval(s.config.HeartbeatInterval)), + } + // Install the pre-dispatch authorization gate only when authz is configured + // (allow-all otherwise, matching the core admission seam's guard). The gate + // re-runs the core admission decision for gated methods so a denied tools/call, + // resources/read, or prompts/get is rejected as HTTP 403 + JSON-RPC 403 before + // the SDK dispatches it, instead of the SDK's 200 result. + if s.authzGateEnabled { + streamableOpts = append(streamableOpts, server.WithCallGate(s.authzCallGate())) + } + streamableServer := server.NewStreamableHTTPServer(s.mcpServer, streamableOpts...) // Create HTTP mux with separated authenticated and unauthenticated routes mux := http.NewServeMux() @@ -581,22 +638,41 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { } // MCP endpoint - apply middleware chain (wrapping order, execution happens in reverse): - // Code wraps: auth → rate-limit → audit → MCP-parsing → telemetry - // Execution order: recovery → body-limit → header-val → auth → - // rate-limit → audit → MCP-parsing → telemetry → handler + // Code wraps: audit → auth → MCP-parsing → telemetry → classification + // Execution order: recovery → body-limit → header-val → audit → auth → + // MCP-parsing → telemetry → classification → handler // // Upstream token refresh failures are detected inside AuthMiddleware itself: // GetAllUpstreamCredentials returns a non-empty failed-provider slice when // any upstream refresh fails, and the middleware short-circuits with - // HTTP 401 + WWW-Authenticate before the request reaches any inner layer. + // HTTP 401 + WWW-Authenticate. Audit wraps auth, so those 401s (and every + // other rejection from the inner chain) still produce an audit event; the + // authenticated identity and parsed MCP data are read back through the + // holder carriers (auth.IdentityHolder, mcp.ParsedRequestHolder) that the + // inner auth/parser middlewares fill. // // The legacy HTTP authz, annotation-enrichment, and discovery layers have all been // removed: every caller now routes through Serve, so authorization is enforced by the // core admission seam (#5438) and capability/health filtering by the core, rather than // by HTTP middleware. + // + // Authorization additionally runs a pre-dispatch CallGate INSIDE the streamable + // server (installed above only when Authz is configured): it re-runs the core + // admission decision for tools/call, resources/read, and prompts/get and rejects a + // denial as HTTP 403 + JSON-RPC 403 before dispatch. The gate sits before session + // validation (403-before-404) and inside the audit middleware, so a denied call is + // audited with outcome "denied". See call_gate.go. var mcpHandler http.Handler = streamableServer + // Classify Modern (2026-07-28) vs Legacy at the decode seam, reject + // malformed Modern requests before dispatch, and route well-formed Modern + // requests to the vMCP core (dispatchModern) instead of the SDK. Applied + // before telemetry (i.e. it runs closer to the handler) so a rejection or a + // dispatcher 403 is still recorded by the telemetry middleware instead of + // bypassing it entirely. + mcpHandler = s.classifyingHandler(mcpHandler) + if s.config.TelemetryProvider != nil { mcpHandler = s.config.TelemetryProvider.Middleware(s.config.Name, "streamable-http")(mcpHandler) slog.Info("telemetry middleware enabled for MCP endpoints") @@ -611,7 +687,21 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { // when auth middleware is nil. mcpHandler = mcpparser.ParsingMiddleware(mcpHandler) - // Apply audit middleware if configured (runs after auth, before discovery) + // Mark health probes from a peer vMCP (peer_probe.go), so session registration + // and lazy injection skip them, and refuse any probe request beyond the + // handshake and capability lists. Runs after authentication. + mcpHandler = peerProbeMiddleware(mcpHandler) + mcpHandler = s.listenStreamGuardMiddleware(mcpHandler) + + // Apply authentication middleware if configured + if s.config.AuthMiddleware != nil { + mcpHandler = s.config.AuthMiddleware(mcpHandler) + slog.Info("authentication middleware enabled for MCP endpoints") + } + + // Apply audit middleware if configured. It wraps authentication so auth + // failures (401) are audited too; the identity for successful requests is + // read back via the auth.IdentityHolder carrier. if s.config.AuditConfig != nil { if err := s.config.AuditConfig.Validate(); err != nil { return nil, fmt.Errorf("invalid audit configuration: %w", err) @@ -627,18 +717,6 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { slog.Info("audit middleware enabled for MCP endpoints") } - // Mark health probes from a peer vMCP (peer_probe.go), so session registration - // and lazy injection skip them, and refuse any probe request beyond the - // handshake and capability lists. Runs after authentication. - mcpHandler = peerProbeMiddleware(mcpHandler) - mcpHandler = s.listenStreamGuardMiddleware(mcpHandler) - - // Apply authentication middleware if configured (runs first in chain) - if s.config.AuthMiddleware != nil { - mcpHandler = s.config.AuthMiddleware(mcpHandler) - slog.Info("authentication middleware enabled for MCP endpoints") - } - mcpHandler = s.applyForwardedHeaderCapture(mcpHandler) // Apply Accept header validation (rejects GET requests without Accept: text/event-stream) @@ -960,7 +1038,7 @@ func (s *Server) SessionManager() *transportsession.Manager { return s.sessionManager } -// MCPServer returns the underlying mark3labs *server.MCPServer instance +// MCPServer returns the underlying mcpcompat *server.MCPServer instance // servicing this vMCP server's /mcp endpoint. // // Intended for embedders that wrap the vMCP composer in their own pipeline and @@ -971,7 +1049,7 @@ func (s *Server) SessionManager() *transportsession.Manager { // // Trust boundary: this accessor is in-process only; the returned pointer is // the same instance for the lifetime of the Server and is safe for concurrent -// use per mark3labs guarantees. +// use per mcpcompat guarantees. // // Safe operations include RequestElicitation against an active session, // registering observability hooks, and reading registered @@ -1023,6 +1101,47 @@ func setSessionResourcesDirect(session server.ClientSession, resources []server. return nil } +// setSessionResourceTemplatesDirect sets resource templates directly on the session via +// the SessionWithResourceTemplates interface, analogous to setSessionResourcesDirect for +// resource templates. Keyed by URI template string. +func setSessionResourceTemplatesDirect(session server.ClientSession, templates []server.ServerResourceTemplate) error { + sessionWithTemplates, ok := session.(server.SessionWithResourceTemplates) + if !ok { + return fmt.Errorf("session does not support per-session resource templates") + } + + existing := sessionWithTemplates.GetSessionResourceTemplates() + templateMap := make(map[string]server.ServerResourceTemplate, len(existing)+len(templates)) + for k, v := range existing { + templateMap[k] = v + } + for _, tmpl := range templates { + templateMap[tmpl.Template.URITemplate] = tmpl + } + sessionWithTemplates.SetSessionResourceTemplates(templateMap) + return nil +} + +// setSessionPromptsDirect sets prompts directly on the session via the SessionWithPrompts +// interface, analogous to setSessionResourcesDirect for prompts. +func setSessionPromptsDirect(session server.ClientSession, prompts []server.ServerPrompt) error { + sessionWithPrompts, ok := session.(server.SessionWithPrompts) + if !ok { + return fmt.Errorf("session does not support per-session prompts") + } + + existing := sessionWithPrompts.GetSessionPrompts() + promptMap := make(map[string]server.ServerPrompt, len(existing)+len(prompts)) + for k, v := range existing { + promptMap[k] = v + } + for _, p := range prompts { + promptMap[p.Prompt.Name] = p + } + sessionWithPrompts.SetSessionPrompts(promptMap) + return nil +} + // setSessionToolsDirect sets tools directly on the session via the SessionWithTools // interface, bypassing MCPServer.AddSessionTools. This avoids sending notifications // through the session's notification channel, which would accumulate as stale @@ -1127,65 +1246,27 @@ func (s *Server) lazyInjectSessionTools(ctx context.Context, wantTool string) { } } -// callerOwnsSession reports whether the caller in ctx matches the identity binding stored for -// sessionID. A missing, terminated or unbound session does not match. +// callerOwnsSession reports whether the caller in ctx matches the identity binding of +// sessionID (enforceSessionBinding, the check every call runs). func (s *Server) callerOwnsSession(ctx context.Context, sessionID string) bool { - binding, err := s.storedBinding(ctx, sessionID) - if err != nil { - return false - } caller, _ := auth.IdentityFromContext(ctx) - return vmcpsession.ValidateCaller(binding, caller) == nil -} - -// lazyInjectSessionResources is the resources counterpart of -// lazyInjectSessionTools: it fills an empty per-session resource store (a session -// that was registered on another pod) from core.ListResources for the request -// identity. The call-time binding check in coreResourceHandler stays the backstop. -func (s *Server) lazyInjectSessionResources(ctx context.Context) { - sess := server.ClientSessionFromContext(ctx) - if sess == nil || isPeerProbe(ctx) { - return - } - sessionWithResources, ok := sess.(server.SessionWithResources) - if !ok || len(sessionWithResources.GetSessionResources()) > 0 { - return - } - sessionID := sess.SessionID() - identity, _ := auth.IdentityFromContext(ctx) - resources, err := s.coreSessionResources(ctx, sessionID, identity) - if err != nil || len(resources) == 0 { - slog.Debug("lazyInjectSessionResources: no resources available for session", "session_id", sessionID) - return - } - if err := setSessionResourcesDirect(sess, resources); err != nil { - slog.Warn("lazyInjectSessionResources: failed to inject resources", "session_id", sessionID, "error", err) - } + return s.enforceSessionBinding(ctx, sessionID, caller) == nil } // handleSessionRegistration processes a new MCP session registration. // It fires AFTER the session is registered in the SDK. -// -// A health probe from a peer vMCP only checks that the MCP handshake works, so -// it neither opens backend connections nor lists capabilities, and it is not -// counted as an active session. A GET registration is never skipped as a probe: -// it always goes through the identity check in attachToBoundSession. func (s *Server) handleSessionRegistration( ctx context.Context, session server.ClientSession, ) { - if isPeerProbe(ctx) && !isListenStreamRegistration(ctx) { + // A health probe from a peer vMCP only checks that the MCP handshake works, so + // it neither opens backend connections nor lists capabilities. + if isPeerProbe(ctx) { return } // Error is logged and handled within handleSessionRegistrationImpl. // The session is terminated on failure; no further action needed here. - // An unregistration while this runs cancels the commit (clientSessionSet). - s.activeClientSessions.begin(session.SessionID()) - if err := s.handleSessionRegistrationImpl(ctx, session); err != nil { - s.activeClientSessions.abort(session.SessionID()) - return - } - s.activeClientSessions.commit(session) + _ = s.handleSessionRegistrationImpl(ctx, session) } // handleSessionRegistrationImpl handles session registration. @@ -1205,31 +1286,10 @@ func (s *Server) handleSessionRegistration( // indexed into the optimizer and only find_tool/call_tool are exposed, // using session-scoped tool handlers. // -// - Prompts: not supported until the SDK adds AddSessionPrompts. +// - Prompts: injected per session via SessionWithPrompts, mirroring resources; +// see injectCoreSessionCapabilities. func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session server.ClientSession) (retErr error) { - s.inFlightRegistrations.Add(1) - defer s.inFlightRegistrations.Add(-1) - - // Only the session's own initialize may create it. mcp-go v0.55 registers a - // session from the initialize POST and from a GET listen stream (or SSE - // reconnect) on a pod that does not hold the session in memory; the latter - // attaches to the bound session or is rejected, and never calls CreateSession. - // It keeps the context of the GET request, so its wait for the binding ends when the - // client goes away. Checked before the cleanup defer below: rejecting a - // re-registration must not terminate the owner's session. - if isListenStreamRegistration(ctx) { - return s.attachToBoundSession(ctx, session) - } - - // The SDK runs this hook with the initialize request's context, whose response - // is already written. A client that closes the connection fast would cancel - // backend initialization half-way. Keep the request values (identity, session), - // drop the cancellation, and bound the work with a server-side timeout. - ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), defaultSessionRegistrationTimeout) - defer cancel() - sessionID := session.SessionID() - slog.Debug("creating session-scoped backends", "session_id", sessionID) // Defer cleanup: if any error occurs, terminate the session and log failures. @@ -1252,15 +1312,29 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv // is no way to retract the session ID after it has been sent. // // NOTE: there is a brief race window between the client receiving the session - // ID and this hook completing. A tools/call that arrives in it gets its tools - // from lazyInjectSessionTools, and enforceSessionBinding waits until - // CreateSession below has written the identity binding. - if _, retErr = s.vmcpSessionMgr.CreateSession(ctx, sessionID); retErr != nil { + // ID and this hook completing. A client that pipelines a tools/call immediately + // after initialize may receive a "tool not found" error before AddSessionTools + // completes. Conforming MCP clients call tools/list before tools/call, so this + // window is expected to be harmless in practice. + // + // The list_changed sink is built here, before CreateSession, so it can be + // threaded through to every backend connector this session opens (#5748). + // Capture BOTH the caller identity AND the per-request forwarded headers from + // the registration context: the asynchronous resync re-enumerates backends + // through the same identity+header-keyed path a live request uses (cache key + // and outbound backend auth are derived from the context, not passed + // explicitly), so it must reconstruct a faithful context or it would + // enumerate unauthenticated. See buildListChangedSink. + identity, _ := auth.IdentityFromContext(ctx) + forwardedHeaders := headerforward.ForwardedHeadersFromContext(ctx) + sink := s.buildListChangedSink(sessionID, session, identity, forwardedHeaders) + if _, retErr = s.vmcpSessionMgr.CreateSession(ctx, sessionID, sink); retErr != nil { slog.Error("failed to create session-scoped backends", "session_id", sessionID, "error", retErr) return retErr } + s.refreshSinks.Store(sessionID, sink) // The core is the single authoritative aggregation: source the advertised tool/resource // set from core.ListTools/ListResources (called once per session here) and install @@ -1270,65 +1344,6 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv return s.injectCoreSessionCapabilities(ctx, session) } -// attachToBoundSession handles a registration that does not come from the -// session's own initialize. mcp-go registers a session again when a GET listen -// stream (or an SSE reconnect) reaches a pod that does not hold it in memory, and -// CreateSession would then build a new session with the caller's identity and -// overwrite the stored binding. Instead, this never creates a session: it gives a -// caller that passes checkListenStreamCaller the session's capabilities. Any -// other caller is rejected and unregistered from this pod; the stored session is -// not changed in any case. -// -// listenStreamGuardMiddleware already runs the same check before mcp-go sees the -// GET, and rejects the request with an HTTP error; this is the fallback for a -// registration that reaches the hook anyway. -// -// Upstream stops the re-bind only from v0.40.0, as a side effect of #5729. Drop -// this guard once the fork is on v0.40 or later and TestCrossPodSessionRebind and -// TestCrossPodGETDuringRegistration pass without it. -func (s *Server) attachToBoundSession(ctx context.Context, session server.ClientSession) error { - sessionID := session.SessionID() - if _, err := s.checkListenStreamCaller(ctx, sessionID); err != nil { - slog.Warn("rejected re-registration of a session", "session_id", sessionID, "error", err) - s.mcpServer.UnregisterSession(ctx, sessionID) - return err - } - return s.injectCoreSessionCapabilities(ctx, session) -} - -// checkListenStreamCaller decides whether the caller in ctx may open a listen -// stream on sessionID. It waits while the creating pod's registration is still -// pending (see waitWhileRegistrationPending), then requires a live session with -// an identity binding (else it returns 404) that matches the caller (else 403). -func (s *Server) checkListenStreamCaller(ctx context.Context, sessionID string) (int, error) { - s.waitWhileRegistrationPending(ctx, sessionID) - binding, err := s.storedBinding(ctx, sessionID) - if err != nil { - return http.StatusNotFound, err - } - caller, _ := auth.IdentityFromContext(ctx) - if err := vmcpsession.ValidateCaller(binding, caller); err != nil { - return http.StatusForbidden, err - } - return http.StatusOK, nil -} - -// storedBinding returns the identity binding of a live stored session record. -func (s *Server) storedBinding(ctx context.Context, sessionID string) (string, error) { - metadata, err := s.sessionDataStorage.Load(ctx, sessionID) - if err != nil { - return "", fmt.Errorf("load session %s: %w", sessionID, err) - } - if metadata[sessionmanager.MetadataKeyTerminated] == sessionmanager.MetadataValTrue { - return "", fmt.Errorf("session %s is terminated", sessionID) - } - binding, bound := metadata[vmcpsession.MetadataKeyIdentityBinding] - if !bound { - return "", errSessionNotRegistered - } - return binding, nil -} - // backendHealth returns the core-owned backend health reporter, or nil when health // monitoring is disabled (the core owns the monitor's lifecycle; #5443 reversal). The // s.core nil guard is defensive: a Serve-built server always has a non-nil core. @@ -1437,101 +1452,27 @@ var notAcceptableBody = []byte( `{"code":-32600,"message":"Not Acceptable: Client must accept text/event-stream"}}`, ) -// sessionIDManagerResolver gives mcp-go the vMCP session manager for every -// request. mcp-go's idle sweeper (WithSessionIdleTTL) resolves the manager with a -// nil request and calls Terminate for each session that is idle on this replica; -// the vMCP manager's Terminate would delete the shared session record for every -// replica, including one where the session is still in use. For the sweeper the -// resolver returns localSweepSessionIDManager, whose Terminate keeps the shared -// record; mcp-go then clears only this replica's SDK state. The node-local -// MultiSession is closed by the session manager's own idle sweep -// (sessionmanager.FactoryConfig.IdleTTL), and the storage TTL expires the shared -// record once no replica uses the session. -type sessionIDManagerResolver struct { - manager server.SessionIdManager -} - -// ResolveSessionIdManager implements server.SessionIdManagerResolver. -func (r sessionIDManagerResolver) ResolveSessionIdManager(req *http.Request) server.SessionIdManager { - if req == nil { - return localSweepSessionIDManager{SessionIdManager: r.manager} - } - return r.manager -} - -// localSweepSessionIDManager is the session ID manager mcp-go's idle sweeper -// uses: Terminate is a no-op, so the sweep never touches the shared session record. -type localSweepSessionIDManager struct { - server.SessionIdManager -} - -// Terminate keeps the shared session record; see sessionIDManagerResolver. -func (localSweepSessionIDManager) Terminate(string) (bool, error) { - return false, nil -} - -// listenStreamKey marks the context of a GET request (the Streamable HTTP listen -// stream) that passed listenStreamGuardMiddleware, for handleSessionRegistrationImpl. -type listenStreamKey struct{} - -// listenStreamGuardMiddleware checks every GET (the Streamable HTTP listen -// stream) before mcp-go sees it. mcp-go's handleGet attaches the GET to the -// session object it holds in memory, or registers a new one for the session ID, -// and then always answers 200: a registration hook cannot fail the request, and -// on a pod that already holds the session no hook runs at all. Without this -// check a caller with someone else's session ID would receive that session's -// notifications and elicitation requests, and on a pod that did not hold the -// session the caller's object would stay in mcp-go for later requests. -// -// A GET without Mcp-Session-Id gets 400: a vMCP session starts with initialize, -// and mcp-go would otherwise generate a session ID that nothing ever binds. A GET -// with a session ID must pass checkListenStreamCaller (404 for a session that is -// missing, terminated or never bound, 403 for a different caller). A GET that -// passes is marked for handleSessionRegistrationImpl, so a registration that -// mcp-go runs for it never creates the session. Runs after authentication. +// listenStreamGuardMiddleware refuses (403) a GET listen stream for a session bound +// to another identity. The SDK opens the stream for any caller that holds the +// session ID, on any pod, and routes the session's server-to-client requests +// (elicitation, sampling) to it. A GET for a missing or ended session is left to +// the SDK, which answers 404 so the client initializes again. func (s *Server) listenStreamGuardMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method != http.MethodGet { - next.ServeHTTP(w, r) - return - } - sessionID := r.Header.Get(server.HeaderKeySessionID) - if sessionID == "" { - writeJSONRPCHTTPError(w, http.StatusBadRequest, "Bad Request: Mcp-Session-Id header is required") - return - } - if status, err := s.checkListenStreamCaller(r.Context(), sessionID); err != nil { - slog.Warn("rejected listen stream", "session_id", sessionID, "status", status, "error", err) - writeJSONRPCHTTPError(w, status, http.StatusText(status)) - return + sessionID := r.Header.Get("Mcp-Session-Id") + if r.Method == http.MethodGet && sessionID != "" { + if sess, ok := s.vmcpSessionMgr.GetMultiSession(r.Context(), sessionID); ok { + binding, _ := sess.GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding) + caller, _ := auth.IdentityFromContext(r.Context()) + if err := vmcpsession.ValidateCaller(binding, caller); err != nil { + slog.Warn("refused a listen stream for a session bound to another caller", "session_id", sessionID) + http.Error(w, "Forbidden", http.StatusForbidden) + return + } + } } - next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), listenStreamKey{}, true))) - }) -} - -// writeJSONRPCHTTPError writes an HTTP error status with a JSON-RPC error body, -// in the same shape as notAcceptableBody. -func writeJSONRPCHTTPError(w http.ResponseWriter, status int, message string) { - body, err := json.Marshal(map[string]any{ - "jsonrpc": "2.0", - "id": "server-error", - "error": map[string]any{"code": -32600, "message": message}, + next.ServeHTTP(w, r) }) - if err != nil { - http.Error(w, message, status) - return - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(status) - if _, err := w.Write(body); err != nil { - slog.Debug("failed to write error response", "error", err) - } -} - -// isListenStreamRegistration reports whether ctx belongs to a GET request. -func isListenStreamRegistration(ctx context.Context) bool { - marked, _ := ctx.Value(listenStreamKey{}).(bool) - return marked } // headerValidatingMiddleware rejects GET requests that do not include diff --git a/pkg/vmcp/server/session_lifecycle_regression_test.go b/pkg/vmcp/server/session_lifecycle_regression_test.go new file mode 100644 index 0000000000..22c75e2924 --- /dev/null +++ b/pkg/vmcp/server/session_lifecycle_regression_test.go @@ -0,0 +1,155 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "io" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// readServeJSONRPC reads an HTTP response body from the Serve-path streamable +// server and unmarshals the JSON-RPC envelope into a map. The streamable server +// is configured with JSONResponse:true, so the body is a single application/json +// document (not an SSE stream). It also returns the raw body bytes for +// substring assertions. +func readServeJSONRPC(t *testing.T, resp *http.Response) (map[string]any, []byte) { + t.Helper() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var env map[string]any + require.NoError(t, json.Unmarshal(body, &env), "body: %s", string(body)) + return env, body +} + +// serveDelete sends an HTTP DELETE to /mcp with the given session ID. +func serveDelete(t *testing.T, baseURL, sessionID string) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext( + context.Background(), http.MethodDelete, baseURL+"/mcp", nil, + ) + require.NoError(t, err) + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// TestRegression_TerminatedSessionRejected verifies that a session terminated +// via the SDK's DELETE path (which drives vmcpSessionMgr.Terminate under the +// hood and forgets the SDK's local session) is no longer accepted for +// subsequent requests: a tools/list POST with the dead session ID must yield a +// 4xx response — the SDK's Validate rejects the unknown/terminated session +// before any handler runs — not a successful JSON-RPC result. +// +// Note: calling vmcpSessionMgr.Terminate alone does NOT immediately evict the +// SDK's in-memory local session (lazy eviction — the entry is closed when the +// cache is later re-checked). Only the DELETE path forgets the SDK session in +// lockstep with storage termination, so the rejection is observable on the +// next request. This mirrors the proven assertion in +// TestIntegration_SessionManagement_Termination. +func TestRegression_TerminatedSessionRejected(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, baseURL := registerServeSession(t, fc) + + // Terminate via the SDK DELETE path: this calls vmcpSessionMgr.Terminate + // (storage delete) AND forgets the SDK's local session, so the next request + // is rejected by Validate rather than served from a stale local copy. + delResp := serveDelete(t, baseURL, sessionID) + defer delResp.Body.Close() + require.Equal(t, http.StatusOK, delResp.StatusCode, "DELETE should return 200") + + // The termination must be reflected in storage before we assert. + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return !ok + }, 2*time.Second, 10*time.Millisecond, "session must be gone from the manager after Terminate") + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + + // A terminated session must NOT be accepted: reject with a 4xx (the SDK's + // Validate returns 404 for an unknown/terminated session). A 200 here would + // mean a terminated session was still serving requests — a security regression. + assert.GreaterOrEqual(t, resp.StatusCode, 400, + "terminated session must be rejected, got status %d", resp.StatusCode) + assert.Less(t, resp.StatusCode, 500, + "rejection should be a client error (4xx), not a server error; got %d", resp.StatusCode) +} + +// TestRegression_DELETE_TerminationEvictsCache verifies that an HTTP DELETE to +// /mcp with a valid Mcp-Session-Id terminates the session: the response is 200 +// (the streamable server rewrites the SDK's 204 to 200 for mcp-go compatibility) +// and the session is subsequently absent from vmcpSessionMgr. +func TestRegression_DELETE_TerminationEvictsCache(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, baseURL := registerServeSession(t, fc) + + delResp := serveDelete(t, baseURL, sessionID) + defer delResp.Body.Close() + + // The streamable server rewrites 204 → 200 and forwards the termination to + // the SessionIdManager, so DELETE returns 200 OK on the happy path. + assert.Equal(t, http.StatusOK, delResp.StatusCode, + "DELETE on a live session should return 200 (got %d)", delResp.StatusCode) + + // The session must be evicted from the vMCP session manager. + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return !ok + }, 2*time.Second, 10*time.Millisecond, "session must be evicted from vmcpSessionMgr after DELETE") +} + +// TestRegression_SessionIdentityBinding_SecondPrincipalRejected is a thin +// regression re-assertion that the Serve call path enforces the session's +// identity binding: a session bound to the anonymous identity rejects a caller +// presenting a different principal. The exhaustive test lives in +// TestServeEnforcesSessionBinding; this guards the regression at the tool-handler +// boundary by invoking coreToolHandler directly with a non-anonymous identity. +func TestRegression_SessionIdentityBinding_SecondPrincipalRejected(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, _ := registerServeSession(t, fc) + + // The registered session is bound to "unauthenticated" (anonymous) by + // newToolSessionFactory. A caller presenting a token is a session-upgrade + // attack and must be rejected before the core is reached. + ctx := auth.WithIdentity(context.Background(), &auth.Identity{Token: "attacker-token"}) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "t", Arguments: map[string]any{}}} + + res, err := srv.coreToolHandler(sessionID, "t", "")(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + assert.True(t, res.IsError, "a second principal must be rejected with an error result") + + body, err := json.Marshal(res) + require.NoError(t, err) + assert.Contains(t, string(body), "Unauthorized", + "the rejection must carry an unauthorized message") + // The core must not have been reached on a binding failure. + assert.Equal(t, int32(0), fc.callToolCalls.Load(), + "core.CallTool must not be reached when the binding check fails") +} diff --git a/pkg/vmcp/server/session_management_integration_test.go b/pkg/vmcp/server/session_management_integration_test.go index 96f247f652..b071d5d217 100644 --- a/pkg/vmcp/server/session_management_integration_test.go +++ b/pkg/vmcp/server/session_management_integration_test.go @@ -15,12 +15,12 @@ import ( "testing" "time" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpsdk "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpsdk "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" @@ -47,8 +47,8 @@ func newNoopMockFactory(t *testing.T) *sessionfactorymocks.MockMultiSessionFacto t.Helper() ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { mock := sessionmocks.NewMockMultiSession(ctrl) mock.EXPECT().ID().Return(id).AnyTimes() mock.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes() @@ -89,8 +89,8 @@ func newMockFactory(t *testing.T, ctrl *gomock.Controller, tools []vmcp.Tool) (* t.Helper() state := &mockFactoryState{} factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { state.makeWithIDCalled.Store(true) // All sessions carry MetadataKeyIdentityBinding so Terminate takes the // Phase 2 (storage.Delete) path. The sentinel value is sufficient for @@ -157,9 +157,19 @@ type serverOptions struct { // session factory), so tests provide it here; buildTestServerWithOptions builds a stub // aggregator from it (and a routing table so composite-tool reachability and call // routing work). - tools []vmcp.Tool + tools []vmcp.Tool + // hiddenTools are put in the routing table but NOT the advertised list, + // reproducing what excludeAllTools / per-workload excludeAll / filter produce + // (aggregator/default_aggregator.go:349): routable so composite workflow steps + // can reach them, absent from tools/list. Used by the hidden-tool regression + // test; leave nil otherwise. + hiddenTools []vmcp.Tool workflowDefs map[string]*composer.WorkflowDefinition optimizerFactory func(context.Context, []mcpsdk.ServerTool) (optimizer.Optimizer, error) + // onBackendCall, when set, is invoked with the tool name of every call that + // reaches BackendClient.CallTool. Lets a test assert a name was never + // forwarded, which a response-level assertion alone cannot prove. + onBackendCall func(toolName string) } // capsFromTools builds the AggregatedCapabilities the stub aggregator returns: the tools @@ -167,11 +177,17 @@ type serverOptions struct { // This mirrors the routing table the legacy mock session built (newMockFactory), so the // core advertises the tools, resolves composite-tool reachability, and routes tools/call // through BackendClient.CallTool. -func capsFromTools(tools []vmcp.Tool) *aggregator.AggregatedCapabilities { - rt := &vmcp.RoutingTable{Tools: make(map[string]*vmcp.BackendTarget, len(tools))} +// +// hidden names land in the routing table only, never in the advertised list — see +// serverOptions.hiddenTools. +func capsFromTools(tools, hidden []vmcp.Tool) *aggregator.AggregatedCapabilities { + rt := &vmcp.RoutingTable{Tools: make(map[string]*vmcp.BackendTarget, len(tools)+len(hidden))} for i := range tools { rt.Tools[tools[i].Name] = &vmcp.BackendTarget{WorkloadID: tools[i].Name} } + for i := range hidden { + rt.Tools[hidden[i].Name] = &vmcp.BackendTarget{WorkloadID: hidden[i].Name} + } return &aggregator.AggregatedCapabilities{Tools: tools, RoutingTable: rt} } @@ -210,10 +226,19 @@ func buildTestServerWithOptions( // Stop is called when the server is stopped (not via httptest but via session manager cleanup). // tools/call routes through core.CallTool → BackendClient.CallTool (the session factory's // own CallTool is bypassed on the Serve path). Return a deterministic result so call - // tests can assert on it. + // tests can assert on it. opts.onBackendCall (when set) records the tool name each + // forwarded call carries, letting a test assert a name never reached the backend at all. mockBackendClient.EXPECT(). - CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - Return(&vmcp.ToolCallResult{Content: []vmcp.Content{{Type: "text", Text: "fake result"}}}, nil). + CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func( + _ context.Context, _ *vmcp.BackendTarget, toolName string, + _ map[string]any, _ map[string]any, _ map[string]string, + ) (*vmcp.ToolCallResult, error) { + if opts.onBackendCall != nil { + opts.onBackendCall(toolName) + } + return &vmcp.ToolCallResult{Content: []vmcp.Content{{Type: "text", Text: "fake result"}}}, nil + }). AnyTimes() rt := router.NewSessionRouter(&vmcp.RoutingTable{}) @@ -226,7 +251,7 @@ func buildTestServerWithOptions( SessionTTL: 5 * time.Minute, SessionFactory: factory, OptimizerFactory: opts.optimizerFactory, - Aggregator: newStubAggregator(capsFromTools(opts.tools)), + Aggregator: newStubAggregator(capsFromTools(opts.tools, opts.hiddenTools)), }, rt, mockBackendClient, diff --git a/pkg/vmcp/server/session_management_realbackend_integration_test.go b/pkg/vmcp/server/session_management_realbackend_integration_test.go index f05a7907fa..bcd265aa74 100644 --- a/pkg/vmcp/server/session_management_realbackend_integration_test.go +++ b/pkg/vmcp/server/session_management_realbackend_integration_test.go @@ -36,7 +36,9 @@ import ( // startRealMCPBackend is defined in testutil_test.go as a shared test utility. // newRealTestHandler builds the full vMCP handler backed by the MCP server at -// backendURL. It is the low-level helper used by newRealTestServer and any test +// backendURL. A well-formed Modern (2026-07-28) request routes through +// classifyingHandler -> dispatchModern; a Legacy request falls through to the +// SDK path. It is the low-level helper used by newRealTestServer and any test // that needs control over the httptest.Server configuration (e.g. WriteTimeout). func newRealTestHandler(t *testing.T, backendURL string) http.Handler { t.Helper() @@ -100,8 +102,10 @@ func newRealTestHandler(t *testing.T, backendURL string) http.Handler { } // newRealTestServer builds a vMCP server with session management and a real -// SessionFactory. The BackendRegistry mock returns the backend at backendURL -// so that CreateSession() opens a real HTTP connection to the MCP server. +// SessionFactory. The BackendRegistry mock returns the backend at backendURL so +// that CreateSession() opens a real HTTP connection to the MCP server. Serves +// both revisions: Legacy through the SDK path, Modern through +// classifyingHandler -> dispatchModern. func newRealTestServer(t *testing.T, backendURL string) *httptest.Server { t.Helper() ts := httptest.NewServer(newRealTestHandler(t, backendURL)) @@ -237,6 +241,119 @@ func TestIntegration_RealBackend_NonSSEGetRejectedWithNotAcceptable(t *testing.T "GET without Accept: text/event-stream must be rejected with 406") } +// TestIntegration_RealBackend_ModernRequestRejectedByClassification verifies +// the classifyingHandler wiring end-to-end through the real chain (not +// just the unit-level table covering classifyingHandler in isolation): +// a request that signals Modern (2026-07-28) via a reserved _meta key, but +// carries a mismatched MCP-Protocol-Version header, is rejected with a +// -32020 (HeaderMismatch) JSON-RPC error before ever reaching the backend. +func TestIntegration_RealBackend_ModernRequestRejectedByClassification(t *testing.T) { + t.Parallel() + + // Rejected by classifyingHandler before dispatch, so no real MCP + // backend is needed. + ts := newRealTestServer(t, "http://127.0.0.1:0") + + body := map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{ + // Presence of a reserved io.modelcontextprotocol/* _meta key is + // itself a Modern signal, regardless of its value (see + // pkg/mcp/revision.go). No protocolVersion is present, so the + // non-empty header below cannot be validated against the body + // and the request is a hard rejection. + "_meta": map[string]any{"io.modelcontextprotocol/clientInfo": map[string]any{"name": "test"}}, + "name": "echo", + "arguments": map[string]any{}, + }, + } + payload, err := json.Marshal(body) + require.NoError(t, err) + + req, err := http.NewRequestWithContext( + context.Background(), http.MethodPost, ts.URL+"/mcp", bytes.NewReader(payload)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", "2025-11-25") + + resp, err := ts.Client().Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusBadRequest, resp.StatusCode, "body: %s", string(respBody)) + + var rpc struct { + Error struct { + Code int64 `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(respBody, &rpc), "body: %s", string(respBody)) + assert.Equal(t, int64(-32020), rpc.Error.Code, "expected CodeHeaderMismatch") +} + +// TestIntegration_RealBackend_ModernRequestRejectedByHeaderMismatch verifies +// the header-consistency path end-to-end through the real chain: a +// well-formed Modern (2026-07-28) request — valid _meta.protocolVersion, +// valid clientCapabilities, and a matching MCP-Protocol-Version header — is +// still rejected with -32020 (HeaderMismatch) when its Mcp-Method HTTP header +// disagrees with the JSON-RPC body's actual method. This exercises the real +// ParsingMiddleware -> classifyingHandler flow (unlike +// TestClassificationMiddleware in classification_test.go, which injects a +// pre-built ParsedMCPRequest and bypasses ParsingMiddleware), and covers a +// genuine Mcp-Method/body mismatch rather than the protocolVersion mismatch +// already covered by TestIntegration_RealBackend_ModernRequestRejectedByClassification. +func TestIntegration_RealBackend_ModernRequestRejectedByHeaderMismatch(t *testing.T) { + t.Parallel() + + // Rejected by classifyingHandler before dispatch, so no real MCP + // backend is needed. + ts := newRealTestServer(t, "http://127.0.0.1:0") + + body := map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{ + "_meta": map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": map[string]any{}, + }, + "name": "echo", + "arguments": map[string]any{}, + }, + } + payload, err := json.Marshal(body) + require.NoError(t, err) + + req, err := http.NewRequestWithContext( + context.Background(), http.MethodPost, ts.URL+"/mcp", bytes.NewReader(payload)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", "2026-07-28") + // Mismatched: the body's actual method is "tools/call". + req.Header.Set("Mcp-Method", "resources/read") + + resp, err := ts.Client().Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusBadRequest, resp.StatusCode, "body: %s", string(respBody)) + + var rpc struct { + Error struct { + Code int64 `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(respBody, &rpc), "body: %s", string(respBody)) + assert.Equal(t, int64(-32020), rpc.Error.Code, "expected CodeHeaderMismatch") +} + // TestIntegration_RealBackend_Termination verifies the session termination path // against a real backend: a DELETE request closes the backend connection, and // subsequent requests with the terminated session ID are rejected. diff --git a/pkg/vmcp/server/session_manager_interface.go b/pkg/vmcp/server/session_manager_interface.go index 70c6ab8ebe..acef22791a 100644 --- a/pkg/vmcp/server/session_manager_interface.go +++ b/pkg/vmcp/server/session_manager_interface.go @@ -6,8 +6,7 @@ package server import ( "context" - mcpserver "github.com/mark3labs/mcp-go/server" - + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" ) @@ -25,7 +24,13 @@ type SessionManager interface { // CreateSession completes Phase 2 of the two-phase session creation pattern. // Called from OnRegisterSession hook once context is available; creates backend // connections and replaces the placeholder with a fully-formed MultiSession. - CreateSession(ctx context.Context, sessionID string) (vmcpsession.MultiSession, error) + // + // sink, when non-nil, is threaded to MultiSessionFactory.MakeSessionWithID and + // reaches every backend connector opened for this session (#5748). Pass nil to + // disable asynchronous list_changed consumption for this session. + CreateSession( + ctx context.Context, sessionID string, sink vmcpsession.ListChangedSink, + ) (vmcpsession.MultiSession, error) // GetMultiSession retrieves the fully-formed MultiSession for the given session ID. // Returns (nil, false) if the session does not exist or is still a placeholder. diff --git a/pkg/vmcp/server/session_registration_context_test.go b/pkg/vmcp/server/session_registration_context_test.go deleted file mode 100644 index 1af326b794..0000000000 --- a/pkg/vmcp/server/session_registration_context_test.go +++ /dev/null @@ -1,84 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "errors" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/stacklok/toolhive/pkg/auth" - transportsession "github.com/stacklok/toolhive/pkg/transport/session" - vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" -) - -type registrationContextKey struct{} - -// registrationContextRecorder records the context CreateSession receives and then -// fails, so handleSessionRegistrationImpl stops before it needs a core. Methods -// not overridden here are never called by the code under test. -type registrationContextRecorder struct { - SessionManager - - ctxErr error - deadline time.Time - hasDeadline bool - subject string - value any - inFlight int64 - srv *Server -} - -var errStopRegistration = errors.New("stop after CreateSession") - -func (r *registrationContextRecorder) CreateSession(ctx context.Context, _ string) (vmcpsession.MultiSession, error) { - r.ctxErr = ctx.Err() - r.deadline, r.hasDeadline = ctx.Deadline() - if identity, ok := auth.IdentityFromContext(ctx); ok { - r.subject = identity.Subject - } - r.value = ctx.Value(registrationContextKey{}) - r.inFlight = r.srv.inFlightRegistrations.Load() - return nil, errStopRegistration -} - -func (*registrationContextRecorder) Terminate(string) (bool, error) { return false, nil } - -func TestHandleSessionRegistrationImpl_DetachedBoundedContext(t *testing.T) { - t.Parallel() - - rec := ®istrationContextRecorder{} - // An empty store: the session has no bound record, so registration goes on to - // CreateSession (see attachToBoundSession). - storage, err := transportsession.NewLocalSessionDataStorage(time.Minute) - require.NoError(t, err) - t.Cleanup(func() { _ = storage.Close() }) - srv := &Server{vmcpSessionMgr: rec, sessionDataStorage: storage} - rec.srv = srv - - ctx := auth.WithIdentity(context.Background(), &auth.Identity{ - PrincipalInfo: auth.PrincipalInfo{Subject: "subject-123"}, - }) - ctx = context.WithValue(ctx, registrationContextKey{}, "kept") - canceledCtx, cancel := context.WithCancel(ctx) - cancel() - - err = srv.handleSessionRegistrationImpl(canceledCtx, &fakeSDKSession{id: "session-123"}) - require.ErrorIs(t, err, errStopRegistration) - - assert.NoError(t, rec.ctxErr, "request cancellation must not reach session creation") - assert.Equal(t, "subject-123", rec.subject, "identity must survive the detach") - assert.Equal(t, "kept", rec.value, "request values must survive the detach") - require.True(t, rec.hasDeadline, "detached context must be bounded") - remaining := time.Until(rec.deadline) - assert.Greater(t, remaining, defaultSessionRegistrationTimeout-5*time.Second) - assert.LessOrEqual(t, remaining, defaultSessionRegistrationTimeout) - - assert.Equal(t, int64(1), rec.inFlight, "registration counts as in flight while it runs") - assert.Equal(t, int64(0), srv.inFlightRegistrations.Load(), "in-flight count drops when it returns") -} diff --git a/pkg/vmcp/server/session_telemetry.go b/pkg/vmcp/server/session_telemetry.go deleted file mode 100644 index 18ad046d1d..0000000000 --- a/pkg/vmcp/server/session_telemetry.go +++ /dev/null @@ -1,65 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "fmt" - - "go.opentelemetry.io/otel/metric" -) - -// sessionTelemetryInstrumentationName matches the vMCP backend telemetry meter. -const sessionTelemetryInstrumentationName = "github.com/stacklok/toolhive/pkg/vmcp" - -// monitorSessionCardinality registers observable gauges for the vMCP session -// layers, so operators can tell transport-session churn apart from live SDK -// session growth and from a post-initialize registration backlog. -func monitorSessionCardinality(meterProvider metric.MeterProvider, srv *Server) error { - meter := meterProvider.Meter(sessionTelemetryInstrumentationName) - - activeTransportSessions, err := meter.Int64ObservableGauge( - "toolhive_vmcp_active_transport_sessions", - metric.WithDescription("Number of active transport-layer MCP sessions tracked by the vMCP"), - ) - if err != nil { - return fmt.Errorf("failed to create active transport sessions gauge: %w", err) - } - - activeSDKSessions, err := meter.Int64ObservableGauge( - "toolhive_vmcp_active_sdk_sessions", - metric.WithDescription("Number of live SDK client sessions currently registered on this vMCP replica"), - ) - if err != nil { - return fmt.Errorf("failed to create active SDK sessions gauge: %w", err) - } - - inFlightRegistrations, err := meter.Int64ObservableGauge( - "toolhive_vmcp_inflight_session_registrations", - metric.WithDescription("Number of session registrations currently executing the post-initialize capability injection hook"), - ) - if err != nil { - return fmt.Errorf("failed to create in-flight session registrations gauge: %w", err) - } - - _, err = meter.RegisterCallback( - func(_ context.Context, observer metric.Observer) error { - observer.ObserveInt64(activeTransportSessions, int64(srv.sessionManager.Count())) - observer.ObserveInt64(activeSDKSessions, srv.activeClientSessionCount()) - observer.ObserveInt64(inFlightRegistrations, srv.inFlightRegistrations.Load()) - return nil - }, - activeTransportSessions, - activeSDKSessions, - inFlightRegistrations, - ) - if err != nil { - return fmt.Errorf("failed to register session cardinality callback: %w", err) - } - return nil -} - -func (s *Server) activeClientSessionCount() int64 { - return s.activeClientSessions.count() -} diff --git a/pkg/vmcp/server/session_telemetry_test.go b/pkg/vmcp/server/session_telemetry_test.go deleted file mode 100644 index d1dfe233c4..0000000000 --- a/pkg/vmcp/server/session_telemetry_test.go +++ /dev/null @@ -1,55 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "context" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - sdkmetric "go.opentelemetry.io/otel/sdk/metric" - "go.opentelemetry.io/otel/sdk/metric/metricdata" - - transportsession "github.com/stacklok/toolhive/pkg/transport/session" -) - -func TestMonitorSessionCardinality(t *testing.T) { - t.Parallel() - - reader := sdkmetric.NewManualReader() - meterProvider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) - - sessionManager := transportsession.NewManager(time.Minute, transportsession.NewStreamableSession) - t.Cleanup(func() { _ = sessionManager.Stop() }) - require.NoError(t, sessionManager.AddWithID("11111111-1111-1111-1111-111111111111")) - require.NoError(t, sessionManager.AddWithID("22222222-2222-2222-2222-222222222222")) - - srv := &Server{sessionManager: sessionManager} - for _, id := range []string{"sdk-session-1", "sdk-session-2", "sdk-session-3"} { - srv.activeClientSessions.begin(id) - srv.activeClientSessions.commit(&fakeSDKSession{id: id}) - } - srv.inFlightRegistrations.Add(1) - - require.NoError(t, monitorSessionCardinality(meterProvider, srv)) - - var rm metricdata.ResourceMetrics - require.NoError(t, reader.Collect(context.Background(), &rm)) - - gauges := map[string]int64{} - for _, sm := range rm.ScopeMetrics { - for _, m := range sm.Metrics { - if g, ok := m.Data.(metricdata.Gauge[int64]); ok && len(g.DataPoints) > 0 { - gauges[m.Name] = g.DataPoints[len(g.DataPoints)-1].Value - } - } - } - assert.Equal(t, map[string]int64{ - "toolhive_vmcp_active_transport_sessions": 2, - "toolhive_vmcp_active_sdk_sessions": 3, - "toolhive_vmcp_inflight_session_registrations": 1, - }, gauges) -} diff --git a/pkg/vmcp/server/sessionmanager/factory.go b/pkg/vmcp/server/sessionmanager/factory.go index 1f88aae7c7..f1eeb9cd8b 100644 --- a/pkg/vmcp/server/sessionmanager/factory.go +++ b/pkg/vmcp/server/sessionmanager/factory.go @@ -11,13 +11,13 @@ import ( "log/slog" "time" - "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/vmcp" @@ -69,24 +69,36 @@ type FactoryConfig struct { // AdvertiseFromCore signals that the advertised capability set is sourced from // the core (the Serve path), not from this factory's per-session aggregation. // - // It is the single switch that selects WHICH layer indexes the shared FTS5 store, - // so the two can never both do it (the AC6 double-index): - // - true: New does NOT install the optimizer decorator on the session factory, - // and exposes the resolved factory via Manager.OptimizerFactory so the - // Serve layer builds a per-session optimizer over the core's tools. - // - false: New installs the optimizer decorator (legacy behavior) and - // Manager.OptimizerFactory returns nil, so a Serve composition root that - // enables the optimizer but forgets this flag gets no Serve-layer - // optimizer rather than a second, divergent upsert into the store. - // Either way New resolves the optimizer factory and owns its store/cleanup. Has no - // effect when the optimizer is disabled. The legacy server.New path leaves this - // false, so its optimizer decorator is unchanged. + // It is required whenever an optimizer is configured: + // - true: New exposes the resolved factory via Manager.OptimizerFactory so + // the Serve layer builds a per-session optimizer over the core's + // tools — the single writer of the shared FTS5 store (the AC6 + // no-double-index guarantee). + // - false: fine without an optimizer; with one, New rejects the config at + // construction, because the Serve layer discards the decorator's + // per-session tools (the optimizer would index the store yet serve + // nobody) and the Modern capability gate would fail open. + // New resolves the optimizer factory and owns its store/cleanup. server.New + // sets this unconditionally (server.go), so every in-tree composition + // advertises from the core; the flag exists for direct-Serve embedders. The + // decorator branch the false case used to select is now unreachable — its + // deletion is tracked in #6103. AdvertiseFromCore bool - // HealthStatusProvider filters the backends a new session connects to: backends - // it reports as unhealthy, unknown or unauthenticated are skipped, the same rule - // the core applies to the advertised set. Nil falls back to the registry status. - HealthStatusProvider health.StatusProvider + // BackendHealth gates which backends a NEW session attempts to connect to. + // Restored sessions are never filtered (see Manager.listAllBackends). + // + // Optional: nil disables health gating and every backend is attempted, which + // is both the pre-#5861 behaviour and the correct fallback when health + // monitoring is switched off. See health.ShouldOpenSession for which statuses + // are skipped. + // + // Because nil is indistinguishable from "monitoring disabled", an embedder that + // calls Serve with a hand-built FactoryConfig and omits this field gets no + // health gating and no warning. server.New always populates it from the core's + // monitor, so in-tree compositions are covered; direct-Serve callers that want + // the #5861 fix must pass it explicitly. + BackendHealth health.StatusProvider // IdleTTL closes a cached MultiSession (and its backend connections) when no // request on this replica has used it for this long. It bounds what a replica @@ -160,6 +172,9 @@ func buildDecoratingFactory( // over the core's advertised set, so the factory's optimizer decorator is skipped // to avoid double-indexing the shared store (see FactoryConfig.AdvertiseFromCore). // Composite tools and their telemetry are owned by the core, not the factory. + // This branch is unreachable: New rejects an optimizer without AdvertiseFromCore, + // so optimizerFactory is nil whenever the flag is false. Deleting the decorator + // path is tracked in #6103. if optimizerFactory != nil && !cfg.AdvertiseFromCore { decorators = append(decorators, optimizerDecoratorFn(optimizerFactory, terminateSession)) } diff --git a/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go b/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go index 102de1fd1e..23fc831a67 100644 --- a/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go +++ b/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go @@ -12,11 +12,11 @@ import ( "time" "github.com/alicebob/miniredis/v2" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" tcredis "github.com/stacklok/toolhive-core/redis" "github.com/stacklok/toolhive/pkg/auth" transportsession "github.com/stacklok/toolhive/pkg/transport/session" @@ -42,9 +42,19 @@ func newUnauthenticatedAuthRegistry(t *testing.T) vmcpauth.OutgoingAuthRegistry return reg } -// newSharedRedisStorage creates a RedisSessionDataStorage pointing at mr. -// The storage is closed via t.Cleanup. +// newSharedRedisStorage creates a RedisSessionDataStorage pointing at mr with +// a long (1h) TTL, suitable for tests that are not exercising TTL expiry +// itself. The storage is closed via t.Cleanup. func newSharedRedisStorage(t *testing.T, mr *miniredis.Miniredis) transportsession.DataStorage { + t.Helper() + return newSharedRedisStorageWithTTL(t, mr, time.Hour) +} + +// newSharedRedisStorageWithTTL is like newSharedRedisStorage but lets the +// caller control the sliding-window TTL, so tests can pin TTL-refresh and +// TTL-expiry behaviour with a short duration combined with mr.FastForward. +// The storage is closed via t.Cleanup. +func newSharedRedisStorageWithTTL(t *testing.T, mr *miniredis.Miniredis, ttl time.Duration) transportsession.DataStorage { t.Helper() storage, err := transportsession.NewRedisSessionDataStorage( context.Background(), @@ -52,7 +62,7 @@ func newSharedRedisStorage(t *testing.T, mr *miniredis.Miniredis) transportsessi Addr: mr.Addr(), }, "test:vmcp:session:", - time.Hour, + ttl, ) require.NoError(t, err) t.Cleanup(func() { _ = storage.Close() }) @@ -89,7 +99,7 @@ func createSession(t *testing.T, sm *Manager, identity *auth.Identity) string { if identity != nil { ctx = auth.WithIdentity(ctx, identity) } - _, err := sm.CreateSession(ctx, sessionID) + _, err := sm.CreateSession(ctx, sessionID, nil) require.NoError(t, err) return sessionID } @@ -388,3 +398,55 @@ func TestHorizontalScaling_InMemoryOnlyMode(t *testing.T) { require.NotNil(t, sess) assert.NotEmpty(t, sess.Tools(), "session on pod A must have tools") } + +// --------------------------------------------------------------------------- +// Regression: cross-pod Validate() must observe a shared-store termination +// --------------------------------------------------------------------------- + +// TestRegression_Validate_TerminatedInStore_ReturnsTerminated pins the real +// U5 fix (toolhive-core v0.0.32, #5742): a session that Manager A ("pod A") +// still holds live in its node-local cache, but which was terminated in the +// SHARED Redis store by a second Manager ("pod B"), must be reported as +// terminated by Manager A's very next Validate() call -- not treated as an +// error (which the streamable transport would map to a retryable 503) and not +// silently reported as still-live. Manager.Validate's documented contract +// collapses "explicitly terminated" and "absent from storage" into the same +// (isTerminated=true, nil) signal; this test exercises the termination +// specifically via a DIFFERENT Manager instance terminating the SAME shared +// store, not via Manager A's own Terminate (which the existing +// "TestSessionManager_Terminate" suite already covers and which additionally +// touches Manager A's local bookkeeping directly). +func TestRegression_Validate_TerminatedInStore_ReturnsTerminated(t *testing.T) { + t.Parallel() + + mr := miniredis.RunT(t) + storage := newSharedRedisStorage(t, mr) + backend := startMCPBackend(t, "backend-alpha", "echo") + + // Manager A ("pod A"): create a full (Phase 2) session. + smA := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend}) + sessionID := createSession(t, smA, nil) + + // Sanity: the session must currently validate as live, or the + // post-termination assertion below would pass vacuously. + isTerminated, err := smA.Validate(sessionID) + require.NoError(t, err) + require.False(t, isTerminated, "a freshly created session must validate as live") + + // Manager B ("pod B"): a second Manager over the SAME shared Redis store + // terminates the session out-of-band. This mutates only the shared store -- + // it does NOT go through Manager A's in-process bookkeeping -- exactly like + // a DELETE served by a different replica. + smB := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend}) + isNotAllowed, err := smB.Terminate(sessionID) + require.NoError(t, err) + require.False(t, isNotAllowed) + + // Manager A must observe the cross-pod termination on the very next + // Validate() call against the shared store. + isTerminated, err = smA.Validate(sessionID) + require.NoError(t, err) + assert.True(t, isTerminated, + "Manager A must report the session as terminated after Manager B terminates it "+ + "in the shared store, even though Manager A never terminated it locally") +} diff --git a/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go b/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go index 25bf9c1c8d..d83834fa6d 100644 --- a/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go +++ b/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go @@ -7,20 +7,22 @@ import ( "context" "testing" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" ) -// TestOptimizerFactoryGatedOnAdvertiseFromCore locks in the AC6 double-index -// guarantee: New surfaces the optimizer factory via OptimizerFactory() ONLY when -// AdvertiseFromCore is true. That makes the session-factory decorator (installed iff -// !AdvertiseFromCore) and the Serve-layer getter mutually exclusive store writers, so a -// Serve composition root that enables the optimizer but forgets the flag gets a nil -// factory (no Serve-layer optimizer) instead of a silent second upsert into the store. +// TestOptimizerFactoryGatedOnAdvertiseFromCore locks in the constructor guard +// behind the AC6 single-writer guarantee and the Modern capability gate's +// enablement signal: New accepts an optimizer only together with +// AdvertiseFromCore, so the Serve layer (via OptimizerFactory()) is the sole +// writer of the shared FTS5 store and a non-nil OptimizerFactory() faithfully +// means "optimizer enabled" (pkg/vmcp/server/modern_gate.go depends on that). +// Optimizer-without-flag is rejected at construction instead of silently +// producing an optimizer that indexes the store but serves nobody. func TestOptimizerFactoryGatedOnAdvertiseFromCore(t *testing.T) { t.Parallel() @@ -31,10 +33,9 @@ func TestOptimizerFactoryGatedOnAdvertiseFromCore(t *testing.T) { tests := []struct { name string advertiseFromCore bool - wantSurfaced bool }{ - {"surfaced to Serve when advertising from core", true, true}, - {"not surfaced on the legacy path (decorator owns it)", false, false}, + {"surfaced to Serve when advertising from core", true}, + {"rejected at construction without AdvertiseFromCore", false}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { @@ -47,16 +48,16 @@ func TestOptimizerFactoryGatedOnAdvertiseFromCore(t *testing.T) { OptimizerFactory: optFactory, AdvertiseFromCore: tc.advertiseFromCore, }, newFakeRegistry()) - require.NoError(t, err) - t.Cleanup(func() { _ = cleanup(context.Background()) }) - if tc.wantSurfaced { - assert.NotNil(t, sm.OptimizerFactory(), - "the factory must be surfaced to the Serve layer when AdvertiseFromCore is set") - } else { - assert.Nil(t, sm.OptimizerFactory(), - "the factory must NOT be surfaced when AdvertiseFromCore is false (decorator owns the store)") + if !tc.advertiseFromCore { + require.ErrorContains(t, err, "AdvertiseFromCore", + "New must reject an optimizer without AdvertiseFromCore at construction") + return } + require.NoError(t, err) + t.Cleanup(func() { _ = cleanup(context.Background()) }) + assert.NotNil(t, sm.OptimizerFactory(), + "the factory must be surfaced to the Serve layer when AdvertiseFromCore is set") }) } } diff --git a/pkg/vmcp/server/sessionmanager/session_cache_close_test.go b/pkg/vmcp/server/sessionmanager/session_cache_close_test.go index bb8ddbd86a..217794081a 100644 --- a/pkg/vmcp/server/sessionmanager/session_cache_close_test.go +++ b/pkg/vmcp/server/sessionmanager/session_cache_close_test.go @@ -34,8 +34,8 @@ func newBoundSessionManager( factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { sess := sessionmocks.NewMockMultiSession(ctrl) sess.EXPECT().ID().Return(id).AnyTimes() sess.EXPECT().GetMetadata().Return(meta).AnyTimes() @@ -58,7 +58,7 @@ func TestTerminate_ClosesCachedSessionWhenStorageIsAlreadyGone(t *testing.T) { sm, storage, closed := newBoundSessionManager(t, 0) sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) // Another replica served the DELETE (or the storage TTL expired). @@ -74,7 +74,7 @@ func TestIdleSweep_ClosesIdleSessionWithoutTouchingStorage(t *testing.T) { sm, storage, closed := newBoundSessionManager(t, 50*time.Millisecond) sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) require.Eventually(t, closed.Load, 5*time.Second, 10*time.Millisecond, diff --git a/pkg/vmcp/server/sessionmanager/session_manager.go b/pkg/vmcp/server/sessionmanager/session_manager.go index 6cd0c0c4d6..fa633a0343 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager.go +++ b/pkg/vmcp/server/sessionmanager/session_manager.go @@ -22,8 +22,8 @@ import ( "time" "github.com/google/uuid" - mcpserver "github.com/mark3labs/mcp-go/server" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/cache" transportsession "github.com/stacklok/toolhive/pkg/transport/session" @@ -45,7 +45,7 @@ const ( ) // Manager bridges the domain session lifecycle (MultiSession / MultiSessionFactory) -// to the mark3labs SDK's SessionIdManager interface. +// to the mcpcompat SDK's SessionIdManager interface. // // It implements a two-phase session-creation pattern: // @@ -79,8 +79,10 @@ type Manager struct { factory vmcpsession.MultiSessionFactory backendReg vmcp.BackendRegistry - // healthStatus filters the backends CreateSession connects to (may be nil). - healthStatus health.StatusProvider + // backendHealth gates which backends a new session connects to, or nil when + // health monitoring is disabled (every backend is attempted). Read-only here; + // the core owns the monitor's lifecycle. See shouldOpenSession. + backendHealth health.StatusProvider // sessions is a node-local cache of live MultiSession objects, separate // from storage because MultiSession contains un-serialisable runtime state @@ -100,6 +102,12 @@ type Manager struct { // registry. It builds the decorating session factory from cfg, wiring the // optimizer and composite tool layers internally. // +// An optimizer (FactoryConfig.OptimizerFactory or OptimizerConfig) requires +// FactoryConfig.AdvertiseFromCore; New rejects the combination otherwise. The +// guard makes a non-nil OptimizerFactory() a faithful "optimizer enabled" +// signal, which the Modern capability gate in pkg/vmcp/server/modern_gate.go +// depends on. +// // The returned cleanup function releases any resources allocated during // construction (e.g. the optimizer's SQLite store). Callers must invoke it // on shutdown. If no cleanup is needed, a no-op function is returned. @@ -114,6 +122,10 @@ func New( if cfg.CacheCapacity < 0 { return nil, nil, fmt.Errorf("sessionmanager.New: CacheCapacity must be >= 0 (got %d)", cfg.CacheCapacity) } + if (cfg.OptimizerFactory != nil || cfg.OptimizerConfig != nil) && !cfg.AdvertiseFromCore { + return nil, nil, fmt.Errorf("sessionmanager.New: the optimizer requires " + + "FactoryConfig.AdvertiseFromCore (the Serve layer builds it over the core's tools)") + } capacity := cfg.CacheCapacity if capacity == 0 { capacity = defaultCacheCapacity @@ -127,19 +139,18 @@ func New( // Build the Manager first so we can reference sm.Terminate and sm.sessions // directly in closures, eliminating the forward-reference variable pattern. sm := &Manager{ - storage: storage, - backendReg: backendRegistry, - healthStatus: cfg.HealthStatusProvider, - } - - // Surface the resolved optimizer factory to the Serve path ONLY when - // AdvertiseFromCore is set. That makes the two store writers mutually exclusive: - // the session-factory decorator runs iff !AdvertiseFromCore (buildDecoratingFactory - // below), and OptimizerFactory() returns a non-nil factory iff AdvertiseFromCore — - // so a Serve composition root that enables the optimizer but forgets the flag gets a - // nil factory (no Serve-layer optimizer) rather than a silent double-index of the - // shared FTS5 store. The legacy server.New path leaves AdvertiseFromCore false and - // never calls OptimizerFactory(), so its decorator is unaffected. + storage: storage, + backendReg: backendRegistry, + backendHealth: cfg.BackendHealth, + } + + // Surface the resolved optimizer factory to the Serve path. The constructor + // guard above rejects an optimizer without AdvertiseFromCore, so the shared + // FTS5 store has exactly one writer (the Serve layer) and OptimizerFactory() + // is non-nil exactly when the optimizer is enabled. server.New sets + // AdvertiseFromCore unconditionally (server.go), so every in-tree composition + // takes this branch; the flag exists for direct-Serve embedders, which New + // rejects when they configure an optimizer without setting it. if cfg.AdvertiseFromCore { sm.optimizerFactory = optimizerFactory } @@ -191,16 +202,17 @@ func (sm *Manager) startIdleSweep(idleTTL time.Duration) func() { return sync.OnceFunc(func() { close(stop) }) } -// OptimizerFactory returns the resolved (telemetry-wrapped) optimizer factory, or -// nil when the optimizer is disabled OR FactoryConfig.AdvertiseFromCore is false. +// OptimizerFactory returns the resolved (telemetry-wrapped) optimizer factory, +// or nil exactly when the optimizer is disabled: New rejects an optimizer +// without FactoryConfig.AdvertiseFromCore, so nil-ness here is a faithful +// "optimizer enabled" signal (the Modern capability gate in +// pkg/vmcp/server/modern_gate.go depends on it). // -// It is consumed by the Serve path (FactoryConfig.AdvertiseFromCore), which builds -// a per-session optimizer over the core's advertised tool set rather than via the -// session decorator. Gating on AdvertiseFromCore makes the decorator and this getter -// mutually exclusive store writers, so the shared FTS5 store can never be double-indexed -// (see New). The optimizer's shared store and its cleanup remain owned by this Manager -// (the cleanup function returned from New). On the legacy server.New path the factory is -// applied internally via the session decorator and this getter is unused. +// It is consumed by the Serve path, which builds a per-session optimizer over +// the core's advertised tool set. The optimizer's shared store and its cleanup +// remain owned by this Manager (the cleanup function returned from New). +// server.New sets AdvertiseFromCore unconditionally, so this getter is live on +// every composition path. func (m *Manager) OptimizerFactory() func(context.Context, []mcpserver.ServerTool) (optimizer.Optimizer, error) { return m.optimizerFactory } @@ -309,10 +321,15 @@ func (sm *Manager) Generate() string { // 4. Persists session metadata to storage and caches the live MultiSession // in the node-local map. // +// sink, when non-nil, is threaded to MultiSessionFactory.MakeSessionWithID so +// every backend connector opened for this session can report an asynchronous +// backend notification (#5748); pass nil to disable that consumption. +// // The returned MultiSession can be retrieved later via GetMultiSession(). func (sm *Manager) CreateSession( ctx context.Context, sessionID string, + sink vmcpsession.ListChangedSink, ) (vmcpsession.MultiSession, error) { if sessionID == "" { return nil, fmt.Errorf("Manager.CreateSession: session ID must not be empty") @@ -346,12 +363,12 @@ func (sm *Manager) CreateSession( // Resolve the caller identity (may be nil for anonymous access). identity, _ := auth.IdentityFromContext(ctx) - // List the registry's backends that are currently eligible for exposure, so a - // new session does not open connections to backends the core does not advertise. - backends := toBackendPointers(health.FilterBackendsForExposure(sm.backendReg.List(ctx), sm.healthStatus)) + // List the backends worth opening connections to, skipping any the health + // monitor already knows are bad (#5861). + backends := sm.listBackendsForNewSession(ctx) // Build the fully-formed MultiSession using the SDK-assigned session ID. - sess, err := sm.factory.MakeSessionWithID(ctx, sessionID, identity, backends) + sess, err := sm.factory.MakeSessionWithID(ctx, sessionID, identity, backends, sink) if err != nil { sm.cleanupFailedPlaceholder(sessionID, placeholder) return nil, fmt.Errorf("Manager.CreateSession: failed to create multi-session: %w", err) @@ -450,10 +467,25 @@ func (sm *Manager) cleanupFailedPlaceholder(sessionID string, metadata map[strin // Validate implements the SDK's SessionIdManager.Validate(). // -// Returns (isTerminated=true, nil) for explicitly terminated sessions. -// Returns (false, error) for unknown sessions — per the SDK interface contract, -// a lookup failure is signalled via err, not via isTerminated. +// Returns (isTerminated=true, nil) for a session that is definitively gone: +// explicitly terminated (placeholder marked terminated=true) OR absent from +// storage (a full MultiSession is deleted on Terminate, TTL-expired, or never +// existed). All of these must reject the request as a hard termination so the +// client re-initializes rather than retrying a dead session. +// +// Returns (false, error) only for a genuine, transient storage error (e.g. the +// backing store is unreachable) — the caller should treat this as retryable and +// NOT drop session state. +// // Returns (false, nil) for valid, active sessions. +// +// This distinction is load-bearing: the streamable transport maps +// (isTerminated=true) to HTTP 404 (-> client ErrSessionTerminated -> re-init) +// and a non-terminated error to HTTP 503 (retryable). Reporting a genuinely-gone +// session as an error would surface as 503 and make a client whose session was +// terminated on another replica retry the dead session forever. Terminate still +// DELETES the key (unchanged), so the resurrection-race guarantee is preserved; +// only the way an absent key is reported here changes. func (sm *Manager) Validate(sessionID string) (isTerminated bool, err error) { if sessionID == "" { return false, fmt.Errorf("Manager.Validate: empty session ID") @@ -464,8 +496,12 @@ func (sm *Manager) Validate(sessionID string) (isTerminated bool, err error) { metadata, err := sm.storage.Load(ctx, sessionID) if errors.Is(err, transportsession.ErrSessionNotFound) { - slog.Debug("Manager.Validate: session not found", "session_id", sessionID) - return false, fmt.Errorf("session not found") + // The session is gone (terminated + deleted, TTL-expired, or never + // existed). Report it as terminated so the transport answers 404 and the + // client re-initializes, rather than as an error (which the transport + // would treat as a transient 503 and the client would retry indefinitely). + slog.Debug("Manager.Validate: session not found; reporting as terminated", "session_id", sessionID) + return true, nil } if err != nil { return false, fmt.Errorf("Manager.Validate: storage error for session %q: %w", sessionID, err) @@ -488,16 +524,17 @@ func (sm *Manager) Validate(sessionID string) (isTerminated bool, err error) { // - MultiSession (Phase 2): the storage key is deleted, then the node-local // cache entry is removed and onEvict closes its backend connections. Caches // on other replicas self-heal on their next Get (checkSession detects -// ErrSessionNotFound) or through the idle sweep. After deletion -// Validate() returns (false, error) — the same response as "never existed". +// ErrSessionNotFound) or through the idle sweep. After deletion Validate() +// reports the absent key as (isTerminated=true, nil), so the next request — +// on this or any other replica — is rejected with a definitive 404 and the +// client re-initializes instead of retrying the dead session. // // - Placeholder (Phase 1): the session is marked terminated=true and left // for TTL cleanup. This prevents CreateSession() from opening backend // connections for an already-terminated session (see fast-fail check in // CreateSession). The terminated flag also lets Validate() return // (isTerminated=true, nil) during the window between termination and TTL -// expiry, allowing the SDK to distinguish "actively terminated" from -// "never existed". +// expiry — the same terminated response the deleted case now gives. // // Returns (isNotAllowed=false, nil) on success; client termination is always permitted. func (sm *Manager) Terminate(sessionID string) (isNotAllowed bool, err error) { @@ -840,7 +877,16 @@ func (sm *Manager) DecorateSession(sessionID string, fn func(sessiontypes.MultiS return nil } -// listAllBackends returns all backends from the registry as a pointer slice. +// listAllBackends returns every backend in the registry as a pointer slice, +// with no health filtering. +// +// Used by the restore path, which must offer the factory the full registry: it +// intersects this list with the session's stored backend IDs +// (session.RestoreSession), so filtering here would silently DROP a backend the +// session already held rather than merely defer connecting to one. Because the +// routing table is rebuilt from whatever reconnects, a dropped backend stays gone +// for the rest of that session's life even after it recovers. New sessions use +// listBackendsForNewSession instead. func (sm *Manager) listAllBackends(ctx context.Context) []*vmcp.Backend { return toBackendPointers(sm.backendReg.List(ctx)) } @@ -852,3 +898,48 @@ func toBackendPointers(raw []vmcp.Backend) []*vmcp.Backend { } return backends } + +// listBackendsForNewSession returns the backends a NEW session should attempt to +// connect to, skipping any the health monitor already knows are bad +// (see shouldOpenSession). +// +// Deliberately not used on the restore path — see listAllBackends. +func (sm *Manager) listBackendsForNewSession(ctx context.Context) []*vmcp.Backend { + raw := sm.backendReg.List(ctx) + backends := make([]*vmcp.Backend, 0, len(raw)) + skipped := 0 + for i := range raw { + if !sm.shouldOpenSession(&raw[i]) { + skipped++ + continue + } + backends = append(backends, &raw[i]) + } + if skipped > 0 { + slog.Debug("skipping backends for session establishment due to health status", + "skipped", skipped, + "attempted", len(backends)) + } + return backends +} + +// shouldOpenSession reports whether a new session should attempt to connect to +// backend, consulting the health monitor when one is wired. +// +// A nil provider means health monitoring is disabled: every backend is +// attempted, preserving the pre-#5861 behaviour. A backend the monitor does not +// track yet falls back to its registry status, which is how a status the registry +// knows but the monitor has not caught up with (e.g. a k8s workload already +// reported unhealthy) is still honoured. Either way only confirmed-bad statuses +// are skipped — see health.ShouldOpenSession for why "not yet classified" must +// fail open rather than closed. +func (sm *Manager) shouldOpenSession(backend *vmcp.Backend) bool { + if sm.backendHealth == nil { + return true + } + status, ok := sm.backendHealth.QueryBackendStatus(backend.ID) + if !ok { + status = backend.HealthStatus + } + return health.ShouldOpenSession(status) +} diff --git a/pkg/vmcp/server/sessionmanager/session_manager_test.go b/pkg/vmcp/server/sessionmanager/session_manager_test.go index 23f1122a02..c5bf65690b 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager_test.go +++ b/pkg/vmcp/server/sessionmanager/session_manager_test.go @@ -65,7 +65,7 @@ func newMockFactory(t *testing.T, ctrl *gomock.Controller, sess vmcpsession.Mult t.Helper() factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil).AnyTimes() return factory } @@ -75,7 +75,7 @@ func newMockFactoryWithError(t *testing.T, ctrl *gomock.Controller, err error) * t.Helper() factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(nil, err).AnyTimes() return factory } @@ -193,15 +193,6 @@ func (r *fakeBackendRegistry) Count() int { return len(r.backends) } -// fakeHealthStatus is a health.StatusProvider backed by a map; backends absent -// from the map are reported as not monitored. -type fakeHealthStatus map[string]vmcp.BackendHealthStatus - -func (f fakeHealthStatus) QueryBackendStatus(id string) (vmcp.BackendHealthStatus, bool) { - status, ok := f[id] - return status, ok -} - // newTestSessionDataStorage creates a LocalSessionDataStorage with a long TTL. // The storage is closed via t.Cleanup. func newTestSessionDataStorage(t *testing.T) transportsession.DataStorage { @@ -309,8 +300,8 @@ func TestSessionManager_CreateSession(t *testing.T) { factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) var createdSess *sessionmocks.MockMultiSession factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { createdSess = newMockSession(t, ctrl, id, tools) return createdSess, nil }).AnyTimes() @@ -323,7 +314,7 @@ func TestSessionManager_CreateSession(t *testing.T) { require.NotEmpty(t, sessionID) // Upgrade to full MultiSession. - multiSess, err := sm.CreateSession(context.Background(), sessionID) + multiSess, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) require.NotNil(t, multiSess) assert.Equal(t, sessionID, multiSess.ID()) @@ -333,37 +324,6 @@ func TestSessionManager_CreateSession(t *testing.T) { assert.NoError(t, loadErr, "session should still exist in storage after CreateSession") }) - t.Run("connects only to backends eligible for exposure", func(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - var gotBackends []string - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, backends []*vmcp.Backend) (vmcpsession.MultiSession, error) { - for _, b := range backends { - gotBackends = append(gotBackends, b.ID) - } - return newMockSession(t, ctrl, id, nil), nil - }).Times(1) - - registry := &fakeBackendRegistry{backends: []vmcp.Backend{ - {ID: "healthy"}, {ID: "unhealthy"}, {ID: "registry-unknown", HealthStatus: vmcp.BackendUnknown}, - }} - health := fakeHealthStatus{"healthy": vmcp.BackendHealthy, "unhealthy": vmcp.BackendUnhealthy} - sm, cleanup, err := New(newTestSessionDataStorage(t), - &FactoryConfig{Base: factory, HealthStatusProvider: health}, registry) - require.NoError(t, err) - t.Cleanup(func() { _ = cleanup(context.Background()) }) - - sessionID := sm.Generate() - require.NotEmpty(t, sessionID) - _, err = sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - assert.Equal(t, []string{"healthy"}, gotBackends) - }) - t.Run("returns error for empty session ID", func(t *testing.T) { t.Parallel() @@ -373,7 +333,7 @@ func TestSessionManager_CreateSession(t *testing.T) { registry := newFakeRegistry() sm, _ := newTestSessionManager(t, factory, registry) - _, err := sm.CreateSession(context.Background(), "") + _, err := sm.CreateSession(context.Background(), "", nil) require.Error(t, err) assert.Contains(t, err.Error(), "session ID must not be empty") }) @@ -392,7 +352,7 @@ func TestSessionManager_CreateSession(t *testing.T) { sessionID := sm.Generate() require.NotEmpty(t, sessionID) - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.Error(t, err) assert.ErrorContains(t, err, "failed to create multi-session") }) @@ -407,8 +367,8 @@ func TestSessionManager_CreateSession(t *testing.T) { factoryCalled := false factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { factoryCalled = true sess := newMockSession(t, ctrl, id, tools) return sess, nil @@ -424,7 +384,7 @@ func TestSessionManager_CreateSession(t *testing.T) { require.NoError(t, storage.Delete(context.Background(), sessionID)) // CreateSession must fail fast before opening any backend connections. - _, createErr := sm.CreateSession(context.Background(), sessionID) + _, createErr := sm.CreateSession(context.Background(), sessionID, nil) require.Error(t, createErr) assert.ErrorContains(t, createErr, "not found") @@ -442,8 +402,8 @@ func TestSessionManager_CreateSession(t *testing.T) { factoryCalled := false factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { factoryCalled = true sess := newMockSession(t, ctrl, id, tools) return sess, nil @@ -462,7 +422,7 @@ func TestSessionManager_CreateSession(t *testing.T) { // CreateSession must fail fast (terminated=true) before opening any // backend connections. - _, createErr := sm.CreateSession(context.Background(), sessionID) + _, createErr := sm.CreateSession(context.Background(), sessionID, nil) require.Error(t, createErr) assert.ErrorContains(t, createErr, "was terminated") @@ -480,8 +440,8 @@ func TestSessionManager_CreateSession(t *testing.T) { var createdSess *sessionmocks.MockMultiSession factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { // Sleep to simulate slow backend initialization, creating a window // where the client can terminate the session after the first check passes. time.Sleep(50 * time.Millisecond) @@ -502,7 +462,7 @@ func TestSessionManager_CreateSession(t *testing.T) { // check and then sleep during MakeSessionWithID. errChan := make(chan error, 1) go func() { - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) errChan <- err }() @@ -544,7 +504,7 @@ func TestSessionManager_Validate(t *testing.T) { assert.Contains(t, err.Error(), "empty session ID") }) - t.Run("returns error for unknown session", func(t *testing.T) { + t.Run("reports unknown session as terminated", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) @@ -553,10 +513,15 @@ func TestSessionManager_Validate(t *testing.T) { registry := newFakeRegistry() sm, _ := newTestSessionManager(t, factory, registry) + // A session that is absent from storage (deleted on Terminate, TTL-expired, + // or never existed) must be reported as terminated (isTerminated=true, nil), + // NOT as an error. The transport maps terminated -> 404 so the client + // re-initializes; reporting an error would surface as a retryable 503 and a + // client whose session was terminated on another replica would retry the + // dead session forever (regression guarded here). isTerminated, err := sm.Validate("non-existent-id") - require.Error(t, err) - assert.False(t, isTerminated) - assert.Contains(t, err.Error(), "session not found") + require.NoError(t, err) + assert.True(t, isTerminated, "an absent session must report as terminated, not as a transient error") }) t.Run("returns false for active session", func(t *testing.T) { @@ -653,8 +618,8 @@ func TestSessionManager_Terminate(t *testing.T) { var closed atomic.Bool factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { createdSess = sessionmocks.NewMockMultiSession(ctrl) createdSess.EXPECT().ID().Return(id).AnyTimes() createdSess.EXPECT().GetMetadata().Return(bindingMeta).AnyTimes() @@ -682,7 +647,7 @@ func TestSessionManager_Terminate(t *testing.T) { sessionID := sm.Generate() require.NotEmpty(t, sessionID) - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) require.NotNil(t, createdSess) @@ -706,8 +671,8 @@ func TestSessionManager_Terminate(t *testing.T) { ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { sess := newMockSession(t, ctrl, id, nil) // Close is called by onEvict when Terminate removes the cache entry. sess.EXPECT().Close().Return(nil).AnyTimes() @@ -720,7 +685,7 @@ func TestSessionManager_Terminate(t *testing.T) { sessionID := sm.Generate() require.NotEmpty(t, sessionID) - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) // Seed MetadataKeyIdentityBinding into storage so Terminate recognises this @@ -887,8 +852,8 @@ func TestSessionManager_GetMultiSession(t *testing.T) { ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { sess := newMockSession(t, ctrl, id, tools) return sess, nil }).Times(1) @@ -899,7 +864,7 @@ func TestSessionManager_GetMultiSession(t *testing.T) { sessionID := sm.Generate() require.NotEmpty(t, sessionID) - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) multiSess, ok := sm.GetMultiSession(t.Context(), sessionID) @@ -942,7 +907,7 @@ func TestSessionManager_GetMultiSession(t *testing.T) { ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) // MakeSessionWithID is only for Phase 2; unused in the restore path. - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Times(0) sessionID := "restore-zero-backend-session" @@ -982,7 +947,7 @@ func TestSessionManager_GetMultiSession(t *testing.T) { // RestoreSession. This test documents that backward-compat behaviour. ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Times(0) sessionID := "restore-legacy-session" @@ -1018,7 +983,7 @@ func TestSessionManager_GetMultiSession(t *testing.T) { // that stale per-backend session IDs do not persist indefinitely in storage. ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Times(0) sessionID := "restore-metadata-persist-session" @@ -1069,7 +1034,7 @@ func TestSessionManager_GetMultiSession(t *testing.T) { // restored session. ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Times(0) sessionID := "restore-concurrent-delete-session" @@ -1115,7 +1080,7 @@ func TestSessionManager_GetMultiSession(t *testing.T) { // metadata drift on the next liveness check and evict if necessary. ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Times(0) sessionID := "restore-update-error-session" @@ -1221,8 +1186,8 @@ func TestSessionManager_DecorateSession(t *testing.T) { ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { return newMockSession(t, ctrl, id, tools), nil }).Times(1) @@ -1231,7 +1196,7 @@ func TestSessionManager_DecorateSession(t *testing.T) { sessionID := sm.Generate() require.NotEmpty(t, sessionID) - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) // Apply a decorator that wraps with an extra tool. @@ -1290,8 +1255,8 @@ func TestSessionManager_DecorateSession(t *testing.T) { // 2. Terminate sees the key and takes the Phase 2 path (storage.Delete). bindingMeta := map[string]string{sessiontypes.MetadataKeyIdentityBinding: "unauthenticated"} factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink) (vmcpsession.MultiSession, error) { sess := sessionmocks.NewMockMultiSession(ctrl) sess.EXPECT().ID().Return(id).AnyTimes() sess.EXPECT().GetMetadata().Return(bindingMeta).AnyTimes() @@ -1314,7 +1279,7 @@ func TestSessionManager_DecorateSession(t *testing.T) { sessionID := sm.Generate() require.NotEmpty(t, sessionID) - _, err := sm.CreateSession(context.Background(), sessionID) + _, err := sm.CreateSession(context.Background(), sessionID, nil) require.NoError(t, err) err = sm.DecorateSession(sessionID, func(sess sessiontypes.MultiSession) sessiontypes.MultiSession { @@ -1344,7 +1309,7 @@ func TestSessionManager_CheckSession(t *testing.T) { t.Helper() ctrl := gomock.NewController(t) f := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - f.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + f.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). AnyTimes().Return(nil, nil) f.EXPECT().RestoreSession(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). AnyTimes().Return(nil, nil) @@ -1534,7 +1499,7 @@ func TestNotifyBackendExpired(t *testing.T) { sm, storage := newTestSessionManager(t, factory, registry) sessionID := sm.Generate() - _, err := sm.CreateSession(t.Context(), sessionID) + _, err := sm.CreateSession(t.Context(), sessionID, nil) require.NoError(t, err) meta := seedBackendMetadata(t, storage, sessionID, @@ -1566,7 +1531,7 @@ func TestNotifyBackendExpired(t *testing.T) { sm, storage := newTestSessionManager(t, factory, registry) sessionID := sm.Generate() - _, err := sm.CreateSession(t.Context(), sessionID) + _, err := sm.CreateSession(t.Context(), sessionID, nil) require.NoError(t, err) meta := seedBackendMetadata(t, storage, sessionID, @@ -1655,7 +1620,7 @@ func TestNotifyBackendExpired(t *testing.T) { sm, storage := newTestSessionManager(t, factory, registry) sessionID := sm.Generate() - _, err := sm.CreateSession(t.Context(), sessionID) + _, err := sm.CreateSession(t.Context(), sessionID, nil) require.NoError(t, err) // Seed MetadataKeyIdentityBinding into storage so Terminate recognises this @@ -1695,7 +1660,7 @@ func TestNotifyBackendExpired(t *testing.T) { sm, storage := newTestSessionManager(t, factory, registry) sessionID := sm.Generate() - _, err := sm.CreateSession(t.Context(), sessionID) + _, err := sm.CreateSession(t.Context(), sessionID, nil) require.NoError(t, err) meta := seedBackendMetadata(t, storage, sessionID, @@ -1729,7 +1694,7 @@ func TestNotifyBackendExpired(t *testing.T) { sm, storage := newTestSessionManager(t, factory, registry) sessionID := sm.Generate() - _, err := sm.CreateSession(t.Context(), sessionID) + _, err := sm.CreateSession(t.Context(), sessionID, nil) require.NoError(t, err) meta := seedBackendMetadata(t, storage, sessionID, @@ -1763,7 +1728,7 @@ func TestNotifyBackendExpired(t *testing.T) { sm, storage := newTestSessionManager(t, factory, registry) sessionID := sm.Generate() - _, err := sm.CreateSession(t.Context(), sessionID) + _, err := sm.CreateSession(t.Context(), sessionID, nil) require.NoError(t, err) // Session must be in cache after CreateSession. @@ -1828,7 +1793,7 @@ func TestLoadSession_Phase2Marker_UsesIdentityBindingKey(t *testing.T) { tools := []vmcp.Tool{{Name: "restored-tool", Description: "a restored tool"}} ctrl := gomock.NewController(t) factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Times(0) sessionID := "identity-binding-session" diff --git a/pkg/vmcp/server/sessionmanager/sliding_ttl_regression_test.go b/pkg/vmcp/server/sessionmanager/sliding_ttl_regression_test.go new file mode 100644 index 0000000000..bbd7b5678a --- /dev/null +++ b/pkg/vmcp/server/sessionmanager/sliding_ttl_regression_test.go @@ -0,0 +1,86 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package sessionmanager + +import ( + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// slidingTTLRegressionTTL is the short sliding-window TTL used by the +// regression tests below, paired with miniredis.FastForward to deterministically +// simulate elapsed time without slowing down the test suite. +const slidingTTLRegressionTTL = 200 * time.Millisecond + +// slidingTTLRegressionEpsilon is added to/subtracted from +// slidingTTLRegressionTTL when fast-forwarding, so assertions are not flaky +// against exact-boundary timing. +const slidingTTLRegressionEpsilon = 50 * time.Millisecond + +// --------------------------------------------------------------------------- +// Regression tests: sliding-TTL renewal on session-manager Validate (§14) +// --------------------------------------------------------------------------- + +// TestRegression_Validate_SlidingTTL_SurvivesUnderActivity pins that +// Manager.Validate's storage.Load (Redis GETEX) renews the session's TTL, so +// a session under regular activity never expires even though the total +// elapsed time far exceeds the configured TTL. This is the positive half of +// the sliding-TTL pair; TestRegression_Validate_NoActivity_ExpiresAfterTTL is +// the negative half proving the fast-forward actually simulates expiry +// (without it, this test could pass vacuously if renewal were broken). +func TestRegression_Validate_SlidingTTL_SurvivesUnderActivity(t *testing.T) { + t.Parallel() + + mr := miniredis.RunT(t) + storage := newSharedRedisStorageWithTTL(t, mr, slidingTTLRegressionTTL) + backend := startMCPBackend(t, "backend-alpha", "echo") + sm := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend}) + + sessionID := createSession(t, sm, nil) + + // Repeatedly validate the session, advancing the fake clock by just under + // the TTL between each call. If Validate's Load renews the TTL (via + // GETEX), the session must never be reported as terminated, even though + // the cumulative elapsed time is many multiples of the configured TTL. + const iterations = 5 + for i := range iterations { + mr.FastForward(slidingTTLRegressionTTL - slidingTTLRegressionEpsilon) + + isTerminated, err := sm.Validate(sessionID) + require.NoError(t, err, "iteration %d: Validate must not error", i) + require.False(t, isTerminated, + "iteration %d: session must survive under repeated activity (sliding TTL renewal)", i) + } +} + +// TestRegression_Validate_NoActivity_ExpiresAfterTTL pins that a session +// which receives no Validate (or other Load) calls does expire once the +// sliding-window TTL elapses. Paired with +// TestRegression_Validate_SlidingTTL_SurvivesUnderActivity so the survival +// test cannot pass vacuously (e.g. if TTL enforcement were disabled +// entirely, the survival test alone would still pass). +func TestRegression_Validate_NoActivity_ExpiresAfterTTL(t *testing.T) { + t.Parallel() + + mr := miniredis.RunT(t) + storage := newSharedRedisStorageWithTTL(t, mr, slidingTTLRegressionTTL) + backend := startMCPBackend(t, "backend-alpha", "echo") + sm := newTestManagerWithSharedStorage(t, storage, []*vmcp.Backend{backend}) + + sessionID := createSession(t, sm, nil) + + // No Validate/Load calls occur here — advance the fake clock straight past + // the TTL. + mr.FastForward(slidingTTLRegressionTTL + slidingTTLRegressionEpsilon) + + isTerminated, err := sm.Validate(sessionID) + require.NoError(t, err) + assert.True(t, isTerminated, "session with no activity must be reported terminated once the TTL elapses") +} diff --git a/pkg/vmcp/server/sessionmanager/slow_backend_regression_test.go b/pkg/vmcp/server/sessionmanager/slow_backend_regression_test.go new file mode 100644 index 0000000000..09cfc57bd7 --- /dev/null +++ b/pkg/vmcp/server/sessionmanager/slow_backend_regression_test.go @@ -0,0 +1,408 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package sessionmanager + +import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/health" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// slowBackendDelay is the per-request delay of the simulated slow-but-working +// backend. It stands in for the real 10-25s latency in #5861, scaled down so +// the suite stays fast while remaining orders of magnitude above the healthy +// backends' sub-millisecond responses. +const slowBackendDelay = 2 * time.Second + +// startSlowMCPBackend starts an in-process MCP backend that delays every +// request by delay before responding. It is deliberately WORKING, not broken: +// #5861 is about a backend that is merely slow, which is why simulating it with +// a connection refusal or a 5xx would not reproduce the failure. +// +// The returned counter reports how many HTTP requests the backend received. It +// is the deterministic signal that the backend was (or was not) contacted: +// asserting it is zero proves the skip directly, where a latency assertion only +// infers it and is vulnerable to CI noise. +func startSlowMCPBackend(t *testing.T, backendID string, delay time.Duration) (*vmcp.Backend, *atomic.Int64) { + t.Helper() + mcpSrv := mcpserver.NewMCPServer(backendID, "1.0.0") + mcpSrv.AddTool( + mcpmcp.NewTool("slow_tool", + mcpmcp.WithDescription("A tool on a slow-but-working backend"), + mcpmcp.WithString("input", mcpmcp.Required()), + ), + func(_ context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + args, _ := req.Params.Arguments.(map[string]any) + input, _ := args["input"].(string) + return &mcpmcp.CallToolResult{ + Content: []mcpmcp.Content{mcpmcp.NewTextContent(input)}, + }, nil + }, + ) + streamableSrv := mcpserver.NewStreamableHTTPServer(mcpSrv) + var requests atomic.Int64 + mux := http.NewServeMux() + // Count and delay every request, including the initialize handshake — the + // delay is what makes this backend slow rather than broken. + mux.Handle("/mcp", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + // Honour cancellation while delaying: on a regression several requests are + // in flight, and an unconditional sleep would block ts.Close in t.Cleanup + // for the full delay on each one. + select { + case <-time.After(delay): + case <-r.Context().Done(): + return + } + streamableSrv.ServeHTTP(w, r) + })) + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + return &vmcp.Backend{ + ID: backendID, + Name: backendID, + BaseURL: ts.URL + "/mcp", + TransportType: "streamable-http", + }, &requests +} + +// healthStub is a health.StatusProvider returning fixed per-backend statuses, +// standing in for a live health.Monitor that has already classified the slow +// backend. #5861's premise is that the monitor ALREADY knows the backend is bad; +// the bug is that session creation never asks. +// +// Named to avoid shadowing the imported health package at call sites. +type healthStub map[string]vmcp.BackendHealthStatus + +func (s healthStub) QueryBackendStatus(backendID string) (vmcp.BackendHealthStatus, bool) { + status, ok := s[backendID] + return status, ok +} + +// newTestManagerWithSharedStorageAndHealth is newTestManagerWithSharedStorage +// with a health.StatusProvider wired, for tests that need the Redis-backed +// restore path (a fresh Manager over shared storage has an empty cache, so +// GetMultiSession restores rather than serving from cache). +func newTestManagerWithSharedStorageAndHealth( + t *testing.T, + storage transportsession.DataStorage, + backends []*vmcp.Backend, + backendHealth health.StatusProvider, +) *Manager { + t.Helper() + backendList := make([]vmcp.Backend, len(backends)) + for i, b := range backends { + backendList[i] = *b + } + registry := vmcp.NewImmutableRegistry(backendList) + factory := vmcpsession.NewSessionFactory(newUnauthenticatedAuthRegistry(t)) + + sm, cleanup, err := New( + storage, + &FactoryConfig{Base: factory, BackendHealth: backendHealth}, + registry, + ) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, cleanup(context.Background())) }) + return sm +} + +// newTestManagerWithHealth creates a Manager over backends with the given +// health.StatusProvider (nil = health monitoring disabled). Uses in-memory +// session storage; the health-gating tests do not exercise Redis. +func newTestManagerWithHealth( + t *testing.T, backends []*vmcp.Backend, backendHealth health.StatusProvider, +) *Manager { + t.Helper() + backendList := make([]vmcp.Backend, len(backends)) + for i, b := range backends { + backendList[i] = *b + } + registry := vmcp.NewImmutableRegistry(backendList) + factory := vmcpsession.NewSessionFactory(newUnauthenticatedAuthRegistry(t)) + + storage, err := transportsession.NewLocalSessionDataStorage(time.Hour) + require.NoError(t, err) + t.Cleanup(func() { _ = storage.Close() }) + + sm, cleanup, err := New( + storage, + &FactoryConfig{Base: factory, BackendHealth: backendHealth}, + registry, + ) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, cleanup(context.Background())) }) + return sm +} + +// --------------------------------------------------------------------------- +// Regression tests: #5861 — a slow backend must not inflate tenant session +// latency once the health monitor has classified it +// --------------------------------------------------------------------------- + +// TestRegression_CreateSession_SkipsKnownBadSlowBackend pins the fix for #5861: +// a backend the health monitor has already classified as bad must not be +// re-attempted during per-session backend init. +// +// Before the fix, the new-session backend list was the registry unfiltered and +// session.makeBaseSession blocked on wg.Wait() for every backend, so a single +// slow backend set the floor for the whole tenant's initialize latency — a +// 10-25s backend against a ~30s client timeout is the reported coin flip. +// Health status gated capability aggregation (core.filterHealthyBackends) but +// never connection establishment, and establishment runs first. +// +// Both halves are asserted over the SAME Manager: the bad backend receives zero +// requests AND the healthy peer still holds a session. Asserting only the former +// would pass if the fix skipped everything, which is the failure mode a +// health-gating change is most likely to introduce. +// +// The request counter is the whole assertion — there is deliberately no +// wall-clock latency bound. Elapsed time only infers the skip, and a ceiling +// tight enough to be meaningful is also tight enough to flake under -race on a +// loaded CI runner. +// +// Degraded is NOT covered here because it is not skipped: see +// health.ShouldOpenSession for why (three producers, only one of which is +// latency) and TestRegression_CreateSession_DegradedIsAttempted. +func TestRegression_CreateSession_SkipsKnownBadSlowBackend(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + status vmcp.BackendHealthStatus + }{ + { + // The state the reporter observed: real latency exceeds the probe + // timeout, so checks fail until the unhealthy threshold is reached. + name: "unhealthy", + status: vmcp.BackendUnhealthy, + }, + { + // Terminal until an operator intervenes; connecting can only burn + // session-establishment latency. + name: "unauthenticated", + status: vmcp.BackendUnauthenticated, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + fast := startMCPBackend(t, "backend-fast", "echo") + slow, slowRequests := startSlowMCPBackend(t, "backend-slow", slowBackendDelay) + + sm := newTestManagerWithHealth(t, []*vmcp.Backend{fast, slow}, healthStub{ + fast.ID: vmcp.BackendHealthy, + slow.ID: tt.status, + }) + + sessionID := createSession(t, sm, nil) + require.NotEmpty(t, sessionID) + + assert.Zero(t, slowRequests.Load(), + "session creation must not contact a backend already known %s "+ + "(#5861); it received %d request(s)", tt.status, slowRequests.Load()) + + sess, ok := sm.GetMultiSession(context.Background(), sessionID) + require.True(t, ok) + require.NotNil(t, sess) + assert.Contains(t, sess.BackendSessions(), fast.ID, + "the healthy backend must still hold a session — skipping the %s "+ + "backend must not drop its healthy peer", tt.status) + }) + } +} + +// TestRegression_CreateSession_DegradedIsAttempted pins that degraded backends +// are still connected, which is the deliberate limit on #5861's fix. +// +// vmcp.BackendDegraded has three producers and only one is about latency. The +// "recovering" producer (statusTracker.RecordSuccess) forces degraded on ANY +// success recorded while consecutiveFailures > 0, overriding the health check's +// own verdict — so a backend that just answered in microseconds carries degraded +// for up to one CheckInterval. Skipping degraded would exclude that fast, working +// backend from every session created in the window right after it recovers. +// +// See health.ShouldOpenSession. If a degradation reason is ever recorded +// alongside the status, the latency producer alone could be skipped and this test +// should be revisited. +func TestRegression_CreateSession_DegradedIsAttempted(t *testing.T) { + t.Parallel() + + backend := startMCPBackend(t, "backend-degraded", "echo") + sm := newTestManagerWithHealth(t, []*vmcp.Backend{backend}, healthStub{ + backend.ID: vmcp.BackendDegraded, + }) + + sessionID := createSession(t, sm, nil) + sess, ok := sm.GetMultiSession(context.Background(), sessionID) + require.True(t, ok) + require.NotNil(t, sess) + assert.Contains(t, sess.BackendSessions(), backend.ID, + "a degraded backend must still be attempted: degraded also means "+ + "'recovering' and 'auth retrying', neither of which is slow") +} + +// TestRegression_CreateSession_MonitorMissHonoursRegistryStatus pins the +// resolution order in Manager.shouldOpenSession: when the monitor does not track +// a backend, the registry's own status decides. +// +// Both directions are covered, because a test suite that only exercises +// admitted statuses would still pass with the registry fallback deleted +// entirely — the zero value is admitted too. +// +// - Untracked + registry says unhealthy: must SKIP. This is the case the +// fallback exists for; a k8s workload already reported unhealthy is honoured +// before the monitor has caught up. +// - Untracked + registry says unknown / unset: must ATTEMPT. "Not yet +// classified" fails open, or a cold monitor strands sessions with zero +// backends during pod startup. +// - Tracked as unknown: must ATTEMPT. The monitor stores Unknown with a +// non-zero failure count for a first failure below the unhealthy threshold +// (statusTracker.RecordFailure); nothing is confirmed until the threshold +// (default 3) is reached. The registry value is deliberately different here +// to prove the monitor's answer wins when it has one. +func TestRegression_CreateSession_MonitorMissHonoursRegistryStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // monitored is the status the health monitor reports; only consulted + // when tracked is true. + monitored vmcp.BackendHealthStatus + tracked bool + // registryStatus is the status the backend carries in the registry. + registryStatus vmcp.BackendHealthStatus + wantConnected bool + }{ + { + name: "untracked by monitor, registry reports unhealthy", + tracked: false, + registryStatus: vmcp.BackendUnhealthy, + wantConnected: false, + }, + { + name: "untracked by monitor, registry reports unknown (k8s Pending)", + tracked: false, + registryStatus: vmcp.BackendUnknown, + wantConnected: true, + }, + { + name: "untracked by monitor, registry status unset", + tracked: false, + registryStatus: "", + wantConnected: true, + }, + { + name: "tracked as unknown overrides an unhealthy registry status", + monitored: vmcp.BackendUnknown, + tracked: true, + registryStatus: vmcp.BackendUnhealthy, + wantConnected: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + backend := startMCPBackend(t, "backend-pending", "echo") + backend.HealthStatus = tt.registryStatus + + stub := healthStub{} + if tt.tracked { + stub[backend.ID] = tt.monitored + } + + sm := newTestManagerWithHealth(t, []*vmcp.Backend{backend}, stub) + sessionID := createSession(t, sm, nil) + + sess, ok := sm.GetMultiSession(context.Background(), sessionID) + require.True(t, ok) + require.NotNil(t, sess) + if tt.wantConnected { + assert.Contains(t, sess.BackendSessions(), backend.ID, + "backend must be attempted: monitor tracked=%v status=%q, "+ + "registry status %q", tt.tracked, tt.monitored, tt.registryStatus) + } else { + assert.NotContains(t, sess.BackendSessions(), backend.ID, + "backend must be skipped: monitor tracked=%v, registry status %q "+ + "is confirmed-bad", tt.tracked, tt.registryStatus) + } + }) + } +} + +// TestRegression_CreateSession_NilHealthProviderConnectsAll pins that disabling +// health monitoring preserves the pre-fix behaviour exactly: every backend is +// attempted. A nil provider must not be read as "everything is unhealthy". +func TestRegression_CreateSession_NilHealthProviderConnectsAll(t *testing.T) { + t.Parallel() + + fast := startMCPBackend(t, "backend-fast", "echo") + sm := newTestManagerWithHealth(t, []*vmcp.Backend{fast}, nil) + + sessionID := createSession(t, sm, nil) + + sess, ok := sm.GetMultiSession(context.Background(), sessionID) + require.True(t, ok) + assert.NotEmpty(t, sess.BackendSessions(), + "with health monitoring disabled every backend must be attempted") +} + +// TestRegression_RestoreSession_IsNotHealthFiltered pins that health gating is +// confined to NEW sessions and never narrows a restored one. +// +// RestoreSession intersects the offered backend list with the session's stored +// backend IDs (session.RestoreSession), so filtering the list on the restore path +// would not defer a connection — it would DROP a backend the session already +// held. Because the routing table is rebuilt from whatever reconnects, that +// backend would stay gone for the rest of the session's life even after it +// recovered, since nothing reconnects an established session's backends. +// +// The backend is healthy when the session is created and unhealthy when it is +// restored, which is exactly the transition that would silently shrink a live +// session's capabilities. +func TestRegression_RestoreSession_IsNotHealthFiltered(t *testing.T) { + t.Parallel() + + mr := miniredis.RunT(t) + storage := newSharedRedisStorage(t, mr) + backend := startMCPBackend(t, "backend-alpha", "echo") + + // Healthy at creation: the session connects and stores the backend ID. + smWriter := newTestManagerWithSharedStorageAndHealth(t, storage, + []*vmcp.Backend{backend}, healthStub{backend.ID: vmcp.BackendHealthy}) + sessionID := createSession(t, smWriter, nil) + + // Unhealthy at restore. A fresh manager has an empty cache, so + // GetMultiSession takes the restore path. + smReader := newTestManagerWithSharedStorageAndHealth(t, storage, + []*vmcp.Backend{backend}, healthStub{backend.ID: vmcp.BackendUnhealthy}) + + sess, ok := smReader.GetMultiSession(t.Context(), sessionID) + require.True(t, ok) + require.NotNil(t, sess) + assert.Contains(t, sess.BackendSessions(), backend.ID, + "restore must not drop a backend the session already held, even when "+ + "health now reports it unhealthy — the drop would be permanent for "+ + "this session") + assert.NotEmpty(t, sess.Tools(), + "the restored session's routing table must still carry the backend's tools") +} diff --git a/pkg/vmcp/server/sessionmanager/telemetry_test.go b/pkg/vmcp/server/sessionmanager/telemetry_test.go index d32acac96d..332b62d73c 100644 --- a/pkg/vmcp/server/sessionmanager/telemetry_test.go +++ b/pkg/vmcp/server/sessionmanager/telemetry_test.go @@ -8,14 +8,14 @@ import ( "fmt" "testing" - "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" sdkmetric "go.opentelemetry.io/otel/sdk/metric" "go.opentelemetry.io/otel/sdk/metric/metricdata" tracenoop "go.opentelemetry.io/otel/trace/noop" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" ) diff --git a/pkg/vmcp/server/sse_keepalive_regression_test.go b/pkg/vmcp/server/sse_keepalive_regression_test.go new file mode 100644 index 0000000000..3f61560513 --- /dev/null +++ b/pkg/vmcp/server/sse_keepalive_regression_test.go @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" +) + +// TestRegression_SSEKeepAlive_PeriodicBytesOnIdleStream verifies that an idle +// SSE stream opened against the Serve-path streamable server keeps producing +// bytes (a keep-alive ping or heartbeat) within a short window. A regression +// that drops the heartbeat wiring will cause this test to fail: no bytes arrive +// on an idle GET stream, so proxies/load balancers eventually close it. +// +// The streamable server mounts the SDK keep-alive over the configured +// HeartbeatInterval, so the test sets a short interval (100ms) and reads with a +// 500ms deadline — generous relative to the interval but tight enough that a +// missing keep-alive is caught rather than hanging the suite. +func TestRegression_SSEKeepAlive_PeriodicBytesOnIdleStream(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + testTool := vmcp.Tool{Name: "keepalive-tool", Description: "keep-alive regression anchor"} + factory, _ := newToolSessionFactory(t, ctrl, []vmcp.Tool{testTool}) + fc := &fakeCore{tools: []vmcp.Tool{testTool}} + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + HeartbeatInterval: 100 * time.Millisecond, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + // Serve through the real Handler wiring rather than a hand-built streamable + // server. This exercises the ServerConfig.HeartbeatInterval → Config → + // Handler → WithHeartbeatInterval path that server.go adds, so removing that + // wiring line in Handler regresses this test (no keep-alive bytes arrive). + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + + // initialize → obtain a session ID for the subsequent GET stream. + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "keepalive-test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode, "initialize should succeed") + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID, "session ID should be returned in Mcp-Session-Id header") + + // Open a long-lived SSE GET stream. The stream must stay open and emit + // keep-alive bytes even though no client request is in flight. + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL+"/mcp", nil) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Mcp-Session-Id", sessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "GET stream should open with 200") + + // Read with a 500ms deadline. The keep-alive interval is 100ms, so at least + // one ping/heartbeat should land within the deadline; a missing keep-alive + // surfaces as a read timeout with zero bytes — the regression. The streamable + // server's body is not a net.Conn, so SetReadDeadline is unavailable; drive + // the read on a goroutine and race it against a timer. + type readResult struct { + data []byte + err error + } + ch := make(chan readResult, 1) + go func() { + buf := make([]byte, 4096) + n, err := resp.Body.Read(buf) + ch <- readResult{data: buf[:n], err: err} + }() + + var got []byte + select { + case r := <-ch: + got = r.data + case <-time.After(500 * time.Millisecond): + } + + // At least some bytes must arrive — the keep-alive ping, an SSE comment + // (`:` prefix), or a JSON-RPC ping request frame. Zero bytes within the + // window means the keep-alive wiring is broken; fail loudly as a regression. + if len(got) == 0 { + t.Fatal("no bytes received on idle SSE stream within 500ms; keep-alive/heartbeat " + + "is not emitting periodic bytes (regression: proxies will close idle streams)") + } + + body := string(got) + t.Logf("received keep-alive bytes on idle SSE stream: %q", body) + // Accept any of: a JSON-RPC ping request, an SSE comment, or any SSE data + // frame. The point is that bytes flow; their exact form is shim-defined. + hasPing := strings.Contains(body, `"method":"ping"`) || + strings.Contains(body, `"method": "ping"`) + hasSSEComment := strings.Contains(body, ":") + hasDataFrame := strings.Contains(body, "data:") + assert.True(t, hasPing || hasSSEComment || hasDataFrame, + "keep-alive bytes should be an SSE comment, a data frame, or a JSON-RPC ping; got %q", body) +} diff --git a/pkg/vmcp/server/status.go b/pkg/vmcp/server/status.go index 84e830e022..ca90c692ae 100644 --- a/pkg/vmcp/server/status.go +++ b/pkg/vmcp/server/status.go @@ -29,13 +29,20 @@ type BackendStatus struct { Health string `json:"health"` // "healthy", "degraded", "unhealthy", "unauthenticated", "unknown" Transport string `json:"transport"` // MCP transport protocol AuthType string `json:"auth_type,omitempty"` // "unauthenticated", "header_injection", "token_exchange" + // MCPRevision is the backend's resolved MCP protocol revision ("2025-11-25" + // or "2026-07-28"), mirroring the value the operator surfaces on + // VirtualMCPServer.status.discoveredBackends[]. Empty when the backend has + // not been probed yet or when health monitoring is disabled, since the + // revision is only observable through the live health state. + MCPRevision string `json:"mcp_revision,omitempty"` } // handleStatus handles /status HTTP requests for operational visibility. // // Security Note: This endpoint is unauthenticated to support operator consumption -// and debugging. It exposes operational metadata (backend names, auth types) -// but NOT secrets, tokens, internal URLs, or request data. In sensitive multi-tenant +// and debugging. It exposes operational metadata (backend names, auth types, +// negotiated MCP protocol revision) but NOT secrets, tokens, internal URLs, or +// request data. In sensitive multi-tenant // deployments, restrict access to this endpoint via network policies. // // For minimal health checking (load balancers), use /health instead. @@ -87,16 +94,22 @@ func (s *Server) buildStatusResponse(ctx context.Context) StatusResponse { hasHealthyBackend := false for _, backend := range backends { // Prefer the live health monitor state over the static registry value. + // The MCP revision has no static counterpart — it is only ever observable + // through the live state, so it stays empty when health monitoring is + // disabled or the backend has not been probed yet. healthStatus := backend.HealthStatus + var mcpRevision string if liveState, ok := liveHealthStates[backend.ID]; ok { healthStatus = liveState.Status + mcpRevision = liveState.MCPRevision } status := BackendStatus{ - Name: backend.Name, - Health: string(healthStatus), - Transport: backend.TransportType, - AuthType: getAuthType(backend.AuthConfig), + Name: backend.Name, + Health: string(healthStatus), + Transport: backend.TransportType, + AuthType: getAuthType(backend.AuthConfig), + MCPRevision: mcpRevision, } backendStatuses = append(backendStatuses, status) diff --git a/pkg/vmcp/server/status_test.go b/pkg/vmcp/server/status_test.go index 4a14d655ce..ac147a9e3e 100644 --- a/pkg/vmcp/server/status_test.go +++ b/pkg/vmcp/server/status_test.go @@ -15,6 +15,7 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/vmcp" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" @@ -34,10 +35,11 @@ type StatusResponse struct { // BackendStatus mirrors the server's backend status structure for test deserialization. type BackendStatus struct { - Name string `json:"name"` - Health string `json:"health"` - Transport string `json:"transport"` - AuthType string `json:"auth_type,omitempty"` + Name string `json:"name"` + Health string `json:"health"` + Transport string `json:"transport"` + AuthType string `json:"auth_type,omitempty"` + MCPRevision string `json:"mcp_revision,omitempty"` } // createTestServerWithBackends creates a test server instance with custom backends @@ -210,12 +212,29 @@ func createTestServerWithHealthMonitor( t.Cleanup(ctrl.Finish) mockBackendClient := mocks.NewMockBackendClient(ctrl) - rt := router.NewSessionRouter(&vmcp.RoutingTable{}) - if setupMock != nil { setupMock(mockBackendClient) } + return createTestServerWithClient(t, backends, monitorCfg, mockBackendClient, groupRef) +} + +// createTestServerWithClient is createTestServerWithHealthMonitor but takes an +// already-built vmcp.BackendClient instead of constructing a bare mock — used +// when a test needs a client satisfying an optional interface (e.g. the +// health monitor's CachedRevision accessor) that the plain mock doesn't +// implement. +func createTestServerWithClient( + t *testing.T, + backends []vmcp.Backend, + monitorCfg health.MonitorConfig, + client vmcp.BackendClient, + groupRef string, +) *server.Server { + t.Helper() + + rt := router.NewSessionRouter(&vmcp.RoutingTable{}) + port := networking.FindAvailable() require.NotZero(t, port, "Failed to find available port") @@ -233,7 +252,7 @@ func createTestServerWithHealthMonitor( GroupRef: groupRef, HealthMonitorConfig: healthMonCfg, SessionFactory: newNoopMockFactory(t), Aggregator: newStubAggregator(nil), - }, rt, mockBackendClient, vmcp.NewImmutableRegistry(backends), nil) + }, rt, client, vmcp.NewImmutableRegistry(backends), nil) require.NoError(t, err) type startResult struct { @@ -367,6 +386,13 @@ func TestStatusEndpoint_ReflectsLiveHealthMonitor_Healthy(t *testing.T) { assert.True(t, status.Healthy) require.Len(t, status.Backends, 1) assert.Equal(t, string(vmcp.BackendHealthy), status.Backends[0].Health) + // The monitor is running here, but this test's client is a plain mock with + // no CachedRevision method — the shape of any BackendClient other than + // httpBackendClient. The revision must stay empty rather than being + // invented, so this is the third distinct state of mcp_revision alongside + // the two dedicated tests below. + assert.Empty(t, status.Backends[0].MCPRevision, + "revision must stay empty for a client that cannot report one") } // TestStatusEndpoint_FallsBackToRegistry_WhenMonitorDisabled confirms the @@ -392,3 +418,67 @@ func TestStatusEndpoint_FallsBackToRegistry_WhenMonitorDisabled(t *testing.T) { assert.Equal(t, string(vmcp.BackendHealthy), healthByName["healthy-backend"]) assert.Equal(t, string(vmcp.BackendUnhealthy), healthByName["unhealthy-backend"]) } + +// revClientStub wraps a MockBackendClient and adds the optional CachedRevision +// accessor the health monitor reads (mirrors pkg/vmcp/health/monitor_test.go's +// revClientStub, unexported there and so not reusable across packages). +type revClientStub struct { + *mocks.MockBackendClient + rev mcpparser.Revision +} + +func (s revClientStub) CachedRevision(string) (mcpparser.Revision, bool) { return s.rev, true } + +// TestStatusEndpoint_MCPRevision_Populated verifies /status surfaces the +// backend's negotiated MCP revision once the health monitor has probed it — +// the field has no static registry counterpart, so this requires a client +// that implements CachedRevision (the bare mock does not). +func TestStatusEndpoint_MCPRevision_Populated(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + client := revClientStub{MockBackendClient: mockClient, rev: mcpparser.RevisionModern} + + backends := []vmcp.Backend{{ID: "b1", Name: "modern-backend", TransportType: "streamable-http"}} + monitorCfg := health.MonitorConfig{ + CheckInterval: 5 * time.Millisecond, + UnhealthyThreshold: 3, + Timeout: time.Second, + } + srv := createTestServerWithClient(t, backends, monitorCfg, client, "") + + require.Eventually(t, func() bool { + status := queryStatus(t, srv) + return len(status.Backends) == 1 && status.Backends[0].MCPRevision == mcpparser.MCPVersionModern + }, 5*time.Second, 20*time.Millisecond, "expected /status to report the negotiated MCP revision") +} + +// TestStatusEndpoint_MCPRevision_AbsentWhenMonitorDisabled mirrors +// TestStatusEndpoint_FallsBackToRegistry_WhenMonitorDisabled: with no health +// monitor (and so no client capable of CachedRevision ever consulted), the +// field must be empty and — because of omitempty — the "mcp_revision" key +// must be absent from the JSON entirely, not merely an empty string. +func TestStatusEndpoint_MCPRevision_AbsentWhenMonitorDisabled(t *testing.T) { + t.Parallel() + + backends := []vmcp.Backend{{ID: "b1", Name: "test-backend", HealthStatus: vmcp.BackendHealthy}} + srv := createTestServerWithBackends(t, backends, "") + + resp, err := http.Get("http://" + srv.Address() + "/status") + require.NoError(t, err) + defer resp.Body.Close() + + var raw struct { + Backends []map[string]any `json:"backends"` + } + require.NoError(t, json.NewDecoder(resp.Body).Decode(&raw)) + require.Len(t, raw.Backends, 1) + _, present := raw.Backends[0]["mcp_revision"] + assert.False(t, present, "mcp_revision must be omitted, not present-empty, when unobservable") +} diff --git a/pkg/vmcp/server/telemetry_integration_test.go b/pkg/vmcp/server/telemetry_integration_test.go index 62ad095cbb..0a23b05e43 100644 --- a/pkg/vmcp/server/telemetry_integration_test.go +++ b/pkg/vmcp/server/telemetry_integration_test.go @@ -83,7 +83,7 @@ func newBackendAwareTestFactory(tools []vmcp.Tool, rt *vmcp.RoutingTable) *backe } func (f *backendAwareTestFactory) MakeSessionWithID( - _ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, + _ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink, ) (vmcpsession.MultiSession, error) { return f.newSession(id), nil } @@ -173,7 +173,7 @@ func TestIntegration_TelemetryMiddleware(t *testing.T) { // Use MinTimes(1) to verify the backend client is actually called during tool execution. // If the tool call doesn't reach the backend client, this will cause a test failure. mockBackendClient.EXPECT(). - CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(&vmcp.ToolCallResult{ StructuredContent: map[string]any{"result": "found"}, Content: []vmcp.Content{}, @@ -386,3 +386,137 @@ func TestIntegration_TelemetryMiddleware(t *testing.T) { cancelServer() } + +// TestIntegration_TelemetryRunsBeforeClassificationRejection is a regression +// guard for the middleware ordering in server.go: classifyingHandler +// must stay closer to the handler than the telemetry middleware, so a +// request that classifyingHandler rejects is still recorded as an +// incoming request instead of being dropped before telemetry ever sees it. +// If classification is ever reordered in front of telemetry, this test +// starts failing because the rejected request would never reach it. +// +// Note: like TestIntegration_TelemetryMiddleware, this does not use +// t.Parallel() since telemetry.NewProvider sets global OTel providers. +// +//nolint:paralleltest // shares global OTel provider state with other telemetry tests +func TestIntegration_TelemetryRunsBeforeClassificationRejection(t *testing.T) { + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + ctx := context.Background() + + telemetryProvider, err := telemetry.NewProvider(ctx, telemetry.Config{ + ServiceName: "vmcp-telemetry-ordering-test", + ServiceVersion: "1.0.0", + EnablePrometheusMetricsPath: true, + }) + require.NoError(t, err) + t.Cleanup(func() { telemetryProvider.Shutdown(ctx) }) + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockBackendClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + + backends := []vmcp.Backend{ + { + ID: "search-svc", + Name: "Search Service", + BaseURL: "http://search-svc:8080", + TransportType: "streamable-http", + HealthStatus: vmcp.BackendHealthy, + }, + } + + rt := router.NewSessionRouter(&vmcp.RoutingTable{}) + agg := aggregator.NewDefaultAggregator( + mockBackendClient, aggregator.NewPrefixConflictResolver("{workload}_"), nil, nil) + factory := newBackendAwareTestFactory(nil, &vmcp.RoutingTable{}) + + srv, err := New(ctx, &Config{ + Name: "telemetry-ordering-vmcp", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, // Random available port + TelemetryProvider: telemetryProvider, + SessionFactory: factory, + Aggregator: agg, + }, rt, mockBackendClient, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err) + + serverCtx, cancelServer := context.WithCancel(ctx) + t.Cleanup(cancelServer) + + serverErrCh := make(chan error, 1) + go func() { + if err := srv.Start(serverCtx); err != nil && !errors.Is(err, context.Canceled) { + serverErrCh <- err + } + }() + + select { + case <-srv.Ready(): + case err := <-serverErrCh: + t.Fatalf("Server failed to start: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("Server timeout waiting for ready") + } + + baseURL := "http://" + srv.Address() + + // Same malformed-Modern rejection payload as + // TestIntegration_RealBackend_ModernRequestRejectedByClassification: a + // reserved _meta key signals Modern, but no valid protocolVersion is + // present and the header names a different (Legacy) version, so + // classifyingHandler rejects with -32020 before dispatch. + body := map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{ + "_meta": map[string]any{"io.modelcontextprotocol/clientInfo": map[string]any{"name": "test"}}, + "name": "echo", + "arguments": map[string]any{}, + }, + } + payload, err := json.Marshal(body) + require.NoError(t, err) + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/mcp", bytes.NewReader(payload)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", "2025-11-25") + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusBadRequest, resp.StatusCode, "body: %s", string(respBody)) + + var rpc struct { + Error struct { + Code int64 `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(respBody, &rpc), "body: %s", string(respBody)) + require.Equal(t, int64(-32020), rpc.Error.Code, "expected CodeHeaderMismatch") + + metricsResp, err := http.Get(baseURL + "/metrics") + require.NoError(t, err) + defer metricsResp.Body.Close() + require.Equal(t, http.StatusOK, metricsResp.StatusCode) + + metricsBody, err := io.ReadAll(metricsResp.Body) + require.NoError(t, err) + metrics := string(metricsBody) + + assert.Contains(t, metrics, "toolhive_mcp_requests", + "telemetry must record the request even though classification rejected it downstream") + assert.Contains(t, metrics, `server="telemetry-ordering-vmcp"`, + "request metrics should identify this vMCP server") + + cancelServer() +} diff --git a/pkg/vmcp/server/testfactory_test.go b/pkg/vmcp/server/testfactory_test.go index 6ad6aff45d..693e301494 100644 --- a/pkg/vmcp/server/testfactory_test.go +++ b/pkg/vmcp/server/testfactory_test.go @@ -28,7 +28,7 @@ type minimalTestFactory struct{} var _ vmcpsession.MultiSessionFactory = (*minimalTestFactory)(nil) func (*minimalTestFactory) MakeSessionWithID( - _ context.Context, _ string, _ *auth.Identity, _ []*vmcp.Backend, + _ context.Context, _ string, _ *auth.Identity, _ []*vmcp.Backend, _ vmcpsession.ListChangedSink, ) (vmcpsession.MultiSession, error) { return nil, fmt.Errorf("minimalTestFactory: MakeSessionWithID not implemented in test helper") } diff --git a/pkg/vmcp/server/testutil_test.go b/pkg/vmcp/server/testutil_test.go index 9a29e20210..cb8fce7950 100644 --- a/pkg/vmcp/server/testutil_test.go +++ b/pkg/vmcp/server/testutil_test.go @@ -11,8 +11,8 @@ import ( "net/http/httptest" "testing" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" ) // startRealMCPBackend creates a real in-process MCP server over streamable-HTTP @@ -51,6 +51,50 @@ func startRealMCPBackend(t *testing.T) string { return ts.URL + "/mcp" } +// startEchoPingBackend is startRealMCPBackend with a second "ping" tool +// (returns "pong"). Two tools let a per-principal projection test give each +// principal a DISTINCT positively-permitted tool, so it can wait for each +// session to register its own permitted tool (proving registration completed +// and the principal resolved) before asserting the ABSENCE of the other +// principal's tool — closing the "empty list because filtered" vs "empty list +// because registration still pending" ambiguity. Returns the /mcp URL. +func startEchoPingBackend(t *testing.T) string { + t.Helper() + + mcpSrv := mcpserver.NewMCPServer("real-backend", "1.0.0") + mcpSrv.AddTool( + mcpmcp.NewTool("echo", + mcpmcp.WithDescription("Echoes the input back"), + mcpmcp.WithString("input", mcpmcp.Required()), + ), + func(_ context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + args, _ := req.Params.Arguments.(map[string]any) + input, _ := args["input"].(string) + return &mcpmcp.CallToolResult{ + Content: []mcpmcp.Content{mcpmcp.NewTextContent(input)}, + }, nil + }, + ) + mcpSrv.AddTool( + mcpmcp.NewTool("ping", + mcpmcp.WithDescription("Returns pong"), + ), + func(_ context.Context, _ mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + return &mcpmcp.CallToolResult{ + Content: []mcpmcp.Content{mcpmcp.NewTextContent("pong")}, + }, nil + }, + ) + + streamableSrv := mcpserver.NewStreamableHTTPServer(mcpSrv) + mux := http.NewServeMux() + mux.Handle("/mcp", streamableSrv) + + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + return ts.URL + "/mcp" +} + // MCPTestClient provides higher-level test utilities for MCP protocol interactions. // This reduces boilerplate and improves test readability by providing semantic methods // instead of manual JSON-RPC construction. @@ -59,6 +103,11 @@ type MCPTestClient struct { sessionID string t *testing.T nextID int + // principal, when non-empty, is sent as the X-Test-Principal header on every + // POST so a test server whose identity middleware honors that header (see + // buildCedarAuthzServer) authenticates this client as a distinct principal. + // Existing callers never set this and are unaffected. + principal string } // NewMCPTestClient creates a new test client for the given server URL. @@ -126,11 +175,48 @@ func (c *MCPTestClient) CallTool(toolName string, args map[string]any) *http.Res return resp } +// ReadResource calls resources/read for the given URI and returns the raw response. +func (c *MCPTestClient) ReadResource(uri string) *http.Response { + c.t.Helper() + + resp := c.postMCP(map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "resources/read", + "params": map[string]any{"uri": uri}, + }, c.sessionID) + c.nextID++ + return resp +} + +// GetPrompt calls prompts/get for the given prompt name and returns the raw response. +func (c *MCPTestClient) GetPrompt(name string) *http.Response { + c.t.Helper() + + resp := c.postMCP(map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "prompts/get", + "params": map[string]any{"name": name}, + }, c.sessionID) + c.nextID++ + return resp +} + // SessionID returns the current session ID (available after InitializeSession). func (c *MCPTestClient) SessionID() string { return c.sessionID } +// WithPrincipal sets the X-Test-Principal header this client sends with every +// subsequent request, so a server whose identity middleware honors that header +// authenticates this client as principal rather than the fixed fallback +// principal. Returns the client for chaining. +func (c *MCPTestClient) WithPrincipal(principal string) *MCPTestClient { + c.principal = principal + return c +} + // Terminate sends a DELETE request to terminate the session. func (c *MCPTestClient) Terminate() *http.Response { c.t.Helper() @@ -172,6 +258,9 @@ func (c *MCPTestClient) postMCP(body map[string]any, sessionID string) *http.Res if sessionID != "" { req.Header.Set("Mcp-Session-Id", sessionID) } + if c.principal != "" { + req.Header.Set("X-Test-Principal", c.principal) + } resp, err := http.DefaultClient.Do(req) if err != nil { diff --git a/pkg/vmcp/server/workflow_converter.go b/pkg/vmcp/server/workflow_converter.go index edbebde9b5..36a13200c1 100644 --- a/pkg/vmcp/server/workflow_converter.go +++ b/pkg/vmcp/server/workflow_converter.go @@ -71,6 +71,7 @@ func ConvertConfigToWorkflowDefinitions( Steps: steps, Timeout: timeout, Output: ct.Output, + Annotations: ct.Annotations, Metadata: make(map[string]string), } diff --git a/pkg/vmcp/server/workflow_converter_test.go b/pkg/vmcp/server/workflow_converter_test.go index 5dcaaa3285..8301e21fe4 100644 --- a/pkg/vmcp/server/workflow_converter_test.go +++ b/pkg/vmcp/server/workflow_converter_test.go @@ -297,3 +297,80 @@ func TestConvertConfigToWorkflowDefinitions_WithOutputConfig(t *testing.T) { }) } } + +// TestConvertConfigToWorkflowDefinitions_WithAnnotations tests that the +// Annotations field is correctly copied from CompositeToolConfig to +// WorkflowDefinition. +func TestConvertConfigToWorkflowDefinitions_WithAnnotations(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input []config.CompositeToolConfig + verify func(t *testing.T, defs map[string]*composer.WorkflowDefinition) + }{ + { + name: "composite tool with annotations", + input: []config.CompositeToolConfig{ + { + Name: "report_tool", + Description: "Read-only report", + Steps: []config.WorkflowStepConfig{ + {ID: "fetch", Type: "tool", Tool: "data.fetch"}, + }, + Annotations: &config.ToolAnnotationsOverride{ + Title: ptrTo("Report Tool"), + ReadOnlyHint: ptrTo(true), + }, + }, + }, + verify: func(t *testing.T, defs map[string]*composer.WorkflowDefinition) { + t.Helper() + require.Len(t, defs, 1) + + def, exists := defs["report_tool"] + require.True(t, exists) + require.NotNil(t, def.Annotations, "Annotations should be set on WorkflowDefinition") + require.NotNil(t, def.Annotations.Title) + assert.Equal(t, "Report Tool", *def.Annotations.Title) + require.NotNil(t, def.Annotations.ReadOnlyHint) + assert.True(t, *def.Annotations.ReadOnlyHint) + assert.Nil(t, def.Annotations.DestructiveHint) + }, + }, + { + name: "composite tool without annotations", + input: []config.CompositeToolConfig{ + { + Name: "plain_tool", + Steps: []config.WorkflowStepConfig{{ID: "s1", Type: "tool", Tool: "data.fetch"}}, + }, + }, + verify: func(t *testing.T, defs map[string]*composer.WorkflowDefinition) { + t.Helper() + require.Len(t, defs, 1) + + def, exists := defs["plain_tool"] + require.True(t, exists) + assert.Nil(t, def.Annotations, "Annotations should be nil when unset (derive at advertise time)") + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + result, err := ConvertConfigToWorkflowDefinitions(tt.input) + require.NoError(t, err) + require.NotNil(t, result) + + if tt.verify != nil { + tt.verify(t, result) + } + }) + } +} + +// ptrTo returns a pointer to v. Used in tests for optional pointer fields. +func ptrTo[T any](v T) *T { return &v } diff --git a/pkg/vmcp/session/connector_integration_test.go b/pkg/vmcp/session/connector_integration_test.go index c66b1bdeaf..cd93f21c07 100644 --- a/pkg/vmcp/session/connector_integration_test.go +++ b/pkg/vmcp/session/connector_integration_test.go @@ -11,11 +11,11 @@ import ( "testing" "github.com/google/uuid" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -114,7 +114,7 @@ func TestSessionFactory_Integration_CapabilityDiscovery(t *testing.T) { } factory := NewSessionFactory(newUnauthenticatedRegistry(t)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) require.NotNil(t, sess) t.Cleanup(func() { require.NoError(t, sess.Close()) }) @@ -143,7 +143,7 @@ func TestSessionFactory_Integration_CallTool(t *testing.T) { } factory := NewSessionFactory(newUnauthenticatedRegistry(t)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, sess.Close()) }) @@ -166,7 +166,7 @@ func TestSessionFactory_Integration_ReadResource(t *testing.T) { } factory := NewSessionFactory(newUnauthenticatedRegistry(t)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, sess.Close()) }) @@ -189,7 +189,7 @@ func TestSessionFactory_Integration_GetPrompt(t *testing.T) { } factory := NewSessionFactory(newUnauthenticatedRegistry(t)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, sess.Close()) }) @@ -218,7 +218,7 @@ func TestSessionFactory_Integration_MultipleBackends(t *testing.T) { } factory := NewSessionFactory(newUnauthenticatedRegistry(t)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, sess.Close()) }) @@ -263,11 +263,11 @@ func TestIdentityBinding_CallerRejection(t *testing.T) { } // No backend connector needed: auth validation fires before any routing. - connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, nil } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), alice, nil) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), alice, nil, nil) require.NoError(t, err) t.Cleanup(func() { _ = sess.Close() }) @@ -328,7 +328,7 @@ func TestIdentityBinding_ReadResource_And_GetPrompt_WithRealBackend(t *testing.T } factory := NewSessionFactory(newUnauthenticatedRegistry(t)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), identity, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), identity, []*vmcp.Backend{backend}, nil) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, sess.Close()) }) @@ -372,11 +372,11 @@ func TestIdentityBinding_RestoreSession_RoundTrip(t *testing.T) { }, } - connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, nil } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), identity, nil) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), identity, nil, nil) require.NoError(t, err) t.Cleanup(func() { _ = sess.Close() }) @@ -478,7 +478,7 @@ func TestSessionFactory_Integration_RestoreSession_SendsStoredSessionHintToBacke // Create the original session — the backend assigns a session ID over // streamable-HTTP and we store it in metadata. - orig, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + orig, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) t.Cleanup(func() { _ = orig.Close() }) diff --git a/pkg/vmcp/session/decorating_factory.go b/pkg/vmcp/session/decorating_factory.go index 2254634a23..729749c07c 100644 --- a/pkg/vmcp/session/decorating_factory.go +++ b/pkg/vmcp/session/decorating_factory.go @@ -52,8 +52,9 @@ func (f *decoratingMultiSessionFactory) MakeSessionWithID( id string, identity *auth.Identity, backends []*vmcp.Backend, + sink ListChangedSink, ) (MultiSession, error) { - sess, err := f.base.MakeSessionWithID(ctx, id, identity, backends) + sess, err := f.base.MakeSessionWithID(ctx, id, identity, backends, sink) if err != nil { return nil, err } diff --git a/pkg/vmcp/session/decorating_factory_test.go b/pkg/vmcp/session/decorating_factory_test.go index ae4776ed85..c1559a1d58 100644 --- a/pkg/vmcp/session/decorating_factory_test.go +++ b/pkg/vmcp/session/decorating_factory_test.go @@ -34,7 +34,7 @@ func TestNewDecoratingFactory_DecoratorsAppliedInOrder(t *testing.T) { sess := sessionmocks.NewMockMultiSession(ctrl) base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) base.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil) var order []int @@ -48,7 +48,7 @@ func TestNewDecoratingFactory_DecoratorsAppliedInOrder(t *testing.T) { } factory := session.NewDecoratingFactory(base, dec1, dec2) - _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil) + _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil, nil) require.NoError(t, err) assert.Equal(t, []int{1, 2}, order) } @@ -62,7 +62,7 @@ func TestNewDecoratingFactory_DecoratorError_ClosesSession(t *testing.T) { base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) base.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil) decErr := errors.New("decorator boom") @@ -70,7 +70,7 @@ func TestNewDecoratingFactory_DecoratorError_ClosesSession(t *testing.T) { return nil, decErr }) - _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil) + _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil, nil) require.ErrorIs(t, err, decErr) } @@ -85,7 +85,7 @@ func TestNewDecoratingFactory_SecondDecoratorError_ClosesCurrentSession(t *testi base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) base.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil) decErr := errors.New("second decorator boom") @@ -93,7 +93,7 @@ func TestNewDecoratingFactory_SecondDecoratorError_ClosesCurrentSession(t *testi dec2 := func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) { return nil, decErr } factory := session.NewDecoratingFactory(base, dec1, dec2) - _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil) + _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil, nil) require.ErrorIs(t, err, decErr) } @@ -106,14 +106,14 @@ func TestNewDecoratingFactory_NilReturnWithNoError_ClosesSession(t *testing.T) { base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) base.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil) factory := session.NewDecoratingFactory(base, func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) { return nil, nil // buggy decorator }) - _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil) + _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil, nil) require.Error(t, err) assert.Contains(t, err.Error(), "nil session") } @@ -127,7 +127,7 @@ func TestNewDecoratingFactory_CloseErrorOnDecoratorFailure_DoesNotSuppressOrigin base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) base.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil) decErr := errors.New("decorator error") @@ -135,7 +135,7 @@ func TestNewDecoratingFactory_CloseErrorOnDecoratorFailure_DoesNotSuppressOrigin return nil, decErr }) - _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil) + _, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil, nil) // The original decorator error, not the close error, is returned. require.ErrorIs(t, err, decErr) } @@ -149,14 +149,14 @@ func TestNewDecoratingFactory_HappyPath_ReturnsFinalSession(t *testing.T) { base := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) base.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(sess, nil) factory := session.NewDecoratingFactory(base, func(_ context.Context, _ session.MultiSession) (session.MultiSession, error) { return finalSess, nil }, ) - got, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil) + got, err := factory.MakeSessionWithID(context.Background(), "id", nil, nil, nil) require.NoError(t, err) assert.Equal(t, finalSess, got) } diff --git a/pkg/vmcp/session/default_session_test.go b/pkg/vmcp/session/default_session_test.go index a2848748e8..5980cfa5b9 100644 --- a/pkg/vmcp/session/default_session_test.go +++ b/pkg/vmcp/session/default_session_test.go @@ -537,7 +537,7 @@ func TestNewSessionFactory_MakeSession(t *testing.T) { } //nolint:unparam // second return is always nil by design in the success-path connector - successConnector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + successConnector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return &mockConnectedBackend{sessID: "bs-1"}, &vmcp.CapabilityList{ Tools: []vmcp.Tool{tool}, Resources: []vmcp.Resource{resource}, @@ -549,7 +549,7 @@ func TestNewSessionFactory_MakeSession(t *testing.T) { t.Parallel() factory := newSessionFactoryWithConnector(successConnector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) require.NotNil(t, sess) @@ -567,9 +567,9 @@ func TestNewSessionFactory_MakeSession(t *testing.T) { t.Parallel() factory := newSessionFactoryWithConnector(successConnector) - s1, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + s1, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) - s2, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + s2, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) assert.NotEqual(t, s1.ID(), s2.ID()) @@ -582,7 +582,7 @@ func TestNewSessionFactory_MakeSession(t *testing.T) { t.Parallel() factory := newSessionFactoryWithConnector(successConnector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, nil) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, nil, nil) require.NoError(t, err) require.NotNil(t, sess) @@ -598,7 +598,7 @@ func TestNewSessionFactory_MakeSession(t *testing.T) { factory := newSessionFactoryWithConnector(successConnector) // Mix of valid and nil entries; nil must not cause a panic. backends := []*vmcp.Backend{nil, backend, nil} - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) require.NotNil(t, sess) @@ -616,7 +616,7 @@ func TestNewSessionFactory_PartialInitialisation(t *testing.T) { {ID: "fail", Name: "fail", BaseURL: "http://fail:9999", TransportType: "streamable-http"}, } - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { if target.WorkloadID == "fail" { return nil, nil, errors.New("backend unavailable") } @@ -626,7 +626,7 @@ func TestNewSessionFactory_PartialInitialisation(t *testing.T) { } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends, nil) require.NoError(t, err, "partial init must not return an error") require.NotNil(t, sess) @@ -650,20 +650,20 @@ func TestNewSessionFactory_ConnectorReturnsNilWithoutError(t *testing.T) { }{ { name: "nil conn with nil caps", - connector: func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector: func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, nil }, }, { name: "nil conn with non-nil caps", - connector: func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector: func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, &vmcp.CapabilityList{}, nil }, }, { name: "non-nil conn with nil caps must close conn to avoid leak", wantConnClose: true, - connector: func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector: func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return &mockConnectedBackend{}, nil, nil }, }, @@ -676,8 +676,8 @@ func TestNewSessionFactory_ConnectorReturnsNilWithoutError(t *testing.T) { // Replace the connector with one that captures the mock so we can // inspect closeCalled after MakeSession returns. var captured *mockConnectedBackend - wrappedConnector := func(ctx context.Context, target *vmcp.BackendTarget, id *auth.Identity, hint string) (internalbk.Session, *vmcp.CapabilityList, error) { - conn, caps, err := tt.connector(ctx, target, id, hint) + wrappedConnector := func(ctx context.Context, target *vmcp.BackendTarget, id *auth.Identity, hint string, sink internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { + conn, caps, err := tt.connector(ctx, target, id, hint, sink) if m, ok := conn.(*mockConnectedBackend); ok { captured = m } @@ -685,7 +685,7 @@ func TestNewSessionFactory_ConnectorReturnsNilWithoutError(t *testing.T) { } factory := newSessionFactoryWithConnector(wrappedConnector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err) require.NotNil(t, sess) assert.Empty(t, sess.Tools()) @@ -707,12 +707,12 @@ func TestNewSessionFactory_ConnectorReturnsConnWithError(t *testing.T) { backend := &vmcp.Backend{ID: "b1", Name: "b1", BaseURL: "http://x:9", TransportType: "streamable-http"} leaked := &mockConnectedBackend{} - connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return leaked, nil, errors.New("init failed but conn was partially opened") } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err, "partial failure must not abort the session") require.NotNil(t, sess) assert.Empty(t, sess.Tools()) @@ -732,7 +732,7 @@ func TestNewSessionFactory_CapabilityNameConflictIsResolvedDeterministically(t * {ID: "alpha", Name: "alpha", BaseURL: "http://alpha:9", TransportType: "streamable-http"}, } - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return &mockConnectedBackend{sessID: target.WorkloadID}, &vmcp.CapabilityList{ Tools: []vmcp.Tool{{Name: "fetch", BackendID: target.WorkloadID}}, Resources: []vmcp.Resource{{URI: "file://data", BackendID: target.WorkloadID}}, @@ -741,7 +741,7 @@ func TestNewSessionFactory_CapabilityNameConflictIsResolvedDeterministically(t * } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) require.NotNil(t, sess) defer func() { require.NoError(t, sess.Close()) }() @@ -766,12 +766,12 @@ func TestNewSessionFactory_AllBackendsFail(t *testing.T) { t.Parallel() backend := &vmcp.Backend{ID: "b1", Name: "b1", BaseURL: "http://x:9", TransportType: "streamable-http"} - connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, errors.New("down") } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err, "all-fail must still return a valid (empty) session") require.NotNil(t, sess) @@ -785,7 +785,7 @@ func TestNewSessionFactory_BackendInitTimeout(t *testing.T) { backend := &vmcp.Backend{ID: "slow", Name: "slow", BaseURL: "http://x:9", TransportType: "streamable-http"} released := make(chan struct{}) - connector := func(ctx context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(ctx context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { select { case <-ctx.Done(): return nil, nil, ctx.Err() @@ -795,7 +795,7 @@ func TestNewSessionFactory_BackendInitTimeout(t *testing.T) { } factory := newSessionFactoryWithConnector(connector, WithBackendInitTimeout(50*time.Millisecond)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, []*vmcp.Backend{backend}, nil) require.NoError(t, err, "timeout is a partial failure, not a hard error") require.NotNil(t, sess) @@ -823,7 +823,7 @@ func TestNewSessionFactory_ParallelInit(t *testing.T) { var mu sync.Mutex var maxConcurrent, current int32 - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { mu.Lock() current++ if current > maxConcurrent { @@ -844,7 +844,7 @@ func TestNewSessionFactory_ParallelInit(t *testing.T) { } factory := newSessionFactoryWithConnector(connector, WithMaxBackendInitConcurrency(3)) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) // All backends must have been initialised. @@ -864,10 +864,10 @@ func TestNewSessionFactory_MakeSession_Metadata(t *testing.T) { backend2 := &vmcp.Backend{ID: "b2", Name: "backend-2", BaseURL: "http://localhost:9002", TransportType: "streamable-http"} //nolint:unparam // error return is always nil by design in the success-path connector - successConnector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + successConnector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return &mockConnectedBackend{}, &vmcp.CapabilityList{}, nil } - failConnector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + failConnector := func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, errors.New("connection refused") } @@ -928,7 +928,7 @@ func TestNewSessionFactory_MakeSession_Metadata(t *testing.T) { t.Parallel() factory := newSessionFactoryWithConnector(tt.connector) - sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), tt.identity, tt.backends) + sess, err := factory.MakeSessionWithID(context.Background(), uuid.New().String(), tt.identity, tt.backends, nil) require.NoError(t, err) require.NotNil(t, sess) defer func() { require.NoError(t, sess.Close()) }() @@ -1138,15 +1138,15 @@ func TestValidateSessionID(t *testing.T) { func TestMakeSessionWithID_InvalidIDReturnsError(t *testing.T) { t.Parallel() - f := newSessionFactoryWithConnector(func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + f := newSessionFactoryWithConnector(func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, nil }) - _, err := f.MakeSessionWithID(context.Background(), "", nil, nil) + _, err := f.MakeSessionWithID(context.Background(), "", nil, nil, nil) require.Error(t, err) assert.Contains(t, err.Error(), "must not be empty") - _, err = f.MakeSessionWithID(context.Background(), "bad id", nil, nil) + _, err = f.MakeSessionWithID(context.Background(), "bad id", nil, nil, nil) require.Error(t, err) assert.Contains(t, err.Error(), "invalid character") } diff --git a/pkg/vmcp/session/factory.go b/pkg/vmcp/session/factory.go index 14e3115a6d..3a64ea3bd7 100644 --- a/pkg/vmcp/session/factory.go +++ b/pkg/vmcp/session/factory.go @@ -9,12 +9,14 @@ import ( "context" "fmt" "log/slog" + "slices" "sort" "strings" "sync" "time" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -55,11 +57,18 @@ type MultiSessionFactory interface { // // All other behaviour (partial initialisation, bounded concurrency, etc.) // is identical to MakeSession. + // + // sink, when non-nil, is threaded to every backend connector for this + // session and invoked (best-effort, on the backend receive-loop goroutine) + // when a backend emits a notification consumed asynchronously — currently + // only notifications/tools/list_changed (#5748). Pass nil to leave that + // consumption disabled for this session. MakeSessionWithID( ctx context.Context, id string, identity *auth.Identity, backends []*vmcp.Backend, + sink ListChangedSink, ) (MultiSession, error) // RestoreSession reconstructs a live MultiSession from persisted metadata. @@ -84,7 +93,8 @@ type MultiSessionFactory interface { } // backendConnector creates a connected, initialised backend Session for use -// within a single MultiSession. It is called once per backend during MakeSession. +// within a single MultiSession. It is called once per backend that is not +// skipped as Modern (see initOneBackend) during MakeSession. // // The connector is responsible for: // 1. Creating and starting the MCP client transport. @@ -96,17 +106,27 @@ type MultiSessionFactory interface { // Mcp-Session-Id hint during Initialize so the backend can resume rather than // re-initialize. Pass an empty string for brand-new sessions. // +// sink, when non-nil, is invoked (best-effort, possibly from a receive-loop +// goroutine) whenever this backend emits a notification the connector consumes +// asynchronously (currently only notifications/tools/list_changed). Pass nil to +// leave that consumption disabled; a connector that does not support it may +// ignore the parameter entirely. See backend.ListChangedSink. +// // The returned backend.Session owns the underlying transport connection and // must be closed when the session ends. The returned CapabilityList is used // to populate the session's routing table and capability lists. // // On error the factory treats the failure as a partial failure: a warning is -// logged and the backend is excluded from the session. +// logged and the backend is excluded from the session — except when the +// backend's revision has meanwhile resolved to Modern, in which case +// initOneBackend logs at DEBUG instead (see initOneBackend's doc comment for +// why a Modern backend's session-establishment sequence fails there). type backendConnector func( ctx context.Context, target *vmcp.BackendTarget, identity *auth.Identity, sessionHint string, + sink backend.ListChangedSink, ) (backend.Session, *vmcp.CapabilityList, error) // defaultMultiSessionFactory is the production MultiSessionFactory implementation. @@ -114,6 +134,7 @@ type defaultMultiSessionFactory struct { connector backendConnector maxConcurrency int backendInitTimeout time.Duration + revisionLookup func(workloadID string) (mcpparser.Revision, bool) } // MultiSessionFactoryOption configures a defaultMultiSessionFactory. @@ -139,6 +160,32 @@ func WithBackendInitTimeout(d time.Duration) MultiSessionFactoryOption { } } +// WithRevisionLookup supplies a function that reports a backend's cached MCP +// revision (Legacy vs. Modern) by workload ID, so initOneBackend can skip +// connecting to a backend known to speak the stateless Modern (2026-07-28) +// revision. The connect attempt fails there anyway, and skipping it avoids +// both the resulting WARN and the wasted work — see initOneBackend's doc +// comment for the exact failure mechanism. +// +// The lookup's second return distinguishes "known Modern" from "unprobed" +// (including RevisionLegacy's zero value): only a confirmed Modern result +// skips the connect. A nil lookup (the default) or a cache miss reproduces +// today's behaviour exactly — every backend gets an unconditional connect +// attempt. +// +// A func value, not an interface, per the option-pattern rule in +// .claude/rules/go-style.md: it stays compile-time safe without exposing +// CachedRevision on vmcp.BackendClient for every implementation to satisfy. +// +// Concurrency: lookup is called from the per-backend init goroutines started +// by makeBaseSession, up to maxConcurrency at once, so it must be safe for +// concurrent use. +func WithRevisionLookup(lookup func(workloadID string) (mcpparser.Revision, bool)) MultiSessionFactoryOption { + return func(f *defaultMultiSessionFactory) { + f.revisionLookup = lookup + } +} + // NewSessionFactory creates a MultiSessionFactory that connects to backends // over HTTP using the given outgoing auth registry. func NewSessionFactory(registry vmcpauth.OutgoingAuthRegistry, opts ...MultiSessionFactoryOption) MultiSessionFactory { @@ -171,28 +218,101 @@ type initResult struct { // It is called from a goroutine inside MakeSession and handles all partial- // initialisation cases: connector errors, and nil conn/caps without an error. // Returns a non-nil *initResult on success, nil when the backend should be -// skipped (failure already logged as a warning). +// skipped. The second return is true when the skip is a deliberate Modern +// (2026-07-28 revision) skip rather than a failure. +// +// Why the connect is worth skipping — the sequence fails, but NOT at a Legacy +// handshake. Measured against a real stateless backend: go-sdk v1.7's Connect +// is Modern-first, so server/discover SUCCEEDS and Modern is negotiated; no +// raw Legacy initialize is ever sent. Connect then opens a +// subscriptions/listen stream, because mcpcompat's Initialize unconditionally +// installs the three list-changed notification handlers +// (mcpcompat/client.installNotificationHandlers) and go-sdk opens that stream +// whenever any of them is set. The stateless server rejects it with +// "session not found", which fails Connect, tears the connection down, and in +// turn fails the follow-up tools/list in initAndQueryCapabilities. The connector +// therefore returns (nil, nil, err). +// +// That rejection is a go-sdk v1.7.0-pre.3 artifact, NOT a spec requirement: the +// Modern revision has no sessions, so nothing in it says subscriptions/listen +// needs one. Do not read this as "Modern requires a session" — it is the SDK's +// stateless server refusing a subscribe it has no session to hang the stream on. +// +// So the wasted work is a whole sequence (discover, a failed subscribe, a +// failed tools/list, teardown), not one round-trip, and it ends in a WARN that +// looks exactly like a genuinely dead backend. Skipping is also correct on its +// own terms: all three purposes of a persistent per-backend connection are +// inapplicable to a Modern backend — list_changed propagation (Modern removed +// server-initiated notifications; its push channel is subscriptions/listen, +// which vMCP does not implement), backend-session-id resume hints (no +// Mcp-Session-Id), and identity-binding metadata. The first of those is +// precisely what makes the connect fail: vMCP asks for a notification stream +// the backend cannot give it. +// +// Known caveat (#5992): the direction that matters here — a backend redeployed +// stateless→stateful, so a cached Modern label is now wrong — self-corrects only +// on error. The Modern call fails, dispatch's isRevisionMismatch fires and +// reclassify flips the cache (pkg/vmcp/client/client.go). (The opposite +// direction, Legacy→Modern, does self-heal on success: legacyInit flips the +// cache when initialize negotiates Modern, added in #5997.) So calls are never +// wrong for long, but until some call has re-probed, this method keeps skipping +// a backend that now wants a connection, and that session holds none for it — +// list_changed propagation is lost until a NEW session is created after the +// reclassification. Closing that window needs a periodic re-probe, which is +// #5992's remaining scope, not this method's. +// +// Cold start: the revision cache is populated by the first backend call +// through dispatch (probeRevision) or, when configured, the health monitor's +// periodic check — neither is ordered before the first client session. So EVERY +// session created before the first backend call still pays the failing sequence +// above for each Modern backend, and its WARN is not suppressed, because the +// cache is still cold when the post-error re-check below runs. N clients that +// connect without calling anything each pay it; there is no per-process bound. +// Sessions created after the cache is warm skip the connect entirely. func (f *defaultMultiSessionFactory) initOneBackend( ctx context.Context, b *vmcp.Backend, identity *auth.Identity, sessionHint string, -) *initResult { + sink backend.ListChangedSink, +) (*initResult, bool) { + target := vmcp.BackendToTarget(b) + + if f.isKnownModern(target.WorkloadID) { + slog.Debug("Skipping backend connection for Modern backend; no session to hold", + "backendID", b.ID, + "backendName", b.Name, + ) + return nil, true + } + bCtx, cancel := context.WithTimeout(ctx, f.backendInitTimeout) defer cancel() - target := vmcp.BackendToTarget(b) - conn, caps, err := f.connector(bCtx, target, identity, sessionHint) + conn, caps, err := f.connector(bCtx, target, identity, sessionHint, sink) if err != nil { if conn != nil { _ = conn.Close() } + // Narrows the WARN to DEBUG when the backend has resolved to Modern + // since the pre-connect check — either an era mismatch that lost the + // cold-start race (the expected case, see this method's doc comment) or + // an unrelated genuine failure (transport, auth) on a backend that + // happens to be Modern. Both are cases where a WARN would misinform. + if f.isKnownModern(target.WorkloadID) { + slog.Debug("Backend failed to initialise and is now known Modern; skipping", + "backendID", b.ID, + "backendName", b.Name, + "error", err, + ) + return nil, true + } slog.Warn("Failed to initialise backend for session; continuing without it", "backendID", b.ID, "backendName", b.Name, "error", err, ) - return nil + return nil, false } if conn == nil || caps == nil { if conn != nil { @@ -202,9 +322,20 @@ func (f *defaultMultiSessionFactory) initOneBackend( "backendID", b.ID, "backendName", b.Name, ) - return nil + return nil, false + } + return &initResult{target: target, conn: conn, caps: caps}, false +} + +// isKnownModern reports whether workloadID's cached revision is confirmed +// Modern. Returns false for an unprobed backend or when no lookup is +// configured — indistinguishable from "attempt the connect" in either case. +func (f *defaultMultiSessionFactory) isKnownModern(workloadID string) bool { + if f.revisionLookup == nil { + return false } - return &initResult{target: target, conn: conn, caps: caps} + rev, known := f.revisionLookup(workloadID) + return known && rev == mcpparser.RevisionModern } // buildRoutingTable populates a RoutingTable and capability lists from a sorted @@ -249,11 +380,12 @@ func (f *defaultMultiSessionFactory) MakeSessionWithID( id string, identity *auth.Identity, backends []*vmcp.Backend, + sink ListChangedSink, ) (MultiSession, error) { if err := validateSessionID(id); err != nil { return nil, err } - return f.makeSession(ctx, id, identity, backends) + return f.makeSession(ctx, id, identity, backends, sink) } // validateSessionID checks that id is non-empty and contains only visible @@ -302,6 +434,7 @@ func (f *defaultMultiSessionFactory) makeBaseSession( identity *auth.Identity, backends []*vmcp.Backend, sessionHints map[string]string, + sink backend.ListChangedSink, ) *defaultMultiSession { filtered := make([]*vmcp.Backend, 0, len(backends)) for _, b := range backends { @@ -314,6 +447,7 @@ func (f *defaultMultiSessionFactory) makeBaseSession( backends = filtered rawResults := make([]*initResult, len(backends)) + modernSkipped := make([]bool, len(backends)) sem := make(chan struct{}, f.maxConcurrency) var wg sync.WaitGroup wg.Add(len(backends)) @@ -322,7 +456,7 @@ func (f *defaultMultiSessionFactory) makeBaseSession( defer wg.Done() sem <- struct{}{} defer func() { <-sem }() - rawResults[i] = f.initOneBackend(ctx, b, identity, sessionHints[b.ID]) + rawResults[i], modernSkipped[i] = f.initOneBackend(ctx, b, identity, sessionHints[b.ID], sink) }(i, b) } wg.Wait() @@ -342,8 +476,14 @@ func (f *defaultMultiSessionFactory) makeBaseSession( return results[i].target.WorkloadID < results[j].target.WorkloadID }) - if len(results) == 0 && len(backends) > 0 { - slog.Warn("All backends failed to initialise; session will have no capabilities", + // Only warn when at least one backend was actually expected to hold a + // session — a set that skipped every backend as Modern is working as + // designed, not a failure. The message deliberately does NOT say the session + // has no capabilities: on the Serve path those come from the core, not from + // these connections, so a session with zero held connections still lists and + // calls tools normally. What is actually lost is list_changed propagation. + if len(results) == 0 && len(backends) > 0 && slices.Contains(modernSkipped, false) { + slog.Warn("No backend held a session for this session; every non-Modern backend failed to initialise", "backendCount", len(backends)) } @@ -377,8 +517,9 @@ func (f *defaultMultiSessionFactory) makeSession( sessID string, identity *auth.Identity, backends []*vmcp.Backend, + sink backend.ListChangedSink, ) (MultiSession, error) { - baseSession := f.makeBaseSession(ctx, sessID, identity, backends, nil) + baseSession := f.makeBaseSession(ctx, sessID, identity, backends, nil, sink) // Apply session binding: extracts the (iss, sub) identity tuple, stores it in // session metadata under MetadataKeyIdentityBinding, and wraps the session with @@ -460,8 +601,12 @@ func (f *defaultMultiSessionFactory) RestoreSession( } // Build the base session (backend connections + routing table) without the - // security wrapper. Pass nil identity — see comment above. - baseSession := f.makeBaseSession(ctx, id, nil, filteredBackends, sessionHints) + // security wrapper. Pass nil identity — see comment above. Pass nil sink: a + // cross-pod restore has no SDK ClientSession to resync (the restoring pod may + // not be the one serving the client's connection), so tools list_changed + // propagation on this path is a follow-up (#5748 scope: the CreateSession + // path only). + baseSession := f.makeBaseSession(ctx, id, nil, filteredBackends, sessionHints, nil) // Restore only the identity-binding key from stored metadata. The other // keys (MetadataKeyBackendIDs, MetadataKeyBackendSessionPrefix.*) are diff --git a/pkg/vmcp/session/factory_metadata_test.go b/pkg/vmcp/session/factory_metadata_test.go index 4d26808cda..c2fecd91d7 100644 --- a/pkg/vmcp/session/factory_metadata_test.go +++ b/pkg/vmcp/session/factory_metadata_test.go @@ -24,7 +24,7 @@ func TestMakeSession_PersistsBackendSessionIDs(t *testing.T) { t.Run("two backends: both session IDs written to metadata", func(t *testing.T) { t.Parallel() - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { ids := map[string]string{ "backend-a": "sess-a", "backend-b": "sess-b", @@ -41,7 +41,7 @@ func TestMakeSession_PersistsBackendSessionIDs(t *testing.T) { {ID: "backend-a"}, {ID: "backend-b"}, } - sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) meta := sess.GetMetadata() @@ -56,7 +56,7 @@ func TestMakeSession_PersistsBackendSessionIDs(t *testing.T) { t.Parallel() factory := newSessionFactoryWithConnector(nilBackendConnector()) - sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, nil) + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, nil, nil) require.NoError(t, err) meta := sess.GetMetadata() @@ -72,7 +72,7 @@ func TestMakeSession_PersistsBackendSessionIDs(t *testing.T) { t.Run("partial failure: only successful backend written", func(t *testing.T) { t.Parallel() - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { if target.WorkloadID == "backend-ok" { return &mockConnectedBackend{sessID: "sess-ok"}, &vmcp.CapabilityList{}, nil } @@ -85,7 +85,7 @@ func TestMakeSession_PersistsBackendSessionIDs(t *testing.T) { {ID: "backend-ok"}, {ID: "backend-fail"}, } - sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends) + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) meta := sess.GetMetadata() @@ -103,7 +103,7 @@ func TestMakeSession_PersistsBackendSessionIDs(t *testing.T) { func TestRestoreSession_FreshlyPopulatesMetadataKeyBackendIDs(t *testing.T) { t.Parallel() - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { ids := map[string]string{ "backend-a": "sess-a", "backend-b": "sess-b", @@ -123,7 +123,7 @@ func TestRestoreSession_FreshlyPopulatesMetadataKeyBackendIDs(t *testing.T) { sessionID := "restore-test-session" // Create the initial session so we have a real token hash in metadata. - original, err := factory.MakeSessionWithID(t.Context(), sessionID, nil, backends) + original, err := factory.MakeSessionWithID(t.Context(), sessionID, nil, backends, nil) require.NoError(t, err) t.Cleanup(func() { _ = original.Close() }) @@ -181,7 +181,7 @@ func TestRestoreSession_PassesStoredSessionHintToConnector(t *testing.T) { // connector records the session hint it receives for each backend. // It always returns a stable session ID so that the original session // has predictable per-backend metadata to store. - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, sessionHint string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, sessionHint string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { mu.Lock() hintsReceived[target.WorkloadID] = sessionHint mu.Unlock() @@ -195,7 +195,7 @@ func TestRestoreSession_PassesStoredSessionHintToConnector(t *testing.T) { } // Create the original session — connector receives empty hints. - original, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends) + original, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends, nil) require.NoError(t, err) t.Cleanup(func() { _ = original.Close() }) @@ -224,6 +224,72 @@ func TestRestoreSession_PassesStoredSessionHintToConnector(t *testing.T) { "RestoreSession must pass stored backend-b session ID as hint to connector") } +// TestMakeSession_ThreadsListChangedSinkToConnector verifies (#5748) that a +// sink passed to MakeSessionWithID reaches the connector for every backend in +// the session, and that a session created WITHOUT a sink argument (the +// pre-existing 4-arg call shape) passes nil — so existing callers are +// unaffected. +func TestMakeSession_ThreadsListChangedSinkToConnector(t *testing.T) { + t.Parallel() + + t.Run("sink threaded to every backend connector", func(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + received := map[string]internalbk.ListChangedSink{} + + connector := func( + _ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, sink internalbk.ListChangedSink, + ) (internalbk.Session, *vmcp.CapabilityList, error) { + mu.Lock() + received[target.WorkloadID] = sink + mu.Unlock() + return &mockConnectedBackend{sessID: "sess-" + target.WorkloadID}, &vmcp.CapabilityList{}, nil + } + + factory := newSessionFactoryWithConnector(connector) + backends := []*vmcp.Backend{{ID: "backend-a"}, {ID: "backend-b"}} + + var sinkCalls int + sink := internalbk.ListChangedSink(func(context.Context, string, internalbk.ChangeKind) { sinkCalls++ }) + + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, backends, sink) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + mu.Lock() + defer mu.Unlock() + require.Len(t, received, 2) + for _, id := range []string{"backend-a", "backend-b"} { + require.NotNil(t, received[id], "connector for %s must receive a non-nil sink", id) + received[id](context.Background(), id, internalbk.KindTools) + } + assert.Equal(t, 2, sinkCalls, "the same sink function must reach every backend's connector call") + }) + + t.Run("no sink argument passes nil (pre-existing call shape unaffected)", func(t *testing.T) { + t.Parallel() + + var receivedSink internalbk.ListChangedSink + var sinkWasNonNil bool + connector := func( + _ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, sink internalbk.ListChangedSink, + ) (internalbk.Session, *vmcp.CapabilityList, error) { + receivedSink = sink + sinkWasNonNil = sink != nil + return &mockConnectedBackend{sessID: "sess"}, &vmcp.CapabilityList{}, nil + } + + factory := newSessionFactoryWithConnector(connector) + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, []*vmcp.Backend{{ID: "backend-a"}}, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + assert.False(t, sinkWasNonNil, "omitting the variadic sink argument must pass nil to the connector") + assert.Nil(t, receivedSink) + }) +} + // TestMakeSession_PassesEmptySessionHintToConnector verifies that MakeSession // (creating a new session, not restoring) passes an empty hint so that the // backend always creates a fresh session. @@ -233,7 +299,7 @@ func TestMakeSession_PassesEmptySessionHintToConnector(t *testing.T) { var mu sync.Mutex hintsReceived := map[string]string{} - connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, sessionHint string) (internalbk.Session, *vmcp.CapabilityList, error) { + connector := func(_ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, sessionHint string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { mu.Lock() hintsReceived[target.WorkloadID] = sessionHint mu.Unlock() @@ -241,7 +307,7 @@ func TestMakeSession_PassesEmptySessionHintToConnector(t *testing.T) { } factory := newSessionFactoryWithConnector(connector) - sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, []*vmcp.Backend{{ID: "backend-a"}}) + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, []*vmcp.Backend{{ID: "backend-a"}}, nil) require.NoError(t, err) t.Cleanup(func() { _ = sess.Close() }) diff --git a/pkg/vmcp/session/factory_revision_test.go b/pkg/vmcp/session/factory_revision_test.go new file mode 100644 index 0000000000..716bcb7b24 --- /dev/null +++ b/pkg/vmcp/session/factory_revision_test.go @@ -0,0 +1,214 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package session + +import ( + "bytes" + "context" + "errors" + "log/slog" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + internalbk "github.com/stacklok/toolhive/pkg/vmcp/session/internal/backend" +) + +// TestInitOneBackend_RevisionAwareSkip verifies that WithRevisionLookup only +// skips the connector for a backend the lookup confirms as Modern; an +// unprobed backend, a known-Legacy backend, or the absence of a lookup at all +// must reproduce today's unconditional connect. +func TestInitOneBackend_RevisionAwareSkip(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + lookup func(string) (mcpparser.Revision, bool) + wantConnectorHit bool + }{ + { + name: "no lookup configured: connector invoked", + wantConnectorHit: true, + }, + { + name: "unprobed backend (cache miss): connector invoked", + lookup: func(string) (mcpparser.Revision, bool) { return mcpparser.RevisionLegacy, false }, + wantConnectorHit: true, + }, + { + name: "known Legacy: connector invoked", + lookup: func(string) (mcpparser.Revision, bool) { return mcpparser.RevisionLegacy, true }, + wantConnectorHit: true, + }, + { + name: "known Modern: connector not invoked", + lookup: func(string) (mcpparser.Revision, bool) { return mcpparser.RevisionModern, true }, + wantConnectorHit: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var hit bool + connector := func( + _ context.Context, target *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink, + ) (internalbk.Session, *vmcp.CapabilityList, error) { + hit = true + return &mockConnectedBackend{sessID: "sess-" + target.WorkloadID}, &vmcp.CapabilityList{}, nil + } + + var opts []MultiSessionFactoryOption + if tt.lookup != nil { + opts = append(opts, WithRevisionLookup(tt.lookup)) + } + factory := newSessionFactoryWithConnector(connector, opts...) + + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, []*vmcp.Backend{{ID: "backend-a"}}, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + assert.Equal(t, tt.wantConnectorHit, hit) + }) + } +} + +// TestInitOneBackend_ColdStartRecheckAvoidsWarning pins the post-error +// re-check on initOneBackend: the backend is unprobed at the pre-connect +// check, the connect is attempted and fails, and only the POST-failure +// re-check resolves Modern. That must log at DEBUG rather than the +// real-failure WARN, and must not attempt a second connect on this call. +// +// This models the cold-start race a real Modern backend loses: its connect +// genuinely does fail (server/discover succeeds, then subscriptions/listen +// cannot be served by a session-less backend — see initOneBackend's doc +// comment), so the error branch here is the ordinary Modern path, not an +// exotic one. The fake connector returns a generic error because this test is +// about the re-check's log-level decision, not about reproducing go-sdk's +// exact failure chain. +// +// The lookup's call-counter is a plain int, not atomic: this session has a +// single backend, so a single per-backend goroutine makes both lookup calls, +// in program order, with no other goroutine reading the counter. The test +// goroutine only reads it after MakeSessionWithID returns, which happens +// after that goroutine completes (via sync.WaitGroup in makeBaseSession) — +// safe under the race detector without extra synchronization. +// +//nolint:paralleltest // setupLogRecorder swaps the global slog default logger; see its doc comment. +func TestInitOneBackend_ColdStartRecheckAvoidsWarning(t *testing.T) { + buf := setupLogRecorder(t) + + var lookupCalls int + lookup := func(string) (mcpparser.Revision, bool) { + lookupCalls++ + if lookupCalls == 1 { + return mcpparser.RevisionLegacy, false // pre-connect check: unprobed + } + return mcpparser.RevisionModern, true // post-failure re-check: resolved Modern + } + + var connectorCalls int + connector := func( + context.Context, *vmcp.BackendTarget, *auth.Identity, string, internalbk.ListChangedSink, + ) (internalbk.Session, *vmcp.CapabilityList, error) { + connectorCalls++ + return nil, nil, errors.New("dial tcp: connection refused") + } + + factory := newSessionFactoryWithConnector(connector, WithRevisionLookup(lookup)) + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, []*vmcp.Backend{{ID: "backend-a"}}, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + assert.Equal(t, 1, connectorCalls, "the connect must still be attempted when the backend is unprobed at the pre-check") + assert.NotContains(t, buf.String(), "Failed to initialise backend", + "a backend that resolves Modern on the post-failure re-check must not log the real-failure WARN") + // Positive control: without this, downgrading the real-failure WARN to DEBUG + // would leave the buffer empty and the assertion above would pass silently. + assert.Contains(t, buf.String(), "is now known Modern", + "the re-check must log the expected-skip DEBUG line, proving the buffer captures at all") +} + +// setupLogRecorder swaps the global slog default logger for a text handler +// writing to the returned buffer, for the duration of the calling test, and +// restores it on cleanup. Matches the established repo pattern (see +// pkg/vmcp/core/core_vmcp_test.go, pkg/vmcp/client/reclassify_test.go). +// +// Not compatible with t.Parallel(): slog's default logger is process-global, +// so a parallel subtest logging concurrently would race this one. Callers +// must not mark themselves (or their subtests) parallel. +func setupLogRecorder(t *testing.T) *bytes.Buffer { + t.Helper() + var buf bytes.Buffer + prev := slog.Default() + // LevelDebug so tests can assert on the DEBUG lines too, not just WARN. + // Without it a test asserting only the ABSENCE of a WARN would still pass + // if that WARN were downgraded to DEBUG — it needs a positive control. + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + return &buf +} + +// allBackendsFailedMsg pins the substring of makeBaseSession's no-session-held +// warning. Kept as the stable prefix rather than the whole line so a wording +// tweak does not silently turn the NotContains assertion below into a vacuous +// pass — the assertion that no warning fired is only meaningful while this +// substring still matches the message the code emits. +const allBackendsFailedMsg = "No backend held a session for this session" + +// TestMakeSession_AllBackendsFailedWarning verifies that the "all backends +// failed" warning fires only when a real failure is present — an all-Modern +// backend set is a deliberate, expected zero-connection outcome, not a +// failure, while mixing a Modern-skip with one genuine failure still warns. +// +//nolint:paralleltest // setupLogRecorder swaps the global slog default logger; see its doc comment. +func TestMakeSession_AllBackendsFailedWarning(t *testing.T) { + tests := []struct { + name string + lookup func(workloadID string) (mcpparser.Revision, bool) + backends []*vmcp.Backend + wantWarn bool + }{ + { + name: "all backends skipped as Modern: no warning", + lookup: func(string) (mcpparser.Revision, bool) { return mcpparser.RevisionModern, true }, + backends: []*vmcp.Backend{{ID: "backend-a"}, {ID: "backend-b"}}, + wantWarn: false, + }, + { + name: "Modern skip mixed with a genuine failure: still warns", + lookup: func(workloadID string) (mcpparser.Revision, bool) { + if workloadID == "backend-modern" { + return mcpparser.RevisionModern, true + } + return mcpparser.RevisionLegacy, false + }, + backends: []*vmcp.Backend{{ID: "backend-modern"}, {ID: "backend-fail"}}, + wantWarn: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + buf := setupLogRecorder(t) + factory := newSessionFactoryWithConnector(nilBackendConnector(), WithRevisionLookup(tt.lookup)) + + sess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), nil, tt.backends, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + if tt.wantWarn { + assert.Contains(t, buf.String(), allBackendsFailedMsg) + } else { + assert.NotContains(t, buf.String(), allBackendsFailedMsg) + } + }) + } +} diff --git a/pkg/vmcp/session/identity_binding_test.go b/pkg/vmcp/session/identity_binding_test.go index 21e41939ae..6fa8489ecd 100644 --- a/pkg/vmcp/session/identity_binding_test.go +++ b/pkg/vmcp/session/identity_binding_test.go @@ -24,7 +24,7 @@ import ( // backend to be skipped during init. This lets us exercise session-metadata // logic without real backend connections. func nilBackendConnector() backendConnector { - return func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string) (internalbk.Session, *vmcp.CapabilityList, error) { + return func(_ context.Context, _ *vmcp.BackendTarget, _ *auth.Identity, _ string, _ internalbk.ListChangedSink) (internalbk.Session, *vmcp.CapabilityList, error) { return nil, nil, nil } } @@ -79,7 +79,7 @@ func TestMakeSession_StoresIdentityBinding(t *testing.T) { factory := newSessionFactoryWithConnector(nilBackendConnector()) sess, err := factory.MakeSessionWithID( - t.Context(), uuid.New().String(), tt.identity, nil, + t.Context(), uuid.New().String(), tt.identity, nil, nil, ) require.NoError(t, err) require.NotNil(t, sess) @@ -104,7 +104,7 @@ func TestMakeSession_RejectsBoundSessionWithoutIdentifyingClaims(t *testing.T) { // Token is present but Claims are empty, so BindSession's extractBindingID fails. identity := identityWithClaims("x", map[string]any{}) - _, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), identity, nil) + _, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), identity, nil, nil) require.Error(t, err, "session creation must fail when bound identity lacks identifying claims") } @@ -226,6 +226,7 @@ func TestRestoreSession_PassesNilIdentityToConnector(t *testing.T) { _ *vmcp.BackendTarget, id *auth.Identity, _ string, + _ internalbk.ListChangedSink, ) (internalbk.Session, *vmcp.CapabilityList, error) { connectorCalled = true capturedIdentity = id @@ -239,7 +240,7 @@ func TestRestoreSession_PassesNilIdentityToConnector(t *testing.T) { }) factory := newSessionFactoryWithConnector(capturingConnector) - multiSess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), originalIdentity, nil) + multiSess, err := factory.MakeSessionWithID(t.Context(), uuid.New().String(), originalIdentity, nil, nil) require.NoError(t, err) t.Cleanup(func() { _ = multiSess.Close() }) diff --git a/pkg/vmcp/session/internal/backend/mcp_session.go b/pkg/vmcp/session/internal/backend/mcp_session.go index 9c3ef2cee1..9a9de722c3 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session.go +++ b/pkg/vmcp/session/internal/backend/mcp_session.go @@ -12,18 +12,21 @@ import ( "net/http" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - mcptransport "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" - + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/secrets" + "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/versions" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" "github.com/stacklok/toolhive/pkg/vmcp/conversion" "github.com/stacklok/toolhive/pkg/vmcp/headerforward" + "github.com/stacklok/toolhive/pkg/vmcp/internal/pagination" ) const ( @@ -38,6 +41,75 @@ const ( defaultBackendRequestTimeout = 30 * time.Second ) +// ChangeKind identifies which capability class a backend reported changed via +// ListChangedSink. Using a typed constant (rather than a bare string) means a +// typo is a compile error at the producer and consumer instead of a silent +// no-op. +type ChangeKind string + +const ( + // KindTools is reported when a backend emits notifications/tools/list_changed. + KindTools ChangeKind = "tools" + + // KindResources is reported when a backend emits + // notifications/resources/list_changed. Per MCP 2025-11-25 there is no + // separate wire method for resource TEMPLATE changes, so this kind also + // covers a resync of the backend's resource templates (see + // resyncSessionResources in pkg/vmcp/server). + KindResources ChangeKind = "resources" + + // KindPrompts is reported when a backend emits + // notifications/prompts/list_changed. + KindPrompts ChangeKind = "prompts" +) + +// ListChangedSink is invoked when a persistent backend connection observes a +// notification this package consumes asynchronously (outside any in-flight +// call): notifications/tools/list_changed (kind=KindTools), +// notifications/resources/list_changed (kind=KindResources, which also covers +// resource templates — MCP 2025-11-25 has no separate wire method for those), +// and notifications/prompts/list_changed (kind=KindPrompts). +// +// The sink is invoked on the mcpcompat client's receive-loop goroutine (see +// Client.dispatch in toolhive-core), so implementations MUST be non-blocking: +// they must hand the work off (e.g. set a dirty flag / signal a worker) and +// return immediately. A sink that does real work (network I/O, cache purges) +// inline stalls that backend's entire notification delivery and lets a +// misbehaving backend amplify one notification into unbounded work. The ctx +// passed here is only for the hand-off; long-lived resync work must run under +// a caller-owned, cancellable context, not this one. +type ListChangedSink func(ctx context.Context, backendWorkloadID string, kind ChangeKind) + +// newListChangedNotificationHandler builds the OnNotification handler +// registered by createMCPClient when sink is non-nil. It dispatches +// notifications/tools/list_changed (kind=KindTools), +// notifications/resources/list_changed (kind=KindResources, also covering +// resource templates), and notifications/prompts/list_changed +// (kind=KindPrompts) to sink, and log-only handles notifications/message; +// every other notification method is ignored — this handler is not the +// mid-call server->client relay (that is pkg/vmcp/client's per-call +// notification forwarder), it only feeds the session-registration resync path. +func newListChangedNotificationHandler(workloadID string, sink ListChangedSink) func(mcp.JSONRPCNotification) { + return func(n mcp.JSONRPCNotification) { + switch n.Method { + case vmcp.MethodToolsListChangedNotification: + sink(context.Background(), workloadID, KindTools) + case vmcp.MethodResourcesListChangedNotification: + sink(context.Background(), workloadID, KindResources) + case vmcp.MethodPromptsListChangedNotification: + sink(context.Background(), workloadID, KindPrompts) + case vmcp.MethodLogNotification: + // Out-of-call backend log messages are not relayed to the downstream + // client on this path; log so the signal is visible rather than + // silently dropped. + slog.Debug("backend log notification received outside call", "backendID", workloadID) + default: + // Other notification types (resource subscription updates, ...) are + // out of scope here. + } + } +} + // httpRoundTripperFunc adapts a plain function to http.RoundTripper. type httpRoundTripperFunc func(*http.Request) (*http.Response, error) @@ -98,7 +170,7 @@ func (i *identityRoundTripper) RoundTrip(req *http.Request) (*http.Response, err // Compile-time assertion: mcpSession must implement Session. var _ Session = (*mcpSession)(nil) -// mcpSession wraps a persistent mark3labs MCP client for one backend. +// mcpSession wraps a persistent mcpcompat MCP client for one backend. // It is created once per backend during MakeSession and closed when the session ends. // // Phase 1 limitation — no reconnection: if the underlying transport drops @@ -131,11 +203,18 @@ func (c *mcpSession) CallTool( slog.Debug("Translating tool name", "clientName", toolName, "backendName", backendName) } + // Strip the reserved io.modelcontextprotocol/* keys before forwarding to this + // Legacy (session-based, stateful) backend: a downstream Modern caller's + // _meta.protocolVersion is only valid on a stateless Modern hop, and go-sdk + // v1.7 hard-rejects ANY _meta.protocolVersion on a stateful server (HTTP 400) + // regardless of its value — see mcpparser.StripReservedMeta. conversion.ToMCPMeta + // below strips too; this stays as the guard at the site that knows the hazard. + meta = mcpparser.StripReservedMeta(meta) result, err := c.client.CallTool(ctx, mcp.CallToolRequest{ Params: mcp.CallToolParams{ Name: backendName, Arguments: arguments, - Meta: conversion.ToMCPMeta(meta), + Meta: conversion.ToMCPMeta(telemetry.MetaWithTraceContext(ctx, meta)), }, }) if err != nil { @@ -168,8 +247,14 @@ func (c *mcpSession) ReadResource( slog.Debug("Translating resource URI", "clientURI", uri, "backendURI", backendURI) } + // Forward-compat: mcpcompat drops Params.Meta on this path today (no-op on + // the wire). See pkg/vmcp/client's TestOutboundMetaTraceContext for + // details/tripwire. result, err := c.client.ReadResource(ctx, mcp.ReadResourceRequest{ - Params: mcp.ReadResourceParams{URI: backendURI}, + Params: mcp.ReadResourceParams{ + URI: backendURI, + Meta: conversion.ToMCPMeta(telemetry.MetaWithTraceContext(ctx, nil)), + }, }) if err != nil { return nil, fmt.Errorf("resource %q read failed on backend %s: %w", uri, c.target.WorkloadID, err) @@ -194,10 +279,14 @@ func (c *mcpSession) GetPrompt( stringArgs := conversion.ConvertPromptArguments(arguments) + // Forward-compat: mcpcompat drops Params.Meta on this path today (no-op on + // the wire). See pkg/vmcp/client's TestOutboundMetaTraceContext for + // details/tripwire. result, err := c.client.GetPrompt(ctx, mcp.GetPromptRequest{ Params: mcp.GetPromptParams{ Name: backendName, Arguments: stringArgs, + Meta: conversion.ToMCPMeta(telemetry.MetaWithTraceContext(ctx, nil)), }, }) if err != nil { @@ -221,11 +310,18 @@ func (c *mcpSession) GetPrompt( // shared across every session it creates; its lifetime matches the connector's. // It is consumed by BuildHeaderForwardTripper to resolve secret-backed entries // in target.HeaderForward. +// +// The returned function's sink parameter, when non-nil, enables persistent +// backend-notification consumption for this backend connection — see +// createMCPClient for what that does and does not enable (nil-sink callers are +// completely unaffected: no OnNotification handler is registered and no +// standalone GET stream is opened). func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( ctx context.Context, target *vmcp.BackendTarget, identity *auth.Identity, sessionHint string, + sink ListChangedSink, ) (Session, *vmcp.CapabilityList, error) { provider := secrets.NewEnvironmentProvider() return func( @@ -233,8 +329,9 @@ func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( target *vmcp.BackendTarget, identity *auth.Identity, sessionHint string, + sink ListChangedSink, ) (Session, *vmcp.CapabilityList, error) { - c, err := createMCPClient(ctx, target, identity, registry, sessionHint, provider) + c, err := createMCPClient(ctx, target, identity, registry, sessionHint, provider, sink) if err != nil { return nil, nil, fmt.Errorf("failed to create MCP client for backend %s: %w", target.WorkloadID, err) } @@ -247,7 +344,7 @@ func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( // Extract the backend-assigned session ID when the transport supports it. // Streamable-HTTP servers send an Mcp-Session-Id response header during - // Initialize; the mark3labs transport captures it internally and exposes + // Initialize; the mcpcompat transport captures it internally and exposes // it via GetSessionId(). SSE transports do not assign a session ID, so // the field remains empty for those backends. var backendSessionID string @@ -259,7 +356,7 @@ func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( } } -// createMCPClient builds and starts a mark3labs MCP client for target. +// createMCPClient builds and starts a mcpcompat MCP client for target. // The transport is started with context.Background() so its lifetime is bound // to client.Close(), not to any caller-supplied init context. // sessionHint, when non-empty, is passed as the initial Mcp-Session-Id for @@ -268,6 +365,19 @@ func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( // ctx is used only to resolve secret-backed entries in target.HeaderForward at // client-creation time; the transport itself is started with context.Background() // as described above. provider supplies values for those secret-backed headers. +// +// sink is gated STRICTLY on non-nil: when nil, this function's behavior is +// byte-for-byte identical to before sink existed (no OnNotification handler is +// registered, and the streamable-HTTP branch does not enable continuous +// listening). When non-nil, an OnNotification handler is registered (before +// c.Start, per the mcpcompat client's registration contract) that invokes sink +// on notifications/tools/list_changed, notifications/resources/list_changed +// (which also covers resource templates), and notifications/prompts/list_changed, +// and log-only handles notifications/message. For the streamable-HTTP transport, +// a non-nil sink also enables mcptransport.WithContinuousListening() — required +// for the backend's asynchronous (outside any in-flight call) notifications to +// reach this client at all — since some backends hang when a standalone GET +// stream is opened against them, this must stay opt-in (#5748 R3). func createMCPClient( ctx context.Context, target *vmcp.BackendTarget, @@ -275,6 +385,7 @@ func createMCPClient( registry vmcpauth.OutgoingAuthRegistry, sessionHint string, provider secrets.Provider, + sink ListChangedSink, ) (*mcpclient.Client, error) { // Resolve and validate the auth strategy once at client creation time. strategyName := authtypes.StrategyTypeUnauthenticated @@ -336,7 +447,7 @@ func createMCPClient( // request/response pair, so a per-response body size limit is safe and // correct. http.Client.Timeout provides a hard wall-clock deadline; // WithHTTPTimeout additionally wraps each SDK request in a - // context.WithTimeout so the mark3labs transport surfaces a descriptive + // context.WithTimeout so the mcpcompat transport surfaces a descriptive // error before the stdlib deadline fires. Both are set to // defaultBackendRequestTimeout: defense-in-depth. sizeLimited := httpRoundTripperFunc(func(req *http.Request) (*http.Response, error) { @@ -353,9 +464,13 @@ func createMCPClient( } return resp, nil }) + // CheckRedirect: keep in sync with pkg/vmcp/client.newBackendHTTPClient — a + // cross-host 30x would otherwise re-inject (via the RoundTripper chain) the + // auth/identity/header-forward credentials at the attacker host. httpClient := &http.Client{ - Transport: sizeLimited, - Timeout: defaultBackendRequestTimeout, + Transport: sizeLimited, + Timeout: defaultBackendRequestTimeout, + CheckRedirect: networking.SameHostRedirectPolicy(), } streamableOpts := []mcptransport.StreamableHTTPCOption{ mcptransport.WithHTTPTimeout(defaultBackendRequestTimeout), @@ -364,6 +479,14 @@ func createMCPClient( if sessionHint != "" { streamableOpts = append(streamableOpts, mcptransport.WithSession(sessionHint)) } + if sink != nil { + // A standalone GET stream is the only way this client can receive a + // notification the backend emits outside an in-flight call (e.g. + // tools/list_changed after a registration-time capability change). + // Gated strictly on sink != nil: some backends hang when this stream + // is opened against them (#5748 R3), so it must stay opt-in. + streamableOpts = append(streamableOpts, mcptransport.WithContinuousListening()) + } c, err = mcpclient.NewStreamableHttpClient(target.BaseURL, streamableOpts...) case "sse": // For SSE, the entire session is delivered as one long-lived HTTP @@ -374,7 +497,9 @@ func createMCPClient( // // http.Client.Timeout is also omitted: it caps the full round-trip // including body reads, which would kill the stream after the timeout. - httpClient := &http.Client{Transport: base} + // CheckRedirect: keep in sync with pkg/vmcp/client.newBackendHTTPClient (see + // the streamable case above). No Timeout: long-lived SSE stream. + httpClient := &http.Client{Transport: base, CheckRedirect: networking.SameHostRedirectPolicy()} c, err = mcpclient.NewSSEMCPClient( target.BaseURL, mcptransport.WithHTTPClient(httpClient), @@ -387,6 +512,18 @@ func createMCPClient( return nil, fmt.Errorf("failed to create %s client: %w", target.TransportType, err) } + // Register the notification handler before Start (mirrors the per-call + // client in pkg/vmcp/client, and matches the mcpcompat client's own + // registration contract): the handler dispatches + // notifications/tools/list_changed, notifications/resources/list_changed, + // and notifications/prompts/list_changed to sink and logs + // notifications/message. Strictly gated on sink != nil so a nil-sink caller + // registers no handler at all — this function's behavior for that caller is + // unchanged. + if sink != nil { + c.OnNotification(newListChangedNotificationHandler(target.WorkloadID, sink)) + } + // Start the transport with context.Background() so that the transport's // lifetime is scoped to the session (terminated by client.Close()) rather // than to the per-backend init timeout context. The init timeout context @@ -425,85 +562,27 @@ func initAndQueryCapabilities( caps := &vmcp.CapabilityList{} if serverCaps.Tools != nil { - toolsResult, listErr := c.ListTools(ctx, mcp.ListToolsRequest{}) - if listErr != nil { - return nil, fmt.Errorf("list tools failed: %w", listErr) - } - for _, t := range toolsResult.Tools { - caps.Tools = append(caps.Tools, vmcp.Tool{ - Name: t.Name, - Description: t.Description, - InputSchema: conversion.ConvertToolInputSchema(t.InputSchema), - OutputSchema: conversion.ConvertToolOutputSchema(t.OutputSchema), - Annotations: conversion.ConvertToolAnnotations(t.Annotations), - BackendID: target.WorkloadID, - }) + tools, err := queryBackendTools(ctx, c, target) + if err != nil { + return nil, err } + caps.Tools = tools } if serverCaps.Resources != nil { - resResult, listErr := c.ListResources(ctx, mcp.ListResourcesRequest{}) - switch { - case errors.Is(listErr, mcp.ErrMethodNotFound): - // Tolerate JSON-RPC -32601 here so a backend that advertises the - // resources capability but does not implement resources/list (e.g. - // Atlassian Rovo, see #5231) still contributes its tools instead of - // being dropped. HTTP-level method absence is intentionally fatal. - slog.Warn("backend advertised resources capability but does not implement resources/list", - "backendID", target.WorkloadID, - "name", target.WorkloadName, - "baseURL", target.BaseURL, - "method", "resources/list", - ) - case listErr != nil: - return nil, fmt.Errorf("list resources failed: %w", listErr) - default: - for _, r := range resResult.Resources { - caps.Resources = append(caps.Resources, vmcp.Resource{ - URI: r.URI, - Name: r.Name, - Description: r.Description, - MimeType: r.MIMEType, - BackendID: target.WorkloadID, - }) - } + resources, err := queryBackendResources(ctx, c, target) + if err != nil { + return nil, err } + caps.Resources = resources } if serverCaps.Prompts != nil { - promptsResult, listErr := c.ListPrompts(ctx, mcp.ListPromptsRequest{}) - switch { - case errors.Is(listErr, mcp.ErrMethodNotFound): - // Tolerate JSON-RPC -32601 here so a backend that advertises the - // prompts capability but does not implement prompts/list (e.g. - // Atlassian Rovo, see #5231) still contributes its tools instead of - // being dropped. HTTP-level method absence is intentionally fatal. - slog.Warn("backend advertised prompts capability but does not implement prompts/list", - "backendID", target.WorkloadID, - "name", target.WorkloadName, - "baseURL", target.BaseURL, - "method", "prompts/list", - ) - case listErr != nil: - return nil, fmt.Errorf("list prompts failed: %w", listErr) - default: - for _, p := range promptsResult.Prompts { - args := make([]vmcp.PromptArgument, len(p.Arguments)) - for j, a := range p.Arguments { - args[j] = vmcp.PromptArgument{ - Name: a.Name, - Description: a.Description, - Required: a.Required, - } - } - caps.Prompts = append(caps.Prompts, vmcp.Prompt{ - Name: p.Name, - Description: p.Description, - Arguments: args, - BackendID: target.WorkloadID, - }) - } + prompts, err := queryBackendPrompts(ctx, c, target) + if err != nil { + return nil, err } + caps.Prompts = prompts } slog.Debug("Backend capabilities", @@ -516,6 +595,113 @@ func initAndQueryCapabilities( return caps, nil } +// queryBackendTools lists the backend's tools, following MCP pagination cursors so a +// backend that paginates (mcpcompat paginates at DefaultPageSize=1000) contributes its +// complete tool set rather than only the first page. +func queryBackendTools(ctx context.Context, c *mcpclient.Client, target *vmcp.BackendTarget) ([]vmcp.Tool, error) { + tools, err := pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Tool, mcp.Cursor, error) { + req := mcp.ListToolsRequest{} + req.Params.Cursor = cursor + result, err := c.ListTools(ctx, req) + if err != nil { + return nil, "", err + } + return result.Tools, result.NextCursor, nil + }) + if err != nil { + return nil, fmt.Errorf("list tools failed: %w", err) + } + out := make([]vmcp.Tool, 0, len(tools)) + for _, t := range tools { + out = append(out, vmcp.Tool{ + Name: t.Name, + Description: t.Description, + InputSchema: conversion.ConvertToolInputSchema(t.InputSchema), + OutputSchema: conversion.ConvertToolOutputSchema(t.OutputSchema), + Annotations: conversion.ConvertToolAnnotations(t.Annotations), + BackendID: target.WorkloadID, + }) + } + return out, nil +} + +// queryBackendResources lists the backend's resources with cursor-following pagination. +// A backend that advertises the resources capability but does not implement +// resources/list (JSON-RPC -32601, e.g. Atlassian Rovo, see #5231) is tolerated: it +// returns no resources instead of failing the whole discovery. HTTP-level method absence +// remains fatal. +func queryBackendResources(ctx context.Context, c *mcpclient.Client, target *vmcp.BackendTarget) ([]vmcp.Resource, error) { + resources, err := pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Resource, mcp.Cursor, error) { + req := mcp.ListResourcesRequest{} + req.Params.Cursor = cursor + result, err := c.ListResources(ctx, req) + if err != nil { + return nil, "", err + } + return result.Resources, result.NextCursor, nil + }) + if errors.Is(err, mcp.ErrMethodNotFound) { + slog.Warn("backend advertised resources capability but does not implement resources/list", + "backendID", target.WorkloadID, "name", target.WorkloadName, "baseURL", target.BaseURL, "method", "resources/list") + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("list resources failed: %w", err) + } + out := make([]vmcp.Resource, 0, len(resources)) + for _, r := range resources { + out = append(out, vmcp.Resource{ + URI: r.URI, + Name: r.Name, + Description: r.Description, + MimeType: r.MIMEType, + BackendID: target.WorkloadID, + }) + } + return out, nil +} + +// queryBackendPrompts lists the backend's prompts with cursor-following pagination. +// Like queryBackendResources, it tolerates a -32601 from a backend that advertises the +// prompts capability without implementing prompts/list. +func queryBackendPrompts(ctx context.Context, c *mcpclient.Client, target *vmcp.BackendTarget) ([]vmcp.Prompt, error) { + prompts, err := pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Prompt, mcp.Cursor, error) { + req := mcp.ListPromptsRequest{} + req.Params.Cursor = cursor + result, err := c.ListPrompts(ctx, req) + if err != nil { + return nil, "", err + } + return result.Prompts, result.NextCursor, nil + }) + if errors.Is(err, mcp.ErrMethodNotFound) { + slog.Warn("backend advertised prompts capability but does not implement prompts/list", + "backendID", target.WorkloadID, "name", target.WorkloadName, "baseURL", target.BaseURL, "method", "prompts/list") + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("list prompts failed: %w", err) + } + out := make([]vmcp.Prompt, 0, len(prompts)) + for _, p := range prompts { + args := make([]vmcp.PromptArgument, len(p.Arguments)) + for j, a := range p.Arguments { + args[j] = vmcp.PromptArgument{ + Name: a.Name, + Description: a.Description, + Required: a.Required, + } + } + out = append(out, vmcp.Prompt{ + Name: p.Name, + Description: p.Description, + Arguments: args, + BackendID: target.WorkloadID, + }) + } + return out, nil +} + // mergeForwardedHeaders returns a HeaderForwardConfig that combines the static // backend configuration (base) with any per-request forwarded headers captured // from the caller's request (see headerforward.CaptureMiddleware). diff --git a/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go b/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go index f1c16f8e80..5832f02664 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go @@ -6,18 +6,20 @@ package backend import ( "context" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" + "strconv" "sync" "testing" - mcpclient "github.com/mark3labs/mcp-go/client" - mcptransport "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -49,6 +51,12 @@ type fakeBackend struct { resources []mcp.Resource prompts []mcp.Prompt + // toolsPageSize, when > 0, makes tools/list paginate: each response returns + // at most toolsPageSize tools plus a NextCursor until the set is exhausted. + // The cursor is the decimal offset of the next page. Used to verify the + // backend connector follows pagination cursors (#5844). + toolsPageSize int + // methodCalls counts how many times each method was invoked, so tests can // assert that (e.g.) tools/list was reached after a recoverable // resources/list failure. @@ -59,6 +67,11 @@ type fakeBackend struct { // method name. Tests asserting transport-chain behavior (e.g. HeaderForward) // use headersFor(method) to inspect the headers a backend actually saw. headersByMethod map[string]http.Header + + // metaByMethod records the raw params._meta keyed by JSON-RPC method name. + // Tests asserting outbound W3C trace context propagation (SEP-414) use + // toolCallMeta() to inspect the _meta a backend actually saw on the wire. + metaByMethod map[string]json.RawMessage } type jsonRPCError struct { @@ -74,6 +87,7 @@ func newFakeBackend(t *testing.T, fb *fakeBackend) string { fb.t = t fb.methodCalls = make(map[string]int) fb.headersByMethod = make(map[string]http.Header) + fb.metaByMethod = make(map[string]json.RawMessage) mux := http.NewServeMux() mux.HandleFunc("/mcp", fb.handle) @@ -108,6 +122,16 @@ func (f *fakeBackend) headersFor(method string) http.Header { return h.Clone() } +// toolCallMeta returns the raw params._meta recorded for the most recent +// tools/call request, or nil if no such request was seen or it carried no +// _meta field. metaByMethod is keyed by method, but every test so far only +// inspects tools/call; add a sibling accessor if another method needs one. +func (f *fakeBackend) toolCallMeta() json.RawMessage { + f.mu.Lock() + defer f.mu.Unlock() + return f.metaByMethod[string(mcp.MethodToolsCall)] +} + // handle implements the JSON-RPC subset needed for backend init. The // streamable-HTTP transport sends POST requests with Accept: // application/json, text/event-stream — we always reply with @@ -133,6 +157,10 @@ func (f *fakeBackend) handle(w http.ResponseWriter, r *http.Request) { JSONRPC string `json:"jsonrpc"` ID json.RawMessage `json:"id"` Method string `json:"method"` + Params struct { + Cursor string `json:"cursor"` + Meta json.RawMessage `json:"_meta"` + } `json:"params"` } if err := json.Unmarshal(body, &msg); err != nil { f.t.Errorf("fakeBackend: decode: %v body=%s", err, string(body)) @@ -143,6 +171,7 @@ func (f *fakeBackend) handle(w http.ResponseWriter, r *http.Request) { f.mu.Lock() f.methodCalls[msg.Method]++ f.headersByMethod[msg.Method] = r.Header.Clone() + f.metaByMethod[msg.Method] = msg.Params.Meta f.mu.Unlock() // Notifications (no id, e.g. notifications/initialized) get an empty 202. @@ -161,6 +190,10 @@ func (f *fakeBackend) handle(w http.ResponseWriter, r *http.Request) { f.writeError(w, msg.ID, f.listToolsErr) return } + if f.toolsPageSize > 0 { + f.writeToolsPage(w, msg.ID, msg.Params.Cursor) + return + } f.writeResult(w, msg.ID, map[string]any{"tools": f.tools}) case string(mcp.MethodResourcesList): if f.listResourcesErr != nil { @@ -207,6 +240,30 @@ func (f *fakeBackend) writeInitializeResult(w http.ResponseWriter, id json.RawMe }) } +// writeToolsPage returns one page of f.tools starting at the offset encoded in +// cursor (empty = first page), plus a NextCursor when more tools remain. The +// cursor is the decimal offset of the next page. +func (f *fakeBackend) writeToolsPage(w http.ResponseWriter, id json.RawMessage, cursor string) { + start := 0 + if cursor != "" { + n, err := strconv.Atoi(cursor) + if err != nil { + f.writeError(w, id, &jsonRPCError{code: mcp.INVALID_PARAMS, message: "bad cursor"}) + return + } + start = n + } + end := start + f.toolsPageSize + if end > len(f.tools) { + end = len(f.tools) + } + result := map[string]any{"tools": f.tools[start:end]} + if end < len(f.tools) { + result["nextCursor"] = strconv.Itoa(end) + } + f.writeResult(w, id, result) +} + func (f *fakeBackend) writeResult(w http.ResponseWriter, id json.RawMessage, result any) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) @@ -234,12 +291,20 @@ func (f *fakeBackend) writeError(w http.ResponseWriter, id json.RawMessage, e *j } } -// newTestClient builds a streamable-HTTP mark3labs client pointing at url and +// newTestClient builds a streamable-HTTP mcpcompat client pointing at url and // runs Start() so the transport is ready for initAndQueryCapabilities. Cleanup // is registered via t.Cleanup. func newTestClient(t *testing.T, url string) *mcpclient.Client { t.Helper() - c, err := mcpclient.NewStreamableHttpClient(url, mcptransport.WithHTTPBasicClient(&http.Client{})) + // Each test client gets its own transport with keep-alive disabled. A bare + // &http.Client{} shares the process-global http.DefaultTransport, so across + // t.Parallel() subtests one client's teardown (t.Cleanup -> Close) can close + // idle connections a sibling is mid-request on. Go surfaces that as + // "connection broken: http: CloseIdleConnections called", which the + // streamable transport reports as a spurious 4xx/legacy-SSE init failure. + c, err := mcpclient.NewStreamableHttpClient(url, mcptransport.WithHTTPBasicClient(&http.Client{ + Transport: &http.Transport{DisableKeepAlives: true}, + })) require.NoError(t, err) require.NoError(t, c.Start(context.Background())) t.Cleanup(func() { _ = c.Close() }) @@ -478,3 +543,52 @@ func TestInitAndQueryCapabilities_FatalErrors(t *testing.T) { }) } } + +// TestInitAndQueryCapabilities_FollowsToolPagination verifies the per-session +// backend connector follows MCP list-pagination cursors (#5844): a backend that +// returns its tools across multiple pages contributes its COMPLETE set, not just +// the first page. Before the fix the connector issued a single tools/list and +// silently dropped every tool beyond the first page. +func TestInitAndQueryCapabilities_FollowsToolPagination(t *testing.T) { + t.Parallel() + + const ( + total = 250 + pageSize = 100 // 3 pages: 100 + 100 + 50 + ) + tools := make([]mcp.Tool, total) + for i := range tools { + tools[i] = mcp.Tool{Name: fmt.Sprintf("tool_%d", i)} + } + fb := &fakeBackend{ + advertiseTools: true, + tools: tools, + toolsPageSize: pageSize, + } + url := newFakeBackend(t, fb) + c := newTestClient(t, url) + target := &vmcp.BackendTarget{ + WorkloadID: "fake-backend", + WorkloadName: "fake-backend", + BaseURL: url, + TransportType: "streamable-http", + } + + caps, err := initAndQueryCapabilities(t.Context(), c, target) + require.NoError(t, err) + require.NotNil(t, caps) + + // Every page accumulated, with no gaps or duplicates across page boundaries. + require.Len(t, caps.Tools, total, "connector must accumulate every page, not just the first") + names := make(map[string]struct{}, total) + for _, tl := range caps.Tools { + names[tl.Name] = struct{}{} + } + assert.Len(t, names, total, "tool names must be distinct across pages (no overlap/duplication)") + assert.Contains(t, names, "tool_0", "first page present") + assert.Contains(t, names, fmt.Sprintf("tool_%d", total-1), "last page present") + + // The cursor loop must issue ceil(total/pageSize) tools/list calls — proof + // pagination actually happened rather than one oversized page. + assert.Equal(t, 3, fb.callCount(string(mcp.MethodToolsList)), "expected one tools/list per page") +} diff --git a/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go b/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go index 9d58de9b76..650c1eba00 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go @@ -9,10 +9,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -112,7 +112,7 @@ func TestHTTPSession_AppliesHeaderForwardToPostInitializeRequests(t *testing.T) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - sess, _, err := connector(ctx, target, nil, "") + sess, _, err := connector(ctx, target, nil, "", nil) require.NoError(t, err) t.Cleanup(func() { _ = sess.Close() }) @@ -155,7 +155,7 @@ func TestHTTPSession_AppliesHeaderForwardToPostInitializeRequests(t *testing.T) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - _, _, err := connector(ctx, target, nil, "") + _, _, err := connector(ctx, target, nil, "", nil) require.Error(t, err) assert.Contains(t, err.Error(), "restricted", "connector must reject restricted header names in HeaderForward config") @@ -216,7 +216,7 @@ func connectAndCallEcho(t *testing.T, target *vmcp.BackendTarget) Session { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) t.Cleanup(cancel) - sess, caps, err := connector(ctx, target, nil, "") + sess, caps, err := connector(ctx, target, nil, "", nil) require.NoError(t, err, "connector must initialise the backend successfully") require.NotNil(t, sess, "connector returned nil session") require.NotNil(t, caps, "connector returned nil capability list") diff --git a/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go b/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go index ac0e8236a6..acbafc907f 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go @@ -8,10 +8,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -85,7 +85,7 @@ func TestHTTPSession_PerRequestIdentity_DrivesUpstreamAuthHeader(t *testing.T) { initCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) t.Cleanup(cancel) - sess, _, err := connector(initCtx, target, staleIdentity, "") + sess, _, err := connector(initCtx, target, staleIdentity, "", nil) require.NoError(t, err, "connector must initialise the backend successfully") t.Cleanup(func() { _ = sess.Close() }) @@ -161,7 +161,7 @@ func TestHTTPSession_FallbackIdentity_UsedWhenContextHasNoIdentity(t *testing.T) initCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) t.Cleanup(cancel) - sess, _, err := connector(initCtx, target, fallbackIdentity, "") + sess, _, err := connector(initCtx, target, fallbackIdentity, "", nil) require.NoError(t, err, "connector must initialise the backend successfully") t.Cleanup(func() { _ = sess.Close() }) diff --git a/pkg/vmcp/session/internal/backend/mcp_session_meta_propagation_test.go b/pkg/vmcp/session/internal/backend/mcp_session_meta_propagation_test.go new file mode 100644 index 0000000000..571b109346 --- /dev/null +++ b/pkg/vmcp/session/internal/backend/mcp_session_meta_propagation_test.go @@ -0,0 +1,196 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package backend + +import ( + "context" + "encoding/json" + "maps" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/propagation" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestHTTPSession_CallTool_InjectsTraceparent asserts that the session-backed +// connector (pkg/vmcp/session/internal/backend) injects the current W3C trace +// context (traceparent) into outbound params._meta (SEP-414) for CallTool, and +// that it is omitted entirely when there is no active span. This mirrors the +// coverage in pkg/vmcp/client's TestOutboundMetaTraceContext for the +// session-side connector. +// +// ReadResource and GetPrompt are not covered here for the same reason +// documented on pkg/vmcp/client's TestOutboundMetaTraceContext: mcpcompat's +// non-resume client does not forward Params.Meta to the wire for those two +// operations (only CallTool does), so there is nothing to observe on fakeBackend +// today. See mcp_session.go's ReadResource/GetPrompt for the code that will +// start working once mcpcompat closes that gap. +// +// This test mutates the global OTEL propagator, so it must NOT run in +// parallel with the rest of this package's tests. +func TestHTTPSession_CallTool_InjectsTraceparent(t *testing.T) { //nolint:paralleltest // Mutates global OTEL propagator + oldPropagator := otel.GetTextMapPropagator() + otel.SetTextMapPropagator(propagation.TraceContext{}) + defer otel.SetTextMapPropagator(oldPropagator) + + fb := &fakeBackend{advertiseTools: true, tools: []mcp.Tool{{Name: "echo"}}} + url := newFakeBackend(t, fb) + + target := &vmcp.BackendTarget{ + WorkloadID: "trace-context-backend", + WorkloadName: "trace-context-backend", + BaseURL: url, + TransportType: "streamable-http", + } + + registry := newTestRegistry(t) + connector := NewHTTPConnector(registry) + + initCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + sess, _, err := connector(initCtx, target, nil, "", nil) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + tp := sdktrace.NewTracerProvider() + defer func() { _ = tp.Shutdown(context.Background()) }() + tracer := tp.Tracer("test") + + //nolint:paralleltest // sequential: shares the fakeBackend/session with the sibling subtests below + t.Run("active span: traceparent on the wire carries this span's TraceID", func(t *testing.T) { + spanCtx, span := tracer.Start(context.Background(), "test-span") + defer span.End() + + _, err := sess.CallTool(spanCtx, "echo", map[string]any{}, nil) + require.NoError(t, err) + + raw := fb.toolCallMeta() + require.NotEmpty(t, raw, "expected params._meta to be sent") + + var meta map[string]any + require.NoError(t, json.Unmarshal(raw, &meta)) + traceparent, ok := meta["traceparent"].(string) + require.True(t, ok, "expected traceparent in captured _meta") + + // Prove span identity, not just presence: a stale or global context would + // still yield a non-empty traceparent. The W3C format is + // "---"; the trace-id field must match + // the span active at the call site. + wantTraceID := span.SpanContext().TraceID().String() + assert.Contains(t, traceparent, wantTraceID, + "traceparent on the wire must carry the active span's TraceID") + }) + + //nolint:paralleltest // sequential: shares the fakeBackend/session with the sibling subtests + t.Run("progressToken and traceparent coexist on the wire", func(t *testing.T) { + spanCtx, span := tracer.Start(context.Background(), "test-span") + defer span.End() + + // progressToken serializes via Meta.ProgressToken while the injected + // traceparent rides Meta.AdditionalFields — assert both survive the round + // trip into the real serialized _meta bytes. progressToken is caller + // _meta (the 4th CallTool arg), not a tool argument. + _, err := sess.CallTool(spanCtx, "echo", map[string]any{}, map[string]any{"progressToken": "tok"}) + require.NoError(t, err) + + raw := fb.toolCallMeta() + require.NotEmpty(t, raw, "expected params._meta to be sent") + + var meta map[string]any + require.NoError(t, json.Unmarshal(raw, &meta)) + + assert.Equal(t, "tok", meta["progressToken"], "progressToken must survive on the wire") + traceparent, ok := meta["traceparent"].(string) + require.True(t, ok, "expected traceparent alongside progressToken in captured _meta") + assert.Contains(t, traceparent, span.SpanContext().TraceID().String(), + "traceparent must carry the active span's TraceID when coexisting with progressToken") + }) + + //nolint:paralleltest // sequential: reuses the "tools/call" capture key populated above + t.Run("no active span: _meta is omitted", func(t *testing.T) { + _, err := sess.CallTool(context.Background(), "echo", map[string]any{}, nil) + require.NoError(t, err) + + raw := fb.toolCallMeta() + assert.Empty(t, raw, "expected no _meta to be sent without an active trace context") + }) +} + +// TestHTTPSession_CallTool_StripsReservedModernMeta pins the session-path +// counterpart to pkg/vmcp/client's TestLegacyCallTool_StripsReservedMeta_RealBackend: +// mcp_session.go's CallTool (mcp_session.go:211) must strip the reserved +// io.modelcontextprotocol/* _meta keys before forwarding a caller's _meta to +// this Legacy (session-based, stateful) backend — go-sdk v1.7 hard-rejects any +// _meta.protocolVersion on a stateful server regardless of its value — while +// non-reserved caller keys (a custom key and progressToken) survive, and the +// caller-supplied map is never mutated. +func TestHTTPSession_CallTool_StripsReservedModernMeta(t *testing.T) { + t.Parallel() + + fb := &fakeBackend{advertiseTools: true, tools: []mcp.Tool{{Name: "echo"}}} + url := newFakeBackend(t, fb) + + target := &vmcp.BackendTarget{ + WorkloadID: "strip-meta-backend", + WorkloadName: "strip-meta-backend", + BaseURL: url, + TransportType: "streamable-http", + } + + registry := newTestRegistry(t) + connector := NewHTTPConnector(registry) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + sess, _, err := connector(ctx, target, nil, "", nil) + require.NoError(t, err) + t.Cleanup(func() { _ = sess.Close() }) + + // Values are scalars (not nested maps/objects) even for the reserved keys: + // their exact shape is irrelevant here (they get stripped either way), and + // keeping the fixture flat lets maps.Clone's shallow copy below prove real + // immutability instead of merely catching top-level key deletion. + callerMeta := map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientInfo": "downstream-modern-client", + "io.modelcontextprotocol/clientCapabilities": "none", + "custom-caller-key": "custom-value", + "progressToken": "tok", + } + // Snapshot before the call to prove the caller's own map is untouched after + // it (go-style.md: copy before mutating caller input). + wantCallerMeta := maps.Clone(callerMeta) + + _, err = sess.CallTool(ctx, "echo", map[string]any{}, callerMeta) + require.NoError(t, err) + + assert.Equal(t, wantCallerMeta, callerMeta, "the caller-supplied _meta map must not be mutated") + + raw := fb.toolCallMeta() + require.NotEmpty(t, raw, "expected params._meta to be sent") + + var gotMeta map[string]any + require.NoError(t, json.Unmarshal(raw, &gotMeta)) + + // Assert on the namespace, not a fixed list: this catches any reserved key, + // including ones added to the fixture later. No passthroughMetaKeys entry is + // in play here, so a blanket prefix check is exact. + for k := range gotMeta { + assert.False(t, strings.HasPrefix(k, mcpparser.ReservedMetaPrefix), + "reserved _meta key %q must be stripped before the Legacy hop", k) + } + assert.Equal(t, "custom-value", gotMeta["custom-caller-key"], "non-reserved caller _meta must survive") + assert.Equal(t, "tok", gotMeta["progressToken"], "progressToken must survive") +} diff --git a/pkg/vmcp/session/internal/backend/mcp_session_redirect_test.go b/pkg/vmcp/session/internal/backend/mcp_session_redirect_test.go new file mode 100644 index 0000000000..e924947efb --- /dev/null +++ b/pkg/vmcp/session/internal/backend/mcp_session_redirect_test.go @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package backend + +import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestHTTPSession_RefusesCrossHostRedirect verifies the session connector installs +// SameHostRedirectPolicy: a backend returning a cross-host 302 on the initialize +// POST is refused before the second hop, so the RoundTripper-injected forwarded +// header never reaches the attacker host (credential-exfil vector, twin of +// pkg/vmcp/client.newBackendHTTPClient). +func TestHTTPSession_RefusesCrossHostRedirect(t *testing.T) { + t.Parallel() + + var attackerHits atomic.Int32 + attacker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attackerHits.Add(1) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(attacker.Close) + + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Redirect(w, &http.Request{}, attacker.URL, http.StatusFound) + })) + t.Cleanup(origin.Close) + + target := &vmcp.BackendTarget{ + WorkloadID: "redirect-backend", + WorkloadName: "redirect-backend", + BaseURL: origin.URL, + TransportType: "streamable-http", + HeaderForward: &vmcp.HeaderForwardConfig{ + AddPlaintextHeaders: map[string]string{"X-Secret": "leak-me"}, + }, + } + + connector := NewHTTPConnector(newTestRegistry(t)) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + t.Cleanup(cancel) + + sess, _, err := connector(ctx, target, nil, "", nil) + if sess != nil { + _ = sess.Close() + } + require.Error(t, err, "a cross-host redirect on initialize must fail the connection") + assert.Zero(t, attackerHits.Load(), "credential-bearing request must not reach the cross-host redirect target") +} diff --git a/pkg/vmcp/session/internal/backend/mcp_session_test.go b/pkg/vmcp/session/internal/backend/mcp_session_test.go index 79ce6ba721..403c432065 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_test.go @@ -5,11 +5,18 @@ package backend import ( "context" + "net/http" + "net/http/httptest" + "sync" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" @@ -31,6 +38,65 @@ func newTestRegistry(t *testing.T) vmcpauth.OutgoingAuthRegistry { // headerforward.MergeForwardedHeaders; full coverage lives in // pkg/vmcp/headerforward/transport_test.go (TestMergeForwardedHeaders). +// TestNewListChangedNotificationHandler_DispatchesByMethod verifies (#5748, +// #5969) newListChangedNotificationHandler's wire-method dispatch: each of +// the three list_changed methods fires sink with the matching ChangeKind, and +// neither notifications/message nor an unknown method fires sink at all. +func TestNewListChangedNotificationHandler_DispatchesByMethod(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + method string + wantFired bool + wantKind ChangeKind + }{ + {"tools/list_changed dispatches KindTools", vmcp.MethodToolsListChangedNotification, true, KindTools}, + {"resources/list_changed dispatches KindResources", vmcp.MethodResourcesListChangedNotification, true, KindResources}, + {"prompts/list_changed dispatches KindPrompts", vmcp.MethodPromptsListChangedNotification, true, KindPrompts}, + {"notifications/message does not dispatch", vmcp.MethodLogNotification, false, ""}, + {"unknown method does not dispatch", "notifications/resources/updated", false, ""}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + type firedCall struct { + workloadID string + kind ChangeKind + } + fired := make(chan firedCall, 1) + sink := func(_ context.Context, workloadID string, kind ChangeKind) { + fired <- firedCall{workloadID: workloadID, kind: kind} + } + + handler := newListChangedNotificationHandler("backend-1", sink) + handler(mcpmcp.JSONRPCNotification{ + JSONRPC: "2.0", + Notification: mcpmcp.Notification{Method: tc.method}, + }) + + if !tc.wantFired { + select { + case got := <-fired: + t.Fatalf("sink must not fire for method %q, but fired with kind %q", tc.method, got.kind) + case <-time.After(50 * time.Millisecond): + } + return + } + + select { + case got := <-fired: + assert.Equal(t, "backend-1", got.workloadID) + assert.Equal(t, tc.wantKind, got.kind) + case <-time.After(5 * time.Second): + t.Fatalf("sink never fired for method %q", tc.method) + } + }) + } +} + func TestCreateMCPClient_UnsupportedTransport(t *testing.T) { t.Parallel() @@ -46,10 +112,438 @@ func TestCreateMCPClient_UnsupportedTransport(t *testing.T) { TransportType: transport, } - _, err := createMCPClient(context.Background(), target, nil, newTestRegistry(t), "", secrets.NewEnvironmentProvider()) + _, err := createMCPClient(context.Background(), target, nil, newTestRegistry(t), "", secrets.NewEnvironmentProvider(), nil) require.Error(t, err) assert.ErrorIs(t, err, vmcp.ErrUnsupportedTransport, "transport %q should return ErrUnsupportedTransport", transport) }) } } + +// listChangedTestBackend is a real mcpcompat MCP server (streamable-HTTP) with +// WithToolCapabilities(true), WithResourceCapabilities(true, true), and +// WithPromptCapabilities(true). It captures its one connected session (via +// OnRegisterSession) and exposes addTool/addResource/addPrompt, which append a +// per-session tool/resource/prompt via SessionWithTools.SetSessionTools / +// SessionWithResources.SetSessionResources / SessionWithPrompts.SetSessionPrompts +// — the same mechanism ToolHive's own vMCP Server uses +// (setSessionToolsDirect/resyncSessionTools and their resource/prompt +// counterparts) — which is what actually drives the underlying go-sdk +// server's real notifications/{tools,resources,prompts}/list_changed +// broadcast. (The mcpcompat-wrapper-level MCPServer.AddTool/AddResource/ +// AddPrompt called AFTER the server starts serving does NOT do this: it only +// mutates the wrapper's pre-serve set, registered once at buildServer time, so +// it has no effect on an already-connected session.) +type listChangedTestBackend struct { + url string + + // ready closes once OnRegisterSession has captured session, decoupling + // addTool/addResource/addPrompt from the client-side race noted below. + ready chan struct{} + + mu sync.Mutex + session mcpserver.ClientSession + tools map[string]mcpserver.ServerTool + resources map[string]mcpserver.ServerResource + prompts map[string]mcpserver.ServerPrompt +} + +func newListChangedTool(name string) mcpserver.ServerTool { + return mcpserver.ServerTool{ + Tool: mcpmcp.NewTool(name, mcpmcp.WithDescription("a list_changed test tool")), + Handler: func(_ context.Context, _ mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + return mcpmcp.NewToolResultText("ok"), nil + }, + } +} + +func newListChangedResource(uri string) mcpserver.ServerResource { + return mcpserver.ServerResource{ + Resource: mcpmcp.Resource{URI: uri, Name: uri}, + Handler: func(context.Context, mcpmcp.ReadResourceRequest) ([]mcpmcp.ResourceContents, error) { + return []mcpmcp.ResourceContents{mcpmcp.TextResourceContents{URI: uri, Text: "ok"}}, nil + }, + } +} + +func newListChangedPrompt(name string) mcpserver.ServerPrompt { + return mcpserver.ServerPrompt{ + Prompt: mcpmcp.NewPrompt(name), + Handler: func(context.Context, mcpmcp.GetPromptRequest) (*mcpmcp.GetPromptResult, error) { + return &mcpmcp.GetPromptResult{Messages: []mcpmcp.PromptMessage{ + mcpmcp.NewPromptMessage(mcpmcp.RoleUser, mcpmcp.NewTextContent("ok")), + }}, nil + }, + } +} + +// newListChangedTestBackend starts the backend; its OnRegisterSession hook +// captures the connecting session and closes ready. +func newListChangedTestBackend(t *testing.T) *listChangedTestBackend { + t.Helper() + b := &listChangedTestBackend{ + ready: make(chan struct{}), + tools: map[string]mcpserver.ServerTool{ + "initial_tool": newListChangedTool("initial_tool"), + }, + resources: map[string]mcpserver.ServerResource{ + "initial_resource": newListChangedResource("initial_resource"), + }, + prompts: map[string]mcpserver.ServerPrompt{ + "initial_prompt": newListChangedPrompt("initial_prompt"), + }, + } + + hooks := &mcpserver.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, session mcpserver.ClientSession) { + b.mu.Lock() + b.session = session + tools := make(map[string]mcpserver.ServerTool, len(b.tools)) + for k, v := range b.tools { + tools[k] = v + } + resources := make(map[string]mcpserver.ServerResource, len(b.resources)) + for k, v := range b.resources { + resources[k] = v + } + prompts := make(map[string]mcpserver.ServerPrompt, len(b.prompts)) + for k, v := range b.prompts { + prompts[k] = v + } + b.mu.Unlock() + + sessionWithTools, ok := session.(mcpserver.SessionWithTools) + if !ok { + t.Errorf("listChangedTestBackend: session does not support per-session tools") + return + } + sessionWithTools.SetSessionTools(tools) + + sessionWithResources, ok := session.(mcpserver.SessionWithResources) + if !ok { + t.Errorf("listChangedTestBackend: session does not support per-session resources") + return + } + sessionWithResources.SetSessionResources(resources) + + sessionWithPrompts, ok := session.(mcpserver.SessionWithPrompts) + if !ok { + t.Errorf("listChangedTestBackend: session does not support per-session prompts") + return + } + sessionWithPrompts.SetSessionPrompts(prompts) + + close(b.ready) + }) + + srv := mcpserver.NewMCPServer("list-changed-backend", "1.0.0", + mcpserver.WithToolCapabilities(true), + mcpserver.WithResourceCapabilities(true, true), + mcpserver.WithPromptCapabilities(true), + mcpserver.WithHooks(hooks), + ) + streamableSrv := mcpserver.NewStreamableHTTPServer(srv) + mux := http.NewServeMux() + mux.Handle("/mcp", streamableSrv) + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + b.url = ts.URL + "/mcp" + return b +} + +// addTool adds name to the connected session's per-session tool set via +// SetSessionTools, reconciling it onto the live go-sdk server and (since +// WithToolCapabilities(true) is set) triggering a real, asynchronous +// notifications/tools/list_changed broadcast — independent of any in-flight +// call. +// +// It waits on b.ready first: the mcpcompat/go-sdk client's Initialize sends +// notifications/initialized and returns as soon as the SEND succeeds, without +// waiting for the SERVER to finish processing it (registerAndSync, which is +// where OnRegisterSession/b.session-capture runs) — so a caller cannot safely +// call addTool immediately after Initialize returns without this explicit +// synchronization. +func (b *listChangedTestBackend) addTool(t *testing.T, name string) { + t.Helper() + select { + case <-b.ready: + case <-time.After(5 * time.Second): + t.Fatal("listChangedTestBackend: timed out waiting for session registration") + } + + b.mu.Lock() + session := b.session + b.tools[name] = newListChangedTool(name) + tools := make(map[string]mcpserver.ServerTool, len(b.tools)) + for k, v := range b.tools { + tools[k] = v + } + b.mu.Unlock() + sessionWithTools, ok := session.(mcpserver.SessionWithTools) + require.True(t, ok, "listChangedTestBackend: session does not support per-session tools") + sessionWithTools.SetSessionTools(tools) +} + +// addResource adds uri to the connected session's per-session resource set via +// SetSessionResources, triggering a real, asynchronous +// notifications/resources/list_changed broadcast. See addTool for the b.ready +// synchronization rationale. +func (b *listChangedTestBackend) addResource(t *testing.T, uri string) { + t.Helper() + select { + case <-b.ready: + case <-time.After(5 * time.Second): + t.Fatal("listChangedTestBackend: timed out waiting for session registration") + } + + b.mu.Lock() + session := b.session + b.resources[uri] = newListChangedResource(uri) + resources := make(map[string]mcpserver.ServerResource, len(b.resources)) + for k, v := range b.resources { + resources[k] = v + } + b.mu.Unlock() + sessionWithResources, ok := session.(mcpserver.SessionWithResources) + require.True(t, ok, "listChangedTestBackend: session does not support per-session resources") + sessionWithResources.SetSessionResources(resources) +} + +// addPrompt adds name to the connected session's per-session prompt set via +// SetSessionPrompts, triggering a real, asynchronous +// notifications/prompts/list_changed broadcast. See addTool for the b.ready +// synchronization rationale. +func (b *listChangedTestBackend) addPrompt(t *testing.T, name string) { + t.Helper() + select { + case <-b.ready: + case <-time.After(5 * time.Second): + t.Fatal("listChangedTestBackend: timed out waiting for session registration") + } + + b.mu.Lock() + session := b.session + b.prompts[name] = newListChangedPrompt(name) + prompts := make(map[string]mcpserver.ServerPrompt, len(b.prompts)) + for k, v := range b.prompts { + prompts[k] = v + } + b.mu.Unlock() + sessionWithPrompts, ok := session.(mcpserver.SessionWithPrompts) + require.True(t, ok, "listChangedTestBackend: session does not support per-session prompts") + sessionWithPrompts.SetSessionPrompts(prompts) +} + +// TestCreateMCPClient_ContinuousListeningGatedOnSink verifies createMCPClient +// only enables the streamable-HTTP standalone GET stream (WithContinuousListening) +// when a non-nil sink is supplied. Some backends hang when this stream is opened +// against them (#5748 R3), so it must stay strictly opt-in. +func TestCreateMCPClient_ContinuousListeningGatedOnSink(t *testing.T) { + t.Parallel() + + b := newListChangedTestBackend(t) + target := &vmcp.BackendTarget{ + WorkloadID: "list-changed-backend", + WorkloadName: "list-changed-backend", + BaseURL: b.url, + TransportType: "streamable-http", + } + + tests := []struct { + name string + sink ListChangedSink + wantContinuous bool + }{ + {name: "nil sink disables continuous listening", sink: nil, wantContinuous: false}, + { + name: "non-nil sink enables continuous listening", + sink: func(context.Context, string, ChangeKind) {}, + wantContinuous: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + c, err := createMCPClient( + context.Background(), target, nil, newTestRegistry(t), "", secrets.NewEnvironmentProvider(), tc.sink, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + sh, ok := c.GetTransport().(*mcptransport.StreamableHTTP) + require.True(t, ok, "expected a streamable-HTTP transport") + assert.Equal(t, tc.wantContinuous, sh.ContinuousListening()) + }) + } +} + +// TestCreateMCPClient_ListChangedSink_FiresOnBackendNotification is the +// end-to-end regression for #5748 (tools) and #5969 (resources, prompts): a +// non-nil sink fires with the backend's WorkloadID and the expected kind when +// the real backend emits notifications/{tools,resources,prompts}/list_changed +// asynchronously (outside any in-flight call), over the same +// continuous-listening connection, and the nil-sink path never panics or +// opens the standalone stream. +func TestCreateMCPClient_ListChangedSink_FiresOnBackendNotification(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + trigger func(t *testing.T, b *listChangedTestBackend) + wantKind ChangeKind + }{ + {"tools", func(t *testing.T, b *listChangedTestBackend) { + t.Helper() + b.addTool(t, "added_tool") + }, KindTools}, + {"resources", func(t *testing.T, b *listChangedTestBackend) { + t.Helper() + b.addResource(t, "added_resource") + }, KindResources}, + {"prompts", func(t *testing.T, b *listChangedTestBackend) { + t.Helper() + b.addPrompt(t, "added_prompt") + }, KindPrompts}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + b := newListChangedTestBackend(t) + target := &vmcp.BackendTarget{ + WorkloadID: "list-changed-backend", + WorkloadName: "list-changed-backend", + BaseURL: b.url, + TransportType: "streamable-http", + } + + type firedCall struct { + backendWorkloadID string + kind ChangeKind + } + // Buffered generously: registration itself sets all three per-session + // stores (tools, resources, prompts), each of which fires its own + // registration-time list_changed broadcast (see serve.go's R1 note) in + // addition to the one this test explicitly triggers below. + fired := make(chan firedCall, 16) + sink := func(_ context.Context, backendWorkloadID string, kind ChangeKind) { + fired <- firedCall{backendWorkloadID: backendWorkloadID, kind: kind} + } + + c, err := createMCPClient( + context.Background(), target, nil, newTestRegistry(t), "", secrets.NewEnvironmentProvider(), sink, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(context.Background(), mcpmcp.InitializeRequest{ + Params: mcpmcp.InitializeParams{ + ProtocolVersion: mcpmcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcpmcp.Implementation{Name: "test-client", Version: "1.0"}, + }, + }) + require.NoError(t, err) + + // Backend asynchronously changes its per-session set — a real + // list_changed trigger, independent of any in-flight call from this + // client. Drain fired notifications (which may also include + // registration-time broadcasts of OTHER kinds, per the buffer comment + // above) until the expected kind is observed. + tc.trigger(t, b) + + deadline := time.After(10 * time.Second) + for { + select { + case got := <-fired: + assert.Equal(t, "list-changed-backend", got.backendWorkloadID) + if got.kind == tc.wantKind { + return + } + case <-deadline: + t.Fatalf("timed out waiting for a %q list_changed notification", tc.wantKind) + } + } + }) + } +} + +// TestCreateMCPClient_ListChangedSink_DoesNotStallInFlightCall is the F9 +// corroborated-concern regression: firing a real backend tools/list_changed — +// whose sink is (deliberately, for the test) slow — must NOT stall an in-flight +// call on the same backend session. The sink runs on the client's +// notification-dispatch goroutine, which is independent of the request path, so +// a CallTool issued while the sink is still working completes promptly. This +// validates the no-deadlock conclusion behind moving the real resync off the +// receive loop (#5748 B-HIGH-2). +func TestCreateMCPClient_ListChangedSink_DoesNotStallInFlightCall(t *testing.T) { + t.Parallel() + + b := newListChangedTestBackend(t) + target := &vmcp.BackendTarget{ + WorkloadID: "list-changed-backend", + WorkloadName: "list-changed-backend", + BaseURL: b.url, + TransportType: "streamable-http", + } + + sinkEntered := make(chan struct{}, 1) + releaseSink := make(chan struct{}) + sink := func(_ context.Context, _ string, kind ChangeKind) { + if kind != KindTools { + return + } + select { + case sinkEntered <- struct{}{}: + default: + } + // Simulate a slow resync body running on the dispatch goroutine to prove + // it does not hold anything the request path needs. + <-releaseSink + } + + c, err := createMCPClient( + context.Background(), target, nil, newTestRegistry(t), "", secrets.NewEnvironmentProvider(), sink, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(context.Background(), mcpmcp.InitializeRequest{ + Params: mcpmcp.InitializeParams{ + ProtocolVersion: mcpmcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcpmcp.Implementation{Name: "test-client", Version: "1.0"}, + }, + }) + require.NoError(t, err) + + // Trigger a real tools/list_changed and wait until the (slow) sink is running. + b.addTool(t, "added_tool") + select { + case <-sinkEntered: + case <-time.After(10 * time.Second): + close(releaseSink) + t.Fatal("sink never fired") + } + + // With the sink still blocked, an in-flight call on the same session must + // complete promptly (independent dispatch vs request paths). + callDone := make(chan error, 1) + go func() { + _, callErr := c.CallTool(context.Background(), mcpmcp.CallToolRequest{ + Params: mcpmcp.CallToolParams{Name: "initial_tool"}, + }) + callDone <- callErr + }() + + select { + case callErr := <-callDone: + require.NoError(t, callErr, "in-flight call must succeed while a resync sink is running") + case <-time.After(10 * time.Second): + close(releaseSink) + t.Fatal("in-flight call stalled while the list_changed sink was running") + } + + close(releaseSink) +} diff --git a/pkg/vmcp/session/mocks/mock_factory.go b/pkg/vmcp/session/mocks/mock_factory.go index 3d850a1eb4..7ec4a60adb 100644 --- a/pkg/vmcp/session/mocks/mock_factory.go +++ b/pkg/vmcp/session/mocks/mock_factory.go @@ -44,18 +44,18 @@ func (m *MockMultiSessionFactory) EXPECT() *MockMultiSessionFactoryMockRecorder } // MakeSessionWithID mocks base method. -func (m *MockMultiSessionFactory) MakeSessionWithID(ctx context.Context, id string, identity *auth.Identity, backends []*vmcp.Backend) (session.MultiSession, error) { +func (m *MockMultiSessionFactory) MakeSessionWithID(ctx context.Context, id string, identity *auth.Identity, backends []*vmcp.Backend, sink session.ListChangedSink) (session.MultiSession, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MakeSessionWithID", ctx, id, identity, backends) + ret := m.ctrl.Call(m, "MakeSessionWithID", ctx, id, identity, backends, sink) ret0, _ := ret[0].(session.MultiSession) ret1, _ := ret[1].(error) return ret0, ret1 } // MakeSessionWithID indicates an expected call of MakeSessionWithID. -func (mr *MockMultiSessionFactoryMockRecorder) MakeSessionWithID(ctx, id, identity, backends any) *gomock.Call { +func (mr *MockMultiSessionFactoryMockRecorder) MakeSessionWithID(ctx, id, identity, backends, sink any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MakeSessionWithID", reflect.TypeOf((*MockMultiSessionFactory)(nil).MakeSessionWithID), ctx, id, identity, backends) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MakeSessionWithID", reflect.TypeOf((*MockMultiSessionFactory)(nil).MakeSessionWithID), ctx, id, identity, backends, sink) } // RestoreSession mocks base method. diff --git a/pkg/vmcp/session/optimizerdec/decorator.go b/pkg/vmcp/session/optimizerdec/decorator.go index b3e4454e88..26a5f393c7 100644 --- a/pkg/vmcp/session/optimizerdec/decorator.go +++ b/pkg/vmcp/session/optimizerdec/decorator.go @@ -10,8 +10,7 @@ import ( "encoding/json" "fmt" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" @@ -25,12 +24,6 @@ const ( FindToolName = "find_tool" // CallToolName is the tool name for routing a call to any backend tool. CallToolName = "call_tool" - // CallToolArgToolName is the JSON argument key for the backend tool name in a call_tool request. - // It must match the json tag on optimizer.CallToolInput.ToolName. - CallToolArgToolName = "tool_name" - // CallToolArgParameters is the JSON argument key for the backend tool parameters in a call_tool request. - // It must match the json tag on optimizer.CallToolInput.Parameters. - CallToolArgParameters = "parameters" ) // Pre-generated schemas for find_tool and call_tool, computed at init time. diff --git a/pkg/vmcp/session/optimizerdec/decorator_test.go b/pkg/vmcp/session/optimizerdec/decorator_test.go index 514c163df4..b119c8fa12 100644 --- a/pkg/vmcp/session/optimizerdec/decorator_test.go +++ b/pkg/vmcp/session/optimizerdec/decorator_test.go @@ -6,14 +6,13 @@ package optimizerdec_test import ( "context" "errors" - "reflect" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" @@ -198,28 +197,3 @@ func TestOptimizerDecorator_CallTool_CallTool(t *testing.T) { require.Error(t, err) }) } - -// TestCallToolArgConstantsMatchStructTags verifies that CallToolArgToolName and -// CallToolArgParameters match the json tags on optimizer.CallToolInput. The middleware -// uses these constants to look up fields from parsed arguments; a mismatch causes an -// authz bypass or parameters being silently dropped. -func TestCallToolArgConstantsMatchStructTags(t *testing.T) { - t.Parallel() - - typ := reflect.TypeOf(optimizer.CallToolInput{}) - - cases := []struct { - field string - constant string - }{ - {"ToolName", optimizerdec.CallToolArgToolName}, - {"Parameters", optimizerdec.CallToolArgParameters}, - } - - for _, tc := range cases { - f, ok := typ.FieldByName(tc.field) - require.True(t, ok, "optimizer.CallToolInput must have a %s field", tc.field) - assert.Equal(t, tc.constant, f.Tag.Get("json"), - "constant for %s must match its json struct tag", tc.field) - } -} diff --git a/pkg/vmcp/session/session.go b/pkg/vmcp/session/session.go index eddf72941b..57e5b92efe 100644 --- a/pkg/vmcp/session/session.go +++ b/pkg/vmcp/session/session.go @@ -5,6 +5,7 @@ package session import ( "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp/session/internal/backend" "github.com/stacklok/toolhive/pkg/vmcp/session/internal/security" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" ) @@ -13,6 +14,25 @@ import ( // backward compatibility and convenience. type MultiSession = sessiontypes.MultiSession +// ListChangedSink is an alias for backend.ListChangedSink, re-exported here so +// callers outside the pkg/vmcp/session/internal/backend package (e.g. +// pkg/vmcp/server, which builds the session-registration sink) can reference it +// without importing the internal package directly. +type ListChangedSink = backend.ListChangedSink + +// ChangeKind is an alias for backend.ChangeKind, re-exported for the same +// reason as ListChangedSink. +type ChangeKind = backend.ChangeKind + +// KindTools, KindResources, and KindPrompts re-export the corresponding +// backend.Kind* constants so out-of-tree callers can match on them without +// importing the internal package. +const ( + KindTools = backend.KindTools + KindResources = backend.KindResources + KindPrompts = backend.KindPrompts +) + // ValidateCaller checks caller against a stored identity-binding string (the // value persisted under MetadataKeyIdentityBinding) and returns nil when the // caller is permitted, or ErrNilCaller / ErrUnauthorizedCaller / ErrSessionOwnerUnknown diff --git a/pkg/vmcp/types.go b/pkg/vmcp/types.go index 19ed4c68b0..1ac7811560 100644 --- a/pkg/vmcp/types.go +++ b/pkg/vmcp/types.go @@ -9,6 +9,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "github.com/stacklok/toolhive/pkg/mcp" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" ) @@ -292,6 +293,11 @@ type DiscoveredBackend struct { // Resets to 0 when the backend becomes healthy again. // +optional ConsecutiveFailures int `json:"consecutiveFailures,omitempty"` + + // MCPRevision is the backend's negotiated MCP protocol revision + // ("2026-07-28" or "2025-11-25"). Empty when the backend has not been probed. + // +optional + MCPRevision string `json:"mcpRevision,omitempty"` } // DeepCopyInto copies the receiver into out. Required for Kubernetes CRD types. @@ -429,6 +435,29 @@ type Resource struct { BackendID string } +// ResourceTemplate represents an MCP resource template capability. +// Unlike a Resource, whose URI identifies a single concrete resource, a +// ResourceTemplate declares an RFC 6570 URI template that expands to a family +// of resource URIs (e.g. "file:///logs/{date}.txt"). Reading an expanded URI is +// served through the existing resource-read path: the router matches the URI +// against these templates when the exact-URI lookup misses. +type ResourceTemplate struct { + // URITemplate is the RFC 6570 URI template (should be globally unique). + URITemplate string + + // Name is a human-readable name. + Name string + + // Description describes the resource template. + Description string + + // MimeType is the MIME type shared by resources matching this template (optional). + MimeType string + + // BackendID identifies the backend that provides this resource template. + BackendID string +} + // Prompt represents an MCP prompt capability. type Prompt struct { // Name is the prompt name (may conflict with other backends). @@ -540,6 +569,11 @@ type ToolCallResult struct { // This includes progressToken, trace context, and custom backend metadata. // Per MCP specification, this field is optional and may be nil. Meta map[string]any + + // BackendID is the logical backend that served the call, set by the core + // after routing. Empty for a composite tool (no single serving backend). + // Audit-only: never serialized to the client. + BackendID string `json:"-"` } // ResourceContent represents a single resource content item, @@ -564,10 +598,13 @@ type ResourceReadResult struct { Contents []ResourceContent // Meta contains protocol-level metadata from the backend (_meta field). - // NOTE: Due to MCP SDK limitations, resources/read handlers cannot forward _meta - // because they return []ResourceContents directly, not a result wrapper. - // This field is preserved for future SDK improvements but may be nil. + // Populated on both the Legacy and Modern client paths; forwarded to the + // client by the vMCP Serve handlers (non-reserved keys only). Meta map[string]any + + // BackendID is the logical backend that served the read, set by the core + // after routing. Audit-only: never serialized to the client. + BackendID string `json:"-"` } // PromptMessage represents a single message in a prompt response, @@ -594,6 +631,47 @@ type PromptGetResult struct { // This includes progressToken, trace context, and custom backend metadata. // Per MCP specification, this field is optional and may be nil. Meta map[string]any + + // BackendID is the logical backend that served the get, set by the core + // after routing. Audit-only: never serialized to the client. + BackendID string `json:"-"` +} + +// Completion reference type constants, matching the MCP completion/complete +// "ref" discriminator values. +const ( + // CompletionRefTypePrompt identifies a prompt-argument completion reference. + CompletionRefTypePrompt = "ref/prompt" + // CompletionRefTypeResource identifies a resource(-template) completion reference. + CompletionRefTypeResource = "ref/resource" +) + +// CompletionRef identifies what a completion/complete request is completing: +// either a prompt argument (Type == CompletionRefTypePrompt, Name set) or a +// resource-template variable (Type == CompletionRefTypeResource, URI set). It is +// the domain equivalent of the MCP PromptReference/ResourceReference shapes, +// following the Anti-Corruption Layer pattern (no mcp-go types cross the core +// boundary, vmcp anti-pattern #5). +type CompletionRef struct { + // Type is the reference discriminator: CompletionRefTypePrompt or + // CompletionRefTypeResource. + Type string + // Name is the prompt name (set only for CompletionRefTypePrompt). + Name string + // URI is the resource URI or URI template (set only for + // CompletionRefTypeResource). + URI string +} + +// CompletionResult holds argument-completion candidates returned by a backend. +// It is the domain equivalent of the MCP CompletionResultDetails. +type CompletionResult struct { + // Values are the completion candidate strings (MCP caps this at 100 items). + Values []string + // Total is the total number of options available, which may exceed len(Values). + Total int + // HasMore indicates additional options exist beyond those in Values. + HasMore bool } // RoutingTable contains the mappings from capability names to backend targets. @@ -611,6 +689,11 @@ type RoutingTable struct { // Resources maps resource URIs to their backend targets. Resources map[string]*BackendTarget + // ResourceTemplates maps resource URI-template strings to their backend targets. + // The router consults this when an exact Resources lookup misses: it matches + // the requested URI against each template and routes to the first match. + ResourceTemplates map[string]*BackendTarget + // Prompts maps prompt names to their backend targets. Prompts map[string]*BackendTarget } @@ -653,8 +736,18 @@ type BackendClient interface { // CallTool invokes a tool on the backend MCP server. // The meta parameter contains _meta fields from the client request that should be forwarded to the backend. // Returns the complete tool result including _meta field from the backend response. + // + // paramHeaders carries the SEP-2243 Mcp-Param-* headers mirrored from the tool's + // x-mcp-header-designated arguments, keyed by full header name. It is supplied by + // the caller because deriving it needs the tool's inputSchema, which lives in the + // aggregated capability view rather than on BackendTarget; passing it explicitly + // keeps this client free of a schema cache and its staleness window. nil or empty + // means nothing to mirror, which is the common case. It applies only to a Modern + // (2026-07-28) backend hop -- SEP-2243 is part of that revision -- and is ignored + // for a Legacy backend. CallTool( - ctx context.Context, target *BackendTarget, toolName string, arguments map[string]any, meta map[string]any, + ctx context.Context, target *BackendTarget, toolName string, arguments map[string]any, + meta map[string]any, paramHeaders map[string]string, ) (*ToolCallResult, error) // ReadResource retrieves a resource from the backend MCP server. @@ -665,11 +758,45 @@ type BackendClient interface { // Returns the complete prompt result including _meta field. GetPrompt(ctx context.Context, target *BackendTarget, name string, arguments map[string]any) (*PromptGetResult, error) + // Complete requests argument-completion candidates from the backend MCP server + // (the completion/complete method). ref identifies the prompt or resource being + // completed; argName/argValue are the argument name and its partial value; + // contextArgs carries previously-resolved argument values for multi-argument + // completion (may be nil). For a CompletionRefTypePrompt ref, the prompt name is + // translated to the backend's own name via target.GetBackendCapabilityName. + // + // If the backend does not advertise the completions capability, an empty (non-nil) + // CompletionResult is returned without error, matching the MCP spec's lenient + // completion semantics. + Complete( + ctx context.Context, target *BackendTarget, ref CompletionRef, + argName, argValue string, contextArgs map[string]string, + ) (*CompletionResult, error) + // ListCapabilities queries a backend for its capabilities. // Returns tools, resources, and prompts exposed by the backend. ListCapabilities(ctx context.Context, target *BackendTarget) (*CapabilityList, error) } +// RevisionReporter reports a backend's resolved MCP protocol revision, keyed by +// workload ID. The second return distinguishes "known" from "not probed yet" — +// callers must not read the Revision when it is false, since RevisionLegacy is +// the zero value. +// +// Deliberately NOT part of [BackendClient]: it is a read-model for +// status/telemetry and for deciding whether a Legacy session connect is worth +// attempting, not a protocol operation, and every BackendClient implementation +// (including the generated mocks) would otherwise have to satisfy it. Consumers +// type-assert to it and degrade gracefully when the assertion fails. +// +// Because that assertion is the failure mode — a wrapper that forgets to forward +// CachedRevision silently disables the behaviour with no compile error — +// implementations should pin themselves with a compile-time assertion, e.g. +// `var _ vmcp.RevisionReporter = (*httpBackendClient)(nil)`. +type RevisionReporter interface { + CachedRevision(workloadID string) (mcp.Revision, bool) +} + // CapabilityList contains the capabilities from a backend's MCP server. // This is returned by BackendClient.ListCapabilities(). type CapabilityList struct { @@ -679,6 +806,9 @@ type CapabilityList struct { // Resources available on this backend. Resources []Resource + // ResourceTemplates available on this backend. + ResourceTemplates []ResourceTemplate + // Prompts available on this backend. Prompts []Prompt diff --git a/pkg/webhook/mutating/middleware.go b/pkg/webhook/mutating/middleware.go index d484f6af98..c1c3d18535 100644 --- a/pkg/webhook/mutating/middleware.go +++ b/pkg/webhook/mutating/middleware.go @@ -129,8 +129,30 @@ func createMutatingHandler(executors []clientExecutor, serverName, transport str return } + // A mutating webhook rewrote the body, so the parse cached by ParsingMiddleware + // now describes a request the backend will not execute. Refresh it before any + // downstream consumer (authorization, audit, telemetry) reads it. + if !bytes.Equal(bodyBytes, mutatedBody) { + republished, err := mcp.RepublishParsedMCPRequest(r, mutatedBody) + if err != nil { + var batchErr *mcp.BatchUnsupportedError + if errors.As(err, &batchErr) { + mcp.WriteBatchUnsupportedError(w) + return + } + slog.Error("Mutating webhook produced an unparsable MCP request", "error", err) + sendErrorResponse(w, http.StatusInternalServerError, "Webhook produced an invalid MCP request", parsedMCP.ID) + return + } + r = republished + } + // Replace the request body with the (potentially mutated) MCP body for downstream handlers. + // ContentLength must be updated alongside it: a patch almost always changes the body + // length, and the reverse proxy rejects a forwarded request whose declared length + // disagrees with the body it can actually read. r.Body = io.NopCloser(bytes.NewBuffer(mutatedBody)) + r.ContentLength = int64(len(mutatedBody)) next.ServeHTTP(w, r) }) } @@ -316,18 +338,20 @@ func readSourceIP(r *http.Request) string { } func sendErrorResponse(w http.ResponseWriter, statusCode int, message string, msgID interface{}) { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(statusCode) - id, err := mcp.ConvertToJSONRPC2ID(msgID) if err != nil { id = jsonrpc2.ID{} // Use empty ID if conversion fails. } + code := mcp.JSONRPCCodeForStatus(statusCode) // Return a JSON-RPC 2.0 error so MCP clients can parse the denial. errResp := &jsonrpc2.Response{ ID: id, - Error: jsonrpc2.NewError(int64(statusCode), message), + Error: jsonrpc2.NewError(code, message), } - _ = json.NewEncoder(w).Encode(errResp) + + // The helper encodes before writing any header, so an encode failure never + // leaves a half-written response (e.g. a status header followed by a + // second write). Nothing to do on a write error for a denial body. + _ = mcp.WriteJSONRPCError(w, statusCode, errResp) } diff --git a/pkg/webhook/mutating/middleware_test.go b/pkg/webhook/mutating/middleware_test.go index 12670dd816..8fe94171d6 100644 --- a/pkg/webhook/mutating/middleware_test.go +++ b/pkg/webhook/mutating/middleware_test.go @@ -50,7 +50,7 @@ func makeMCPRequest(tb testing.TB, body []byte) *http.Request { req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) parsedMCP := &mcp.ParsedMCPRequest{ Method: "tools/call", - ID: float64(1), + ID: int64(1), } ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, parsedMCP) req = req.WithContext(ctx) @@ -58,6 +58,17 @@ func makeMCPRequest(tb testing.TB, body []byte) *http.Request { return req } +// newUnparsedMCPRequest builds a POST request with a JSON body and no parsed +// request in its context, so mcp.ParsingMiddleware will actually parse it +// (see shouldParseMCPRequest) rather than short-circuiting on an existing parse. +func newUnparsedMCPRequest(tb testing.TB, body []byte) *http.Request { + tb.Helper() + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.RemoteAddr = "192.168.1.1:1234" + return req +} + //nolint:paralleltest // Shares mock server state func TestMutatingMiddleware_AllowedWithPatch(t *testing.T) { const reqBody = `{"jsonrpc":"2.0","method":"tools/call","id":1,"params":{"name":"db","arguments":{"query":"SELECT *"}}}` @@ -244,6 +255,14 @@ func TestMutatingMiddleware_HTTP422AlwaysDenies(t *testing.T) { assert.False(t, nextCalled) assert.Equal(t, http.StatusUnprocessableEntity, rr.Code) assert.Contains(t, rr.Body.String(), "Request denied by webhook policy") + + // An HTTP 422 always-deny must carry the same JSON-RPC denial code as + // every other denial (mcp.JSONRPCCodeDenied), not the raw 422 status. + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + errObj, ok := errResp["error"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(mcp.JSONRPCCodeDenied), errObj["code"]) }) } } @@ -277,6 +296,12 @@ func TestMutatingMiddleware_ScopeViolation_FailPolicy(t *testing.T) { assert.False(t, nextCalled) assert.Equal(t, http.StatusInternalServerError, rr.Code) + + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + errObj, ok := errResp["error"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(mcp.CodeInternalError), errObj["code"]) } func TestMutatingMiddleware_ScopeViolation_IgnorePolicy(t *testing.T) { @@ -389,6 +414,210 @@ func TestMutatingMiddleware_ChainedMutations(t *testing.T) { assert.Equal(t, "SELECT *", args["query"], "original query should be preserved") } +// renameToolReqBody and renameToolPatch model a mutating webhook that renames +// the called tool. Shared by the tests that pin the tool-rename invariant from +// both the authz (context) and audit (holder) observation points. +const renameToolReqBody = `{"jsonrpc":"2.0","method":"tools/call","id":1,"params":{"name":"read_file","arguments":{"path":"/etc/passwd"}}}` + +var renameToolPatch = []JSONPatchOp{ + {Op: "replace", Path: "/mcp_request/params/name", Value: json.RawMessage(`"delete_file"`)}, +} + +// newMutatingWebhookServer starts a fake mutating webhook that always allows +// the request and returns the given JSON Patch. Callers must not mutate patch +// afterwards: it's read by the handler on every request, including from +// parallel tests that share a patch value. +func newMutatingWebhookServer(t *testing.T, patch []JSONPatchOp) *httptest.Server { + t.Helper() + patchJSON, err := json.Marshal(patch) + require.NoError(t, err) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req webhook.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + + resp := webhook.MutatingResponse{ + Response: webhook.Response{Version: webhook.APIVersion, UID: req.UID, Allowed: true}, + PatchType: patchTypeJSONPatch, + Patch: patchJSON, + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + })) + t.Cleanup(server.Close) + return server +} + +// TestMutatingMiddleware_ParsedRequestReflectsRenamedTool pins the invariant +// that authorization (which reads mcp.GetParsedMCPRequest) depends on: the +// cached parse in the request context must always describe the SAME tool +// name that ends up in the body the backend receives, even after a mutating +// webhook renames the tool. The chain here uses the real mcp.ParsingMiddleware +// (not a hand-built ParsedMCPRequest) so the cached parse is produced exactly +// as it is in production. +func TestMutatingMiddleware_ParsedRequestReflectsRenamedTool(t *testing.T) { + t.Parallel() + + server := newMutatingWebhookServer(t, renameToolPatch) + mw := createMutatingHandler(makeExecutors(t, []webhook.Config{makeConfig(server.URL, webhook.FailurePolicyFail)}), "srv", "stdio") + + var nextCalled bool + handler := mcp.ParsingMiddleware(mw(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + nextCalled = true + parsed := mcp.GetParsedMCPRequest(r.Context()) + require.NotNil(t, parsed) + + bodyBytes, err := io.ReadAll(r.Body) + require.NoError(t, err) + assert.EqualValues(t, len(bodyBytes), r.ContentLength, "ContentLength must match the mutated body actually readable from r.Body") + + var body map[string]interface{} + require.NoError(t, json.Unmarshal(bodyBytes, &body)) + params := body["params"].(map[string]interface{}) + + assert.Equal(t, "delete_file", parsed.ResourceID, "cached parse must reflect the renamed tool") + assert.Equal(t, params["name"], parsed.ResourceID, "cached parse and mutated body must agree on the tool name") + }))) + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, newUnparsedMCPRequest(t, []byte(renameToolReqBody))) + + require.True(t, nextCalled, "next must be called for the request to be observed") + require.Equal(t, http.StatusOK, rr.Code) +} + +// TestMutatingMiddleware_ParsedRequestReflectsPatchedArgument pins the same +// invariant as the tool-rename case but for an argument value: a webhook that +// replaces an existing argument (not one that merely adds a new field) must be +// reflected in the cached parse's Arguments, since authorization/audit read +// Arguments from the parse rather than re-reading the body. +func TestMutatingMiddleware_ParsedRequestReflectsPatchedArgument(t *testing.T) { + t.Parallel() + + const reqBody = `{"jsonrpc":"2.0","method":"tools/call","id":1,"params":{"name":"db","arguments":{"query":"SELECT *"}}}` + patch := []JSONPatchOp{ + {Op: "replace", Path: "/mcp_request/params/arguments/query", Value: json.RawMessage(`"DROP TABLE users"`)}, + } + + server := newMutatingWebhookServer(t, patch) + mw := createMutatingHandler(makeExecutors(t, []webhook.Config{makeConfig(server.URL, webhook.FailurePolicyFail)}), "srv", "stdio") + + var nextCalled bool + handler := mcp.ParsingMiddleware(mw(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + nextCalled = true + parsed := mcp.GetParsedMCPRequest(r.Context()) + require.NotNil(t, parsed) + assert.Equal(t, "DROP TABLE users", parsed.Arguments["query"], "cached parse must reflect the patched argument") + }))) + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, newUnparsedMCPRequest(t, []byte(reqBody))) + + require.True(t, nextCalled, "next must be called for the request to be observed") + require.Equal(t, http.StatusOK, rr.Code) +} + +// TestMutatingMiddleware_HolderReflectsMutatedResourceID pins the invariant +// that audit (pkg/audit/auditor.go) depends on: it wraps the chain, injects an +// mcp.ParsedRequestHolder before calling in, and reads the holder after +// ServeHTTP returns. That holder must reflect the SAME mutated tool name that +// ends up in the body the backend receives, not the pre-mutation name the +// parser first observed. +func TestMutatingMiddleware_HolderReflectsMutatedResourceID(t *testing.T) { + t.Parallel() + + server := newMutatingWebhookServer(t, renameToolPatch) + mw := createMutatingHandler(makeExecutors(t, []webhook.Config{makeConfig(server.URL, webhook.FailurePolicyFail)}), "srv", "stdio") + handler := mcp.ParsingMiddleware(mw(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}))) + + holder := &mcp.ParsedRequestHolder{} + req := newUnparsedMCPRequest(t, []byte(renameToolReqBody)) + req = req.WithContext(mcp.WithParsedRequestHolder(req.Context(), holder)) + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + require.Equal(t, http.StatusOK, rr.Code) + require.NotNil(t, holder.Parsed, "audit holder must be filled by the parser") + assert.Equal(t, "delete_file", holder.Parsed.ResourceID, "audit holder must reflect the mutated tool name") +} + +// TestMutatingMiddleware_PatchRemovesMethod_ReturnsConformantError pins the +// republish failure path: a webhook patch can produce a body that is valid +// JSON but no longer a JSON-RPC request (e.g. the "method" field is removed). +// RepublishParsedMCPRequest fails to parse it, so the middleware must fail +// closed with a 500 and a conformant JSON-RPC error, never reaching next. +func TestMutatingMiddleware_PatchRemovesMethod_ReturnsConformantError(t *testing.T) { + t.Parallel() + + patch := []JSONPatchOp{ + {Op: "remove", Path: "/mcp_request/method"}, + } + server := newMutatingWebhookServer(t, patch) + mw := createMutatingHandler(makeExecutors(t, []webhook.Config{makeConfig(server.URL, webhook.FailurePolicyFail)}), "srv", "stdio") + + var nextCalled bool + handler := mcp.ParsingMiddleware(mw(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + nextCalled = true + }))) + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, newUnparsedMCPRequest(t, []byte(renameToolReqBody))) + + assert.False(t, nextCalled, "next must not be called when the mutated body fails to re-parse") + assert.Equal(t, http.StatusInternalServerError, rr.Code) + assert.Equal(t, "application/json", rr.Header().Get("Content-Type")) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &decoded)) + assert.Equal(t, "2.0", decoded["jsonrpc"]) + assert.Equal(t, float64(1), decoded["id"]) + errObj, ok := decoded["error"].(map[string]interface{}) + require.True(t, ok, "response must contain a JSON-RPC error object") + assert.NotEmpty(t, errObj["message"]) + _, hasResult := decoded["result"] + assert.False(t, hasResult, "error response must not contain a result key") +} + +// TestMutatingMiddleware_NoMutation_ParseUnchangedAndNextCalled pins the +// bytes.Equal guard: when a webhook allows the request without a patch, the +// body handed to RepublishParsedMCPRequest would be byte-identical to what +// ParsingMiddleware already parsed, so no republish should happen, and next +// must still be called with the original, correct parse intact. +func TestMutatingMiddleware_NoMutation_ParseUnchangedAndNextCalled(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + resp := webhook.MutatingResponse{ + Response: webhook.Response{Version: webhook.APIVersion, UID: "uid", Allowed: true}, + // No patch: the body is unchanged. + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + })) + t.Cleanup(server.Close) + + mw := createMutatingHandler(makeExecutors(t, []webhook.Config{makeConfig(server.URL, webhook.FailurePolicyFail)}), "srv", "stdio") + + var nextCalled bool + handler := mcp.ParsingMiddleware(mw(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + nextCalled = true + parsed := mcp.GetParsedMCPRequest(r.Context()) + require.NotNil(t, parsed) + assert.Equal(t, "read_file", parsed.ResourceID, "parse must reflect the original, unmutated tool name") + + bodyBytes, err := io.ReadAll(r.Body) + require.NoError(t, err) + assert.EqualValues(t, len(bodyBytes), r.ContentLength, "ContentLength must match the unchanged body on the no-mutation path") + }))) + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, newUnparsedMCPRequest(t, []byte(renameToolReqBody))) + + assert.True(t, nextCalled, "next must be called on the no-mutation path") + assert.Equal(t, http.StatusOK, rr.Code) +} + func TestMutatingMiddleware_SkipNonMCPRequests(t *testing.T) { t.Parallel() mw := createMutatingHandler(nil, "srv", "stdio") @@ -405,6 +634,65 @@ func TestMutatingMiddleware_SkipNonMCPRequests(t *testing.T) { assert.Equal(t, http.StatusOK, rr.Code) } +// TestMutatingMiddlewareDenialIsConformantJSONRPC pins the sendErrorResponse envelope on +// the wire: lowercase "jsonrpc"/"id"/"error" keys, no "id" key for a notification, and no +// "result" key. assert.JSONEq is case-sensitive, so a single assertion catches "Error"/"ID" +// substituted for "error"/"id" (the historical bug) alongside a missing "jsonrpc" tag. +func TestMutatingMiddlewareDenialIsConformantJSONRPC(t *testing.T) { + t.Parallel() + + // Deterministic denial: the webhook always disallows, driving sendErrorResponse + // through the real middleware chain rather than constructing the response by hand. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(webhook.MutatingResponse{ + Response: webhook.Response{Version: webhook.APIVersion, UID: "resp-uid", Allowed: false}, + }) + })) + t.Cleanup(server.Close) + + cfg := makeConfig(server.URL, webhook.FailurePolicyFail) + mw := createMutatingHandler(makeExecutors(t, []webhook.Config{cfg}), "srv", "stdio") + + const wantErr = `"error":{"code":403,"message":"Request denied by webhook policy"}` + + testCases := []struct { + name string + id interface{} // ParsedMCPRequest.ID; nil for a notification + wantBody string + }{ + {name: "int64 id", id: int64(42), wantBody: `{"jsonrpc":"2.0","id":42,` + wantErr + `}`}, + {name: "string id", id: "abc", wantBody: `{"jsonrpc":"2.0","id":"abc",` + wantErr + `}`}, + {name: "nil id (notification) omits id entirely", id: nil, wantBody: `{"jsonrpc":"2.0",` + wantErr + `}`}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"tools/call"}`))) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: tc.id}) + req = req.WithContext(ctx) + + rr := httptest.NewRecorder() + mw(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("next handler must not be called for a denied request") + })).ServeHTTP(rr, req) + + require.Equal(t, http.StatusForbidden, rr.Code) + assert.Equal(t, "application/json", rr.Header().Get("Content-Type")) + assert.JSONEq(t, tc.wantBody, rr.Body.String()) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &decoded)) + _, hasID := decoded["id"] + assert.Equal(t, tc.id != nil, hasID, "id key presence mismatch") + _, hasResult := decoded["result"] + assert.False(t, hasResult, "denied response must not contain a result key") + }) + } +} + // makeJSONBodyOfSize builds a syntactically-valid JSON body of exactly `size` // bytes by padding the value of a "data" field with ASCII characters. The // resulting bytes are valid JSON-RPC for use as an MCP request body. @@ -435,7 +723,7 @@ func TestMutatingMiddleware_RequestBodySizeLimit(t *testing.T) { body := makeJSONBodyOfSize(t, webhook.MaxRequestSize+1) req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) - ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: 1}) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: int64(1)}) req = req.WithContext(ctx) var nextCalled bool @@ -449,9 +737,9 @@ func TestMutatingMiddleware_RequestBodySizeLimit(t *testing.T) { var errResp map[string]interface{} require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) - errObj, ok := errResp["Error"].(map[string]interface{}) + errObj, ok := errResp["error"].(map[string]interface{}) require.True(t, ok) - assert.Equal(t, float64(http.StatusRequestEntityTooLarge), errObj["code"]) + assert.Equal(t, float64(mcp.CodeInvalidRequest), errObj["code"]) assert.Equal(t, "Request body exceeds maximum size", errObj["message"]) }) @@ -466,7 +754,7 @@ func TestMutatingMiddleware_RequestBodySizeLimit(t *testing.T) { body := makeJSONBodyOfSize(t, webhook.MaxRequestSize) req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) - ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: 1}) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: int64(1)}) req = req.WithContext(ctx) var nextCalled bool @@ -720,40 +1008,6 @@ func TestMutatingMiddleware_MalformedPatchJSON(t *testing.T) { assert.Equal(t, http.StatusInternalServerError, rr.Code) } -//nolint:paralleltest -func TestMutatingMiddleware_StringRequestID(t *testing.T) { - // Tests that the middleware correctly handles a string JSON-RPC ID. - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - resp := webhook.MutatingResponse{ - Response: webhook.Response{Version: webhook.APIVersion, UID: "uid", Allowed: false}, - } - w.Header().Set("Content-Type", "application/json") - _ = json.NewEncoder(w).Encode(resp) - })) - defer server.Close() - - cfg := makeConfig(server.URL, webhook.FailurePolicyFail) - mw := createMutatingHandler(makeExecutors(t, []webhook.Config{cfg}), "srv", "stdio") - - reqBody := []byte(`{"jsonrpc":"2.0","method":"tools/call","id":"string-id"}`) - req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(reqBody)) - // Use string ID in parsedMCP. - parsedMCP := &mcp.ParsedMCPRequest{Method: "tools/call", ID: "string-id"} - ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, parsedMCP) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - mw(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})).ServeHTTP(rr, req) - - assert.Equal(t, http.StatusForbidden, rr.Code) - assert.Contains(t, rr.Body.String(), "Request denied by webhook policy") - - // Confirm JSON-RPC error has the string ID. - var errResp map[string]interface{} - require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) - require.NotNil(t, errResp["ID"]) -} - //nolint:paralleltest func TestMutatingMiddleware_InvalidPatchOp_FailPolicy(t *testing.T) { // Returns a well-formed JSON array but with an invalid op type, so diff --git a/pkg/webhook/validating/middleware.go b/pkg/webhook/validating/middleware.go index e787bb6540..12ca23235f 100644 --- a/pkg/webhook/validating/middleware.go +++ b/pkg/webhook/validating/middleware.go @@ -93,7 +93,7 @@ func createValidatingHandler(executors []clientExecutor, serverName, transport s return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Skip if it's not a parsed MCP request or if this is a client response. parsedMCP := mcp.GetParsedMCPRequest(r.Context()) - if parsedMCP == nil || !parsedMCP.IsRequest { + if parsedMCP == nil || parsedMCP.IsClientResponse() { next.ServeHTTP(w, r) return } @@ -181,19 +181,21 @@ func readSourceIP(r *http.Request) string { } func sendErrorResponse(w http.ResponseWriter, statusCode int, message string, msgID interface{}) { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(statusCode) - id, err := mcp.ConvertToJSONRPC2ID(msgID) if err != nil { id = jsonrpc2.ID{} // Use empty ID if conversion fails } + code := mcp.JSONRPCCodeForStatus(statusCode) // Return a JSON-RPC 2.0 error so MCP clients can parse the denial. // The HTTP status code signals the error at the transport level; the JSON-RPC body carries the detail. errResp := &jsonrpc2.Response{ ID: id, - Error: jsonrpc2.NewError(int64(statusCode), message), + Error: jsonrpc2.NewError(code, message), } - _ = json.NewEncoder(w).Encode(errResp) + + // The helper encodes before writing any header, so an encode failure never + // leaves a half-written response (e.g. a status header followed by a + // second write). Nothing to do on a write error for a denial body. + _ = mcp.WriteJSONRPCError(w, statusCode, errResp) } diff --git a/pkg/webhook/validating/middleware_test.go b/pkg/webhook/validating/middleware_test.go index 45247884f9..20b2164a11 100644 --- a/pkg/webhook/validating/middleware_test.go +++ b/pkg/webhook/validating/middleware_test.go @@ -7,6 +7,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "testing" @@ -77,7 +78,7 @@ func TestValidatingMiddleware(t *testing.T) { // Add parsed MCP request and auth identity to context parsedMCP := &mcp.ParsedMCPRequest{ Method: "tools/call", - ID: 1, + ID: int64(1), IsRequest: true, } ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, parsedMCP) @@ -170,9 +171,9 @@ func TestValidatingMiddleware(t *testing.T) { err := json.Unmarshal(rr.Body.Bytes(), &errResp) require.NoError(t, err) - errObj, ok := errResp["Error"].(map[string]interface{}) + errObj, ok := errResp["error"].(map[string]interface{}) require.True(t, ok) - assert.Equal(t, float64(http.StatusForbidden), errObj["code"]) + assert.Equal(t, float64(mcp.JSONRPCCodeDenied), errObj["code"]) assert.Equal(t, "Request denied by policy", errObj["message"]) }) @@ -292,6 +293,76 @@ func TestValidatingMiddleware(t *testing.T) { }) } +// TestValidatingMiddlewareDenialIsConformantJSONRPC pins the sendErrorResponse envelope on +// the wire: lowercase "jsonrpc"/"id"/"error" keys, no "id" key for a notification, and no +// "result" key. assert.JSONEq is case-sensitive, so a single assertion catches "Error"/"ID" +// substituted for "error"/"id" (the historical bug) alongside a missing "jsonrpc" tag. +func TestValidatingMiddlewareDenialIsConformantJSONRPC(t *testing.T) { + t.Parallel() + + // Deterministic denial: the webhook always disallows, driving sendErrorResponse + // through the real middleware chain rather than constructing the response by hand. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(webhook.Response{ + Version: webhook.APIVersion, + UID: "resp-uid", + Allowed: false, + Message: "ignored", // sendErrorResponse ignores the webhook's message to avoid leaking info + }) + })) + t.Cleanup(server.Close) + + cfg := webhook.Config{ + Name: "test-webhook", + URL: server.URL, + Timeout: webhook.DefaultTimeout, + FailurePolicy: webhook.FailurePolicyFail, + TLSConfig: &webhook.TLSConfig{InsecureSkipVerify: true}, + } + client, err := webhook.NewClient(cfg, webhook.TypeValidating, nil) + require.NoError(t, err) + mw := createValidatingHandler([]clientExecutor{{client: client, config: cfg}}, "test-server", "stdio") + + const wantErr = `"error":{"code":403,"message":"Request denied by policy"}` + + testCases := []struct { + name string + id interface{} // ParsedMCPRequest.ID; nil for a notification + wantBody string + }{ + {name: "int64 id", id: int64(42), wantBody: `{"jsonrpc":"2.0","id":42,` + wantErr + `}`}, + {name: "string id", id: "abc", wantBody: `{"jsonrpc":"2.0","id":"abc",` + wantErr + `}`}, + {name: "nil id (notification) omits id entirely", id: nil, wantBody: `{"jsonrpc":"2.0",` + wantErr + `}`}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"tools/call"}`))) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: tc.id}) + req = req.WithContext(ctx) + + rr := httptest.NewRecorder() + mw(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("next handler must not be called for a denied request") + })).ServeHTTP(rr, req) + + require.Equal(t, http.StatusForbidden, rr.Code) + assert.Equal(t, "application/json", rr.Header().Get("Content-Type")) + assert.JSONEq(t, tc.wantBody, rr.Body.String()) + + var decoded map[string]any + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &decoded)) + _, hasID := decoded["id"] + assert.Equal(t, tc.id != nil, hasID, "id key presence mismatch") + _, hasResult := decoded["result"] + assert.False(t, hasResult, "denied response must not contain a result key") + }) + } +} + // makeJSONBodyOfSize builds a syntactically-valid JSON body of exactly `size` // bytes by padding the value of a "data" field with ASCII characters. The // resulting bytes are valid JSON-RPC for use as an MCP request body. @@ -334,7 +405,7 @@ func TestValidatingMiddleware_RequestBodySizeLimit(t *testing.T) { body := makeJSONBodyOfSize(t, webhook.MaxRequestSize+1) req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) - ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: 1, IsRequest: true}) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: int64(1), IsRequest: true}) req = req.WithContext(ctx) var nextCalled bool @@ -349,9 +420,9 @@ func TestValidatingMiddleware_RequestBodySizeLimit(t *testing.T) { // The error response is a JSON-RPC envelope. var errResp map[string]interface{} require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) - errObj, ok := errResp["Error"].(map[string]interface{}) + errObj, ok := errResp["error"].(map[string]interface{}) require.True(t, ok) - assert.Equal(t, float64(http.StatusRequestEntityTooLarge), errObj["code"]) + assert.Equal(t, float64(mcp.CodeInvalidRequest), errObj["code"]) assert.Equal(t, "Request body exceeds maximum size", errObj["message"]) }) @@ -374,7 +445,7 @@ func TestValidatingMiddleware_RequestBodySizeLimit(t *testing.T) { body := makeJSONBodyOfSize(t, webhook.MaxRequestSize) req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body)) - ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: 1, IsRequest: true}) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: int64(1), IsRequest: true}) req = req.WithContext(ctx) var nextCalled bool @@ -388,6 +459,42 @@ func TestValidatingMiddleware_RequestBodySizeLimit(t *testing.T) { }) } +// errReader always fails, for exercising the io.ReadAll error path in +// createValidatingHandler that is distinct from an http.MaxBytesError (the +// oversized-body case already covered by TestValidatingMiddleware_RequestBodySizeLimit). +type errReader struct{ err error } + +func (r errReader) Read([]byte) (int, error) { return 0, r.err } + +// TestValidatingMiddleware_BodyReadError_MapsToInternalErrorCode pins the JSON-RPC +// code on the "Failed to read request body" HTTP 500 path (sendErrorResponse's +// default case in mcp.JSONRPCCodeForStatus), the validating-package equivalent of +// TestMutatingMiddleware_ScopeViolation_FailPolicy's HTTP 500 assertion. +func TestValidatingMiddleware_BodyReadError_MapsToInternalErrorCode(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodPost, "/", errReader{err: errors.New("boom")}) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: int64(1)}) + req = req.WithContext(ctx) + + mw := createValidatingHandler(nil, "srv", "stdio") + + var nextCalled bool + nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { nextCalled = true }) + + rr := httptest.NewRecorder() + mw(nextHandler).ServeHTTP(rr, req) + + assert.False(t, nextCalled) + assert.Equal(t, http.StatusInternalServerError, rr.Code) + + var errResp map[string]interface{} + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + errObj, ok := errResp["error"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(mcp.CodeInternalError), errObj["code"]) +} + func TestMiddlewareParams_Validate(t *testing.T) { t.Parallel() tests := []struct { @@ -542,7 +649,7 @@ func TestCreateMiddleware_ResolvesHMACSecret(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"tools/call","id":1}`))) req = req.WithContext(context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{ Method: "tools/call", - ID: 1, + ID: int64(1), IsRequest: true, })) @@ -597,7 +704,7 @@ func TestValidatingMiddleware_HTTP422AlwaysDenies(t *testing.T) { reqBody := []byte(`{"jsonrpc":"2.0","method":"tools/call","id":1}`) req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(reqBody)) - ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: 1, IsRequest: true}) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, &mcp.ParsedMCPRequest{Method: "tools/call", ID: int64(1), IsRequest: true}) req = req.WithContext(ctx) var nextCalled bool @@ -707,8 +814,9 @@ func TestMultiWebhookChain(t *testing.T) { // Verify error response var errResp map[string]interface{} - _ = json.Unmarshal(rr.Body.Bytes(), &errResp) - errObj := errResp["Error"].(map[string]interface{}) + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + errObj, ok := errResp["error"].(map[string]interface{}) + require.True(t, ok) assert.Equal(t, "Request denied by policy", errObj["message"]) }) @@ -732,8 +840,9 @@ func TestMultiWebhookChain(t *testing.T) { // Verify error response var errResp map[string]interface{} - _ = json.Unmarshal(rr.Body.Bytes(), &errResp) - errObj := errResp["Error"].(map[string]interface{}) + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &errResp)) + errObj, ok := errResp["error"].(map[string]interface{}) + require.True(t, ok) assert.Equal(t, "Request denied by policy", errObj["message"]) }) diff --git a/pkg/workloads/manager.go b/pkg/workloads/manager.go index f208d96168..644abef2fb 100644 --- a/pkg/workloads/manager.go +++ b/pkg/workloads/manager.go @@ -793,7 +793,7 @@ func (d *DefaultManager) spawnDetached(ctx context.Context, runConfig *runner.Ru } if err := detachedCmd.Start(); err != nil { - if statusErr := d.statuses.SetWorkloadStatus(ctx, runConfig.BaseName, rt.WorkloadStatusError, ""); statusErr != nil { + if statusErr := d.statuses.SetWorkloadStatus(ctx, runConfig.BaseName, rt.WorkloadStatusError, err.Error()); statusErr != nil { slog.Warn("Failed to set workload status to error", "workload", runConfig.BaseName, "error", statusErr) } return 0, err @@ -1554,8 +1554,12 @@ func (d *DefaultManager) startWorkload(ctx context.Context, name string, mcpRunn } if err != nil { - // If we could not start the workload, set the status to error before returning - if statusErr := d.statuses.SetWorkloadStatus(ctx, name, rt.WorkloadStatusError, ""); statusErr != nil { + // If we could not start the workload, set the status to error before + // returning. Record err itself as the status context: RunWorkload's + // retry loop has already written the failure reason there, and an + // empty context here would clobber it — CI debug output then shows a + // bare status "error" with no way to tell what actually failed. + if statusErr := d.statuses.SetWorkloadStatus(ctx, name, rt.WorkloadStatusError, err.Error()); statusErr != nil { slog.Warn("Failed to set workload status to error", "workload", name, "error", statusErr) } } diff --git a/pkg/workloads/statuses/file_status.go b/pkg/workloads/statuses/file_status.go index b5f77e4f9c..d118476759 100644 --- a/pkg/workloads/statuses/file_status.go +++ b/pkg/workloads/statuses/file_status.go @@ -279,16 +279,26 @@ func (f *fileStatusManager) ListWorkloads(ctx context.Context, listAll bool, lab return nil, fmt.Errorf("failed to parse label filters: %w", err) } - // Get workloads from runtime - runtimeContainers, err := f.runtime.ListWorkloads(ctx) + // Get workloads from files BEFORE snapshotting the runtime. The order is + // load-bearing: a workload's status file only flips to `running` after its + // container has been created and started, so a file that says `running` in + // a snapshot taken at time T implies the container existed before T. If a + // later runtime snapshot lacks it, the container is genuinely gone and + // handleRuntimeMissing's unhealthy verdict is correct. With the snapshots + // the other way round, a workload that starts between the runtime snapshot + // and the file read looks like "file running, runtime missing" and gets its + // status file permanently overwritten as unhealthy — a healthy, freshly + // started workload poisoned by a concurrent `thv list` (see #4432 for the + // broader write-on-read hazard in this path). + fileWorkloadsWithPID, err := f.getWorkloadsFromFiles() if err != nil { - return nil, fmt.Errorf("failed to list workloads from runtime: %w", err) + return nil, fmt.Errorf("failed to get workloads from files: %w", err) } - // Get workloads from files - fileWorkloadsWithPID, err := f.getWorkloadsFromFiles() + // Get workloads from runtime + runtimeContainers, err := f.runtime.ListWorkloads(ctx) if err != nil { - return nil, fmt.Errorf("failed to get workloads from files: %w", err) + return nil, fmt.Errorf("failed to list workloads from runtime: %w", err) } // TODO: Fetch the runconfig if present to populate additional fields like package, tool type, group etc. @@ -973,9 +983,12 @@ func (f *fileStatusManager) handleRuntimeMissing( if fileWorkload.Status == rt.WorkloadStatusRunning || fileWorkload.Status == rt.WorkloadStatusStopped { // The workload cannot be running or stopped if the runtime container is not found contextMsg := fmt.Sprintf("workload %s not found in runtime, marking as unhealthy", workloadName) - if err := f.SetWorkloadStatus(ctx, workloadName, rt.WorkloadStatusUnhealthy, contextMsg); err != nil { + if err := f.setUnhealthyIfStillPresent(ctx, workloadName, contextMsg); err != nil { return core.Workload{}, err } + // The returned view reflects what this list observed regardless of + // whether the guarded write persisted anything (the file may have been + // removed or repurposed by a concurrent operation in the meantime). fileWorkload.Status = rt.WorkloadStatusUnhealthy } @@ -983,6 +996,41 @@ func (f *fileStatusManager) handleRuntimeMissing( return fileWorkload, nil } +// setUnhealthyIfStillPresent marks a workload unhealthy only if, re-checked +// under the file lock, its status file still exists and still claims running +// or stopped. handleRuntimeMissing decides to mark unhealthy from a snapshot +// that is stale by the time the write happens; an unconditional +// SetWorkloadStatus would recreate the file of a workload that `thv rm` +// deleted in the meantime (SetWorkloadStatus creates missing files), leaving +// a ghost "unhealthy" entry behind, or clobber a fresher status such as +// `removing`. Skipping in those cases is not an error: the persisted state is +// already newer than the observation that triggered the write. +func (f *fileStatusManager) setUnhealthyIfStillPresent(ctx context.Context, workloadName, contextMsg string) error { + return f.withFileLock(ctx, workloadName, func(statusFilePath string) error { + if _, err := os.Stat(statusFilePath); os.IsNotExist(err) { + slog.Debug("skipping unhealthy mark: status file no longer exists", "workload", workloadName) + return nil + } else if err != nil { + return fmt.Errorf("failed to check status file for workload %s: %w", workloadName, err) + } + + statusFile, err := f.readStatusFile(statusFilePath) + if err != nil { + return fmt.Errorf("failed to read existing status for workload %s: %w", workloadName, err) + } + if statusFile.Status != rt.WorkloadStatusRunning && statusFile.Status != rt.WorkloadStatusStopped { + slog.Debug("skipping unhealthy mark: status changed since observation", + "workload", workloadName, "status", statusFile.Status) + return nil + } + + statusFile.Status = rt.WorkloadStatusUnhealthy + statusFile.StatusContext = contextMsg + statusFile.UpdatedAt = time.Now() + return f.writeStatusFile(statusFilePath, *statusFile) + }) +} + // isProxyUnhealthy checks if the proxy process is running for the workload. // Returns (unhealthyWorkload, true) if proxy is not running, (emptyWorkload, false) if proxy is healthy or not applicable. func (f *fileStatusManager) isProxyUnhealthy( diff --git a/pkg/workloads/statuses/file_status_test.go b/pkg/workloads/statuses/file_status_test.go index 883ee0ba5e..d9546af162 100644 --- a/pkg/workloads/statuses/file_status_test.go +++ b/pkg/workloads/statuses/file_status_test.go @@ -1204,6 +1204,129 @@ func TestFileStatusManager_ListWorkloads_WithValidation(t *testing.T) { assert.True(t, isProxyValidated, "Proxy down workload should be detected as unhealthy") } +// TestFileStatusManager_ListWorkloads_DoesNotPoisonWorkloadStartingDuringList +// pins the snapshot ordering inside ListWorkloads: status files must be read +// before the runtime is queried. With the inverted order, a workload whose +// status file flips starting → running between the two snapshots (a `thv run` +// finishing concurrently) is seen as "file running, runtime missing", and +// handleRuntimeMissing permanently overwrites its status file as unhealthy — +// poisoning a healthy, freshly started workload. The mock flips the file to +// running from inside the runtime-list call, which under the correct ordering +// is the later of the two snapshots. +func TestFileStatusManager_ListWorkloads_DoesNotPoisonWorkloadStartingDuringList(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + manager, mockRuntime, mockRunConfigStore := newTestFileStatusManager(t, ctrl) + ctx := context.Background() + + mockRunConfigStore.EXPECT().Exists(gomock.Any(), gomock.Any()).Return(false, nil).AnyTimes() + + const workloadName = "starting-workload" + require.NoError(t, manager.SetWorkloadStatus(ctx, workloadName, rt.WorkloadStatusStarting, "container starting")) + + // Simulate the concurrent startup completing between the file snapshot and + // the runtime snapshot: flip the file to running and return a runtime view + // from before the container existed. + mockRuntime.EXPECT().ListWorkloads(gomock.Any()).DoAndReturn(func(callCtx context.Context) ([]rt.ContainerInfo, error) { + require.NoError(t, manager.SetWorkloadStatus(callCtx, workloadName, rt.WorkloadStatusRunning, "")) + return []rt.ContainerInfo{}, nil + }) + + workloads, err := manager.ListWorkloads(ctx, true, nil) + require.NoError(t, err) + require.Len(t, workloads, 1) + assert.Equal(t, workloadName, workloads[0].Name) + assert.Equal(t, rt.WorkloadStatusStarting, workloads[0].Status, + "a workload still starting in the file snapshot must be reported as starting, not poisoned") + + // The persisted status file must still say running: ListWorkloads must not + // have overwritten the concurrent writer's state with unhealthy. + data, err := os.ReadFile(manager.getStatusFilePath(workloadName)) + require.NoError(t, err) + var persisted workloadStatusFile + require.NoError(t, json.Unmarshal(data, &persisted)) + assert.Equal(t, rt.WorkloadStatusRunning, persisted.Status, + "ListWorkloads must not overwrite the status file of a workload that finished starting mid-list") +} + +// TestFileStatusManager_ListWorkloads_UnhealthyMarkIsGuarded pins the guarded +// write in handleRuntimeMissing: by the time the merge decides to mark a +// workload unhealthy, its snapshot is stale, and a concurrent `thv rm` may +// have deleted the status file (an unconditional SetWorkloadStatus would +// recreate it as a ghost) or moved it to `removing` (which must not be +// clobbered by a stale unhealthy verdict). Both mutations are injected from +// inside the runtime-list call, which is the later snapshot. +func TestFileStatusManager_ListWorkloads_UnhealthyMarkIsGuarded(t *testing.T) { + t.Parallel() + + const workloadName = "guarded-workload" + + tests := []struct { + name string + mutate func(ctx context.Context, m *fileStatusManager) + wantFileExists bool + wantFileStatus rt.WorkloadStatus + wantFileContext string + }{ + { + name: "file deleted mid-list is not resurrected", + mutate: func(ctx context.Context, m *fileStatusManager) { + require.NoError(t, m.DeleteWorkloadStatus(ctx, workloadName)) + }, + wantFileExists: false, + }, + { + name: "status moved to removing mid-list is not clobbered", + mutate: func(ctx context.Context, m *fileStatusManager) { + require.NoError(t, m.SetWorkloadStatus(ctx, workloadName, rt.WorkloadStatusRemoving, "being removed")) + }, + wantFileExists: true, + wantFileStatus: rt.WorkloadStatusRemoving, + wantFileContext: "being removed", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + manager, mockRuntime, mockRunConfigStore := newTestFileStatusManager(t, ctrl) + ctx := context.Background() + + mockRunConfigStore.EXPECT().Exists(gomock.Any(), gomock.Any()).Return(false, nil).AnyTimes() + + require.NoError(t, manager.SetWorkloadStatus(ctx, workloadName, rt.WorkloadStatusRunning, "container started")) + + mockRuntime.EXPECT().ListWorkloads(gomock.Any()).DoAndReturn(func(callCtx context.Context) ([]rt.ContainerInfo, error) { + tt.mutate(callCtx, manager) + return []rt.ContainerInfo{}, nil + }) + + _, err := manager.ListWorkloads(ctx, true, nil) + require.NoError(t, err) + + data, err := os.ReadFile(manager.getStatusFilePath(workloadName)) + if !tt.wantFileExists { + require.True(t, os.IsNotExist(err), + "status file deleted by a concurrent rm must not be recreated by ListWorkloads") + return + } + require.NoError(t, err) + var persisted workloadStatusFile + require.NoError(t, json.Unmarshal(data, &persisted)) + assert.Equal(t, tt.wantFileStatus, persisted.Status, + "a status written after the merge's observation must not be clobbered") + assert.Equal(t, tt.wantFileContext, persisted.StatusContext) + }) + } +} + func TestFileStatusManager_GetWorkload_vs_ListWorkloads_Consistency(t *testing.T) { t.Parallel() diff --git a/pkg/workloads/upgrade/applier.go b/pkg/workloads/upgrade/applier.go index 1c394d5ca6..a60b9d10d7 100644 --- a/pkg/workloads/upgrade/applier.go +++ b/pkg/workloads/upgrade/applier.go @@ -301,9 +301,13 @@ func (a *Applier) buildUpgradedConfig( runner.WithHost(old.Host), runner.WithVolumes(slices.Clone(old.Volumes)), runner.WithSecrets(mergedSecrets), + // Intentionally NOT paired with WithNetworkIsolationExplicit: a legacy + // config with isolation + a host/none network mode should degrade-and-warn + // on upgrade, not hard-fail. See #5775. runner.WithNetworkIsolation(old.IsolateNetwork), runner.WithAllowDockerGateway(old.AllowDockerGateway), runner.WithTrustProxyHeaders(old.TrustProxyHeaders), + runner.WithStrictProtocolValidation(old.StrictProtocolValidation), runner.WithProxyMode(old.ProxyMode), runner.WithCmdArgs(slices.Clone(old.CmdArgs)), runner.WithStateless(old.Stateless), diff --git a/pkg/workloads/upgrade/applier_test.go b/pkg/workloads/upgrade/applier_test.go index cc11fc1a05..52e03a4ed9 100644 --- a/pkg/workloads/upgrade/applier_test.go +++ b/pkg/workloads/upgrade/applier_test.go @@ -507,6 +507,7 @@ func fullyConfiguredOld(t *testing.T) *runner.RunConfig { PermissionProfileNameOrPath: "none", IsolateNetwork: true, TrustProxyHeaders: true, + StrictProtocolValidation: true, Stateless: true, EndpointPrefix: "/mcp", SessionTTL: "30m", @@ -569,6 +570,7 @@ func TestApplier_Apply_PreservesFullUserConfig(t *testing.T) { assert.Equal(t, old.PermissionProfileNameOrPath, got.PermissionProfileNameOrPath, "named profile preserved") assert.Equal(t, old.IsolateNetwork, got.IsolateNetwork) assert.Equal(t, old.TrustProxyHeaders, got.TrustProxyHeaders) + assert.Equal(t, old.StrictProtocolValidation, got.StrictProtocolValidation) assert.Equal(t, old.Stateless, got.Stateless) assert.Equal(t, old.EndpointPrefix, got.EndpointPrefix) assert.Equal(t, old.SessionTTL, got.SessionTTL) @@ -620,3 +622,46 @@ func TestApplier_Apply_PreservesHTTPTransportPorts(t *testing.T) { assert.Equal(t, 21000, h.updatedConfig.Port, "existing proxy port preserved over candidate's") assert.Equal(t, 22000, h.updatedConfig.TargetPort, "existing target port preserved over candidate's") } + +// TestApplier_Apply_DegradesLegacyHostIsolation verifies that upgrading a legacy +// config that persisted network isolation together with a host network mode +// degrades (isolation dropped) rather than failing: the upgrade path +// intentionally does not pass WithNetworkIsolationExplicit. See #5775. +func TestApplier_Apply_DegradesLegacyHostIsolation(t *testing.T) { + t.Parallel() + h := newApplierHarness(t) + + old := &runner.RunConfig{ + Name: "wl", + Image: "ghcr.io/example/server:1.0.0", + RegistryServerName: applyServerName, + Group: "default", + Transport: transporttypes.TransportTypeStreamableHTTP, + ProxyMode: transporttypes.ProxyModeStreamableHTTP, + Port: 21000, + TargetPort: 22000, + IsolateNetwork: true, + PermissionProfile: &permissions.Profile{ + Network: &permissions.NetworkPermissions{Mode: "host"}, + }, + EnvVars: map[string]string{}, + } + h.loadConfig = old + h.provider.EXPECT(). + GetServer(applyServerName). + Return(®types.ImageMetadata{Image: "ghcr.io/example/server:1.2.0"}, nil) + candidate := ®types.ImageMetadata{Image: "ghcr.io/example/server:1.2.0"} + candidate.Transport = string(transporttypes.TransportTypeStreamableHTTP) + h.resolveImageURL = "ghcr.io/example/server:1.2.0" + h.resolveMeta = candidate + h.configMock.EXPECT().GetConfig().Return(&appconfig.Config{}).AnyTimes() + h.expectUpdate() + + _, err := h.applier.Apply(context.Background(), "wl", ApplyOptions{ + EnvVarValidator: &runner.DetachedEnvVarValidator{}, + }) + require.NoError(t, err, "legacy host+isolation config must degrade, not fail, on upgrade") + require.NotNil(t, h.updatedConfig) + assert.False(t, h.updatedConfig.IsolateNetwork, + "isolation must be dropped when upgrading a host-network workload") +} diff --git a/test/conformance/README.md b/test/conformance/README.md index ee3d319f58..7e8e8c33bf 100644 --- a/test/conformance/README.md +++ b/test/conformance/README.md @@ -42,9 +42,30 @@ Lists scenarios that are known to fail and accepted. CI fails only on failures and also fails CI, so keep the list accurate. The CLI (`thv run`) proxy is transparent, so the list is currently empty. +## Targets + +`run-conformance.sh` supports two targets via `CONFORMANCE_TARGET`: + +- **`proxy`** (default, `task conformance`) — runs the reference server through + `thv run` and tests the transparent CLI proxy endpoint. This path is + transport-transparent and does **not** exercise ToolHive's own MCP serving + code, so its baseline (`expected-failures.yaml`) is empty. +- **`vmcp`** (`task conformance-vmcp`) — runs the reference server as a backend + in a group, aggregates it through `thv vmcp serve`, and tests the vMCP + endpoint. vMCP terminates and re-serves MCP using the migrated + `mcpcompat`/go-sdk serving surface, so this target is what actually exercises + that code. The generated config is patched from the `prefix` conflict strategy + to `priority` so backend tool names are preserved (unprefixed) for the suite's + fixtures, and the run asserts `serverInfo.name` identifies vMCP before testing. + Its baseline is `expected-failures-vmcp.yaml`. + ## Scope -Covers the CLI (`thv run`) proxy path with the `active` suite. Not covered: +Covers the CLI (`thv run`) proxy path and the vMCP aggregator path +(`thv vmcp serve`) with the `active` suite. The vmcp target is what exercises +ToolHive's `mcpcompat`/go-sdk serving surface; its baseline lives in +`expected-failures-vmcp.yaml` (sampling, completions, and resource subscribe are +expected to fail until vMCP supports them). Not covered: - **Kubernetes operator proxy** — it rewrites the outbound `Host` header to the upstream Service identity, which breaks MCP servers that validate `Host` (per diff --git a/test/conformance/expected-failures-vmcp.yaml b/test/conformance/expected-failures-vmcp.yaml new file mode 100644 index 0000000000..1a07dcef92 --- /dev/null +++ b/test/conformance/expected-failures-vmcp.yaml @@ -0,0 +1,16 @@ +# Conformance baseline for the ToolHive vMCP aggregator (`thv vmcp serve`). +# +# vMCP terminates the MCP protocol and re-serves an aggregated view of its +# backends. This PR makes it conformant — completions, resource templates, +# resource subscribe/unsubscribe (ack-level), and mid-call forwarding of +# elicitation/sampling/progress/logging (and embedded-resource decode via +# toolhive-core v0.0.32) — so the previously-failing scenarios are expected to +# PASS now, and this list is empty. +# +# List a scenario name here only if it is KNOWN to fail and accepted. The suite +# fails CI on any failure NOT listed here; an entry that starts passing again is +# reported as STALE and also fails CI, so keep this accurate. If the CI vMCP leg +# surfaces a genuinely-unfixed scenario (e.g. resource subscribe requiring an +# actual notifications/resources/updated rather than just the ack), add it here +# with a comment pointing to the follow-up. +server: [] diff --git a/test/conformance/expected-failures.yaml b/test/conformance/expected-failures.yaml index 6116d79337..31cf8e7d6a 100644 --- a/test/conformance/expected-failures.yaml +++ b/test/conformance/expected-failures.yaml @@ -6,4 +6,10 @@ # # The CLI proxy is transport-transparent, so this list is empty: every scenario # in the `active` suite is expected to pass. +# +# NOTE: the server-INITIATED-request scenarios tools-call-sampling and +# tools-call-elicitation flake intermittently due to an upstream reference-server +# bug (conformance#407), NOT a ToolHive defect, so they are quarantined in the +# gating logic in run-conformance.sh rather than listed here — a flaky entry here +# would stale-fail the runs where it passes. server: [] diff --git a/test/conformance/run-conformance.sh b/test/conformance/run-conformance.sh index 5749fc6cd9..07ad488cea 100755 --- a/test/conformance/run-conformance.sh +++ b/test/conformance/run-conformance.sh @@ -2,39 +2,86 @@ # Run the MCP conformance suite against a ToolHive deployment. # # Packages the conformance repo's reference server (examples/servers/typescript) -# dynamically at a pinned version, runs it through `thv run` (streamable-http), -# and points `npx @modelcontextprotocol/conformance` at the proxied endpoint. +# dynamically at a pinned version and points `npx @modelcontextprotocol/conformance` +# at a ToolHive-exposed endpoint. Two targets are supported via CONFORMANCE_TARGET: +# +# proxy (default) Runs the reference server through `thv run` (streamable-http) +# and tests the transparent CLI proxy endpoint. This path does +# not exercise the migrated mcpcompat/go-sdk serving code. +# vmcp Runs the reference server as a backend in a group, aggregates +# it through `thv vmcp serve`, and tests the vMCP endpoint. This +# path DOES exercise the mcpcompat/go-sdk serving surface. # # A single CONFORMANCE_VERSION pins both the npm tool and the git server source # (repo tags vX.Y.Z map 1:1 to npm versions), keeping fixtures and tool in sync. # # Env: -# CONFORMANCE_VERSION conformance tool + server version (default: 0.1.16) -# CONFORMANCE_SUITE suite to run: active|all|pending (default: active) -# THV_BINARY path to the thv binary (default: thv) +# CONFORMANCE_TARGET deployment under test: proxy|vmcp (default: proxy) +# CONFORMANCE_VERSION conformance tool + server version (default: 0.1.16) +# CONFORMANCE_SUITE suite to run: active|all|pending (default: active) +# THV_BINARY path to the thv binary (default: thv) +# VMCP_PORT port for `thv vmcp serve` (vmcp target) (default: 8099) set -euo pipefail +CONFORMANCE_TARGET="${CONFORMANCE_TARGET:-proxy}" CONFORMANCE_VERSION="${CONFORMANCE_VERSION:-0.1.16}" CONFORMANCE_SUITE="${CONFORMANCE_SUITE:-active}" THV_BINARY="${THV_BINARY:-thv}" +VMCP_PORT="${VMCP_PORT:-8099}" SERVER_NAME="conf-sut-ci" +GROUP_NAME="conf-vmcp-group" IMAGE="mcp-conformance-server:ci" SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" RESULTS_DIR="${SCRIPT_DIR}/results" -EXPECTED_FAILURES="${SCRIPT_DIR}/expected-failures.yaml" # thv proxying to a container is a local-only, trusted-dev workflow. export TOOLHIVE_DEV=true export TOOLHIVE_SKIP_DESKTOP_CHECK=1 +# Shared JSON-RPC initialize payload used by the readiness probe and the vMCP +# serverInfo sanity check. +INIT_PAYLOAD='{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"probe","version":"1"}}}' + CLONE_DIR="" +VMCP_PID="" cleanup() { - "${THV_BINARY}" rm "${SERVER_NAME}" >/dev/null 2>&1 || true + # Diagnostics: capture the ToolHive proxy logs before teardown so a CI failure + # (e.g. the sampling round-trip timing out) is debuggable from the uploaded + # results artifact, not just the client-side checks.json. + mkdir -p "${RESULTS_DIR}" + "${THV_BINARY}" logs --proxy "${SERVER_NAME}" > "${RESULTS_DIR}/proxy-logs.txt" 2>&1 || true + if [ -n "${VMCP_PID}" ]; then + kill "${VMCP_PID}" >/dev/null 2>&1 || true + fi + "${THV_BINARY}" rm -f "${SERVER_NAME}" >/dev/null 2>&1 || true + if [ "${CONFORMANCE_TARGET}" = "vmcp" ]; then + # `group rm` prompts for confirmation only when the group still has + # workloads; feed "y" so the non-interactive EXIT trap can never block. + echo y | "${THV_BINARY}" group rm "${GROUP_NAME}" >/dev/null 2>&1 || true + fi [ -n "${CLONE_DIR}" ] && rm -rf "${CLONE_DIR}" || true } trap cleanup EXIT +# resolve_workload_url prints the proxy URL of a running workload, polling until +# it appears (or the retry budget is exhausted, in which case it prints nothing). +resolve_workload_url() { + local name="$1" url="" + for _ in $(seq 1 30); do + url="$("${THV_BINARY}" list --format json 2>/dev/null \ + | python3 -c "import sys,json +d=json.load(sys.stdin) +print(next((w['url'] for w in d if w.get('name')=='${name}' and w.get('status')=='running' and w.get('url')), ''))" \ + 2>/dev/null || true)" + [ -n "${url}" ] && break + sleep 2 + done + printf '%s' "${url}" +} + +echo "==> Target: ${CONFORMANCE_TARGET}" + echo "==> Cloning conformance repo v${CONFORMANCE_VERSION}" CLONE_DIR="$(mktemp -d)" for attempt in 1 2 3; do @@ -58,39 +105,194 @@ echo "==> Building reference server image" cp "${SCRIPT_DIR}/Dockerfile" "${SERVER_DIR}/Dockerfile" docker build -t "${IMAGE}" "${SERVER_DIR}" -echo "==> Starting server through ToolHive" -"${THV_BINARY}" rm -f "${SERVER_NAME}" >/dev/null 2>&1 || true -"${THV_BINARY}" run "${IMAGE}" \ - --transport streamable-http --target-port 3000 --name "${SERVER_NAME}" +if [ "${CONFORMANCE_TARGET}" = "proxy" ]; then + # --------------------------------------------------------------------------- + # proxy target: test the transparent `thv run` CLI proxy endpoint directly. + # --------------------------------------------------------------------------- + EXPECTED_FAILURES="${SCRIPT_DIR}/expected-failures.yaml" -echo "==> Resolving proxy URL" -URL="" -for _ in $(seq 1 30); do - URL="$("${THV_BINARY}" list --format json 2>/dev/null \ - | python3 -c "import sys,json; d=json.load(sys.stdin) -print(next((w['url'] for w in d if w.get('name')=='${SERVER_NAME}' and w.get('status')=='running' and w.get('url')), ''))" \ - 2>/dev/null || true)" - [ -n "${URL}" ] && break - sleep 2 -done -[ -z "${URL}" ] && { echo "ERROR: could not resolve ${SERVER_NAME} URL" >&2; "${THV_BINARY}" list || true; exit 1; } -echo " URL: ${URL}" + echo "==> Starting server through ToolHive" + "${THV_BINARY}" rm -f "${SERVER_NAME}" >/dev/null 2>&1 || true + "${THV_BINARY}" run "${IMAGE}" \ + --transport streamable-http --target-port 3000 --name "${SERVER_NAME}" --debug + + echo "==> Resolving proxy URL" + URL="$(resolve_workload_url "${SERVER_NAME}")" + [ -z "${URL}" ] && { echo "ERROR: could not resolve ${SERVER_NAME} URL" >&2; "${THV_BINARY}" list || true; exit 1; } + echo " URL: ${URL}" + +elif [ "${CONFORMANCE_TARGET}" = "vmcp" ]; then + # --------------------------------------------------------------------------- + # vmcp target: run the reference server as a group backend, aggregate it + # through vMCP, and test the vMCP endpoint (exercises mcpcompat/go-sdk). + # --------------------------------------------------------------------------- + EXPECTED_FAILURES="${SCRIPT_DIR}/expected-failures-vmcp.yaml" + VMCP_CONFIG="${CLONE_DIR}/vmcp.yaml" + + echo "==> Creating group ${GROUP_NAME}" + echo y | "${THV_BINARY}" group rm "${GROUP_NAME}" >/dev/null 2>&1 || true + "${THV_BINARY}" group create "${GROUP_NAME}" + + echo "==> Starting backend workload through ToolHive" + "${THV_BINARY}" rm -f "${SERVER_NAME}" >/dev/null 2>&1 || true + "${THV_BINARY}" run "${IMAGE}" \ + --transport streamable-http --target-port 3000 \ + --name "${SERVER_NAME}" --group "${GROUP_NAME}" + + echo "==> Waiting for backend to become running (so vmcp init can discover it)" + BACKEND_URL="$(resolve_workload_url "${SERVER_NAME}")" + [ -z "${BACKEND_URL}" ] && { echo "ERROR: backend ${SERVER_NAME} never became running" >&2; "${THV_BINARY}" list || true; exit 1; } + echo " backend URL: ${BACKEND_URL}" + + echo "==> Generating vMCP config with thv vmcp init" + "${THV_BINARY}" vmcp init --group "${GROUP_NAME}" --config "${VMCP_CONFIG}" + + # Patch the generated config so tool names are PRESERVED (not prefixed with the + # backend name), otherwise the conformance suite's tool fixtures won't match. + # `thv vmcp init` emits: + # aggregation: + # conflictResolution: prefix + # conflictResolutionConfig: + # prefixFormat: "{workload}_" + # We rewrite it to the "priority" strategy with a single-backend priorityOrder: + # aggregation: + # conflictResolution: priority + # conflictResolutionConfig: + # priorityOrder: + # - conf-sut-ci + # PyYAML may be absent, so this is a plain-text substitution (no yaml module). + echo "==> Patching aggregation strategy prefix -> priority" + python3 - "${VMCP_CONFIG}" "${SERVER_NAME}" <<'PY' +import sys +path, server = sys.argv[1], sys.argv[2] +with open(path, encoding="utf-8") as f: + text = f.read() +text = text.replace( + " conflictResolution: prefix", + " conflictResolution: priority", +) +text = text.replace( + ' prefixFormat: "{workload}_"', + " priorityOrder:\n - " + server, +) +if " conflictResolution: priority" not in text: + sys.exit("ERROR: could not patch conflictResolution to priority (template changed?)") +if "prefixFormat" in text: + sys.exit("ERROR: prefixFormat still present after patch (template changed?)") +with open(path, "w", encoding="utf-8") as f: + f.write(text) +PY + + echo "==> Patched vMCP config:" + echo "-----------------------------------------------------------------" + cat "${VMCP_CONFIG}" + echo "-----------------------------------------------------------------" + + echo "==> Validating patched config" + "${THV_BINARY}" vmcp validate --config "${VMCP_CONFIG}" + + echo "==> Starting thv vmcp serve on port ${VMCP_PORT}" + "${THV_BINARY}" vmcp serve --config "${VMCP_CONFIG}" --port "${VMCP_PORT}" & + VMCP_PID=$! + echo " vmcp serve PID: ${VMCP_PID}" + + URL="http://127.0.0.1:${VMCP_PORT}/mcp" + echo " URL: ${URL}" + +else + echo "ERROR: unknown CONFORMANCE_TARGET '${CONFORMANCE_TARGET}' (expected proxy|vmcp)" >&2 + exit 1 +fi echo "==> Waiting for endpoint to be ready" +code="" for _ in $(seq 1 30); do code="$(curl -s -o /dev/null -w '%{http_code}' "${URL}" -X POST \ -H 'content-type: application/json' \ -H 'accept: application/json, text/event-stream' \ - -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"probe","version":"1"}}}' || true)" + -d "${INIT_PAYLOAD}" || true)" [ "${code}" = "200" ] && break sleep 2 done [ "${code}" = "200" ] || { echo "ERROR: endpoint never became ready (last HTTP ${code})" >&2; exit 1; } -echo "==> Running conformance suite (${CONFORMANCE_SUITE})" +# For the vmcp target, prove we actually hit vMCP and not a mis-resolved URL by +# asserting the initialize response identifies vMCP in result.serverInfo.name. +if [ "${CONFORMANCE_TARGET}" = "vmcp" ]; then + echo "==> Sanity check: verifying serverInfo identifies vMCP" + INIT_RESP="$(curl -s "${URL}" -X POST \ + -H 'content-type: application/json' \ + -H 'accept: application/json, text/event-stream' \ + -d "${INIT_PAYLOAD}" || true)" + SERVER_INFO_NAME="$(printf '%s' "${INIT_RESP}" | python3 -c " +import sys, json +raw = sys.stdin.read() +# streamable-http may frame the response as SSE ('data: {json}' lines). +chunks = [l[5:].strip() for l in raw.splitlines() if l.strip().startswith('data:')] +for c in (chunks or [raw]): + try: + d = json.loads(c) + except Exception: + continue + info = (d.get('result') or {}).get('serverInfo') or {} + if info.get('name'): + print(info['name']) + break +" 2>/dev/null || true)" + echo " serverInfo.name: ${SERVER_INFO_NAME:-}" + case "$(printf '%s' "${SERVER_INFO_NAME}" | tr '[:upper:]' '[:lower:]')" in + *vmcp*|*virtual*) + echo " OK: endpoint identifies as vMCP" + ;; + *) + echo "ERROR: serverInfo.name '${SERVER_INFO_NAME}' does not identify vMCP;" \ + "the URL may not point at the vMCP server" >&2 + echo " raw initialize response: ${INIT_RESP}" >&2 + exit 1 + ;; + esac +fi + +echo "==> Running conformance suite (${CONFORMANCE_SUITE}) against ${CONFORMANCE_TARGET}" +echo " expected-failures: ${EXPECTED_FAILURES}" rm -rf "${RESULTS_DIR}" +mkdir -p "${RESULTS_DIR}" +set +e npx -y "@modelcontextprotocol/conformance@${CONFORMANCE_VERSION}" server \ --url "${URL}" \ --suite "${CONFORMANCE_SUITE}" \ --expected-failures "${EXPECTED_FAILURES}" \ - -o "${RESULTS_DIR}" + -o "${RESULTS_DIR}" 2>&1 | tee "${RESULTS_DIR}/suite-output.log" +suite_rc=${PIPESTATUS[0]} +set -e + +# Quarantine: the scenarios that drive a server->client request flake +# intermittently (~1/3) due to an upstream reference-server bug, NOT a ToolHive +# defect. The everything-server sends its server->client request +# (sampling/createMessage or elicitation/create) on the client's standalone GET +# SSE stream without a relatedRequestId, so it races the tools/call handler; +# when the handler wins, the request is dropped and the client times out at 60s. +# In suite v0.1.16 exactly four scenarios register client-side handlers for +# server->client requests (src/scenarios/server/tools.ts and +# elicitation-{defaults,enums}.ts) and are therefore exposed to the race: +# tools-call-sampling, tools-call-elicitation, +# elicitation-sep1034-defaults, elicitation-sep1330-enums +# See: +# upstream: https://github.com/modelcontextprotocol/conformance/issues/407 +# toolhive: #5886 (ingress Squid cold-start mitigation, reduces but can't +# eliminate the client-side ordering race) +# We can't use expected-failures.yaml (a flaky entry stale-fails the ~2/3 of runs +# where it passes), so ignore ONLY these scenarios in the gating here; every +# other scenario still fails the job. Remove this block once #407 is fixed +# upstream and CONFORMANCE_VERSION is bumped to a release that contains the fix. +quarantined='tools-call-sampling|tools-call-elicitation|elicitation-sep1034-defaults|elicitation-sep1330-enums' +if [ "${suite_rc}" -ne 0 ]; then + unexpected="$(sed -n '/Unexpected failures (not in baseline):/,$p' "${RESULTS_DIR}/suite-output.log" \ + | grep -oE '✗ [a-z0-9-]+' | sed 's/✗ //' | sort -u)" + others="$(printf '%s\n' "${unexpected}" | grep -vE "^(${quarantined})?\$" || true)" + if [ -n "${unexpected}" ] && [ -z "${others}" ]; then + echo "==> Ignoring quarantined flaky server->client scenarios (${quarantined}) per upstream conformance#407 / toolhive#5886; no other unexpected failures." + exit 0 + fi + exit "${suite_rc}" +fi diff --git a/test/e2e/api_helpers.go b/test/e2e/api_helpers.go index 7b509a2032..3ab58503fa 100644 --- a/test/e2e/api_helpers.go +++ b/test/e2e/api_helpers.go @@ -27,6 +27,10 @@ type ServerConfig struct { StartTimeout time.Duration RequestTimeout time.Duration DebugMode bool + // ExtraEnv adds "KEY=VALUE" environment variables to the `thv serve` + // subprocess, for tests that need to opt into experimental or + // feature-gated behavior without changing it for every other E2E test. + ExtraEnv []string } // NewServerConfig creates a new API server configuration with defaults @@ -53,6 +57,31 @@ type Server struct { stdout *strings.Builder } +// RetryOnPortConflict runs attempt with port and returns its result. The +// port is only probed free at selection time (networking.FindOrUsePort); +// the real bind happens later, in a subprocess this test doesn't control, so +// another process can steal the port before that bind occurs. If attempt +// fails with an address-already-in-use signal, RetryOnPortConflict +// allocates a fresh port and retries attempt exactly once -- mirroring +// startEraBackendOnPort's handling of the same TOCTOU window for +// Docker-bound ports in test/e2e/vmcp_dual_era_helpers_test.go. It returns +// the port actually used (the original, or the fresh one on retry) along +// with attempt's result. +func RetryOnPortConflict(port int, attempt func(port int) error) (int, error) { + err := attempt(port) + if err == nil || !strings.Contains(err.Error(), "address already in use") { + return port, err + } + + newPort, ferr := networking.FindOrUsePort(0) + if ferr != nil { + GinkgoWriter.Printf("failed to find a fresh port after port %d lost a bind race: %v\n", port, ferr) + return port, err + } + GinkgoWriter.Printf("port %d lost a bind race, retrying once with fresh port %d: %v\n", port, newPort, err) + return newPort, attempt(newPort) +} + // NewServer creates and starts a new API server instance by running `thv serve` as a subprocess. func NewServer(config *ServerConfig) (*Server, error) { testConfig := NewTestConfig() @@ -83,54 +112,63 @@ func NewServer(config *ServerConfig) (*Server, error) { // would fail because no client config paths exist in the temp home dir. _ = os.WriteFile(filepath.Join(tempHome, ".claude.json"), []byte("{}"), 0600) - ctx, cancel := context.WithCancel(context.Background()) - - // Create string builders to capture output - var stdout, stderr strings.Builder - - // Create the command: thv serve --host 127.0.0.1 --port - //nolint:gosec // Intentional for e2e testing - cmd := exec.CommandContext( - ctx, - testConfig.THVBinary, - "serve", - "--host", - "127.0.0.1", - "--port", - strconv.Itoa(port), - ) - // Set environment variables including temporary config paths - cmd.Env = append([]string{ - "TOOLHIVE_DEV=true", - fmt.Sprintf("XDG_CONFIG_HOME=%s", tempXdgConfigHome), - fmt.Sprintf("HOME=%s", tempHome), - }, cmd.Env...) - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - // Start the server process - if err := cmd.Start(); err != nil { - cancel() - return nil, fmt.Errorf("failed to start thv serve: %w", err) - } + var server *Server + // The returned port is discarded here: the closure below assigns the + // real port onto server.port via the Server struct literal. + _, err = RetryOnPortConflict(port, func(p int) error { + ctx, cancel := context.WithCancel(context.Background()) + + // Create string builders to capture output + var stdout, stderr strings.Builder + + // Create the command: thv serve --host 127.0.0.1 --port + //nolint:gosec // Intentional for e2e testing + cmd := exec.CommandContext( + ctx, + testConfig.THVBinary, + "serve", + "--host", + "127.0.0.1", + "--port", + strconv.Itoa(p), + ) + // Set environment variables including temporary config paths + cmd.Env = append([]string{ + "TOOLHIVE_DEV=true", + fmt.Sprintf("XDG_CONFIG_HOME=%s", tempXdgConfigHome), + fmt.Sprintf("HOME=%s", tempHome), + }, config.ExtraEnv...) + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + // Start the server process + if err := cmd.Start(); err != nil { + cancel() + return fmt.Errorf("failed to start thv serve: %w", err) + } - server := &Server{ - config: config, - baseURL: fmt.Sprintf("http://127.0.0.1:%d", port), - cmd: cmd, - ctx: ctx, - cancel: cancel, - httpClient: &http.Client{ - Timeout: config.RequestTimeout, - }, - port: port, - stdout: &stdout, - stderr: &stderr, - } + server = &Server{ + config: config, + baseURL: fmt.Sprintf("http://127.0.0.1:%d", p), + cmd: cmd, + ctx: ctx, + cancel: cancel, + httpClient: &http.Client{ + Timeout: config.RequestTimeout, + }, + port: p, + stdout: &stdout, + stderr: &stderr, + } - // Wait for server to be ready - if err := server.WaitForReady(); err != nil { - _ = server.Stop() + // Wait for server to be ready + if err := server.WaitForReady(); err != nil { + _ = server.Stop() + return err + } + return nil + }) + if err != nil { return nil, err } diff --git a/test/e2e/api_skills_test.go b/test/e2e/api_skills_test.go index 9ca6d33c1a..3dbd668dc8 100644 --- a/test/e2e/api_skills_test.go +++ b/test/e2e/api_skills_test.go @@ -9,6 +9,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "net/url" "os" "path/filepath" "time" @@ -17,6 +18,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/stacklok/toolhive/pkg/skills/lockfile" "github.com/stacklok/toolhive/test/e2e" ) @@ -47,13 +49,14 @@ type skillMetadataResponse struct { } type installSkillRequest struct { - Name string `json:"name"` - Version string `json:"version,omitempty"` - Scope string `json:"scope,omitempty"` - ProjectRoot string `json:"project_root,omitempty"` - Client string `json:"client,omitempty"` - Force bool `json:"force,omitempty"` - Group string `json:"group,omitempty"` + Name string `json:"name"` + Version string `json:"version,omitempty"` + Scope string `json:"scope,omitempty"` + ProjectRoot string `json:"project_root,omitempty"` + Client string `json:"client,omitempty"` + Force bool `json:"force,omitempty"` + Group string `json:"group,omitempty"` + AllowUnsigned bool `json:"allow_unsigned,omitempty"` } type installSkillResponse struct { @@ -1048,3 +1051,336 @@ var _ = Describe("Skills API", Label("api", "api-registry", "skills", "e2e"), fu }) }) }) + +// makeE2EProjectRoot creates a resolved temp dir with a .git directory, +// satisfying the project-root validation the lock file requires. +func makeE2EProjectRoot() string { + dir := GinkgoT().TempDir() + resolved, err := filepath.EvalSymlinks(dir) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + ExpectWithOffset(1, os.MkdirAll(filepath.Join(resolved, ".git"), 0o755)).To(Succeed()) + return resolved +} + +func uninstallScopedSkill(server *e2e.Server, name, projectRoot string) *http.Response { + u := fmt.Sprintf("%s/api/v1beta/skills/%s?scope=project&project_root=%s", + server.BaseURL(), name, url.QueryEscape(projectRoot)) + req, err := http.NewRequest(http.MethodDelete, u, nil) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + resp, err := http.DefaultClient.Do(req) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + return resp +} + +// buildAndPushSkill creates a skill directory, builds it, and pushes it to +// ociRegistry, returning the OCI reference. +func buildAndPushSkill(server *e2e.Server, ociRegistry *httptest.Server, skillName, description string) string { + ociRef := fmt.Sprintf("%s/e2e-test/%s:v0.1.0", ociRegistry.Listener.Addr().String(), skillName) + + skillDir := createTestSkillDir(skillName, description) + buildResp := buildSkill(server, skillDir, ociRef) + defer buildResp.Body.Close() + ExpectWithOffset(1, buildResp.StatusCode).To(Equal(http.StatusOK)) + + pushResp := pushSkill(server, ociRef) + defer pushResp.Body.Close() + ExpectWithOffset(1, pushResp.StatusCode).To(Equal(http.StatusNoContent)) + + return ociRef +} + +// This RFC THV-0080 feature is gated behind TOOLHIVE_SKILLS_LOCK_ENABLED +// while it lands across a stack of PRs (see skills.LockFileFeatureEnabled), +// so this Describe block runs its own server with the gate turned on rather +// than sharing the "Skills API" block's default-off server above. +var _ = Describe("Project-scope skills lock file (RFC THV-0080)", Label("api", "api-registry", "skills", "skills-lock", "e2e"), func() { + var ( + config *e2e.ServerConfig + apiServer *e2e.Server + ) + + BeforeEach(func() { + config = e2e.NewServerConfig() + config.ExtraEnv = []string{"TOOLHIVE_SKILLS_LOCK_ENABLED=true"} + apiServer = e2e.StartServer(config) + }) + + It("records an installed skill in the project's toolhive.lock.yaml", func() { + projectRoot := makeE2EProjectRoot() + skillName := "lock-e2e-skill" + + By("Starting an in-process OCI registry and pushing a test skill") + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A skill for lock file E2E testing") + + By("Installing the skill into the project") + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, AllowUnsigned: true, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(http.StatusCreated)) + var installResult installSkillResponse + Expect(json.NewDecoder(installResp.Body).Decode(&installResult)).To(Succeed()) + + By("Verifying the lock file was written with the expected entry") + root, err := lockfile.OpenRoot(projectRoot) + Expect(err).ToNot(HaveOccurred()) + lf, err := lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + entry, ok := lf.Get(skillName) + Expect(ok).To(BeTrue(), "expected a lock entry for %q", skillName) + Expect(entry.Source).To(Equal(ociRef)) + Expect(entry.Digest).ToNot(BeEmpty()) + Expect(entry.ContentDigest).ToNot(BeEmpty()) + Expect(entry.Explicit).To(BeTrue()) + Expect(entry.Unsigned).To(BeTrue(), "the unsigned exception must be recorded in the lock entry") + + By("Cleaning up") + cleanupResp := uninstallScopedSkill(apiServer, skillName, projectRoot) + defer cleanupResp.Body.Close() + }) + + It("rejects an unsigned project-scoped install without allow_unsigned", func() { + projectRoot := makeE2EProjectRoot() + skillName := "lock-e2e-unsigned-rejected" + + By("Starting an in-process OCI registry and pushing an unsigned test skill") + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "An unsigned skill that must be rejected") + + By("Installing without allow_unsigned and expecting a 403") + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(http.StatusForbidden)) + + By("Verifying no lock entry was written") + root, err := lockfile.OpenRoot(projectRoot) + Expect(err).ToNot(HaveOccurred()) + lf, err := lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + _, ok := lf.Get(skillName) + Expect(ok).To(BeFalse(), "a rejected install must not write a lock entry") + }) + + It("removes the lock entry when the skill is uninstalled", func() { + projectRoot := makeE2EProjectRoot() + skillName := "lock-e2e-uninstall-skill" + + By("Starting an in-process OCI registry and pushing a test skill") + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A skill for lock file uninstall E2E testing") + + By("Installing the skill into the project") + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, AllowUnsigned: true, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(http.StatusCreated)) + + root, err := lockfile.OpenRoot(projectRoot) + Expect(err).ToNot(HaveOccurred()) + lf, err := lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + _, ok := lf.Get(skillName) + Expect(ok).To(BeTrue(), "expected a lock entry before uninstall") + + By("Uninstalling the skill") + uninstallResp := uninstallScopedSkill(apiServer, skillName, projectRoot) + defer uninstallResp.Body.Close() + Expect(uninstallResp.StatusCode).To(Equal(http.StatusNoContent)) + + By("Verifying the lock entry was removed") + lf, err = lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + _, ok = lf.Get(skillName) + Expect(ok).To(BeFalse(), "expected the lock entry to be removed after uninstall") + }) + + It("restores a deleted skill's files via sync", func() { + projectRoot := makeE2EProjectRoot() + skillName := "lock-e2e-sync-skill" + + By("Starting an in-process OCI registry and pushing a test skill") + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A skill for sync E2E testing") + + By("Installing the skill into the project") + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, AllowUnsigned: true, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(http.StatusCreated)) + + By("Deleting the installed files without touching the lock file or DB record") + skillDir := filepath.Join(projectRoot, ".claude", "skills", skillName) + Expect(os.RemoveAll(skillDir)).To(Succeed()) + _, statErr := os.Stat(filepath.Join(skillDir, "SKILL.md")) + Expect(os.IsNotExist(statErr)).To(BeTrue()) + + By("Running sync --check — it must report drift without restoring anything") + checkResp := syncSkills(apiServer, syncSkillsRequest{ProjectRoot: projectRoot, Check: true}) + defer checkResp.Body.Close() + Expect(checkResp.StatusCode).To(Equal(http.StatusOK)) + var checkResult syncResultResponse + Expect(json.NewDecoder(checkResp.Body).Decode(&checkResult)).To(Succeed()) + Expect(checkResult.Drifted).To(ContainElement(skillName)) + _, statErr = os.Stat(filepath.Join(skillDir, "SKILL.md")) + Expect(os.IsNotExist(statErr)).To(BeTrue(), "--check must not restore files") + + By("Running sync — it must restore the missing files") + syncResp := syncSkills(apiServer, syncSkillsRequest{ProjectRoot: projectRoot}) + defer syncResp.Body.Close() + Expect(syncResp.StatusCode).To(Equal(http.StatusOK)) + var syncResult syncResultResponse + Expect(json.NewDecoder(syncResp.Body).Decode(&syncResult)).To(Succeed()) + Expect(syncResult.Installed).To(ContainElement(skillName)) + + content, readErr := os.ReadFile(filepath.Join(skillDir, "SKILL.md")) + Expect(readErr).ToNot(HaveOccurred()) + Expect(string(content)).To(ContainSubstring(skillName)) + + By("Cleaning up") + cleanupResp := uninstallScopedSkill(apiServer, skillName, projectRoot) + defer cleanupResp.Body.Close() + }) + + It("upgrades a skill when newer content is republished at the same reference", func() { + projectRoot := makeE2EProjectRoot() + skillName := "lock-e2e-upgrade-skill" + + By("Starting an in-process OCI registry and pushing the initial version") + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "The original description") + + By("Installing the skill into the project") + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, AllowUnsigned: true, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(http.StatusCreated)) + + root, err := lockfile.OpenRoot(projectRoot) + Expect(err).ToNot(HaveOccurred()) + lf, err := lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + before, ok := lf.Get(skillName) + Expect(ok).To(BeTrue()) + + By("Republishing newer content at the same OCI reference") + newSkillDir := createTestSkillDir(skillName, "The updated description") + rebuildResp := buildSkill(apiServer, newSkillDir, ociRef) + defer rebuildResp.Body.Close() + Expect(rebuildResp.StatusCode).To(Equal(http.StatusOK)) + repushResp := pushSkill(apiServer, ociRef) + defer repushResp.Body.Close() + Expect(repushResp.StatusCode).To(Equal(http.StatusNoContent)) + + By("Previewing the upgrade — it must report the change without installing it") + previewResp := upgradeSkills(apiServer, upgradeSkillsRequest{ProjectRoot: projectRoot, Preview: true}) + defer previewResp.Body.Close() + Expect(previewResp.StatusCode).To(Equal(http.StatusOK)) + var previewResult upgradeResultResponse + Expect(json.NewDecoder(previewResp.Body).Decode(&previewResult)).To(Succeed()) + Expect(previewResult.Outcomes).To(HaveLen(1)) + Expect(previewResult.Outcomes[0].Status).To(Equal("upgraded")) + + lf, err = lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + afterPreview, ok := lf.Get(skillName) + Expect(ok).To(BeTrue()) + Expect(afterPreview.Digest).To(Equal(before.Digest), "preview must not rewrite the lock file") + + By("Applying the upgrade") + applyResp := upgradeSkills(apiServer, upgradeSkillsRequest{ProjectRoot: projectRoot}) + defer applyResp.Body.Close() + Expect(applyResp.StatusCode).To(Equal(http.StatusOK)) + var applyResult upgradeResultResponse + Expect(json.NewDecoder(applyResp.Body).Decode(&applyResult)).To(Succeed()) + Expect(applyResult.Outcomes).To(HaveLen(1)) + Expect(applyResult.Outcomes[0].Status).To(Equal("upgraded")) + + lf, err = lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + after, ok := lf.Get(skillName) + Expect(ok).To(BeTrue()) + Expect(after.Digest).ToNot(Equal(before.Digest), "the lock file must record the new digest") + Expect(after.Source).To(Equal(before.Source), "Source must never be rewritten by upgrade") + + By("Cleaning up") + cleanupResp := uninstallScopedSkill(apiServer, skillName, projectRoot) + defer cleanupResp.Body.Close() + }) +}) + +type upgradeSkillsRequest struct { + ProjectRoot string `json:"project_root"` + Names []string `json:"names,omitempty"` + Preview bool `json:"preview,omitempty"` + FailOnChanges bool `json:"fail_on_changes,omitempty"` + AllowRefChange bool `json:"allow_ref_change,omitempty"` + Clients []string `json:"clients,omitempty"` +} + +type upgradeOutcomeResponse struct { + Name string `json:"name"` + Status string `json:"status"` + OldDigest string `json:"old_digest,omitempty"` + NewDigest string `json:"new_digest,omitempty"` + NewResolvedReference string `json:"new_resolved_reference,omitempty"` + Reason string `json:"reason,omitempty"` + Error string `json:"error,omitempty"` +} + +type upgradeResultResponse struct { + Outcomes []upgradeOutcomeResponse `json:"outcomes"` +} + +func upgradeSkills(server *e2e.Server, req upgradeSkillsRequest) *http.Response { + jsonData, err := json.Marshal(req) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + + resp, err := http.Post( + server.BaseURL()+"/api/v1beta/skills/upgrade", + "application/json", + bytes.NewBuffer(jsonData), + ) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + return resp +} + +type syncSkillsRequest struct { + ProjectRoot string `json:"project_root"` + Clients []string `json:"clients,omitempty"` + Prune bool `json:"prune,omitempty"` + Check bool `json:"check,omitempty"` + Adopt bool `json:"adopt,omitempty"` +} + +type syncResultResponse struct { + Installed []string `json:"installed,omitempty"` + Drifted []string `json:"drifted,omitempty"` + AlreadyCurrent []string `json:"already_current,omitempty"` + NeverManaged []string `json:"never_managed,omitempty"` + RemovedFromLock []string `json:"removed_from_lock,omitempty"` + Pruned []string `json:"pruned,omitempty"` +} + +func syncSkills(server *e2e.Server, req syncSkillsRequest) *http.Response { + jsonData, err := json.Marshal(req) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + + resp, err := http.Post( + server.BaseURL()+"/api/v1beta/skills/sync", + "application/json", + bytes.NewBuffer(jsonData), + ) + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + return resp +} diff --git a/test/e2e/chainsaw/operator/multi-tenancy/setup/chainsaw-test.yaml b/test/e2e/chainsaw/operator/multi-tenancy/setup/chainsaw-test.yaml index 4aabcf830a..d7b2cd8877 100644 --- a/test/e2e/chainsaw/operator/multi-tenancy/setup/chainsaw-test.yaml +++ b/test/e2e/chainsaw/operator/multi-tenancy/setup/chainsaw-test.yaml @@ -42,6 +42,10 @@ spec: - operator.rbac.scope=namespace - --set - operator.rbac.allowedNamespaces={toolhive-system,test-namespace,toolhive-test-ns-1,toolhive-test-ns-2} + # The StorageVersionMigrator watches cluster-scoped CRDs, so it is + # unsupported (and chart-rejected) in namespace scope. Opt out here. + - --set + - operator.features.storageVersionMigrator=false - assert: file: assert-operator-ready.yaml - assert: diff --git a/test/e2e/chainsaw/operator/single-tenancy/test-scenarios/mcpwebhookconfig/chainsaw-test.yaml b/test/e2e/chainsaw/operator/single-tenancy/test-scenarios/mcpwebhookconfig/chainsaw-test.yaml index c26e217f50..7c49e40db3 100644 --- a/test/e2e/chainsaw/operator/single-tenancy/test-scenarios/mcpwebhookconfig/chainsaw-test.yaml +++ b/test/e2e/chainsaw/operator/single-tenancy/test-scenarios/mcpwebhookconfig/chainsaw-test.yaml @@ -116,17 +116,6 @@ spec: echo "MCPServer did not get a webhookConfigHash" kubectl get mcpserver missing-webhook-server -n toolhive-system -o yaml exit 1 - - assert: - resource: - apiVersion: toolhive.stacklok.dev/v1alpha1 - kind: MCPWebhookConfig - metadata: - name: non-existent-webhook-config - namespace: toolhive-system - status: - referencingWorkloads: - - kind: MCPServer - name: missing-webhook-server # Delete in dependency order so the MCPWebhookConfig finalizer can clear before Chainsaw cleanup. - name: cleanup-webhook-config-dependencies diff --git a/test/e2e/cli_llm_all_clients_test.go b/test/e2e/cli_llm_all_clients_test.go index 8a704af58d..faaae45342 100644 --- a/test/e2e/cli_llm_all_clients_test.go +++ b/test/e2e/cli_llm_all_clients_test.go @@ -254,12 +254,9 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", binDir, err = e2e.CreateFakeBrowserDir(tempDir) Expect(err).ToNot(HaveOccurred()) - // Allocate a free port for the OIDC mock server. - oidcPort, err = networking.FindOrUsePort(0) - Expect(err).ToNot(HaveOccurred()) - - oidcServer, err = e2e.NewOIDCMockServer(oidcPort, clientID, clientSecret) + oidcServer, err = e2e.NewOIDCMockServer(0, clientID, clientSecret) Expect(err).ToNot(HaveOccurred()) + oidcPort = oidcServer.Port() oidcServer.EnableAutoComplete() Expect(oidcServer.Start()).To(Succeed()) @@ -541,39 +538,13 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", proxyPort, portErr := networking.FindOrUsePort(0) Expect(portErr).ToNot(HaveOccurred()) - By(fmt.Sprintf("setting proxy port to %d", proxyPort)) - thvCmd("llm", "config", "set", "--proxy-port", fmt.Sprintf("%d", proxyPort)).ExpectSuccess() - - By("starting the proxy in a goroutine") - type proxyResult struct { - stdout, stderr string - err error - } - done := make(chan proxyResult, 1) - proxyCmd := thvCmd("llm", "proxy", "start") - go func() { - out, serr, rerr := proxyCmd.RunWithTimeout(15 * time.Second) - done <- proxyResult{out, serr, rerr} - }() - DeferCleanup(func() { - _ = proxyCmd.Interrupt() - select { - case <-done: - case <-time.After(5 * time.Second): - } + By(fmt.Sprintf("starting the proxy on port %d", proxyPort)) + _, portErr = e2e.RetryOnPortConflict(proxyPort, func(port int) error { + return startLLMProxy(thvCmd, func() *e2e.THVCommand { + return thvCmd("llm", "proxy", "start") + }, port, 15*time.Second) }) - - By(fmt.Sprintf("waiting for proxy to listen on port %d", proxyPort)) - proxyAddr := fmt.Sprintf("127.0.0.1:%d", proxyPort) - Eventually(func() error { - conn, err := net.DialTimeout("tcp", proxyAddr, 200*time.Millisecond) - if err != nil { - return err - } - _ = conn.Close() - return nil - }, 10*time.Second, 300*time.Millisecond).Should(Succeed(), - "proxy should be listening on %s", proxyAddr) + Expect(portErr).ToNot(HaveOccurred()) }) }) @@ -630,9 +601,7 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", // Start a local HTTPS mock gateway so the proxy can forward requests // quickly rather than timing out on DNS resolution for a fake domain. - gwPort, portErr := networking.FindOrUsePort(0) - Expect(portErr).ToNot(HaveOccurred()) - gw, gwErr := e2e.NewLLMGatewayMock(gwPort) + gw, gwErr := e2e.NewLLMGatewayMock(0) Expect(gwErr).ToNot(HaveOccurred()) Expect(gw.Start()).To(Succeed()) defer func() { _ = gw.Stop() }() @@ -660,36 +629,17 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", proxyPort, portErr2 := networking.FindOrUsePort(0) Expect(portErr2).ToNot(HaveOccurred()) - By(fmt.Sprintf("setting proxy port to %d and starting proxy", proxyPort)) - thvCmd("llm", "config", "set", "--proxy-port", fmt.Sprintf("%d", proxyPort)).ExpectSuccess() - - done := make(chan struct{}) - proxyCmd := thvCmd("llm", "proxy", "start").WithEnv( - "SSL_CERT_FILE="+gwCertFile, - rebindEnvKey+"="+rebindToken, - ) - go func() { - defer close(done) - _, _, _ = proxyCmd.RunWithTimeout(15 * time.Second) - }() - DeferCleanup(func() { - _ = proxyCmd.Interrupt() - select { - case <-done: - case <-time.After(5 * time.Second): - } + By(fmt.Sprintf("starting the proxy on port %d", proxyPort)) + proxyPort, portErr2 = e2e.RetryOnPortConflict(proxyPort, func(port int) error { + return startLLMProxy(thvCmd, func() *e2e.THVCommand { + return thvCmd("llm", "proxy", "start").WithEnv( + "SSL_CERT_FILE="+gwCertFile, + rebindEnvKey+"="+rebindToken, + ) + }, port, 15*time.Second) }) - + Expect(portErr2).ToNot(HaveOccurred()) proxyAddr := fmt.Sprintf("127.0.0.1:%d", proxyPort) - Eventually(func() error { - conn, dialErr := net.DialTimeout("tcp", proxyAddr, 200*time.Millisecond) - if dialErr != nil { - return dialErr - } - _ = conn.Close() - return nil - }, 10*time.Second, 300*time.Millisecond).Should(Succeed(), - "proxy should be listening on %s", proxyAddr) rebindClient := &http.Client{Timeout: 10 * time.Second} @@ -765,16 +715,15 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", Skip(fmt.Sprintf("client %q not supported on %s", clientTC.name, runtime.GOOS)) } - // Allocate ports for the mock gateway and the proxy. - gatewayPort, portErr := networking.FindOrUsePort(0) - Expect(portErr).ToNot(HaveOccurred()) + // Allocate a port for the proxy. proxyPort, portErr := networking.FindOrUsePort(0) Expect(portErr).ToNot(HaveOccurred()) // Start the mock LLM gateway (HTTPS with self-signed cert). - By(fmt.Sprintf("[%s] starting mock LLM gateway on port %d", clientTC.name, gatewayPort)) - gateway, gwErr := e2e.NewLLMGatewayMock(gatewayPort) + gateway, gwErr := e2e.NewLLMGatewayMock(0) Expect(gwErr).ToNot(HaveOccurred()) + gatewayPort := gateway.Port() + By(fmt.Sprintf("[%s] starting mock LLM gateway on port %d", clientTC.name, gatewayPort)) Expect(gateway.Start()).To(Succeed()) defer func() { _ = gateway.Stop() }() @@ -822,34 +771,14 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", } // Configure the proxy port and start it. - By(fmt.Sprintf("[%s] setting proxy port to %d", clientTC.name, proxyPort)) - thvCmd("llm", "config", "set", "--proxy-port", fmt.Sprintf("%d", proxyPort)).ExpectSuccess() - - By(fmt.Sprintf("[%s] starting the proxy", clientTC.name)) - done := make(chan struct{}) - proxyCmd := thvCmdWithToken("llm", "proxy", "start") - go func() { - defer close(done) - _, _, _ = proxyCmd.RunWithTimeout(20 * time.Second) - }() - DeferCleanup(func() { - _ = proxyCmd.Interrupt() - select { - case <-done: - case <-time.After(5 * time.Second): - } + By(fmt.Sprintf("[%s] starting the proxy on port %d", clientTC.name, proxyPort)) + proxyPort, portErr = e2e.RetryOnPortConflict(proxyPort, func(port int) error { + return startLLMProxy(thvCmd, func() *e2e.THVCommand { + return thvCmdWithToken("llm", "proxy", "start") + }, port, 20*time.Second) }) - + Expect(portErr).ToNot(HaveOccurred()) proxyAddr := fmt.Sprintf("127.0.0.1:%d", proxyPort) - Eventually(func() error { - conn, dialErr := net.DialTimeout("tcp", proxyAddr, 200*time.Millisecond) - if dialErr != nil { - return dialErr - } - _ = conn.Close() - return nil - }, 10*time.Second, 300*time.Millisecond).Should(Succeed(), - "proxy should be listening on %s", proxyAddr) // Send requests through the proxy and verify the gateway received them // with the correct Bearer token. Use a client with an explicit timeout @@ -1275,6 +1204,75 @@ var _ = Describe("thv llm — all-client matrix", Label("cli", "llm", "clients", // Helpers // ───────────────────────────────────────────────────────────────────────────── +// llmProxyResult captures a `thv llm proxy start` subprocess's outcome, so +// startLLMProxy can inspect it for an address-already-in-use signal if the +// subprocess exits before its listener comes up. +type llmProxyResult struct { + stdout, stderr string + err error +} + +// startLLMProxy sets the proxy port via `thv llm config set --proxy-port`, +// starts `thv llm proxy start` (built by buildProxyCmd, so each call site +// can inject its own env/token wiring) in the background, registers a +// DeferCleanup to interrupt it, and waits for it to accept TCP connections +// on 127.0.0.1:port. The port is only probed free at selection time +// (networking.FindOrUsePort); the real bind happens later in this +// subprocess, so another process can steal it in between. On that race, +// this returns an error containing "address already in use" so callers can +// retry via e2e.RetryOnPortConflict. +func startLLMProxy( + thvCmd func(args ...string) *e2e.THVCommand, + buildProxyCmd func() *e2e.THVCommand, + port int, + runTimeout time.Duration, +) error { + thvCmd("llm", "config", "set", "--proxy-port", fmt.Sprintf("%d", port)).ExpectSuccess() + + proxyCmd := buildProxyCmd() + done := make(chan llmProxyResult, 1) + go func() { + out, serr, rerr := proxyCmd.RunWithTimeout(runTimeout) + done <- llmProxyResult{out, serr, rerr} + }() + DeferCleanup(func() { + _ = proxyCmd.Interrupt() + select { + case <-done: + case <-time.After(5 * time.Second): + } + }) + + addr := fmt.Sprintf("127.0.0.1:%d", port) + deadline := time.After(10 * time.Second) + ticker := time.NewTicker(300 * time.Millisecond) + defer ticker.Stop() + for { + select { + case res := <-done: + return fmt.Errorf("proxy exited before listening on %s: err=%v stderr=%q", addr, res.err, res.stderr) + case <-deadline: + return fmt.Errorf("timed out waiting for proxy to listen on %s", addr) + case <-ticker.C: + conn, dialErr := net.DialTimeout("tcp", addr, 200*time.Millisecond) + if dialErr != nil { + continue + } + _ = conn.Close() + // A successful dial doesn't prove OUR subprocess is the listener -- if it + // lost the bind race, another process already listening on this exact + // port would also accept the connection. Give the subprocess's own exit + // (if any) a moment to land on `done` before trusting the dial. + select { + case res := <-done: + return fmt.Errorf("proxy exited before listening on %s: err=%v stderr=%q", addr, res.err, res.stderr) + case <-time.After(100 * time.Millisecond): + return nil + } + } + } +} + // createFakeBinary writes a minimal no-op shell script named `name` in dir. // This satisfies the LLMBinaryName check in DetectedLLMGatewayClients. func createFakeBinary(dir, name string) error { diff --git a/test/e2e/cli_llm_setup_test.go b/test/e2e/cli_llm_setup_test.go index 196ce7e6a3..56d6de13c6 100644 --- a/test/e2e/cli_llm_setup_test.go +++ b/test/e2e/cli_llm_setup_test.go @@ -18,7 +18,6 @@ import ( "gopkg.in/yaml.v3" "github.com/stacklok/toolhive/pkg/llm" - "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/test/e2e" ) @@ -160,15 +159,12 @@ var _ = Describe("thv llm setup / teardown", Label("cli", "llm", "setup", "e2e") By("Configuring environment secrets provider") thvCmd("secret", "provider", "environment").ExpectSuccess() - // Allocate a free port for the OIDC mock server. + // Create and start the mock OIDC server. var err error - oidcPort, err = networking.FindOrUsePort(0) + oidcServer, err = e2e.NewOIDCMockServer(0, clientID, clientSecret) Expect(err).ToNot(HaveOccurred()) - - // Create and start the mock OIDC server. + oidcPort = oidcServer.Port() By(fmt.Sprintf("Starting OIDC mock server on port %d", oidcPort)) - oidcServer, err = e2e.NewOIDCMockServer(oidcPort, clientID, clientSecret) - Expect(err).ToNot(HaveOccurred()) oidcServer.EnableAutoComplete() Expect(oidcServer.Start()).To(Succeed()) diff --git a/test/e2e/cli_skills_lock_test.go b/test/e2e/cli_skills_lock_test.go new file mode 100644 index 0000000000..66ca0983f2 --- /dev/null +++ b/test/e2e/cli_skills_lock_test.go @@ -0,0 +1,213 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "errors" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + + "github.com/google/go-containerregistry/pkg/registry" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/pkg/skills/lockfile" + "github.com/stacklok/toolhive/test/e2e" +) + +// This RFC THV-0080 feature is gated behind TOOLHIVE_SKILLS_LOCK_ENABLED +// while it lands across a stack of PRs (see skills.LockFileFeatureEnabled), +// so this Describe block runs its own server with the gate turned on rather +// than sharing the default-off server other CLI skills tests use. +var _ = Describe("Skills CLI lock file exit codes (RFC THV-0080)", Label("api", "cli", "skills", "skills-lock", "e2e"), func() { + var ( + config *e2e.ServerConfig + apiServer *e2e.Server + thvConfig *e2e.TestConfig + ) + + BeforeEach(func() { + config = e2e.NewServerConfig() + config.ExtraEnv = []string{"TOOLHIVE_SKILLS_LOCK_ENABLED=true"} + apiServer = e2e.StartServer(config) + thvConfig = e2e.NewTestConfig() + }) + + thvSkillCmd := func(args ...string) *e2e.THVCommand { + fullArgs := append([]string{"skill"}, args...) + return e2e.NewTHVCommand(thvConfig, fullArgs...). + WithEnv("TOOLHIVE_API_URL=" + apiServer.BaseURL()) + } + + exitCodeOf := func(err error) int { + var exitErr *exec.ExitError + ExpectWithOffset(1, errors.As(err, &exitErr)).To(BeTrue(), "expected an *exec.ExitError, got %T: %v", err, err) + return exitErr.ExitCode() + } + + Describe("thv skill sync --check", func() { + It("exits 0 when the project matches its lock file", func() { + projectRoot := makeE2EProjectRoot() + skillName := "cli-lock-clean-skill" + + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A clean skill for CLI exit code testing") + + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(201)) + + stdout, _ := thvSkillCmd("sync", "--check", "--project-root", projectRoot).ExpectSuccess() + Expect(stdout).To(ContainSubstring("Up to date")) + Expect(stdout).To(ContainSubstring(skillName)) + }) + + It("exits 2 when the project has drifted from its lock file", func() { + projectRoot := makeE2EProjectRoot() + skillName := "cli-lock-drifted-skill" + + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A drifted skill for CLI exit code testing") + + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(201)) + + By("Deleting the installed files so the project drifts from the lock file") + skillDir := filepath.Join(projectRoot, ".claude", "skills", skillName) + Expect(os.RemoveAll(skillDir)).To(Succeed()) + + _, _, err := thvSkillCmd("sync", "--check", "--project-root", projectRoot).Run() + Expect(err).To(HaveOccurred()) + Expect(exitCodeOf(err)).To(Equal(2)) + }) + }) + + Describe("thv skill sync without --yes", func() { + It("exits 4 when running non-interactively without --yes", func() { + projectRoot := makeE2EProjectRoot() + + _, _, err := thvSkillCmd("sync", "--project-root", projectRoot).Run() + Expect(err).To(HaveOccurred(), "a non-interactive sync without --yes must refuse rather than proceed silently") + Expect(exitCodeOf(err)).To(Equal(4)) + }) + }) + + Describe("thv skill sync --yes", func() { + It("exits 0 and proceeds without prompting", func() { + projectRoot := makeE2EProjectRoot() + + stdout, _ := thvSkillCmd("sync", "--yes", "--project-root", projectRoot).ExpectSuccess() + Expect(stdout).To(ContainSubstring("Nothing to sync")) + }) + }) + + Describe("thv skill sync partial failure", func() { + It("exits 3 when a pinned skill cannot be reinstalled", func() { + projectRoot := makeE2EProjectRoot() + skillName := "cli-lock-exit3-skill" + + ociRegistry := httptest.NewServer(registry.New()) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A skill whose registry will vanish") + + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(201)) + + By("Tampering with the installed files and killing the registry, so the repair reinstall must fail") + skillMD := filepath.Join(projectRoot, ".claude", "skills", skillName, "SKILL.md") + Expect(os.WriteFile(skillMD, []byte("tampered"), 0o644)).To(Succeed()) + ociRegistry.Close() + + _, _, err := thvSkillCmd("sync", "--yes", "--project-root", projectRoot).Run() + Expect(err).To(HaveOccurred(), "a failed repair must not exit 0") + Expect(exitCodeOf(err)).To(Equal(3), "an operational failure is exit 3, not a freshness signal") + }) + }) + + Describe("thv skill sync --check on a fresh clone", func() { + It("exits 2 when the lock file has entries but nothing is installed", func() { + projectRoot := makeE2EProjectRoot() + skillName := "cli-lock-freshclone-skill" + + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "A skill for the fresh-clone gate") + + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(201)) + + By("Simulating a fresh clone: the committed lock file survives, local install state does not") + lockPath := filepath.Join(projectRoot, "toolhive.lock.yaml") + lockBytes, err := os.ReadFile(lockPath) //nolint:gosec // fixed test path + Expect(err).ToNot(HaveOccurred()) + uninstallResp := uninstallScopedSkill(apiServer, skillName, projectRoot) + defer uninstallResp.Body.Close() + Expect(uninstallResp.StatusCode).To(Equal(204)) + Expect(os.WriteFile(lockPath, lockBytes, 0o644)).To(Succeed()) + + _, _, cmdErr := thvSkillCmd("sync", "--check", "--project-root", projectRoot).Run() + Expect(cmdErr).To(HaveOccurred(), "the CI gate must not green-light a checkout with nothing installed") + Expect(exitCodeOf(cmdErr)).To(Equal(2)) + }) + }) + + Describe("thv skill upgrade --fail-on-changes", func() { + It("exits 2 when a skill would change, without installing it", func() { + projectRoot := makeE2EProjectRoot() + skillName := "cli-lock-upgrade-fail-on-changes-skill" + + ociRegistry := httptest.NewServer(registry.New()) + DeferCleanup(ociRegistry.Close) + ociRef := buildAndPushSkill(apiServer, ociRegistry, skillName, "The original description") + + installResp := installSkill(apiServer, installSkillRequest{ + Name: ociRef, Scope: "project", ProjectRoot: projectRoot, + }) + defer installResp.Body.Close() + Expect(installResp.StatusCode).To(Equal(201)) + + By("Republishing newer content at the same OCI reference") + newSkillDir := createTestSkillDir(skillName, "The updated description") + rebuildResp := buildSkill(apiServer, newSkillDir, ociRef) + defer rebuildResp.Body.Close() + Expect(rebuildResp.StatusCode).To(Equal(200)) + repushResp := pushSkill(apiServer, ociRef) + defer repushResp.Body.Close() + Expect(repushResp.StatusCode).To(Equal(204)) + + root, err := lockfile.OpenRoot(projectRoot) + Expect(err).ToNot(HaveOccurred()) + before, err := lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + beforeEntry, ok := before.Get(skillName) + Expect(ok).To(BeTrue()) + + _, _, err = thvSkillCmd("upgrade", "--yes", "--fail-on-changes", "--project-root", projectRoot).Run() + Expect(err).To(HaveOccurred()) + Expect(exitCodeOf(err)).To(Equal(2)) + + By("Verifying nothing was actually installed before the conflict was reported") + after, err := lockfile.Load(root) + Expect(err).ToNot(HaveOccurred()) + afterEntry, ok := after.Get(skillName) + Expect(ok).To(BeTrue()) + Expect(afterEntry.Digest).To(Equal(beforeEntry.Digest), + "--fail-on-changes must not install a changed skill before reporting the conflict") + }) + }) +}) diff --git a/test/e2e/dual_era_functional_test.go b/test/e2e/dual_era_functional_test.go new file mode 100644 index 0000000000..6132611de3 --- /dev/null +++ b/test/e2e/dual_era_functional_test.go @@ -0,0 +1,452 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// This spec is the Tier-1 functional round-trip tier of the e2e suite (issue +// #5837, Commit 2). Modern (2026-07-28) fidelity is driven by the real, +// independent SDK client (yardstick-client) against the streamable proxy -- +// call-tool succeeds and -action info confirms it actually negotiated +// 2026-07-28 with an empty session id, not a silent downgrade -- plus one +// raw-client wire check against the transparent proxy for the +// no-Mcp-Session-Id property (that proxy's own, distinct Modern code path). +// Legacy (2025-11-25) is yardstick-client negotiating down: per Step 2.1's +// discovery, yardstick-client emits genuine Modern traffic on its own +// (server/discover, no initialize) whenever the backend's own discover +// response advertises 2026-07-28 support, so the Legacy leg deliberately +// uses a NON-stateless yardstick-server (its own StreamableHTTPOptions.Stateless +// is false), whose discover response omits 2026-07-28, forcing the SDK +// client to fall back to a real initialize handshake -- confirmed live +// before writing this (see the two legs below). + +// execTimeout bounds go install and every yardstick-client invocation below, +// so a wedged module proxy or stuck client fails fast instead of riding out +// the suite's global timeout. +const execTimeout = 90 * time.Second + +// yardstickClientVersion is derived from images.YardstickServerImage's tag +// (the part after ":") so a future image bump can't silently pair a +// mismatched go-sdk version between yardstick-server and yardstick-client. +var yardstickClientVersion = "v" + strings.SplitN(images.YardstickServerImage, ":", 2)[1] + +var ( + yardstickClientOnce sync.Once + yardstickClientPath string + yardstickClientErr error +) + +// ensureYardstickClientBinary go-installs the real yardstick-client module +// (github.com/stackloklabs/yardstick, a separate Go module -- this does NOT +// add go-sdk v1.7 to this repo's go.mod) into a scratch GOBIN, once per test +// run, and returns its path. +// +// ponytail: the scratch dir is not removed -- it's shared (via sync.Once) +// across every spec in this file, so a DeferCleanup tied to whichever spec +// happens to trigger the install first would delete it out from under later +// specs. One leaked temp dir per suite run is a cheap tradeoff; the OS temp +// dir is reaped independently. Upgrade to a suite-level AfterSuite removal +// if that ever stops being true (e.g. a long-lived CI temp volume). +func ensureYardstickClientBinary() (string, error) { + yardstickClientOnce.Do(func() { + dir, err := os.MkdirTemp("", "yardstick-client-") + if err != nil { + yardstickClientErr = fmt.Errorf("create scratch GOBIN: %w", err) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), execTimeout) + defer cancel() + cmd := exec.CommandContext(ctx, "go", "install", + "github.com/stackloklabs/yardstick/cmd/yardstick-client@"+yardstickClientVersion) + cmd.Env = append(os.Environ(), "GOBIN="+dir) + if out, err := cmd.CombinedOutput(); err != nil { + yardstickClientErr = fmt.Errorf("go install yardstick-client@%s: %w: %s", yardstickClientVersion, err, out) + return + } + yardstickClientPath = filepath.Join(dir, "yardstick-client") + }) + return yardstickClientPath, yardstickClientErr +} + +// runYardstickClient runs yardstick-client with a bounded timeout and returns +// its combined output. +func runYardstickClient(clientPath string, args ...string) (string, error) { + ctx, cancel := context.WithTimeout(context.Background(), execTimeout) + defer cancel() + //nolint:gosec // fixed action/flags, test-controlled address/port + out, err := exec.CommandContext(ctx, clientPath, args...).CombinedOutput() + return string(out), err +} + +// proxyPort extracts the port from a proxy URL like "http://127.0.0.1:PORT/mcp". +func proxyPort(proxyURL string) string { + rest := strings.TrimPrefix(proxyURL, "http://") + hostPort := strings.SplitN(rest, "/", 2)[0] + return strings.SplitN(hostPort, ":", 2)[1] +} + +// sessionIDFromInfo parses yardstick-client's `-action info` output and +// returns the value of its "Session ID: " line (empty for Modern, a real id +// for Legacy). Parsing the value, rather than substring-matching the whole +// line, survives a trailing space or relabel and can't pass vacuously. +func sessionIDFromInfo(infoOutput string) string { + for _, line := range strings.Split(infoOutput, "\n") { + if rest, ok := strings.CutPrefix(line, " Session ID:"); ok { + return strings.TrimSpace(rest) + } + } + return "" +} + +var _ = Describe("Modern Functional Round-Trip — Streamable Proxy", Label("proxy", "stateless", "dual-era", "e2e"), Serial, func() { + var ( + config *e2e.TestConfig + serverName string + client *e2e.RawMCPClient + proxyURL string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("func-streamable") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + client, err = e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + e2e.NewTHVCommand(config, "run", + "--name", serverName, + "--transport", "stdio", + "-e", "BACKEND_MODE=echo", + images.YardstickServerImage, + ).ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("round-trips through a real independent SDK client (yardstick-client) as Modern", func() { + // Wire-level "no Mcp-Session-Id for Modern" is already covered by the + // crown jewel + mixing specs; this proves functional interop with a + // real go-sdk v1.7 client instead of re-asserting that. + clientPath, err := ensureYardstickClientBinary() + Expect(err).ToNot(HaveOccurred(), "yardstick-client must be installable") + port := proxyPort(proxyURL) + + out, err := runYardstickClient(clientPath, + "-transport", "streamable-http", "-address", "127.0.0.1", "-port", port, + "-action", "call-tool", "-tool", "echo", "-args", `{"input":"modernfunctional"}`, + ) + Expect(err).ToNot(HaveOccurred(), "yardstick-client call-tool should succeed: %s", out) + Expect(out).To(ContainSubstring("modernfunctional")) + + // This backend's own transport is stdio (unaffected by any Stateless + // gating -- see the file doc comment), so yardstick-client's discover + // negotiates Modern here. info confirms it, not a silent downgrade. + infoOut, err := runYardstickClient(clientPath, + "-transport", "streamable-http", "-address", "127.0.0.1", "-port", port, "-action", "info", + ) + Expect(err).ToNot(HaveOccurred(), "yardstick-client info should succeed: %s", infoOut) + Expect(infoOut).To(ContainSubstring("Protocol Version: 2026-07-28")) + Expect(sessionIDFromInfo(infoOut)).To(BeEmpty(), "Modern must have no session id (inverse of the Legacy leg)") + }) + + It("reaches server/discover (baseline: no authz middleware wired at all)", func() { + // As of #5953, server/discover is always-allowed in + // pkg/authz/middleware.go's MCPMethodToFeatureOperation map, so this + // no-authz baseline and the "under authz" case in the + // "server/discover Authorization Guard" Describe below now assert + // the same 200 outcome -- kept as two separate specs because they + // exercise different code paths: this one never wires the authz + // middleware at all (RunConfig.AuthzConfig == nil, + // pkg/runner/middleware.go:210), the other wires it and relies on + // discover's always-allowed entry in that map. + req, err := e2e.NewModernRequest("server/discover", nil) + Expect(err).ToNot(HaveOccurred()) + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(200)) + Expect(resp.Error).To(BeNil()) + }) +}) + +// No second yardstick-client leg here: SDK negotiation (discover/initialize) +// is backend-side logic, identical regardless of which ToolHive proxy fronts +// it -- the streamable-proxy leg above already proves it end-to-end. This +// spec keeps the raw-client no-session-id check, which does exercise the +// transparent proxy's own (distinct) Modern code path. +var _ = Describe("Modern Functional Round-Trip — Transparent Proxy", Label("proxy", "stateless", "dual-era", "e2e"), Serial, func() { + var ( + config *e2e.TestConfig + serverName string + mockServer *statelessMockMCPServer + client *e2e.RawMCPClient + proxyURL string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("func-transparent") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + mockServer, err = newStatelessMockMCPServer() + Expect(err).ToNot(HaveOccurred()) + + client, err = e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + e2e.NewTHVCommand(config, "run", "--name", serverName, "--stateless", mockServer.URL()+"/mcp").ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if mockServer != nil { + mockServer.Stop() + mockServer = nil + } + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("round-trips a Modern request with no Mcp-Session-Id", func() { + req, err := e2e.NewModernRequest("tools/list", nil) + Expect(err).ToNot(HaveOccurred()) + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(200)) + Expect(resp.Error).To(BeNil()) + Expect(resp.Headers.Get(e2e.HeaderMCPSessionID)).To(BeEmpty(), "Modern must never carry Mcp-Session-Id") + }) +}) + +var _ = Describe("Legacy Functional Round-Trip — real SDK client (yardstick-client)", Label("proxy", "stateless", "dual-era", "e2e"), Serial, func() { + var ( + config *e2e.TestConfig + backendName string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + backendName = e2e.GenerateUniqueServerName("func-legacy-backend") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + // Deliberately NOT stateless: yardstick-server's own discover response + // then omits 2026-07-28 (StreamableHTTPOptions.Stateless gates it -- + // confirmed live), so yardstick-client's own SDK negotiation falls + // back to a real Legacy initialize handshake. --group "default" is + // the always-present built-in group; no group create/cleanup needed. + startYardstickOnPort(config, "default", backendName, allocateVMCPPort()) + }) + + AfterEach(func() { + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, backendName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("completes a Legacy tool call end-to-end through the transparent proxy", func() { + clientPath, err := ensureYardstickClientBinary() + Expect(err).ToNot(HaveOccurred(), "yardstick-client must be installable") + + proxyURL, err := e2e.GetMCPServerURL(config, backendName) + Expect(err).ToNot(HaveOccurred()) + port := proxyPort(proxyURL) + + out, err := runYardstickClient(clientPath, + "-transport", "streamable-http", + "-address", "127.0.0.1", + "-port", port, + "-action", "call-tool", + "-tool", "echo", + "-args", `{"input":"legacyfunctional"}`, + ) + Expect(err).ToNot(HaveOccurred(), "yardstick-client call-tool should succeed: %s", out) + Expect(out).To(ContainSubstring("legacyfunctional")) + + // info action reports the negotiated protocol version and session id; + // confirms this really was the Legacy leg, not a silent Modern upgrade. + infoOut, err := runYardstickClient(clientPath, + "-transport", "streamable-http", "-address", "127.0.0.1", "-port", port, "-action", "info") + Expect(err).ToNot(HaveOccurred(), "yardstick-client info should succeed: %s", infoOut) + Expect(infoOut).To(ContainSubstring("Protocol Version: 2025-11-25")) + Expect(sessionIDFromInfo(infoOut)).ToNot(BeEmpty(), "Legacy must have a real session id") + }) +}) + +// server/discover used to be excluded from pkg/authz/middleware.go's +// MCPMethodToFeatureOperation map (default-deny, 403) to keep its +// unfiltered tool/resource descriptors away from ResponseFilteringWriter +// (which only filters tools/list, prompts/list, resources/list, and +// find_tool). #5953 ("Serve MCP 2026-07-28 Modern stateless requests through +// vMCP") flipped this to always-allowed ({Feature:"", Operation:""} -- +// initialize parity, same rationale as "initialize" itself already being +// always-allowed): discover carries the same Capabilities/Instructions shape +// initialize does, so it adds no new exposure class, and #5953's own +// middleware comment notes discover would "just pass through unfiltered" +// under the response-filter either way -- a deliberate, documented tradeoff, +// not an oversight. +// +// Consequence for dual-era: Modern-through-authz no longer downgrades. +// Previously a real Modern client's discover got 403'd under authz, so its +// fallback-on-non-modern-error path kicked in and it silently negotiated +// down to a Legacy 2025-11-25 initialize handshake. Now discover succeeds, +// so the client stays Modern -- this guard now positively proves that. +var _ = Describe("server/discover Authorization Guard", Label("proxy", "stateless", "dual-era", "e2e", "authz"), Serial, func() { + var ( + config *e2e.TestConfig + serverName string + client *e2e.RawMCPClient + proxyURL string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("func-discover-authz") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + client, err = e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + // Minimal Cedar policy allowing tools/call+tools/list (so the backend + // is otherwise usable) -- mirrors osv_authz_test.go's pattern. + authzConfig := `{ + "version": "1.0", + "type": "cedarv1", + "cedar": { + "policies": [ + "permit(principal, action == Action::\"call_tool\", resource == Tool::\"echo\");", + "permit(principal, action == Action::\"list_tools\", resource);" + ], + "entities_json": "[]" + } +}` + tempDir, err := os.MkdirTemp("", "func-discover-authz") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = os.RemoveAll(tempDir) }) + authzConfigPath := filepath.Join(tempDir, "authz-config.json") + Expect(os.WriteFile(authzConfigPath, []byte(authzConfig), 0600)).To(Succeed()) + + e2e.NewTHVCommand(config, "run", + "--name", serverName, + "--transport", "stdio", + "-e", "BACKEND_MODE=echo", + "--authz-config", authzConfigPath, + images.YardstickServerImage, + ).ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("server/discover is allowed under authz (Modern's initialize replacement)", func() { + req, err := e2e.NewModernRequest("server/discover", nil) + Expect(err).ToNot(HaveOccurred()) + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(200)) + Expect(resp.Error).To(BeNil()) + + By("confirming authz still allows the policy-permitted echo tool call") + callReq, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "stillworks"}, + }) + Expect(err).ToNot(HaveOccurred()) + callResp, err := client.Send(context.Background(), proxyURL, callReq) + Expect(err).ToNot(HaveOccurred()) + Expect(callResp.StatusCode).To(Equal(200)) + Expect(callResp.Error).To(BeNil()) + }) + + It("still denies with 403 a method the Cedar policy does not permit, with a conformant JSON-RPC envelope", func() { + // server/discover's always-allowed flip doesn't touch the general + // default-deny rule (MCPMethodToFeatureOperation's doc comment: + // "Methods not in this map are denied by default"). tools/call IS in + // the map and routes through Cedar; the policy in this BeforeEach + // only permits resource == Tool::"echo", so calling any other tool + // name still gets denied -- keeps the 403/denial path covered by this + // suite. + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "notpermitted", + "arguments": map[string]any{"input": "shouldbedenied"}, + }) + Expect(err).ToNot(HaveOccurred()) + req.WithID(int64(9001)) + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(403)) + Expect(resp.Error).ToNot(BeNil()) + Expect(resp.Error.Code).To(Equal(int64(403))) + // The denial envelope must be conformant JSON-RPC: correct version tag + // and the request's id echoed back. + Expect(resp.JSONRPC).To(Equal("2.0")) + Expect(resp.ID).To(Equal(json.Number("9001"))) + }) +}) diff --git a/test/e2e/dual_era_mixing_test.go b/test/e2e/dual_era_mixing_test.go new file mode 100644 index 0000000000..9337bb72cc --- /dev/null +++ b/test/e2e/dual_era_mixing_test.go @@ -0,0 +1,285 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "context" + "fmt" + "strings" + "sync" + "time" + + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// This spec is the dual-era mixing tier of the e2e suite (issue #5837, tier +// 2): Legacy (2025-11-25, session-based) and Modern (2026-07-28, stateless) +// raw clients hit the same streamable proxy instance concurrently. This +// tier exercises concurrent mixed-era *coexistence* -- it does NOT force a +// simultaneity collision like the crown jewel (dual_era_proxy_test.go) does +// with a barrier backend. What it proves is response-body delivery +// correctness (no client ever receives another client's response body), +// bounded by the echo backend having no observable session-scoped state -- +// it does not prove session-metadata isolation inside the proxy itself. +// +// Cross-delivery correlation reuses the crown jewel's mechanism +// (dual_era_proxy_test.go's extractNonce): every request, Legacy or Modern, +// carries a unique nonce in a custom _meta key (not one of the reserved +// io.modelcontextprotocol/* keys, so it never flips a Legacy request's +// classification -- see pkg/mcp/revision.go's hasModernSignal), and +// yardstick's echo tool reflects whatever _meta it received regardless of +// era. A response with the wrong nonce means this request's response body +// was delivered to (or received from) a different client. +// +// Confirmed live (thv run + curl) before writing this: Legacy initialize +// returns Mcp-Session-Id; a Legacy tools/call reusing that session id and a +// Modern tools/call both echo a custom _meta.nonce; neither ordinary +// tools/call response (Legacy or Modern) carries Mcp-Session-Id -- only +// initialize does; and a Legacy tools/call with a bogus/unregistered +// Mcp-Session-Id gets HTTP 404 with JSON-RPC -32001 "Session not found", +// proving the proxy actually enforces the session rather than merely +// minting one (see the negative probe below). So era-distinctness isn't a +// per-response header check; it's that Legacy requires and enforces a +// session id issued by initialize while Modern never sends one at all (see +// runMixedRound's assertions). + +// Deliberately NOT setting "Accept: application/json, text/event-stream" +// (mcp_raw_client.go's Send doc comment floats it for a live streamable-HTTP +// endpoint): verified empirically that doing so makes THIS proxy switch to +// its SSE response path (handleSingleRequestSSE in streamable_proxy.go, +// triggered by any Accept header containing "text/event-stream"), and the +// raw client's populateEnvelope only parses a plain JSON object body, not an +// SSE-framed one -- so nonce extraction silently returns "". Without the +// header the proxy returns a plain JSON response, which is what every other +// spec in this suite (and this one) relies on. +var _ = Describe("Dual-Era Concurrent Mixing", Label("proxy", "stateless", "dual-era", "e2e"), Serial, func() { + const ( + legacySessions = 3 + modernPerRound = 3 + rounds = 3 + clientTimeout = 15 * time.Second + // joinTimeout bounds one round: comfortably above clientTimeout (the + // longest a single request should take) so a stuck proxy fails fast + // instead of hanging until the suite's global timeout. + joinTimeout = clientTimeout + 10*time.Second + ) + + var ( + config *e2e.TestConfig + serverName string + client *e2e.RawMCPClient + proxyURL string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("dual-era-mix") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + client, err = e2e.NewRawMCPClient(clientTimeout) + Expect(err).ToNot(HaveOccurred()) + + By("Starting the yardstick backend under --transport stdio with BACKEND_MODE=echo") + e2e.NewTHVCommand(config, "run", + "--name", serverName, + "--transport", "stdio", + "-e", "BACKEND_MODE=echo", + images.YardstickServerImage, + ).ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("mixes concurrent Legacy and Modern traffic without cross-era delivery", func() { + // Standing gate: Legacy initialize assigns a session id (existing + // Legacy behavioral assertion, not a byte-for-byte golden -- no + // committed Legacy golden fixture exists to diff against; see + // mcp_raw_client_golden_test.go, which only covers Modern). + sessionIDs := make([]string, legacySessions) + for i := range sessionIDs { + initReq := e2e.NewLegacyInitializeRequest(fmt.Sprintf("legacy-client-%d", i), "1.0") + resp, err := client.Send(context.Background(), proxyURL, initReq) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(200)) + sessionIDs[i] = resp.Headers.Get(e2e.HeaderMCPSessionID) + Expect(sessionIDs[i]).ToNot(BeEmpty(), "initialize must assign a session id") + } + + // Negative probe: Legacy must actually ENFORCE the session, not just + // mint one. The echo tool is session-agnostic, so a Legacy path that + // silently degraded to stateless would still echo correctly and this + // suite's conservation checks alone wouldn't catch it. Confirmed live + // (thv run + curl) before writing this: a bogus/unregistered session + // id gets HTTP 404, JSON-RPC -32001 "Session not found". + bogusReq, err := e2e.NewLegacyRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "shouldneverrun"}, + }) + Expect(err).ToNot(HaveOccurred()) + bogusReq.WithSessionID("bogus-" + uuid.NewString()) + resp, err := client.Send(context.Background(), proxyURL, bogusReq) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).ToNot(Equal(200), "a bogus/unregistered Mcp-Session-Id must be rejected") + Expect(resp.Error).ToNot(BeNil(), "rejection must be a JSON-RPC error") + + for round := 0; round < rounds; round++ { + runMixedRound(client, proxyURL, round, sessionIDs, modernPerRound, joinTimeout) + } + }) +}) + +// mixResult is one goroutine's outcome within a single mixed round. +type mixResult struct { + era string // "legacy" or "modern", for failure messages + nonce string + gotNonce string + err error + statusCode int + rpcErr *e2e.RawRPCError + sessionHdr string // response's Mcp-Session-Id header, if any + start, end time.Time +} + +// runMixedRound fires len(sessionIDs) Legacy requests (each reusing its own +// already-initialized session) and modernPerRound stateless Modern requests +// concurrently -- a start channel releases every goroutine together, so both +// eras are genuinely in flight at once, not sequential per era (the crown +// jewel's anti-theater fire pattern). Every request carries a unique nonce +// in _meta; a wrong echoed nonce is a cross-era or cross-client leak. +func runMixedRound( + client *e2e.RawMCPClient, + proxyURL string, + round int, + sessionIDs []string, + modernPerRound int, + joinTimeout time.Duration, +) { + total := len(sessionIDs) + modernPerRound + results := make([]mixResult, total) + var wg sync.WaitGroup + start := make(chan struct{}) + + // Derived from joinTimeout (not context.Background()) so a stuck round + // tears down its in-flight requests promptly, mirroring the crown jewel. + ctx, cancel := context.WithTimeout(context.Background(), joinTimeout) + defer cancel() + + fire := func(i int, era string, build func() (*e2e.RawRequest, error)) { + wg.Add(1) + go func() { + defer wg.Done() + nonce := fmt.Sprintf("round-%d-%s-%d", round, era, i) + <-start + + req, err := build() + if err != nil { + results[i] = mixResult{era: era, nonce: nonce, err: err} + return + } + req.SetMeta("nonce", nonce) + + reqStart := time.Now() + resp, sendErr := client.Send(ctx, proxyURL, req) + reqEnd := time.Now() + if sendErr != nil { + results[i] = mixResult{era: era, nonce: nonce, err: sendErr, start: reqStart, end: reqEnd} + return + } + results[i] = mixResult{ + era: era, + nonce: nonce, + statusCode: resp.StatusCode, + rpcErr: resp.Error, + gotNonce: extractNonce(resp), + sessionHdr: resp.Headers.Get(e2e.HeaderMCPSessionID), + start: reqStart, + end: reqEnd, + } + }() + } + + for i, sid := range sessionIDs { + sid := sid + fire(i, "legacy", func() (*e2e.RawRequest, error) { + req, err := e2e.NewLegacyRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "mixcheck"}, + }) + if err != nil { + return nil, err + } + return req.WithSessionID(sid), nil + }) + } + for i := 0; i < modernPerRound; i++ { + fire(len(sessionIDs)+i, "modern", func() (*e2e.RawRequest, error) { + return e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "mixcheck"}, + }) + }) + } + + close(start) // release every goroutine (both eras) together + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(joinTimeout): + Fail(fmt.Sprintf("round %d: timed out waiting for %d mixed requests to finish", round, total)) + } + + distinctNonces := map[string]bool{} + var legacyOverlapsModern bool + for i, r := range results { + ExpectWithOffset(1, r.err).ToNot(HaveOccurred(), "round %d %s request %d", round, r.era, i) + ExpectWithOffset(1, r.rpcErr).To(BeNil(), "round %d %s request %d: %+v", round, r.era, i, r.rpcErr) + ExpectWithOffset(1, r.statusCode).To(Equal(200), "round %d %s request %d", round, r.era, i) + ExpectWithOffset(1, r.gotNonce).To(Equal(r.nonce), + "round %d %s request %d: echoed nonce must match -- a mismatch is a cross-era/cross-client leak", + round, r.era, i) + if r.era == "modern" { + ExpectWithOffset(1, r.sessionHdr).To(BeEmpty(), + "round %d modern request %d: stateless Modern must never carry Mcp-Session-Id", round, i) + } + distinctNonces[r.gotNonce] = true + + // Anti-theater: both eras must have genuinely overlapped in flight at + // least once this round, not fired sequentially era-by-era. + for _, other := range results { + if r.era != other.era && r.start.Before(other.end) && other.start.Before(r.end) { + legacyOverlapsModern = true + } + } + } + ExpectWithOffset(1, distinctNonces).To(HaveLen(total), "round %d: must see %d distinct correlated nonces", round, total) + ExpectWithOffset(1, legacyOverlapsModern).To(BeTrue(), + "round %d: no Legacy request overlapped in time with a Modern request -- "+ + "this round fired the eras sequentially instead of concurrently", round) +} diff --git a/test/e2e/dual_era_observability_test.go b/test/e2e/dual_era_observability_test.go new file mode 100644 index 0000000000..00528c74ae --- /dev/null +++ b/test/e2e/dual_era_observability_test.go @@ -0,0 +1,183 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "bytes" + "compress/gzip" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + coltracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1" + "google.golang.org/protobuf/proto" + + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// This spec is the observability tier of the e2e suite (issue #5837, Commit +// 2, Step 2.3): a real OTLP/HTTP collector, not a log-grep, receives actual +// spans from a live `thv run` process and asserts real attribute values -- +// telemetry_middleware_e2e_test.go only greps proxy log lines and never +// inspects a span. mcp.protocol.version and mcp.client.name are set in +// pkg/telemetry/middleware.go from the incoming MCP-Protocol-Version header +// and _meta.clientInfo respectively (Modern's only source of client +// attribution, since Modern has no initialize request -- see +// addMethodSpecificAttributes's doc comment). + +// otlpSpan is a minimal, test-only projection of a received OTLP span. +type otlpSpan struct { + Name string + Attributes map[string]string +} + +// otlpTraceCollector is an in-process OTLP/HTTP trace receiver: it decodes +// real ExportTraceServiceRequest protobuf (the wire format +// otlptracehttp actually sends -- no compression by default, see +// pkg/telemetry/providers/otlp/tracing.go, now toolhive-core/telemetry/providers/otlp) and records every span's name and +// string attributes for assertions. +type otlpTraceCollector struct { + server *httptest.Server + mu sync.Mutex + spans []otlpSpan +} + +func newOTLPTraceCollector() *otlpTraceCollector { + c := &otlpTraceCollector{} + c.server = httptest.NewServer(http.HandlerFunc(c.handle)) + return c +} + +func (c *otlpTraceCollector) handle(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + if r.Header.Get("Content-Encoding") == "gzip" { + gz, gzErr := gzip.NewReader(bytes.NewReader(body)) + if gzErr != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + defer gz.Close() + if body, err = io.ReadAll(gz); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + } + + var req coltracepb.ExportTraceServiceRequest + if err := proto.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + c.mu.Lock() + defer c.mu.Unlock() + for _, rs := range req.GetResourceSpans() { + for _, ss := range rs.GetScopeSpans() { + for _, s := range ss.GetSpans() { + attrs := make(map[string]string, len(s.GetAttributes())) + for _, kv := range s.GetAttributes() { + attrs[kv.GetKey()] = kv.GetValue().GetStringValue() + } + c.spans = append(c.spans, otlpSpan{Name: s.GetName(), Attributes: attrs}) + } + } + } + w.WriteHeader(http.StatusOK) +} + +// Spans returns a snapshot of every span received so far. +func (c *otlpTraceCollector) Spans() []otlpSpan { + c.mu.Lock() + defer c.mu.Unlock() + return append([]otlpSpan(nil), c.spans...) +} + +func (c *otlpTraceCollector) URL() string { return c.server.URL } +func (c *otlpTraceCollector) Stop() { c.server.Close() } + +var _ = Describe("Dual-Era Observability — OTLP Spans", Label("proxy", "stateless", "dual-era", "e2e", "telemetry"), Serial, func() { + var ( + config *e2e.TestConfig + serverName string + client *e2e.RawMCPClient + proxyURL string + collector *otlpTraceCollector + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("func-otel") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + collector = newOTLPTraceCollector() + + client, err = e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + e2e.NewTHVCommand(config, "run", + "--name", serverName, + "--transport", "stdio", + "-e", "BACKEND_MODE=echo", + "--otel-endpoint", collector.URL(), + "--otel-insecure", + "--otel-sampling-rate", "1.0", // always-sample: assertions must not race the default 10% ratio + images.YardstickServerImage, + ).ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if collector != nil { + collector.Stop() + } + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("emits mcp.protocol.version and mcp.client.name on a Modern span", func() { + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "otelcheck"}, + }) + Expect(err).ToNot(HaveOccurred()) + req.WithClientInfo("otel-e2e-client", "1.0") + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(200)) + Expect(resp.Error).To(BeNil()) + + // The SDK's BatchSpanProcessor flushes on its own timer (default 5s), + // not on every span -- poll rather than assert immediately. + Eventually(func() []otlpSpan { + return collector.Spans() + }, 20*time.Second, 500*time.Millisecond).Should(ContainElement(SatisfyAll( + HaveField("Attributes", HaveKeyWithValue("mcp.protocol.version", e2e.MCPVersionModern)), + HaveField("Attributes", HaveKeyWithValue("mcp.client.name", "otel-e2e-client")), + )), "expected a span carrying both mcp.protocol.version and mcp.client.name") + }) +}) diff --git a/test/e2e/dual_era_proxy_test.go b/test/e2e/dual_era_proxy_test.go new file mode 100644 index 0000000000..7311361860 --- /dev/null +++ b/test/e2e/dual_era_proxy_test.go @@ -0,0 +1,264 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// This spec is the crown-jewel regression test for the streamable HTTP +// proxy's per-request routing (pkg/transport/proxy/streamable): a Modern +// (2026-07-28) request always gets a fresh UUID routing token and never +// falls through to the client-supplied Mcp-Session-Id, even when many +// concurrent Modern requests share the same (session, id) pair. If that +// guarantee regresses, one client's response can be delivered to another +// client waiting on the identical composite key -- a confidentiality bug, +// see resolveSessionForRequest's doc comment in streamable_proxy.go. +// +// yardstick-server's barrier mode (BACKEND_MODE=barrier) buffers BARRIER_N +// non-lifecycle requests and releases all their responses at once, forcing +// BARRIER_N requests to be simultaneously in-flight at the proxy -- the +// precondition the leak needs. +// +// This spec loops internally (rounds), and CI additionally runs it repeated +// (ginkgo --repeat / go test -count) plus nightly, since a race window this +// narrow does not reliably reproduce on a single run. +var _ = Describe("Dual-Era Cross-Delivery Proxy", Label("proxy", "stateless", "dual-era", "e2e"), Serial, func() { + const ( + barrierN = 8 + rounds = 20 + barrierTimeout = 3 * time.Second + // clientTimeout bounds a single HTTP round trip: comfortably above + // barrierTimeout so a legitimate barrier release is never mistaken + // for a stuck proxy, but still fails the test (instead of the + // suite's global timeout) if the proxy hangs. + clientTimeout = barrierTimeout + 10*time.Second + // wellBeforeBarrierTimeout bounds how long any single response may + // take. Barrier mode's safety valve flushes a *partial* buffer on + // starvation, so "N responses came back" alone doesn't prove the + // barrier filled to N -- a response arriving near/after + // barrierTimeout means the round never actually forced all N + // requests to collide, and must fail rather than pass. + wellBeforeBarrierTimeout = barrierTimeout / 2 + + // All requests in a round deliberately share this session and id -- + // the exact leak precondition -- and differ only by nonce. + crossDeliverySessionID = "cross-delivery-probe" + crossDeliveryRequestID = int64(1) + ) + + var ( + config *e2e.TestConfig + serverName string + client *e2e.RawMCPClient + proxyURL string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("dual-era") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + client, err = e2e.NewRawMCPClient(clientTimeout) + Expect(err).ToNot(HaveOccurred()) + + By("Starting the yardstick backend under --transport stdio with BACKEND_MODE=barrier") + e2e.NewTHVCommand(config, "run", + "--name", serverName, + "--transport", "stdio", + "-e", "BACKEND_MODE=barrier", + "-e", fmt.Sprintf("BARRIER_N=%d", barrierN), + "-e", fmt.Sprintf("BARRIER_TIMEOUT_SECONDS=%d", int(barrierTimeout.Seconds())), + images.YardstickServerImage, + ).ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("never delivers one client's response to another under forced request collisions", func() { + // Positive control, run once: barrierN-1 requests can never fill the + // barrier, so they must be held until BARRIER_TIMEOUT_SECONDS's safety + // flush (~barrierTimeout), not answered immediately. Without this, a + // barrier that silently degraded into a no-op (answering every request + // right away) would be indistinguishable from a real barrier: both + // produce all-N-fast in the rounds below, so the crown jewel would stay + // green while forcing zero collisions. + By("Verifying the barrier actually withholds a short-of-N batch until its safety timeout") + results := fireCrossDeliveryBatch(client, proxyURL, "control", barrierN-1, barrierTimeout, + crossDeliverySessionID, crossDeliveryRequestID) + assertCorrelated(results, barrierN-1, "control") + for i, r := range results { + ExpectWithOffset(1, r.elapsed).To(BeNumerically(">=", wellBeforeBarrierTimeout), + "control request %d: returned too fast (%v) for a %d-of-%d batch -- the barrier is not "+ + "withholding responses below N, so it cannot be trusted to force the collision the rounds below rely on", + i, r.elapsed, barrierN-1, barrierN) + } + + for round := 0; round < rounds; round++ { + label := fmt.Sprintf("round-%d", round) + results := fireCrossDeliveryBatch(client, proxyURL, label, barrierN, barrierTimeout, + crossDeliverySessionID, crossDeliveryRequestID) + assertCorrelated(results, barrierN, label) + for i, r := range results { + ExpectWithOffset(1, r.elapsed).To(BeNumerically("<", wellBeforeBarrierTimeout), + "%s request %d: response arrived too close to BARRIER_TIMEOUT_SECONDS (%v) -- the barrier "+ + "likely flushed a partial buffer via its safety timeout instead of filling to N, meaning this "+ + "round never forced the collision", label, i, barrierTimeout) + } + } + }) +}) + +// crossDeliveryResult is one goroutine's outcome within a single barrier batch. +type crossDeliveryResult struct { + nonce string // the nonce this goroutine sent + err error + statusCode int + elapsed time.Duration + rpcErr *e2e.RawRPCError + gotNonce string // the nonce reflected back in the response, if any +} + +// fireCrossDeliveryBatch fires count concurrent Modern tools/call requests -- +// all sharing the same foreign Mcp-Session-Id and JSON-RPC id, differing only +// by a per-request nonce carried in _meta -- and returns each goroutine's +// outcome, indexed by request number. label distinguishes nonces/failure +// messages across batches (a round number, or "control"). +func fireCrossDeliveryBatch( + client *e2e.RawMCPClient, + proxyURL, label string, + count int, + barrierTimeout time.Duration, + sessionID string, + requestID int64, +) []crossDeliveryResult { + // joinTimeout bounds the whole batch: comfortably above barrierTimeout (the + // longest a request should ever be held) so a stuck proxy fails fast + // instead of hanging until the suite's global timeout. + joinTimeout := barrierTimeout + 15*time.Second + ctx, cancel := context.WithTimeout(context.Background(), joinTimeout) + defer cancel() + + results := make([]crossDeliveryResult, count) + var wg sync.WaitGroup + start := make(chan struct{}) + + for i := 0; i < count; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + nonce := fmt.Sprintf("%s-req-%d", label, i) + <-start // released together -- sequential firing deadlocks the barrier + + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "crossdelivery"}, + }) + if err != nil { + results[i] = crossDeliveryResult{nonce: nonce, err: err} + return + } + req.WithID(requestID).WithSessionID(sessionID).SetMeta("nonce", nonce) + + reqStart := time.Now() + resp, sendErr := client.Send(ctx, proxyURL, req) + elapsed := time.Since(reqStart) + if sendErr != nil { + results[i] = crossDeliveryResult{nonce: nonce, err: sendErr, elapsed: elapsed} + return + } + results[i] = crossDeliveryResult{ + nonce: nonce, + statusCode: resp.StatusCode, + elapsed: elapsed, + rpcErr: resp.Error, + gotNonce: extractNonce(resp), + } + }(i) + } + + close(start) // fire all requests concurrently + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(joinTimeout): + Fail(fmt.Sprintf("%s: timed out waiting for %d goroutines to finish", label, count)) + } + + return results +} + +// assertCorrelated checks the universal cross-delivery invariants for a batch: +// every request succeeded with no error/JSON-RPC error/non-200, and every +// response's nonce matches the request that goroutine sent (a mismatch is a +// cross-delivery leak), with exactly wantCount distinct nonces observed. +func assertCorrelated(results []crossDeliveryResult, wantCount int, label string) { + distinctNonces := map[string]bool{} + for i, r := range results { + ExpectWithOffset(2, r.err).ToNot(HaveOccurred(), + "%s request %d: must not error or time out", label, i) + ExpectWithOffset(2, r.rpcErr).To(BeNil(), + "%s request %d: must not get a JSON-RPC error: %+v", label, i, r.rpcErr) + ExpectWithOffset(2, r.statusCode).To(Equal(http.StatusOK), + "%s request %d: must return HTTP 200", label, i) + ExpectWithOffset(2, r.gotNonce).To(Equal(r.nonce), + "%s request %d: response nonce must match the request this goroutine sent -- "+ + "a mismatch is a cross-delivery leak", label, i) + distinctNonces[r.gotNonce] = true + } + ExpectWithOffset(2, distinctNonces).To(HaveLen(wantCount), + "%s: must see exactly %d distinct correlated nonces", label, wantCount) +} + +// extractNonce pulls the echoed "nonce" key out of a tools/call response's +// result._meta, per yardstick-server's echo tool: "Also echoes back any +// _meta field from the request for testing metadata propagation." Returns +// "" if absent or unparsable, which the caller's nonce-equality assertion +// then reports as a mismatch. +func extractNonce(resp *e2e.RawResponse) string { + if resp.Result == nil { + return "" + } + var result struct { + Meta map[string]any `json:"_meta"` + } + if err := json.Unmarshal(resp.Result, &result); err != nil { + return "" + } + nonce, _ := result.Meta["nonce"].(string) + return nonce +} diff --git a/test/e2e/dual_era_resilience_test.go b/test/e2e/dual_era_resilience_test.go new file mode 100644 index 0000000000..dd8fea2c3e --- /dev/null +++ b/test/e2e/dual_era_resilience_test.go @@ -0,0 +1,247 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "context" + "fmt" + "net/http" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// This spec is the fault-injection tier of the e2e suite (issue #5837, +// Commit 3, Step 3.1): the streamable proxy must degrade gracefully when its +// backend hangs or crashes mid-request -- a clean bounded error to the +// caller, no proxy-wide hang, and other traffic unaffected. +// +// Confirmed live before writing this (yardstick-server 1.2.0's +// BACKEND_MODE=hang/crash, via a real `thv run` instance): +// +// - Hang: HANG_AFTER_N counts requests; the Nth one never gets a backend +// response. The proxy's own TOOLHIVE_PROXY_REQUEST_TIMEOUT bounds it (a +// plain-text 504 "Timeout waiting for response from container" via +// writeHTTPError -- not a JSON-RPC body, so RawResponse.Error stays nil; +// assert on StatusCode instead). Concurrent unrelated requests are NOT +// blocked while the hung one waits (verified: a concurrent request +// completed in 14ms while the hung one was still waiting on its own +// goroutine, and the hung one's later 504 arrived on schedule). The +// backend process itself survives a hang -- no restart, a request sent +// afterward succeeds immediately. +// - Crash: the in-flight request at CRASH_AFTER_N gets the same clean +// writeHTTPError response (a 5xx -- observed 504, "server error" is the +// load-bearing property, not the exact code) well within a few seconds -- the +// stdio transport detects the closed pipe faster than any configured +// timeout, it doesn't wait for TOOLHIVE_PROXY_REQUEST_TIMEOUT to elapse. +// Recovery is real (`pkg/workloads/manager.go`'s RunWorkload retries +// with backoff) but SLOW and layered in this environment: Docker's own +// restart policy fires first (~13s), races a doomed thv re-attach +// (409 conflict, observed), then thv's own workload-level restart +// kicks in only after re-attach gives up (~48s total in one measured +// run) -- and the proxy port is unreachable for the whole window. That +// is far outside what a per-spec e2e timeout should eat, and getting it +// reliably non-flaky would mean re-resolving the proxy URL and polling +// through repeated container recreation, which is disproportionate to +// this step. Per the plan's explicit either/or, this spec asserts only +// the clean-error window for crash, not full recovery -- documented +// here as a deliberate scope decision, not an oversight. +// +// Redis-unavailable -> 503 (case 3) is OUT OF SCOPE for this file: `thv run` +// has NO CLI flag exposing Redis-backed session storage at all -- +// RunConfig.ScalingConfig.SessionRedis (pkg/runner/config.go) is set only via +// a hand-built RunConfig (e.g. `--from-config`) or the Kubernetes operator; +// there is no `thv run --session-storage redis` equivalent to stand up in a +// local e2e test. Standing this up would mean hand-authoring a full +// RunConfig JSON with no existing local-e2e precedent (the one Redis +// fixture in this suite, vmcp_infra_features_test.go, wires it into `thv +// vmcp serve`, a different proxy entirely). Flagging per the plan's +// "STOP and report" clause rather than building something disproportionate. +var _ = Describe("Dual-Era Resilience — Fault Injection", Label("proxy", "stateless", "dual-era", "e2e"), Serial, func() { + var ( + config *e2e.TestConfig + serverName string + client *e2e.RawMCPClient + proxyURL string + ) + + AfterEach(func() { + if config == nil || !config.CleanupAfter { + return + } + // The crash spec's backend may still be mid Docker-restart / thv + // re-attach (see the file doc comment -- ~48s, one measured run) when + // this fires; retry rather than hard-failing the spec on a transient + // "container is restarting" error masking the real test result. + Eventually(func() error { + return e2e.StopAndRemoveMCPServer(config, serverName) + }, 60*time.Second, 2*time.Second).Should(Succeed(), "should be able to stop and remove server") + }) + + startFaultBackend := func(mode string, extraEnv ...string) { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("resilience-" + mode) + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + client, err = e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + args := []string{ + "run", + "--name", serverName, + "--transport", "stdio", + "-e", "BACKEND_MODE=" + mode, + } + for _, e := range extraEnv { + args = append(args, "-e", e) + } + args = append(args, images.YardstickServerImage) + + // Bound the request timeout so a hang fails fast in the test, not + // after the streamable proxy's 60s default. + e2e.NewTHVCommand(config, args...).WithEnv("TOOLHIVE_PROXY_REQUEST_TIMEOUT=5s").ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + } + + echoRequest := func(input string) *e2e.RawRequest { + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": input}, + }) + Expect(err).ToNot(HaveOccurred()) + return req + } + + Context("backend hang", func() { + BeforeEach(func() { + // The 2nd request hangs; the 1st is a healthy baseline. + startFaultBackend("hang", "HANG_AFTER_N=2") + }) + + It("isolates a hung request from concurrent traffic, and the backend stays usable", func() { + By("a baseline request succeeds") + resp, err := client.Send(context.Background(), proxyURL, echoRequest("baseline")) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + + By("firing the hang-triggering request concurrently with an unrelated one") + var wg sync.WaitGroup + var hungErr, normalErr error + var hungStatus, normalStatus int + var hungElapsed, normalElapsed time.Duration + + // Deterministic ordering: launch the hang-trigger first and give + // it a head start to actually reach the backend as request #2, + // THEN launch "normal" -- racing both from a shared start gate + // left which request HANG_AFTER_N=2 counted as a coin flip (MoE + // flake finding). "normal" still overlaps the hang in time (fired + // while the hang-trigger is still in flight), which is exactly + // the isolation property under test. + wg.Add(1) + go func() { + defer wg.Done() + t0 := time.Now() + r, sendErr := client.Send(context.Background(), proxyURL, echoRequest("hangtrigger")) + hungElapsed = time.Since(t0) + hungErr = sendErr + if r != nil { + hungStatus = r.StatusCode + } + }() + time.Sleep(300 * time.Millisecond) + + wg.Add(1) + go func() { + defer wg.Done() + t0 := time.Now() + r, sendErr := client.Send(context.Background(), proxyURL, echoRequest("normal")) + normalElapsed = time.Since(t0) + normalErr = sendErr + if r != nil { + normalStatus = r.StatusCode + } + }() + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(20 * time.Second): + Fail("timed out waiting for concurrent hang/normal requests to finish") + } + + Expect(hungErr).ToNot(HaveOccurred()) + Expect(normalErr).ToNot(HaveOccurred()) + Expect(normalStatus).To(Equal(http.StatusOK), "the unrelated concurrent request must succeed") + Expect(normalElapsed).To(BeNumerically("<", 5*time.Second), + "the unrelated request must not be blocked behind the hung one") + Expect(hungStatus).To(Equal(http.StatusGatewayTimeout), "the hung request must get a clean bounded error") + Expect(hungElapsed).To(BeNumerically(">=", 4*time.Second), + "sanity: the hang must actually have been waited out by the proxy timeout, not failed for some other reason") + + By("Legacy recovers: a fresh initialize + call after the hang still works") + initReq := e2e.NewLegacyInitializeRequest("resilience-client", "1.0") + initResp, err := client.Send(context.Background(), proxyURL, initReq) + Expect(err).ToNot(HaveOccurred()) + Expect(initResp.StatusCode).To(Equal(http.StatusOK)) + sessionID := initResp.Headers.Get(e2e.HeaderMCPSessionID) + Expect(sessionID).ToNot(BeEmpty()) + + legacyReq, err := e2e.NewLegacyRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "postH"}, + }) + Expect(err).ToNot(HaveOccurred()) + legacyReq.WithSessionID(sessionID) + legacyResp, err := client.Send(context.Background(), proxyURL, legacyReq) + Expect(err).ToNot(HaveOccurred()) + Expect(legacyResp.StatusCode).To(Equal(http.StatusOK)) + Expect(legacyResp.Error).To(BeNil()) + }) + }) + + Context("backend crash", func() { + BeforeEach(func() { + // The 2nd request crashes the backend; the 1st is a healthy baseline. + startFaultBackend("crash", "CRASH_AFTER_N=2") + }) + + It("returns a clean error for the crashing request instead of hanging or panicking the proxy", func() { + By("a baseline request succeeds") + resp, err := client.Send(context.Background(), proxyURL, echoRequest("baseline")) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + + By("the crash-triggering request gets a clean, bounded error") + t0 := time.Now() + crashResp, err := client.Send(context.Background(), proxyURL, echoRequest("crashtrigger")) + elapsed := time.Since(t0) + Expect(err).ToNot(HaveOccurred(), "the client must get an HTTP response, not a connection error/hang") + Expect(crashResp.StatusCode).To(BeNumerically(">=", 500), "a crash must surface as a server error, not a silent 200") + // 10s: comfortably above the observed ~3-5s crash-detection, far below + // the ~48s full-recovery latency -- tight enough to actually fail if + // crash-handling regressed to riding out recovery instead of detecting + // the closed pipe. (The 15s client.Send timeout would fire first and + // fail the Expect above, so bounding by that alone is a tautology.) + Expect(elapsed).To(BeNumerically("<", 10*time.Second), + fmt.Sprintf("must fail fast (took %s) -- a crash must not ride out recovery", elapsed)) + }) + }) +}) diff --git a/test/e2e/export_test.go b/test/e2e/export_test.go index ab92ed3d0b..55c9794e46 100644 --- a/test/e2e/export_test.go +++ b/test/e2e/export_test.go @@ -5,7 +5,6 @@ package e2e_test import ( "encoding/json" - "fmt" "os" "path/filepath" "time" @@ -212,5 +211,5 @@ var _ = Describe("Export Command", Label("core", "export", "e2e"), func() { // generateExportTestServerName creates a unique server name for export tests func generateExportTestServerName(prefix string) string { - return fmt.Sprintf("%s-%d", prefix, GinkgoRandomSeed()) + return e2e.GenerateUniqueServerName(prefix) } diff --git a/test/e2e/group_rm_test.go b/test/e2e/group_rm_test.go index 09130caa7c..5ab622a0aa 100644 --- a/test/e2e/group_rm_test.go +++ b/test/e2e/group_rm_test.go @@ -4,10 +4,8 @@ package e2e_test import ( - "fmt" "os/exec" "strings" - "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -25,9 +23,8 @@ var _ = Describe("Group RM E2E Tests", Label("core", "groups", "e2e"), func() { BeforeEach(func() { config = e2e.NewTestConfig() - // Use a shared timestamp for all workload names in this test - groupName = fmt.Sprintf("group-rm-cancel-group-%d", time.Now().UnixNano()) - secondGroupName = fmt.Sprintf("group-rm-cancel-group-2-%d", time.Now().UnixNano()) + groupName = e2e.GenerateUniqueServerName("group-rm-cancel-group") + secondGroupName = e2e.GenerateUniqueServerName("group-rm-cancel-group-2") createdWorkloads = []string{} // Check if thv binary is available @@ -61,7 +58,7 @@ var _ = Describe("Group RM E2E Tests", Label("core", "groups", "e2e"), func() { Describe("thv group rm command", func() { It("should return error when group does not exist", func() { - groupName := fmt.Sprintf("group-rm-non-existent-group-%d", time.Now().UnixNano()) + groupName := e2e.GenerateUniqueServerName("group-rm-non-existent-group") _, stderr, err := e2e.NewTHVCommand(config, "group", "rm", groupName).ExpectFailure() Expect(err).To(HaveOccurred()) Expect(stderr).To(ContainSubstring("does not exist")) @@ -69,12 +66,11 @@ var _ = Describe("Group RM E2E Tests", Label("core", "groups", "e2e"), func() { It("should cancel deletion when user does not confirm", func() { // Add a workload to the group - workloadName := fmt.Sprintf("group-rm-test-workload-%d", time.Now().UnixNano()) + workloadName := e2e.GenerateUniqueServerName("group-rm-test-workload") createWorkloadInGroup(workloadName, groupName) // Verify the workload is running - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) + e2e.ExpectMCPServersRunning(config, workloadName) // Try to delete the group but provide 'n' for no cmd := exec.Command(config.THVBinary, "group", "rm", groupName) @@ -105,13 +101,18 @@ var _ = Describe("Group RM E2E Tests", Label("core", "groups", "e2e"), func() { }) It("should delete group with workloads", func() { - // Create workloads in the group - groupWorkload1 := fmt.Sprintf("group-rm-group-workload-1-%d", GinkgoRandomSeed()) - groupWorkload2 := fmt.Sprintf("group-rm-group-workload-2-%d", GinkgoRandomSeed()) + // Two members and two non-members are what make "rm keeps every + // member and leaves every non-member running" an assertion rather + // than an anecdote, so the setup stays at four. The workload count + // is not what makes this spec flaky: the same readiness wait has + // timed out in the two-workload specs here and in + // list_group_e2e_test.go (#6040). + groupWorkload1 := e2e.GenerateUniqueServerName("group-rm-group-workload-1") + groupWorkload2 := e2e.GenerateUniqueServerName("group-rm-group-workload-2") // Create workloads not in the group - nonGroupWorkload1 := fmt.Sprintf("group-rm-non-group-workload-1-%d", GinkgoRandomSeed()) - nonGroupWorkload2 := fmt.Sprintf("group-rm-non-group-workload-2-%d", GinkgoRandomSeed()) + nonGroupWorkload1 := e2e.GenerateUniqueServerName("group-rm-non-group-workload-1") + nonGroupWorkload2 := e2e.GenerateUniqueServerName("group-rm-non-group-workload-2") createWorkloadInGroup(groupWorkload1, groupName) createWorkloadInGroup(groupWorkload2, groupName) @@ -119,10 +120,7 @@ var _ = Describe("Group RM E2E Tests", Label("core", "groups", "e2e"), func() { createWorkloadInGroup(nonGroupWorkload2, secondGroupName) // Verify all workloads are running - for _, workloadName := range []string{groupWorkload1, groupWorkload2, nonGroupWorkload1, nonGroupWorkload2} { - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) - } + e2e.ExpectMCPServersRunning(config, groupWorkload1, groupWorkload2, nonGroupWorkload1, nonGroupWorkload2) // Delete the group (provide confirmation) cmd := exec.Command(config.THVBinary, "group", "rm", groupName) @@ -147,17 +145,14 @@ var _ = Describe("Group RM E2E Tests", Label("core", "groups", "e2e"), func() { It("should delete group and workloads with --with-workloads flag", func() { // Create multiple workloads in the group - workload1 := fmt.Sprintf("group-rm-with-workloads-1-%d", GinkgoRandomSeed()) - workload2 := fmt.Sprintf("group-rm-with-workloads-2-%d", GinkgoRandomSeed()) + workload1 := e2e.GenerateUniqueServerName("group-rm-with-workloads-1") + workload2 := e2e.GenerateUniqueServerName("group-rm-with-workloads-2") createWorkloadInGroup(workload1, groupName) createWorkloadInGroup(workload2, groupName) // Verify all workloads are running - for _, workloadName := range []string{workload1, workload2} { - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) - } + e2e.ExpectMCPServersRunning(config, workload1, workload2) // Delete the group with --with-workloads flag (provide confirmation) cmd := exec.Command(config.THVBinary, "group", "rm", groupName, "--with-workloads") diff --git a/test/e2e/group_test.go b/test/e2e/group_test.go index 546b2355e3..6401229cf5 100644 --- a/test/e2e/group_test.go +++ b/test/e2e/group_test.go @@ -35,7 +35,7 @@ var _ = Describe("Group", Label("core", "groups", "e2e"), func() { createWorkloadInGroup := func(workloadName, groupName string) { e2e.NewTHVCommand(config, "run", "fetch", "--group", groupName, "--name", workloadName).ExpectSuccess() createdWorkloads = append(createdWorkloads, workloadName) - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) + err := e2e.WaitForMCPServer(config, workloadName, e2e.ServerReadyTimeout()) Expect(err).ToNot(HaveOccurred()) } @@ -192,7 +192,7 @@ var _ = Describe("Group", Label("core", "groups", "e2e"), func() { _, _ = e2e.NewTHVCommand(config, "restart", workloadName).ExpectSuccess() By("Verifying the workload is still in the correct group") - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) + err := e2e.WaitForMCPServer(config, workloadName, e2e.ServerReadyTimeout()) Expect(err).ToNot(HaveOccurred()) workloadGroupName, err := getWorkloadGroup(workloadName) diff --git a/test/e2e/helpers.go b/test/e2e/helpers.go index cbf40af20d..1dccff283a 100644 --- a/test/e2e/helpers.go +++ b/test/e2e/helpers.go @@ -6,16 +6,21 @@ package e2e import ( "context" + "encoding/json" "fmt" "os" "os/exec" "path/filepath" "strings" + "sync" "syscall" "time" . "github.com/onsi/ginkgo/v2" //nolint:staticcheck // Standard practice for Ginkgo . "github.com/onsi/gomega" //nolint:staticcheck // Standard practice for Gomega + + rt "github.com/stacklok/toolhive/pkg/container/runtime" + "github.com/stacklok/toolhive/pkg/core" ) // GenerateUniqueServerName creates a unique server name for tests @@ -133,7 +138,45 @@ func (c *THVCommand) ExpectFailure() (string, string, error) { return stdout, stderr, err } -// WaitForMCPServer waits for an MCP server to be running +// ServerReadyTimeoutEnv overrides how long specs wait for a freshly started +// workload to reach the running state. +const ServerReadyTimeoutEnv = "TOOLHIVE_E2E_SERVER_READY_TIMEOUT" + +// defaultServerReadyTimeout is the readiness budget used when +// ServerReadyTimeoutEnv is unset. +// +// Reaching `running` covers more than starting a container: the detached worker +// also polls the workload's own MCP endpoint until initialize succeeds before it +// flips the status (pkg/runner.waitForInitializeSuccess, which the product is +// willing to wait 5 minutes for). Specs that start several workloads at once +// need a budget well above container-start time, otherwise they fail while the +// workload is still legitimately starting on a loaded CI runner. +const defaultServerReadyTimeout = 2 * time.Minute + +// ServerReadyTimeout returns the budget a spec should give a freshly started +// workload to reach the running state. Set ServerReadyTimeoutEnv to any Go +// duration (e.g. "3m") to raise it on a slow machine. +func ServerReadyTimeout() time.Duration { + raw := os.Getenv(ServerReadyTimeoutEnv) + if raw == "" { + return defaultServerReadyTimeout + } + + timeout, err := time.ParseDuration(raw) + if err != nil || timeout <= 0 { + Fail(fmt.Sprintf("%s must be a positive Go duration (e.g. \"3m\"), got %q", ServerReadyTimeoutEnv, raw)) + } + return timeout +} + +// WaitForMCPServer waits for an MCP server to be running. +// +// The status is read from the named workload's own record in `thv list --all +// --format json`. Substring-matching the text table instead would accept "this +// name appears somewhere and some workload is running", which can return while +// the named workload is still starting. A timeout reports the last status +// observed and dumps the server state, so a CI log shows whether the workload +// was still starting, had errored, or never appeared at all. func WaitForMCPServer(config *TestConfig, serverName string, timeout time.Duration) error { ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() @@ -141,28 +184,96 @@ func WaitForMCPServer(config *TestConfig, serverName string, timeout time.Durati ticker := time.NewTicker(1 * time.Second) defer ticker.Stop() + var lastObserved string for { + workload, err := findWorkload(config, serverName) + switch { + case err != nil: + // `thv list` can fail transiently while other workloads are + // starting, so keep polling and surface the failure only if the + // timeout expires. + lastObserved = fmt.Sprintf("thv list failed: %v", err) + case workload == nil: + lastObserved = "not listed" + case workload.Status == rt.WorkloadStatusRunning: + return nil + default: + lastObserved = fmt.Sprintf("status %q", workload.Status) + if workload.StatusContext != "" { + lastObserved += fmt.Sprintf(" (%s)", workload.StatusContext) + } + } + select { case <-ctx.Done(): - return fmt.Errorf("timeout waiting for MCP server %s to be running", serverName) + DebugServerState(config, serverName) + return fmt.Errorf("timeout waiting for MCP server %s to be running after %v; last observed: %s", + serverName, timeout, lastObserved) case <-ticker.C: - stdout, _, err := NewTHVCommand(config, "list").Run() - if err != nil { - continue - } - - // Check if the server is listed and running - if strings.Contains(stdout, serverName) && strings.Contains(stdout, "running") { - return nil - } } } } +// ExpectMCPServersRunning waits for each named workload to reach the running +// state within ServerReadyTimeout, and fails naming the workload that did not. +// +// All workloads are polled concurrently so that the total wait is bounded by +// the slowest workload rather than the sum of all waits. +func ExpectMCPServersRunning(config *TestConfig, serverNames ...string) { + errs := make([]error, len(serverNames)) + var wg sync.WaitGroup + for i, name := range serverNames { + wg.Add(1) + go func(i int, name string) { + defer wg.Done() + errs[i] = WaitForMCPServer(config, name, ServerReadyTimeout()) + }(i, name) + } + wg.Wait() + for i, serverName := range serverNames { + ExpectWithOffset(1, errs[i]).ToNot(HaveOccurred(), + "workload %s (of %v) should reach the running state", serverName, serverNames) + } +} + // IsServerRunning checks if an MCP server is running func IsServerRunning(config *TestConfig, serverName string) bool { - stdout, _ := NewTHVCommand(config, "list").ExpectSuccess() - return strings.Contains(stdout, serverName) && strings.Contains(stdout, "running") + workload, err := findWorkload(config, serverName) + ExpectWithOffset(1, err).ToNot(HaveOccurred(), "should be able to list workloads") + return workload != nil && workload.Status == rt.WorkloadStatusRunning +} + +// listTimeout bounds a single `thv list` call. It is well above how long a list +// takes even on a loaded runner, and well below any readiness budget: callers +// poll, so a list that overruns is retried rather than fatal. Run() would apply +// TestConfig.TestTimeout instead, letting one hung list outlast the wait it is +// being polled for. +const listTimeout = 30 * time.Second + +// FindWorkload returns the named workload's record as reported by `thv list`, or +// nil if it is not listed. newCmd builds each list invocation, so a spec that +// runs thv under an isolated config/home/data env passes its own builder and +// observes the workloads created in that state rather than the real config. +// +// --all is required so that workloads which have not reached running yet +// (starting, error) are visible rather than filtered out. +func FindWorkload(newCmd func(args ...string) *THVCommand, serverName string) (*core.Workload, error) { + stdout, stderr, err := newCmd("list", "--all", "--format", "json").RunWithTimeout(listTimeout) + if err != nil { + return nil, fmt.Errorf("thv list: %w (stderr: %s)", err, strings.TrimSpace(stderr)) + } + + var workloads []core.Workload + if err := json.Unmarshal([]byte(stdout), &workloads); err != nil { + return nil, fmt.Errorf("failed to parse thv list output %q: %w", stdout, err) + } + + for i := range workloads { + if workloads[i].Name == serverName { + return &workloads[i], nil + } + } + return nil, nil } // StopAndRemoveMCPServer stops and removes an MCP server @@ -238,6 +349,18 @@ func DebugServerState(config *TestConfig, serverName string) { GinkgoWriter.Printf("Server logs:\n%s\n", logs) } + // The container log alone has repeatedly been useless for readiness + // timeouts: a workload stuck in `starting` or flipped to `error` fails + // inside the detached supervisor (transport start, readiness probe, + // restart loop), whose output goes to the proxy log file — not to the + // container. Dump it too so a CI failure names the actual blocker. + proxyLogs, stderr, err := NewTHVCommand(config, "logs", serverName, "--proxy").Run() + if err != nil { + GinkgoWriter.Printf("Failed to get proxy logs: %v\nStderr: %s\n", err, stderr) + } else { + GinkgoWriter.Printf("Proxy logs:\n%s\n", proxyLogs) + } + GinkgoWriter.Printf("=== End debugging for %s ===\n", serverName) } @@ -355,3 +478,11 @@ func CreateFakeBrowserDir(tempDir string) (string, error) { } return dir, nil } + +// findWorkload returns the named workload's record as reported by `thv list`, +// or nil if it is not listed. +func findWorkload(config *TestConfig, serverName string) (*core.Workload, error) { + return FindWorkload(func(args ...string) *THVCommand { + return NewTHVCommand(config, args...) + }, serverName) +} diff --git a/test/e2e/hostile_input_proxy_test.go b/test/e2e/hostile_input_proxy_test.go new file mode 100644 index 0000000000..0f6bbeca91 --- /dev/null +++ b/test/e2e/hostile_input_proxy_test.go @@ -0,0 +1,278 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/test/e2e" +) + +// This spec is the hostile-input tier of the dual-era proxy e2e suite (issue +// #5837, tier 2): a hand-rolled raw client drives malformed/adversarial +// Modern (2026-07-28) traffic at the transparent proxy (remote MCP +// passthrough, session-aware -- no --stateless) fronting the in-process +// statelessMockMCPServer (see stateless_proxy_test.go), and asserts each +// rejection fires at the edge -- exact error.code, error.data, and echoed +// id -- without the backend ever being contacted. +// +// Known gaps, deliberately not asserted here (see the design plan and +// pkg/mcp/revision.go's ClassifyRevision doc comment): Mcp-Method/Mcp-Name +// header validation (unimplemented), batch/client-response smuggling +// (handled as Legacy today, deferred), and header-driven GET/DELETE 405 +// (that gate is stateless-mode-driven, not classifier-driven). Also a known +// spec divergence: ClassifyRevision is transport-agnostic and cannot tell +// "stdio, no header concept" (-32602 is correct) apart from "HTTP, header +// omitted" (where the draft's Server Validation rules make a missing +// MCP-Protocol-Version header itself a -32020 condition) -- see the TODO in +// revision.go. This suite exercises the HTTP transport but the classifier +// currently returns -32602 either way. +var _ = Describe("Hostile Input Proxy", Label("proxy", "stateless", "hostile-input", "e2e"), Serial, func() { + var ( + config *e2e.TestConfig + serverName string + mockServer *statelessMockMCPServer + client *e2e.RawMCPClient + proxyURL string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + serverName = e2e.GenerateUniqueServerName("hostile-input") + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + mockServer, err = newStatelessMockMCPServer() + Expect(err).ToNot(HaveOccurred(), "should be able to start mock server") + + client, err = e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + By("Starting thv against the in-process mock (transparent proxy, session-aware)") + e2e.NewTHVCommand(config, "run", "--name", serverName, mockServer.URL()+"/mcp").ExpectSuccess() + + err = e2e.WaitForMCPServer(config, serverName, 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "server should be running within 60 seconds") + + proxyURL, err = e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to get proxy URL") + if !strings.HasSuffix(proxyURL, "/mcp") { + proxyURL += "/mcp" + } + }) + + AfterEach(func() { + if mockServer != nil { + mockServer.Stop() + mockServer = nil + } + if config.CleanupAfter { + err := e2e.StopAndRemoveMCPServer(config, serverName) + Expect(err).ToNot(HaveOccurred(), "should be able to stop and remove server") + } + }) + + It("rejects hostile Modern inputs at the edge without contacting the backend", func() { + // A manual loop over cases, not DescribeTable: BeforeEach's thv run + // setup (~seconds) is shared once across all cases here instead of + // paying it 5x, one per DescribeTable entry. + for _, tc := range hostileRejectionCases() { + By(tc.name) + countBefore := mockServer.GetCount() + + req, err := e2e.NewModernRequest("tools/list", nil) + Expect(err).ToNot(HaveOccurred(), tc.name) + req.WithID(tc.id) + tc.mutate(req) + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred(), tc.name) + Expect(resp.StatusCode).To(Equal(tc.wantStatus), tc.name) + // Every rejection path here echoes the id, session-not-found + // included: #5945 removed the one asymmetry (see below). + Expect(resp.ID).To(Equal(json.Number(fmt.Sprintf("%d", tc.id))), "%s: echoed id", tc.name) + Expect(resp.Error).ToNot(BeNil(), tc.name) + Expect(resp.Error.Code).To(Equal(tc.wantCode), tc.name) + if tc.wantData != nil { + var gotData map[string]any + Expect(json.Unmarshal(resp.Error.Data, &gotData)).To(Succeed(), "%s: error.data must decode", tc.name) + Expect(gotData).To(Equal(tc.wantData), "%s: error.data", tc.name) + } else { + Expect(resp.Error.Data).To(BeEmpty(), "%s: error.data should be absent", tc.name) + } + Expect(mockServer.GetCount()).To(Equal(countBefore), + "%s: backend must not be contacted for a rejected request", tc.name) + } + }) + + It("forwards a well-formed Modern request with no session id (200)", func() { + countBefore := mockServer.GetCount() + + req, err := e2e.NewModernRequest("tools/list", nil) + Expect(err).ToNot(HaveOccurred()) + req.WithID(int64(2001)) + + resp, err := client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(resp.Error).To(BeNil()) + Expect(resp.ID).To(Equal(json.Number("2001"))) + // "exactly once", not just ">= 1": proves the proxy forwards a + // well-formed sessionless Modern request to the backend a single + // time -- no retry, no double-send. + Expect(mockServer.GetCount()).To(Equal(countBefore+1), "backend must be contacted exactly once") + }) + + It("handles oversized and truncated bodies cleanly, without a crash or hang", func() { + By("an oversized body is rejected with 413, not forwarded") + countBefore := mockServer.GetCount() + oversized := []byte(fmt.Sprintf( + `{"jsonrpc":"2.0","id":3001,"method":"tools/list","params":{"padding":"%s"}}`, + strings.Repeat("a", 9<<20))) // 9 MiB, above bodylimit.DefaultMaxRequestBodySize (8 MiB) + resp, err := client.SendRaw(context.Background(), proxyURL, + map[string]string{"Content-Type": "application/json"}, oversized) + Expect(err).ToNot(HaveOccurred(), "must not hang or crash the proxy") + Expect(resp.StatusCode).To(Equal(http.StatusRequestEntityTooLarge)) + Expect(mockServer.GetCount()).To(Equal(countBefore), + "oversized body must be rejected by bodylimit before the backend is contacted") + + // Unlike the classification cases above, an unparsable body is NOT + // rejected at the edge: parseRPCRequest fails to decode it as either a + // single request or a batch, so RoundTrip never calls ClassifyRevision + // and forwards the body unclassified (Legacy passthrough) -- the proxy + // does not validate JSON syntax before forwarding. The 400 below comes + // from the *backend's* own json.Unmarshal failure, after it has + // already been contacted. + By("a truncated body is forwarded as Legacy passthrough; the backend rejects it") + countBefore = mockServer.GetCount() + truncated := []byte(`{"jsonrpc":"2.0","id":3002,"method":"tools/list","para`) // deliberately cut mid-object + resp, err = client.SendRaw(context.Background(), proxyURL, + map[string]string{"Content-Type": "application/json"}, truncated) + Expect(err).ToNot(HaveOccurred(), "must not hang or crash the proxy") + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + Expect(mockServer.GetCount()).To(Equal(countBefore+1), + "an unparsable body skips classification and IS forwarded to the backend") + + By("the proxy still serves a well-formed request after the hostile inputs above") + req, err := e2e.NewModernRequest("tools/list", nil) + Expect(err).ToNot(HaveOccurred()) + req.WithID(int64(3003)) + resp, err = client.Send(context.Background(), proxyURL, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusOK), "proxy must survive the oversized/truncated requests above") + }) +}) + +// hostileCase is one adversarial Modern request: start from a well-formed +// NewModernRequest, apply mutate to break exactly one classification/session +// rule, and expect the proxy to reject it before the backend is contacted. +type hostileCase struct { + name string + id int64 + mutate func(req *e2e.RawRequest) + wantStatus int + wantCode int64 + wantData map[string]any // nil means the error must carry no "data" field at all +} + +// hostileRejectionCases enumerates the classification-error and session-guard +// rejections the transparent proxy must produce for adversarial Modern +// traffic. Each starts from NewModernRequest's well-formed default (header +// MCP-Protocol-Version + _meta.protocolVersion both "2026-07-28", +// _meta.clientCapabilities {}) and mutates exactly the one thing that breaks it. +func hostileRejectionCases() []hostileCase { + return []hostileCase{ + { + name: "header/body protocolVersion mismatch -> -32020", + id: 1001, + mutate: func(req *e2e.RawRequest) { + // Body still claims Modern (NewModernRequest's default _meta); + // the header disagrees. Neither value wins -- hard rejection. + req.SetHeader(e2e.HeaderMCPProtocolVersion, "2025-11-25") + }, + wantStatus: http.StatusBadRequest, + wantCode: mcp.CodeHeaderMismatch, + wantData: map[string]any{"header": "2025-11-25", "body": e2e.MCPVersionModern}, + }, + { + name: "unsupported version in _meta.protocolVersion -> -32022", + id: 1002, + mutate: func(req *e2e.RawRequest) { + req.SetMeta(e2e.MetaKeyProtocolVersion, "2099-01-01") + }, + wantStatus: http.StatusBadRequest, + wantCode: mcp.CodeUnsupportedProtocolVersion, + wantData: map[string]any{"supported": []any{e2e.MCPVersionModern}, "requested": "2099-01-01"}, + }, + { + name: "missing clientCapabilities -> -32021", + id: 1003, + mutate: func(req *e2e.RawRequest) { + req.DeleteMeta(e2e.MetaKeyClientCapabilities) + }, + wantStatus: http.StatusBadRequest, + wantCode: mcp.CodeMissingClientCapability, + wantData: map[string]any{"requiredCapabilities": map[string]any{}}, + }, + { + name: "reserved _meta key, no header, no valid version -> -32602", + id: 1004, + mutate: func(req *e2e.RawRequest) { + // clientCapabilities alone remains as the reserved-key Modern + // signal; with no header and no body protocolVersion, the draft + // spec defines no dedicated code, so this falls back to the + // standard JSON-RPC Invalid Params code. + req.DeleteHeader(e2e.HeaderMCPProtocolVersion) + req.DeleteMeta(e2e.MetaKeyProtocolVersion) + }, + wantStatus: http.StatusBadRequest, + wantCode: mcp.CodeInvalidParams, + wantData: nil, + }, + { + name: "foreign Mcp-Session-Id -> 404 session-not-found", + id: 1005, + mutate: func(req *e2e.RawRequest) { + // Otherwise well-formed Modern request, but with a session id + // never registered with this proxy. The 2026-07-28 draft spec + // says Modern SHOULD ignore session id -- and the *streamable* + // proxy does (resolveSessionForRequest mints a fresh routing + // token unconditionally, see the crown jewel). The transparent + // proxy deliberately deviates: here Mcp-Session-Id has + // downstream authority -- it's looked up in the shared session + // store, rewritten to the backend's assigned SID, and forwarded + // to a stateful backend (see RoundTrip's guard + rewrite in + // transparent_proxy.go). Branching that guard on the + // client-forgeable classified revision instead of header + // presence would be fail-open: a forged Modern _meta could + // carry a downstream-authoritative SID past session validation. + // Keying on presence alone is the fail-safe choice, and it is + // the security invariant this suite exists to pin down -- + // see TestGuardUnknownSessionFiresDespiteForgedModernRevision + // in revision_guard_regression_test.go. + req.WithSessionID("foreign-" + uuid.NewString()) + }, + // Like the classification rejections above, this 404 echoes the + // request id: session.NotFoundResponse used to hardcode a nil id + // while ClassificationErrorResponse on the same RoundTrip echoed + // it, an asymmetry #5945 removed by threading the parsed id through + // to the guard. + wantStatus: http.StatusNotFound, + wantCode: session.CodeSessionNotFound, + wantData: nil, + }, + } +} diff --git a/test/e2e/images/images.go b/test/e2e/images/images.go index 8b112dd45a..0f2e9a2f7e 100644 --- a/test/e2e/images/images.go +++ b/test/e2e/images/images.go @@ -12,9 +12,13 @@ package images const ( yardstickServerImageURL = "ghcr.io/stackloklabs/yardstick/yardstick-server" - yardstickServerImageTag = "1.1.1" - // YardstickServerImage is used in operator tests across multiple transport protocols - // (stdio, SSE, streamable-http) and tenancy modes. + yardstickServerImageTag = "1.2.0" + // YardstickServerImage (go-sdk v1.7) is the yardstick backend used across the + // operator, vMCP, and dual-era e2e tests, over multiple transport protocols (stdio, + // SSE, streamable-http) and tenancy modes. A session-mode deployment negotiates down + // to the Legacy (2025-11-25) data plane and the vMCP classifies it Legacy from + // server/discover's supportedVersions (#5993), so this single image serves both + // eras -- the dual-era tests drive it Modern (stateless), the rest drive it Legacy. // Note: This image is also referenced in 8 YAML fixture files under // test/e2e/chainsaw/operator/. Those files are declarative Kubernetes manifests // and cannot import Go constants directly. @@ -86,8 +90,22 @@ const ( // Reference MCP test server providing ~16 diverse example tools. EverythingMCPServerImage = everythingMCPServerImageURL + ":" + everythingMCPServerImageTag - idaProMCPServerImageURL = "ghcr.io/stacklok/dockyard/uvx/ida-pro-mcp" - idaProMCPServerImageTag = "1.4.0" + idaProMCPServerImageURL = "ghcr.io/stacklok/dockyard/uvx/ida-pro-mcp" + idaProMCPServerImageTag = "1.4.0" + // idaProMCPServerImageDigest pins the 2026-07-27 build of tag 1.4.0. + // + // The tag is mutable: dockyard periodically rebuilds it and resolves Python + // dependencies fresh each time. ida-pro-mcp 1.4.0 declares an unbounded + // "mcp>=1.16.0" and imports mcp.server.fastmcp, which was removed in the + // mcp Python SDK 2.0.0 (released 2026-07-28). The 2026-07-31 rebuild picked + // up mcp 2.0.0, so the container now crashloops on startup with + // "ModuleNotFoundError: No module named 'mcp.server.fastmcp'". + // + // 1.4.0 is the newest ida-pro-mcp release on PyPI, so there is nothing to + // bump to. Pin the last build that resolved a 1.x mcp (1.28.1) until + // upstream caps the constraint or adopts the 2.x API. Every other backend + // in the optimizer test is unaffected -- pagerduty-mcp, the only other + // dockyard uvx image, caps "mcp[cli]~=1.8". idaProMCPServerImageDigest = "sha256:5a596d965fe05052d615a41152213f93ebc72bf46b07304718de948833d73c70" // IDAProMCPServerImage is used for testing multi-backend optimizer scenarios. // Provides ~47 IDA Pro reverse engineering tools (decompile, disassemble, rename, etc.). @@ -99,6 +117,33 @@ const ( // Provides ~64 PagerDuty incident management tools (incidents, services, schedules, etc.). PagerDutyMCPServerImage = pagerdutyMCPServerImageURL + ":" + pagerdutyMCPServerImageTag + timeMCPServerImageURL = "ghcr.io/stacklok/dockyard/uvx/mcp-server-time" + timeMCPServerImageTag = "2026.7.10" + // timeMCPServerImageDigest pins the 2026-07-27 build of tag 2026.7.10. + // + // Same failure mode as idaProMCPServerImageDigest above, different symptom: + // the tag is mutable, dockyard rebuilds it and resolves Python dependencies + // fresh, and the 2026-07-31 rebuild picked up the mcp Python SDK 2.0.0 + // released 2026-07-28. mcp-server-time imports McpError from + // mcp.shared.exceptions, which 2.x moved, so the container now exits on + // startup with "ImportError: cannot import name 'McpError' from + // 'mcp.shared.exceptions'". The proxy suites then fail with + // "timeout waiting for MCP server to be ready". + // + // Pin the last build that resolved a 1.x mcp until upstream caps the + // constraint or adopts the 2.x API. Verified by driving both builds over + // stdio: the pinned digest answers initialize, tag 2026.7.10 does not. + timeMCPServerImageDigest = "sha256:5cca77dec3fefbacad35e3008ab1660f8725ef246213afb24231504fc73c999d" + // TimeMCPServerImage is the stdio backend used by the proxy e2e suites + // (proxy_stdio_test.go, stdio_proxy_over_streamable_http_mcp_server_test.go). + // Provides get_current_time / convert_time. + // + // These suites reference this constant rather than the "time" registry entry + // so a mutable upstream tag cannot break CI: the registry resolves "time" to + // the moving :2026.7.10 tag, which is currently broken. Registry-name + // resolution itself is still covered by the suites that run "osv". + TimeMCPServerImage = timeMCPServerImageURL + ":" + timeMCPServerImageTag + "@" + timeMCPServerImageDigest + redisImageURL = "redis" redisImageTag = "7-alpine" // RedisImage is used for Redis-backed session storage in scaling tests. diff --git a/test/e2e/list_group_e2e_test.go b/test/e2e/list_group_e2e_test.go index 7a3423c198..c7fb0e48d7 100644 --- a/test/e2e/list_group_e2e_test.go +++ b/test/e2e/list_group_e2e_test.go @@ -70,8 +70,7 @@ var _ = Describe("List Group", Label("core", "groups", "e2e"), func() { _, _ = e2e.NewTHVCommand(config, "run", "fetch", "--group", groupName, "--name", workloadName).ExpectSuccess() // Wait for workload to be fully registered - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) + e2e.ExpectMCPServersRunning(config, workloadName) }) AfterEach(func() { @@ -131,8 +130,7 @@ var _ = Describe("List Group", Label("core", "groups", "e2e"), func() { _, _ = e2e.NewTHVCommand(config, "run", "fetch", "--group", groupName, "--name", workloadName, "--label", "test=value").ExpectSuccess() // Wait for workload to be fully registered - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) + e2e.ExpectMCPServersRunning(config, workloadName) }) AfterEach(func() { @@ -194,10 +192,7 @@ var _ = Describe("List Group", Label("core", "groups", "e2e"), func() { _, _ = e2e.NewTHVCommand(config, "run", "fetch", "--group", secondGroupName, "--name", workload2Name).ExpectSuccess() // Wait for workloads to be fully registered - err := e2e.WaitForMCPServer(config, workload1Name, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) - err = e2e.WaitForMCPServer(config, workload2Name, 60*time.Second) - Expect(err).ToNot(HaveOccurred()) + e2e.ExpectMCPServersRunning(config, workload1Name, workload2Name) }) AfterEach(func() { diff --git a/test/e2e/llm_gateway_mock.go b/test/e2e/llm_gateway_mock.go index f7fb27f208..a920103450 100644 --- a/test/e2e/llm_gateway_mock.go +++ b/test/e2e/llm_gateway_mock.go @@ -13,6 +13,7 @@ import ( "crypto/x509/pkix" "encoding/json" "encoding/pem" + "errors" "fmt" "math/big" "net" @@ -38,10 +39,11 @@ type GatewayRequest struct { // for plain HTTP. The HTTP variant is simpler for e2e tests where the thv // subprocess cannot be easily configured to trust a self-signed cert. type LLMGatewayMock struct { - server *http.Server - port int - useTLS bool - certPEM []byte // non-nil only when useTLS is true + server *http.Server + listener net.Listener + port int + useTLS bool + certPEM []byte // non-nil only when useTLS is true mu sync.Mutex requests []GatewayRequest @@ -50,18 +52,31 @@ type LLMGatewayMock struct { // NewLLMGatewayMock creates a mock LLM gateway that serves HTTPS with a // self-signed certificate. Use CertPEM / TLSClientConfig to build a trusting // HTTP client. Call Start to begin serving. +// +// The listener is bound before anything else, closing the window between +// port selection and bind that a separate "pick a port, then bind later" +// step would leave open. Pass port 0 to let the OS choose an ephemeral port; +// use Port() to read back the port actually bound. func NewLLMGatewayMock(port int) (*LLMGatewayMock, error) { + listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + return nil, fmt.Errorf("failed to listen: %w", err) + } + realPort := listener.Addr().(*net.TCPAddr).Port + certPEM, keyPEM, err := generateGatewayCert() if err != nil { + _ = listener.Close() return nil, err } cert, err := tls.X509KeyPair(certPEM, keyPEM) if err != nil { + _ = listener.Close() return nil, fmt.Errorf("loading TLS key pair: %w", err) } - m := &LLMGatewayMock{port: port, useTLS: true, certPEM: certPEM} + m := &LLMGatewayMock{listener: listener, port: realPort, useTLS: true, certPEM: certPEM} m.server = m.newServer() m.server.TLSConfig = &tls.Config{ Certificates: []tls.Certificate{cert}, @@ -71,6 +86,11 @@ func NewLLMGatewayMock(port int) (*LLMGatewayMock, error) { return m, nil } +// Port returns the port the mock gateway is bound to. +func (m *LLMGatewayMock) Port() int { + return m.port +} + func (m *LLMGatewayMock) newServer() *http.Server { mux := http.NewServeMux() mux.HandleFunc("/v1/models", m.handleModels) @@ -84,7 +104,6 @@ func (m *LLMGatewayMock) newServer() *http.Server { _, _ = w.Write([]byte("OK")) }) return &http.Server{ - Addr: fmt.Sprintf(":%d", m.port), Handler: mux, ReadHeaderTimeout: 10 * time.Second, } @@ -96,11 +115,11 @@ func (m *LLMGatewayMock) Start() error { go func() { var err error if m.useTLS { - err = m.server.ListenAndServeTLS("", "") + err = m.server.ServeTLS(m.listener, "", "") } else { - err = m.server.ListenAndServe() + err = m.server.Serve(m.listener) } - if err != nil && err != http.ErrServerClosed { + if err != nil && !errors.Is(err, http.ErrServerClosed) { errCh <- err } }() diff --git a/test/e2e/mcp_client_helpers.go b/test/e2e/mcp_client_helpers.go index f08b8a2e10..1b989e05e7 100644 --- a/test/e2e/mcp_client_helpers.go +++ b/test/e2e/mcp_client_helpers.go @@ -11,11 +11,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" //nolint:staticcheck // Standard practice for Ginkgo . "github.com/onsi/gomega" //nolint:staticcheck // Standard practice for Gomega + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // MCPClientHelper provides high-level MCP client operations for e2e tests diff --git a/test/e2e/mcp_raw_client.go b/test/e2e/mcp_raw_client.go new file mode 100644 index 0000000000..455ea166cd --- /dev/null +++ b/test/e2e/mcp_raw_client.go @@ -0,0 +1,480 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "maps" + "net/http" + "strings" + "sync/atomic" + "time" +) + +// The following mirror the reserved wire constants defined in +// pkg/mcp/revision.go. They are duplicated here rather than imported: this +// client must stay independent of any MCP library (production or SDK) so it +// can emit traffic no conforming implementation would ever produce. +const ( + // MCPVersionModern is the Modern (stateless) MCP protocol version. + MCPVersionModern = "2026-07-28" + + // MetaKeyProtocolVersion is the reserved _meta key carrying the per-request + // protocol version on Modern requests. + MetaKeyProtocolVersion = "io.modelcontextprotocol/protocolVersion" + // MetaKeyClientInfo is the reserved _meta key carrying client implementation + // info on Modern requests. Optional: a Modern request may omit it. + MetaKeyClientInfo = "io.modelcontextprotocol/clientInfo" + // MetaKeyClientCapabilities is the reserved _meta key carrying client + // capabilities on Modern requests. + MetaKeyClientCapabilities = "io.modelcontextprotocol/clientCapabilities" + + // HeaderMCPProtocolVersion is the HTTP header carrying the Modern protocol + // version. + HeaderMCPProtocolVersion = "MCP-Protocol-Version" + // HeaderMCPSessionID is the HTTP header carrying the Legacy session id. + HeaderMCPSessionID = "Mcp-Session-Id" + // HeaderMCPMethod is the HTTP header carrying the JSON-RPC method on Modern + // requests. A real go-sdk v1.7 streamable-HTTP server requires it on every + // Modern request (see pkg/mcp/revision.go ValidateHeaderConsistency); the + // ToolHive single-server proxy does not read it, but a Modern backend does. + HeaderMCPMethod = "Mcp-Method" + // HeaderMCPName is the HTTP header naming the request's target + // tool/resource/prompt on name-bearing Modern methods (tools/call, + // resources/read, prompts/get). Sent as the plain identifier; the draft + // spec also permits a base64 sentinel encoding, which we do not use. + HeaderMCPName = "Mcp-Name" +) + +// requestIDCounter assigns default JSON-RPC ids so distinct RawRequests get +// distinct ids unless a test overrides one via WithID (e.g. to collide two +// ids on purpose). +var requestIDCounter atomic.Int64 + +// RawRequest builds a single JSON-RPC request for MCP-over-HTTP, or carries +// a raw pre-built body that bypasses the builder entirely. Construct one +// with NewLegacyRequest, NewLegacyInitializeRequest, or NewModernRequest, +// then chain the With*/Set*/Delete* methods to adjust it. The zero value is +// not usable. +// +// Finish building before sending: Send/SendRaw only read a RawRequest, so a +// fully-built one is safe to share read-only across goroutines (e.g. one +// RawRequest fanned out to many concurrent senders), but mutating it with +// Set*/With* concurrently with a send is a data race. +type RawRequest struct { + method string + id any + params map[string]any + meta map[string]any // nil means no _meta object at all + headers map[string]string + body []byte // when set, overrides method/id/params/meta entirely +} + +// NewLegacyRequest builds a Legacy (2025-11-25) JSON-RPC request for the +// given method. params is cloned; the caller's map is never mutated by the +// returned RawRequest. +func NewLegacyRequest(method string, params map[string]any) (*RawRequest, error) { + return newRawRequest(method, params) +} + +// NewLegacyInitializeRequest builds a standard Legacy "initialize" request. +func NewLegacyInitializeRequest(clientName, clientVersion string) *RawRequest { + r, _ := newRawRequest("initialize", map[string]any{ // method is a fixed non-empty literal; newRawRequest cannot fail here. + "protocolVersion": "2025-11-25", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{ + "name": clientName, + "version": clientVersion, + }, + }) + return r +} + +// NewModernRequest builds a Modern (2026-07-28) JSON-RPC request emitting the +// full conformant wire shape a real go-sdk v1.7 client produces: the +// MCP-Protocol-Version header, the reserved _meta keys protocolVersion and +// clientCapabilities (clientInfo is optional per the draft schema and is +// omitted by default), the Mcp-Method header (always), and the Mcp-Name header +// for name-bearing methods (tools/call, resources/read, prompts/get). A Modern +// backend rejects a request missing these headers (-32020); the ToolHive +// single-server proxy ignores them. Each field can be independently overridden +// or removed via SetHeader/DeleteHeader and SetMeta/DeleteMeta. +func NewModernRequest(method string, params map[string]any) (*RawRequest, error) { + r, err := newRawRequest(method, params) + if err != nil { + return nil, err + } + r.headers[HeaderMCPProtocolVersion] = MCPVersionModern + r.headers[HeaderMCPMethod] = method + if name, ok := modernNameHeader(method, params); ok { + r.headers[HeaderMCPName] = name + } + r.meta = map[string]any{ + MetaKeyProtocolVersion: MCPVersionModern, + MetaKeyClientCapabilities: map[string]any{}, + } + return r, nil +} + +// modernNameHeader returns the plain Mcp-Name header value for a name-bearing +// Modern method, mirroring go-sdk v1.7's streamable-HTTP encoding and +// pkg/mcp's ValidateHeaderConsistency: tools/call and prompts/get name the +// "name" param, resources/read names the "uri" param. Any other method (e.g. +// tools/list, server/discover) carries no Mcp-Name, and sending one there +// would fail the backend's consistency check, so it is omitted. +func modernNameHeader(method string, params map[string]any) (string, bool) { + key := "" + switch method { + case "tools/call", "prompts/get": + key = "name" + case "resources/read": + key = "uri" + } + if key == "" { + return "", false + } + v, ok := params[key].(string) + return v, ok && v != "" +} + +func newRawRequest(method string, params map[string]any) (*RawRequest, error) { + if method == "" { + return nil, errors.New("mcp_raw_client: method must not be empty") + } + return &RawRequest{ + method: method, + id: requestIDCounter.Add(1), + params: maps.Clone(params), + headers: map[string]string{}, + }, nil +} + +// WithID overrides the JSON-RPC id, e.g. to collide two requests' ids or to +// use a value outside the int64-safe float64 range. +func (r *RawRequest) WithID(id any) *RawRequest { + r.id = id + return r +} + +// WithSessionID sets the Mcp-Session-Id header, including a foreign/spoofed +// session id. +func (r *RawRequest) WithSessionID(sessionID string) *RawRequest { + return r.SetHeader(HeaderMCPSessionID, sessionID) +} + +// WithClientInfo sets the _meta clientInfo object. +func (r *RawRequest) WithClientInfo(name, version string) *RawRequest { + return r.SetMeta(MetaKeyClientInfo, map[string]any{"name": name, "version": version}) +} + +// WithStreamableAccept sets "Accept: application/json, text/event-stream". +// A real go-sdk streamable-HTTP server rejects a POST without this header +// (HTTP 400), so requests bound for a real backend (e.g. the k8s tier) must +// set it. Requests to the ToolHive proxy still omit it by convention: the +// proxy does not require it, and omitting it keeps responses plain JSON (the +// client parses SSE responses too -- see sseResponsePayload -- and flipping +// proxy-bound requests to the conformant Accept is tracked in #6104). +func (r *RawRequest) WithStreamableAccept() *RawRequest { + return r.SetHeader("Accept", "application/json, text/event-stream") +} + +// WithRawBody replaces the entire request body with body, sent verbatim. +// Use this for truncated, oversized, or garbage bodies. body is cloned; +// mutating the caller's slice afterward has no effect on the request. +func (r *RawRequest) WithRawBody(body []byte) *RawRequest { + r.body = bytes.Clone(body) + return r +} + +// SetHeader sets an arbitrary HTTP header on the request. +func (r *RawRequest) SetHeader(key, value string) *RawRequest { + r.headers[key] = value + return r +} + +// DeleteHeader removes a header, e.g. to omit MCP-Protocol-Version entirely. +func (r *RawRequest) DeleteHeader(key string) *RawRequest { + delete(r.headers, key) + return r +} + +// SetMeta sets a params._meta key to an arbitrary value, including a +// non-object value for a key the spec requires to be an object (to trigger +// error paths). It lazily creates the _meta object if none exists yet. +func (r *RawRequest) SetMeta(key string, value any) *RawRequest { + if r.meta == nil { + r.meta = map[string]any{} + } + r.meta[key] = value + return r +} + +// DeleteMeta removes a params._meta key, e.g. to make clientInfo absent. +func (r *RawRequest) DeleteMeta(key string) *RawRequest { + delete(r.meta, key) + return r +} + +// marshal renders the JSON-RPC wire body for r, or returns r.body verbatim +// if WithRawBody was used. +func (r *RawRequest) marshal() ([]byte, error) { + if r.body != nil { + return r.body, nil + } + + wire := map[string]any{ + "jsonrpc": "2.0", + "method": r.method, + } + if r.id != nil { + wire["id"] = r.id + } + + params := maps.Clone(r.params) + if r.meta != nil { + if params == nil { + params = map[string]any{} + } + params["_meta"] = r.meta + } + if params != nil { + wire["params"] = params + } + + return json.Marshal(wire) +} + +// NewBatchBody renders a JSON-RPC batch: a JSON array of the wire bodies of +// reqs, in order. Combine with RawMCPClient.SendRaw to send it. Each element +// is validated the same as a standalone RawRequest; a deliberately malformed +// batch payload (bad JSON, wrong array shape) must be hand-built and sent +// via SendRaw directly. +func NewBatchBody(reqs ...*RawRequest) ([]byte, error) { + parts := make([]json.RawMessage, 0, len(reqs)) + for i, r := range reqs { + b, err := r.marshal() + if err != nil { + return nil, fmt.Errorf("mcp_raw_client: marshal batch element %d: %w", i, err) + } + parts = append(parts, b) + } + return json.Marshal(parts) +} + +// RawRPCError is the JSON-RPC "error" object of a RawResponse. +type RawRPCError struct { + Code int64 + Message string + Data json.RawMessage +} + +// RawResponse is the parsed result of sending a single (non-batch) JSON-RPC +// request over MCP-over-HTTP. For batch responses (a JSON array), inspect +// Body directly instead of JSONRPC/ID/Result/Error. +type RawResponse struct { + StatusCode int + Headers http.Header + Body []byte // raw response body, always populated + JSONRPC string // echoed "jsonrpc" version tag; "" if absent or unparsable + ID any // echoed "id"; nil if absent, JSON null, or unparsable + Result json.RawMessage + Error *RawRPCError +} + +// RawMCPClient sends raw MCP-over-HTTP JSON-RPC requests with byte-level +// control over the wire shape. It wraps a single *http.Client (and its +// connection pool), shared and safe for concurrent use across goroutines — +// construct one RawMCPClient per test suite, not one per request. +type RawMCPClient struct { + httpClient *http.Client +} + +// NewRawMCPClient creates a RawMCPClient. Every request made through it is +// bounded by timeout, so a stuck proxy fails fast instead of hanging the +// test. +func NewRawMCPClient(timeout time.Duration) (*RawMCPClient, error) { + if timeout <= 0 { + return nil, fmt.Errorf("mcp_raw_client: timeout must be positive, got %s", timeout) + } + return &RawMCPClient{ + httpClient: &http.Client{ + Timeout: timeout, + // Raised from the net/http default of 2 idle conns/host: adversarial/ + // concurrency tests fan many goroutines out to the same proxy host, + // and the default would force constant reconnection. + Transport: &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 100, + IdleConnTimeout: 90 * time.Second, + }, + }, + }, nil +} + +// Send marshals req and POSTs it to url. No Accept header is set by default: +// the ToolHive proxy returns a plain JSON body without it, and switches to an +// SSE response body whenever Accept lists text/event-stream. A request bound +// for a REAL streamable-HTTP MCP backend -- which rejects a POST lacking that +// Accept with HTTP 400 -- should opt in via req.WithStreamableAccept(). +// +// Either way the JSON-RPC envelope is parsed into the RawResponse: a +// Content-Type: text/event-stream body has its response event extracted first +// (see sseResponsePayload). Some servers answer a POST as SSE regardless of +// the request's Accept header -- mcpcompat's cross-replica rehydration path +// does, because go-sdk v1.7.0-pre.3 does not export the transport-level +// JSONResponse knob (StreamableServerTransport.jsonResponse) that its +// top-level handler options set -- so this parsing cannot be opt-in. +func (c *RawMCPClient) Send(ctx context.Context, url string, req *RawRequest) (*RawResponse, error) { + body, err := req.marshal() + if err != nil { + return nil, fmt.Errorf("mcp_raw_client: marshal request: %w", err) + } + return c.SendRaw(ctx, url, req.headers, body) +} + +// SendRaw POSTs body verbatim to url with the given headers — for batches +// built with NewBatchBody, or any payload (truncated, oversized, garbage) +// not expressible via RawRequest. +func (c *RawMCPClient) SendRaw(ctx context.Context, url string, headers map[string]string, body []byte) (*RawResponse, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("mcp_raw_client: build request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + for k, v := range headers { + httpReq.Header.Set(k, v) + } + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("mcp_raw_client: do request: %w", err) + } + defer func() { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("mcp_raw_client: read response body: %w", err) + } + + result := &RawResponse{ + StatusCode: resp.StatusCode, + Headers: resp.Header.Clone(), + Body: respBody, + } + envelope := respBody + if strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") { + envelope = sseResponsePayload(respBody) + } + populateEnvelope(envelope, result) + return result, nil +} + +// wireError is the shape of a JSON-RPC "error" object on the wire. +type wireError struct { + Code int64 `json:"code"` + Message string `json:"message"` + Data json.RawMessage `json:"data"` +} + +// wireResponse is the shape of a single JSON-RPC response object. ID is kept +// as raw JSON so decodeID can preserve large integers exactly. +type wireResponse struct { + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *wireError `json:"error"` +} + +// populateEnvelope best-effort parses body as a single JSON-RPC response +// object into out. A malformed or non-object body (e.g. a batch array, or +// deliberately garbled test input) leaves JSONRPC/ID/Result/Error at their +// zero values; out.Body still carries the raw bytes for the caller to inspect. +func populateEnvelope(body []byte, out *RawResponse) { + var w wireResponse + if err := json.Unmarshal(body, &w); err != nil { + return + } + // encoding/json struct tags match case-insensitively as a fallback, which + // would let a miscased "JSONRPC" or "JsonRpc" key satisfy `json:"jsonrpc"` + // -- exactly the kind of malformed envelope this field exists to catch + // (see #5950). Read it from a raw map instead, which is case-sensitive. + var raw map[string]json.RawMessage + if err := json.Unmarshal(body, &raw); err == nil { + _ = json.Unmarshal(raw["jsonrpc"], &out.JSONRPC) + } + out.ID = decodeID(w.ID) + out.Result = w.Result + if w.Error != nil { + out.Error = &RawRPCError{Code: w.Error.Code, Message: w.Error.Message, Data: w.Error.Data} + } +} + +// sseResponsePayload extracts the JSON-RPC response object from an SSE-framed +// (text/event-stream) POST response body. A streamable-HTTP server may +// interleave notifications on the same stream before the response, so the +// payload returned is the LAST event whose data is a JSON object carrying a +// top-level "result" or "error" key -- the one JSON-RPC response terminating +// the stream -- not simply the last event. Multi-line data fields are joined +// with newlines per the SSE spec. Returns nil when no event carries a +// response, leaving the envelope unpopulated (Body still holds the raw +// stream for the caller to inspect). +func sseResponsePayload(body []byte) []byte { + var response []byte + var data [][]byte + flush := func() { + if len(data) == 0 { + return + } + payload := bytes.Join(data, []byte("\n")) + data = nil + var obj map[string]json.RawMessage + if err := json.Unmarshal(payload, &obj); err != nil { + return + } + if _, ok := obj["result"]; ok { + response = payload + return + } + if _, ok := obj["error"]; ok { + response = payload + } + } + for _, line := range bytes.Split(body, []byte("\n")) { + line = bytes.TrimSuffix(line, []byte("\r")) + if len(line) == 0 { // blank line: event boundary + flush() + continue + } + if value, ok := bytes.CutPrefix(line, []byte("data:")); ok { + data = append(data, bytes.TrimPrefix(value, []byte(" "))) + } + // Other fields (event:, id:, retry:, comments) don't carry payload. + } + flush() // stream may end without a trailing blank line + return response +} + +// decodeID decodes a raw JSON-RPC id, returning a json.Number for numeric +// ids (so large int64 values round-trip exactly, unlike the float64 that +// encoding/json's default decodes numbers into) or a string for string ids. +// An absent or null id decodes to nil. +func decodeID(raw json.RawMessage) any { + if len(raw) == 0 || string(raw) == "null" { + return nil + } + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.UseNumber() + var v any + if err := dec.Decode(&v); err != nil { + return nil + } + return v +} diff --git a/test/e2e/mcp_raw_client_golden_test.go b/test/e2e/mcp_raw_client_golden_test.go new file mode 100644 index 0000000000..c990eec97c --- /dev/null +++ b/test/e2e/mcp_raw_client_golden_test.go @@ -0,0 +1,121 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e + +import ( + "bytes" + "encoding/json" + "fmt" + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +// goldenModernRequestFixture is a Modern (2026-07-28) "tools/list" request +// emitted by a real github.com/modelcontextprotocol/go-sdk v1.7.0-pre.3 client. +// It exists so NewModernRequest's _meta shape is checked against an independent +// oracle, not just against assertions written by the same author who wrote the +// encoder. To regenerate: point any go-sdk v1.7 client at +// `yardstick-server --stateless` (>= v1.2.0, itself on go-sdk v1.7.0-pre.3) and +// capture the raw tools/list request body it sends. +const goldenModernRequestFixture = "testdata/golden_modern_request.json" + +// extractMeta pulls params._meta out of a decoded wire object as a +// map[string]any, failing the test with a clear message if the shape isn't +// what's expected. +func extractMeta(t *testing.T, wire map[string]any) map[string]any { + t.Helper() + params, ok := wire["params"].(map[string]any) + require.True(t, ok, "wire request must have an object params field: %#v", wire["params"]) + meta, ok := params["_meta"].(map[string]any) + require.True(t, ok, "params must have an object _meta field: %#v", params["_meta"]) + return meta +} + +// jsonKind reports the JSON type of a value decoded by encoding/json into +// an any (object/array/string/number/bool/null), discarding its content. +func jsonKind(v any) string { + switch v.(type) { + case map[string]any: + return "object" + case []any: + return "array" + case string: + return "string" + case float64: + return "number" + case bool: + return "bool" + case nil: + return "null" + default: + return fmt.Sprintf("%T", v) // unreachable for encoding/json output + } +} + +// metaShape reduces a decoded _meta map to its keys and each value's JSON +// kind, discarding nested content. NewModernRequest deliberately sends a +// minimal empty clientCapabilities object rather than mirroring the go-sdk +// client's specific declared capabilities (roots, sampling, ...), so nested +// capability fields are never expected to match value-for-value — only that +// the same top-level _meta keys exist and each is the same kind of JSON +// value ("shape/keys", per the acceptance criteria). +func metaShape(meta map[string]any) map[string]string { + shape := make(map[string]string, len(meta)) + for k, v := range meta { + shape[k] = jsonKind(v) + } + return shape +} + +// TestModernRequestMatchesGoldenSDKFixture checks that NewModernRequest's +// wire output byte-matches the golden go-sdk fixture's _meta shape and keys. +// +// A literal full-envelope byte comparison would not hold even for two +// semantically identical requests: RawRequest.marshal builds the body from a +// map[string]any (encoding/json sorts map keys alphabetically), while the +// go-sdk marshals a Go struct in its own field order. So this test scopes +// the byte-level comparison to what the two sides can actually agree on: the +// _meta object's keys and each value's kind, independently re-marshaled on +// each side. Since encoding/json deterministically sorts map keys, comparing +// the re-marshaled bytes of two map[string]string shapes is still a genuine +// byte-for-byte check — just correctly scoped to the part of the AC that is +// comparable byte-for-byte ("_meta shape/keys", not deep capability content +// that NewModernRequest never claims to reproduce). +func TestModernRequestMatchesGoldenSDKFixture(t *testing.T) { + t.Parallel() + + goldenBytes, err := os.ReadFile(goldenModernRequestFixture) + require.NoError(t, err, "golden fixture must exist; see the goldenModernRequestFixture doc comment to regenerate it") + var golden map[string]any + require.NoError(t, json.Unmarshal(goldenBytes, &golden)) + require.Equal(t, "tools/list", golden["method"], "golden fixture must have been captured for tools/list to match NewModernRequest below") + goldenMeta := extractMeta(t, golden) + + req, err := NewModernRequest("tools/list", nil) + require.NoError(t, err) + rawMeta := extractMeta(t, wireOf(t, req)) + + require.Equal(t, goldenMeta[MetaKeyProtocolVersion], rawMeta[MetaKeyProtocolVersion], + "protocolVersion must match exactly, not just in shape") + + // clientInfo is present in the SDK-captured fixture (the SDK client has an + // Implementation configured) but omitted by NewModernRequest per spec + // ("SHOULD", not "MUST" — see NewModernRequest's doc comment); exclude it + // from the shape comparison below rather than requiring both sides to + // agree on an optional key. + require.Contains(t, goldenMeta, MetaKeyClientInfo, "golden fixture is expected to carry optional clientInfo") + require.NotContains(t, rawMeta, MetaKeyClientInfo, "NewModernRequest omits clientInfo by default") + delete(goldenMeta, MetaKeyClientInfo) + + goldenShapeJSON, err := json.Marshal(metaShape(goldenMeta)) + require.NoError(t, err) + rawShapeJSON, err := json.Marshal(metaShape(rawMeta)) + require.NoError(t, err) + + require.True(t, bytes.Equal(goldenShapeJSON, rawShapeJSON), + "raw client _meta must byte-match the golden go-sdk fixture's _meta shape/keys (excluding optional clientInfo)\ngolden shape: %s\nraw shape: %s", + goldenShapeJSON, rawShapeJSON) +} diff --git a/test/e2e/mcp_raw_client_test.go b/test/e2e/mcp_raw_client_test.go new file mode 100644 index 0000000000..77d01123a3 --- /dev/null +++ b/test/e2e/mcp_raw_client_test.go @@ -0,0 +1,498 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e + +import ( + "context" + "encoding/json" + "io" + "math" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// wireOf marshals req and decodes it into a generic map for field assertions. +func wireOf(t *testing.T, req *RawRequest) map[string]any { + t.Helper() + b, err := req.marshal() + require.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal(b, &m)) + return m +} + +func TestRawClientBuilders(t *testing.T) { + t.Parallel() + + t.Run("ModernRequest sets header and default meta keys, clientInfo omitted", func(t *testing.T) { + t.Parallel() + req, err := NewModernRequest("tools/list", nil) + require.NoError(t, err) + + require.Equal(t, MCPVersionModern, req.headers[HeaderMCPProtocolVersion]) + require.Equal(t, "tools/list", req.headers[HeaderMCPMethod], "Mcp-Method is required on every Modern request") + require.NotContains(t, req.headers, HeaderMCPName, "tools/list is not name-bearing, so no Mcp-Name") + + wire := wireOf(t, req) + meta := wire["params"].(map[string]any)["_meta"].(map[string]any) + require.Equal(t, MCPVersionModern, meta[MetaKeyProtocolVersion]) + require.Equal(t, map[string]any{}, meta[MetaKeyClientCapabilities]) + require.NotContains(t, meta, MetaKeyClientInfo) + }) + + t.Run("ModernRequest sets Mcp-Name to the plain target name for name-bearing methods", func(t *testing.T) { + t.Parallel() + req, err := NewModernRequest("tools/call", map[string]any{"name": "echo"}) + require.NoError(t, err) + require.Equal(t, "tools/call", req.headers[HeaderMCPMethod]) + require.Equal(t, "echo", req.headers[HeaderMCPName], "tools/call must name its tool in Mcp-Name, matching body params.name") + }) + + t.Run("ModernRequest fields are independently overridable", func(t *testing.T) { + t.Parallel() + req, err := NewModernRequest("tools/list", nil) + require.NoError(t, err) + + req.SetMeta(MetaKeyProtocolVersion, "bogus-version"). + SetMeta(MetaKeyClientCapabilities, "not-an-object"). // non-object value, to trigger error paths + WithClientInfo("probe", "1.0"). + DeleteHeader(HeaderMCPProtocolVersion) + + _, hasHeader := req.headers[HeaderMCPProtocolVersion] + require.False(t, hasHeader, "header should be omittable") + + wire := wireOf(t, req) + meta := wire["params"].(map[string]any)["_meta"].(map[string]any) + require.Equal(t, "bogus-version", meta[MetaKeyProtocolVersion]) + require.Equal(t, "not-an-object", meta[MetaKeyClientCapabilities]) + require.Equal(t, map[string]any{"name": "probe", "version": "1.0"}, meta[MetaKeyClientInfo]) + + req.DeleteMeta(MetaKeyClientInfo) + wire = wireOf(t, req) + meta = wire["params"].(map[string]any)["_meta"].(map[string]any) + require.NotContains(t, meta, MetaKeyClientInfo) + }) + + t.Run("LegacyInitializeRequest carries no _meta and the legacy protocol version", func(t *testing.T) { + t.Parallel() + req := NewLegacyInitializeRequest("probe", "1.0") + wire := wireOf(t, req) + + require.Equal(t, "initialize", wire["method"]) + params := wire["params"].(map[string]any) + require.Equal(t, "2025-11-25", params["protocolVersion"]) + require.NotContains(t, params, "_meta") + }) + + t.Run("LegacyRequest WithSessionID sets the session header, including a spoofed one", func(t *testing.T) { + t.Parallel() + req, err := NewLegacyRequest("tools/call", map[string]any{"name": "echo"}) + require.NoError(t, err) + req.WithSessionID("not-my-session") + + require.Equal(t, "not-my-session", req.headers[HeaderMCPSessionID]) + }) + + t.Run("NewLegacyRequest rejects an empty method", func(t *testing.T) { + t.Parallel() + _, err := NewLegacyRequest("", nil) + require.Error(t, err) + }) + + t.Run("WithID preserves a large int64 and allows colliding ids", func(t *testing.T) { + t.Parallel() + const big = int64(9223372036854775807) // math.MaxInt64; beyond float64's exact integer range + + reqA, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + reqA.WithID(big) + reqB, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + reqB.WithID(big) // deliberate id collision with reqA + + for _, req := range []*RawRequest{reqA, reqB} { + b, err := req.marshal() + require.NoError(t, err) + require.Contains(t, string(b), `"id":9223372036854775807`) + } + }) + + t.Run("WithRawBody sends the body verbatim, ignoring method and params", func(t *testing.T) { + t.Parallel() + req, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + garbage := []byte(`{not even json`) + req.WithRawBody(garbage) + + b, err := req.marshal() + require.NoError(t, err) + require.Equal(t, garbage, b) + + // Mutating the caller's slice afterward must not affect the request. + garbage[0] = 'X' + b, err = req.marshal() + require.NoError(t, err) + require.Equal(t, []byte(`{not even json`), b) + }) + + t.Run("NewBatchBody produces a JSON array of the wire bodies in order", func(t *testing.T) { + t.Parallel() + reqA, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + reqA.WithID(int64(1)) + reqB, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + reqB.WithID(int64(2)) + + batch, err := NewBatchBody(reqA, reqB) + require.NoError(t, err) + + var arr []map[string]any + require.NoError(t, json.Unmarshal(batch, &arr)) + require.Len(t, arr, 2) + require.Equal(t, "ping", arr[0]["method"]) + require.InDelta(t, float64(1), arr[0]["id"], 0) + require.Equal(t, "tools/list", arr[1]["method"]) + require.InDelta(t, float64(2), arr[1]["id"], 0) + }) + + t.Run("params map passed by the caller is not mutated", func(t *testing.T) { + t.Parallel() + callerParams := map[string]any{"name": "echo"} + req, err := NewModernRequest("tools/call", callerParams) + require.NoError(t, err) + req.SetMeta("extra", "value") + _, err = req.marshal() + require.NoError(t, err) + + require.Equal(t, map[string]any{"name": "echo"}, callerParams, "caller's map must be untouched") + }) +} + +// capturedRequest snapshots what the sender actually transmitted. +type capturedRequest struct { + headers http.Header + body []byte +} + +// newCapturingServer starts an httptest.Server that records the last request +// it received (via the returned *capturedRequest pointer, refreshed on every +// call) and replies with responseBody. +func newCapturingServer(t *testing.T, responseBody string) (*httptest.Server, *capturedRequest) { + t.Helper() + captured := &capturedRequest{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + captured.headers = r.Header.Clone() + captured.body = body + + w.Header().Set("Content-Type", "application/json") + w.Header().Set(HeaderMCPSessionID, "server-assigned-session") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(responseBody)) + })) + t.Cleanup(server.Close) + return server, captured +} + +func TestRawClientSend(t *testing.T) { + t.Parallel() + + client, err := NewRawMCPClient(5 * time.Second) + require.NoError(t, err) + + t.Run("transmits headers and body, and parses result/id/response headers", func(t *testing.T) { + t.Parallel() + server, captured := newCapturingServer(t, `{"jsonrpc":"2.0","id":42,"result":{"ok":true}}`) + + req := NewLegacyInitializeRequest("probe", "1.0").WithSessionID("client-picked-session") + req.WithID(int64(42)) + + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Equal(t, "client-picked-session", captured.headers.Get(HeaderMCPSessionID), + "the client-set session header must reach the server verbatim") + + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "server-assigned-session", resp.Headers.Get(HeaderMCPSessionID), + "response headers must be surfaced so tests can assert session-id echoing") + require.Equal(t, json.Number("42"), resp.ID) + require.JSONEq(t, `{"ok":true}`, string(resp.Result)) + require.Nil(t, resp.Error) + require.Equal(t, "2.0", resp.JSONRPC) + }) + + t.Run("pre-fix denial shape (capitalized keys, no jsonrpc tag) leaves JSONRPC empty", func(t *testing.T) { + t.Parallel() + // This is the exact shape a jsonrpc2.Response marshaled with + // encoding/json's default reflection produces (no json tags, unexported + // ID field): pins that the harness reads it as having no version tag, + // rather than defaulting or being fooled into "2.0" by the stdlib's + // case-insensitive key matching. The case above is what actually fails + // if JSONRPC stops being populated. + server, _ := newCapturingServer(t, `{"Result":null,"Error":{"code":403,"message":"x"},"ID":{}}`) + + req, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Empty(t, resp.JSONRPC) + }) + + t.Run("a miscased jsonrpc key does not satisfy the version-tag check", func(t *testing.T) { + t.Parallel() + // encoding/json struct tags match case-insensitively as a fallback, so + // a naive `json:"jsonrpc"` unmarshal would let this non-conformant key + // through. resp.JSONRPC must only ever reflect the exact lowercase key. + server, _ := newCapturingServer(t, `{"JSONRPC":"2.0","id":1,"result":{}}`) + + req, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Empty(t, resp.JSONRPC) + }) + + t.Run("parses error.code, error.data and a large-int64 id without precision loss", func(t *testing.T) { + t.Parallel() + const bigID = "9223372036854775807" + server, _ := newCapturingServer(t, + `{"jsonrpc":"2.0","id":`+bigID+`,"error":{"code":-32020,"message":"mismatch","data":{"header":"2025-11-25"}}}`) + + req, err := NewModernRequest("tools/list", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Equal(t, json.Number(bigID), resp.ID) + id, err := resp.ID.(json.Number).Int64() + require.NoError(t, err) + require.Equal(t, int64(math.MaxInt64), id) + + require.NotNil(t, resp.Error) + require.Equal(t, int64(-32020), resp.Error.Code) + require.Equal(t, "mismatch", resp.Error.Message) + require.JSONEq(t, `{"header":"2025-11-25"}`, string(resp.Error.Data)) + }) + + t.Run("SendRaw transmits a batch body and leaves envelope fields zero on a non-object response", func(t *testing.T) { + t.Parallel() + server, captured := newCapturingServer(t, `[{"jsonrpc":"2.0","id":1,"result":{}},{"jsonrpc":"2.0","id":2,"result":{}}]`) + + reqA, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + reqA.WithID(int64(1)) + reqB, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + reqB.WithID(int64(2)) + batch, err := NewBatchBody(reqA, reqB) + require.NoError(t, err) + + resp, err := client.SendRaw(context.Background(), server.URL, nil, batch) + require.NoError(t, err) + + var sentArr []map[string]any + require.NoError(t, json.Unmarshal(captured.body, &sentArr)) + require.Len(t, sentArr, 2, "the server must receive a two-element JSON array") + + require.Nil(t, resp.ID, "a batch (array) response is not a single envelope") + require.Nil(t, resp.Error) + require.JSONEq(t, `[{"jsonrpc":"2.0","id":1,"result":{}},{"jsonrpc":"2.0","id":2,"result":{}}]`, string(resp.Body)) + }) + + t.Run("garbage body is sent verbatim and a non-JSON response leaves the envelope zero but Body populated", func(t *testing.T) { + t.Parallel() + server, captured := newCapturingServer(t, `not json at all`) + + resp, err := client.SendRaw(context.Background(), server.URL, map[string]string{"X-Test": "1"}, []byte(`{truncated`)) + require.NoError(t, err) + + require.Equal(t, []byte(`{truncated`), captured.body) + require.Equal(t, "1", captured.headers.Get("X-Test")) + + require.Nil(t, resp.ID) + require.Nil(t, resp.Error) + require.Equal(t, []byte(`not json at all`), resp.Body) + }) + + t.Run("no Accept header by default, opt-in via WithStreamableAccept", func(t *testing.T) { + t.Parallel() + server, captured := newCapturingServer(t, `{"jsonrpc":"2.0","id":1,"result":{}}`) + + req, err := NewModernRequest("tools/list", nil) + require.NoError(t, err) + _, err = client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + require.Empty(t, captured.headers.Get("Accept"), + "no Accept by default: keeps proxy responses plain JSON; conformant-Accept flip tracked in #6104") + + withAccept, err := NewModernRequest("tools/list", nil) + require.NoError(t, err) + withAccept.WithStreamableAccept() + _, err = client.Send(context.Background(), server.URL, withAccept) + require.NoError(t, err) + require.Equal(t, "application/json, text/event-stream", captured.headers.Get("Accept"), + "WithStreamableAccept opts into the Accept a real streamable-HTTP backend requires") + }) +} + +// newSSEServer serves body verbatim as a text/event-stream POST response — +// the shape mcpcompat's cross-replica rehydration path produces regardless of +// the request's Accept header (go-sdk v1.7.0-pre.3 does not export the +// transport-level JSONResponse knob its top-level handler options set). +func newSSEServer(t *testing.T, body string) *httptest.Server { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body)) + })) + t.Cleanup(server.Close) + return server +} + +// TestRawClientSSEResponse pins that an SSE-framed POST response body has its +// JSON-RPC response event parsed into the envelope, so era-pinned raw-client +// specs work against a server that answers SSE unasked (the rehydrated +// cross-pod session path) exactly as they do against a plain-JSON one. +func TestRawClientSSEResponse(t *testing.T) { + t.Parallel() + + client, err := NewRawMCPClient(5 * time.Second) + require.NoError(t, err) + + t.Run("parses the response event, skipping preceding notifications", func(t *testing.T) { + t.Parallel() + server := newSSEServer(t, "event: message\n"+ + `data: {"jsonrpc":"2.0","method":"notifications/message","params":{}}`+"\n\n"+ + "event: message\n"+ + `data: {"jsonrpc":"2.0","id":7,"result":{"tools":[{"name":"echo"}]}}`+"\n\n") + + req, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "2.0", resp.JSONRPC) + require.Equal(t, json.Number("7"), resp.ID) + require.JSONEq(t, `{"tools":[{"name":"echo"}]}`, string(resp.Result)) + require.Nil(t, resp.Error) + }) + + t.Run("parses an error response event, including multi-line data", func(t *testing.T) { + t.Parallel() + // The SSE spec joins consecutive data: lines with \n; a JSON payload + // split across them must still parse. + server := newSSEServer(t, "event: message\n"+ + `data: {"jsonrpc":"2.0","id":3,`+"\n"+ + `data: "error":{"code":429,"message":"Rate limit exceeded"}}`+"\n\n") + + req, err := NewLegacyRequest("tools/call", map[string]any{"name": "echo"}) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.NotNil(t, resp.Error) + require.EqualValues(t, 429, resp.Error.Code) + require.Equal(t, "Rate limit exceeded", resp.Error.Message) + }) + + t.Run("a stream with no response event leaves the envelope empty, Body raw", func(t *testing.T) { + t.Parallel() + body := "event: message\n" + + `data: {"jsonrpc":"2.0","method":"notifications/message","params":{}}` + "\n\n" + server := newSSEServer(t, body) + + req, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Empty(t, resp.JSONRPC) + require.Nil(t, resp.Result) + require.Nil(t, resp.Error) + require.Equal(t, body, string(resp.Body), "Body must keep the raw stream for the caller") + }) + + t.Run("a stream without a trailing blank line still yields the response", func(t *testing.T) { + t.Parallel() + server := newSSEServer(t, "event: message\n"+ + `data: {"jsonrpc":"2.0","id":1,"result":{}}`) + + req, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Equal(t, json.Number("1"), resp.ID) + require.JSONEq(t, `{}`, string(resp.Result)) + }) + + t.Run("CRLF line endings parse identically to LF", func(t *testing.T) { + t.Parallel() + server := newSSEServer(t, "event: message\r\n"+ + `data: {"jsonrpc":"2.0","id":5,"result":{"tools":[{"name":"echo"}]}}`+"\r\n\r\n") + + req, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Equal(t, "2.0", resp.JSONRPC) + require.Equal(t, json.Number("5"), resp.ID) + require.JSONEq(t, `{"tools":[{"name":"echo"}]}`, string(resp.Result)) + require.Nil(t, resp.Error) + }) + + t.Run("data: with no leading space still parses", func(t *testing.T) { + t.Parallel() + server := newSSEServer(t, "event: message\n"+ + `data:{"jsonrpc":"2.0","id":9,"result":{}}`+"\n\n") + + req, err := NewLegacyRequest("ping", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Equal(t, json.Number("9"), resp.ID) + require.JSONEq(t, `{}`, string(resp.Result)) + }) + + t.Run("a batch (array) data payload leaves the envelope empty, Body raw", func(t *testing.T) { + t.Parallel() + // A JSON array is not a single response object, so it never populates + // the envelope — matching RawResponse's batch contract (inspect Body). + body := "event: message\n" + + `data: [{"jsonrpc":"2.0","id":1,"result":{}},{"jsonrpc":"2.0","id":2,"result":{}}]` + "\n\n" + server := newSSEServer(t, body) + + req, err := NewLegacyRequest("tools/list", nil) + require.NoError(t, err) + resp, err := client.Send(context.Background(), server.URL, req) + require.NoError(t, err) + + require.Empty(t, resp.JSONRPC) + require.Nil(t, resp.ID) + require.Nil(t, resp.Result) + require.Nil(t, resp.Error) + require.Equal(t, body, string(resp.Body), "Body must keep the raw stream for the caller") + }) +} + +func TestNewRawMCPClientRejectsNonPositiveTimeout(t *testing.T) { + t.Parallel() + _, err := NewRawMCPClient(0) + require.Error(t, err) + _, err = NewRawMCPClient(-time.Second) + require.Error(t, err) +} diff --git a/test/e2e/network_isolation_envoy_test.go b/test/e2e/network_isolation_envoy_test.go new file mode 100644 index 0000000000..2121934fda --- /dev/null +++ b/test/e2e/network_isolation_envoy_test.go @@ -0,0 +1,307 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +// NetworkIsolationEnvoy contains e2e tests for the Envoy network proxy backend +// (TOOLHIVE_NETWORK_PROXY=envoy). The tests cover the same behavioural contract +// as the Squid suite in network_isolation_test.go — egress allowlist, inbound +// Host gating, docker-gateway deny/allow — plus Envoy-specific guards such as +// the route-timeout regression test. + +import ( + "context" + "fmt" + "io" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/test/e2e" +) + +var _ = Describe("NetworkIsolationEnvoy", Label("proxy", "network", "isolation", "envoy", "e2e"), func() { + var ( + config *e2e.TestConfig + permissionProfileDir string + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + + var err error + permissionProfileDir, err = os.MkdirTemp("", "network-isolation-envoy-profiles-*") + Expect(err).ToNot(HaveOccurred()) + + err = e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + }) + + AfterEach(func() { + if permissionProfileDir != "" { + os.RemoveAll(permissionProfileDir) + } + }) + + // envoyRun returns a thv run command with TOOLHIVE_NETWORK_PROXY=envoy injected. + envoyRun := func(config *e2e.TestConfig, args ...string) *e2e.THVCommand { + return e2e.NewTHVCommand(config, args...). + WithEnv("TOOLHIVE_NETWORK_PROXY=envoy") + } + + // startFetchServer starts the fetch MCP server under Envoy network isolation + // with the given extra run args and returns the workload name. + startFetchServer := func(nameSuffix, profileJSON string, extraRunArgs ...string) string { + serverName := fmt.Sprintf("nie-%s-%d", nameSuffix, GinkgoRandomSeed()) + DeferCleanup(func() { + if config.CleanupAfter { + _ = e2e.StopAndRemoveMCPServer(config, serverName) + } + }) + + runArgs := append([]string{"run", "--name", serverName}, extraRunArgs...) + if profileJSON != "" { + profilePath := filepath.Join(permissionProfileDir, nameSuffix+".json") + Expect(os.WriteFile(profilePath, []byte(profileJSON), 0644)).To(Succeed()) + runArgs = append(runArgs, "--permission-profile", profilePath) + } + runArgs = append(runArgs, "fetch") + + envoyRun(config, runArgs...).ExpectSuccess() + Expect(e2e.WaitForMCPServer(config, serverName, 180*time.Second)). + To(Succeed(), "Envoy server should be running within 180 seconds") + return serverName + } + + // fetchThrough drives the fetch tool against the given server and returns + // the tool result. + fetchThrough := func(serverName, targetURL string, callTimeout time.Duration) *mcp.CallToolResult { + serverURL, err := e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred()) + Expect(e2e.WaitForMCPServerReady(config, serverURL, "streamable-http", 60*time.Second)). + To(Succeed()) + + mcpClient, err := e2e.NewMCPClientForStreamableHTTP(config, serverURL) + Expect(err).ToNot(HaveOccurred()) + defer mcpClient.Close() + + ctx, cancel := context.WithTimeout(context.Background(), callTimeout) + defer cancel() + Expect(mcpClient.Initialize(ctx)).To(Succeed()) + + result, err := mcpClient.CallTool(ctx, "fetch", map[string]interface{}{"url": targetURL}) + Expect(err).ToNot(HaveOccurred(), "CallTool should complete without transport error") + return result + } + + // resultText concatenates the text content of a tool result. + resultText := func(result *mcp.CallToolResult) string { + var sb strings.Builder + for _, c := range result.Content { + if tc, ok := mcp.AsTextContent(c); ok { + sb.WriteString(tc.Text) + } + } + return sb.String() + } + + // dockerBridgeGatewayIP returns the bridge gateway IP from the default Docker + // bridge network. + dockerBridgeGatewayIP := func() string { + //nolint:gosec // fixed, test-controlled arguments + out, err := exec.Command("docker", "network", "inspect", "bridge", + "-f", "{{range .IPAM.Config}}{{.Gateway}}{{end}}").Output() + Expect(err).ToNot(HaveOccurred()) + return strings.TrimSpace(string(out)) + } + + // ── Behavioural parity with Squid ──────────────────────────────────────── + + Describe("Outbound and inbound restrictions", func() { + verifyNetworkRestrictions := func(nameSuffix string, extraRunArgs ...string) { + serverName := fmt.Sprintf("nie-%s-%d", nameSuffix, GinkgoRandomSeed()) + DeferCleanup(func() { + if config.CleanupAfter { + Expect(e2e.StopAndRemoveMCPServer(config, serverName)).To(Succeed()) + } + }) + + profile := `{ + "name": "test-network-isolation", + "network": { + "inbound": { "allow_host": ["localhost", "127.0.0.1"] }, + "outbound": { "insecure_allow_all": false, "allow_host": ["example.com"], "allow_port": [80, 443] } + } + }` + profilePath := filepath.Join(permissionProfileDir, nameSuffix+".json") + Expect(os.WriteFile(profilePath, []byte(profile), 0644)).To(Succeed()) + + runArgs := append([]string{"run", "--name", serverName}, extraRunArgs...) + runArgs = append(runArgs, "--permission-profile", profilePath, "fetch") + envoyRun(config, runArgs...).ExpectSuccess() + + Expect(e2e.WaitForMCPServer(config, serverName, 180*time.Second)).To(Succeed()) + serverURL, err := e2e.GetMCPServerURL(config, serverName) + Expect(err).ToNot(HaveOccurred()) + Expect(e2e.WaitForMCPServerReady(config, serverURL, "streamable-http", 60*time.Second)).To(Succeed()) + + mcpClient, err := e2e.NewMCPClientForStreamableHTTP(config, serverURL) + Expect(err).ToNot(HaveOccurred()) + defer mcpClient.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + Expect(mcpClient.Initialize(ctx)).To(Succeed()) + + By("Blocked URL must fail outbound") + result, err := mcpClient.CallTool(ctx, "fetch", map[string]interface{}{"url": "https://google.com"}) + Expect(err).ToNot(HaveOccurred()) + Expect(result.IsError).To(BeTrue(), "blocked URL must return an error result") + + By("Non-allowed Host header must be rejected inbound") + req, err := http.NewRequest("GET", serverURL, nil) + Expect(err).ToNot(HaveOccurred()) + req.Host = "example.org" + resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req) + Expect(err).ToNot(HaveOccurred()) + defer resp.Body.Close() + Expect(resp.StatusCode).ToNot(Equal(http.StatusOK)) + } + + It("should enforce restrictions when --isolate-network is set", func() { + verifyNetworkRestrictions("explicit", "--isolate-network") + }) + + It("should enforce restrictions with default isolation", func() { + verifyNetworkRestrictions("default") + }) + }) + + // ── Gateway deny / allow ────────────────────────────────────────────────── + + Describe("Docker-gateway deny and --allow-docker-gateway", func() { + It("denies the bridge gateway IP by default and reaches it with the flag", func() { + listener, err := net.Listen("tcp", ":0") //nolint:gosec // ephemeral test port + Expect(err).ToNot(HaveOccurred()) + port := listener.Addr().(*net.TCPAddr).Port + srv := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, "host-ok") }), + ReadHeaderTimeout: 5 * time.Second, + } + go func() { _ = srv.Serve(listener) }() + DeferCleanup(func() { _ = srv.Close() }) + + gatewayIP := dockerBridgeGatewayIP() + target := fmt.Sprintf("http://%s:%d/", gatewayIP, port) + + By("Gateway IP must be blocked by default") + deny := startFetchServer("deny", "") + Expect(fetchThrough(deny, target, 30*time.Second).IsError).To(BeTrue(), + "egress must deny the bridge gateway IP without --allow-docker-gateway") + + By("Gateway IP must be reachable with --allow-docker-gateway") + allow := startFetchServer("allow", "", "--allow-docker-gateway") + result := fetchThrough(allow, target, 30*time.Second) + if result.IsError { + Skip("docker bridge gateway is not routable to the host in this environment") + } + Expect(resultText(result)).To(ContainSubstring("host-ok")) + }) + }) + + // ── Envoy-specific regression guards ───────────────────────────────────── + + // ── AllowPort enforcement ───────────────────────────────────────────────── + + Describe("AllowPort enforcement", func() { + // This test proves that AllowPort is honoured end-to-end through a real + // Envoy container, not just in the generated config. It uses a restrictive + // profile that allows example.com on port 443 only and asserts: + // - HTTPS (port 443) succeeds + // - HTTP (port 80) is blocked + // This was a parity gap vs Squid (see #5915): before the fix, Envoy + // ignored AllowPort and the HTTP request would have succeeded. + It("blocks a request on a non-allowed port while permitting the allowed port", func() { + profile := `{ + "name": "port-test", + "network": { + "outbound": { + "insecure_allow_all": false, + "allow_host": ["example.com"], + "allow_port": [443] + } + } + }` + server := startFetchServer("port", profile) + + By("HTTPS request (port 443) must succeed") + // The allowed leg needs a working DNS → TLS → example.com chain through + // Envoy's dynamic_forward_proxy. The very first egress request after + // startup can race the isolation stack's own warm-up (DFP DNS cache, + // the per-workload DNS container), failing instantly with a local 503 + // long before anything has actually reached the network — observed in + // CI as an IsError result within ~40ms of readiness. The property this + // spec pins is steady-state reachability of an allowed port, not + // first-request-after-boot success (which Envoy does not promise), so + // retry until the deadline and surface the tool's error text — the one + // diagnostic that distinguishes a warm-up race from a broken AllowPort. + var lastHTTPSError string + Eventually(func() bool { + httpsResult := fetchThrough(server, "https://example.com", 30*time.Second) + if httpsResult.IsError { + lastHTTPSError = resultText(httpsResult) + return false + } + return true + }, 60*time.Second, 2*time.Second).Should(BeTrue(), func() string { + return "https://example.com must be allowed (port 443 is in AllowPort); last fetch error: " + lastHTTPSError + }) + + By("HTTP request (port 80) must be blocked") + httpResult := fetchThrough(server, "http://example.com", 30*time.Second) + Expect(httpResult.IsError).To(BeTrue(), + "http://example.com must be blocked (port 80 is not in AllowPort)") + }) + }) + + Describe("Envoy route timeout disabled for long-lived streams", func() { + // Guards the timeout:"0s" fix. Envoy's default RouteAction.timeout is 15s; + // this test proves it was disabled by having the upstream sleep 16s (past + // the cap) and asserting the response completes. The 30s CallTool context + // leaves ~14s of MCP overhead headroom after the 16s upstream delay. + It("does not truncate a response that takes longer than Envoy's 15s default", func() { + listener, err := net.Listen("tcp", ":0") //nolint:gosec // ephemeral test port + Expect(err).ToNot(HaveOccurred()) + port := listener.Addr().(*net.TCPAddr).Port + srv := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(16 * time.Second) + _, _ = io.WriteString(w, "slow-ok") + }), + ReadHeaderTimeout: 5 * time.Second, + } + go func() { _ = srv.Serve(listener) }() + DeferCleanup(func() { _ = srv.Close() }) + + target := fmt.Sprintf("http://%s:%d/", dockerBridgeGatewayIP(), port) + server := startFetchServer("slow", "", "--allow-docker-gateway") + + // 60s: 16s upstream sleep + generous MCP client round-trip headroom. + // The 30s default was too tight under CI infrastructure load. + result := fetchThrough(server, target, 60*time.Second) + if result.IsError { + Skip("docker bridge gateway is not routable to the host in this environment") + } + Expect(resultText(result)).To(ContainSubstring("slow-ok"), + "a 16s response must complete — if it doesn't, the 15s route timeout was not disabled") + }) + }) +}) diff --git a/test/e2e/network_isolation_test.go b/test/e2e/network_isolation_test.go index 278a9c113d..7ea78487e5 100644 --- a/test/e2e/network_isolation_test.go +++ b/test/e2e/network_isolation_test.go @@ -15,10 +15,10 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/test/e2e" ) diff --git a/test/e2e/oidc_mock.go b/test/e2e/oidc_mock.go index af670d9f62..b7e3a1e4ad 100644 --- a/test/e2e/oidc_mock.go +++ b/test/e2e/oidc_mock.go @@ -11,6 +11,7 @@ import ( "encoding/json" "errors" "fmt" + "net" "net/http" "time" @@ -26,6 +27,7 @@ import ( // OIDCMockServer represents a lightweight OIDC server using Ory Fosite type OIDCMockServer struct { server *http.Server + listener net.Listener provider fosite.OAuth2Provider store *storage.MemoryStore port int @@ -71,14 +73,31 @@ func WithClientAudience(audiences ...string) OIDCMockOption { // NewOIDCMockServer creates a new OIDC mock server using Ory Fosite. // Use WithClientAudience to set client-level options and WithAccessTokenLifespan // for Fosite-level settings. Both option kinds may be mixed in a single call. +// +// The listener is bound before anything else, closing the window between +// port selection and bind that a separate "pick a port, then bind later" +// step would leave open. Pass port 0 to let the OS choose an ephemeral port; +// use Port() to read back the port actually bound. func NewOIDCMockServer(port int, clientID, clientSecret string, opts ...OIDCMockOption) (*OIDCMockServer, error) { - config := defaultFositeConfig(port) + listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + return nil, fmt.Errorf("failed to listen: %w", err) + } + realPort := listener.Addr().(*net.TCPAddr).Port + + config := defaultFositeConfig(realPort) for _, opt := range opts { if opt.fositeOpt != nil { opt.fositeOpt(config) } } - return newOIDCMockServer(port, clientID, clientSecret, config, opts...) + mockServer, err := newOIDCMockServer(realPort, clientID, clientSecret, config, opts...) + if err != nil { + _ = listener.Close() + return nil, err + } + mockServer.listener = listener + return mockServer, nil } // defaultFositeConfig returns the standard Fosite config for the mock server. @@ -174,7 +193,6 @@ func newOIDCMockServer( mockServer.setupRoutes(mux) mockServer.server = &http.Server{ - Addr: fmt.Sprintf(":%d", port), Handler: mux, ReadHeaderTimeout: 10 * time.Second, // Prevent Slowloris attacks } @@ -430,7 +448,7 @@ func (m *OIDCMockServer) handleJWKS(w http.ResponseWriter, _ *http.Request) { // Start starts the OIDC mock server func (m *OIDCMockServer) Start() error { go func() { - if err := m.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + if err := m.server.Serve(m.listener); err != nil && !errors.Is(err, http.ErrServerClosed) { fmt.Printf("OIDC mock server error: %v\n", err) } }() @@ -440,6 +458,11 @@ func (m *OIDCMockServer) Start() error { return nil } +// Port returns the port the mock server is bound to. +func (m *OIDCMockServer) Port() int { + return m.port +} + // Stop stops the OIDC mock server func (m *OIDCMockServer) Stop() error { if m.server != nil { diff --git a/test/e2e/protocol_builds_e2e_test.go b/test/e2e/protocol_builds_e2e_test.go index 9926066deb..7e08329411 100644 --- a/test/e2e/protocol_builds_e2e_test.go +++ b/test/e2e/protocol_builds_e2e_test.go @@ -207,9 +207,15 @@ var _ = Describe("Protocol Builds E2E", Label("mcp", "mcp-protocol", "protocols" It("should build and start successfully and provide arxiv tools [Serial]", func() { By("Starting the ArXiv MCP server using uvx:// protocol") + // --build-with pins the transitive Python mcp SDK below 2.0: + // arxiv-mcp-server declares an unbounded mcp>=1.27.0 and + // import-crashes under the mcp 2.0.0 major (see #6108). The + // constraint also exercises the uvx builder's --build-with + // plumbing end-to-end. stdout, stderr := e2e.NewTHVCommand(config, "run", "--name", serverName, "--transport", "stdio", + "--build-with", "mcp<2", "uvx://arxiv-mcp-server").ExpectSuccess() // The command should indicate success and show build process diff --git a/test/e2e/proxy_oauth_test.go b/test/e2e/proxy_oauth_test.go index 31d8bd793f..47bff7f8f3 100644 --- a/test/e2e/proxy_oauth_test.go +++ b/test/e2e/proxy_oauth_test.go @@ -15,12 +15,12 @@ import ( "regexp" "strconv" "strings" + "sync" "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/test/e2e" ) @@ -54,27 +54,25 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" osvServerName = generateUniqueOIDCServerName("osv-oauth-target") proxyServerName = generateUniqueOIDCServerName("proxy-oauth-test") - // Find available ports for our mock servers using networking utilities - mockOIDCPort, err = networking.FindOrUsePort(0) - Expect(err).ToNot(HaveOccurred()) - - proxyPort, err = networking.FindOrUsePort(0) - Expect(err).ToNot(HaveOccurred()) - - mockOIDCBaseURL = fmt.Sprintf("http://localhost:%d", mockOIDCPort) + // proxyPort is discovered per-It from the `thv proxy` subprocess's own + // stdout once it binds -- see discoverProxyPort. It's started with + // port 0 rather than a pre-selected port to close the find-then-bind + // TOCTOU window. // Start mock OIDC server using Ory Fosite By("Starting mock OIDC server") specReport := CurrentSpecReport() if strings.Contains(specReport.FullText(), "Proxy OAuth Authentication E2E") { mockOIDCServer, err = e2e.NewOIDCMockServer( - mockOIDCPort, clientID, clientSecret, + 0, clientID, clientSecret, e2e.WithAccessTokenLifespan(2*time.Second), ) } else { - mockOIDCServer, err = e2e.NewOIDCMockServer(mockOIDCPort, clientID, clientSecret) + mockOIDCServer, err = e2e.NewOIDCMockServer(0, clientID, clientSecret) } Expect(err).ToNot(HaveOccurred()) + mockOIDCPort = mockOIDCServer.Port() + mockOIDCBaseURL = fmt.Sprintf("http://localhost:%d", mockOIDCPort) // Enable auto-complete for MCP tests mockOIDCServer.EnableAutoComplete() @@ -137,15 +135,20 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" base := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host) By("Starting the proxy with OAuth configuration") - proxyCmd = startProxyWithOAuth( - config, - proxyServerName, - base, - proxyPort, - mockOIDCBaseURL, - clientID, - clientSecret, - ) + var out *syncBuffer + proxyCmd, out = startProxyWithOAuth(config, proxyServerName, base, 0, mockOIDCBaseURL, clientID, clientSecret) + // This test never drives the OAuth flow to completion (no + // WaitForAuthRequest/CompleteAuthRequest call), so the proxy may + // exit on --remote-auth-timeout before it ever binds a listener + // and prints the port line -- pre-existing behavior, unrelated to + // port selection. Discovery failure here is expected and + // harmless: the assertions below already tolerate the port being + // unset (proxyPort stays 0) or the proxy having exited. + if p, portErr := discoverProxyPort(out, 10*time.Second); portErr != nil { + GinkgoWriter.Printf("proxy port not discovered (OAuth flow may not have completed): %v\n", portErr) + } else { + proxyPort = p + } // Give the proxy some time to start and potentially complete OAuth flow time.Sleep(10 * time.Second) @@ -180,14 +183,11 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" base := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host) By("Starting the proxy with OAuth auto-detection") - proxyCmd = startProxyWithOAuthDetection( - config, - proxyServerName, - base, - proxyPort, - clientID, - clientSecret, - ) + var out *syncBuffer + proxyCmd, out, proxyPort, err = startProxyAndDiscoverPort(func() (*exec.Cmd, *syncBuffer) { + return startProxyWithOAuthDetection(config, proxyServerName, base, 0, clientID, clientSecret) + }, 5*time.Second) + Expect(err).ToNot(HaveOccurred(), "proxy output: %s", out.String()) // Give the proxy time to start time.Sleep(5 * time.Second) @@ -212,15 +212,15 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" base := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host) By("Starting the proxy with invalid OAuth credentials") - proxyCmd = startProxyWithOAuth( - config, - proxyServerName, - base, - proxyPort, - mockOIDCBaseURL, - "invalid-client", - "invalid-secret", - ) + var out *syncBuffer + proxyCmd, out = startProxyWithOAuth(config, proxyServerName, base, 0, mockOIDCBaseURL, "invalid-client", "invalid-secret") + // The proxy is expected to exit on the OAuth failure below before it + // ever binds a listener, so it never prints the port line -- this + // test doesn't use proxyPort, so a discovery failure is expected + // and harmless. + if _, portErr := discoverProxyPort(out, 2*time.Second); portErr != nil { + GinkgoWriter.Printf("proxy port not discovered (expected, proxy should exit on OAuth failure): %v\n", portErr) + } By("Verifying the proxy process exits due to OAuth failure") // The proxy should exit when OAuth fails due to invalid client credentials @@ -255,15 +255,16 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" base := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host) By("Starting the proxy with missing OAuth issuer but remote-auth enabled") - proxyCmd = startProxyWithOAuth( - config, - proxyServerName, - base, - proxyPort, - "", // Empty issuer - clientID, - clientSecret, - ) + const emptyIssuer = "" + var out *syncBuffer + proxyCmd, out = startProxyWithOAuth(config, proxyServerName, base, 0, emptyIssuer, clientID, clientSecret) + // The proxy is expected to exit immediately below due to the + // missing issuer, before it ever binds a listener, so it never + // prints the port line -- this test doesn't use proxyPort, so a + // discovery failure is expected and harmless. + if _, portErr := discoverProxyPort(out, 2*time.Second); portErr != nil { + GinkgoWriter.Printf("proxy port not discovered (expected, proxy should exit due to missing issuer): %v\n", portErr) + } By("Verifying the proxy process exits due to missing issuer") // The proxy should exit immediately when --remote-auth is enabled but issuer is missing @@ -298,14 +299,11 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" base := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host) By("Starting the proxy with auto-detection (no --remote-auth flag)") - proxyCmd = startProxyWithAutoDetection( - config, - proxyServerName, - base, - proxyPort, - clientID, - clientSecret, - ) + var out *syncBuffer + proxyCmd, out, proxyPort, err = startProxyAndDiscoverPort(func() (*exec.Cmd, *syncBuffer) { + return startProxyWithAutoDetection(config, proxyServerName, base, 0, clientID, clientSecret) + }, 5*time.Second) + Expect(err).ToNot(HaveOccurred(), "proxy output: %s", out.String()) // Give the proxy time to try auto-detection time.Sleep(5 * time.Second) @@ -330,50 +328,49 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" GinkgoWriter.Printf("Base URL for proxy: %s\n", baseURL) By("Starting the proxy with OAuth configuration and longer timeout") - var outputBuffer *bytes.Buffer - proxyCmd, outputBuffer = startProxyWithOAuthForMCP( - config, - proxyServerName, - baseURL, // Use base URL instead of full URL - proxyPort, - mockOIDCBaseURL, - clientID, - clientSecret, - ) - - By("Extracting OAuth URL from proxy output and completing the flow") - // Give the proxy a moment to start and display the OAuth URL - time.Sleep(5 * time.Second) - - // Extract OAuth URL from captured output - output := outputBuffer.String() - GinkgoWriter.Printf("Captured proxy output: %s\n", output) - - // Use regex to extract the OAuth URL - // Pattern: "Please open this URL in your browser: " - urlPattern := regexp.MustCompile(`Please open this URL in your browser: (https?://[^\s"]+)`) - matches := urlPattern.FindStringSubmatch(output) - - var authURL string - if len(matches) >= 2 { - authURL = matches[1] - GinkgoWriter.Printf("Extracted OAuth URL from buffer: %s\n", authURL) - } else { - // Fallback: construct the URL from what we know - // We can see the URL in the logs, so let's construct it - authURL = fmt.Sprintf("%s/auth?client_id=%s&response_type=code&scope=openid+profile+email", mockOIDCBaseURL, clientID) - GinkgoWriter.Printf("Using constructed OAuth URL: %s\n", authURL) - } - - // Complete the OAuth flow by visiting the URL with auto_complete parameter - err = completeOAuthFlow(authURL) - if err != nil { - GinkgoWriter.Printf("Failed to complete OAuth flow: %v\n", err) - Skip("Skipping MCP test due to OAuth flow completion failure") - } - - // Wait for proxy to complete OAuth and start - time.Sleep(5 * time.Second) + var outputBuffer *syncBuffer + // startProxyAndDiscoverPort retries the whole closure -- fresh subprocess, + // fresh port, fresh OAuth exchange -- if the previous attempt lost the + // real port-bind race (see the doc comment on discoverProxyPort). + proxyCmd, outputBuffer, proxyPort, err = startProxyAndDiscoverPort(func() (*exec.Cmd, *syncBuffer) { + // baseURL, not the full server URL -- the transparent proxy needs the + // base URL (see the comment above where it's derived). + cmd, out := startProxyWithOAuthForMCP(config, proxyServerName, baseURL, 0, mockOIDCBaseURL, clientID, clientSecret) + + By("Extracting OAuth URL from proxy output and completing the flow") + // Give the proxy a moment to start and display the OAuth URL + time.Sleep(5 * time.Second) + + // Extract OAuth URL from captured output + output := out.String() + GinkgoWriter.Printf("Captured proxy output: %s\n", output) + + // Use regex to extract the OAuth URL + // Pattern: "Please open this URL in your browser: " + urlPattern := regexp.MustCompile(`Please open this URL in your browser: (https?://[^\s"]+)`) + matches := urlPattern.FindStringSubmatch(output) + + var authURL string + if len(matches) >= 2 { + authURL = matches[1] + GinkgoWriter.Printf("Extracted OAuth URL from buffer: %s\n", authURL) + } else { + // Fallback: construct the URL from what we know + // We can see the URL in the logs, so let's construct it + authURL = fmt.Sprintf("%s/auth?client_id=%s&response_type=code&scope=openid+profile+email", mockOIDCBaseURL, clientID) + GinkgoWriter.Printf("Using constructed OAuth URL: %s\n", authURL) + } + + // Complete the OAuth flow by visiting the URL with auto_complete parameter + if flowErr := completeOAuthFlow(authURL); flowErr != nil { + GinkgoWriter.Printf("Failed to complete OAuth flow: %v\n", flowErr) + Skip("Skipping MCP test due to OAuth flow completion failure") + } + + return cmd, out + }, 15*time.Second) + By("Waiting for the proxy to complete OAuth and bind its listener") + Expect(err).ToNot(HaveOccurred(), "proxy output: %s", outputBuffer.String()) By("Testing MCP connection through proxy") proxyURL := fmt.Sprintf("http://localhost:%d/mcp", proxyPort) @@ -417,29 +414,27 @@ var _ = Describe("Proxy OAuth Authentication E2E", Label("proxy", "oauth", "e2e" GinkgoWriter.Printf("Base URL for proxy: %s\n", baseURL) By("Starting the proxy with OAuth-enabled MCP support") - var outputBuffer *bytes.Buffer - proxyCmd, outputBuffer = startProxyWithOAuthForMCP( - config, - proxyServerName, - baseURL, - proxyPort, - mockOIDCBaseURL, - clientID, - clientSecret, - ) - - By("Completing the initial OAuth flow") - Eventually(outputBuffer.String, 5*time.Second, 500*time.Millisecond). - Should(ContainSubstring("Please open this URL")) - - matches := regexp.MustCompile(`Please open this URL in your browser: (https?://[^\s"]+)`). - FindStringSubmatch(outputBuffer.String()) - Expect(matches).To(HaveLen(2)) - authURL := matches[1] - Expect(completeOAuthFlow(authURL)).To(Succeed()) - - By("Giving proxy time to finish OAuth exchange") - time.Sleep(2 * time.Second) + var outputBuffer *syncBuffer + // startProxyAndDiscoverPort retries the whole closure -- fresh subprocess, + // fresh port, fresh OAuth exchange -- if the previous attempt lost the + // real port-bind race (see the doc comment on discoverProxyPort). + proxyCmd, outputBuffer, proxyPort, err = startProxyAndDiscoverPort(func() (*exec.Cmd, *syncBuffer) { + cmd, out := startProxyWithOAuthForMCP(config, proxyServerName, baseURL, 0, mockOIDCBaseURL, clientID, clientSecret) + + By("Completing the initial OAuth flow") + Eventually(out.String, 5*time.Second, 500*time.Millisecond). + Should(ContainSubstring("Please open this URL")) + + matches := regexp.MustCompile(`Please open this URL in your browser: (https?://[^\s"]+)`). + FindStringSubmatch(out.String()) + Expect(matches).To(HaveLen(2)) + authURL := matches[1] + Expect(completeOAuthFlow(authURL)).To(Succeed()) + + return cmd, out + }, 10*time.Second) + By("Waiting for the proxy to finish the OAuth exchange and bind its listener") + Expect(err).ToNot(HaveOccurred(), "proxy output: %s", outputBuffer.String()) By("Waiting for access token to expire") time.Sleep(3 * time.Second) // longer than the 2s lifespan @@ -483,7 +478,108 @@ func checkServerHealth(healthUrl string) error { return fmt.Errorf("server not healthy, status: %d", resp.StatusCode) } -func startProxyWithOAuth(config *e2e.TestConfig, serverName, targetURL string, port int, issuer, clientID, clientSecret string) *exec.Cmd { +// syncBuffer is a bytes.Buffer guarded by a mutex, safe for concurrent +// writes (exec.Cmd's internal stdout/stderr copier goroutines) and reads +// (a polling goroutine like discoverProxyPort, or Eventually matchers) +// while the subprocess is still running. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (s *syncBuffer) Write(p []byte) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.buf.Write(p) +} + +func (s *syncBuffer) String() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.buf.String() +} + +// startLongRunningWithCapture is like e2e.StartLongRunningTHVCommand, but +// also captures stdout/stderr into a buffer. Callers inspect the buffer +// (e.g. via discoverProxyPort) rather than reaping the process via cmd.Wait +// -- cmd.Wait must never be called concurrently, and each test in this file +// that cares about the process's exit decides independently whether and +// when to call it. +func startLongRunningWithCapture(config *e2e.TestConfig, args ...string) (*exec.Cmd, *syncBuffer) { + cmd := exec.Command(config.THVBinary, args...) //nolint:gosec // Intentional for e2e testing + cmd.Env = os.Environ() + + out := &syncBuffer{} + multiWriter := io.MultiWriter(out, GinkgoWriter) + cmd.Stdout = multiWriter + cmd.Stderr = multiWriter + + Expect(cmd.Start()).To(Succeed()) + return cmd, out +} + +// proxyBoundPortPattern matches the port number `thv proxy` prints once its +// listener is actually bound, e.g. "... on port 54321 -> ..." (see +// cmd/thv/app/proxy.go's fmt.Printf after proxy.Start()). +var proxyBoundPortPattern = regexp.MustCompile(`on port (\d+)`) + +// discoverProxyPort polls out (the captured stdout/stderr of a `thv proxy` +// subprocess started with --port 0) for the line the subprocess prints once +// it has actually bound its listener, and returns the port it chose. +// +// Passing port 0 and reading the real port back afterwards -- rather than +// pre-selecting a port and handing it to the subprocess -- narrows the +// find-then-bind TOCTOU window but does not close it: FindOrUsePort(0) +// (cmd/thv/app/proxy.go) still does the classic probe-then-release, and the +// real bind happens later, inside transparent.NewTransparentProxy(...).Start(), +// after the OAuth exchange (handleOutgoingAuthentication) completes. Another +// process can still steal the port during that window; the window is just +// usually much shorter than the original pre-selected-port bug. See +// startProxyAndDiscoverPort for the retry that closes it at call sites that +// use it. +// +// It's expected (and safe to ignore) for this to time out on tests where the +// proxy process exits before ever binding, e.g. on an OAuth failure. +func discoverProxyPort(out *syncBuffer, timeout time.Duration) (int, error) { + deadline := time.Now().Add(timeout) + for { + if m := proxyBoundPortPattern.FindStringSubmatch(out.String()); m != nil { + return strconv.Atoi(m[1]) + } + if time.Now().After(deadline) { + return 0, fmt.Errorf("proxy port not found in output within %s", timeout) + } + time.Sleep(150 * time.Millisecond) + } +} + +// startProxyAndDiscoverPort starts a `thv proxy` subprocess via start, waits +// for it to report its real bound port via discoverProxyPort, and retries +// once (a fresh subprocess, drawing a fresh random port) if the failure was +// a lost port-bind race -- detected by the captured output containing +// "address already in use". A discovery failure for any OTHER reason (e.g. +// this call site's proxy is expected to exit before ever binding, such as +// on an OAuth-config failure) is returned as-is, uncorrected: the caller +// decides whether that's fatal or tolerable. +func startProxyAndDiscoverPort( + start func() (*exec.Cmd, *syncBuffer), + timeout time.Duration, +) (*exec.Cmd, *syncBuffer, int, error) { + cmd, out := start() + port, err := discoverProxyPort(out, timeout) + if err == nil { + return cmd, out, port, nil + } + if !strings.Contains(out.String(), "address already in use") { + return cmd, out, 0, err + } + GinkgoWriter.Printf("proxy lost a port-bind race, retrying once: %v\n", err) + cmd, out = start() + port, err = discoverProxyPort(out, timeout) + return cmd, out, port, err +} + +func startProxyWithOAuth(config *e2e.TestConfig, serverName, targetURL string, port int, issuer, clientID, clientSecret string) (*exec.Cmd, *syncBuffer) { args := []string{ "proxy", "--host", "localhost", @@ -513,10 +609,10 @@ func startProxyWithOAuth(config *e2e.TestConfig, serverName, targetURL string, p // Log the command for debugging GinkgoWriter.Printf("Starting proxy with args: %v\n", args) - return e2e.StartLongRunningTHVCommand(config, args...) + return startLongRunningWithCapture(config, args...) } -func startProxyWithOAuthDetection(config *e2e.TestConfig, serverName, targetURL string, port int, clientID, clientSecret string) *exec.Cmd { +func startProxyWithOAuthDetection(config *e2e.TestConfig, serverName, targetURL string, port int, clientID, clientSecret string) (*exec.Cmd, *syncBuffer) { args := []string{ "proxy", "--host", "localhost", @@ -528,10 +624,10 @@ func startProxyWithOAuthDetection(config *e2e.TestConfig, serverName, targetURL serverName, } - return e2e.StartLongRunningTHVCommand(config, args...) + return startLongRunningWithCapture(config, args...) } -func startProxyWithAutoDetection(config *e2e.TestConfig, serverName, targetURL string, port int, clientID, clientSecret string) *exec.Cmd { +func startProxyWithAutoDetection(config *e2e.TestConfig, serverName, targetURL string, port int, clientID, clientSecret string) (*exec.Cmd, *syncBuffer) { args := []string{ "proxy", "--host", "localhost", @@ -546,10 +642,10 @@ func startProxyWithAutoDetection(config *e2e.TestConfig, serverName, targetURL s // Log the command for debugging GinkgoWriter.Printf("Starting proxy with auto-detection args: %v\n", args) - return e2e.StartLongRunningTHVCommand(config, args...) + return startLongRunningWithCapture(config, args...) } -func startProxyWithOAuthForMCP(config *e2e.TestConfig, serverName, targetURL string, port int, issuer, clientID, clientSecret string) (*exec.Cmd, *bytes.Buffer) { +func startProxyWithOAuthForMCP(config *e2e.TestConfig, serverName, targetURL string, port int, issuer, clientID, clientSecret string) (*exec.Cmd, *syncBuffer) { args := []string{ "proxy", "--host", "localhost", @@ -567,23 +663,7 @@ func startProxyWithOAuthForMCP(config *e2e.TestConfig, serverName, targetURL str // Log the command for debugging GinkgoWriter.Printf("Starting proxy with OAuth for MCP args: %v\n", args) - // Create command - cmd := exec.Command(config.THVBinary, args...) - cmd.Env = os.Environ() - - // Create buffer to capture output (capture both stdout and stderr) - var outputBuffer bytes.Buffer - - // Use MultiWriter to write to both buffer and GinkgoWriter - multiWriter := io.MultiWriter(&outputBuffer, GinkgoWriter) - cmd.Stdout = multiWriter - cmd.Stderr = multiWriter // Capture stderr too since logger might write there - - // Start the command - err := cmd.Start() - Expect(err).ToNot(HaveOccurred()) - - return cmd, &outputBuffer + return startLongRunningWithCapture(config, args...) } // completeOAuthFlow programmatically completes the OAuth flow by visiting the authorization URL diff --git a/test/e2e/proxy_stdio_test.go b/test/e2e/proxy_stdio_test.go index 3833b394a7..cd281e0b0f 100644 --- a/test/e2e/proxy_stdio_test.go +++ b/test/e2e/proxy_stdio_test.go @@ -17,6 +17,7 @@ import ( "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" ) const ( @@ -27,8 +28,6 @@ func generateUniqueProxyStdioServerName(prefix string) string { return fmt.Sprintf("%s-%d-%d-%d", prefix, os.Getpid(), time.Now().UnixNano(), GinkgoRandomSeed()) } -const timeMCPServerImage = "ghcr.io/stacklok/dockyard/uvx/mcp-server-time:2026.7.10@sha256:5cca77dec3fefbacad35e3008ab1660f8725ef246213afb24231504fc73c999d" - var _ = Describe("Proxy Stdio E2E", Label("proxy", "stdio", "e2e"), Serial, func() { var ( config *e2e.TestConfig @@ -185,7 +184,7 @@ var _ = Describe("Proxy Stdio E2E", Label("proxy", "stdio", "e2e"), Serial, func BeforeEach(func() { transportType = types.TransportTypeStdio proxyMode = "sse" - mcpServerName = timeMCPServerImage + mcpServerName = images.TimeMCPServerImage }) It("should proxy MCP requests successfully", func() { By("Getting time server URL") @@ -236,7 +235,7 @@ var _ = Describe("Proxy Stdio E2E", Label("proxy", "stdio", "e2e"), Serial, func BeforeEach(func() { transportType = types.TransportTypeStdio proxyMode = "streamable-http" - mcpServerName = timeMCPServerImage + mcpServerName = images.TimeMCPServerImage }) It("should proxy MCP requests successfully", func() { By("Getting time server URL") diff --git a/test/e2e/restart_test.go b/test/e2e/restart_test.go index d2463b1c7c..db11126f74 100644 --- a/test/e2e/restart_test.go +++ b/test/e2e/restart_test.go @@ -4,7 +4,6 @@ package e2e_test import ( - "fmt" "strings" "time" @@ -203,5 +202,5 @@ var _ = Describe("Server Restart", Label("core", "restart", "e2e"), func() { // generateTestServerName creates a unique server name for restart tests func generateTestServerName(prefix string) string { - return fmt.Sprintf("%s-%d", prefix, GinkgoRandomSeed()) + return e2e.GenerateUniqueServerName(prefix) } diff --git a/test/e2e/rm_group_test.go b/test/e2e/rm_group_test.go index 67c25cc28e..000e24aeab 100644 --- a/test/e2e/rm_group_test.go +++ b/test/e2e/rm_group_test.go @@ -4,9 +4,6 @@ package e2e_test import ( - "fmt" - "time" - . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -23,8 +20,8 @@ var _ = Describe("Group Remove E2E Tests", Label("core", "groups", "e2e"), func( BeforeEach(func() { config = e2e.NewTestConfig() - testGroupName = fmt.Sprintf("rm-test-group-%d-%d", GinkgoRandomSeed(), time.Now().UnixNano()) - secondGroupName = fmt.Sprintf("rm-test-group-2-%d-%d", GinkgoRandomSeed(), time.Now().UnixNano()) + testGroupName = e2e.GenerateUniqueServerName("rm-test-group") + secondGroupName = e2e.GenerateUniqueServerName("rm-test-group-2") createdWorkloads = []string{} // Check if thv binary is available @@ -59,7 +56,7 @@ var _ = Describe("Group Remove E2E Tests", Label("core", "groups", "e2e"), func( Describe("thv rm --group command", func() { It("should return error when group does not exist", func() { - groupName := fmt.Sprintf("rm-non-existent-group-%d", GinkgoRandomSeed()) + groupName := e2e.GenerateUniqueServerName("rm-non-existent-group") _, stderr, err := e2e.NewTHVCommand(config, "rm", "--group", groupName).ExpectFailure() Expect(err).To(HaveOccurred()) Expect(stderr).To(ContainSubstring("does not exist")) @@ -72,20 +69,17 @@ var _ = Describe("Group Remove E2E Tests", Label("core", "groups", "e2e"), func( }) It("should remove workloads from group", func() { - groupWorkload1 := fmt.Sprintf("rm-group-workload-1-%d", GinkgoRandomSeed()) - groupWorkload2 := fmt.Sprintf("rm-group-workload-2-%d", GinkgoRandomSeed()) - nonGroupWorkload1 := fmt.Sprintf("rm-non-group-workload-1-%d", GinkgoRandomSeed()) - nonGroupWorkload2 := fmt.Sprintf("rm-non-group-workload-2-%d", GinkgoRandomSeed()) + groupWorkload1 := e2e.GenerateUniqueServerName("rm-group-workload-1") + groupWorkload2 := e2e.GenerateUniqueServerName("rm-group-workload-2") + nonGroupWorkload1 := e2e.GenerateUniqueServerName("rm-non-group-workload-1") + nonGroupWorkload2 := e2e.GenerateUniqueServerName("rm-non-group-workload-2") createWorkloadInGroup(groupWorkload1, testGroupName) createWorkloadInGroup(groupWorkload2, testGroupName) createWorkloadInGroup(nonGroupWorkload1, secondGroupName) createWorkloadInGroup(nonGroupWorkload2, secondGroupName) // Wait for the workloads to appear in thv list - for _, workloadName := range []string{groupWorkload1, groupWorkload2, nonGroupWorkload1, nonGroupWorkload2} { - err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second) - Expect(err).NotTo(HaveOccurred()) - } + e2e.ExpectMCPServersRunning(config, groupWorkload1, groupWorkload2, nonGroupWorkload1, nonGroupWorkload2) // Remove all workloads in the group e2e.NewTHVCommand(config, "rm", "--group", testGroupName).ExpectSuccess() diff --git a/test/e2e/run_tests.sh b/test/e2e/run_tests.sh index 97f0a6f76f..e948aae9da 100755 --- a/test/e2e/run_tests.sh +++ b/test/e2e/run_tests.sh @@ -45,6 +45,10 @@ fi TEST_TIMEOUT="${TEST_TIMEOUT:-20m}" echo -e "${GREEN}✓${NC} Test timeout: $TEST_TIMEOUT" +# Set number of parallel Ginkgo processes (default 1 = sequential) +PROCS="${PROCS:-1}" +echo -e "${GREEN}✓${NC} Ginkgo procs: $PROCS" + # Export environment variables for tests export THV_BINARY export TEST_TIMEOUT @@ -57,8 +61,9 @@ echo "" cd "$(dirname "$0")" # Build ginkgo command with conditional GitHub output flag -GINKGO_CMD="ginkgo run --timeout=\"$TEST_TIMEOUT\"" +GINKGO_CMD="ginkgo run --timeout=\"$TEST_TIMEOUT\" --procs=$PROCS" GINKGO_CMD="$GINKGO_CMD --junit-report=junit-report.xml --output-dir=." +GINKGO_CMD="$GINKGO_CMD --silence-skips" if [ -n "$GITHUB_ACTIONS" ]; then echo -e "${GREEN}✓${NC} GitHub Actions detected, enabling GitHub output format" GINKGO_CMD="$GINKGO_CMD --github-output --vv" diff --git a/test/e2e/status_test.go b/test/e2e/status_test.go index 01f58012e3..93efba2cc1 100644 --- a/test/e2e/status_test.go +++ b/test/e2e/status_test.go @@ -5,7 +5,6 @@ package e2e_test import ( "encoding/json" - "fmt" "strings" "time" @@ -24,7 +23,7 @@ var _ = Describe("Status Command", Label("core", "status", "e2e"), func() { BeforeEach(func() { config = e2e.NewTestConfig() - serverName = fmt.Sprintf("status-test-%d", GinkgoRandomSeed()) + serverName = e2e.GenerateUniqueServerName("status-test") // Check if thv binary is available err := e2e.CheckTHVBinaryAvailable(config) diff --git a/test/e2e/stdio_proxy_over_streamable_http_mcp_server_test.go b/test/e2e/stdio_proxy_over_streamable_http_mcp_server_test.go index b435d17672..a1c9c610c0 100644 --- a/test/e2e/stdio_proxy_over_streamable_http_mcp_server_test.go +++ b/test/e2e/stdio_proxy_over_streamable_http_mcp_server_test.go @@ -7,6 +7,7 @@ import ( "bytes" "context" "encoding/json" + "io" "net/http" "time" @@ -14,6 +15,7 @@ import ( . "github.com/onsi/gomega" "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" ) var _ = Describe("TimeStreamableHttpMcpServer", Label("proxy", "streamable-http", "e2e"), Serial, func() { @@ -39,13 +41,13 @@ var _ = Describe("TimeStreamableHttpMcpServer", Label("proxy", "streamable-http" } }) - It("should respond to a single get_current_time request and a batch request", func() { + It("should respond to a single get_current_time request and reject a batch request", func() { By("Starting the time MCP server with streamable-http proxy") e2e.NewTHVCommand(config, "run", "--name", serverName, "--transport", "stdio", "--proxy-mode", "streamable-http", - timeMCPServerImage).ExpectSuccess() + images.TimeMCPServerImage).ExpectSuccess() By("Waiting for the server to be running") err := e2e.WaitForMCPServer(config, serverName, 60*time.Second) @@ -77,7 +79,12 @@ var _ = Describe("TimeStreamableHttpMcpServer", Label("proxy", "streamable-http" result := mcpClient.ExpectToolCall(ctx, "get_current_time", arguments) Expect(result.Content).ToNot(BeEmpty(), "Should return the current time") - By("Sending a batch JSON-RPC request") + By("Sending a batch JSON-RPC request, which must be rejected") + // JSON-RPC batching was removed in MCP revision 2025-06-18. ToolHive + // serves only 2025-11-25 and 2026-07-28, so the proxy rejects any + // batch (a top-level JSON array) with an HTTP 400 / JSON-RPC -32600 + // "Invalid Request" and a null id, rather than forwarding its nested + // calls past authz/audit/tool-filtering (see #5745, #5931). batch := []map[string]interface{}{ { "method": "tools/call", @@ -116,23 +123,18 @@ var _ = Describe("TimeStreamableHttpMcpServer", Label("proxy", "streamable-http" resp, err := client.Do(req) Expect(err).ToNot(HaveOccurred()) defer resp.Body.Close() - Expect(resp.StatusCode).To(Equal(200)) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest), "Batch requests must be rejected") - var responses []map[string]interface{} - decoder := json.NewDecoder(resp.Body) - err = decoder.Decode(&responses) + body, err := io.ReadAll(resp.Body) Expect(err).ToNot(HaveOccurred()) - Expect(responses).To(HaveLen(2)) - - ids := map[float64]bool{4: false, 5: false} - for _, r := range responses { - id, ok := r["id"].(float64) - Expect(ok).To(BeTrue(), "Each response should have an id") - ids[id] = true - Expect(r["result"]).ToNot(BeNil(), "Each response should have a result") - } - Expect(ids[4]).To(BeTrue()) - Expect(ids[5]).To(BeTrue()) + Expect(string(body)).To(ContainSubstring(`"code":-32600`), "Should carry a JSON-RPC Invalid Request error") + + // A batch has no single request id to echo. MCP encodes that by + // omitting the "id" key entirely (schema/2025-11-25 types the + // error response id as optional, not nullable), never as null. + var parsed map[string]any + Expect(json.Unmarshal(body, &parsed)).To(Succeed()) + Expect(parsed).ToNot(HaveKey("id"), "Batch rejection omits the JSON-RPC id key rather than emitting null") }) }) }) diff --git a/test/e2e/testdata/golden_modern_request.json b/test/e2e/testdata/golden_modern_request.json new file mode 100644 index 0000000000..3576d57959 --- /dev/null +++ b/test/e2e/testdata/golden_modern_request.json @@ -0,0 +1,19 @@ +{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": { + "_meta": { + "io.modelcontextprotocol/clientCapabilities": { + "roots": { + "listChanged": true + } + }, + "io.modelcontextprotocol/clientInfo": { + "name": "golden-fixture-capture", + "version": "0.0.0" + }, + "io.modelcontextprotocol/protocolVersion": "2026-07-28" + } + } +} diff --git a/test/e2e/thv-operator/acceptance_tests/dual_era_k8s_test.go b/test/e2e/thv-operator/acceptance_tests/dual_era_k8s_test.go new file mode 100644 index 0000000000..855f7add00 --- /dev/null +++ b/test/e2e/thv-operator/acceptance_tests/dual_era_k8s_test.go @@ -0,0 +1,308 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package acceptancetests + +import ( + "context" + "fmt" + "net/http" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/meta" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/utils/ptr" + "sigs.k8s.io/controller-runtime/pkg/client" + + mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" + ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil" + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" + "github.com/stacklok/toolhive/test/e2e/thv-operator/testutil" +) + +// This is Commit 4, Steps 4.1+4.2 of the dual-era e2e plan (issue #5837). +// AUTHORED FOR CI (test-e2e-lifecycle.yml) -- NOT run locally. Per the +// Step 4.0 discovery, this worktree's environment only has a shared, +// long-lived kind cluster (unrelated ngrok-operator/xaa-phase1 namespaces, +// an operator build that predates this merge) that must not be touched, and +// it lacks the NodePort host-port range this harness needs anyway. Step 4.1's +// core (BeforeAll + the pods-ready/Redis-wired It) WAS verified live against +// that cluster in an isolated, self-cleaning throwaway namespace before this +// rewrite -- see the git history of this file / the Step 4.1 report. Every +// addition since (NodePort access, traffic-driving, pod eviction, Redis +// kill) is new and verified only by go build/vet/lint, not a live run. +// +// Runs in its own namespace (not "default") to keep its Redis and MCPServer +// resources isolated from the concurrent ratelimit_test.go Ordered block +// under --procs=8. +// +// SCOPE (deliberate): this single-server MCPServer tier covers the MODERN +// (2026-07-28, stateless) path only -- Modern multi-replica routing and +// pod-eviction self-heal. It deliberately does NOT cover mixed Legacy+Modern +// traffic or a session-store (Redis) outage: both need a Legacy (session) +// client to open a session against this STATELESS backend, and a single +// go-sdk backend cannot be both stateless (to serve Modern per-request +// clients, spec below) and session-issuing (for Legacy clients) at once. +// Serving both eras over one backend is cross-generation bridging, which is a +// vMCP capability (epic #5743, design #5756), out of scope for a single-server +// MCPServer -- the proxy here does per-request version discrimination, not +// generation bridging. When that bridge is implemented, add the mixed-era and +// session-store-outage coverage as a vMCP e2e test, not here. +// +// Confirmed live (Step 4.1): the MCPServer controller creates TWO separate +// workloads -- a Deployment named after the CR (the proxy runner, one +// "toolhive" container, selected by app.kubernetes.io/name=mcpserver + +// app.kubernetes.io/instance=, sized by spec.replicas) and a +// StatefulSet with the same name (the actual MCP server backend, one "mcp" +// container running the image, selected by app=, sized by +// spec.backendReplicas). Both reached Ready with spec.sessionAffinity: None +// and Redis-backed spec.sessionStorage; the CR's Ready condition surfaces +// session storage as a (non-)warning: type SessionStorageWarning, status +// False, reason SessionStorageConfigured once Redis is wired correctly. +// +// [MoE finding, resolved] The original design tried to prove "no pod-pinning" +// by tagging each Modern request with a unique _meta.nonce and grepping +// which backend pod's log (yardstick's echoHandler does +// `log.Printf("echo tool called with metadata: %+v", req.Params.Meta)`, +// cmd/yardstick-server/main.go:69 -- confirmed from source, not a run) +// contained it, then asserting >=2 distinct pods appear. That assertion is +// UNSOUND and was dropped: kube-proxy balances per TCP CONNECTION, not per +// HTTP request, and BOTH hops (raw-client-to-proxy AND the proxy's own +// outbound connection to the single backend ClusterIP) reuse keep-alive +// connections outside this test's control. With only 2 backend replicas, +// "all N requests land in one pod's log" is an ordinary, EXPECTED outcome +// of connection reuse under perfectly correct (non-pinned) behavior -- the +// test's signal cannot distinguish "working as intended" from "pinned", +// which is the definition of an unsound assertion, not a real check. There +// is no viable alternative signal either: the proxy itself logs nothing +// per-request that's nonce-correlatable (only structured slog.Debug lines +// with no request-identifying payload), so attributing to PROXY pods +// (rather than backend pods) has the same problem from the other direction. +// +// What replaced it, per the plan's own pre-approved fallback: the +// deterministic, sound properties. (1) The k8s Service actually has +// SessionAffinity: None set (a direct field read of the live object, not an +// inference from traffic) -- this is what spec.sessionAffinity actually +// configures, and is what would silently regress to ClientIP if the CRD +// wiring broke. (2) Many independent Modern requests all succeed against +// the multi-replica deployment (functional correctness under load, still a +// real check, just not a distinct-pod count). +// +// [MoE-relevant] Separately NOT asserting the literal wording "a Modern +// request carrying a stale/foreign Mcp-Session-Id is still served, not +// 404'd": that contradicts the security-hardened, MoE-reviewed behavior +// established in Step 1.5/the hostile-input suite +// (revision_guard_regression_test.go's +// TestGuardUnknownSessionFiresDespiteForgedModernRevision) -- the session +// guard deliberately fires on Mcp-Session-Id header PRESENCE regardless of +// classified revision, precisely so a forged Modern signal can't bypass +// session validation. A foreign session id on Modern getting 404 there is +// correct, not a bug. +var _ = Describe("Dual-Era Multi-Replica Backend", Ordered, func() { + const ( + testNamespace = "dual-era-k8s" + serverName = "dual-era-echo" + // Replicas drives the proxy runner Deployment, BackendReplicas the backend + // StatefulSet -- independent counts, and Status.ReadyReplicas reports the proxy + // one only ("number of ready proxy replicas", mcpserver_types.go). + proxyReplicas = 2 + backendReplicas = 2 + timeout = 3 * time.Minute + pollingInterval = 2 * time.Second + ) + + var nodePort int32 + + backendPodLabels := map[string]string{"app": serverName} + + proxyURL := func() string { + return fmt.Sprintf("http://localhost:%d/mcp", nodePort) + } + + // backendPods lists the StatefulSet-backed backend pods for this MCPServer. + backendPods := func() []corev1.Pod { + podList := &corev1.PodList{} + Expect(k8sClient.List(ctx, podList, client.InNamespace(testNamespace), client.MatchingLabels(backendPodLabels))).To(Succeed()) + return podList.Items + } + + BeforeAll(func() { + By("Creating a dedicated namespace") + ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: testNamespace}} + Expect(client.IgnoreAlreadyExists(k8sClient.Create(ctx, ns))).To(Succeed()) + + By("Deploying Redis for session storage") + EnsureRedis(ctx, k8sClient, testNamespace, timeout, pollingInterval) + + By("Creating a multi-replica MCPServer (proxy + backend) with Redis session storage") + server := &mcpv1beta1.MCPServer{ + ObjectMeta: metav1.ObjectMeta{ + Name: serverName, + Namespace: testNamespace, + }, + Spec: mcpv1beta1.MCPServerSpec{ + Image: images.YardstickServerImage, + Transport: "streamable-http", + ProxyPort: 8080, + MCPPort: 8080, + Replicas: ptr.To(int32(proxyReplicas)), + BackendReplicas: ptr.To(int32(backendReplicas)), + SessionAffinity: "None", + Env: []mcpv1beta1.EnvVar{ + {Name: "TRANSPORT", Value: "streamable-http"}, + {Name: "STATELESS", Value: "true"}, // Modern (2026-07-28) capable backend + {Name: "BACKEND_MODE", Value: "echo"}, + }, + SessionStorage: &mcpv1beta1.SessionStorageConfig{ + Provider: "redis", + Address: "redis." + testNamespace + ".svc.cluster.local:6379", + }, + }, + } + Expect(k8sClient.Create(ctx, server)).To(Succeed()) + + // Ping-free readiness: the operator's Deployment readiness probe is a + // plain HTTP GET /health (confirmed live in Step 4.0/4.1, reading the + // Deployment spec), never an MCP "ping" (which Modern removed). This + // comment describes the OPERATOR's own probe; the test's own /health + // gate below is separate, added because kube-proxy route programming + // can lag pod-Ready. + By("Waiting for MCPServer to be running") + testutil.WaitForMCPServerRunning(ctx, k8sClient, serverName, testNamespace, timeout, pollingInterval) + + By("Creating NodePort service for MCPServer proxy") + testutil.CreateNodePortService(ctx, k8sClient, serverName, testNamespace) + nodePort = testutil.GetNodePort(ctx, k8sClient, serverName+"-nodeport", testNamespace, timeout, pollingInterval) + + By("Waiting for the proxy to actually be reachable on the NodePort") + httpClient := &http.Client{Timeout: 5 * time.Second} + Eventually(func() error { + resp, err := httpClient.Get(fmt.Sprintf("http://localhost:%d/health", nodePort)) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("health check returned %d", resp.StatusCode) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + }) + + AfterAll(func() { + By("Deleting the dedicated namespace (cascades: MCPServer, Redis, NodePort service)") + _ = k8sClient.Delete(ctx, &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: testNamespace}}) + }) + + It("brings up both proxy and backend replicas ready, with Redis wired", func() { + // Gate on the full replica count, not just "something is up": CheckPodsReady + // passes on a single ready pod, so it would clear here with 1 of 2 replicas and + // leave the multi-replica premise of this whole suite unverified. + By(fmt.Sprintf("all %d proxy runner pods are ready", proxyReplicas)) + Eventually(func() error { + return testutil.CheckPodsReadyCount(ctx, k8sClient, testNamespace, map[string]string{ + "app.kubernetes.io/name": "mcpserver", + "app.kubernetes.io/instance": serverName, + }, proxyReplicas) + }, timeout, pollingInterval).Should(Succeed()) + + By(fmt.Sprintf("all %d backend (yardstick) pods are ready", backendReplicas)) + Eventually(func() error { + return testutil.CheckPodsReadyCount(ctx, k8sClient, testNamespace, backendPodLabels, backendReplicas) + }, timeout, pollingInterval).Should(Succeed()) + + By("Redis session storage is wired without a warning, with both replicas ready in status") + // Status is written by the reconciler, so it trails the pods above by up to one + // reconcile even now that the gates wait for the full count. Poll rather than + // reading once. + server := &mcpv1beta1.MCPServer{} + Eventually(func() error { + if err := k8sClient.Get(ctx, client.ObjectKey{Name: serverName, Namespace: testNamespace}, server); err != nil { + return err + } + cond := meta.FindStatusCondition(server.Status.Conditions, "SessionStorageWarning") + if cond == nil { + return fmt.Errorf("expected a SessionStorageWarning condition once sessionStorage is set") + } + if cond.Status != metav1.ConditionFalse { + return fmt.Errorf("Redis session storage should be configured cleanly: %s", cond.Message) + } + if server.Status.ReadyReplicas != proxyReplicas { + return fmt.Errorf("status.readyReplicas is %d, want %d", server.Status.ReadyReplicas, proxyReplicas) + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + }) + + It("configures SessionAffinity:None and serves many Modern requests across the multi-replica deployment", func() { + By("the proxy Service actually has SessionAffinity: None (not the ClientIP default)") + svc := &corev1.Service{} + Expect(k8sClient.Get(ctx, client.ObjectKey{ + Name: ctrlutil.CreateProxyServiceName(serverName), + Namespace: testNamespace, + }, svc)).To(Succeed()) + Expect(svc.Spec.SessionAffinity).To(Equal(corev1.ServiceAffinityNone)) + + By("many independent sessionless Modern requests all succeed") + rawClient, err := e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + for i := range 8 { + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "anypod"}, + }) + Expect(err).ToNot(HaveOccurred()) + // The backend is a real go-sdk streamable-HTTP server, which 400s a + // POST without this Accept; the ToolHive proxy tiers must NOT set it. + req.WithStreamableAccept() + resp, err := rawClient.Send(context.Background(), proxyURL(), req) + Expect(err).ToNot(HaveOccurred(), "request %d", i) + Expect(resp.StatusCode).To(Equal(http.StatusOK), "request %d", i) + Expect(resp.Error).To(BeNil(), "request %d", i) + Expect(resp.Headers.Get(e2e.HeaderMCPSessionID)).To(BeEmpty(), "Modern must never carry Mcp-Session-Id") + } + }) + + It("degrades cleanly and self-heals when a backend pod is evicted mid-traffic", func() { + rawClient, err := e2e.NewRawMCPClient(15 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + before := backendPods() + Expect(before).To(HaveLen(backendReplicas), "precondition: all backend replicas present before evicting") + victim := before[0].Name + + By(fmt.Sprintf("deleting backend pod %s", victim)) + Expect(k8sClient.Delete(ctx, &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: victim, Namespace: testNamespace}, + })).To(Succeed()) + + By("a fresh Modern request still succeeds -- served by a surviving or replacement pod") + Eventually(func() (int, error) { + req, buildErr := e2e.NewModernRequest("tools/call", map[string]any{ + "name": "echo", + "arguments": map[string]any{"input": "aftereviction"}, + }) + if buildErr != nil { + return 0, buildErr + } + req.WithStreamableAccept() // real go-sdk backend requires Accept: text/event-stream + resp, sendErr := rawClient.Send(context.Background(), proxyURL(), req) + if sendErr != nil { + return 0, sendErr + } + return resp.StatusCode, nil + }, timeout, pollingInterval).Should(Equal(http.StatusOK)) + + // CheckPodsReady only requires ONE ready pod, so it would pass the instant the + // surviving pod is observed -- it never proves the StatefulSet actually replaced + // the evicted one. Assert the full count is back AND all of them are Ready. + By("the StatefulSet replaces the evicted pod -- full replica count, all Ready") + Eventually(func() error { + return testutil.CheckPodsReadyCount(ctx, k8sClient, testNamespace, backendPodLabels, backendReplicas) + }, timeout, pollingInterval).Should(Succeed()) + }) +}) diff --git a/test/e2e/thv-operator/acceptance_tests/ratelimit_test.go b/test/e2e/thv-operator/acceptance_tests/ratelimit_test.go index c1fc7372d7..438e95b326 100644 --- a/test/e2e/thv-operator/acceptance_tests/ratelimit_test.go +++ b/test/e2e/thv-operator/acceptance_tests/ratelimit_test.go @@ -122,13 +122,13 @@ var _ = Describe("MCPServer Rate Limiting", Ordered, func() { Expect(status).To(Equal(http.StatusTooManyRequests), "4th request should be rate limited, body: %s", string(body)) - By("Verifying JSON-RPC error code -32029") + By("Verifying JSON-RPC error code 429") var resp map[string]any Expect(json.Unmarshal(body, &resp)).To(Succeed()) errObj, ok := resp["error"].(map[string]any) Expect(ok).To(BeTrue(), "response should have error object") - Expect(errObj["code"]).To(BeEquivalentTo(-32029)) + Expect(errObj["code"]).To(BeEquivalentTo(429)) Expect(errObj["message"]).To(Equal("Rate limit exceeded")) data, ok := errObj["data"].(map[string]any) @@ -314,13 +314,13 @@ var _ = Describe("MCPServer Rate Limiting", Ordered, func() { By("Verifying Retry-After header is present (AC12)") Expect(retryAfter).ToNot(BeEmpty(), "Retry-After header should be set on 429 response") - By("Verifying JSON-RPC error code -32029 with retryAfterSeconds") + By("Verifying JSON-RPC error code 429 with retryAfterSeconds") var resp map[string]any Expect(json.Unmarshal(body, &resp)).To(Succeed()) errObj, ok := resp["error"].(map[string]any) Expect(ok).To(BeTrue(), "response should have error object") - Expect(errObj["code"]).To(BeEquivalentTo(-32029)) + Expect(errObj["code"]).To(BeEquivalentTo(429)) data, ok := errObj["data"].(map[string]any) Expect(ok).To(BeTrue(), "error should have data object") diff --git a/test/e2e/thv-operator/testutil/k8s.go b/test/e2e/thv-operator/testutil/k8s.go index 17fe03adaa..fa6031113b 100644 --- a/test/e2e/thv-operator/testutil/k8s.go +++ b/test/e2e/thv-operator/testutil/k8s.go @@ -26,32 +26,40 @@ import ( ) // CheckPodsReady checks that at least one pod matching the given labels is running and ready. +// +// It passes the instant a single pod is observed ready, so it cannot gate on a multi-replica +// rollout -- use CheckPodsReadyCount when the workload has a known replica count. func CheckPodsReady(ctx context.Context, c client.Client, namespace string, labels map[string]string) error { - podList := &corev1.PodList{} - if err := c.List(ctx, podList, - client.InNamespace(namespace), - client.MatchingLabels(labels)); err != nil { - return fmt.Errorf("failed to list pods: %w", err) + ready, total, err := countReadyPods(ctx, c, namespace, labels) + if err != nil { + return err } - - if len(podList.Items) == 0 { + if total == 0 { return fmt.Errorf("no pods found with labels %v", labels) } + if ready == 0 { + return fmt.Errorf("no running and ready pods found with labels %v", labels) + } + return nil +} - runningPods := 0 - for _, pod := range podList.Items { - if pod.Status.Phase != corev1.PodRunning { - continue - } - for _, cond := range pod.Status.Conditions { - if cond.Type == corev1.PodReady && cond.Status == corev1.ConditionTrue { - runningPods++ - break - } - } +// CheckPodsReadyCount checks that exactly want pods matching the given labels are running and +// ready. Use this over CheckPodsReady whenever the expected replica count is known, so a gate +// on "the workload is up" cannot be satisfied by a partial rollout. +func CheckPodsReadyCount( + ctx context.Context, + c client.Client, + namespace string, + labels map[string]string, + want int, +) error { + ready, total, err := countReadyPods(ctx, c, namespace, labels) + if err != nil { + return err } - if runningPods == 0 { - return fmt.Errorf("no running and ready pods found with labels %v", labels) + if ready != want { + return fmt.Errorf("expected %d running and ready pods with labels %v, got %d ready of %d", + want, labels, ready, total) } return nil } @@ -170,3 +178,35 @@ func GetNodePort( return nodePort } + +// countReadyPods returns how many pods matching labels are both Running and Ready, alongside the +// total number matched. Terminating pods count towards the total but never towards ready: they can +// still report Ready while shutting down, and the operator excludes them from its own status counts +// for that reason (categorizePodStatus in mcpserver_controller.go, issue #4498), so counting them +// would let a gate pass on a pod that is going away. +func countReadyPods( + ctx context.Context, + c client.Client, + namespace string, + labels map[string]string, +) (ready, total int, err error) { + podList := &corev1.PodList{} + if err := c.List(ctx, podList, + client.InNamespace(namespace), + client.MatchingLabels(labels)); err != nil { + return 0, 0, fmt.Errorf("failed to list pods: %w", err) + } + + for _, pod := range podList.Items { + if pod.DeletionTimestamp != nil || pod.Status.Phase != corev1.PodRunning { + continue + } + for _, cond := range pod.Status.Conditions { + if cond.Type == corev1.PodReady && cond.Status == corev1.ConditionTrue { + ready++ + break + } + } + } + return ready, len(podList.Items), nil +} diff --git a/test/e2e/thv-operator/virtualmcp/helpers.go b/test/e2e/thv-operator/virtualmcp/helpers.go index 4f9f67420d..a14bbb8a1c 100644 --- a/test/e2e/thv-operator/virtualmcp/helpers.go +++ b/test/e2e/thv-operator/virtualmcp/helpers.go @@ -21,9 +21,6 @@ import ( "time" "github.com/go-logr/logr" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" appsv1 "k8s.io/api/apps/v1" @@ -35,6 +32,9 @@ import ( "k8s.io/utils/ptr" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/test/e2e/images" "github.com/stacklok/toolhive/test/e2e/thv-operator/testutil" @@ -112,6 +112,18 @@ func (c *InitializedMCPClient) Close() { // CreateInitializedMCPClient creates an MCP client, starts the transport, and initializes // the connection with the given client name. Returns an InitializedMCPClient that should // be closed when done using defer client.Close(). +// +// DO NOT use this (or any mcpcompat-built client) in a spec that asserts +// SESSION semantics against a vMCP. The underlying go-sdk client is +// Modern-first — it probes server/discover before initialize and upgrades to +// 2026-07-28 whenever the server advertises it — and cannot be pinned to +// Legacy (#5911). Against a Modern-serving vMCP it gets no session, and +// GetSessionId() is empty. Session specs must pin the era explicitly with the +// raw Legacy primitives in legacy_session_helpers_test.go instead. (Specs that +// only assert tool/resource behavior are era-agnostic and may use this freely; +// clients of the thv-proxyrunner transparent proxy negotiate against the +// BACKEND, so MCPServer/MCPRemoteProxy session specs are unaffected as long as +// the backend itself is Legacy.) func CreateInitializedMCPClient(nodePort int32, clientName string, timeout time.Duration) (*InitializedMCPClient, error) { serverURL := fmt.Sprintf("http://localhost:%d/mcp", nodePort) mcpClient, err := mcpclient.NewStreamableHttpClient(serverURL) @@ -1278,10 +1290,11 @@ type VMCPStatusResponse struct { // VMCPBackendStatus mirrors server.BackendStatus // (pkg/vmcp/server/status.go) for test deserialization. type VMCPBackendStatus struct { - Name string `json:"name"` - Health string `json:"health"` // "healthy", "degraded", "unhealthy", "unknown" - Transport string `json:"transport"` - AuthType string `json:"auth_type,omitempty"` + Name string `json:"name"` + Health string `json:"health"` // "healthy", "degraded", "unhealthy", "unknown" + Transport string `json:"transport"` + AuthType string `json:"auth_type,omitempty"` + MCPRevision string `json:"mcp_revision,omitempty"` } // VMCPBackendsHealthResponse mirrors BackendHealthResponse diff --git a/test/e2e/thv-operator/virtualmcp/legacy_session_helpers_test.go b/test/e2e/thv-operator/virtualmcp/legacy_session_helpers_test.go new file mode 100644 index 0000000000..d5a50e969a --- /dev/null +++ b/test/e2e/thv-operator/virtualmcp/legacy_session_helpers_test.go @@ -0,0 +1,146 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package virtualmcp + +import ( + "context" + "encoding/json" + "fmt" + + coremcp "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/test/e2e" +) + +// This file holds the era-pinned Legacy (2025-11-25) session primitives shared +// by the specs that assert Legacy-session semantics — Redis-backed cross-pod +// reconstruction, pod-restart recovery, lazy termination eviction, +// upstreamInject identity restore, and the dual-era Redis suite. +// +// They exist because such specs must NOT use CreateInitializedMCPClient (or +// any mcpcompat-built client): mcpcompat wraps go-sdk, go-sdk v1.7's Connect +// is Modern-first — it probes server/discover before initialize and upgrades +// to 2026-07-28 whenever the server advertises it — and the shim cannot pin a +// protocol version (#5911). Against a Modern-serving vMCP such a client +// negotiates Modern, which has no sessions, and every session assertion fails. +// The raw client pins the era explicitly per request instead, which is the +// point: a spec about sessions declares, at each call site, that it depends on +// the Legacy revision (the #6051 convention). +// +// Both helpers return errors rather than asserting, so they are safe inside +// Eventually/Consistently retry loops and on goroutines without GinkgoRecover. + +// legacySessionInit performs the Legacy initialize + notifications/initialized +// handshake against url and returns the session ID vMCP assigned. extraHeaders +// (e.g. Authorization for OIDC-protected instances) are applied to every +// request of the handshake. +func legacySessionInit( + client *e2e.RawMCPClient, url, clientName string, extraHeaders map[string]string, +) (string, error) { + ctx := context.Background() + req := e2e.NewLegacyInitializeRequest(clientName, "1.0") + for k, v := range extraHeaders { + req.SetHeader(k, v) + } + resp, err := client.Send(ctx, url, req) + if err != nil { + return "", fmt.Errorf("initialize: %w", err) + } + if resp.StatusCode != 200 { + return "", fmt.Errorf("initialize: status %d, body: %s", resp.StatusCode, resp.Body) + } + sessionID := resp.Headers.Get(e2e.HeaderMCPSessionID) + if sessionID == "" { + return "", fmt.Errorf("initialize did not assign a session id") + } + + notifyReq, err := e2e.NewLegacyRequest("notifications/initialized", nil) + if err != nil { + return "", fmt.Errorf("build notifications/initialized: %w", err) + } + // WithID(nil) omits "id" entirely, making this a true JSON-RPC notification. + notifyReq.WithID(nil).WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, coremcp.MCPVersionLegacy) + for k, v := range extraHeaders { + notifyReq.SetHeader(k, v) + } + notifyResp, err := client.Send(ctx, url, notifyReq) + if err != nil { + return "", fmt.Errorf("notifications/initialized: %w", err) + } + if notifyResp.StatusCode != 202 { + return "", fmt.Errorf("notifications/initialized: status %d, body: %s", notifyResp.StatusCode, notifyResp.Body) + } + + return sessionID, nil +} + +// legacySessionCallTool sends a Legacy tools/call for toolName on sessionID +// and returns the raw response. It deliberately does NOT inspect +// resp.StatusCode/Error — a non-2xx response is a normal, successfully-received +// HTTP response some callers assert on (e.g. rejection specs); only an actual +// transport failure from Send is surfaced as the returned error. For a plain +// "the call worked and echoed X" assertion, pass the response to +// dualEraEchoErr(resp, wantOutput, "") — the empty resultType is what a Legacy +// client's envelope carries. +func legacySessionCallTool( + client *e2e.RawMCPClient, url, sessionID, toolName string, + args map[string]any, extraHeaders map[string]string, +) (*e2e.RawResponse, error) { + req, err := e2e.NewLegacyRequest("tools/call", map[string]any{ + "name": toolName, + "arguments": args, + }) + if err != nil { + return nil, fmt.Errorf("build tools/call: %w", err) + } + req.WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, coremcp.MCPVersionLegacy) + for k, v := range extraHeaders { + req.SetHeader(k, v) + } + return client.Send(context.Background(), url, req) +} + +// legacySessionListTools sends a Legacy tools/list on sessionID and returns +// the aggregated tool names. Sending it to a pod that has never seen sessionID +// is exactly what triggers a Redis-backed session reconstruction there, so the +// cross-pod specs call this against pod B with pod A's session. extraHeaders +// are applied to the request. +func legacySessionListTools( + client *e2e.RawMCPClient, url, sessionID string, extraHeaders map[string]string, +) ([]string, error) { + req, err := e2e.NewLegacyRequest("tools/list", nil) + if err != nil { + return nil, fmt.Errorf("build tools/list: %w", err) + } + req.WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, coremcp.MCPVersionLegacy) + for k, v := range extraHeaders { + req.SetHeader(k, v) + } + resp, err := client.Send(context.Background(), url, req) + if err != nil { + return nil, fmt.Errorf("tools/list: %w", err) + } + if resp.StatusCode != 200 { + return nil, fmt.Errorf("tools/list: status %d, body: %s", resp.StatusCode, resp.Body) + } + // A schema-rejected tools/list comes back as a JSON-RPC error inside a 200 + // (unlike tools/call, where yardstick reports validation failure via + // isError:true), so this check is meaningful. + if resp.Error != nil { + return nil, fmt.Errorf("tools/list: JSON-RPC error: %+v", resp.Error) + } + + var result struct { + Tools []struct { + Name string `json:"name"` + } `json:"tools"` + } + if err := json.Unmarshal(resp.Result, &result); err != nil { + return nil, fmt.Errorf("tools/list: unmarshal result: %w, raw: %s", err, resp.Result) + } + names := make([]string, len(result.Tools)) + for i, t := range result.Tools { + names[i] = t.Name + } + return names, nil +} diff --git a/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go b/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go index f226d99760..139587b817 100644 --- a/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go +++ b/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go @@ -9,9 +9,6 @@ import ( "fmt" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" @@ -20,6 +17,9 @@ import ( "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" "github.com/stacklok/toolhive/test/e2e/images" diff --git a/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go b/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go index 8d10d954ef..9b3a32a63c 100644 --- a/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go +++ b/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go @@ -13,9 +13,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" "github.com/redis/go-redis/v9" @@ -27,6 +24,9 @@ import ( "k8s.io/apimachinery/pkg/util/intstr" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/test/e2e/images" "github.com/stacklok/toolhive/test/e2e/thv-operator/testutil" @@ -116,6 +116,24 @@ func cleanupRedis(name string) { }) } +// scaleRedis sets the named Redis Deployment's replica count — 0 to simulate a +// session-store outage, 1 to restore. Parameterized by name (unlike the +// acceptance-tier helper removed in 5fa67123b, which hardcoded a shared "redis" +// Deployment) so scaling one spec's Redis cannot disturb another's under +// --procs=8. +// +// Callers should wrap this in Eventually: it is a read-modify-write, and +// kube-controller-manager writes .status on the same object, so an isolated +// Update can lose a race and return a 409 conflict. +func scaleRedis(name string, replicas int32) error { + deploy := &appsv1.Deployment{} + if err := k8sClient.Get(ctx, types.NamespacedName{Name: name, Namespace: defaultNamespace}, deploy); err != nil { + return err + } + deploy.Spec.Replicas = int32Ptr(replicas) + return k8sClient.Update(ctx, deploy) +} + // getReadyMCPServerPods returns all Running+Ready pods for an MCPServer. // //nolint:unparam // namespace kept as parameter for reusability across test contexts diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go index a7ceca6611..8da104d667 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go index 56cd88f0ec..a49a887e03 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go index 42e851d8a8..f0fc6a13c3 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go @@ -12,9 +12,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" @@ -24,6 +21,9 @@ import ( "k8s.io/apimachinery/pkg/util/intstr" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go index 980fa5db84..26933ab179 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go @@ -9,11 +9,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go index 0d6de43cf0..d5d948f12c 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go @@ -7,11 +7,11 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go index c7eac4cc92..b07b02595f 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go index 3f8ed12e73..5dc1f6fe96 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go index 38c40e03fb..f568c9168a 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" @@ -84,6 +84,13 @@ var _ = Describe("VirtualMCPServer Composite Referenced Workflow", Ordered, func }, }, Timeout: vmcpconfig.Duration(30 * time.Second), + // Explicit annotations on the referenced definition. More + // conservative than any read-only floor, so they pass the + // safety-floor guardrail and reach ListTools. + Annotations: &vmcpconfig.ToolAnnotationsOverride{ + Title: stringPtr("Referenced Echo Twice"), + DestructiveHint: boolPtr(true), + }, }, }, } @@ -230,6 +237,36 @@ var _ = Describe("VirtualMCPServer Composite Referenced Workflow", Ordered, func // Second echo: echoes the result of first echo GinkgoWriter.Printf("Referenced composite tool result: %+v\n", result.Content) }) + + It("should advertise the referenced definition's annotations in tool listing", func() { + By("Creating and initializing MCP client for VirtualMCPServer") + mcpClient, err := CreateInitializedMCPClient(vmcpNodePort, "toolhive-composite-ref-test", 30*time.Second) + Expect(err).ToNot(HaveOccurred()) + defer mcpClient.Close() + + By("Listing tools from VirtualMCPServer") + listRequest := mcp.ListToolsRequest{} + tools, err := mcpClient.Client.ListTools(mcpClient.Ctx, listRequest) + Expect(err).ToNot(HaveOccurred()) + + By("Verifying the referenced composite tool carries its explicit annotations") + var foundComposite bool + for _, tool := range tools.Tools { + if tool.Name == compositeToolName { + foundComposite = true + GinkgoWriter.Printf(" Tool: %s annotations=%+v\n", tool.Name, tool.Annotations) + Expect(tool.Annotations.Title).To(Equal("Referenced Echo Twice")) + Expect(tool.Annotations.DestructiveHint).ToNot(BeNil()) + Expect(*tool.Annotations.DestructiveHint).To(BeTrue()) + // The yardstick echo backend does not declare readOnlyHint, + // so the derived floor is not read-only. + Expect(tool.Annotations.ReadOnlyHint).ToNot(BeNil()) + Expect(*tool.Annotations.ReadOnlyHint).To(BeFalse()) + break + } + } + Expect(foundComposite).To(BeTrue(), "Should find referenced composite tool: %s", compositeToolName) + }) }) Context("when verifying referenced composite tool configuration", func() { diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go index 8dc05d4ce3..e6d11fd242 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" @@ -31,7 +31,9 @@ var _ = Describe("VirtualMCPServer Composite Sequential Workflow", Ordered, func vmcpNodePort int32 // Composite tool names - compositeToolName = "echo_twice" + compositeToolName = "echo_twice" + annotatedToolName = "echo_annotated" + contradictingToolName = "echo_contradicting" ) BeforeAll(func() { @@ -87,6 +89,72 @@ var _ = Describe("VirtualMCPServer Composite Sequential Workflow", Ordered, func }, }, }, + // Composite tool with explicit annotations. An explicit hint may be + // MORE conservative than the derived floor, so this passes through. + { + Name: annotatedToolName, + Description: "Echoes the input with explicit conservative annotations", + Parameters: thvjson.NewMap(map[string]any{ + "type": "object", + "properties": map[string]any{ + "message": map[string]any{ + "type": "string", + "description": "The message to echo", + }, + }, + "required": []any{"message"}, + }), + Annotations: &vmcpconfig.ToolAnnotationsOverride{ + Title: stringPtr("Annotated Echo"), + ReadOnlyHint: boolPtr(false), + DestructiveHint: boolPtr(true), + IdempotentHint: boolPtr(true), + }, + Timeout: vmcpconfig.Duration(30 * time.Second), + Steps: []vmcpconfig.WorkflowStepConfig{ + { + ID: "echo", + Type: "tool", + Tool: fmt.Sprintf("%s.echo", backendName), + Arguments: thvjson.NewMap(map[string]any{ + "input": "{{ .params.message }}", + }), + }, + }, + }, + // Composite tool whose explicit annotations CONTRADICT the derived + // safety floor: readOnlyHint=true claims the tool does not modify + // its environment, but the yardstick echo tool does not declare + // readOnlyHint, so the floor is not read-only. The guardrail drops + // this tool from tools/list at runtime while the server stays Ready. + { + Name: contradictingToolName, + Description: "Contradicting annotations: dropped by the safety-floor guardrail", + Parameters: thvjson.NewMap(map[string]any{ + "type": "object", + "properties": map[string]any{ + "message": map[string]any{ + "type": "string", + "description": "The message to echo", + }, + }, + "required": []any{"message"}, + }), + Annotations: &vmcpconfig.ToolAnnotationsOverride{ + ReadOnlyHint: boolPtr(true), + }, + Timeout: vmcpconfig.Duration(30 * time.Second), + Steps: []vmcpconfig.WorkflowStepConfig{ + { + ID: "echo", + Type: "tool", + Tool: fmt.Sprintf("%s.echo", backendName), + Arguments: thvjson.NewMap(map[string]any{ + "input": "{{ .params.message }}", + }), + }, + }, + }, }, }), v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ @@ -195,6 +263,77 @@ var _ = Describe("VirtualMCPServer Composite Sequential Workflow", Ordered, func // Second echo: echoes the result of first echo GinkgoWriter.Printf("Composite tool result: %+v\n", result.Content) }) + + It("should advertise annotations on composite tools", func() { + By("Creating and initializing MCP client for VirtualMCPServer") + mcpClient, err := CreateInitializedMCPClient(vmcpNodePort, "toolhive-composite-test", 30*time.Second) + Expect(err).ToNot(HaveOccurred()) + defer mcpClient.Close() + + By("Listing tools from VirtualMCPServer") + listRequest := mcp.ListToolsRequest{} + tools, err := mcpClient.Client.ListTools(mcpClient.Ctx, listRequest) + Expect(err).ToNot(HaveOccurred()) + + toolAnnotations := make(map[string]mcp.ToolAnnotation, len(tools.Tools)) + for _, tool := range tools.Tools { + toolAnnotations[tool.Name] = tool.Annotations + GinkgoWriter.Printf(" Tool: %s annotations=%+v\n", tool.Name, tool.Annotations) + } + + By("Verifying derived annotations on the composite tool") + // The yardstick echo backend declares NO annotations, so derivation is + // fail-closed: echo_twice (no explicit annotations) advertises the + // conservative floor — readOnlyHint=false, destructiveHint=true, + // openWorldHint=true. + compositeAnn, found := toolAnnotations[compositeToolName] + Expect(found).To(BeTrue(), "Should find composite tool: %s", compositeToolName) + Expect(compositeAnn.ReadOnlyHint).ToNot(BeNil(), "Composite tool should advertise a derived readOnlyHint") + Expect(*compositeAnn.ReadOnlyHint).To(BeFalse(), + "Derived readOnlyHint should be false: the yardstick echo tool does not declare itself read-only") + Expect(compositeAnn.DestructiveHint).ToNot(BeNil(), "Composite tool should advertise a derived destructiveHint") + Expect(*compositeAnn.DestructiveHint).To(BeTrue(), + "Derived destructiveHint should be true: a step whose annotations are unknown taints the floor") + Expect(compositeAnn.OpenWorldHint).ToNot(BeNil(), "Composite tool should advertise a derived openWorldHint") + Expect(*compositeAnn.OpenWorldHint).To(BeTrue(), + "Derived openWorldHint should be true: a step whose annotations are unknown taints the floor") + + By("Verifying explicit annotations pass through when more conservative than the floor") + annotatedAnn, found := toolAnnotations[annotatedToolName] + Expect(found).To(BeTrue(), "Should find explicitly annotated composite tool: %s", annotatedToolName) + Expect(annotatedAnn.Title).To(Equal("Annotated Echo")) + Expect(annotatedAnn.ReadOnlyHint).ToNot(BeNil()) + Expect(*annotatedAnn.ReadOnlyHint).To(BeFalse()) + Expect(annotatedAnn.DestructiveHint).ToNot(BeNil()) + Expect(*annotatedAnn.DestructiveHint).To(BeTrue()) + Expect(annotatedAnn.IdempotentHint).ToNot(BeNil()) + Expect(*annotatedAnn.IdempotentHint).To(BeTrue()) + + By("Verifying the contradicting composite tool is dropped") + _, found = toolAnnotations[contradictingToolName] + Expect(found).To(BeFalse(), + "Composite tool with annotations contradicting the safety floor should be dropped: %s", contradictingToolName) + + By("Verifying the contradicting composite tool is also uncallable") + callRequest := mcp.CallToolRequest{} + callRequest.Params.Name = contradictingToolName + callRequest.Params.Arguments = map[string]any{"message": "should-not-run"} + _, err = mcpClient.Client.CallTool(mcpClient.Ctx, callRequest) + Expect(err).To(HaveOccurred(), + "CallTool on a dropped contradicting composite must fail; advertised equals executed") + }) + + It("should stay Ready while dropping the contradicting composite tool", func() { + vmcpServer := &mcpv1beta1.VirtualMCPServer{} + err := k8sClient.Get(ctx, types.NamespacedName{ + Name: vmcpServerName, + Namespace: testNamespace, + }, vmcpServer) + Expect(err).ToNot(HaveOccurred()) + + Expect(vmcpServer.Status.Phase).To(Equal(mcpv1beta1.VirtualMCPServerPhaseReady), + "VirtualMCPServer should stay Ready when the annotation guardrail drops a composite tool") + }) }) Context("when verifying composite tool configuration", func() { @@ -206,7 +345,7 @@ var _ = Describe("VirtualMCPServer Composite Sequential Workflow", Ordered, func }, vmcpServer) Expect(err).ToNot(HaveOccurred()) - Expect(vmcpServer.Spec.Config.CompositeTools).To(HaveLen(1)) + Expect(vmcpServer.Spec.Config.CompositeTools).ToNot(BeEmpty()) compositeTool := vmcpServer.Spec.Config.CompositeTools[0] Expect(compositeTool.Name).To(Equal(compositeToolName)) @@ -230,3 +369,9 @@ var _ = Describe("VirtualMCPServer Composite Sequential Workflow", Ordered, func }) }) }) + +// boolPtr returns a pointer to b. Used for optional *bool annotation fields. +func boolPtr(b bool) *bool { return &b } + +// stringPtr returns a pointer to s. Used for optional *string annotation fields. +func stringPtr(s string) *string { return &s } diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go index 743ae9fb1b..c2bb5fbd37 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcp "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go index b1b9e064b9..b94c39110b 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go @@ -12,13 +12,13 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_backends_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_backends_test.go new file mode 100644 index 0000000000..16ad841eca --- /dev/null +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_backends_test.go @@ -0,0 +1,243 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package virtualmcp + +import ( + "fmt" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" + "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" + coremcp "github.com/stacklok/toolhive/pkg/mcp" + vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// This spec closes #5979: unlike virtualmcp_discovered_mode_test.go (Legacy +// backends only, gofetch/osv), it mixes a Legacy backend (gofetch) with a +// Modern (2026-07-28, stateless) backend -- yardstick-server run with +// TRANSPORT=streamable-http, STATELESS=true, BACKEND_MODE=echo, the env +// combination confirmed at acceptance_tests/dual_era_k8s_test.go:151-155 -- +// in the same discovered-mode aggregation. That's what actually exercises +// probeRevision/modernEnumerate/dispatch resolving BOTH eras behind one vMCP; +// the existing spec never resolves a Modern revision at all. +// +// OUT OF SCOPE (deliberate): reclassify / mid-run revision-flip coverage. +// #5979 lists it as a stretch goal; it's hard to trigger deterministically +// against a live backend, and pkg/vmcp/client/reclassify_test.go already +// covers it deterministically as a unit test. +// +// mcpRevision is only observable through the LIVE health monitor state +// (pkg/vmcp/health/status.go), on both surfaces this spec checks +// (VirtualMCPServer.status.discoveredBackends[] and the /status endpoint). +// Health monitoring is NOT the operator's default -- it's gated on +// Operational.FailureHandling.HealthCheckInterval > 0 (pkg/vmcp/cli/serve.go) +// -- so the VirtualMCPServer config below sets FailureHandling explicitly; +// without it both mcpRevision assertions would time out empty. +var _ = Describe("VirtualMCPServer Dual-Era Backends", Ordered, func() { + var ( + testNamespace = "default" + mcpGroupName = "test-dual-era-group" + vmcpServerName = "test-vmcp-dual-era" + legacyBackend = "backend-dual-era-legacy" + modernBackend = "backend-dual-era-modern" + echoToolName = "echo" + timeout = 3 * time.Minute + pollingInterval = 2 * time.Second + vmcpNodePort int32 + ) + + // assertRevisions checks that got[modernBackend] and got[legacyBackend] hold + // the expected MCP revision, returning a descriptive error otherwise. Shared + // by the CRD-status and /status assertions below, which differ only in how + // they build the map. + assertRevisions := func(got map[string]string) error { + if rev := got[modernBackend]; rev != coremcp.MCPVersionModern { + return fmt.Errorf("modern backend %s: want revision %q, got %q (all: %v)", + modernBackend, coremcp.MCPVersionModern, rev, got) + } + if rev := got[legacyBackend]; rev != coremcp.MCPVersionLegacy { + return fmt.Errorf("legacy backend %s: want revision %q, got %q (all: %v)", + legacyBackend, coremcp.MCPVersionLegacy, rev, got) + } + return nil + } + + BeforeAll(func() { + By("Creating MCPGroup") + CreateMCPGroupAndWait(ctx, k8sClient, mcpGroupName, testNamespace, + "Test MCP Group for VirtualMCP dual-era backend E2E tests", timeout, pollingInterval) + + By("Creating Legacy (gofetch) and Modern (yardstick, stateless echo) backends in parallel") + CreateMultipleMCPServersInParallel(ctx, k8sClient, []BackendConfig{ + { + Name: legacyBackend, + Namespace: testNamespace, + GroupRef: mcpGroupName, + Image: images.GofetchServerImage, + }, + { + Name: modernBackend, + Namespace: testNamespace, + GroupRef: mcpGroupName, + // Same image as the Legacy backends above: since #6004 it is a single + // go-sdk v1.7 build serving both eras, and STATELESS is what decides + // which. In session mode it negotiates down and vMCP classifies it + // Legacy from server/discover's supportedVersions, so a forgotten + // STATELESS here would silently produce an all-Legacy backend set -- + // which is why the mcpRevision assertions below check both eras. + Image: images.YardstickServerImage, + Env: []mcpv1beta1.EnvVar{ + {Name: "STATELESS", Value: "true"}, // Modern (2026-07-28) capable backend + {Name: "BACKEND_MODE", Value: "echo"}, + }, + }, + }, timeout, pollingInterval) + + By("Creating VirtualMCPServer in discovered mode over the mixed backend set") + vmcpServer := v1beta1test.NewVirtualMCPServer(vmcpServerName, testNamespace, + v1beta1test.WithVMCPGroupRef(mcpGroupName), + v1beta1test.WithVMCPConfig(vmcpconfig.Config{ + Group: mcpGroupName, + Aggregation: &vmcpconfig.AggregationConfig{ + ConflictResolution: "prefix", // avoid tool-name collisions across backends + }, + // Health monitoring is what resolves and caches MCPRevision; it is + // opt-in (see the header comment), so it must be turned on explicitly + // for the mcpRevision assertions below to ever see a non-empty value. + Operational: &vmcpconfig.OperationalConfig{ + FailureHandling: &vmcpconfig.FailureHandlingConfig{ + HealthCheckInterval: vmcpconfig.Duration(10 * time.Second), + // Must be set explicitly and be strictly LESS than the interval. + // Supplying the failureHandling object at all makes kubebuilder + // apply its siblings' CRD defaults, and healthCheckTimeout + // defaults to 10s -- equal to the interval above, which + // pkg/vmcp/config/validator.go rejects ("must be less than + // healthCheckInterval to prevent checks from queuing up"). The + // vmcp pod then crash-loops on config validation and the + // VirtualMCPServer never leaves Pending, so every assertion here + // fails on a readiness timeout with no hint of the real cause. + HealthCheckTimeout: vmcpconfig.Duration(2 * time.Second), + UnhealthyThreshold: 3, + StatusReportingInterval: vmcpconfig.Duration(5 * time.Second), + }, + }, + }), + v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{ + Type: "anonymous", + }), + v1beta1test.MutateVMCP(func(v *mcpv1beta1.VirtualMCPServer) { + v.Spec.ServiceType = "NodePort" + }), + ) + Expect(k8sClient.Create(ctx, vmcpServer)).To(Succeed()) + + By("Waiting for VirtualMCPServer to be ready") + WaitForVirtualMCPServerReady(ctx, k8sClient, vmcpServerName, testNamespace, timeout, pollingInterval) + + By("Waiting for VirtualMCPServer to discover both backends") + WaitForCondition(ctx, k8sClient, vmcpServerName, testNamespace, "BackendsDiscovered", "True", timeout, pollingInterval) + + By("Getting NodePort for VirtualMCPServer") + vmcpNodePort = GetVMCPNodePort(ctx, k8sClient, vmcpServerName, testNamespace, timeout, pollingInterval) + }) + + AfterAll(func() { + By("Cleaning up VirtualMCPServer") + vmcpServer := v1beta1test.NewVirtualMCPServer(vmcpServerName, testNamespace) + if err := k8sClient.Delete(ctx, vmcpServer); err != nil { + GinkgoWriter.Printf("Warning: failed to delete VirtualMCPServer: %v\n", err) + } + + By("Cleaning up backend MCPServers") + for _, name := range []string{legacyBackend, modernBackend} { + backend := &mcpv1beta1.MCPServer{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: testNamespace}, + } + if err := k8sClient.Delete(ctx, backend); err != nil { + GinkgoWriter.Printf("Warning: failed to delete backend %s: %v\n", name, err) + } + } + + By("Cleaning up MCPGroup") + mcpGroup := &mcpv1beta1.MCPGroup{ + ObjectMeta: metav1.ObjectMeta{Name: mcpGroupName, Namespace: testNamespace}, + } + if err := k8sClient.Delete(ctx, mcpGroup); err != nil { + GinkgoWriter.Printf("Warning: failed to delete MCPGroup: %v\n", err) + } + }) + + Context("when aggregating a mixed Legacy+Modern backend set", func() { + It("aggregates tools from both the Legacy and Modern backends", func() { + // This is the real coverage gap the discovered-mode test leaves open: + // it only exercises Legacy backends, so it never resolves a Modern + // revision or calls modernEnumerate. Retrying with fresh sessions + // (WaitForExpectedTools) avoids flaking on first-probe races. + WaitForExpectedTools(vmcpNodePort, "toolhive-e2e-dual-era-list", func(tools []mcp.Tool) error { + var hasFetch, hasEcho bool + for _, t := range tools { + if strings.Contains(t.Name, fetchToolName) { + hasFetch = true + } + if strings.Contains(t.Name, echoToolName) { + hasEcho = true + } + } + if !hasFetch { + return fmt.Errorf("missing Legacy backend's %q tool in aggregated list: %v", fetchToolName, toolNames(tools)) + } + if !hasEcho { + return fmt.Errorf("missing Modern backend's %q tool in aggregated list: %v", echoToolName, toolNames(tools)) + } + return nil + }, timeout) + }) + + It("routes a tool call to the Modern backend and it succeeds", func() { + // Covers dispatch's Modern call path with real traffic, not just a + // readiness/count check (per .claude/rules/testing.md's E2E Test + // Coverage rule). + TestToolListingAndCall(vmcpNodePort, "toolhive-e2e-dual-era-call", echoToolName, "dual-era-modern-check") + }) + }) + + Context("when surfacing per-backend MCP revision", func() { + It("surfaces the correct mcpRevision in VirtualMCPServer.status.discoveredBackends for each backend", func() { + Eventually(func() error { + status, err := GetVirtualMCPServerStatus(ctx, k8sClient, vmcpServerName, testNamespace) + if err != nil { + return fmt.Errorf("failed to get VirtualMCPServer status: %w", err) + } + + got := make(map[string]string, len(status.DiscoveredBackends)) + for _, b := range status.DiscoveredBackends { + got[b.Name] = b.MCPRevision + } + return assertRevisions(got) + }, timeout, pollingInterval).Should(Succeed()) + }) + + It("surfaces the correct mcp_revision in the /status endpoint's backends[] for each backend", func() { + Eventually(func() error { + status, err := GetVMCPStatus(vmcpNodePort) + if err != nil { + return fmt.Errorf("failed to get /status: %w", err) + } + + got := make(map[string]string, len(status.Backends)) + for _, b := range status.Backends { + got[b.Name] = b.MCPRevision + } + return assertRevisions(got) + }, timeout, pollingInterval).Should(Succeed()) + }) + }) +}) diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_redis_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_redis_test.go new file mode 100644 index 0000000000..08ed3e3fb4 --- /dev/null +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_dual_era_redis_test.go @@ -0,0 +1,793 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package virtualmcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + + mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" + "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" + coremcp "github.com/stacklok/toolhive/pkg/mcp" + vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// dualEraRedisKeyPrefix is the vMCP session-storage key prefix used by this +// spec. It must end in ':' (pkg/transport/session/session_data_storage_redis.go +// requires it to avoid key collisions across tenants), and is shared as a +// package-level constant because the concurrency and Redis-outage specs that +// build on this one (added in later steps of this file's plan) key their own +// direct Redis reads/writes off the exact same prefix. +const dualEraRedisKeyPrefix = "thv:vmcp:dualera:" + +// This spec is the Redis-backed counterpart of virtualmcp_dual_era_backends_test.go: +// same mixed Legacy+Modern backend set behind one vMCP, but with session storage +// backed by Redis instead of the in-memory default. It exists as the shared +// fixture for two follow-on specs (not yet added): concurrent Legacy+Modern +// traffic against the Redis-backed session store, and a Redis-outage spec that +// scales Redis to 0 replicas mid-test. +// +// This file is also the end-to-end assertion that Redis session storage does +// NOT close vMCP's Modern capability gate (modernDispatchBlockers, +// pkg/vmcp/server/modern_gate.go): Legacy clients keep their shared, +// reconstructible sessions while Modern clients are served statelessly and +// store nothing. Redis therefore must never be added to the gate's blocker +// list — the Modern legs below fail with -32022 if it is. +// +// Both backends run the SAME yardstick-server image and set BACKEND_MODE=echo +// explicitly so that STATELESS is the ONLY difference between them. (TRANSPORT +// is not set here at all: CreateMultipleMCPServersInParallel (helpers.go +// ~:538-541, ~:562-564) already defaults Transport to "streamable-http" and +// prepends a TRANSPORT env entry itself -- see the Legacy backend's Env +// comment below.) Setting BACKEND_MODE explicitly is deliberate: since #6004 +// yardstick-server is a single go-sdk v1.7 build serving both MCP eras, and a +// forgotten STATELESS would silently collapse this "mixed" backend set into +// two Legacy backends -- every assertion in this file would still pass while +// testing one era twice. +// Using two yardsticks (rather than the gofetch+yardstick pairing in +// virtualmcp_dual_era_backends_test.go) is also deliberate: both eras must +// expose the SAME echo tool so a value can be round-tripped on either client +// edge, which the concurrency spec (Step 3) needs to prove no cross-delivery +// between principals occurs. +// +// IMPORTANT -- this spec must NOT use CreateInitializedMCPClient, WaitForExpectedTools, +// or ToolsContainAll (all in helpers.go / wait_for_tools_helpers.go). Those build a +// go-sdk/mcpcompat client, and go-sdk v1.7's Connect is Modern-first: it sends +// server/discover before anything else. This vMCP serves Modern (Redis session +// storage does not close the capability gate -- see the header above), so +// classifyingHandler routes server/discover to dispatchModern instead of the +// SDK's stateful fallback (pkg/vmcp/server/classification.go), and such a client +// always negotiates Modern and gets no session -- there is no way to make it open +// a Legacy session against this vMCP. Use *e2e.RawMCPClient instead: it pins +// the era explicitly per request (Legacy via e2e.NewLegacyRequest + the +// MCP-Protocol-Version header), which is what actually gets a Legacy session +// from a Modern-serving vMCP. +var _ = Describe("VirtualMCPServer Dual-Era Backends over Redis Session Storage", Ordered, func() { + var ( + timeout = 5 * time.Minute + pollingInterval = 2 * time.Second + + redisName string + mcpGroupName string + legacyBackend string + modernBackend string + vmcpServerName string + vmcpNodePort int32 + rawClient *e2e.RawMCPClient + ) + + // assertRevisions checks that got[legacyBackend] and got[modernBackend] hold + // the expected MCP revision, returning a descriptive error (listing the whole + // map) otherwise. Mirrors assertRevisions in virtualmcp_dual_era_backends_test.go. + assertRevisions := func(got map[string]string) error { + if rev := got[legacyBackend]; rev != coremcp.MCPVersionLegacy { + return fmt.Errorf("legacy backend %s: want revision %q, got %q (all: %v)", + legacyBackend, coremcp.MCPVersionLegacy, rev, got) + } + if rev := got[modernBackend]; rev != coremcp.MCPVersionModern { + return fmt.Errorf("modern backend %s: want revision %q, got %q (all: %v)", + modernBackend, coremcp.MCPVersionModern, rev, got) + } + return nil + } + + BeforeAll(func() { + ts := time.Now().UnixNano() + redisName = fmt.Sprintf("dualera-redis-%d", ts) + mcpGroupName = fmt.Sprintf("dualera-group-%d", ts) + legacyBackend = fmt.Sprintf("dualera-legacy-%d", ts) + modernBackend = fmt.Sprintf("dualera-modern-%d", ts) + vmcpServerName = fmt.Sprintf("dualera-vmcp-%d", ts) + + var err error + rawClient, err = e2e.NewRawMCPClient(20 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + By("Deploying Redis for vMCP session storage") + deployRedis(redisName) + + By("Creating MCPGroup") + CreateMCPGroupAndWait(ctx, k8sClient, mcpGroupName, defaultNamespace, + "Test MCP Group for VirtualMCP dual-era Redis session E2E tests", timeout, pollingInterval) + + By("Creating Legacy and Modern (stateless) yardstick backends in parallel") + CreateMultipleMCPServersInParallel(ctx, k8sClient, []BackendConfig{ + { + Name: legacyBackend, + Namespace: defaultNamespace, + GroupRef: mcpGroupName, + Image: images.YardstickServerImage, + // No STATELESS: this backend negotiates down and vMCP classifies it + // Legacy from server/discover's supportedVersions. BACKEND_MODE is set + // explicitly (see the header comment) so STATELESS is the only + // difference between the two backends below. No TRANSPORT either -- + // CreateMultipleMCPServersInParallel already sets it from + // BackendConfig.Transport (defaulting to "streamable-http"), so adding + // it here would produce a duplicate Env entry. + Env: []mcpv1beta1.EnvVar{ + {Name: "BACKEND_MODE", Value: "echo"}, + }, + }, + { + Name: modernBackend, + Namespace: defaultNamespace, + GroupRef: mcpGroupName, + Image: images.YardstickServerImage, + Env: []mcpv1beta1.EnvVar{ + {Name: "STATELESS", Value: "true"}, // Modern (2026-07-28) capable backend + {Name: "BACKEND_MODE", Value: "echo"}, + }, + }, + }, timeout, pollingInterval) + + By("Creating VirtualMCPServer over the mixed backend set with Redis session storage") + vmcpServer := v1beta1test.NewVirtualMCPServer(vmcpServerName, defaultNamespace, + v1beta1test.WithVMCPGroupRef(mcpGroupName), + v1beta1test.WithVMCPConfig(vmcpconfig.Config{ + Group: mcpGroupName, + Aggregation: &vmcpconfig.AggregationConfig{ + ConflictResolution: "prefix", // tool names become "_echo" + }, + // Health monitoring is what resolves and caches MCPRevision; it is + // opt-in (see the header comment), so it must be turned on explicitly + // for the era-gate check below to ever see a non-empty value. + Operational: &vmcpconfig.OperationalConfig{ + FailureHandling: &vmcpconfig.FailureHandlingConfig{ + HealthCheckInterval: vmcpconfig.Duration(10 * time.Second), + // Must be set explicitly and be strictly LESS than the interval. + // Supplying the failureHandling object at all makes kubebuilder + // apply its siblings' CRD defaults, and healthCheckTimeout + // defaults to 10s -- equal to the interval above, which + // pkg/vmcp/config/validator.go rejects ("must be less than + // healthCheckInterval to prevent checks from queuing up"). The + // vmcp pod then crash-loops on config validation and the + // VirtualMCPServer never leaves Pending, so every assertion here + // fails on a readiness timeout with no hint of the real cause. + HealthCheckTimeout: vmcpconfig.Duration(2 * time.Second), + UnhealthyThreshold: 3, + StatusReportingInterval: vmcpconfig.Duration(5 * time.Second), + }, + }, + }), + v1beta1test.WithVMCPIncomingAuth(&mcpv1beta1.IncomingAuthConfig{Type: "anonymous"}), + v1beta1test.WithVMCPSessionStorage(&mcpv1beta1.SessionStorageConfig{ + Provider: mcpv1beta1.SessionStorageProviderRedis, + Address: fmt.Sprintf("%s.%s.svc.cluster.local:6379", redisName, defaultNamespace), + KeyPrefix: dualEraRedisKeyPrefix, + }), + // One replica, for three reasons: + // 1. Every request takes the local-session Validate branch + // (WithSessionIdManager is wired unconditionally, server.go:570), + // which is what will make the outage spec's 503 deterministic. At two + // replicas, a request landing on the non-owning pod without a cached + // rehydration takes the rehydration path instead, whose transient-error + // branch returns 404 rather than 503 (session_manager.go:437-443 states + // this mapping explicitly). + // 2. GetVMCPStatus and readRedisSessionBackendIDs (Step 3) each observe a + // single pod's view of the world, so the era gate below and Step 3's + // direct Redis read are deterministic only when there is exactly one + // pod to ask. + // 3. SessionAffinity becomes moot -- there is only one pod to be sticky + // to -- so it is deliberately left unset rather than set to a value + // that would do nothing. + v1beta1test.WithVMCPReplicas(1), + v1beta1test.MutateVMCP(func(v *mcpv1beta1.VirtualMCPServer) { + v.Spec.ServiceType = "NodePort" + }), + ) + Expect(k8sClient.Create(ctx, vmcpServer)).To(Succeed()) + + By("Waiting for VirtualMCPServer to be ready") + WaitForVirtualMCPServerReady(ctx, k8sClient, vmcpServerName, defaultNamespace, timeout, pollingInterval) + + // Fail-fast: without this, a dropped/rejected spec.sessionStorage (bad + // Address, bad KeyPrefix) still leaves the vMCP Ready on its in-memory + // fallback, and every assertion below is satisfied byte-for-byte by that + // fallback -- the only symptom would be an unexplained Redis GET miss much + // later. + By("Confirming the operator wired Redis session storage without a warning") + WaitForCondition(ctx, k8sClient, vmcpServerName, defaultNamespace, + mcpv1beta1.ConditionSessionStorageWarning, "False", timeout, pollingInterval) + + By("Waiting for VirtualMCPServer to discover both backends") + WaitForCondition(ctx, k8sClient, vmcpServerName, defaultNamespace, "BackendsDiscovered", "True", timeout, pollingInterval) + + By("Getting NodePort for VirtualMCPServer") + vmcpNodePort = GetVMCPNodePort(ctx, k8sClient, vmcpServerName, defaultNamespace, timeout, pollingInterval) + + // Era gate: this belongs here, as a BeforeAll precondition, rather than as + // its own It. As a separate spec it would duplicate + // virtualmcp_dual_era_backends_test.go:222-253 and trip the Redundancy rule + // in .claude/rules/testing.md; as a precondition it is the same guard with + // no duplicate spec (the CLI tier does exactly this at + // vmcp_dual_era_test.go:108-112). It guards against a forgotten STATELESS + // silently collapsing the mixed backend set to all-Legacy -- every + // assertion below would still pass while testing one era twice. + By("Waiting for vMCP to classify each backend's MCP revision correctly") + Eventually(func() error { + status, err := GetVMCPStatus(vmcpNodePort) + if err != nil { + return fmt.Errorf("failed to get /status: %w", err) + } + got := make(map[string]string, len(status.Backends)) + for _, b := range status.Backends { + got[b.Name] = b.MCPRevision + } + return assertRevisions(got) + }, timeout, pollingInterval).Should(Succeed()) + + By("Waiting for both backends' echo tools to be aggregated") + // Final warm-up: confirms tool aggregation itself is stable (a concern + // independent of the revision classification checked above) before any + // It in this file runs. This CANNOT use WaitForExpectedTools -- see the + // file header comment -- so it is reimplemented here on the raw client: a + // fresh Legacy session on every attempt, then a Legacy tools/list. The + // fresh-session-per-attempt part is the point (not incidental): it is what + // WaitForExpectedTools does too, and for the same reason -- session-scoped + // tool discovery silently skips a backend that is not yet ready, so retrying + // on the SAME session would just re-observe the same stale, incomplete list. + vmcpURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + Eventually(func() error { + sessionID, err := dualEraLegacyInitErr(rawClient, vmcpURL) + if err != nil { + return fmt.Errorf("legacy init: %w", err) + } + names, err := dualEraLegacyListTools(rawClient, vmcpURL, sessionID) + if err != nil { + return fmt.Errorf("legacy tools/list: %w", err) + } + got := make(map[string]bool, len(names)) + for _, n := range names { + got[n] = true + } + for _, want := range []string{legacyBackend + "_echo", modernBackend + "_echo"} { + if !got[want] { + return fmt.Errorf("missing expected tool %q; got %v", want, names) + } + } + return nil + }, timeout, pollingInterval).Should(Succeed()) + }) + + AfterAll(func() { + By("Cleaning up VirtualMCPServer") + if err := k8sClient.Delete(ctx, v1beta1test.NewVirtualMCPServer(vmcpServerName, defaultNamespace)); err != nil { + GinkgoWriter.Printf("Warning: failed to delete VirtualMCPServer: %v\n", err) + } + + By("Cleaning up backend MCPServers") + for _, name := range []string{legacyBackend, modernBackend} { + backend := &mcpv1beta1.MCPServer{ + ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: defaultNamespace}, + } + if err := k8sClient.Delete(ctx, backend); err != nil { + GinkgoWriter.Printf("Warning: failed to delete backend %s: %v\n", name, err) + } + } + + By("Cleaning up MCPGroup") + mcpGroup := &mcpv1beta1.MCPGroup{ + ObjectMeta: metav1.ObjectMeta{Name: mcpGroupName, Namespace: defaultNamespace}, + } + if err := k8sClient.Delete(ctx, mcpGroup); err != nil { + GinkgoWriter.Printf("Warning: failed to delete MCPGroup: %v\n", err) + } + + // Redis is torn down last (vMCP before Redis) so the vmcp pod never + // outlives its session store. + By("Cleaning up Redis") + cleanupRedis(redisName) + + Eventually(func() bool { + err := k8sClient.Get(ctx, types.NamespacedName{Name: vmcpServerName, Namespace: defaultNamespace}, &mcpv1beta1.VirtualMCPServer{}) + return apierrors.IsNotFound(err) + }, timeout, pollingInterval).Should(BeTrue()) + }) + + It("round-trips a tools/call on both the Legacy and the Modern client edge", func() { + vmcpURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + + // Prove both client-facing helpers actually reach a working backend, not + // just that the tool names are listed (per .claude/rules/testing.md's E2E + // Test Coverage rule: functional behavior, not just infrastructure state). + // Tool aggregation itself is already proven by the BeforeAll warm-up -- + // repeating a tools/list here would just be the same assertion twice. + // + // Inputs must match yardstick's echo tool's input schema, "^[a-zA-Z0-9]+$" + // (no hyphens/underscores/dots) -- see the doc comments on dualEraLegacyCall + // and dualEraModernCall. + // + // This single-request-per-era shape is a strict subset of what the + // concurrency spec below also exercises, and is kept deliberately rather + // than left as unexplained overlap: it isolates plain Legacy/Modern + // routing on one request each, so a red concurrency spec (many requests + // racing) can be told apart from a red routing spec (this one) -- the + // former points at concurrency handling, the latter at routing itself. + legacySessionID := dualEraLegacyInit(rawClient, vmcpURL) + legacyResp := dualEraLegacyCall(rawClient, vmcpURL, legacySessionID, legacyBackend+"_echo", "dualeraredislegacy") + expectDualEraEcho(legacyResp, "dualeraredislegacy", "") + + modernResp := dualEraModernCall(rawClient, vmcpURL, modernBackend+"_echo", "dualeraredismodern") + expectDualEraEcho(modernResp, "dualeraredismodern", "complete") + }) + + It("serves concurrent Legacy-session and Modern-stateless traffic over a shared Redis session store, storing session state only for the Legacy backend", func() { + vmcpURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + + // One Legacy session, established up front and reused by every Legacy + // request below and by the final survival check at the end. + sessionID := dualEraLegacyInit(rawClient, vmcpURL) + + const perEra = 4 + const rounds = 3 // covers same-era interleaving too, not just one batch + + // dualEraConcurrentResult pairs a response (and any Send error) with + // the input THIS goroutine sent, so the post-join assertion can check + // that this request got back its own echo rather than another + // goroutine's -- that mismatch is the cross-delivery failure this spec + // exists to catch. err is asserted before resp is touched (see the + // post-join loop below): the goroutines below call only the + // error-returning dualEra*CallErr helpers and never assert on the + // result themselves, so a nil resp from a failed send must never + // reach expectDualEraEcho on the main goroutine. + type dualEraConcurrentResult struct { + resp *e2e.RawResponse + err error + input string + legacy bool + } + + for round := 0; round < rounds; round++ { + total := 2 * perEra + // Pre-sized slice indexed by goroutine number -- no shared map, + // no append from goroutines. + results := make([]dualEraConcurrentResult, total) + start := make(chan struct{}) + var wg sync.WaitGroup + + // The goroutines below call ONLY dualEraLegacyCallErr / + // dualEraModernCallErr -- the error-returning variants -- and + // store (resp, err) in results without asserting on either. A + // Fail() (what a failing Expect does) inside a goroutine with no + // GinkgoRecover on that goroutine's stack panics the whole test + // process, not just this spec, so no Expect/Fail of ours may run + // here. Every substantive assertion -- the error check first, + // then expectDualEraEcho, then the session-id header check -- + // happens below, after the join, on the main goroutine. + for i := 0; i < perEra; i++ { + idx, input := i, fmt.Sprintf("round%dlegacy%d", round, i) + wg.Add(1) + go func() { + defer wg.Done() + <-start + resp, err := dualEraLegacyCallErr(rawClient, vmcpURL, sessionID, legacyBackend+"_echo", input) + results[idx] = dualEraConcurrentResult{resp: resp, err: err, input: input, legacy: true} + }() + } + for i := 0; i < perEra; i++ { + idx, input := perEra+i, fmt.Sprintf("round%dmodern%d", round, i) + wg.Add(1) + go func() { + defer wg.Done() + <-start + resp, err := dualEraModernCallErr(rawClient, vmcpURL, modernBackend+"_echo", input) + results[idx] = dualEraConcurrentResult{resp: resp, err: err, input: input, legacy: false} + }() + } + + close(start) // release every goroutine together, so requests are genuinely in flight simultaneously + + // Timeout barrier per .claude/rules/testing.md's "Concurrent Tests: + // Always Add Timeouts to Blocking Barriers" -- a panicking + // goroutine or a routing regression that drops a response must + // not hang the whole suite. + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(60 * time.Second): + Fail(fmt.Sprintf("round %d: timed out waiting for %d concurrent requests", round, total)) + } + + for _, r := range results { + Expect(r.err).ToNot(HaveOccurred(), "input %q", r.input) + if r.legacy { + expectDualEraEcho(r.resp, r.input, "") + continue + } + expectDualEraEcho(r.resp, r.input, "complete") + Expect(r.resp.Headers.Get(e2e.HeaderMCPSessionID)).To(BeEmpty(), + "Modern edge must not mint a session") + } + } + + By("Confirming Redis stores session state only for the Legacy backend") + backendIDs, err := readRedisSessionBackendIDs(redisName, dualEraRedisKeyPrefix, sessionID) + Expect(err).ToNot(HaveOccurred()) + // True regardless of the health monitor's MCPRevision cache state: + // - The actual state here is warm: this file's BeforeAll era gate already + // forced classification of both backends before any It runs. On the + // warm path, initOneBackend's isKnownModern check short-circuits and + // returns (nil, true) for the Modern backend without ever attempting a + // connection to it (pkg/vmcp/session/factory.go:281-287). + // WithRevisionLookup is wired whenever the backend client implements + // vmcp.RevisionReporter, which the default client does + // (pkg/vmcp/cli/serve.go:361-363). + // - And even off that path (a cold cache), populateBackendMetadata only + // writes MetadataKeyBackendSessionPrefix+workloadID when the backend's + // connection reports a non-empty session ID (factory.go:414-425: + // "if sessID := r.conn.SessionID(); sessID != "" { ... }"). A stateless + // Modern backend never mints an Mcp-Session-Id, so no + // vmcp.backend.session. key can be written either way. + // Backend ID is the MCPServer resource name + // (pkg/vmcp/k8s/backend_reconciler.go:41, pkg/vmcp/registry.go:363-368). + Expect(backendIDs).To(HaveKey("vmcp.backend.session." + legacyBackend)) + Expect(backendIDs).ToNot(HaveKey("vmcp.backend.session." + modernBackend)) + + By("Confirming the Legacy session still works after the direct Redis read above") + // Rounds 1 and 2 already reused sessionID after round 0's concurrent + // Modern traffic, so "Modern traffic did not disturb the Legacy session" + // was already established before this call runs. What this call actually + // adds: the session still works after the direct Redis GET immediately + // above, i.e. that read did not consume or invalidate anything, at the + // very end of the spec. + finalResp := dualEraLegacyCall(rawClient, vmcpURL, sessionID, legacyBackend+"_echo", "aftersurvive") + expectDualEraEcho(finalResp, "aftersurvive", "") + }) + + // This It is deliberately last: it destroys Redis's data (scaling it to 0 + // wipes every key, since the fixture has no persistence) and must not run + // before the two specs above, which depend on Redis being up throughout. + // It restores Redis itself partway through (see the "Restoring Redis" step + // below), so AfterAll does not need to -- and must not skip Redis cleanup + // on the assumption this It always reaches that step. + It("returns 503 on the Legacy session path during a Redis outage while the Modern path stays unaffected, then recovers with a fresh session", func() { + vmcpURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + vmcpLabels := map[string]string{ + "app.kubernetes.io/name": "virtualmcpserver", + "app.kubernetes.io/instance": vmcpServerName, + } + + By("Recording the vmcp pod's restart count before the outage") + podsBefore, err := GetVirtualMCPServerPods(ctx, k8sClient, vmcpServerName, defaultNamespace) + Expect(err).ToNot(HaveOccurred()) + Expect(podsBefore.Items).To(HaveLen(1), "replicas=1 fixture: exactly one vmcp pod expected") + restartCountBefore := vmcpContainerRestartCount(&podsBefore.Items[0]) + + // Load-bearing: the 503 asserted below requires a session established + // BEFORE Redis goes down -- see that assertion's comment. + By("Establishing a Legacy session and one successful call before the outage") + sessionID := dualEraLegacyInit(rawClient, vmcpURL) + preOutageResp := dualEraLegacyCall(rawClient, vmcpURL, sessionID, legacyBackend+"_echo", "beforeoutage") + expectDualEraEcho(preOutageResp, "beforeoutage", "") + + By("Scaling Redis to 0 replicas to simulate a session-store outage") + Eventually(func() error { return scaleRedis(redisName, 0) }, timeout, pollingInterval).Should(Succeed()) + Eventually(func() (int, error) { + podList := &corev1.PodList{} + if err := k8sClient.List(ctx, podList, client.InNamespace(defaultNamespace), + client.MatchingLabels{"app": redisName}); err != nil { + return -1, err + } + return len(podList.Items), nil + }, timeout, pollingInterval).Should(Equal(0), "Redis pod should be gone after scaling to 0") + + By("Confirming the Legacy path 503s on the pre-outage session") + // With replicas=1 (see the BeforeAll comment on WithVMCPReplicas), + // every request takes sessionmanager.Manager.Validate's local-session + // branch, which returns (false, err) on Redis's transient error; the + // transport maps that to 503 (session_manager.go:444-471, + // toolhive-core@v0.0.34/mcpcompat/server/transports.go:462-496). A + // client that instead runs "initialize" DURING the outage gets a 200 + // with NO Mcp-Session-Id -- Manager.Generate returns "" after two + // failed storage.Create attempts (session_manager.go:238-267), and + // go-sdk hands out an unusable ephemeral session, never a 503. That + // quirk is NOT asserted here; this pins the intended contract for a + // session that predates the outage, not that edge case. + // + // Wrapped in Eventually: there is a short window right after the + // Deployment scales to 0 where the Redis Service may still have stale + // endpoints. dualEraLegacyCallErr (rather than dualEraLegacyCall) lets + // a transient Send error be retried instead of aborting the spec, + // mirroring how dualEraLegacyInitErr exists alongside dualEraLegacyInit. + var outageResp *e2e.RawResponse + Eventually(func() (int, error) { + resp, err := dualEraLegacyCallErr(rawClient, vmcpURL, sessionID, legacyBackend+"_echo", "duringoutage") + if err != nil { + return 0, err + } + outageResp = resp + return resp.StatusCode, nil + }, timeout, pollingInterval).Should(Equal(http.StatusServiceUnavailable)) + // The 503 body is written with http.Error, so it is plain text, not a + // JSON-RPC envelope -- assert on status + body substring, not resp.Error. + Expect(string(outageResp.Body)).To(ContainSubstring("session validation unavailable")) + + By("Confirming the Modern path is unaffected while Redis is down") + // dispatchModern (pkg/vmcp/server/modern_dispatch.go) holds no + // session-store reference at all, which is why both calls below keep + // succeeding through the whole outage. The second call is the bridge + // cell -- vMCP synthesizes a per-request backend session for the + // Legacy backend reached via the Modern client edge -- and it + // succeeding during an outage is the sharpest proof that the Modern + // path touches the session store zero times. Consistently (not one + // call) so a single lucky success cannot pass. + Consistently(func() error { + modernResp := dualEraModernCall(rawClient, vmcpURL, modernBackend+"_echo", "duringoutage") + if err := dualEraEchoErr(modernResp, "duringoutage", "complete"); err != nil { + return fmt.Errorf("modern backend: %w", err) + } + bridgeResp := dualEraModernCall(rawClient, vmcpURL, legacyBackend+"_echo", "duringoutage") + if err := dualEraEchoErr(bridgeResp, "duringoutage", "complete"); err != nil { + return fmt.Errorf("bridge cell (legacy backend via modern edge): %w", err) + } + return nil + }, 5*time.Second, pollingInterval).Should(Succeed()) + + By("Confirming the outage did not restart or un-ready the vmcp pod") + // Liveness /health is an unconditional 200 and readiness /readyz gates + // only on the watcher cache sync (pkg/vmcp/server/server.go:915,:955), + // so a Redis outage that starts AFTER the pod is already running must + // not restart it. This is distinct from Redis being down at pod + // START, which WOULD fail startup -- NewRedisSessionDataStorage pings + // Redis on construction -- which is exactly why the ordering in this + // spec (Redis up when the pod started, taken down mid-run) matters. + podsAfter, err := GetVirtualMCPServerPods(ctx, k8sClient, vmcpServerName, defaultNamespace) + Expect(err).ToNot(HaveOccurred()) + Expect(podsAfter.Items).To(HaveLen(1)) + Expect(vmcpContainerRestartCount(&podsAfter.Items[0])).To(Equal(restartCountBefore)) + WaitForPodsReady(ctx, k8sClient, defaultNamespace, vmcpLabels, timeout, pollingInterval) + + By("Restoring Redis") + Eventually(func() error { return scaleRedis(redisName, 1) }, timeout, pollingInterval).Should(Succeed()) + WaitForPodsReady(ctx, k8sClient, defaultNamespace, map[string]string{"app": redisName}, timeout, pollingInterval) + + By("Confirming a fresh Legacy session works after recovery") + var newSessionID string + Eventually(func() error { + id, err := dualEraLegacyInitErr(rawClient, vmcpURL) + if err != nil { + return err + } + newSessionID = id + return nil + }, timeout, pollingInterval).Should(Succeed()) + Expect(newSessionID).ToNot(BeEmpty()) + recoveredResp := dualEraLegacyCall(rawClient, vmcpURL, newSessionID, legacyBackend+"_echo", "afterrecovery") + expectDualEraEcho(recoveredResp, "afterrecovery", "") + + By("Confirming the stale pre-outage session was cleared by the outage, not merely surviving it") + // Redis here is a plain in-memory Deployment with no persistence, so + // scaling it to 0 earlier wiped every key; Validate now sees + // redis.Nil -> ErrSessionNotFound -> (true, nil) -> 404 + // (session_manager.go:444-471, transports.go:462-496). This is NOT + // asserting session survival -- it proves the outage really cleared + // state, so the recovery above is genuine functional recovery, not a + // session that was never gone. + var staleResp *e2e.RawResponse + Eventually(func() (int, error) { + resp, err := dualEraLegacyCallErr(rawClient, vmcpURL, sessionID, legacyBackend+"_echo", "stalesession") + if err != nil { + return 0, err + } + staleResp = resp + return resp.StatusCode, nil + }, timeout, pollingInterval).Should(Equal(http.StatusNotFound)) + Expect(string(staleResp.Body)).To(ContainSubstring("Session terminated")) + }) +}) + +// dualEraLegacyInit performs the Legacy (2025-11-25) initialize handshake +// against url on client: an "initialize" request followed by the +// notifications/initialized notification the spec requires a client to send +// immediately after, before any other request on the session. Mirrors +// legacyInitialize in vmcp_dual_era_test.go. client is taken as a parameter +// rather than captured, for consistency with the other dualEra* helpers below +// (this one is only ever called from the main spec goroutine, not from the +// concurrency spec's fan-out). +// +// Aborts the spec via Expect on failure. Use dualEraLegacyInitErr instead +// inside an Eventually retry loop (e.g. the BeforeAll warm-up), where a +// transient failure must be retried rather than failing the spec outright. +func dualEraLegacyInit(mcpClient *e2e.RawMCPClient, url string) string { + GinkgoHelper() + sessionID, err := dualEraLegacyInitErr(mcpClient, url) + Expect(err).ToNot(HaveOccurred()) + return sessionID +} + +// dualEraLegacyInitErr is the error-returning variant of dualEraLegacyInit, +// for use inside an Eventually retry loop. It delegates to the shared +// legacySessionInit primitive (legacy_session_helpers_test.go). +func dualEraLegacyInitErr(mcpClient *e2e.RawMCPClient, url string) (string, error) { + return legacySessionInit(mcpClient, url, "dual-era-redis-legacy-client", nil) +} + +// dualEraLegacyListTools sends a Legacy tools/list on sessionID and returns +// the aggregated tool names. It delegates to the shared legacySessionListTools +// primitive (legacy_session_helpers_test.go). +// +// Returns an error rather than using Expect/Ginkgo assertions so it is safe +// to call from inside an Eventually retry loop (e.g. the BeforeAll warm-up). +func dualEraLegacyListTools(mcpClient *e2e.RawMCPClient, url, sessionID string) ([]string, error) { + return legacySessionListTools(mcpClient, url, sessionID, nil) +} + +// dualEraLegacyCall sends a Legacy tools/call for toolName on sessionID, +// carrying the MCP-Protocol-Version header the 2025-11-25 spec requires on +// every post-initialize request. Mirrors legacyToolCall in +// vmcp_dual_era_test.go. +// +// input MUST match yardstick's echo tool's input schema, "^[a-zA-Z0-9]+$" -- +// no hyphens, underscores, or dots. A non-matching input is NOT rejected as a +// JSON-RPC error: yardstick returns it as a successful (HTTP 200, resp.Error +// == nil) result with "isError": true and a validation-failure message in the +// content, so resp.Error being nil does NOT prove the call worked -- callers +// must inspect the result body (see expectDualEraEcho). Build any per-call +// unique nonce as e.g. fmt.Sprintf("round%dreq%d", round, i), never with a +// separator character. +func dualEraLegacyCall(mcpClient *e2e.RawMCPClient, url, sessionID, toolName, input string) *e2e.RawResponse { + GinkgoHelper() + resp, err := dualEraLegacyCallErr(mcpClient, url, sessionID, toolName, input) + Expect(err).ToNot(HaveOccurred()) + return resp +} + +// dualEraLegacyCallErr is the error-returning variant of dualEraLegacyCall, +// for use inside an Eventually/Consistently retry loop (e.g. the Redis-outage +// spec's 503 poll), where a transient Send failure must be retried rather +// than aborting the spec, and inside the concurrency spec's goroutines, where +// a failing Expect would panic the whole test process rather than fail the +// spec (no GinkgoRecover on that goroutine's stack). mcpClient is taken as a +// parameter rather than captured, which is what makes calling this from +// concurrent goroutines safe in the first place (RawMCPClient documents +// itself as safe for concurrent use, mcp_raw_client.go ~:287-290). It +// deliberately does NOT inspect resp.StatusCode/Error -- a non-2xx response +// (503, 404, ...) is a normal, successfully-received HTTP response, not a Go +// error; only an actual transport failure from Send (e.g. connection +// refused) is surfaced as the returned error. +func dualEraLegacyCallErr(mcpClient *e2e.RawMCPClient, url, sessionID, toolName, input string) (*e2e.RawResponse, error) { + return legacySessionCallTool(mcpClient, url, sessionID, toolName, map[string]any{"input": input}, nil) +} + +// dualEraModernCall sends a stateless Modern tools/call for toolName. Mirrors +// modernToolCall in vmcp_dual_era_test.go. +// +// input is subject to the same "^[a-zA-Z0-9]+$" echo-tool constraint as +// dualEraLegacyCall's input -- see that doc comment for the nonce-building +// guidance and the isError-but-200 gotcha. +// +// CRITICAL: neither this nor dualEraLegacyCall/dualEraLegacyInit ever calls +// WithStreamableAccept(). vMCP does not require that Accept header, and +// setting it switches vMCP to an SSE response body that this client's +// plain-JSON parser cannot read (mcp_raw_client.go); e2e.NewModernRequest +// already emits the full conformant Modern wire shape otherwise. +// +// Aborts the spec via Expect on failure. Use dualEraModernCallErr instead +// inside the concurrency spec's goroutines, where a failing Expect would +// panic the process rather than fail the spec. +func dualEraModernCall(mcpClient *e2e.RawMCPClient, url, toolName, input string) *e2e.RawResponse { + GinkgoHelper() + resp, err := dualEraModernCallErr(mcpClient, url, toolName, input) + Expect(err).ToNot(HaveOccurred()) + return resp +} + +// dualEraModernCallErr is the error-returning variant of dualEraModernCall, +// for use inside the concurrency spec's goroutines: a Fail() (what a failing +// Expect does) on a goroutine with no GinkgoRecover on its stack panics the +// whole test process, not just this spec, so those goroutines must call this +// instead and let the main goroutine assert on the returned error after the +// join. Mirrors dualEraLegacyCallErr's split from dualEraLegacyCall. +func dualEraModernCallErr(mcpClient *e2e.RawMCPClient, url, toolName, input string) (*e2e.RawResponse, error) { + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": toolName, + "arguments": map[string]any{"input": input}, + }) + if err != nil { + return nil, fmt.Errorf("build tools/call: %w", err) + } + return mcpClient.Send(context.Background(), url, req) +} + +// expectDualEraEcho asserts a tools/call response succeeded, round-tripped +// wantOutput, and carries the era-correct resultType. resultType is present +// ONLY on a Modern client's own response envelope ("complete"); it is empty +// for a Legacy client, regardless of which era backend actually served the +// call. Mirrors assertBridgedCall + echoResult in vmcp_dual_era_test.go. +// +// This proves a round-trip to *a* backend with the era-correct envelope; it +// does NOT by itself prove WHICH backend answered -- both backends run the +// same yardstick echo tool, so a misrouted call would echo the same value +// back. Backend attribution is proven separately, by the era-gate /status +// check in this file's BeforeAll. +func expectDualEraEcho(resp *e2e.RawResponse, wantOutput, wantResultType string) { + GinkgoHelper() + Expect(dualEraEchoErr(resp, wantOutput, wantResultType)).To(Succeed()) +} + +// dualEraEchoErr is the error-returning counterpart of expectDualEraEcho, for +// use inside a Consistently/Eventually poll (e.g. the Redis-outage spec's +// Modern-path check) where a mismatch must be reported back to the poller +// instead of aborting the whole spec via Expect. +func dualEraEchoErr(resp *e2e.RawResponse, wantOutput, wantResultType string) error { + if resp.Error != nil { + return fmt.Errorf("unexpected JSON-RPC error: %+v", resp.Error) + } + if resp.StatusCode != 200 { + return fmt.Errorf("unexpected status %d, body: %s", resp.StatusCode, resp.Body) + } + + var result struct { + ResultType string `json:"resultType"` + StructuredContent struct { + Output string `json:"output"` + } `json:"structuredContent"` + } + if err := json.Unmarshal(resp.Result, &result); err != nil { + return fmt.Errorf("unmarshal result: %w, raw: %s", err, resp.Result) + } + // Include the raw result in both messages: an empty Output is the common + // failure and says nothing on its own about whether the call reached the + // wrong tool, echoed into a different field, or returned no content at all. + if result.StructuredContent.Output != wantOutput { + return fmt.Errorf("tools/call round-trip did not echo the expected output %q, got %q; raw result: %s", + wantOutput, result.StructuredContent.Output, resp.Result) + } + if result.ResultType != wantResultType { + return fmt.Errorf("unexpected resultType %q, want %q; raw result: %s", + result.ResultType, wantResultType, resp.Result) + } + return nil +} + +// vmcpContainerRestartCount returns the restart count of pod's "vmcp" +// container -- the container name this file's BeforeAll pod-template env +// injection targets. Fails the spec loudly (rather than returning a zero +// value) if no such container status exists, since that would silently +// defeat the Redis-outage spec's restart-count assertion. +func vmcpContainerRestartCount(pod *corev1.Pod) int32 { + GinkgoHelper() + for _, cs := range pod.Status.ContainerStatuses { + if cs.Name == "vmcp" { + return cs.RestartCount + } + } + Fail(fmt.Sprintf("pod %s has no %q container status", pod.Name, "vmcp")) + return 0 +} diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go index 22b191fe30..04259511e1 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go index 190f9f98fc..ae59b5113b 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go @@ -8,13 +8,13 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go index 701502cc90..f3c6611da3 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go @@ -8,11 +8,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go index d2b564ddd1..38d8a74e87 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go @@ -8,13 +8,13 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcp "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go index 7772ed1f13..8e0281da51 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go @@ -8,11 +8,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go index 75a8f1f2ed..96363c4422 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go @@ -12,9 +12,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" dto "github.com/prometheus/client_model/go" @@ -23,10 +20,14 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" "github.com/stacklok/toolhive/pkg/ratelimit" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/test/e2e" "github.com/stacklok/toolhive/test/e2e/images" ) @@ -214,18 +215,42 @@ var _ = ginkgo.Describe("VirtualMCPServer Rate Limiting", ginkgo.Ordered, func() _, err = mcpClient.CallTool(ctx, req) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - result, err := mcpClient.CallTool(ctx, req) + // The mcpcompat client is go-sdk-backed and Modern-first: against a + // Modern-serving vMCP it negotiates 2026-07-28, where a rate-limited + // call surfaces as a real JSON-RPC error object carrying the domain + // code (writeModernCodedError) — not as the Legacy SDK seam's IsError + // tool result with the code smuggled into structuredContent + // (conversion.CodedErrorResult). + _, err = mcpClient.CallTool(ctx, req) + gomega.Expect(err).To(gomega.HaveOccurred(), "second tools/call must be rejected by the per-user limit") + gomega.Expect(err.Error()).To(gomega.ContainSubstring(ratelimit.MessageRateLimited)) + + // Pin the full Modern envelope with an era-pinned raw request (the + // #6051 convention): the rate limiter's 429 with data.retryAfterSeconds + // must survive the Modern dispatch path, which the go-sdk client above + // can only surface as an error string. + ginkgo.By("Verifying the Modern rate-limit error envelope") + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(result.IsError).To(gomega.BeTrue()) - - structured, ok := result.StructuredContent.(map[string]any) - gomega.Expect(ok).To(gomega.BeTrue()) - gomega.Expect(structured["code"]).To(gomega.BeNumerically("==", ratelimit.CodeRateLimited)) - gomega.Expect(structured["message"]).To(gomega.Equal(ratelimit.MessageRateLimited)) - - data, ok := structured["data"].(map[string]any) - gomega.Expect(ok).To(gomega.BeTrue()) - gomega.Expect(data["retryAfterSeconds"]).To(gomega.BeNumerically(">", 0)) + rawReq, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": toolName, + "arguments": map[string]any{"input": "ratelimittest"}, + }) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + rawReq.SetHeader("Authorization", "Bearer "+token) + rawResp, err := rawClient.Send(ctx, fmt.Sprintf("http://localhost:%d/mcp", vmcpLocalPort), rawReq) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + gomega.Expect(rawResp.StatusCode).To(gomega.Equal(http.StatusOK), + "a rate-limited Modern call rides HTTP 200 (429 is in go-sdk's transient retry set), body: %s", rawResp.Body) + gomega.Expect(rawResp.Error).ToNot(gomega.BeNil(), "body: %s", rawResp.Body) + gomega.Expect(rawResp.Error.Code).To(gomega.Equal(ratelimit.CodeRateLimited)) + gomega.Expect(rawResp.Error.Message).To(gomega.Equal(ratelimit.MessageRateLimited)) + var retryData struct { + RetryAfterSeconds float64 `json:"retryAfterSeconds"` + } + gomega.Expect(json.Unmarshal(rawResp.Error.Data, &retryData)).To(gomega.Succeed(), + "error data: %s", rawResp.Error.Data) + gomega.Expect(retryData.RetryAfterSeconds).To(gomega.BeNumerically(">", 0)) ginkgo.By("Scraping non-zero rate limit metrics") gomega.Eventually(func() error { diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go index 97f48641cb..a2785b6ac5 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go @@ -11,14 +11,10 @@ import ( "crypto/x509" "encoding/json" "encoding/pem" - "errors" "fmt" "net/http" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" @@ -29,6 +25,8 @@ import ( mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" + coremcp "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/test/e2e" "github.com/stacklok/toolhive/test/e2e/images" ) @@ -278,18 +276,22 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() gomega.Expect(err).NotTo(gomega.HaveOccurred()) defer cleanupB() - ginkgo.By("Initializing session on pod A") - clientA, err := CreateInitializedMCPClient(int32(localPortA), "e2e-redis-test", 30*time.Second) + // Era-pinned raw client: this spec asserts Legacy-session semantics, so + // it must pin 2025-11-25 per request (see legacy_session_helpers_test.go + // for why the mcpcompat client cannot be used here). + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - defer clientA.Close() + serverURLA := fmt.Sprintf("http://localhost:%d/mcp", localPortA) + serverURLB := fmt.Sprintf("http://localhost:%d/mcp", localPortB) - sessionID := clientA.Client.GetSessionId() - gomega.Expect(sessionID).NotTo(gomega.BeEmpty(), "session ID must be assigned after Initialize") + ginkgo.By("Initializing a Legacy session on pod A") + sessionID, err := legacySessionInit(rawClient, serverURLA, "e2e-redis-test", nil) + gomega.Expect(err).NotTo(gomega.HaveOccurred(), "session ID must be assigned after Initialize") ginkgo.By("Listing tools on pod A") - toolsA, err := clientA.Client.ListTools(clientA.Ctx, mcp.ListToolsRequest{}) + toolsA, err := legacySessionListTools(rawClient, serverURLA, sessionID, nil) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - gomega.Expect(toolsA.Tools).NotTo(gomega.BeEmpty(), "pod A must return tools") + gomega.Expect(toolsA).NotTo(gomega.BeEmpty(), "pod A must return tools") ginkgo.By("Verifying pod A stored backend session IDs in Redis") backendIDsBeforeRestore, err := readRedisSessionBackendIDs(redisName, "thv:vmcp:e2e:", sessionID) @@ -297,22 +299,11 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() gomega.Expect(backendIDsBeforeRestore).NotTo(gomega.BeEmpty(), "pod A must have written per-backend session IDs to Redis so pod B can use them as hints") - ginkgo.By(fmt.Sprintf("Connecting to pod B (%s) with the same session ID", podB.Name)) - serverURLB := fmt.Sprintf("http://localhost:%d/mcp", localPortB) - clientB, err := mcpclient.NewStreamableHttpClient(serverURLB, transport.WithSession(sessionID)) - gomega.Expect(err).NotTo(gomega.HaveOccurred()) - defer func() { _ = clientB.Close() }() - - startCtx, startCancel := context.WithTimeout(context.Background(), 30*time.Second) - defer startCancel() - gomega.Expect(clientB.Start(startCtx)).To(gomega.Succeed()) - - ginkgo.By("Listing tools on pod B using the session from pod A") - listCtx, listCancel := context.WithTimeout(context.Background(), 30*time.Second) - defer listCancel() - toolsB, err := clientB.ListTools(listCtx, mcp.ListToolsRequest{}) - gomega.Expect(err).NotTo(gomega.HaveOccurred()) - gomega.Expect(toolsB.Tools).NotTo(gomega.BeEmpty(), "pod B must return tools via Redis-reconstructed session") + ginkgo.By(fmt.Sprintf("Listing tools on pod B (%s) with the session from pod A", podB.Name)) + toolsB, err := legacySessionListTools(rawClient, serverURLB, sessionID, nil) + gomega.Expect(err).NotTo(gomega.HaveOccurred(), + "pod B must serve the session it has never seen by reconstructing it from Redis") + gomega.Expect(toolsB).NotTo(gomega.BeEmpty(), "pod B must return tools via Redis-reconstructed session") ginkgo.By("Verifying backend session IDs in Redis are the same hints pod B received") backendIDsAfterRestore, err := readRedisSessionBackendIDs(redisName, "thv:vmcp:e2e:", sessionID) @@ -322,7 +313,7 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() "pod B used them as hints and the IDs must be stable") ginkgo.By("Verifying both pods return the same tool count") - gomega.Expect(toolsB.Tools).To(gomega.HaveLen(len(toolsA.Tools)), + gomega.Expect(toolsB).To(gomega.HaveLen(len(toolsA)), "pod B must see same session state as pod A") }) }) @@ -419,25 +410,25 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() ginkgo.By("Getting the NodePort for the VirtualMCPServer") vmcpNodePort := GetVMCPNodePort(ctx, k8sClient, vmcpName, defaultNamespace, timeout, pollInterval) - ginkgo.By("Initializing an MCP session") - mcpClientA, err := CreateInitializedMCPClient(vmcpNodePort, "e2e-redis-restart-test", 30*time.Second) + // Era-pinned raw client: this spec asserts Legacy-session semantics, so + // it must pin 2025-11-25 per request (see legacy_session_helpers_test.go). + // The raw client is also what makes the "pod killed, not clean client + // disconnect" simulation below trivial: it holds no connection state and + // never sends the DELETE /mcp a real client's Close() would, which would + // terminate the session in Redis before the pod restart and defeat the + // purpose of this test. Simply not sending anything models the kill. + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - sessionID := mcpClientA.Client.GetSessionId() - gomega.Expect(sessionID).NotTo(gomega.BeEmpty(), "session ID must be assigned after Initialize") + serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + + ginkgo.By("Initializing a Legacy MCP session") + sessionID, err := legacySessionInit(rawClient, serverURL, "e2e-redis-restart-test", nil) + gomega.Expect(err).NotTo(gomega.HaveOccurred(), "session ID must be assigned after Initialize") ginkgo.By("Verifying tools are available before pod restart") - toolsBefore, err := mcpClientA.Client.ListTools(mcpClientA.Ctx, mcp.ListToolsRequest{}) + toolsBefore, err := legacySessionListTools(rawClient, serverURL, sessionID, nil) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - gomega.Expect(toolsBefore.Tools).NotTo(gomega.BeEmpty()) - - // Cancel context to stop in-flight requests without sending DELETE. - // This simulates the pod being killed, not a clean client disconnect. - // We intentionally skip Client.Close() here because Close() sends a - // DELETE /mcp request that would terminate the session in Redis before - // the pod is restarted — defeating the purpose of this test. - // The transport's background goroutine (started by Start()) selects on - // ctx.Done(), so Cancel() is sufficient to stop it without leaking. - mcpClientA.Cancel() + gomega.Expect(toolsBefore).NotTo(gomega.BeEmpty()) ginkgo.By("Getting the running pod name before restart") var pods []corev1.Pod @@ -489,27 +480,15 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() return checkHTTPHealthReady(vmcpNodePort) }, timeout, pollInterval).Should(gomega.Succeed()) - ginkgo.By("Creating a new client with the SAME session ID") - serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) - newClient, err := mcpclient.NewStreamableHttpClient(serverURL, transport.WithSession(sessionID)) - gomega.Expect(err).NotTo(gomega.HaveOccurred()) - defer func() { _ = newClient.Close() }() - - startCtx, startCancel := context.WithTimeout(context.Background(), 30*time.Second) - defer startCancel() - gomega.Expect(newClient.Start(startCtx)).To(gomega.Succeed()) - - // Send 5 requests to give confidence the fix holds: without Redis-backed - // session reconstruction, each request would fail because the new pod's - // in-memory cache is cold. + // Send 5 requests with the SAME session ID to give confidence the fix + // holds: without Redis-backed session reconstruction, each request would + // fail because the new pod's in-memory cache is cold. ginkgo.By("Sending 5 requests to verify the session is recovered from Redis on the new pod") for i := range 5 { - listCtx, listCancel := context.WithTimeout(context.Background(), 30*time.Second) - toolsAfter, listErr := newClient.ListTools(listCtx, mcp.ListToolsRequest{}) - listCancel() + toolsAfter, listErr := legacySessionListTools(rawClient, serverURL, sessionID, nil) gomega.Expect(listErr).NotTo(gomega.HaveOccurred(), "Request %d/5 should succeed after pod restart — session must be recovered from Redis", i+1) - gomega.Expect(toolsAfter.Tools).To(gomega.HaveLen(len(toolsBefore.Tools)), + gomega.Expect(toolsAfter).To(gomega.HaveLen(len(toolsBefore)), "Request %d/5 should return the same tools as before restart", i+1) } }) @@ -644,42 +623,40 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() gomega.Expect(err).NotTo(gomega.HaveOccurred()) defer cleanupB() - ginkgo.By("Initializing a session on pod A") - clientA, err := CreateInitializedMCPClient(int32(localPortA), "e2e-redis-term-test", 30*time.Second) + // Era-pinned raw client: this spec asserts Legacy-session semantics, so + // it must pin 2025-11-25 per request (see legacy_session_helpers_test.go). + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - sessionID := clientA.Client.GetSessionId() - gomega.Expect(sessionID).NotTo(gomega.BeEmpty(), "session ID must be assigned after Initialize") + serverURLA := fmt.Sprintf("http://localhost:%d/mcp", localPortA) + serverURLB := fmt.Sprintf("http://localhost:%d/mcp", localPortB) + + ginkgo.By("Initializing a Legacy session on pod A") + var sessionID string + gomega.Eventually(func() error { + var initErr error + sessionID, initErr = legacySessionInit(rawClient, serverURLA, "e2e-redis-term-test", nil) + return initErr + }, timeout, pollInterval).Should(gomega.Succeed(), + "pod A should accept session initialization once backend routing is ready") ginkgo.By("Verifying the session is usable on pod A") - toolsA, err := clientA.Client.ListTools(clientA.Ctx, mcp.ListToolsRequest{}) + toolsA, err := legacySessionListTools(rawClient, serverURLA, sessionID, nil) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - gomega.Expect(toolsA.Tools).NotTo(gomega.BeEmpty()) + gomega.Expect(toolsA).NotTo(gomega.BeEmpty()) ginkgo.By(fmt.Sprintf("Reconstructing the session on pod B (%s) via Redis", podB.Name)) - serverURLB := fmt.Sprintf("http://localhost:%d/mcp", localPortB) - clientB, err := mcpclient.NewStreamableHttpClient(serverURLB, transport.WithSession(sessionID)) - gomega.Expect(err).NotTo(gomega.HaveOccurred()) - defer func() { _ = clientB.Close() }() - - startCtx, startCancel := context.WithTimeout(context.Background(), 30*time.Second) - defer startCancel() - gomega.Expect(clientB.Start(startCtx)).To(gomega.Succeed()) - - listCtxB, listCancelB := context.WithTimeout(context.Background(), 30*time.Second) - toolsB, err := clientB.ListTools(listCtxB, mcp.ListToolsRequest{}) - listCancelB() + toolsB, err := legacySessionListTools(rawClient, serverURLB, sessionID, nil) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - gomega.Expect(toolsB.Tools).NotTo(gomega.BeEmpty(), + gomega.Expect(toolsB).NotTo(gomega.BeEmpty(), "pod B should serve the session before termination") - // Terminate the session on pod A by sending DELETE /mcp directly. - // We do this via raw HTTP rather than clientA.Close() to avoid the - // context-cancellation ordering in InitializedMCPClient.Close(). + // Terminate the session on pod A by sending DELETE /mcp directly (the + // raw client only speaks POST JSON-RPC, and DELETE is a plain HTTP + // method, not a JSON-RPC request). ginkgo.By("Terminating the session on pod A via DELETE /mcp") - deleteURL := fmt.Sprintf("http://localhost:%d/mcp", localPortA) deleteCtx, deleteCancel := context.WithTimeout(context.Background(), 10*time.Second) defer deleteCancel() - req, err := http.NewRequestWithContext(deleteCtx, http.MethodDelete, deleteURL, nil) + req, err := http.NewRequestWithContext(deleteCtx, http.MethodDelete, serverURLA, nil) gomega.Expect(err).NotTo(gomega.HaveOccurred()) req.Header.Set("Mcp-Session-Id", sessionID) resp, err := http.DefaultClient.Do(req) @@ -687,22 +664,29 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() _ = resp.Body.Close() gomega.Expect(resp.StatusCode).To(gomega.BeElementOf(http.StatusOK, http.StatusNoContent), "DELETE /mcp should return 200 or 204") - clientA.Cancel() // Pod B's in-memory cache still holds the session, but the ValidatingCache's // checkSession callback will find the key absent in Redis (deleted by the // Terminate call above) and return ErrExpired, triggering lazy eviction. - // The next request from pod B should therefore fail with a session-not-found error. + // The next request from pod B should therefore fail with HTTP 404 + // (session not found — what mcpcompat surfaces as ErrSessionTerminated). ginkgo.By("Verifying pod B rejects subsequent requests for the terminated session") gomega.Eventually(func() error { - listCtx, listCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer listCancel() - _, listErr := clientB.ListTools(listCtx, mcp.ListToolsRequest{}) - if listErr == nil { + listReq, buildErr := e2e.NewLegacyRequest("tools/list", nil) + if buildErr != nil { + return fmt.Errorf("build tools/list: %w", buildErr) + } + listReq.WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, coremcp.MCPVersionLegacy) + listResp, listErr := rawClient.Send(context.Background(), serverURLB, listReq) + if listErr != nil { + return fmt.Errorf("tools/list: %w", listErr) + } + if listResp.StatusCode == http.StatusOK { return fmt.Errorf("expected pod B to reject the terminated session, but request succeeded") } - if !errors.Is(listErr, transport.ErrSessionTerminated) { - return fmt.Errorf("expected ErrSessionTerminated (404), got: %w", listErr) + if listResp.StatusCode != http.StatusNotFound { + return fmt.Errorf("expected 404 (session not found) for the terminated session, got %d, body: %s", + listResp.StatusCode, listResp.Body) } return nil }, timeout, pollInterval).Should(gomega.Succeed(), @@ -1075,36 +1059,23 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() "Authorization": fmt.Sprintf("Bearer %s", asToken), } - ginkgo.By("Initializing MCP session on pod A with embedded AS token") - serverURLPodA := fmt.Sprintf("http://localhost:%d/mcp", localPortA) - clientA, err := mcpclient.NewStreamableHttpClient(serverURLPodA, - transport.WithHTTPHeaders(authHeader)) + // Era-pinned raw client: this spec asserts Legacy-session semantics, so + // it must pin 2025-11-25 per request (see legacy_session_helpers_test.go). + // The embedded AS JWT rides along as an extra header on every request. + rawClient, err := e2e.NewRawMCPClient(60 * time.Second) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - defer func() { _ = clientA.Close() }() - - initCtxA, initCancelA := context.WithTimeout(context.Background(), 60*time.Second) - defer initCancelA() - gomega.Expect(clientA.Start(initCtxA)).To(gomega.Succeed(), - "MCP client should start on pod A — check embedded AS JWT validity and OIDC config") - - initRequest := mcp.InitializeRequest{} - initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION - initRequest.Params.ClientInfo = mcp.Implementation{ - Name: "e2e-upstreamInject-test", - Version: "1.0.0", - } - _, err = clientA.Initialize(initCtxA, initRequest) - gomega.Expect(err).NotTo(gomega.HaveOccurred(), "Initialize on pod A should succeed") + serverURLPodA := fmt.Sprintf("http://localhost:%d/mcp", localPortA) - sessionID := clientA.GetSessionId() - gomega.Expect(sessionID).NotTo(gomega.BeEmpty(), "session ID must be assigned after Initialize") + ginkgo.By("Initializing a Legacy MCP session on pod A with embedded AS token") + sessionID, err := legacySessionInit(rawClient, serverURLPodA, "e2e-upstreamInject-test", authHeader) + gomega.Expect(err).NotTo(gomega.HaveOccurred(), + "Initialize on pod A should succeed and assign a session ID — "+ + "check embedded AS JWT validity and OIDC config") ginkgo.By("Listing tools on pod A to verify the session is functional") - toolsCtxA, toolsCancelA := context.WithTimeout(context.Background(), 30*time.Second) - defer toolsCancelA() - toolsA, err := clientA.ListTools(toolsCtxA, mcp.ListToolsRequest{}) + toolsA, err := legacySessionListTools(rawClient, serverURLPodA, sessionID, authHeader) gomega.Expect(err).NotTo(gomega.HaveOccurred()) - gomega.Expect(toolsA.Tools).NotTo(gomega.BeEmpty(), + gomega.Expect(toolsA).NotTo(gomega.BeEmpty(), "pod A must return tools after Initialize") ginkgo.By("Verifying backend received a Bearer token on pod A's Initialize (upstreamInject fired)") @@ -1122,27 +1093,14 @@ var _ = ginkgo.Describe("VirtualMCPServer Redis-Backed Session Sharing", func() }, 30*time.Second, 2*time.Second).Should(gomega.BeNumerically(">=", 1), "backend must have received at least one Initialize call from pod A") - ginkgo.By(fmt.Sprintf("Connecting to pod B (%s) with the same session ID (triggers RestoreSession)", podB.Name)) + ginkgo.By(fmt.Sprintf("Listing tools on pod B (%s) with the same session ID (triggers RestoreSession)", podB.Name)) serverURLPodB := fmt.Sprintf("http://localhost:%d/mcp", localPortB) - clientB, err := mcpclient.NewStreamableHttpClient(serverURLPodB, - transport.WithSession(sessionID), - transport.WithHTTPHeaders(authHeader)) - gomega.Expect(err).NotTo(gomega.HaveOccurred()) - defer func() { _ = clientB.Close() }() - - startCtxB, startCancelB := context.WithTimeout(context.Background(), 60*time.Second) - defer startCancelB() - gomega.Expect(clientB.Start(startCtxB)).To(gomega.Succeed()) - - ginkgo.By("Listing tools on pod B using the restored session") - listCtxB, listCancelB := context.WithTimeout(context.Background(), 30*time.Second) - defer listCancelB() - toolsB, err := clientB.ListTools(listCtxB, mcp.ListToolsRequest{}) + toolsB, err := legacySessionListTools(rawClient, serverURLPodB, sessionID, authHeader) gomega.Expect(err).NotTo(gomega.HaveOccurred(), "pod B must not return an error — with the context fix, the identity "+ "(including UpstreamTokens) is propagated to RestoreSession so the "+ "backend Initialize is authenticated; without the fix this 401s") - gomega.Expect(toolsB.Tools).NotTo(gomega.BeEmpty(), + gomega.Expect(toolsB).NotTo(gomega.BeEmpty(), "pod B must return tools via the Redis-reconstructed session") ginkgo.By("Verifying backend received a second Initialize call from pod B's RestoreSession") diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go index 936d94be5d..229a65090f 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go @@ -13,9 +13,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" appsv1 "k8s.io/api/apps/v1" @@ -27,6 +24,7 @@ import ( mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" + "github.com/stacklok/toolhive/test/e2e" "github.com/stacklok/toolhive/test/e2e/images" ) @@ -222,61 +220,84 @@ var _ = ginkgo.Describe("VirtualMCPServer Session Management", func() { gomega.Expect(hmacSecretEnvVar.ValueFrom.SecretKeyRef.Key).To(gomega.Equal("hmac-secret")) }) + // These specs assert Legacy-session semantics, so they pin 2025-11-25 per + // request via the raw client (see legacy_session_helpers_test.go for why + // the mcpcompat client cannot be used for session specs). ginkgo.It("Should allow multiple clients to connect with independent sessions", func() { - ginkgo.By("Creating first client") - firstClient, err := CreateInitializedMCPClient(vmcpNodePort, "client-first", 30*time.Second) + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - defer firstClient.Close() - - sessionIDFirst := firstClient.Client.GetSessionId() - gomega.Expect(sessionIDFirst).NotTo(gomega.BeEmpty(), "first client should have a session ID") + serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) - ginkgo.By("Creating second client") - secondClient, err := CreateInitializedMCPClient(vmcpNodePort, "client-second", 30*time.Second) - gomega.Expect(err).ToNot(gomega.HaveOccurred()) - defer secondClient.Close() + ginkgo.By("Establishing the first Legacy session") + sessionIDFirst, err := legacySessionInit(rawClient, serverURL, "client-first", nil) + gomega.Expect(err).ToNot(gomega.HaveOccurred(), "first client should get a session ID") - sessionIDSecond := secondClient.Client.GetSessionId() - gomega.Expect(sessionIDSecond).NotTo(gomega.BeEmpty(), "second client should have a session ID") + ginkgo.By("Establishing the second Legacy session") + sessionIDSecond, err := legacySessionInit(rawClient, serverURL, "client-second", nil) + gomega.Expect(err).ToNot(gomega.HaveOccurred(), "second client should get a session ID") ginkgo.By("Verifying sessions are independent (different IDs)") gomega.Expect(sessionIDFirst).NotTo(gomega.Equal(sessionIDSecond)) - ginkgo.By("Both clients can list tools from the backend") - toolsFirst, err := firstClient.Client.ListTools(firstClient.Ctx, mcp.ListToolsRequest{}) + ginkgo.By("Both sessions can list tools from the backend") + toolsFirst, err := legacySessionListTools(rawClient, serverURL, sessionIDFirst, nil) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(toolsFirst.Tools).NotTo(gomega.BeEmpty()) + gomega.Expect(toolsFirst).NotTo(gomega.BeEmpty()) - toolsSecond, err := secondClient.Client.ListTools(secondClient.Ctx, mcp.ListToolsRequest{}) + toolsSecond, err := legacySessionListTools(rawClient, serverURL, sessionIDSecond, nil) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(toolsSecond.Tools).NotTo(gomega.BeEmpty()) + gomega.Expect(toolsSecond).NotTo(gomega.BeEmpty()) - ginkgo.By("Both clients see the same tool catalog") - gomega.Expect(toolsFirst.Tools).To(gomega.HaveLen(len(toolsSecond.Tools))) + ginkgo.By("Both sessions see the same tool catalog") + gomega.Expect(toolsFirst).To(gomega.HaveLen(len(toolsSecond))) }) ginkgo.It("Should allow a client to make multiple calls on the same session", func() { - client, err := CreateInitializedMCPClient(vmcpNodePort, "multi-call-client", 30*time.Second) + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - defer client.Close() + serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) - sessionID := client.Client.GetSessionId() - gomega.Expect(sessionID).NotTo(gomega.BeEmpty()) + sessionID, err := legacySessionInit(rawClient, serverURL, "multi-call-client", nil) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + // Every request below pins the SAME session ID explicitly, so "the + // session stays stable across calls" is asserted by each call being + // served (a dead or rotated session would 404). ginkgo.By("Listing tools multiple times on the same session") for i := range 3 { - tools, err := client.Client.ListTools(client.Ctx, mcp.ListToolsRequest{}) + tools, err := legacySessionListTools(rawClient, serverURL, sessionID, nil) gomega.Expect(err).ToNot(gomega.HaveOccurred(), "call %d should succeed", i+1) - gomega.Expect(tools.Tools).NotTo(gomega.BeEmpty()) - // Session ID must remain stable across calls - gomega.Expect(client.Client.GetSessionId()).To(gomega.Equal(sessionID)) + gomega.Expect(tools).NotTo(gomega.BeEmpty()) } }) ginkgo.It("Should route tool calls through the session to the backend", func() { - // TestToolListingAndCall discovers the actual (possibly-prefixed) tool name via - // ListTools and calls it with alphanumeric-only input (yardstick requirement). - TestToolListingAndCall(vmcpNodePort, "tool-call-client", "echo", "sessiontest") + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + + sessionID, err := legacySessionInit(rawClient, serverURL, "tool-call-client", nil) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + + // Discover the actual (possibly-prefixed) echo tool name via tools/list, + // then call it on the session with alphanumeric-only input (yardstick + // requirement). + tools, err := legacySessionListTools(rawClient, serverURL, sessionID, nil) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + var echoToolName string + for _, name := range tools { + if strings.Contains(name, "echo") { + echoToolName = name + break + } + } + gomega.Expect(echoToolName).NotTo(gomega.BeEmpty(), "should find an echo tool in the tool list") + + resp, err := legacySessionCallTool(rawClient, serverURL, sessionID, echoToolName, + map[string]any{"input": "sessiontest"}, nil) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) + gomega.Expect(dualEraEchoErr(resp, "sessiontest", "")).To(gomega.Succeed(), + "the session-routed tools/call must echo the input back") }) }) @@ -395,21 +416,27 @@ var _ = ginkgo.Describe("VirtualMCPServer Session Management", func() { return tokenResp.AccessToken } - // newAuthHTTPClient wraps an HTTP client that adds Bearer token to every request. - newAuthHTTPClient := func(token string) *http.Client { - return &http.Client{ - Transport: &authRoundTripper{token: token, transport: http.DefaultTransport}, - Timeout: 30 * time.Second, - } + // bearerHeader builds the extra-header map that carries a JWT on every + // raw Legacy request. Session token binding is Legacy-session semantics, + // so these specs pin 2025-11-25 per request via the raw client (see + // legacy_session_helpers_test.go for why the mcpcompat client cannot be + // used for session specs). + bearerHeader := func(token string) map[string]string { + return map[string]string{"Authorization": "Bearer " + token} } - // connectWithToken initialises an MCP client authenticated with the given JWT. - connectWithToken := func(serverURL, token string) *mcpclient.Client { - httpClient := newAuthHTTPClient(token) - mc := InitializeMCPClientWithRetries(serverURL, 2*time.Minute, - transport.WithHTTPBasicClient(httpClient), - ) - return mc + // legacySessionInitWithRetries establishes a Legacy session, retrying + // while the OIDC middleware and backend routing warm up (the role + // InitializeMCPClientWithRetries played for the mcpcompat client). + legacySessionInitWithRetries := func(rawClient *e2e.RawMCPClient, serverURL, clientName, token string) string { + var sessionID string + gomega.Eventually(func() error { + var initErr error + sessionID, initErr = legacySessionInit(rawClient, serverURL, clientName, bearerHeader(token)) + return initErr + }, 2*time.Minute, pollInterval).Should(gomega.Succeed(), + "vMCP should accept an authenticated session initialization once ready") + return sessionID } ginkgo.BeforeAll(func() { @@ -503,13 +530,12 @@ var _ = ginkgo.Describe("VirtualMCPServer Session Management", func() { ginkgo.It("Client using another client's session ID with a different token is rejected", func() { serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) ginkgo.By("Alice establishes a session") aliceToken := getJWTForSubject("alice") - aliceClient := connectWithToken(serverURL, aliceToken) - defer aliceClient.Close() - - aliceSessionID := aliceClient.GetSessionId() + aliceSessionID := legacySessionInitWithRetries(rawClient, serverURL, "alice-client", aliceToken) gomega.Expect(aliceSessionID).NotTo(gomega.BeEmpty()) ginkgo.By("Bob gets a different JWT") @@ -565,56 +591,51 @@ var _ = ginkgo.Describe("VirtualMCPServer Session Management", func() { ginkgo.It("Each client gets their own independent session", func() { serverURL := fmt.Sprintf("http://localhost:%d/mcp", vmcpNodePort) + rawClient, err := e2e.NewRawMCPClient(30 * time.Second) + gomega.Expect(err).ToNot(gomega.HaveOccurred()) - ginkgo.By("Alice and Bob each connect with their own token") + ginkgo.By("Alice and Bob each establish a session with their own token") aliceToken := getJWTForSubject("alice") bobToken := getJWTForSubject("bob") - aliceClient := connectWithToken(serverURL, aliceToken) - defer aliceClient.Close() - - bobClient := connectWithToken(serverURL, bobToken) - defer bobClient.Close() + aliceSessionID := legacySessionInitWithRetries(rawClient, serverURL, "alice-client", aliceToken) + bobSessionID := legacySessionInitWithRetries(rawClient, serverURL, "bob-client", bobToken) ginkgo.By("Verifying they have distinct session IDs") - aliceSessionID := aliceClient.GetSessionId() - bobSessionID := bobClient.GetSessionId() gomega.Expect(aliceSessionID).NotTo(gomega.BeEmpty()) gomega.Expect(bobSessionID).NotTo(gomega.BeEmpty()) gomega.Expect(aliceSessionID).NotTo(gomega.Equal(bobSessionID)) - ginkgo.By("Both clients can independently list tools") - toolsA, err := aliceClient.ListTools(ctx, mcp.ListToolsRequest{}) + ginkgo.By("Both sessions can independently list tools") + toolsA, err := legacySessionListTools(rawClient, serverURL, aliceSessionID, bearerHeader(aliceToken)) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(toolsA.Tools).NotTo(gomega.BeEmpty()) + gomega.Expect(toolsA).NotTo(gomega.BeEmpty()) - toolsB, err := bobClient.ListTools(ctx, mcp.ListToolsRequest{}) + toolsB, err := legacySessionListTools(rawClient, serverURL, bobSessionID, bearerHeader(bobToken)) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(toolsB.Tools).NotTo(gomega.BeEmpty()) + gomega.Expect(toolsB).NotTo(gomega.BeEmpty()) - ginkgo.By("Both clients can independently call tools on their own sessions") + ginkgo.By("Both sessions can independently call tools") // Discover the real tool name (may be prefixed as backendName_echo). // Use alphanumeric-only input — yardstick rejects values with hyphens. var echoToolName string - for _, tool := range toolsA.Tools { - if strings.Contains(tool.Name, "echo") { - echoToolName = tool.Name + for _, name := range toolsA { + if strings.Contains(name, "echo") { + echoToolName = name break } } gomega.Expect(echoToolName).NotTo(gomega.BeEmpty(), "should find an echo tool in the tool list") - callReq := mcp.CallToolRequest{} - callReq.Params.Name = echoToolName - callReq.Params.Arguments = map[string]any{"input": "aliceindependentcall"} - aliceResult, err := aliceClient.CallTool(ctx, callReq) + aliceResp, err := legacySessionCallTool(rawClient, serverURL, aliceSessionID, echoToolName, + map[string]any{"input": "aliceindependentcall"}, bearerHeader(aliceToken)) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(aliceResult.IsError).To(gomega.BeFalse()) + gomega.Expect(dualEraEchoErr(aliceResp, "aliceindependentcall", "")).To(gomega.Succeed()) - callReq.Params.Arguments = map[string]any{"input": "bobindependentcall"} - bobResult, err := bobClient.CallTool(ctx, callReq) + bobResp, err := legacySessionCallTool(rawClient, serverURL, bobSessionID, echoToolName, + map[string]any{"input": "bobindependentcall"}, bearerHeader(bobToken)) gomega.Expect(err).ToNot(gomega.HaveOccurred()) - gomega.Expect(bobResult.IsError).To(gomega.BeFalse()) + gomega.Expect(dualEraEchoErr(bobResp, "bobindependentcall", "")).To(gomega.Succeed()) }) }) diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go index 41ba4a3856..33699557b3 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go index d6121d6920..f17bb90e0b 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go @@ -10,12 +10,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go b/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go index 9699e21fc2..7b612625a6 100644 --- a/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go +++ b/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go @@ -9,11 +9,12 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" + + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // WaitForExpectedTools creates MCP sessions with retry until the validateTools diff --git a/test/e2e/thvignore_test.go b/test/e2e/thvignore_test.go index 66a9cd894f..03616baca4 100644 --- a/test/e2e/thvignore_test.go +++ b/test/e2e/thvignore_test.go @@ -25,7 +25,7 @@ var _ = Describe("THVIgnore E2E Tests", Label("core", "thvignore", "e2e"), func( BeforeEach(func() { config = e2e.NewTestConfig() - serverName = fmt.Sprintf("thvignore-test-%d", GinkgoRandomSeed()) + serverName = e2e.GenerateUniqueServerName("thvignore-test") // Check if thv binary is available err := e2e.CheckTHVBinaryAvailable(config) diff --git a/test/e2e/upgrade_e2e_test.go b/test/e2e/upgrade_e2e_test.go index 5e6e4eb42a..b48cedf2b5 100644 --- a/test/e2e/upgrade_e2e_test.go +++ b/test/e2e/upgrade_e2e_test.go @@ -5,14 +5,15 @@ package e2e_test import ( "encoding/json" + "fmt" "os" "path/filepath" - "strings" "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/stacklok/toolhive/pkg/container/runtime" "github.com/stacklok/toolhive/pkg/runner" "github.com/stacklok/toolhive/pkg/workloads/upgrade" "github.com/stacklok/toolhive/test/e2e" @@ -107,7 +108,7 @@ var _ = Describe("Upgrade Command", Label("core", "upgrade", "e2e"), func() { thvCmd("run", "--name", serverName, rawOSVImage).ExpectSuccess() By("Waiting for the server to be running") - waitForIsolatedMCPServer(thvCmd, serverName, 60*time.Second) + waitForIsolatedMCPServer(thvCmd, serverName, e2e.ServerReadyTimeout()) By("Checking the workload for an available upgrade") stdout, _ := thvCmd("upgrade", "check", serverName).ExpectSuccess() @@ -123,7 +124,7 @@ var _ = Describe("Upgrade Command", Label("core", "upgrade", "e2e"), func() { thvCmd("run", "--name", serverName, rawOSVImage).ExpectSuccess() By("Waiting for the server to be running") - waitForIsolatedMCPServer(thvCmd, serverName, 60*time.Second) + waitForIsolatedMCPServer(thvCmd, serverName, e2e.ServerReadyTimeout()) By("Checking the workload for an available upgrade in JSON format") stdout, _ := thvCmd("upgrade", "check", serverName, "--format", "json").ExpectSuccess() @@ -240,19 +241,43 @@ var _ = Describe("Upgrade Command", Label("core", "upgrade", "e2e"), func() { }) }) -// waitForIsolatedMCPServer polls `thv list` (through the supplied isolated-env -// command builder) until the named workload reports running, or fails the spec -// on timeout. It mirrors e2e.WaitForMCPServer but runs every poll under the same -// isolated config/home/data env as the rest of the spec, so it observes the -// workload created in that isolated state rather than the real ToolHive config. +// waitForIsolatedMCPServer polls until the named workload reports running, or +// fails the spec on timeout. It exists alongside e2e.WaitForMCPServer because it +// runs every poll through the supplied isolated-env command builder, so it +// observes the workload created in that isolated config/home/data state rather +// than the real ToolHive config. The lookup itself is shared, so the status is +// read from the named workload's own record. func waitForIsolatedMCPServer(thvCmd func(args ...string) *e2e.THVCommand, serverName string, timeout time.Duration) { GinkgoHelper() - Eventually(func() bool { - stdout, _, err := thvCmd("list").Run() - if err != nil { + var lastObserved string + ok := func() bool { + workload, err := e2e.FindWorkload(thvCmd, serverName) + switch { + case err != nil: + lastObserved = fmt.Sprintf("thv list failed: %v", err) + return false + case workload == nil: + lastObserved = "not listed" + return false + case workload.Status == runtime.WorkloadStatusRunning: + return true + default: + lastObserved = fmt.Sprintf("status %q (%s)", workload.Status, workload.StatusContext) return false } - return strings.Contains(stdout, serverName) && strings.Contains(stdout, "running") - }, timeout, 1*time.Second).Should(BeTrue(), - "workload %q should be running within %s", serverName, timeout) + } + // On timeout, dump state through the SAME isolated env before failing — + // e2e.DebugServerState would query the real ToolHive config and show + // nothing. Without this dump a CI failure here is undiagnosable: it reads + // as a bare "not running" with no way to tell a stuck startup from an + // errored workload (which is exactly what happened on main run 30335474873). + Eventually(ok, timeout, 1*time.Second).Should(BeTrue(), func() string { + stdout, stderr, err := thvCmd("list", "--all").Run() + GinkgoWriter.Printf("isolated thv list output:\nStdout: %s\nStderr: %s\nError: %v\n", stdout, stderr, err) + logs, stderr, err := thvCmd("logs", serverName).Run() + GinkgoWriter.Printf("Server logs:\n%s\nStderr: %s\nError: %v\n", logs, stderr, err) + proxyLogs, stderr, err := thvCmd("logs", serverName, "--proxy").Run() + GinkgoWriter.Printf("Proxy logs:\n%s\nStderr: %s\nError: %v\n", proxyLogs, stderr, err) + return fmt.Sprintf("workload %q should be running within %s; last observed: %s", serverName, timeout, lastObserved) + }) } diff --git a/test/e2e/vmcp_cli_features_test.go b/test/e2e/vmcp_cli_features_test.go index ea0c2348c7..41dff597fb 100644 --- a/test/e2e/vmcp_cli_features_test.go +++ b/test/e2e/vmcp_cli_features_test.go @@ -16,11 +16,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "gopkg.in/yaml.v3" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvjson "github.com/stacklok/toolhive/pkg/json" "github.com/stacklok/toolhive/pkg/vmcp" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/vmcp_dual_era_helpers_test.go b/test/e2e/vmcp_dual_era_helpers_test.go new file mode 100644 index 0000000000..0a89a55bb4 --- /dev/null +++ b/test/e2e/vmcp_dual_era_helpers_test.go @@ -0,0 +1,213 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "encoding/json" + "fmt" + "net/http" + "os" + "os/exec" + "strconv" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/stacklok/toolhive/test/e2e" + "github.com/stacklok/toolhive/test/e2e/images" +) + +// launchYardstickLegacyOnPort starts a non-stateless (Legacy-only) yardstick +// backend on the given port in the given group. Deliberately structured as a +// near-duplicate of launchYardstickModernOnPort below rather than reusing the +// shared launchYardstickOnPort (vmcp_cli_helpers_test.go), so that everything +// differing between the two backends this suite stands up is visible side by +// side. Both explicitly set BACKEND_MODE=echo rather than one relying on +// yardstick-server's default, leaving STATELESS as the only difference. +func launchYardstickLegacyOnPort(config *e2e.TestConfig, groupName, backendName string, port int) { + portStr := strconv.Itoa(port) + e2e.NewTHVCommand(config, + "run", images.YardstickServerImage, + "--name", backendName, + "--group", groupName, + "--transport", "streamable-http", + "--isolate-network=false", + "--target-port", portStr, + "--env", "TRANSPORT=streamable-http", + "--env", "BACKEND_MODE=echo", + "--", "-port", portStr, "-transport", "streamable-http", + ).ExpectSuccess() +} + +// startEraBackendOnPort launches a backend via launch and waits for it to +// become ready, retrying once with a freshly allocated port if the container +// lost a race binding its host port. allocateVMCPPort only probes that a port +// is free at call time and closes its own listener immediately -- the actual +// bind happens later, asynchronously, inside the container Docker starts for +// the detached `thv run` process, so another workload's container can grab +// the same ephemeral port in between (observed as "address already in use"). +func startEraBackendOnPort(config *e2e.TestConfig, backendName string, port int, launch func(port int)) { + launch(port) + err := e2e.WaitForMCPServer(config, backendName, 120*time.Second) + if err != nil && strings.Contains(err.Error(), "address already in use") { + GinkgoWriter.Printf("%s lost a port-bind race on %d, retrying with a fresh port: %v\n", backendName, port, err) + e2e.NewTHVCommand(config, "rm", backendName).ExpectSuccess() + port = allocateVMCPPort() + launch(port) + err = e2e.WaitForMCPServer(config, backendName, 120*time.Second) + } + Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("yardstick backend %q should become running", backendName)) +} + +// startYardstickLegacyOnPort runs launchYardstickLegacyOnPort and waits for +// the backend to become ready. +func startYardstickLegacyOnPort(config *e2e.TestConfig, groupName, backendName string, port int) { + startEraBackendOnPort(config, backendName, port, func(p int) { + launchYardstickLegacyOnPort(config, groupName, backendName, p) + }) +} + +// launchYardstickModernOnPort starts a stateless (Modern-capable) yardstick +// backend on the given port in the given group. It differs from +// launchYardstickLegacyOnPort above only in STATELESS=true, which is what makes +// yardstick's own discover response advertise 2026-07-28 support (confirmed at +// test/e2e/thv-operator/acceptance_tests/dual_era_k8s_test.go:151-155). +// +// Both backends run the same image, which since #6004 is a single go-sdk v1.7 +// build serving both eras: in session mode it negotiates down and vMCP +// classifies it Legacy from server/discover's supportedVersions. That is why +// the suite's BeforeEach asserts each backend's negotiated revision via +// /status rather than trusting the env. One forgotten STATELESS would +// otherwise leave both backends Legacy, and every cell here would still pass +// while testing the Legacy backend edge twice. +func launchYardstickModernOnPort(config *e2e.TestConfig, groupName, backendName string, port int) { + portStr := strconv.Itoa(port) + e2e.NewTHVCommand(config, + "run", images.YardstickServerImage, + "--name", backendName, + "--group", groupName, + "--transport", "streamable-http", + "--isolate-network=false", + "--target-port", portStr, + "--env", "TRANSPORT=streamable-http", + "--env", "STATELESS=true", + "--env", "BACKEND_MODE=echo", + "--", "-port", portStr, "-transport", "streamable-http", + ).ExpectSuccess() +} + +// startYardstickModernOnPort runs launchYardstickModernOnPort and waits for +// the backend to become ready. +func startYardstickModernOnPort(config *e2e.TestConfig, groupName, backendName string, port int) { + startEraBackendOnPort(config, backendName, port, func(p int) { + launchYardstickModernOnPort(config, groupName, backendName, p) + }) +} + +// healthCheckConfigYAML is appended to a `thv vmcp init`-generated config file +// to turn on health monitoring. This is required for /status to ever report +// mcp_revision: pkg/vmcp/server/status.go's BackendStatus.MCPRevision doc +// notes it "is only observable through the live health state", and +// pkg/vmcp/cli/serve.go only builds a health.MonitorConfig when +// Operational.FailureHandling.HealthCheckInterval > 0. Quick mode (--group) +// generates a config with no Operational section at all and has no CLI flag +// to add one (see generateQuickModeConfig), so this suite uses config-file +// mode instead, purely to make mcp_revision observable. +const healthCheckConfigYAML = ` +operational: + failureHandling: + healthCheckInterval: 1s + unhealthyThreshold: 1 +` + +// appendHealthCheckConfig appends healthCheckConfigYAML to the config file at +// path (a fresh top-level YAML key; thv vmcp init never emits an +// "operational:" section, so there is nothing to overwrite). +func appendHealthCheckConfig(path string) { + GinkgoHelper() + f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0) + Expect(err).ToNot(HaveOccurred()) + defer f.Close() + _, err = f.WriteString(healthCheckConfigYAML) + Expect(err).ToNot(HaveOccurred()) +} + +// startDualEraVMCP starts `thv vmcp serve`, using --config configPath when +// configPath is non-empty, else quick mode (--group groupName). vMCP serves both +// MCP revisions whenever the capability gate is open (modern_gate.go), and these +// specs configure no blockers, so this only assembles the arguments — +// it no longer overrides the inherited environment, and so delegates the process +// start to e2e.StartLongRunningTHVCommand. +func startDualEraVMCP(config *e2e.TestConfig, groupName, configPath string, port int) *exec.Cmd { + GinkgoHelper() + args := []string{"vmcp", "serve", "--port", strconv.Itoa(port)} + if configPath != "" { + args = append(args, "--config", configPath) + } else { + args = append(args, "--group", groupName) + } + return e2e.StartLongRunningTHVCommand(config, args...) +} + +// vmcpStatusResponse decodes the subset of pkg/vmcp/server.StatusResponse +// this suite needs. Decoded independently rather than by importing +// pkg/vmcp/server: that package is a server implementation detail a +// black-box e2e client should not depend on. +type vmcpStatusResponse struct { + Backends []struct { + Name string `json:"name"` + MCPRevision string `json:"mcp_revision"` + } `json:"backends"` +} + +// vmcpStatusURL returns the /status endpoint URL for a vMCP serve process +// listening on the given port. +func vmcpStatusURL(port int) string { + return fmt.Sprintf("http://127.0.0.1:%d/status", port) +} + +// fetchVMCPStatus GETs statusURL and decodes it as a vmcpStatusResponse. +func fetchVMCPStatus(statusURL string) (vmcpStatusResponse, error) { + httpClient := http.Client{Timeout: 5 * time.Second} + resp, err := httpClient.Get(statusURL) //nolint:gosec // test-controlled URL + if err != nil { + return vmcpStatusResponse{}, err + } + defer resp.Body.Close() + + var out vmcpStatusResponse + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return vmcpStatusResponse{}, err + } + return out, nil +} + +// waitForBackendRevision polls statusURL until it reports backendName's +// mcp_revision as wantRevision, or fails the spec after timeout. +// +// This is the only positive proof in this suite that vMCP classified a +// backend's MCP edge correctly. Every other assertion here (tool call +// succeeds, echoes the right value) is satisfied identically whether vMCP +// classified the "Modern" backend as Modern or silently fell back to Legacy +// -- both backends run the same yardstick echo tool, so a misclassification +// would not change any response body. Without this check the 2x2 matrix this +// suite claims to cover could silently collapse to 1x2. +func waitForBackendRevision(statusURL, backendName, wantRevision string) { + GinkgoHelper() + Eventually(func() (string, error) { + status, err := fetchVMCPStatus(statusURL) + if err != nil { + return "", err + } + for _, b := range status.Backends { + if b.Name == backendName { + return b.MCPRevision, nil + } + } + return "", nil + }, 30*time.Second, 500*time.Millisecond).Should(Equal(wantRevision), + "vMCP never classified backend %q as %q via /status", backendName, wantRevision) +} diff --git a/test/e2e/vmcp_dual_era_test.go b/test/e2e/vmcp_dual_era_test.go new file mode 100644 index 0000000000..ec1fc83941 --- /dev/null +++ b/test/e2e/vmcp_dual_era_test.go @@ -0,0 +1,388 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package e2e_test + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/test/e2e" +) + +// This suite is the live e2e tier proving vMCP bridges all four client x +// backend MCP-revision combinations (issue #5912). vMCP is an aggregating +// gateway: it terminates the client connection and re-originates calls to +// each backend, classifying the two edges independently -- the client edge +// per request (Legacy 2025-11-25 vs. Modern 2026-07-28, cached per +// server/discover probe on the backend edge). Existing coverage +// (pkg/vmcp/server/bridge_regression_test.go) pins this at the unit/ +// integration level with mocked backends; this file proves it with two real +// yardstick-server backends and a real thv vmcp serve process. +// +// De-risked live before writing this: a Modern server/discover probe sent +// BY vMCP through the thv proxy fronting a stateless yardstick backend is +// correctly classified Modern (its tool appears in vMCP's aggregated +// tools/list under the "{backend}_" prefix, and a tools/call routed to it +// succeeds) -- confirming the one assumption this whole suite rests on. +// +// Also confirmed live: an ordinary (non-initialize) tools/call response -- +// through vMCP, on either client edge -- never carries Mcp-Session-Id; only +// the Legacy initialize response does. So "session id present for Legacy" +// is asserted on the initialize handshake, not on every subsequent call +// (mirrors the same finding already documented in dual_era_mixing_test.go +// for the single-server transparent proxy). +// +// Deliberate harness choice: no request in this file sets +// "Accept: application/json, text/event-stream" (a MUST on both revisions). +// e2e.RawMCPClient now parses SSE-framed POST responses (sseResponsePayload, +// mcp_raw_client.go), but these specs still omit the header so Legacy +// responses stay plain JSON; flipping them to the conformant Accept is +// tracked in #6104. Until then the bridge is proven here only under a +// non-conformant Accept header. +var _ = Describe("vMCP Dual-Era Bridge", Label("vmcp", "dual-era", "e2e"), Serial, func() { + Context("one Legacy and one Modern backend in the same group", func() { + var ( + config *e2e.TestConfig + groupName string + legacyBackendName string + modernBackendName string + legacyToolName string + modernToolName string + vMCPCmd *exec.Cmd + vMCPPort int + vMCPURL string + rawClient *e2e.RawMCPClient + ) + + BeforeEach(func() { + config = e2e.NewTestConfig() + groupName = e2e.GenerateUniqueServerName("vmcp-dual-era-group") + legacyBackendName = e2e.GenerateUniqueServerName("vmcp-dual-era-legacy") + modernBackendName = e2e.GenerateUniqueServerName("vmcp-dual-era-modern") + legacyToolName = legacyBackendName + "_echo" + modernToolName = modernBackendName + "_echo" + vMCPCmd = nil + vMCPPort = allocateVMCPPort() + vMCPURL = vmcpEndpointURL(vMCPPort) + + err := e2e.CheckTHVBinaryAvailable(config) + Expect(err).ToNot(HaveOccurred(), "thv binary should be available") + + rawClient, err = e2e.NewRawMCPClient(20 * time.Second) + Expect(err).ToNot(HaveOccurred()) + + By("creating a group with one Legacy and one Modern-capable backend") + e2e.NewTHVCommand(config, "group", "create", groupName).ExpectSuccess() + // allocateVMCPPort just allocates any free 127.0.0.1 port; reused here + // for the backends too. + startYardstickLegacyOnPort(config, groupName, legacyBackendName, allocateVMCPPort()) + startYardstickModernOnPort(config, groupName, modernBackendName, allocateVMCPPort()) + + By("generating a vMCP config with health monitoring enabled") + // Config-file mode (not quick mode) solely so health monitoring can be + // turned on -- see appendHealthCheckConfig's doc comment -- which is what + // makes /status's mcp_revision observable below. + tmpDir, err := os.MkdirTemp("", "vmcp-dual-era-config-*") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = os.RemoveAll(tmpDir) }) + configFilePath := filepath.Join(tmpDir, "vmcp.yaml") + initVMCPConfig(config, groupName, configFilePath) + appendHealthCheckConfig(configFilePath) + + By("starting thv vmcp serve") + vMCPCmd = startDualEraVMCP(config, groupName, configFilePath, vMCPPort) + err = e2e.WaitForMCPServerReady(config, vMCPURL, "streamable-http", 60*time.Second) + Expect(err).ToNot(HaveOccurred(), "vMCP server should become ready") + + By("confirming vMCP classified each backend's MCP revision correctly") + // The positive, direct proof that this suite's 2x2 matrix has not + // silently collapsed to 1x2 -- see waitForBackendRevision's doc comment. + statusURL := vmcpStatusURL(vMCPPort) + waitForBackendRevision(statusURL, legacyBackendName, mcpparser.MCPVersionLegacy) + waitForBackendRevision(statusURL, modernBackendName, mcpparser.MCPVersionModern) + }) + + AfterEach(func() { + stopVMCPProcess(vMCPCmd) + vMCPCmd = nil + + if config.CleanupAfter { + _ = e2e.StopAndRemoveMCPServer(config, legacyBackendName) + _ = e2e.StopAndRemoveMCPServer(config, modernBackendName) + _ = e2e.RemoveGroup(config, groupName) + } + }) + + It("bridges Legacy client to Legacy backend (baseline)", func() { + ctx := context.Background() + sessionID := legacyInitialize(ctx, rawClient, vMCPURL) + + resp := legacyToolCall(ctx, rawClient, vMCPURL, sessionID, legacyToolName, "legacytolegacy") + assertBridgedCall(resp, "legacytolegacy", "") + }) + + It("bridges Modern client to Modern backend (both edges stateless)", func() { + ctx := context.Background() + resp := modernToolCall(ctx, rawClient, vMCPURL, modernToolName, "moderntomodern") + assertBridgedCall(resp, "moderntomodern", "complete") + Expect(resp.Headers.Get(e2e.HeaderMCPSessionID)).To(BeEmpty(), "Modern responses must never carry Mcp-Session-Id") + }) + + It("bridges Modern client to Legacy backend (vMCP synthesizes a per-request backend session)", func() { + ctx := context.Background() + resp := modernToolCall(ctx, rawClient, vMCPURL, legacyToolName, "moderntolegacy") + assertBridgedCall(resp, "moderntolegacy", "complete") + Expect(resp.Headers.Get(e2e.HeaderMCPSessionID)).To(BeEmpty(), "Modern responses must never carry Mcp-Session-Id") + }) + + It("bridges Legacy client to Modern backend (vMCP terminates the client session, re-originates stateless)", func() { + ctx := context.Background() + sessionID := legacyInitialize(ctx, rawClient, vMCPURL) + + resp := legacyToolCall(ctx, rawClient, vMCPURL, sessionID, modernToolName, "legacytomodern") + assertBridgedCall(resp, "legacytomodern", "") + }) + + It("never cross-delivers between two concurrent principals across mismatched backends", func() { + // Principal A: a Legacy session calling the Legacy backend, perEra + // times per round. Principal B: a stateless Modern client calling the + // Modern backend, also perEra times per round. All 2*perEra requests + // fire concurrently, each carrying a unique nonce -- a wrong echoed + // nonce means vMCP delivered one request's response to another. Firing + // several requests per era (not just one) also covers same-era + // same-backend interleaving, not only a cross-era swap. + const rounds = 5 + const perEra = 4 + sessionID := legacyInitialize(context.Background(), rawClient, vMCPURL) + + for round := 0; round < rounds; round++ { + label := fmt.Sprintf("round-%d", round) + results := fireConcurrentBridgedBatch(rawClient, vMCPURL, sessionID, legacyToolName, modernToolName, perEra, round) + assertCorrelated(results, 2*perEra, label) + } + }) + + It("never cross-delivers when both BRIDGE cells are in flight together", func() { + // The spec above races the two SAME-era cells (Legacy client→Legacy + // backend, Modern client→Modern backend). That leaves the cells this + // suite exists for untested under concurrency: here the tool names are + // crossed, so the Legacy session calls the MODERN backend and the + // stateless Modern client calls the LEGACY backend -- both bridge cells, + // concurrently, with distinct principals. + // + // This is the configuration where a cross-delivery would matter most. + // Each bridge cell re-originates through a different egress (the Legacy + // session's calls reach a Modern backend via the stateless modernCall + // path; the Modern client's calls reach a Legacy backend via a + // per-request initialize+call+close), and both are in flight at once. A + // wrong echoed nonce means vMCP crossed a response between two + // principals whose requests took *different* egress paths -- the failure + // the per-request client construction is what prevents, and which + // RFC-0083 D6's shared connection pool is the change most likely to + // break. The pkg/vmcp/server bridge pins cover that at the unit level + // for the Modern→Legacy direction; this is the live counterpart for + // both directions at once. + const rounds = 5 + const perEra = 4 + sessionID := legacyInitialize(context.Background(), rawClient, vMCPURL) + + for round := 0; round < rounds; round++ { + label := fmt.Sprintf("bridge-round-%d", round) + // Crossed on purpose: legacy CLIENT calls the modern tool, modern + // CLIENT calls the legacy tool. + results := fireConcurrentBridgedBatch(rawClient, vMCPURL, sessionID, modernToolName, legacyToolName, perEra, round) + assertCorrelated(results, 2*perEra, label) + } + }) + }) +}) + +// legacyInitialize sends a Legacy initialize request, followed by the +// notifications/initialized notification the 2025-11-25 spec requires a +// client to send immediately after (before any other request on the +// session) -- a real client always sends it, so this suite should too rather +// than only ever exercising the (also-valid, but unrepresentative) omitted +// case. Returns the assigned Mcp-Session-Id, failing the spec if the +// handshake does not succeed, omits the session id, or the notification is +// not accepted. +func legacyInitialize(ctx context.Context, client *e2e.RawMCPClient, url string) string { + GinkgoHelper() + req := e2e.NewLegacyInitializeRequest("dual-era-legacy-client", "1.0") + resp, err := client.Send(ctx, url, req) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(200), "body: %s", resp.Body) + sessionID := resp.Headers.Get(e2e.HeaderMCPSessionID) + Expect(sessionID).ToNot(BeEmpty(), "Legacy initialize must assign a session id") + + notifyReq, err := e2e.NewLegacyRequest("notifications/initialized", nil) + Expect(err).ToNot(HaveOccurred()) + // WithID(nil) omits "id" entirely, making this a true JSON-RPC notification. + notifyReq.WithID(nil).WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, mcpparser.MCPVersionLegacy) + notifyResp, err := client.Send(ctx, url, notifyReq) + Expect(err).ToNot(HaveOccurred()) + Expect(notifyResp.StatusCode).To(Equal(202), "body: %s", notifyResp.Body) + + return sessionID +} + +// legacyToolCall sends a Legacy tools/call for toolName with sessionID, +// carrying the MCP-Protocol-Version header the 2025-11-25 spec requires on +// every post-initialize request. +func legacyToolCall( + ctx context.Context, client *e2e.RawMCPClient, url, sessionID, toolName, input string, +) *e2e.RawResponse { + GinkgoHelper() + req, err := e2e.NewLegacyRequest("tools/call", map[string]any{ + "name": toolName, + "arguments": map[string]any{"input": input}, + }) + Expect(err).ToNot(HaveOccurred()) + req.WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, mcpparser.MCPVersionLegacy) + resp, err := client.Send(ctx, url, req) + Expect(err).ToNot(HaveOccurred()) + return resp +} + +// modernToolCall sends a stateless Modern tools/call for toolName. +func modernToolCall(ctx context.Context, client *e2e.RawMCPClient, url, toolName, input string) *e2e.RawResponse { + GinkgoHelper() + req, err := e2e.NewModernRequest("tools/call", map[string]any{ + "name": toolName, + "arguments": map[string]any{"input": input}, + }) + Expect(err).ToNot(HaveOccurred()) + resp, err := client.Send(ctx, url, req) + Expect(err).ToNot(HaveOccurred()) + return resp +} + +// fireConcurrentBridgedBatch fires perEra concurrent Legacy tools/call +// requests (principal A, sessionID, legacyClientTool) and perEra concurrent +// Modern tools/call requests (principal B, stateless, modernClientTool) -- +// 2*perEra requests total, all released together via a start channel so they +// are genuinely in flight at once (not sequential per era) -- each carrying a +// round-unique nonce in _meta. Reuses crossDeliveryResult/extractNonce from +// dual_era_proxy_test.go (same package); the caller asserts on the returned +// results with assertCorrelated. +// +// The two tool names decide which CELLS run concurrently, and the caller picks: +// passing (legacyTool, modernTool) races the two SAME-era cells, while passing +// them crossed (modernTool, legacyTool) races the two BRIDGE cells. Both matter +// -- see the callers. +func fireConcurrentBridgedBatch( + client *e2e.RawMCPClient, url, sessionID, legacyClientTool, modernClientTool string, perEra, round int, +) []crossDeliveryResult { + const joinTimeout = 20 * time.Second + ctx, cancel := context.WithTimeout(context.Background(), joinTimeout) + defer cancel() + + total := 2 * perEra + results := make([]crossDeliveryResult, total) + var wg sync.WaitGroup + start := make(chan struct{}) + + fire := func(i int, build func() (*e2e.RawRequest, error)) { + wg.Add(1) + go func() { + defer wg.Done() + nonce := fmt.Sprintf("round-%d-req-%d", round, i) + <-start + + req, err := build() + if err != nil { + results[i] = crossDeliveryResult{nonce: nonce, err: err} + return + } + req.SetMeta("nonce", nonce) + + reqStart := time.Now() + resp, sendErr := client.Send(ctx, url, req) + elapsed := time.Since(reqStart) + if sendErr != nil { + results[i] = crossDeliveryResult{nonce: nonce, err: sendErr, elapsed: elapsed} + return + } + results[i] = crossDeliveryResult{ + nonce: nonce, + statusCode: resp.StatusCode, + elapsed: elapsed, + rpcErr: resp.Error, + gotNonce: extractNonce(resp), + } + }() + } + + for i := 0; i < perEra; i++ { + fire(i, func() (*e2e.RawRequest, error) { + req, err := e2e.NewLegacyRequest("tools/call", map[string]any{ + "name": legacyClientTool, + "arguments": map[string]any{"input": "concurrencycheck"}, + }) + if err != nil { + return nil, err + } + return req.WithSessionID(sessionID).SetHeader(e2e.HeaderMCPProtocolVersion, mcpparser.MCPVersionLegacy), nil + }) + } + for i := 0; i < perEra; i++ { + fire(perEra+i, func() (*e2e.RawRequest, error) { + return e2e.NewModernRequest("tools/call", map[string]any{ + "name": modernClientTool, + "arguments": map[string]any{"input": "concurrencycheck"}, + }) + }) + } + + close(start) // release every goroutine together + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(joinTimeout): + Fail(fmt.Sprintf("round %d: timed out waiting for %d concurrent requests", round, total)) + } + + return results +} + +// echoResult is the decoded shape of yardstick's echo tool result. +// ResultType is only present in a Modern client's response envelope +// (empty for Legacy) -- see this file's doc comment. +type echoResult struct { + ResultType string `json:"resultType"` + StructuredContent struct { + Output string `json:"output"` + } `json:"structuredContent"` +} + +// assertBridgedCall asserts a tools/call response succeeded, round-tripped +// wantOutput, and carries the era-correct resultType (only a Modern client's +// own response is wrapped with resultType, regardless of which era backend +// served it). This proves a round-trip to *a* backend with the era-correct +// envelope; it does NOT by itself prove WHICH backend answered -- both +// backends run the same yardstick echo tool, so a misrouted call would echo +// the same value back. Backend attribution is proven separately, by the +// waitForBackendRevision /status check in this context's BeforeEach. +func assertBridgedCall(resp *e2e.RawResponse, wantOutput, wantResultType string) { + GinkgoHelper() + Expect(resp.Error).To(BeNil(), "unexpected JSON-RPC error: %+v", resp.Error) + Expect(resp.StatusCode).To(Equal(200), "body: %s", resp.Body) + + var result echoResult + Expect(json.Unmarshal(resp.Result, &result)).To(Succeed(), "result: %s", resp.Result) + Expect(result.StructuredContent.Output).To(Equal(wantOutput), + "tools/call round-trip did not echo the expected output") + Expect(result.ResultType).To(Equal(wantResultType)) +} diff --git a/test/e2e/vmcp_infra_features_test.go b/test/e2e/vmcp_infra_features_test.go index 3e5c7415b1..8450402324 100644 --- a/test/e2e/vmcp_infra_features_test.go +++ b/test/e2e/vmcp_infra_features_test.go @@ -104,12 +104,12 @@ var _ = Describe("vMCP infra features", Label("vmcp", "e2e", "infra"), func() { BeforeEach(func() { fx.setup("vmcp-auth-oidc", "vmcp-auth-oidc-*") - oidcPort = allocateVMCPPort() var err error - oidcServer, err = e2e.NewOIDCMockServer(oidcPort, "test-client", "test-secret", + oidcServer, err = e2e.NewOIDCMockServer(0, "test-client", "test-secret", e2e.WithClientAudience("vmcp-e2e-test"), ) Expect(err).ToNot(HaveOccurred()) + oidcPort = oidcServer.Port() Expect(oidcServer.Start()).To(Succeed()) discoveryURL := fmt.Sprintf("http://localhost:%d/.well-known/openid-configuration", oidcPort) Eventually(func() error { diff --git a/test/e2e/vmcp_optimizer_test.go b/test/e2e/vmcp_optimizer_test.go index e715c8413f..a4d13f7ddf 100644 --- a/test/e2e/vmcp_optimizer_test.go +++ b/test/e2e/vmcp_optimizer_test.go @@ -8,7 +8,8 @@ // This file adds deeper coverage for RFC THV-0059 Phase 4: // // - Tier-1 find→call round-trip: verifies that find_tool locates the yardstick -// echo tool by description and call_tool invokes it end-to-end. +// echo tool by description and call_tool invokes it end-to-end, both with +// tool_name alongside parameters and with it nested inside them. // - Tier-1 two-backend with conflict resolution: verifies that optimizer // discovers tools from both backends when prefix conflict resolution is active. // - Tier-1 composite + optimizer: verifies that composite tools are indexed by @@ -27,10 +28,10 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvjson "github.com/stacklok/toolhive/pkg/json" vmcp "github.com/stacklok/toolhive/pkg/vmcp" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" @@ -93,7 +94,7 @@ var _ = Describe("vMCP optimizer", Label("vmcp", "e2e", "optimizer"), func() { BeforeEach(func() { fx.setup("vmcp-opt-roundtrip", "") }) AfterEach(func() { fx.teardown() }) - It("find_tool locates the echo tool and call_tool invokes it end-to-end", func() { + It("find_tool locates the echo tool and call_tool invokes it, flat and nested", func() { By("starting thv vmcp serve with --optimizer") fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg, "vmcp", "serve", @@ -133,6 +134,20 @@ var _ = Describe("vMCP optimizer", Label("vmcp", "e2e", "optimizer"), func() { Expect(callResult.Content).ToNot(BeEmpty(), "call_tool must return content") Expect(mcp.GetTextFromContent(callResult.Content[0])).To(ContainSubstring("hellooptimizer"), "echo tool must return the input message") + + By("invoking the same tool with tool_name nested inside parameters") + nestedResult, err := mcpClient.CallTool(ctx, "call_tool", map[string]any{ + "parameters": map[string]any{ + "tool_name": echoToolName, + "input": "hellonested", + }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(nestedResult.IsError).To(BeFalse(), + "call_tool must accept a tool_name nested inside parameters") + Expect(nestedResult.Content).ToNot(BeEmpty(), "call_tool must return content") + Expect(mcp.GetTextFromContent(nestedResult.Content[0])).To(ContainSubstring("hellonested"), + "the nested tool_name must be hoisted and the remaining parameters passed through") }) }) diff --git a/test/integration/vmcp/helpers/backend.go b/test/integration/vmcp/helpers/backend.go index c5594b625b..bef9376419 100644 --- a/test/integration/vmcp/helpers/backend.go +++ b/test/integration/vmcp/helpers/backend.go @@ -10,8 +10,8 @@ import ( "net/http/httptest" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" ) // BackendTool defines a tool for MCP backend servers. @@ -180,7 +180,7 @@ func WithCaptureHeaders() BackendServerOption { } } -// CreateBackendServer creates an MCP backend server using the mark3labs/mcp-go SDK. +// CreateBackendServer creates an MCP backend server using the stacklok/toolhive-core/mcpcompat SDK. // It returns an *httptest.Server ready to accept streamable-HTTP connections. // // The server automatically registers all provided tools with proper closure handling diff --git a/test/integration/vmcp/helpers/helpers_test.go b/test/integration/vmcp/helpers/helpers_test.go index 3d186c0ee5..28c19b3bb3 100644 --- a/test/integration/vmcp/helpers/helpers_test.go +++ b/test/integration/vmcp/helpers/helpers_test.go @@ -6,8 +6,9 @@ package helpers import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // TestGetToolNames tests the GetToolNames helper function. diff --git a/test/integration/vmcp/helpers/mcp_client.go b/test/integration/vmcp/helpers/mcp_client.go index 75167cbdf6..4bbfdc072c 100644 --- a/test/integration/vmcp/helpers/mcp_client.go +++ b/test/integration/vmcp/helpers/mcp_client.go @@ -10,14 +10,15 @@ import ( "sync" "testing" - "github.com/mark3labs/mcp-go/client" - mcptransport "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) -// MCPClient wraps the mark3labs MCP client with test-friendly methods. +// MCPClient wraps the mcpcompat MCP client with test-friendly methods. // It automatically handles initialization and provides semantic assertion helpers // that integrate with Go's testing.TB interface. // @@ -173,12 +174,41 @@ func (c *MCPClient) Close() error { func (c *MCPClient) ListTools(ctx context.Context) *mcp.ListToolsResult { c.tb.Helper() + // Follow MCP pagination cursors so the returned set is the complete tool + // list across every page, not just the first. This is spec-conforming and + // harmless to callers that only assert containment. + var allTools []mcp.Tool + var cursor mcp.Cursor + for { + tools, next := c.ListToolsPage(ctx, cursor) + allTools = append(allTools, tools...) + if next == "" { + break + } + cursor = next + } + + c.tb.Logf("Listed %d tools from MCP server", len(allTools)) + return &mcp.ListToolsResult{Tools: allTools} +} + +// ListToolsPage lists a single page of tools starting at the given cursor +// (pass an empty cursor for the first page). It returns that page's tools and +// the NextCursor to fetch the following page (empty when no more pages remain). +// +// Use this when a test needs to assert page-level behavior (e.g. that the first +// vMCP page is bounded by the server page size and carries a NextCursor, +// proving pagination actually occurred). Most tests should use ListTools, which +// accumulates across pages. +func (c *MCPClient) ListToolsPage(ctx context.Context, cursor mcp.Cursor) ([]mcp.Tool, mcp.Cursor) { + c.tb.Helper() + request := mcp.ListToolsRequest{} + request.Params.Cursor = cursor result, err := c.client.ListTools(ctx, request) require.NoError(c.tb, err, "failed to list tools") - c.tb.Logf("Listed %d tools from MCP server", len(result.Tools)) - return result + return result.Tools, result.NextCursor } // CallTool calls the specified tool with the given arguments. diff --git a/test/integration/vmcp/helpers/vmcp_server.go b/test/integration/vmcp/helpers/vmcp_server.go index aecb2f4cad..3416b7c65c 100644 --- a/test/integration/vmcp/helpers/vmcp_server.go +++ b/test/integration/vmcp/helpers/vmcp_server.go @@ -5,7 +5,6 @@ package helpers import ( "context" - "net" "testing" "time" @@ -75,6 +74,7 @@ type vmcpServerConfig struct { workflowDefs map[string]*composer.WorkflowDefinition telemetryProvider *telemetry.Provider passthroughHeaders []string + sessionTTL time.Duration } // WithPrefixConflictResolution configures prefix-based conflict resolution. @@ -112,25 +112,27 @@ func WithPassthroughHeaders(headers ...string) VMCPServerOption { } } -// getFreePort returns an available TCP port on localhost. -func getFreePort(tb testing.TB) int { - tb.Helper() - - listener, err := net.Listen("tcp", "127.0.0.1:0") - require.NoError(tb, err, "failed to get free port") - defer func() { _ = listener.Close() }() - - addr, ok := listener.Addr().(*net.TCPAddr) - if !ok { - tb.Fatalf("failed to get TCP address from listener") +// WithSessionTTL overrides the server's session time-to-live (default 30m). +// A short TTL is useful for sliding-TTL / eviction regression tests. A zero +// value leaves the default in place. +func WithSessionTTL(ttl time.Duration) VMCPServerOption { + return func(c *vmcpServerConfig) { + c.sessionTTL = ttl } - return addr.Port } // NewVMCPServer creates a vMCP server for testing using the Serve path (core.New + // Serve). The server is automatically started and ready when this function returns. // Use functional options to customize behavior. // +// The server binds an OS-assigned port (ServerConfig.Port is 0) and holds that +// listener for its lifetime, so callers MUST read the address back from the +// returned server's Address() rather than assuming a port. Address() reports the +// bound port only once the server is ready, which this function waits for before +// returning. Do not pre-select a port here by binding and closing a probe +// listener: that releases the port, and any of the parallel tests in this suite +// can claim it before the server rebinds (#6034). +// // The Serve path is used (rather than the legacy server.New path) so that // integration tests exercise the production code path that routes tool/resource // calls through core.VMCP. This is required for testing per-request @@ -187,7 +189,7 @@ func NewVMCPServer( Name: "test-vmcp", Version: "1.0.0", Host: "127.0.0.1", - Port: getFreePort(tb), + Port: 0, // OS-assigned; see the note on NewVMCPServer EndpointPath: "/mcp", SessionTTL: 30 * time.Minute, AuthMiddleware: auth.AnonymousMiddleware, @@ -198,6 +200,9 @@ func NewVMCPServer( Base: sessionFactory, }, } + if config.sessionTTL > 0 { + serverCfg.SessionTTL = config.sessionTTL + } vmcpServer, err := vmcpserver.Serve(ctx, coreVMCP, serverCfg) require.NoError(tb, err, "failed to create vMCP server") diff --git a/test/integration/vmcp/pagination_regression_test.go b/test/integration/vmcp/pagination_regression_test.go new file mode 100644 index 0000000000..5dc9e83387 --- /dev/null +++ b/test/integration/vmcp/pagination_regression_test.go @@ -0,0 +1,104 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// TestRegression_Over1000Tools_CompleteSetReceived is a regression anchor for +// the vMCP pagination gap: a backend exposing more than the MCP page size +// (mcpcompat paginates at DefaultPageSize=1000) must surface its complete tool +// set across pagination cursors. +// +// vMCP follows list-pagination cursors in two independent code paths, and this +// test gates both: +// - the session-projection path (pkg/vmcp/session/.../mcp_session.go), which +// the aggregated ListTools count exercises; and +// - the aggregator/routing-table path (pkg/vmcp/client), which CallTool on a +// tool living beyond the first page exercises. +// +// Before the fix, each path issued a single ListTools with no cursor loop, so a +// paginating backend returned only its first page and every tool beyond it was +// silently dropped. +func TestRegression_Over1000Tools_CompleteSetReceived(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + // Generate 1000+ tools on a single backend to exceed the MCP page size. + // mcpcompat's default page size is 1000, so 1100 tools force pagination + // into a first full page plus a short tail page. + const toolCount = 1100 + tools := make([]helpers.BackendTool, 0, toolCount) + for i := 0; i < toolCount; i++ { + name := fmt.Sprintf("tool_%04d", i) + tools = append(tools, helpers.NewBackendTool( + name, + fmt.Sprintf("backend tool number %d", i), + func(_ context.Context, _ map[string]any) string { + return `{"ok": true}` + }, + )) + } + + backend := helpers.CreateBackendServer(t, tools, helpers.WithBackendName("many-tools")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("many-tools", + helpers.WithURL(backend.URL+"/mcp"), + helpers.WithMetadata("group", "pagination-test"), + ), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + ) + + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + client := helpers.NewMCPClient(ctx, t, vmcpURL) + defer client.Close() + + // The first vMCP page must be bounded by the server page size and carry a + // non-empty NextCursor. This proves pagination actually occurred rather + // than the backend returning everything in one page (which would make the + // count assertions below pass even without cursor following). + firstPage, firstCursor := client.ListToolsPage(ctx, "") + assert.LessOrEqual(t, len(firstPage), 1000, + "the first vMCP page must not exceed the server page size (1000)") + assert.NotEmpty(t, firstCursor, + "a non-empty NextCursor must be returned when tools remain beyond the first page") + + // The complete tool set must be received across pages. Assert an exact + // count AND that many DISTINCT names: a plain count could be satisfied by + // page overlap/duplication, which distinct-name counting catches. + result := client.ListTools(ctx) + require.Len(t, result.Tools, toolCount, + "vMCP must surface the complete backend tool set across pagination cursors") + + names := helpers.GetToolNames(result) + distinct := make(map[string]struct{}, len(names)) + for _, n := range names { + distinct[n] = struct{}{} + } + assert.Len(t, distinct, toolCount, + "every tool name must be distinct; duplicates indicate page overlap") + + // Calling the highest-index tool (which lives on the tail page, beyond page + // 1) must succeed. This exercises the aggregator/routing-table path in + // pkg/vmcp/client, which the session-projection ListTools count alone does + // not cover: if that path dropped the tail, the tool would be unroutable. + lastTool := fmt.Sprintf("many-tools_tool_%04d", toolCount-1) + callResult := client.CallTool(ctx, lastTool, map[string]any{}) + helpers.AssertToolCallSuccess(t, callResult) +} diff --git a/test/integration/vmcp/per_session_projection_regression_test.go b/test/integration/vmcp/per_session_projection_regression_test.go new file mode 100644 index 0000000000..6f4ac07141 --- /dev/null +++ b/test/integration/vmcp/per_session_projection_regression_test.go @@ -0,0 +1,424 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sort" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/env" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/authz/authorizers" + "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth/factory" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// principalHeader is the request header the test identity middleware reads to +// determine the authenticated principal. The Cedar authorizer resolves the +// principal from the identity's "sub" claim, so the middleware sets both +// Subject and Claims["sub"] to the header value. +const principalHeader = "X-Test-Principal" + +// newCedarVMCPServer builds a Cedar-authz-backed vMCP server (via server.New, +// mirroring pkg/vmcp/server/authz_integration_test.go) backed by a real MCP +// backend at backendURL. An identity middleware injects a principal derived +// from the X-Test-Principal request header so each session binds to its +// caller and the Cedar admission seam can resolve Client::"". +// +// A priority conflict resolver is used so tool names are NOT prefixed — the +// Cedar policies can name tools by their raw backend names (e.g. +// Tool::"secret-tool"). +func newCedarVMCPServer(t *testing.T, backendURL string, policies ...string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: backendURL, + TransportType: "streamable-http", + } + mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + authReg, err := vmcpauth.NewOutgoingAuthRegistry(context.Background(), &env.OSReader{}) + require.NoError(t, err) + factory := vmcpsession.NewSessionFactory(authReg) + + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + // Priority resolver keeps raw tool names so Cedar policies can name + // Tool::"secret-tool" rather than a prefixed variant. + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + // Identity middleware: derive the principal from the X-Test-Principal + // header so two sessions (alice, bob) bind to different identities. The + // Cedar authorizer resolves the principal from the "sub" claim, so both + // Subject and Claims["sub"] are set to the header value. + identityMiddleware := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + principal := r.Header.Get(principalHeader) + if principal == "" { + principal = "anonymous" + } + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: principal, + Name: principal, + Claims: map[string]any{"sub": principal, "name": principal}, + }} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) + } + + authzCfg, err := authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{Policies: policies, EntitiesJSON: "[]"}, + }) + require.NoError(t, err) + + srv, err := server.New( + context.Background(), + &server.Config{ + Name: "test-vmcp", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: factory, + Aggregator: agg, + AuthMiddleware: identityMiddleware, + Authz: authzCfg, + }, + router.NewSessionRouter(&vmcp.RoutingTable{}), + backendClient, + mockBackendRegistry, + nil, + ) + require.NoError(t, err) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// rawClient is a minimal JSON-RPC-over-HTTP client for the Cedar authz tests. +// It is needed because the helpers.MCPClient.CallTest helper uses require.NoError +// on the SDK error, which would abort the test on a policy-rejected call. This +// client surfaces the raw HTTP response so tests can inspect status codes and +// error bodies directly. +type rawClient struct { + baseURL string + sessionID string + nextID int +} + +func newRawClient(baseURL string) *rawClient { + return &rawClient{baseURL: baseURL, nextID: 1} +} + +func (c *rawClient) initialize(t *testing.T, principal string) { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + c.nextID++ + defer resp.Body.Close() + c.sessionID = resp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, c.sessionID, "initialize response missing Mcp-Session-Id header") + + // Send the initialized notification (no id, no response expected). + notif := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "method": "notifications/initialized", + }, c.sessionID) + notif.Body.Close() +} + +func (c *rawClient) listTools(t *testing.T, principal string) map[string]any { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "tools/list", + "params": map[string]any{}, + }, c.sessionID) + c.nextID++ + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "tools/list failed: %s", string(body)) + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + return parsed +} + +func (c *rawClient) callTool(t *testing.T, principal, name string, args map[string]any) (int, map[string]any) { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "tools/call", + "params": map[string]any{ + "name": name, + "arguments": args, + }, + }, c.sessionID) + c.nextID++ + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) // best-effort; body may be plain text on 404 + return resp.StatusCode, parsed +} + +func (c *rawClient) postMCP(t *testing.T, principal string, body map[string]any, sessionID string) *http.Response { + t.Helper() + rawBody, err := json.Marshal(body) + require.NoError(t, err) + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, c.baseURL+"/mcp", bytes.NewReader(rawBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(principalHeader, principal) + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// toolNamesFromResult extracts the tool names from a tools/list JSON-RPC result. +func toolNamesFromResult(t *testing.T, result map[string]any) []string { + t.Helper() + rawResult, ok := result["result"].(map[string]any) + require.True(t, ok, "response missing result object: %v", result) + toolsRaw, ok := rawResult["tools"].([]any) + require.True(t, ok, "result missing tools array: %v", rawResult) + names := make([]string, 0, len(toolsRaw)) + for _, t2 := range toolsRaw { + tool, ok := t2.(map[string]any) + require.True(t, ok) + name, ok := tool["name"].(string) + require.True(t, ok) + names = append(names, name) + } + sort.Strings(names) + return names +} + +// startAuthzBackend starts a real MCP backend exposing the named tools. Each +// tool echoes back its name and the input so a caller can confirm the backend +// was reached. +func startAuthzBackend(t *testing.T, toolNames ...string) string { + t.Helper() + tools := make([]helpers.BackendTool, 0, len(toolNames)) + for _, name := range toolNames { + tools = append(tools, helpers.NewBackendTool( + name, + "test tool "+name, + func(_ context.Context, args map[string]any) string { + return `{"tool":` + name + `,"args":` + toJSON(args) + `}` + }, + )) + } + backend := helpers.CreateBackendServer(t, tools, helpers.WithBackendName("real-backend")) + t.Cleanup(backend.Close) + return backend.URL + "/mcp" +} + +// toJSON is a tiny marshal helper that never fails the test (returns "null"). +func toJSON(v map[string]any) string { + b, err := json.Marshal(v) + if err != nil { + return "null" + } + return string(b) +} + +// TestRegression_TwoSessions_DifferentAuthz_DifferentToolsList proves the +// core admission seam projects a per-identity tool list: two sessions bound to +// different principals see different advertised tool sets under the same Cedar +// policy. Alice (permitted for secret-tool) sees both secret-tool and +// public-tool; Bob (only permitted for public-tool) sees only public-tool. +func TestRegression_TwoSessions_DifferentAuthz_DifferentToolsList(t *testing.T) { + t.Parallel() + + backendURL := startAuthzBackend(t, "secret-tool", "public-tool") + + ts := newCedarVMCPServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"secret-tool");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"public-tool");`, + ) + + // Alice's session. + alice := newRawClient(ts.URL) + alice.initialize(t, "alice") + aliceTools := toolNamesFromResult(t, alice.listTools(t, "alice")) + assert.Equal(t, []string{"public-tool", "secret-tool"}, aliceTools, + "alice (permitted for secret-tool) must see both tools") + + // Bob's session. + bob := newRawClient(ts.URL) + bob.initialize(t, "bob") + bobTools := toolNamesFromResult(t, bob.listTools(t, "bob")) + assert.Equal(t, []string{"public-tool"}, bobTools, + "bob (not permitted for secret-tool) must see only public-tool") +} + +// TestRegression_FilteredToolUncallableInSessionA_CallableInSessionB proves +// the list-filter/call-deny pairing: a tool filtered out of a session (bob's) +// is rejected as unknown when called, while the same tool is callable in a +// session where it was advertised (alice's). +func TestRegression_FilteredToolUncallableInSessionA_CallableInSessionB(t *testing.T) { + t.Parallel() + + backendURL := startAuthzBackend(t, "secret-tool", "public-tool") + + ts := newCedarVMCPServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"secret-tool");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"public-tool");`, + ) + + // Alice can call secret-tool (it is in her session's advertised set). + alice := newRawClient(ts.URL) + alice.initialize(t, "alice") + require.Contains(t, toolNamesFromResult(t, alice.listTools(t, "alice")), "secret-tool") + + status, result := alice.callTool(t, "alice", "secret-tool", map[string]any{"input": "hi"}) + require.Equal(t, http.StatusOK, status, "alice's permitted call must succeed: %v", result) + assert.False(t, isToolResultError(result), "alice's call to secret-tool must not be an error: %v", result) + + // Bob cannot call secret-tool: it was filtered out of his session, so the + // SDK rejects it as an unknown tool. Poll briefly because the session's + // tool injection runs in a hook that fires after initialize returns. + bob := newRawClient(ts.URL) + bob.initialize(t, "bob") + require.NotContains(t, toolNamesFromResult(t, bob.listTools(t, "bob")), "secret-tool") + + var lastResult map[string]any + require.Eventually(t, func() bool { + var s int + s, lastResult = bob.callTool(t, "bob", "secret-tool", map[string]any{"input": "hi"}) + if s != http.StatusOK { + return true // 404/400 also indicates rejection + } + return isToolResultError(lastResult) + }, 5*time.Second, 50*time.Millisecond, + "bob's call to secret-tool (filtered out) must be rejected") +} + +// TestRegression_SetSessionTools_FixedAtInitialize_NoMidSessionReconciliation +// asserts that a session's advertised tool set is fixed at initialize time +// and is not re-aggregated on each tools/list call. Calling tools/list twice +// on the same session must return the same set. +func TestRegression_SetSessionTools_FixedAtInitialize_NoMidSessionReconciliation(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + backend := helpers.CreateBackendServer(t, []helpers.BackendTool{ + helpers.NewBackendTool("tool-a", "tool A", func(_ context.Context, _ map[string]any) string { + return `{"tool":"a"}` + }), + helpers.NewBackendTool("tool-b", "tool B", func(_ context.Context, _ map[string]any) string { + return `{"tool":"b"}` + }), + }, helpers.WithBackendName("ab-backend")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("ab-backend", helpers.WithURL(backend.URL+"/mcp")), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + ) + + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + client := helpers.NewMCPClient(ctx, t, vmcpURL) + defer client.Close() + + first := helpers.GetToolNames(client.ListTools(ctx)) + second := helpers.GetToolNames(client.ListTools(ctx)) + + sort.Strings(first) + sort.Strings(second) + + assert.Equal(t, []string{"ab-backend_tool-a", "ab-backend_tool-b"}, first, + "initial tools/list must show both tools") + assert.Equal(t, first, second, + "the session's tool set is fixed at initialize; a second tools/list must return the same set") +} + +// isToolResultError returns true when a tools/call JSON-RPC result has +// IsError=true (the SDK rejected the call) or the response carries a +// JSON-RPC error object. +func isToolResultError(result map[string]any) bool { + if _, ok := result["error"]; ok { + return true + } + if r, ok := result["result"].(map[string]any); ok { + if isErr, _ := r["isError"].(bool); isErr { + return true + } + } + return false +} + +// mustParseJSONRPC reads resp.Body and unmarshals it as a JSON-RPC envelope. +// If the body is not valid JSON (e.g. a plain-text 404), it returns an empty +// map so callers can fall back to status-code checks. +func mustParseJSONRPC(t *testing.T, resp *http.Response) map[string]any { + t.Helper() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) // best-effort; body may be plain text on 4xx + return parsed +} + +// Compile-time check that the mcpcompat imports are used (not mark3labs/mcp-go). +var ( + _ = mcp.LATEST_PROTOCOL_VERSION + _ = mcpserver.NewMCPServer +) diff --git a/test/integration/vmcp/sliding_ttl_regression_test.go b/test/integration/vmcp/sliding_ttl_regression_test.go new file mode 100644 index 0000000000..6a5d7a4cd9 --- /dev/null +++ b/test/integration/vmcp/sliding_ttl_regression_test.go @@ -0,0 +1,146 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// TestRegression_SlidingSessionTTL_TrafficKeepsSessionAlive proves the +// session TTL is sliding: a session that receives traffic at least once per +// TTL window stays alive indefinitely, while an idle session is evicted once +// its TTL elapses. +// +// The transport session storage (LocalSessionDataStorage.Load) refreshes the +// last-access timestamp on every read, and the SDK's SessionIdManager.Validate +// is called on every request — so active traffic keeps the session alive while +// a session that goes idle for longer than the TTL is rejected on its next +// request. +// +// This test is timing-sensitive and must NOT run in parallel. +// +//nolint:paralleltest // timing-sensitive: relies on real TTL expiry and background cleanup +func TestRegression_SlidingSessionTTL_TrafficKeepsSessionAlive(t *testing.T) { + const sessionTTL = 2 * time.Second + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + backend := helpers.CreateBackendServer(t, []helpers.BackendTool{ + helpers.NewBackendTool("ping", "ping tool", func(_ context.Context, _ map[string]any) string { + return `{"pong":true}` + }), + }, helpers.WithBackendName("ping-backend")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("ping-backend", helpers.WithURL(backend.URL+"/mcp")), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + helpers.WithSessionTTL(sessionTTL), + ) + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + + // ── Active session: traffic every 500ms for 4s (2x the TTL) ────────────── + // + // If the TTL were fixed (not sliding), the session would expire at the 2s + // mark and the next tools/call would fail. A sliding TTL refreshes + // last-access on each request, so all calls must succeed. + // + // We use tools/call (not tools/list) here because tools/list returns the + // SDK's cached tool set even after the vMCP session storage has evicted the + // session (the go-sdk transport map is not TTL-gated), so it would pass + // whether the TTL slides or not. tools/call routes through + // enforceSessionBinding → storage.Load, so it fails once the session is gone + // — making it the assertion that actually distinguishes sliding from fixed TTL. + activeClient := helpers.NewMCPClient(ctx, t, vmcpURL) + defer activeClient.Close() + + const activeTool = "ping-backend_ping" + + const ( + tickInterval = 500 * time.Millisecond + totalWindow = 4 * time.Second + ) + ticks := int(totalWindow / tickInterval) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + ticker := time.NewTicker(tickInterval) + defer ticker.Stop() + for i := 0; i < ticks; i++ { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + // tools/call routes through enforceSessionBinding → storage.Load, + // which refreshes last-access in the transport session storage. + _ = activeClient.CallTool(ctx, activeTool, map[string]any{}) + } + }() + wg.Wait() + + // One final call after the full window: the session must still be alive. + // CallTool fails the test if the session was evicted (the SDK surfaces the + // missing session as a transport error or an error result). + final := activeClient.CallTool(ctx, activeTool, map[string]any{}) + assert.False(t, final.IsError, + "active session must survive past its TTL while receiving traffic; got error result %v", final.Content) + + // ── Idle session: no traffic for 3s (> TTL) must be evicted ────────────── + // + // A second session that goes idle for longer than the TTL must be rejected + // on its next tool call. tools/list is NOT a reliable eviction signal here: + // the go-sdk StreamableHTTPHandler keeps its per-session transport in an + // in-memory map that is not TTL-gated, so tools/list continues to return + // the cached tool set even after the vMCP session storage has evicted the + // session. The eviction IS observable on tools/call, whose handler runs + // enforceSessionBinding → GetMultiSession → checkSession → storage.Load, + // which returns ErrSessionNotFound once the TTL has elapsed and the + // background cleanup sweep has removed the entry. + idleClient := newRawClient(vmcpURL) + idleClient.initialize(t, "") + + // Let the session go idle well past the TTL. The cleanup goroutine runs on + // a ttl/2 ticker, so waiting 2x the TTL after the idle period guarantees the + // background sweep has evicted the session. + idleFor := 3 * sessionTTL + time.Sleep(idleFor) + + // A tool call on the idle session must be rejected: enforceSessionBinding + // fails because the session is gone from storage. The SDK surfaces this as + // either an HTTP error (4xx) or a JSON-RPC error result — assert both. + resp := idleClient.postMCP(t, "", map[string]any{ + "jsonrpc": "2.0", + "id": idleClient.nextID, + "method": "tools/call", + "params": map[string]any{ + "name": "ping-backend_ping", + "arguments": map[string]any{}, + }, + }, idleClient.sessionID) + idleClient.nextID++ + defer resp.Body.Close() + + rejected := resp.StatusCode >= 400 || + isToolResultError(mustParseJSONRPC(t, resp)) + assert.True(t, rejected, + "idle session must be evicted after exceeding the TTL (status %d)", resp.StatusCode) +} + +// Compile-time check that this file stays in the vmcp_test package. +var _ vmcp.Backend diff --git a/test/integration/vmcp/vmcp_integration_test.go b/test/integration/vmcp/vmcp_integration_test.go index e8b37a6da4..91e296b810 100644 --- a/test/integration/vmcp/vmcp_integration_test.go +++ b/test/integration/vmcp/vmcp_integration_test.go @@ -12,10 +12,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvjson "github.com/stacklok/toolhive/pkg/json" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/test/integration/vmcp/vmcp_typing_integration_test.go b/test/integration/vmcp/vmcp_typing_integration_test.go index d973cf738a..f34abf635f 100644 --- a/test/integration/vmcp/vmcp_typing_integration_test.go +++ b/test/integration/vmcp/vmcp_typing_integration_test.go @@ -9,9 +9,9 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/composer" "github.com/stacklok/toolhive/test/integration/vmcp/helpers"