From 678a7a6664dd96d69cc592bd548890537ee2d36b Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 23 Jul 2026 20:22:34 +0200 Subject: [PATCH 1/9] Add inert outbound Modern MCP client shim vMCP re-originates calls to each backend, so Phase 3 must let it speak the 2026-07-28 ("Modern") revision to backends as a client: stateless, no initialize, no Mcp-Session-Id, server/discover. mcpcompat/go-sdk v1.6.1 cannot express such a request (its only no-initialize primitive is the private, Legacy-shaped resumeCall), so a hand-rolled shim is required. Add pkg/vmcp/client/modern.go with modernCall: a raw JSON-RPC-over-HTTP requester that sends the Modern headers/_meta, bounds the response body, decodes the Modern envelope, and maps a valid -32601 to method-not-found while flagging non-Modern responses (errWrongEra) and input_required. Extract buildBackendRoundTripper from defaultClientFactory so the shim and the Legacy mcp-go client share one auth+identity+header-forward+trace chain (order and per-call ctx preserved). Add mcp.Revision.String() and mcp.ModernRequestMeta as the single source of truth for the reserved _meta. Inert: nothing wires modernCall into dispatch yet (later steps do). Co-Authored-By: Claude Opus 4.8 --- pkg/mcp/revision.go | 46 +++- pkg/vmcp/client/client.go | 40 ++- pkg/vmcp/client/modern.go | 271 +++++++++++++++++++ pkg/vmcp/client/modern_integration_test.go | 164 ++++++++++++ pkg/vmcp/client/modern_test.go | 286 +++++++++++++++++++++ 5 files changed, 796 insertions(+), 11 deletions(-) create mode 100644 pkg/vmcp/client/modern.go create mode 100644 pkg/vmcp/client/modern_integration_test.go create mode 100644 pkg/vmcp/client/modern_test.go diff --git a/pkg/mcp/revision.go b/pkg/mcp/revision.go index 969020d9ed..ea044bccf7 100644 --- a/pkg/mcp/revision.go +++ b/pkg/mcp/revision.go @@ -26,6 +26,20 @@ const ( // 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. @@ -41,12 +55,30 @@ const metaKeyClientInfo = "io.modelcontextprotocol/clientInfo" // schema's RequestMetaObject. const metaKeyClientCapabilities = "io.modelcontextprotocol/clientCapabilities" -// reservedModernMetaKeys are the _meta keys a Legacy client never sets. The +// ReservedModernMetaKeys are the _meta keys a Legacy client never sets. The // presence of any one of them — independent of whether its 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 these keys turns into a rejection, never a downgrade. -var reservedModernMetaKeys = []string{metaKeyProtocolVersion, metaKeyClientInfo, metaKeyClientCapabilities} +// +// Exported so a Modern client (which sets these keys, the mirror of the +// classifier that reads them) can strip a caller's copies before overlaying +// its own authoritative values — see ModernRequestMeta. +var ReservedModernMetaKeys = []string{metaKeyProtocolVersion, metaKeyClientInfo, metaKeyClientCapabilities} + +// 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. They are @@ -325,6 +357,14 @@ var nameRequiredMethods = map[string]bool{ "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). @@ -412,7 +452,7 @@ func hasModernSignal(meta map[string]any, protoHeader string) bool { if protoHeader == MCPVersionModern { return true } - for _, key := range reservedModernMetaKeys { + for _, key := range ReservedModernMetaKeys { if _, ok := meta[key]; ok { return true } diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 119de0c829..112c04d106 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -497,14 +497,27 @@ func (*httpBackendClient) newSSEClient( return c, nil } -func (h *httpBackendClient) defaultClientFactory( - ctx context.Context, target *vmcp.BackendTarget, forwarding bool, -) (*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. +// 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) @@ -595,6 +608,17 @@ func (h *httpBackendClient) defaultClientFactory( 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 diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go new file mode 100644 index 0000000000..3dfbe0b6d3 --- /dev/null +++ b/pkg/vmcp/client/modern.go @@ -0,0 +1,271 @@ +// 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" + "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" + +// 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") + +// errModernInputRequired is returned when a Modern envelope decodes with a +// resultType other than "complete" (e.g. "input_required"). Multi-round tool +// retrieval is deferred; this shim detects and errors rather than returning a +// blank success. +var errModernInputRequired = errors.New("Modern response requires additional input (multi-round retrieval unsupported)") + +// 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/go-sdk v1.6.1 cannot express +// (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. +// +// Errors: errWrongEra when the peer is not Modern, mcp.ErrMethodNotFound for a +// valid -32601 error body, errModernInputRequired for a non-"complete" envelope, +// and a wrapped call error for any other JSON-RPC error. +func modernCall( + ctx context.Context, + hc *http.Client, + endpoint, method string, + params map[string]any, + name string, + out any, +) error { + id := modernRequestID.Add(1) + + reqParams := maps.Clone(params) + if reqParams == nil { + reqParams = map[string]any{} + } + reqParams["_meta"] = mergeModernMeta(params["_meta"]) + + 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) + } + 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", name) + } + // Mcp-Session-Id is deliberately never set: Modern is stateless. + + resp, err := hc.Do(req) + if err != nil { + return fmt.Errorf("sending %s request: %w", 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. + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + }() + + result, rpcErr, err := readModernEnvelope(resp, id) + if err != nil { + return err + } + if rpcErr != nil { + if rpcErr.Code == jsonRPCCodeMethodNotFound { + return fmt.Errorf("%w: %s", mcp.ErrMethodNotFound, rpcErr.Message) + } + 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). + // A result with no resultType is a Legacy-shaped body: wrong era. + var envelope struct { + ResultType string `json:"resultType"` + } + if json.Unmarshal(result, &envelope) != nil { + return errWrongEra + } + switch envelope.ResultType { + case modernResultTypeComplete: + // proceed to decode + case "": + return errWrongEra + default: + return fmt.Errorf("%w: resultType=%q", errModernInputRequired, envelope.ResultType) + } + + 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 (maps.Clone). +func mergeModernMeta(callerMeta any) map[string]any { + meta := map[string]any{} + if m, ok := callerMeta.(map[string]any); ok { + meta = maps.Clone(m) + for _, k := range mcpparser.ReservedModernMetaKeys { + delete(meta, k) + } + } + for k, v := range mcpparser.ModernRequestMeta(modernClientName, versions.Version) { + meta[k] = v + } + return meta +} + +// modernRPCError is the JSON-RPC error object. +type modernRPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// modernRPCEnvelope is the outer JSON-RPC response envelope. Method is set only +// on server->client requests/notifications interleaved on an SSE stream, which a +// single-shot client ignores. +type modernRPCEnvelope struct { + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *modernRPCError `json:"error"` + Method string `json:"method"` +} + +// 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. +func readModernEnvelope(resp *http.Response, wantID int64) (json.RawMessage, *modernRPCError, error) { + body := io.LimitReader(resp.Body, maxResponseSize) + + if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") { + return readModernSSE(body, wantID) + } + + data, err := io.ReadAll(body) + if err != nil { + return nil, nil, fmt.Errorf("reading response body: %w", 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, +// consuming (ignoring) any server->client requests/notifications interleaved on +// the stream. A stream that ends without a matching response yields errWrongEra. +func readModernSSE(body io.Reader, wantID int64) (json.RawMessage, *modernRPCError, error) { + sc := bufio.NewScanner(body) + sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + 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 != "" { + continue // server->client request/notification; not our response + } + 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("reading SSE stream: %w", 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 +} diff --git a/pkg/vmcp/client/modern_integration_test.go b/pkg/vmcp/client/modern_integration_test.go new file mode 100644 index 0000000000..0ca3948827 --- /dev/null +++ b/pkg/vmcp/client/modern_integration_test.go @@ -0,0 +1,164 @@ +// 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" + + "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, + // main re-added this kill-switch (default off, #5959); the harness + // must opt in so a well-formed Modern request reaches dispatchModern + // rather than falling through to the Legacy SDK path. + ModernDispatchEnabled: true, + }, + 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, "", &out) + 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") +} diff --git a/pkg/vmcp/client/modern_test.go b/pkg/vmcp/client/modern_test.go new file mode 100644 index 0000000000..7e27db2ace --- /dev/null +++ b/pkg/vmcp/client/modern_test.go @@ -0,0 +1,286 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "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. +func completeEnvelope(t *testing.T, id any, 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": id, "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, 1, 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) + 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, 1, 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)) + + 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, 1, 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, "", &out)) + 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, "", &out)) + assert.Equal(t, "complete", out["resultType"]) + assert.Equal(t, true, out["ok"]) +} + +// 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 wrong-era", + status: http.StatusOK, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}`, + wantErr: errWrongEra, + }, + { + 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", + }, + } + + 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) + require.Error(t, err) + if tt.wantErr != nil { + assert.ErrorIs(t, err, tt.wantErr) + 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) + }) + } +} + +// 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) +} From b27fb3169de782d34d8101162d55b6c9c27ebe6e Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 23 Jul 2026 20:42:49 +0200 Subject: [PATCH 2/9] Detect and cache each backend's MCP revision vMCP always spoke Legacy (initialize) to backends, so a Modern-only backend that removed initialize was unreachable. Teach the shared backend client to resolve each backend's revision once and remember it. Add a per-backend revisions cache (sync.Map keyed by workload ID) and probeRevision: it attempts server/discover Modern-first and classifies a backend Modern only on a clean discover result or a recognized Modern protocol-error code; every other outcome (bare 404/400/-32600/405, a -32601 for discover, empty/non-JSON body, a Legacy-shaped 200, or a transport failure) falls back to the unchanged Legacy initialize path, so a Legacy backend is never stranded. ListCapabilities consults the cache and probes on a miss; Legacy stays byte-for-byte unchanged and pays only one discover probe on first contact. Modern enumeration is a documented seam for the next step, so a detected Modern backend currently aggregates zero capabilities (inert, not broken). Also bound modernCall's drain-before-close now that discover is a live path. Co-Authored-By: Claude Opus 4.8 --- .../auth_error_mapping_regression_test.go | 9 + pkg/vmcp/client/client.go | 156 ++++++++++++- pkg/vmcp/client/client_test.go | 6 + pkg/vmcp/client/modern.go | 32 ++- pkg/vmcp/client/revision_test.go | 207 ++++++++++++++++++ 5 files changed, 407 insertions(+), 3 deletions(-) create mode 100644 pkg/vmcp/client/revision_test.go diff --git a/pkg/vmcp/client/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go index ea4c0ef0f0..6d86d3eecb 100644 --- a/pkg/vmcp/client/auth_error_mapping_regression_test.go +++ b/pkg/vmcp/client/auth_error_mapping_regression_test.go @@ -20,6 +20,7 @@ import ( "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" ) @@ -81,6 +82,10 @@ func TestRegression_401_MapsToErrAuthenticationFailed(t *testing.T) { 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), @@ -132,6 +137,8 @@ func TestRegression_403OnInitialize_LegacySSEFallback(t *testing.T) { 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), @@ -183,6 +190,8 @@ func TestRegression_403OnInitialize_MatchesSentinel(t *testing.T) { TransportType: "streamable-http", } + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + _, err := h.ListCapabilities(context.Background(), target) require.Error(t, err) diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 112c04d106..7113fb1f06 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -18,6 +18,7 @@ import ( "net" "net/http" "os" + "sync" "sync/atomic" "syscall" "time" @@ -29,6 +30,7 @@ import ( "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/secrets" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/versions" @@ -139,6 +141,18 @@ type httpBackendClient struct { // 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. + // + // ponytail: never evicted — a transient failure on the FIRST probe pins a + // backend to RevisionLegacy for the process lifetime. Recovery comes from + // Step 4's re-classification-on-error (re-probe + flip); add a TTL/re-probe + // only if flapping backends surface. + revisions sync.Map // map[string]mcpparser.Revision } // NewHTTPBackendClient creates a new HTTP-based backend client. @@ -907,12 +921,152 @@ func queryPrompts(ctx context.Context, c *client.Client, supported bool, backend return &mcp.ListPromptsResult{Prompts: []mcp.Prompt{}}, nil } +// 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 +} + +// setRevision records a backend's resolved MCP revision. +func (h *httpBackendClient) setRevision(workloadID string, rev mcpparser.Revision) { + h.revisions.Store(workloadID, rev) +} + +// 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, err + } + return &http.Client{Transport: rt, Timeout: 30 * time.Second}, nil +} + +// 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. +func (h *httpBackendClient) modernDiscover( + ctx context.Context, target *vmcp.BackendTarget, +) (*mcp.ServerCapabilities, error) { + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, err + } + var discover struct { + Capabilities mcp.ServerCapabilities `json:"capabilities"` + } + if err := modernCall(ctx, hc, target.BaseURL, "server/discover", nil, "", &discover); err != nil { + return nil, err + } + return &discover.Capabilities, nil +} + +// probeRevision resolves and caches a backend's MCP revision, Modern-first. +// +// It attempts a Modern server/discover and classifies MODERN only on (a) a clean +// discover result, or (b) a Modern-specific protocol error (-3202x), which proves +// the peer validated our Modern headers/_meta. EVERY other outcome — +// errWrongEra, a -32601 (discover is mandatory for Modern, so its absence means +// not Modern), a generic JSON-RPC error (-32600/-32603), a bare 404/400/405, an +// empty/non-JSON body, a 200-with-Legacy-result, an input_required envelope, or a +// timeout — falls back to LEGACY. This never strands a Legacy backend on a probe +// hiccup. +// +// A hard error is returned only when the backend transport cannot be built at all +// (e.g. invalid auth/CA config); that is a genuine misconfiguration, not a +// revision signal. +func (h *httpBackendClient) probeRevision( + ctx context.Context, target *vmcp.BackendTarget, +) (mcpparser.Revision, *mcp.ServerCapabilities, error) { + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return 0, nil, fmt.Errorf("failed to build transport for backend %s: %w", target.WorkloadID, err) + } + + var discover struct { + Capabilities mcp.ServerCapabilities `json:"capabilities"` + } + err = modernCall(ctx, hc, target.BaseURL, "server/discover", nil, "", &discover) + switch { + case err == nil: + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + return mcpparser.RevisionModern, &discover.Capabilities, nil + case errors.Is(err, errModernProtocolError): + // The peer validated our Modern protocol metadata and rejected it: it IS + // Modern, discover just failed application-side. No usable caps. + // ponytail: first probe yields an empty-but-successful capability list + // here (nil caps); later cache-hit modernDiscover re-surfaces this error. + // Reconciled in Step 2b when enumeration replaces the discover-only list. + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + return mcpparser.RevisionModern, nil, nil + default: + 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, nil + } +} + +// modernCapabilityList builds the discover-level CapabilityList for a Modern +// backend from its server/discover capability flags. +func (*httpBackendClient) modernCapabilityList( + target *vmcp.BackendTarget, caps *mcp.ServerCapabilities, +) *vmcp.CapabilityList { + slog.Debug("backend speaks Modern; discover capability flags", + "backend", target.WorkloadID, + "tools", caps != nil && caps.Tools != nil, + "resources", caps != nil && caps.Resources != nil, + "prompts", caps != nil && caps.Prompts != nil) + // ponytail: Modern tools/resources/prompts enumeration lands in Step 2b (#5911). + // For now discover only reports presence; the enumerations stay empty. + return &vmcp.CapabilityList{ + Tools: []vmcp.Tool{}, + Resources: []vmcp.Resource{}, + ResourceTemplates: []vmcp.ResourceTemplate{}, + Prompts: []vmcp.Prompt{}, + } +} + // 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. +// +// On the first call for a backend it probes the MCP revision Modern-first +// (probeRevision) and caches it. A Modern backend returns the discover-level +// capability list; a Legacy backend takes the unchanged initialize+enumerate +// path below. Subsequent calls read the cached revision and skip the probe. 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) + rev, cached := h.cachedRevision(target.WorkloadID) + switch { + case !cached: + probed, modernCaps, err := h.probeRevision(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "probe revision") + } + if probed == mcpparser.RevisionModern { + return h.modernCapabilityList(target, modernCaps), nil + } + // Legacy: fall through to the initialize+enumerate path below. + case rev == mcpparser.RevisionModern: + // Known Modern: one discover round-trip, no Legacy fallback. + modernCaps, err := h.modernDiscover(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "modern discover") + } + return h.modernCapabilityList(target, modernCaps), nil + } + // Create a client for this backend (not yet initialized) c, err := h.clientFactory(ctx, target, false) if err != nil { diff --git a/pkg/vmcp/client/client_test.go b/pkg/vmcp/client/client_test.go index edea258f7b..aadaa7faba 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -43,6 +43,7 @@ import ( "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/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/auth" authmocks "github.com/stacklok/toolhive/pkg/vmcp/auth/mocks" @@ -73,6 +74,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) diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go index 3dfbe0b6d3..1bd8385195 100644 --- a/pkg/vmcp/client/modern.go +++ b/pkg/vmcp/client/modern.go @@ -54,6 +54,14 @@ var errWrongEra = errors.New("backend response is not a Modern (2026-07-28) MCP // blank success. 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/-32022): 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") + // 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. @@ -125,8 +133,10 @@ func modernCall( } 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. - _, _ = io.Copy(io.Discard, resp.Body) + // 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() }() @@ -138,6 +148,9 @@ func modernCall( 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) + } return fmt.Errorf("modern %s: rpc error %d: %s", method, rpcErr.Code, rpcErr.Message) } @@ -269,3 +282,18 @@ func modernIDMatches(raw json.RawMessage, wantID int64) bool { var n int64 return json.Unmarshal(raw, &n) == nil && n == wantID } + +// isModernProtocolCode reports whether code is one of the Modern-specific +// JSON-RPC error codes (header/meta validation, -3202x). Only these prove the +// peer is Modern; generic JSON-RPC codes (-32600/-32601/-32603) do not, since a +// Legacy backend also returns them. +func isModernProtocolCode(code int) bool { + switch int64(code) { + case mcpparser.CodeHeaderMismatch, + mcpparser.CodeMissingClientCapability, + mcpparser.CodeUnsupportedProtocolVersion: + return true + default: + return false + } +} diff --git a/pkg/vmcp/client/revision_test.go b/pkg/vmcp/client/revision_test.go new file mode 100644 index 0000000000..d04d16fc40 --- /dev/null +++ b/pkg/vmcp/client/revision_test.go @@ -0,0 +1,207 @@ +// 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" + + 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" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// 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. +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{}}, + }, + }) + 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 -> Modern", + handler: func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(discoverEnvelope(t, r)) + }, + wantRev: mcpparser.RevisionModern, + }, + { + name: "recognized Modern protocol error (-32022) -> 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"}}`)) + }, + 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_TimeoutFallsBackToLegacy verifies a dead backend (connection +// refused) classifies Legacy rather than erroring. +func TestProbeRevision_TimeoutFallsBackToLegacy(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"} + + rev, caps, err := h.probeRevision(context.Background(), target) + require.NoError(t, err) + assert.Equal(t, mcpparser.RevisionLegacy, rev) + assert.Nil(t, caps) +} + +// TestListCapabilities_ModernServedFromCache verifies the cache: a Modern +// backend is probed once, and a second ListCapabilities is served from the +// cached revision (one discover round-trip, no re-probe fallback ladder). +func TestListCapabilities_ModernServedFromCache(t *testing.T) { + t.Parallel() + + var discoverCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + discoverCalls.Add(1) + assert.Equal(t, "server/discover", r.Header.Get("Mcp-Method")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(discoverEnvelope(t, r)) + })) + 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.NotNil(t, caps1) + + rev, ok := h.cachedRevision("modern") + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev) + + caps2, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.NotNil(t, caps2) + + // Step 2a is discover-only: enumerations are empty (Step 2b fills them). + assert.Empty(t, caps2.Tools) + + // Two ListCapabilities calls => exactly two discover round-trips (one probe, + // one cache-hit discover). If the cache were ignored, a re-probe would still + // be two, so the real signal is that no OTHER method was ever called and the + // backend never received an initialize handshake. + assert.EqualValues(t, 2, discoverCalls.Load()) +} From d099745086ff36ef3e104ca3ed33c2ac0bea46b1 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 23 Jul 2026 20:58:03 +0200 Subject: [PATCH 3/9] Enumerate Modern backend capabilities Step 2a detected a Modern backend but aggregated none of its capabilities. Enumerate them so a Modern backend contributes real tools/resources/prompts. Add modernEnumerate: it lists tools, resources, resource templates, and prompts over the Modern shim, gated by the discover capability flags (only querying what the backend advertised) and following nextCursor across pages via the shared pagination.ListAll helper (same cycle/iteration safety as the Legacy path, #5851). Results feed the shared newCapabilityListFromMCP, so the domain capability shape cannot diverge from Legacy. resources/templates/list degrades to empty on -32601, mirroring the Legacy tolerance. Both the first-probe and cache-hit Modern paths now enumerate consistently, resolving the transient empty-vs-error asymmetry from the previous step. Co-Authored-By: Claude Opus 4.8 --- pkg/vmcp/client/client.go | 269 ++++++++++++++------- pkg/vmcp/client/modern_enumerate_test.go | 139 +++++++++++ pkg/vmcp/client/modern_integration_test.go | 34 +++ pkg/vmcp/client/revision_test.go | 37 +-- 4 files changed, 380 insertions(+), 99 deletions(-) create mode 100644 pkg/vmcp/client/modern_enumerate_test.go diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 7113fb1f06..28be19eab1 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -1004,9 +1004,10 @@ func (h *httpBackendClient) probeRevision( case errors.Is(err, errModernProtocolError): // The peer validated our Modern protocol metadata and rejected it: it IS // Modern, discover just failed application-side. No usable caps. - // ponytail: first probe yields an empty-but-successful capability list - // here (nil caps); later cache-hit modernDiscover re-surfaces this error. - // Reconciled in Step 2b when enumeration replaces the discover-only list. + // nil caps => modernEnumerate returns an empty list. The cache-hit path + // tolerates the same -3202x error to nil caps, so both yield an empty + // list consistently (a Modern backend that rejects our discover exposes + // no enumerable capabilities). h.setRevision(target.WorkloadID, mcpparser.RevisionModern) return mcpparser.RevisionModern, nil, nil default: @@ -1017,24 +1018,178 @@ func (h *httpBackendClient) probeRevision( } } -// modernCapabilityList builds the discover-level CapabilityList for a Modern -// backend from its server/discover capability flags. -func (*httpBackendClient) modernCapabilityList( - target *vmcp.BackendTarget, caps *mcp.ServerCapabilities, +// 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, "create client") + } + endpoint := target.BaseURL + + var tools []mcp.Tool + if caps != nil && caps.Tools != nil { + tools, err = pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Tool, mcp.Cursor, error) { + var page struct { + Tools []mcp.Tool `json:"tools"` + NextCursor mcp.Cursor `json:"nextCursor"` + } + if err := modernCall(ctx, hc, endpoint, "tools/list", cursorParams(cursor), "", &page); err != nil { + return nil, "", err + } + return page.Tools, page.NextCursor, nil + }) + 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 = pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Resource, mcp.Cursor, error) { + var page struct { + Resources []mcp.Resource `json:"resources"` + NextCursor mcp.Cursor `json:"nextCursor"` + } + if err := modernCall(ctx, hc, endpoint, "resources/list", cursorParams(cursor), "", &page); err != nil { + return nil, "", err + } + return page.Resources, page.NextCursor, nil + }) + 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 = pagination.ListAll( + ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.ResourceTemplate, mcp.Cursor, error) { + var page struct { + ResourceTemplates []mcp.ResourceTemplate `json:"resourceTemplates"` + NextCursor mcp.Cursor `json:"nextCursor"` + } + if err := modernCall(ctx, hc, endpoint, "resources/templates/list", cursorParams(cursor), "", &page); err != nil { + return nil, "", err + } + return page.ResourceTemplates, page.NextCursor, nil + }) + switch { + case errors.Is(err, mcp.ErrMethodNotFound): + templates = nil + case err != nil: + return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") + } + } + + var prompts []mcp.Prompt + if caps != nil && caps.Prompts != nil { + prompts, err = pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Prompt, mcp.Cursor, error) { + var page struct { + Prompts []mcp.Prompt `json:"prompts"` + NextCursor mcp.Cursor `json:"nextCursor"` + } + if err := modernCall(ctx, hc, endpoint, "prompts/list", cursorParams(cursor), "", &page); err != nil { + return nil, "", err + } + return page.Prompts, page.NextCursor, nil + }) + 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)} +} + +// 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 { - slog.Debug("backend speaks Modern; discover capability flags", - "backend", target.WorkloadID, - "tools", caps != nil && caps.Tools != nil, - "resources", caps != nil && caps.Resources != nil, - "prompts", caps != nil && caps.Prompts != nil) - // ponytail: Modern tools/resources/prompts enumeration lands in Step 2b (#5911). - // For now discover only reports presence; the enumerations stay empty. - return &vmcp.CapabilityList{ - Tools: []vmcp.Tool{}, - Resources: []vmcp.Resource{}, - ResourceTemplates: []vmcp.ResourceTemplate{}, - Prompts: []vmcp.Prompt{}, + capabilities := &vmcp.CapabilityList{ + Tools: make([]vmcp.Tool, len(tools)), + Resources: make([]vmcp.Resource, len(resources)), + ResourceTemplates: make([]vmcp.ResourceTemplate, len(templates)), + Prompts: make([]vmcp.Prompt, len(prompts)), + } + + for i, tool := range tools { + capabilities.Tools[i] = vmcp.Tool{ + Name: tool.Name, + Description: tool.Description, + InputSchema: conversion.ConvertToolInputSchema(tool.InputSchema), + OutputSchema: conversion.ConvertToolOutputSchema(tool.OutputSchema), + Annotations: conversion.ConvertToolAnnotations(tool.Annotations), + BackendID: backendID, + } + } + + for i, resource := range resources { + capabilities.Resources[i] = vmcp.Resource{ + URI: resource.URI, + Name: resource.Name, + Description: resource.Description, + MimeType: resource.MIMEType, + BackendID: backendID, + } + } + + // 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{ + Name: arg.Name, + Description: arg.Description, + Required: arg.Required, + } + } + capabilities.Prompts[i] = vmcp.Prompt{ + Name: prompt.Name, + Description: prompt.Description, + Arguments: args, + BackendID: backendID, + } + } + + return capabilities } // ListCapabilities queries a backend for its MCP capabilities. @@ -1055,16 +1210,19 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B return nil, wrapBackendError(err, target.WorkloadID, "probe revision") } if probed == mcpparser.RevisionModern { - return h.modernCapabilityList(target, modernCaps), nil + return h.modernEnumerate(ctx, target, modernCaps) } // Legacy: fall through to the initialize+enumerate path below. case rev == mcpparser.RevisionModern: - // Known Modern: one discover round-trip, no Legacy fallback. + // Known Modern: one discover round-trip, no Legacy fallback. A -3202x + // protocol error is tolerated as nil caps so this cache-hit path yields + // the same empty enumeration as the first probe (probeRevision), rather + // than surfacing an error only on subsequent calls. modernCaps, err := h.modernDiscover(ctx, target) - if err != nil { + if err != nil && !errors.Is(err, errModernProtocolError) { return nil, wrapBackendError(err, target.WorkloadID, "modern discover") } - return h.modernCapabilityList(target, modernCaps), nil + return h.modernEnumerate(ctx, target, modernCaps) } // Create a client for this backend (not yet initialized) @@ -1113,66 +1271,11 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B return nil, wrapBackendError(err, target.WorkloadID, "list prompts") } - // Convert MCP types to vmcp types - capabilities := &vmcp.CapabilityList{ - Tools: make([]vmcp.Tool, len(toolsResp.Tools)), - Resources: make([]vmcp.Resource, len(resourcesResp.Resources)), - ResourceTemplates: make([]vmcp.ResourceTemplate, len(resourceTemplatesResp.ResourceTemplates)), - Prompts: make([]vmcp.Prompt, len(promptsResp.Prompts)), - } - - // Convert tools - for i, tool := range toolsResp.Tools { - capabilities.Tools[i] = vmcp.Tool{ - Name: tool.Name, - Description: tool.Description, - InputSchema: conversion.ConvertToolInputSchema(tool.InputSchema), - OutputSchema: conversion.ConvertToolOutputSchema(tool.OutputSchema), - Annotations: conversion.ConvertToolAnnotations(tool.Annotations), - BackendID: target.WorkloadID, - } - } - - // Convert resources - for i, resource := range resourcesResp.Resources { - capabilities.Resources[i] = vmcp.Resource{ - URI: resource.URI, - Name: resource.Name, - Description: resource.Description, - MimeType: resource.MIMEType, - BackendID: target.WorkloadID, - } - } - - // Convert resource templates (pass-through: no URI-template rewriting, like resources) - for i, template := range resourceTemplatesResp.ResourceTemplates { - capabilities.ResourceTemplates[i] = vmcp.ResourceTemplate{ - URITemplate: template.URITemplate, - Name: template.Name, - Description: template.Description, - MimeType: template.MIMEType, - BackendID: target.WorkloadID, - } - } - - // Convert prompts - for i, prompt := range promptsResp.Prompts { - args := make([]vmcp.PromptArgument, len(prompt.Arguments)) - for j, arg := range prompt.Arguments { - args[j] = vmcp.PromptArgument{ - Name: arg.Name, - Description: arg.Description, - Required: arg.Required, - } - } - - capabilities.Prompts[i] = vmcp.Prompt{ - Name: prompt.Name, - Description: prompt.Description, - Arguments: args, - BackendID: target.WorkloadID, - } - } + // 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 diff --git a/pkg/vmcp/client/modern_enumerate_test.go b/pkg/vmcp/client/modern_enumerate_test.go new file mode 100644 index 0000000000..71ce2e3000 --- /dev/null +++ b/pkg/vmcp/client/modern_enumerate_test.go @@ -0,0 +1,139 @@ +// 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 index 0ca3948827..a4d20f9234 100644 --- a/pkg/vmcp/client/modern_integration_test.go +++ b/pkg/vmcp/client/modern_integration_test.go @@ -18,6 +18,7 @@ import ( 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" @@ -162,3 +163,36 @@ func TestIntegration_ModernCall_Discover(t *testing.T) { 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/revision_test.go b/pkg/vmcp/client/revision_test.go index d04d16fc40..2cefaa5fa9 100644 --- a/pkg/vmcp/client/revision_test.go +++ b/pkg/vmcp/client/revision_test.go @@ -168,16 +168,25 @@ func TestProbeRevision_TimeoutFallsBackToLegacy(t *testing.T) { // TestListCapabilities_ModernServedFromCache verifies the cache: a Modern // backend is probed once, and a second ListCapabilities is served from the -// cached revision (one discover round-trip, no re-probe fallback ladder). +// cached revision (discover + enumerate, never a Legacy initialize handshake). func TestListCapabilities_ModernServedFromCache(t *testing.T) { t.Parallel() - var discoverCalls atomic.Int32 + var initializeCalls atomic.Int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - discoverCalls.Add(1) - assert.Equal(t, "server/discover", r.Header.Get("Mcp-Method")) - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write(discoverEnvelope(t, r)) + 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{}}}) + 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) @@ -186,22 +195,18 @@ func TestListCapabilities_ModernServedFromCache(t *testing.T) { caps1, err := h.ListCapabilities(context.Background(), target) require.NoError(t, err) - require.NotNil(t, caps1) + 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.NotNil(t, caps2) + require.Len(t, caps2.Tools, 1) + assert.Equal(t, "echo", caps2.Tools[0].Name) - // Step 2a is discover-only: enumerations are empty (Step 2b fills them). - assert.Empty(t, caps2.Tools) - - // Two ListCapabilities calls => exactly two discover round-trips (one probe, - // one cache-hit discover). If the cache were ignored, a re-probe would still - // be two, so the real signal is that no OTHER method was ever called and the - // backend never received an initialize handshake. - assert.EqualValues(t, 2, discoverCalls.Load()) + assert.Zero(t, initializeCalls.Load(), "a Modern backend must never receive a Legacy initialize") } From 7318a2c99dd7403a49df1d969de354a6f6a1e967 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 23 Jul 2026 21:18:54 +0200 Subject: [PATCH 4/9] Dispatch backend calls per negotiated revision With detection and enumeration in place, route the actual calls to each backend over the protocol it negotiated, so Modern backends become callable and Legacy backends are untouched. Add a dispatch seam that resolves a backend's cached revision (probing on a miss) and runs the call under it, then route CallTool, ReadResource, GetPrompt, and Complete through it. The Modern branch issues tools/call, resources/read, prompts/get, and completion/complete over the shim, applying the same advertised-to-backend capability-name translation as Legacy to both the Mcp-Name header and the body identifier, forwarding the caller _meta, and mapping the result (content, structuredContent, isError, result _meta) back through the shared converters. The closure dispatches on the revision argument, so a later re-probe-and-retry can re-run it under a corrected revision without touching the call sites. Co-Authored-By: Claude Opus 4.8 --- .../auth_error_mapping_regression_test.go | 2 + pkg/vmcp/client/client.go | 321 ++++++++++++++++-- pkg/vmcp/client/client_test.go | 3 + pkg/vmcp/client/modern_calls_test.go | 203 +++++++++++ 4 files changed, 509 insertions(+), 20 deletions(-) create mode 100644 pkg/vmcp/client/modern_calls_test.go diff --git a/pkg/vmcp/client/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go index 6d86d3eecb..cea0a3f2d3 100644 --- a/pkg/vmcp/client/auth_error_mapping_regression_test.go +++ b/pkg/vmcp/client/auth_error_mapping_regression_test.go @@ -293,6 +293,8 @@ func TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure(t *testin TransportType: "streamable-http", } + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 28be19eab1..530064cf7e 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -11,6 +11,7 @@ import ( "context" "crypto/tls" "crypto/x509" + "encoding/json" "errors" "fmt" "io" @@ -149,9 +150,10 @@ type httpBackendClient struct { // the Modern-first discover probe. // // ponytail: never evicted — a transient failure on the FIRST probe pins a - // backend to RevisionLegacy for the process lifetime. Recovery comes from - // Step 4's re-classification-on-error (re-probe + flip); add a TTL/re-probe - // only if flapping backends surface. + // backend to RevisionLegacy for the process lifetime. Recovery depends on + // re-classification-on-error (re-probe and flip the cached revision when a + // call reveals the other era); a TTL/periodic re-probe is deferred until + // flapping backends surface. revisions sync.Map // map[string]mcpparser.Revision } @@ -1087,6 +1089,7 @@ func (h *httpBackendClient) modernEnumerate( switch { case errors.Is(err, mcp.ErrMethodNotFound): templates = nil + err = nil // clear so a later reader can't mistake it for a live error case err != nil: return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") } @@ -1290,10 +1293,31 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B return capabilities, nil } -// CallTool invokes a tool on the backend MCP server. -// Returns the complete tool result including _meta field. +// dispatch resolves the backend's MCP revision (cache hit, else probeRevision) +// and runs fn exactly once with it. Every call verb routes through this seam so +// revision selection lives in one place. // -//nolint:gocyclo // this function is complex because it handles tool calls with various content types and error handling. +// ponytail: retry-on-revision-mismatch is deferred — when the probe classifies +// one era but the call reveals the other, fn would re-probe and run again. For +// now fn runs exactly once, no re-probe. +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 { + return wrapBackendError(err, target.WorkloadID, "probe revision") + } + rev = probed + } + return fn(ctx, rev) +} + +// 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, @@ -1302,7 +1326,52 @@ func (h *httpBackendClient) CallTool( meta map[string]any, ) (*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) + } else { + 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. +func (h *httpBackendClient) modernCallTool( + ctx context.Context, target *vmcp.BackendTarget, toolName string, arguments, meta map[string]any, +) (*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") + } + params := map[string]any{"name": backendToolName, "arguments": arguments} + if len(meta) > 0 { + params["_meta"] = meta + } + var result mcp.CallToolResult + if err := modernCall(ctx, hc, target.BaseURL, "tools/call", params, backendToolName, &result); 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), 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, true) if err != nil { @@ -1352,6 +1421,13 @@ func (h *httpBackendClient) CallTool( // dropped. See drainServerToClientNotifications for the lost-notification race. h.drainServerToClientNotifications(ctx, c) + return toolResultFromMCP(result, toolName, target.WorkloadID), 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) *vmcp.ToolCallResult { // Extract _meta field from backend response responseMeta := conversion.FromMCPMeta(result.Meta) @@ -1371,10 +1447,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 } @@ -1388,12 +1464,12 @@ func (h *httpBackendClient) CallTool( var structuredContent map[string]any if result.StructuredContent != nil { if structuredMap, ok := result.StructuredContent.(map[string]any); ok { - slog.Debug("using structured content from tool", "tool", toolName, "backend", target.WorkloadID) + slog.Debug("using structured content from tool", "tool", toolName, "backend", backendID) structuredContent = structuredMap } else { // StructuredContent is not an object - fall through to Content processing slog.Debug("structuredContent is not an object, falling back to Content", - "tool", toolName, "backend", target.WorkloadID) + "tool", toolName, "backend", backendID) } } @@ -1409,16 +1485,78 @@ 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") + } + // 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"` + } + if err := modernCall(ctx, hc, target.BaseURL, "resources/read", map[string]any{"uri": backendURI}, backendURI, &res); 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, false) if err != nil { @@ -1465,8 +1603,8 @@ func (h *httpBackendClient) ReadResource( }, 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, @@ -1474,7 +1612,68 @@ 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") + } + 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, &res); 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, false) if err != nil { @@ -1520,9 +1719,10 @@ func (h *httpBackendClient) GetPrompt( }, nil } -// Complete requests argument-completion candidates from the backend MCP server. -// Returns an empty (non-nil) result when the backend does not advertise the -// completions capability, matching the MCP spec's lenient completion semantics. +// 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, @@ -1530,9 +1730,90 @@ func (h *httpBackendClient) Complete( argName, argValue string, contextArgs map[string]string, ) (*vmcp.CompletionResult, error) { - slog.Debug("requesting completion from backend", - "ref_type", ref.Type, "backend", target.WorkloadName) + 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") + } + 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, "", &res) + 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 { diff --git a/pkg/vmcp/client/client_test.go b/pkg/vmcp/client/client_test.go index aadaa7faba..3c121a9864 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -372,6 +372,7 @@ func TestHTTPBackendClient_CallTool_WithMockFactory(t *testing.T) { TransportType: "streamable-http", } + backendClient.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) result, err := backendClient.CallTool(context.Background(), target, "test_tool", map[string]any{}, nil) require.Error(t, err) @@ -402,6 +403,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) @@ -432,6 +434,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) diff --git a/pkg/vmcp/client/modern_calls_test.go b/pkg/vmcp/client/modern_calls_test.go new file mode 100644 index 0000000000..7570295f2f --- /dev/null +++ b/pkg/vmcp/client/modern_calls_test.go @@ -0,0 +1,203 @@ +// 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"}) + 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") +} + +// 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) + 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) +} From cd49d213910a1cb647eb073c39f4bbe2e9b4c33a Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 23 Jul 2026 21:55:57 +0200 Subject: [PATCH 5/9] Re-probe and retry on backend revision mismatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cached revision can be wrong: a first-probe blip may pin a backend to Legacy, or a backend may be redeployed onto the other revision. Recover without waiting for a restart, but never at the cost of running a tool twice. When a call fails with a signal that proves the cached revision is wrong, dispatch re-probes the backend authoritatively, updates the cache, and retries the call once under the corrected revision. Recovery is gated on safety: retry happens only when the backend provably did NOT execute the request — a transport/protocol rejection (errWrongEra) or a Legacy initialize-step failure (errLegacyInitFailed). A well-formed Legacy-shaped 200 success to a Modern request (errLegacyResponseBody) means a lenient Legacy backend may already have run the tool, so the cache is corrected but the call is NOT retried. The Legacy mismatch signal is scoped to the initialize step so a genuine data-plane -32601 (e.g. an unimplemented completion/complete) never triggers a re-probe. ListCapabilities self-corrects through the same seam. A real revision change emits a WARN and increments a reclassification counter. Co-Authored-By: Claude Opus 4.8 --- .../auth_error_mapping_regression_test.go | 146 +++++----- pkg/vmcp/client/client.go | 206 ++++++++++--- pkg/vmcp/client/modern.go | 13 +- pkg/vmcp/client/modern_test.go | 4 +- pkg/vmcp/client/reclassify_test.go | 271 ++++++++++++++++++ .../backendtelemetry/backendtelemetry.go | 29 ++ .../backendtelemetry/backendtelemetry_test.go | 11 + 7 files changed, 557 insertions(+), 123 deletions(-) create mode 100644 pkg/vmcp/client/reclassify_test.go diff --git a/pkg/vmcp/client/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go index cea0a3f2d3..3cb0ca9f6e 100644 --- a/pkg/vmcp/client/auth_error_mapping_regression_test.go +++ b/pkg/vmcp/client/auth_error_mapping_regression_test.go @@ -59,20 +59,19 @@ func TestRegression_401_MapsToErrAuthenticationFailed(t *testing.T) { })) t.Cleanup(srv.Close) - h := &httpBackendClient{ - 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 - }, + 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{ @@ -95,12 +94,12 @@ func TestRegression_401_MapsToErrAuthenticationFailed(t *testing.T) { // TestRegression_403OnInitialize_LegacySSEFallback verifies that a backend // returning HTTP 403 on initialize is classified as ErrBackendUnavailable. // -// NOTE: The mcp-go streamable-HTTP transport returns a generic HTTP error for -// 403 ("request failed with status 403"), not transport.ErrLegacySSEServer. -// The "legacy SSE" hint in wrapBackendError is only added when the origin error -// IS transport.ErrLegacySSEServer (returned by SSE transport, not streamable-HTTP). -// For streamable-HTTP, 403 falls through to string-based classification and -// correctly maps to ErrBackendUnavailable, but without the SSE-specific message. +// 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() @@ -114,20 +113,19 @@ func TestRegression_403OnInitialize_LegacySSEFallback(t *testing.T) { })) t.Cleanup(srv.Close) - h := &httpBackendClient{ - 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 - }, + 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{ @@ -147,14 +145,14 @@ func TestRegression_403OnInitialize_LegacySSEFallback(t *testing.T) { "error message should reference 403 status, got: %v", err) } -// TestRegression_403OnInitialize_MatchesSentinel verifies that -// transport.ErrLegacySSEServer is NOT in the error chain for 403 on -// initialize, because wrapBackendError uses %v (not %w) for the -// original error, AND the mcp-go streamable-HTTP transport does not -// return ErrLegacySSEServer for 403 (it returns a generic HTTP error). -// Regardless of which error type is at the origin, the sentinel should -// never be in the chain. -func TestRegression_403OnInitialize_MatchesSentinel(t *testing.T) { +// 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) { @@ -167,20 +165,19 @@ func TestRegression_403OnInitialize_MatchesSentinel(t *testing.T) { })) t.Cleanup(srv.Close) - h := &httpBackendClient{ - 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 - }, + 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{ @@ -195,10 +192,12 @@ func TestRegression_403OnInitialize_MatchesSentinel(t *testing.T) { _, err := h.ListCapabilities(context.Background(), target) require.Error(t, err) - // wrapBackendError uses %v for the original error, so - // transport.ErrLegacySSEServer is NOT in the chain. - assert.False(t, errors.Is(err, mcptransport.ErrLegacySSEServer), - "transport.ErrLegacySSEServer should NOT be in the error chain (wrapBackendError uses %v)") + 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 @@ -270,20 +269,19 @@ func TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure(t *testin })) t.Cleanup(srv.Close) - h := &httpBackendClient{ - 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 - }, + 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{ diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 530064cf7e..2e26aeb6b6 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -41,6 +41,7 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/conversion" "github.com/stacklok/toolhive/pkg/vmcp/headerforward" 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" ) @@ -754,7 +755,10 @@ func wrapBackendError(err error, backendID string, operation string) error { // so we surface a clear message and classify as backend unavailable to allow recovery. 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) } @@ -788,8 +792,11 @@ 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) } @@ -1195,39 +1202,46 @@ func newCapabilityListFromMCP( return capabilities } -// ListCapabilities queries a backend for its MCP capabilities. -// Returns tools, resources, and prompts exposed by the backend. +// 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. // -// On the first call for a backend it probes the MCP revision Modern-first -// (probeRevision) and caches it. A Modern backend returns the discover-level -// capability list; a Legacy backend takes the unchanged initialize+enumerate -// path below. Subsequent calls read the cached revision and skip the probe. +// 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) - - rev, cached := h.cachedRevision(target.WorkloadID) - switch { - case !cached: - probed, modernCaps, err := h.probeRevision(ctx, target) - if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "probe revision") - } - if probed == mcpparser.RevisionModern { - return h.modernEnumerate(ctx, target, modernCaps) - } - // Legacy: fall through to the initialize+enumerate path below. - case rev == mcpparser.RevisionModern: - // Known Modern: one discover round-trip, no Legacy fallback. A -3202x - // protocol error is tolerated as nil caps so this cache-hit path yields - // the same empty enumeration as the first probe (probeRevision), rather - // than surfacing an error only on subsequent calls. - modernCaps, err := h.modernDiscover(ctx, target) - if err != nil && !errors.Is(err, errModernProtocolError) { - return nil, wrapBackendError(err, target.WorkloadID, "modern discover") + 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 h.modernEnumerate(ctx, target, modernCaps) + 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 { @@ -1240,9 +1254,9 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B }() // Initialize the client and get server capabilities - serverCaps, err := initializeClient(ctx, c) + serverCaps, err := h.legacyInit(ctx, c, target.WorkloadID) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + return nil, err } slog.Debug("backend capabilities", @@ -1293,13 +1307,41 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B return capabilities, nil } +// 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. +func (*httpBackendClient) legacyInit( + ctx context.Context, c *client.Client, backendID string, +) (*mcp.ServerCapabilities, error) { + caps, err := initializeClient(ctx, c) + if err != nil { + return nil, fmt.Errorf("%w: %w", errLegacyInitFailed, wrapBackendError(err, backendID, "initialize client")) + } + return caps, nil +} + // dispatch resolves the backend's MCP revision (cache hit, else probeRevision) -// and runs fn exactly once with it. Every call verb routes through this seam so -// revision selection lives in one place. +// and runs fn with it. Every call verb AND ListCapabilities route through this +// seam so revision selection — and self-correction — lives in one place. +// +// 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. // -// ponytail: retry-on-revision-mismatch is deferred — when the probe classifies -// one era but the call reveals the other, fn would re-probe and run again. For -// now fn runs exactly once, no re-probe. +// When the revision was just probed in THIS call (uncached), a re-probe would +// return the same answer, so a mismatch is surfaced directly without re-probing. func (h *httpBackendClient) dispatch( ctx context.Context, target *vmcp.BackendTarget, fn func(ctx context.Context, rev mcpparser.Revision) error, @@ -1312,7 +1354,79 @@ func (h *httpBackendClient) dispatch( } rev = probed } - return fn(ctx, rev) + + 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) or errLegacyResponseBody +// (a Legacy-shaped success body). 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) + } + 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 @@ -1384,9 +1498,9 @@ func (h *httpBackendClient) legacyCallTool( }() // Initialize the client and capture the backend's advertised capabilities. - serverCaps, err := initializeClient(ctx, c) + serverCaps, err := h.legacyInit(ctx, c, target.WorkloadID) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + return nil, err } // When forwarders are bound and the backend advertises logging, request debug @@ -1569,8 +1683,8 @@ func (h *httpBackendClient) legacyReadResource( }() // 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.WorkloadID); err != nil { + return nil, err } // Read the resource using the original URI from the backend's perspective. @@ -1686,8 +1800,8 @@ func (h *httpBackendClient) legacyGetPrompt( }() // 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.WorkloadID); err != nil { + return nil, err } // Get the prompt using the original prompt name from the backend's perspective. @@ -1826,9 +1940,9 @@ func (h *httpBackendClient) legacyComplete( }() // Initialize the client and capture the backend's advertised capabilities. - serverCaps, err := initializeClient(ctx, c) + serverCaps, err := h.legacyInit(ctx, c, target.WorkloadID) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + return nil, err } // Backends that do not advertise completions cannot serve completion/complete; diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go index 1bd8385195..07e7e5aee7 100644 --- a/pkg/vmcp/client/modern.go +++ b/pkg/vmcp/client/modern.go @@ -48,6 +48,14 @@ const jsonRPCCodeMethodNotFound = -32601 // 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 success result (no resultType); it may have executed the request") + // errModernInputRequired is returned when a Modern envelope decodes with a // resultType other than "complete" (e.g. "input_required"). Multi-round tool // retrieval is deferred; this shim detects and errors rather than returning a @@ -166,7 +174,10 @@ func modernCall( case modernResultTypeComplete: // proceed to decode case "": - return errWrongEra + // 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: return fmt.Errorf("%w: resultType=%q", errModernInputRequired, envelope.ResultType) } diff --git a/pkg/vmcp/client/modern_test.go b/pkg/vmcp/client/modern_test.go index 7e27db2ace..d63495ba12 100644 --- a/pkg/vmcp/client/modern_test.go +++ b/pkg/vmcp/client/modern_test.go @@ -237,11 +237,11 @@ func TestModernCall_ErrorMapping(t *testing.T) { wantErr: errWrongEra, }, { - name: "200 with Legacy-shaped result (no resultType) is wrong-era", + 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: errWrongEra, + wantErr: errLegacyResponseBody, }, { name: "other JSON-RPC error surfaces as a call error, not wrong-era", diff --git a/pkg/vmcp/client/reclassify_test.go b/pkg/vmcp/client/reclassify_test.go new file mode 100644 index 0000000000..e341d6353b --- /dev/null +++ b/pkg/vmcp/client/reclassify_test.go @@ -0,0 +1,271 @@ +// 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}, + {"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) so a +// re-probe classifies Modern, and rejects Legacy requests (no Mcp-Method header) +// with 404 so a Legacy initialize surfaces transport.ErrLegacySSEServer. +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{}}}) + case "tools/list": + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "echo", "inputSchema": map[string]any{"type": "object"}}}, + }) + 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) + 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") +} + +// TestReclassify_WarnsOnlyOnActualChange captures slog to confirm the WARN (which +// gates the reclassification counter in the same branch) fires only when the +// revision actually changes. +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_SelfCorrectsMisCachedBackend: a backend pinned Legacy by a +// stale cache recovers on the next ListCapabilities via one re-probe. +func TestListCapabilities_SelfCorrectsMisCachedBackend(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) + require.Len(t, caps.Tools, 1, "Modern enumeration recovered after re-probe") + assert.Equal(t, "echo", caps.Tools[0].Name) + + rev, _ := h.cachedRevision(target.WorkloadID) + assert.Equal(t, mcpparser.RevisionModern, rev) +} diff --git a/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go b/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go index dae74b5f67..051990e83a 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go @@ -13,8 +13,10 @@ 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" @@ -30,6 +32,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. +// +// ponytail: 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( diff --git a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go index 0c96135357..64713fb525 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go @@ -4,9 +4,20 @@ package backendtelemetry import ( + "context" "testing" ) +// 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() From b1ad4b9331822f1330a7088c307f78ba28470325 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Thu, 23 Jul 2026 22:11:37 +0200 Subject: [PATCH 6/9] Surface each backend's negotiated MCP revision Operators running a mixed Legacy/Modern fleet need to see which revision each backend negotiated, to debug era mismatches and plan upgrades. Expose the client's cached revision through a small optional revisionReporter accessor (kept off the BackendClient interface, so no mock churn), copy it into the health State read-model during each probe, and surface it as MCPRevision on DiscoveredBackend (flowing to VirtualMCPServer status via the operator's type alias) and as an mcp.protocol.revision label on the per-backend metrics. The client stays the source of truth; status only reads. Regenerates the CRD schema; no DeepCopy regen is needed for the scalar field. Co-Authored-By: Claude Opus 4.8 --- ...olhive.stacklok.dev_virtualmcpservers.yaml | 10 +++ ...olhive.stacklok.dev_virtualmcpservers.yaml | 10 +++ pkg/vmcp/client/client.go | 8 +++ pkg/vmcp/health/monitor.go | 28 ++++++++ pkg/vmcp/health/monitor_test.go | 66 +++++++++++++++++++ pkg/vmcp/health/status.go | 24 +++++++ pkg/vmcp/health/status_test.go | 19 ++++++ .../backendtelemetry/backendtelemetry.go | 31 +++++++++ .../backendtelemetry/backendtelemetry_test.go | 40 +++++++++++ pkg/vmcp/types.go | 5 ++ 10 files changed, 241 insertions(+) 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 289512fedd..ffd1118f61 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 @@ -3405,6 +3405,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 @@ -6847,6 +6852,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_virtualmcpservers.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml index c6b863f913..5b79959d56 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml @@ -3408,6 +3408,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 @@ -6850,6 +6855,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/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 2e26aeb6b6..a3e0d7747e 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -946,6 +946,14 @@ func (h *httpBackendClient) setRevision(workloadID string, rev mcpparser.Revisio h.revisions.Store(workloadID, rev) } +// 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 and health +// monitor can read it via an optional interface 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 diff --git a/pkg/vmcp/health/monitor.go b/pkg/vmcp/health/monitor.go index 7c5f3cead4..4fda3d8b06 100644 --- a/pkg/vmcp/health/monitor.go +++ b/pkg/vmcp/health/monitor.go @@ -13,10 +13,19 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" healthcontext "github.com/stacklok/toolhive/pkg/vmcp/health/context" ) +// revisionReporter is the optional accessor for a backend's cached MCP revision. +// The backend client (directly or through the telemetry decorator) implements it; +// it is NOT part of vmcp.BackendClient, so the monitor reaches it via a type +// assertion and simply reports no revision when absent. +type revisionReporter interface { + CachedRevision(workloadID string) (mcpparser.Revision, bool) +} + // WithHealthCheckMarker marks a context as a health check request. // Authentication layers can use IsHealthCheck to identify and skip authentication // for health check requests. @@ -119,6 +128,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 revisionReporter. + revisions revisionReporter + // statusTracker tracks health status for all backends. statusTracker *statusTracker @@ -252,8 +265,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.(revisionReporter) + return &Monitor{ checker: checker, + revisions: revisions, statusTracker: statusTracker, checkInterval: config.CheckInterval, backends: backends, @@ -487,6 +505,14 @@ func (m *Monitor) performHealthCheck(ctx context.Context, backend *vmcp.Backend) slog.Debug("health check succeeded for backend", "backend", backend.Name, "status", status) 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()) + } + } } // GetBackendStatus returns the current health status for a backend. @@ -735,6 +761,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 } @@ -752,6 +779,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 d03f305395..e91f3ca189 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" ) @@ -1256,3 +1257,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/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 051990e83a..7a683386ee 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go @@ -23,11 +23,20 @@ import ( "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" ) +// revisionReporter is the optional accessor the concrete backend client exposes +// for its cached MCP revision (see client.CachedRevision). It is NOT part of +// vmcp.BackendClient, so it is reached via a type assertion — a client that does +// not implement it simply reports no revision. +type revisionReporter interface { + CachedRevision(workloadID string) (mcpparser.Revision, bool) +} + const ( instrumentationName = "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -132,6 +141,26 @@ type telemetryBackendClient struct { var _ vmcp.BackendClient = telemetryBackendClient{} +// CachedRevision forwards to the wrapped client's optional 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.(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 { @@ -182,6 +211,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...) diff --git a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go index 64713fb525..5a43f7605e 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go @@ -6,8 +6,48 @@ 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 diff --git a/pkg/vmcp/types.go b/pkg/vmcp/types.go index 95e43f63e6..ff2bd9d425 100644 --- a/pkg/vmcp/types.go +++ b/pkg/vmcp/types.go @@ -286,6 +286,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. From 80d0e9017fc42d6d09e850958de6bc028b1665b5 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Fri, 24 Jul 2026 10:33:27 +0200 Subject: [PATCH 7/9] Harden the Modern backend client path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address three review findings on the outbound Modern client. Status-blind era signals could poison the revision cache: an empty or non-JSON response was read as "wrong era" regardless of HTTP status, so a transient 401/403/407/5xx on a Modern backend was treated as a revision mismatch and flipped it to Legacy. readModernEnvelope now classifies by status first — 400/404/405 stay a genuine not-Modern signal, auth (401/403/ 407) and transient (408/429/5xx, transport failures) become distinct errors that never drive reclassification — and probeRevision leaves a backend unprobed on an inconclusive result instead of caching Legacy on a blip. The SSE reader capped a single event at 4 MiB while the JSON branch and the docstring allowed 100 MiB, failing large tool/resource results over SSE; the scanner token limit now matches maxResponseSize (still bounded by the outer LimitReader). Each Modern call built a fresh transport whose idle connections lingered for ~90s; a plain CloseIdleConnections would silently no-op through the wrapper chain, so the auth/identity/header-forward/trace RoundTrippers now forward it to the underlying transport and each call site closes idle connections when it returns. No connection pooling is introduced. Co-Authored-By: Claude Opus 4.8 --- pkg/vmcp/client/client.go | 63 +++++++++++++++++---- pkg/vmcp/client/client_test.go | 26 +++++++++ pkg/vmcp/client/modern.go | 57 ++++++++++++++++--- pkg/vmcp/client/modern_test.go | 72 ++++++++++++++++++++++++ pkg/vmcp/client/reclassify_test.go | 32 +++++++++++ pkg/vmcp/client/revision_test.go | 17 +++--- pkg/vmcp/headerforward/transport.go | 9 +++ pkg/vmcp/headerforward/transport_test.go | 19 +++++++ 8 files changed, 269 insertions(+), 26 deletions(-) diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index a3e0d7747e..6ce3492461 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -371,6 +371,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 @@ -387,6 +395,14 @@ func (t *tracePropagatingRoundTripper) RoundTrip(req *http.Request) (*http.Respo return t.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. @@ -412,6 +428,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 @@ -979,6 +1003,7 @@ func (h *httpBackendClient) modernDiscover( if err != nil { return nil, err } + defer hc.CloseIdleConnections() var discover struct { Capabilities mcp.ServerCapabilities `json:"capabilities"` } @@ -988,20 +1013,24 @@ func (h *httpBackendClient) modernDiscover( return &discover.Capabilities, nil } -// probeRevision resolves and caches a backend's MCP revision, Modern-first. +// probeRevision resolves a backend's MCP revision, Modern-first. // -// It attempts a Modern server/discover and classifies MODERN only on (a) a clean -// discover result, or (b) a Modern-specific protocol error (-3202x), which proves -// the peer validated our Modern headers/_meta. EVERY other outcome — -// errWrongEra, a -32601 (discover is mandatory for Modern, so its absence means -// not Modern), a generic JSON-RPC error (-32600/-32603), a bare 404/400/405, an -// empty/non-JSON body, a 200-with-Legacy-result, an input_required envelope, or a -// timeout — falls back to LEGACY. This never strands a Legacy backend on a probe -// hiccup. +// It attempts a Modern server/discover and: +// - classifies MODERN (and caches it) on a clean discover result, or on a +// Modern-specific protocol error (-3202x) which proves the peer validated our +// Modern headers/_meta; +// - 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 an input_required envelope; +// - 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 returned only when the backend transport cannot be built at all -// (e.g. invalid auth/CA config); that is a genuine misconfiguration, not a -// revision signal. +// 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. func (h *httpBackendClient) probeRevision( ctx context.Context, target *vmcp.BackendTarget, ) (mcpparser.Revision, *mcp.ServerCapabilities, error) { @@ -1009,6 +1038,7 @@ func (h *httpBackendClient) probeRevision( if err != nil { return 0, nil, fmt.Errorf("failed to build transport for backend %s: %w", target.WorkloadID, err) } + defer hc.CloseIdleConnections() var discover struct { Capabilities mcp.ServerCapabilities `json:"capabilities"` @@ -1027,6 +1057,10 @@ func (h *httpBackendClient) probeRevision( // no enumerable capabilities). h.setRevision(target.WorkloadID, mcpparser.RevisionModern) return mcpparser.RevisionModern, nil, 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, nil, err default: slog.Debug("backend is not Modern; falling back to Legacy", "backend", target.WorkloadID, "probe_error", err) @@ -1051,6 +1085,7 @@ func (h *httpBackendClient) modernEnumerate( if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "create client") } + defer hc.CloseIdleConnections() endpoint := target.BaseURL var tools []mcp.Tool @@ -1477,6 +1512,7 @@ func (h *httpBackendClient) modernCallTool( 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 @@ -1644,6 +1680,7 @@ func (h *httpBackendClient) modernReadResource( 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), @@ -1763,6 +1800,7 @@ func (h *httpBackendClient) modernGetPrompt( if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "create client") } + defer hc.CloseIdleConnections() params := map[string]any{ "name": backendPromptName, "arguments": conversion.ConvertPromptArguments(arguments), @@ -1878,6 +1916,7 @@ func (h *httpBackendClient) modernComplete( 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 diff --git a/pkg/vmcp/client/client_test.go b/pkg/vmcp/client/client_test.go index 3c121a9864..3f9836f652 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -1779,3 +1779,29 @@ 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") +} diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go index 07e7e5aee7..fca4143db0 100644 --- a/pkg/vmcp/client/modern.go +++ b/pkg/vmcp/client/modern.go @@ -70,6 +70,19 @@ var errModernInputRequired = errors.New("Modern response requires additional inp // 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") + // 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. @@ -92,9 +105,18 @@ var modernRequestID atomic.Int64 // header-forward/trace chain (see buildBackendRoundTripper); modernCall adds no // transport concerns of its own. // -// Errors: errWrongEra when the peer is not Modern, mcp.ErrMethodNotFound for a -// valid -32601 error body, errModernInputRequired for a non-"complete" envelope, -// and a wrapped call error for any other JSON-RPC error. +// 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. func modernCall( ctx context.Context, hc *http.Client, @@ -137,7 +159,9 @@ func modernCall( resp, err := hc.Do(req) if err != nil { - return fmt.Errorf("sending %s request: %w", method, err) + // 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 @@ -229,7 +253,23 @@ type modernRPCEnvelope struct { // 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) (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") { @@ -238,7 +278,7 @@ func readModernEnvelope(resp *http.Response, wantID int64) (json.RawMessage, *mo data, err := io.ReadAll(body) if err != nil { - return nil, nil, fmt.Errorf("reading response body: %w", err) + return nil, nil, fmt.Errorf("%w: reading response body: %w", errModernTransient, err) } if len(bytes.TrimSpace(data)) == 0 { return nil, nil, errWrongEra @@ -258,7 +298,10 @@ func readModernEnvelope(resp *http.Response, wantID int64) (json.RawMessage, *mo // the stream. A stream that ends without a matching response yields errWrongEra. func readModernSSE(body io.Reader, wantID int64) (json.RawMessage, *modernRPCError, error) { sc := bufio.NewScanner(body) - sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + // 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 { @@ -280,7 +323,7 @@ func readModernSSE(body io.Reader, wantID int64) (json.RawMessage, *modernRPCErr return env.Result, env.Error, nil } if err := sc.Err(); err != nil { - return nil, nil, fmt.Errorf("reading SSE stream: %w", err) + return nil, nil, fmt.Errorf("%w: reading SSE stream: %w", errModernTransient, err) } return nil, nil, errWrongEra } diff --git a/pkg/vmcp/client/modern_test.go b/pkg/vmcp/client/modern_test.go index d63495ba12..bd367efd6e 100644 --- a/pkg/vmcp/client/modern_test.go +++ b/pkg/vmcp/client/modern_test.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/stretchr/testify/assert" @@ -250,6 +251,48 @@ func TestModernCall_ErrorMapping(t *testing.T) { 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 { @@ -268,6 +311,12 @@ func TestModernCall_ErrorMapping(t *testing.T) { 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 @@ -279,6 +328,29 @@ func TestModernCall_ErrorMapping(t *testing.T) { } } +// 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, "", &got)) + 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() diff --git a/pkg/vmcp/client/reclassify_test.go b/pkg/vmcp/client/reclassify_test.go index e341d6353b..0816e27671 100644 --- a/pkg/vmcp/client/reclassify_test.go +++ b/pkg/vmcp/client/reclassify_test.go @@ -225,6 +225,38 @@ func TestDispatch_NoReprobeOnLegacyDataPlaneMethodNotFound(t *testing.T) { 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. diff --git a/pkg/vmcp/client/revision_test.go b/pkg/vmcp/client/revision_test.go index 2cefaa5fa9..f2380f11d8 100644 --- a/pkg/vmcp/client/revision_test.go +++ b/pkg/vmcp/client/revision_test.go @@ -147,9 +147,11 @@ func TestProbeRevision_TruthTable(t *testing.T) { } } -// TestProbeRevision_TimeoutFallsBackToLegacy verifies a dead backend (connection -// refused) classifies Legacy rather than erroring. -func TestProbeRevision_TimeoutFallsBackToLegacy(t *testing.T) { +// 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. @@ -160,10 +162,11 @@ func TestProbeRevision_TimeoutFallsBackToLegacy(t *testing.T) { h := newProbeClient(t) target := &vmcp.BackendTarget{WorkloadID: "dead", BaseURL: url, TransportType: "streamable-http"} - rev, caps, err := h.probeRevision(context.Background(), target) - require.NoError(t, err) - assert.Equal(t, mcpparser.RevisionLegacy, rev) - assert.Nil(t, caps) + _, _, 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 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) +} From 97a65aaaf22eaa3ce462fafa85725ce4bd161ed8 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Fri, 24 Jul 2026 20:52:44 +0200 Subject: [PATCH 8/9] Constrain backend redirects; refine Modern probe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address an external security/correctness review of the backend client path. A malicious backend returning a cross-host 30x could exfiltrate the injected upstream credential: auth/identity/header-forward are applied in a RoundTripper that re-runs on every redirect hop, so Go's built-in Authorization stripping does not protect them. Install networking.SameHostRedirectPolicy on every backend HTTP client — via a single newBackendHTTPClient choke point for the per-call client (streamable, SSE, Modern shim) and directly on the session connector's two clients in mcp_session.go — so a cross-host redirect is refused before any credential leaves for the wrong host. Same-host redirects still follow. Also classify a Modern server/discover reply that decodes to an input_required envelope as Modern (not Legacy): a decoded resultType envelope proves the peer is Modern, so it now shares the protocol-error case (nil caps, empty enumeration) instead of caching Legacy. Document the two deliberate deferrals confirmed by review: outbound Mcp-Name is sent raw (non-ASCII identifiers unspecified until sentinel encoding lands), and the revision cache is last-writer-wins-safe (idempotent for deterministic backends; a flapping backend self-heals via reclassify), so no singleflight. Co-Authored-By: Claude Opus 4.8 --- pkg/vmcp/client/client.go | 58 +++++++++++------ pkg/vmcp/client/client_test.go | 64 +++++++++++++++++++ pkg/vmcp/client/modern.go | 4 ++ pkg/vmcp/client/revision_test.go | 10 +++ .../session/internal/backend/mcp_session.go | 13 +++- .../backend/mcp_session_redirect_test.go | 60 +++++++++++++++++ 6 files changed, 185 insertions(+), 24 deletions(-) create mode 100644 pkg/vmcp/session/internal/backend/mcp_session_redirect_test.go diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 6ce3492461..bf55c26736 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -32,6 +32,7 @@ import ( "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" @@ -150,10 +151,12 @@ type httpBackendClient struct { // first ListCapabilities for a backend and read on subsequent calls to skip // the Modern-first discover probe. // - // ponytail: never evicted — a transient failure on the FIRST probe pins a - // backend to RevisionLegacy for the process lifetime. Recovery depends on - // re-classification-on-error (re-probe and flip the cached revision when a - // call reveals the other era); a TTL/periodic re-probe is deferred until + // ponytail: never evicted, no singleflight/CAS. Concurrent first-probes are + // last-writer-wins-safe: writes are idempotent for a deterministic backend + // (every probe agrees), and a flapping backend self-heals via + // dispatch->reclassify (a call revealing the other era re-probes and flips). + // A transient probe failure caches nothing (probeRevision returns uncached), + // so a blip cannot pin a revision; a TTL/periodic re-probe is deferred until // flapping backends surface. revisions sync.Map // map[string]mcpparser.Revision } @@ -483,10 +486,7 @@ func (*httpBackendClient) newStreamableHTTPClient( } return resp, nil }) - httpClient := &http.Client{ - Transport: sizeLimitedTransport, - Timeout: 30 * time.Second, - } + httpClient := newBackendHTTPClient(sizeLimitedTransport, 30*time.Second) transportOpts := []transport.StreamableHTTPCOption{ transport.WithHTTPTimeout(30 * time.Second), transport.WithHTTPBasicClient(httpClient), @@ -522,7 +522,8 @@ func (*httpBackendClient) newSSEClient( ) (*client.Client, error) { c, err := client.NewSSEMCPClient( target.BaseURL, - transport.WithHTTPClient(&http.Client{Transport: baseTransport}), + // 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) @@ -990,7 +991,21 @@ func (h *httpBackendClient) buildModernHTTPClient(ctx context.Context, target *v if err != nil { return nil, err } - return &http.Client{Transport: rt, Timeout: 30 * time.Second}, nil + return newBackendHTTPClient(rt, 30*time.Second), nil +} + +// 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 @@ -1016,13 +1031,14 @@ func (h *httpBackendClient) modernDiscover( // probeRevision resolves a backend's MCP revision, Modern-first. // // It attempts a Modern server/discover and: -// - classifies MODERN (and caches it) on a clean discover result, or on a +// - classifies MODERN (and caches it) on a clean discover result, on a // Modern-specific protocol error (-3202x) which proves the peer validated our -// Modern headers/_meta; +// 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 an input_required envelope; +// JSON-RPC error, a bare 404/400/405, an empty/non-JSON body, or a +// 200-with-Legacy-result; // - 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). @@ -1048,13 +1064,13 @@ func (h *httpBackendClient) probeRevision( case err == nil: h.setRevision(target.WorkloadID, mcpparser.RevisionModern) return mcpparser.RevisionModern, &discover.Capabilities, nil - case errors.Is(err, errModernProtocolError): - // The peer validated our Modern protocol metadata and rejected it: it IS - // Modern, discover just failed application-side. No usable caps. - // nil caps => modernEnumerate returns an empty list. The cache-hit path - // tolerates the same -3202x error to nil caps, so both yield an empty - // list consistently (a Modern backend that rejects our discover exposes - // no enumerable capabilities). + 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"). Discover just failed + // application-side, so there are no usable caps — nil caps => modernEnumerate + // returns an empty list, and the cache-hit path tolerates the same to nil + // caps, so both yield an empty list consistently. h.setRevision(target.WorkloadID, mcpparser.RevisionModern) return mcpparser.RevisionModern, nil, nil case errors.Is(err, errModernAuth), errors.Is(err, errModernTransient): diff --git a/pkg/vmcp/client/client_test.go b/pkg/vmcp/client/client_test.go index 3f9836f652..4969267d4c 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -44,6 +44,7 @@ import ( 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" @@ -1805,3 +1806,66 @@ func TestModernChain_CloseIdleConnectionsForwards(t *testing.T) { 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/modern.go b/pkg/vmcp/client/modern.go index fca4143db0..c6d2abd41f 100644 --- a/pkg/vmcp/client/modern.go +++ b/pkg/vmcp/client/modern.go @@ -153,6 +153,10 @@ func modernCall( // Mcp-Method is required on EVERY Modern request (ValidateHeaderConsistency). req.Header.Set("Mcp-Method", method) if name != "" && mcpparser.IsNameRequiredMethod(method) { + // ponytail: sent raw; non-ASCII identifiers are not sentinel-encoded yet. + // URIs and ASCII names are safe; a non-ASCII name is unspecified behavior + // (a strict peer MAY reject or misinterpret the header). Add =?base64?..?= + // encoding if such names appear. req.Header.Set("Mcp-Name", name) } // Mcp-Session-Id is deliberately never set: Modern is stateless. diff --git a/pkg/vmcp/client/revision_test.go b/pkg/vmcp/client/revision_test.go index f2380f11d8..e4342a5d99 100644 --- a/pkg/vmcp/client/revision_test.go +++ b/pkg/vmcp/client/revision_test.go @@ -81,6 +81,16 @@ func TestProbeRevision_TruthTable(t *testing.T) { }, 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) { diff --git a/pkg/vmcp/session/internal/backend/mcp_session.go b/pkg/vmcp/session/internal/backend/mcp_session.go index 109a727ed2..95558bb909 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session.go +++ b/pkg/vmcp/session/internal/backend/mcp_session.go @@ -16,6 +16,7 @@ import ( mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/networking" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/versions" @@ -459,9 +460,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), @@ -488,7 +493,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), 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") +} From d24494c4c96b6d8b6b9160f76bb7c7d4e034c45c Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Fri, 24 Jul 2026 21:28:06 +0200 Subject: [PATCH 9/9] Satisfy golangci-lint on the Modern client path CI lint (not run locally) flagged the new code. Behavior-preserving cleanups: extract helpers to bring modernCall and modernEnumerate under the gocyclo limit; drop probeRevision's unused *ServerCapabilities return (unparam) and update callers; lowercase a sentinel error string (ST1005); wrap two over-long lines (lll); add t.Parallel to a reclassify test (paralleltest); remove an ineffectual err reset (ineffassign); fix import grouping (gci). Co-Authored-By: Claude Opus 4.8 --- pkg/vmcp/client/client.go | 104 +++++++++------------ pkg/vmcp/client/modern.go | 12 ++- pkg/vmcp/client/modern_enumerate_test.go | 1 - pkg/vmcp/client/modern_integration_test.go | 1 - pkg/vmcp/client/reclassify_test.go | 3 +- pkg/vmcp/client/revision_test.go | 9 +- 6 files changed, 60 insertions(+), 70 deletions(-) diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index bf55c26736..c5163600e7 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -1047,41 +1047,38 @@ func (h *httpBackendClient) modernDiscover( // // 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. +// The resolved capabilities are intentionally not returned: callers that need +// them (ListCapabilities) re-fetch via modernDiscover, so probeRevision only +// classifies and caches the revision. func (h *httpBackendClient) probeRevision( ctx context.Context, target *vmcp.BackendTarget, -) (mcpparser.Revision, *mcp.ServerCapabilities, error) { +) (mcpparser.Revision, error) { hc, err := h.buildModernHTTPClient(ctx, target) if err != nil { - return 0, nil, fmt.Errorf("failed to build transport for backend %s: %w", target.WorkloadID, err) + return 0, fmt.Errorf("failed to build transport for backend %s: %w", target.WorkloadID, err) } defer hc.CloseIdleConnections() - var discover struct { - Capabilities mcp.ServerCapabilities `json:"capabilities"` - } - err = modernCall(ctx, hc, target.BaseURL, "server/discover", nil, "", &discover) + err = modernCall(ctx, hc, target.BaseURL, "server/discover", nil, "", nil) switch { case err == nil: h.setRevision(target.WorkloadID, mcpparser.RevisionModern) - return mcpparser.RevisionModern, &discover.Capabilities, nil + 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"). Discover just failed - // application-side, so there are no usable caps — nil caps => modernEnumerate - // returns an empty list, and the cache-hit path tolerates the same to nil - // caps, so both yield an empty list consistently. + // Modern envelope (resultType present, != "complete"). h.setRevision(target.WorkloadID, mcpparser.RevisionModern) - return mcpparser.RevisionModern, nil, nil + 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, nil, err + return 0, err default: 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, nil + return mcpparser.RevisionLegacy, nil } } @@ -1106,16 +1103,7 @@ func (h *httpBackendClient) modernEnumerate( var tools []mcp.Tool if caps != nil && caps.Tools != nil { - tools, err = pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Tool, mcp.Cursor, error) { - var page struct { - Tools []mcp.Tool `json:"tools"` - NextCursor mcp.Cursor `json:"nextCursor"` - } - if err := modernCall(ctx, hc, endpoint, "tools/list", cursorParams(cursor), "", &page); err != nil { - return nil, "", err - } - return page.Tools, page.NextCursor, nil - }) + tools, err = modernListAll[mcp.Tool](ctx, hc, endpoint, "tools/list", "tools") if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "list tools") } @@ -1124,16 +1112,7 @@ func (h *httpBackendClient) modernEnumerate( var resources []mcp.Resource var templates []mcp.ResourceTemplate if caps != nil && caps.Resources != nil { - resources, err = pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Resource, mcp.Cursor, error) { - var page struct { - Resources []mcp.Resource `json:"resources"` - NextCursor mcp.Cursor `json:"nextCursor"` - } - if err := modernCall(ctx, hc, endpoint, "resources/list", cursorParams(cursor), "", &page); err != nil { - return nil, "", err - } - return page.Resources, page.NextCursor, nil - }) + resources, err = modernListAll[mcp.Resource](ctx, hc, endpoint, "resources/list", "resources") if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "list resources") } @@ -1141,21 +1120,10 @@ func (h *httpBackendClient) modernEnumerate( // 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 = pagination.ListAll( - ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.ResourceTemplate, mcp.Cursor, error) { - var page struct { - ResourceTemplates []mcp.ResourceTemplate `json:"resourceTemplates"` - NextCursor mcp.Cursor `json:"nextCursor"` - } - if err := modernCall(ctx, hc, endpoint, "resources/templates/list", cursorParams(cursor), "", &page); err != nil { - return nil, "", err - } - return page.ResourceTemplates, page.NextCursor, nil - }) + templates, err = modernListAll[mcp.ResourceTemplate](ctx, hc, endpoint, "resources/templates/list", "resourceTemplates") switch { case errors.Is(err, mcp.ErrMethodNotFound): templates = nil - err = nil // clear so a later reader can't mistake it for a live error case err != nil: return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") } @@ -1163,16 +1131,7 @@ func (h *httpBackendClient) modernEnumerate( var prompts []mcp.Prompt if caps != nil && caps.Prompts != nil { - prompts, err = pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]mcp.Prompt, mcp.Cursor, error) { - var page struct { - Prompts []mcp.Prompt `json:"prompts"` - NextCursor mcp.Cursor `json:"nextCursor"` - } - if err := modernCall(ctx, hc, endpoint, "prompts/list", cursorParams(cursor), "", &page); err != nil { - return nil, "", err - } - return page.Prompts, page.NextCursor, nil - }) + prompts, err = modernListAll[mcp.Prompt](ctx, hc, endpoint, "prompts/list", "prompts") if err != nil { return nil, wrapBackendError(err, target.WorkloadID, "list prompts") } @@ -1194,6 +1153,32 @@ func cursorParams(cursor mcp.Cursor) map[string]any { 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), "", &page); 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 @@ -1407,7 +1392,7 @@ func (h *httpBackendClient) dispatch( ) error { rev, cached := h.cachedRevision(target.WorkloadID) if !cached { - probed, _, err := h.probeRevision(ctx, target) + probed, err := h.probeRevision(ctx, target) if err != nil { return wrapBackendError(err, target.WorkloadID, "probe revision") } @@ -1444,7 +1429,7 @@ func (h *httpBackendClient) dispatch( func (h *httpBackendClient) reclassify( ctx context.Context, target *vmcp.BackendTarget, prev mcpparser.Revision, ) mcpparser.Revision { - corrected, _, err := h.probeRevision(ctx, target) + corrected, err := h.probeRevision(ctx, target) if err != nil { // Transport couldn't even be built to re-probe; keep the prior revision. return prev @@ -1711,7 +1696,8 @@ func (h *httpBackendClient) modernReadResource( } `json:"contents"` Meta map[string]any `json:"_meta"` } - if err := modernCall(ctx, hc, target.BaseURL, "resources/read", map[string]any{"uri": backendURI}, backendURI, &res); err != nil { + params := map[string]any{"uri": backendURI} + if err := modernCall(ctx, hc, target.BaseURL, "resources/read", params, backendURI, &res); err != nil { return nil, fmt.Errorf("resource read failed on backend %s: %w", target.WorkloadID, err) } mcpContents := make([]mcp.ResourceContents, len(res.Contents)) diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go index c6d2abd41f..28a071637d 100644 --- a/pkg/vmcp/client/modern.go +++ b/pkg/vmcp/client/modern.go @@ -54,13 +54,13 @@ var errWrongEra = errors.New("backend response is not a Modern (2026-07-28) MCP // 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 success result (no resultType); it may have executed the request") +var errLegacyResponseBody = errors.New("backend returned a Legacy-shaped body (no resultType); it may have executed") // errModernInputRequired is returned when a Modern envelope decodes with a // resultType other than "complete" (e.g. "input_required"). Multi-round tool // retrieval is deferred; this shim detects and errors rather than returning a // blank success. -var errModernInputRequired = errors.New("Modern response requires additional input (multi-round retrieval unsupported)") +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/-32022): the peer validated our @@ -180,6 +180,13 @@ func modernCall( 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) @@ -191,7 +198,6 @@ func modernCall( } // The Modern result is an envelope keyed by resultType (modern_envelope.go). - // A result with no resultType is a Legacy-shaped body: wrong era. var envelope struct { ResultType string `json:"resultType"` } diff --git a/pkg/vmcp/client/modern_enumerate_test.go b/pkg/vmcp/client/modern_enumerate_test.go index 71ce2e3000..14c25c1da6 100644 --- a/pkg/vmcp/client/modern_enumerate_test.go +++ b/pkg/vmcp/client/modern_enumerate_test.go @@ -15,7 +15,6 @@ import ( "github.com/stretchr/testify/require" mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" - "github.com/stacklok/toolhive/pkg/vmcp" ) diff --git a/pkg/vmcp/client/modern_integration_test.go b/pkg/vmcp/client/modern_integration_test.go index a4d20f9234..e4371467de 100644 --- a/pkg/vmcp/client/modern_integration_test.go +++ b/pkg/vmcp/client/modern_integration_test.go @@ -17,7 +17,6 @@ import ( 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" diff --git a/pkg/vmcp/client/reclassify_test.go b/pkg/vmcp/client/reclassify_test.go index 0816e27671..9e2105752e 100644 --- a/pkg/vmcp/client/reclassify_test.go +++ b/pkg/vmcp/client/reclassify_test.go @@ -19,7 +19,6 @@ import ( "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" @@ -260,6 +259,8 @@ func TestDispatch_TransientDoesNotReclassify(t *testing.T) { // 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 diff --git a/pkg/vmcp/client/revision_test.go b/pkg/vmcp/client/revision_test.go index e4342a5d99..260e6c92b3 100644 --- a/pkg/vmcp/client/revision_test.go +++ b/pkg/vmcp/client/revision_test.go @@ -14,12 +14,11 @@ import ( "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" - - mcpparser "github.com/stacklok/toolhive/pkg/mcp" - "github.com/stacklok/toolhive/pkg/vmcp" ) // newProbeClient builds a real httpBackendClient with an unauthenticated @@ -145,7 +144,7 @@ func TestProbeRevision_TruthTable(t *testing.T) { h := newProbeClient(t) target := &vmcp.BackendTarget{WorkloadID: "b1", BaseURL: srv.URL, TransportType: "streamable-http"} - rev, _, err := h.probeRevision(context.Background(), target) + rev, err := h.probeRevision(context.Background(), target) require.NoError(t, err) assert.Equal(t, tt.wantRev, rev) @@ -172,7 +171,7 @@ func TestProbeRevision_TransientLeavesUnprobed(t *testing.T) { h := newProbeClient(t) target := &vmcp.BackendTarget{WorkloadID: "dead", BaseURL: url, TransportType: "streamable-http"} - _, _, err := h.probeRevision(context.Background(), target) + _, err := h.probeRevision(context.Background(), target) require.Error(t, err) _, ok := h.cachedRevision("dead")