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/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/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go index ea4c0ef0f0..3cb0ca9f6e 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" ) @@ -58,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{ @@ -81,6 +81,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), @@ -90,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() @@ -109,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{ @@ -132,6 +135,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), @@ -140,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) { @@ -160,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{ @@ -183,13 +187,17 @@ 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) - // 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 @@ -261,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{ @@ -284,6 +291,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 119de0c829..c5163600e7 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" @@ -18,6 +19,7 @@ import ( "net" "net/http" "os" + "sync" "sync/atomic" "syscall" "time" @@ -29,6 +31,8 @@ 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/networking" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/versions" @@ -38,6 +42,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" ) @@ -139,6 +144,21 @@ 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, 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 } // NewHTTPBackendClient creates a new HTTP-based backend client. @@ -354,6 +374,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 @@ -370,6 +398,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. @@ -395,6 +431,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 @@ -442,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), @@ -481,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) @@ -497,14 +539,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 +650,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 @@ -714,7 +780,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) } @@ -748,8 +817,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) } @@ -883,102 +955,278 @@ func queryPrompts(ctx context.Context, c *client.Client, supported bool, backend return &mcp.ListPromptsResult{Prompts: []mcp.Prompt{}}, nil } -// ListCapabilities queries a backend for its MCP capabilities. -// Returns tools, resources, and prompts exposed by the backend. -// Only queries capabilities that the server advertises during initialization. -func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { - slog.Debug("querying capabilities from backend", "backend", target.WorkloadName, "url", target.BaseURL) - - // Create a client for this backend (not yet initialized) - c, err := h.clientFactory(ctx, target, false) - if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "create client") +// 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 } - defer func() { - if err := c.Close(); err != nil { - slog.Debug("failed to close client", "error", err) - } - }() + return v.(mcpparser.Revision), true +} - // Initialize the client and get server capabilities - serverCaps, err := initializeClient(ctx, c) +// setRevision records a backend's resolved MCP revision. +func (h *httpBackendClient) setRevision(workloadID string, rev mcpparser.Revision) { + 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 +// discover probe must carry the same security controls as every other backend +// call, so it must NOT use a bare http.Client. A 30s timeout matches the +// streamable-HTTP client; the response body is bounded inside modernCall +// (io.LimitReader), so no size-limit transport wrapper is needed here. +func (h *httpBackendClient) buildModernHTTPClient(ctx context.Context, target *vmcp.BackendTarget) (*http.Client, error) { + rt, err := h.buildBackendRoundTripper(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + return nil, err } + return newBackendHTTPClient(rt, 30*time.Second), nil +} - slog.Debug("backend capabilities", - "backend", target.WorkloadID, - "tools", serverCaps.Tools != nil, - "resources", serverCaps.Resources != nil, - "prompts", serverCaps.Prompts != 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(), + } +} - // Query each capability type based on server advertisement - // Check for nil BEFORE passing to functions to avoid interface{} nil pointer issues - toolsResp, err := queryTools(ctx, c, serverCaps.Tools != nil, target.WorkloadID) +// 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, wrapBackendError(err, target.WorkloadID, "list tools") + return nil, err + } + defer hc.CloseIdleConnections() + 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 +} - resourcesResp, err := queryResources(ctx, c, serverCaps.Resources != nil, target.WorkloadID) +// 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, on a +// Modern-specific protocol error (-3202x) which proves the peer validated our +// Modern headers/_meta, or on an input_required envelope (only returned after +// decoding a valid Modern envelope, so it too proves Modern); +// - classifies LEGACY (and caches it) on a GENUINE not-Modern signal — +// errWrongEra, a -32601 (discover is mandatory for Modern), a generic +// JSON-RPC error, a bare 404/400/405, an empty/non-JSON body, 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). +// Caching on these would flip a Modern backend to Legacy on a brief outage; +// leaving it unprobed lets the next call re-probe once the outage clears. +// +// A hard error is also returned when the backend transport cannot be built at all +// (e.g. invalid auth/CA config); that is a genuine misconfiguration. +// 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, error) { + hc, err := h.buildModernHTTPClient(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "list resources") + return 0, fmt.Errorf("failed to build transport for backend %s: %w", target.WorkloadID, err) } + defer hc.CloseIdleConnections() - // Resource templates share the same server capability advertisement as resources. - resourceTemplatesResp, err := queryResourceTemplates(ctx, c, serverCaps.Resources != nil, target.WorkloadID) - if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") + err = modernCall(ctx, hc, target.BaseURL, "server/discover", nil, "", nil) + switch { + case err == nil: + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + return mcpparser.RevisionModern, nil + case errors.Is(err, errModernProtocolError), errors.Is(err, errModernInputRequired): + // Both prove the peer is Modern: -3202x means it validated our Modern + // protocol metadata; input_required is only returned after decoding a valid + // Modern envelope (resultType present, != "complete"). + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + return mcpparser.RevisionModern, nil + case errors.Is(err, errModernAuth), errors.Is(err, errModernTransient): + // Inconclusive: an auth blip or a transient outage tells us nothing about + // the revision. Leave the backend unprobed so the next call re-probes. + return 0, err + default: + slog.Debug("backend is not Modern; falling back to Legacy", + "backend", target.WorkloadID, "probe_error", err) + h.setRevision(target.WorkloadID, mcpparser.RevisionLegacy) + return mcpparser.RevisionLegacy, nil } +} - promptsResp, err := queryPrompts(ctx, c, serverCaps.Prompts != nil, target.WorkloadID) +// modernEnumerate builds a backend's CapabilityList by enumerating each +// capability the discover flags advertise, via the Modern (2026-07-28) shim. +// Each list is gated on the matching discover flag (mirroring the Legacy +// initialize-path gating) and follows nextCursor across pages (#5851). +// +// caps may be nil: a Modern backend that rejected our discover with a -3202x +// protocol error (errModernProtocolError) is still classified Modern but yields +// no capability flags, so nothing is enumerated and an empty list is returned — +// the same outcome on the first probe and on later cache hits. +func (h *httpBackendClient) modernEnumerate( + ctx context.Context, target *vmcp.BackendTarget, caps *mcp.ServerCapabilities, +) (*vmcp.CapabilityList, error) { + hc, err := h.buildModernHTTPClient(ctx, target) if err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "list prompts") + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + endpoint := target.BaseURL + + var tools []mcp.Tool + if caps != nil && caps.Tools != nil { + tools, err = modernListAll[mcp.Tool](ctx, hc, endpoint, "tools/list", "tools") + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list tools") + } + } + + var resources []mcp.Resource + var templates []mcp.ResourceTemplate + if caps != nil && caps.Resources != nil { + resources, err = modernListAll[mcp.Resource](ctx, hc, endpoint, "resources/list", "resources") + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list resources") + } + + // Resource templates share the resources capability flag. A backend that + // does not implement resources/templates/list (-32601) degrades to an + // empty template list, mirroring the Legacy queryResourceTemplates path. + templates, err = modernListAll[mcp.ResourceTemplate](ctx, hc, endpoint, "resources/templates/list", "resourceTemplates") + switch { + case errors.Is(err, mcp.ErrMethodNotFound): + templates = nil + case err != nil: + return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") + } } - // Convert MCP types to vmcp types + var prompts []mcp.Prompt + if caps != nil && caps.Prompts != nil { + prompts, err = modernListAll[mcp.Prompt](ctx, hc, endpoint, "prompts/list", "prompts") + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list prompts") + } + } + + slog.Debug("backend capabilities queried (modern)", + "backend", target.WorkloadName, + "tools", len(tools), "resources", len(resources), + "resource_templates", len(templates), "prompts", len(prompts)) + return newCapabilityListFromMCP(target.WorkloadID, tools, resources, templates, prompts), nil +} + +// cursorParams builds the Modern list request params carrying a pagination +// cursor, or nil for the first page. +func cursorParams(cursor mcp.Cursor) map[string]any { + if cursor == "" { + return nil + } + return map[string]any{"cursor": string(cursor)} +} + +// modernListAll fetches every page of a Modern */list method, following +// nextCursor (#5851). itemsField is the result key holding the item array +// ("tools", "resources", "resourceTemplates", "prompts"); the envelope is decoded +// into a raw-message map so one helper serves all four list shapes. +func modernListAll[T any]( + ctx context.Context, hc *http.Client, endpoint, method, itemsField string, +) ([]T, error) { + return pagination.ListAll(ctx, func(ctx context.Context, cursor mcp.Cursor) ([]T, mcp.Cursor, error) { + var page map[string]json.RawMessage + if err := modernCall(ctx, hc, endpoint, method, cursorParams(cursor), "", &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 +// identical domain shapes. +func newCapabilityListFromMCP( + backendID string, + tools []mcp.Tool, resources []mcp.Resource, templates []mcp.ResourceTemplate, prompts []mcp.Prompt, +) *vmcp.CapabilityList { capabilities := &vmcp.CapabilityList{ - Tools: make([]vmcp.Tool, len(toolsResp.Tools)), - Resources: make([]vmcp.Resource, len(resourcesResp.Resources)), - ResourceTemplates: make([]vmcp.ResourceTemplate, len(resourceTemplatesResp.ResourceTemplates)), - Prompts: make([]vmcp.Prompt, len(promptsResp.Prompts)), + Tools: make([]vmcp.Tool, len(tools)), + Resources: make([]vmcp.Resource, len(resources)), + ResourceTemplates: make([]vmcp.ResourceTemplate, len(templates)), + Prompts: make([]vmcp.Prompt, len(prompts)), } - // Convert tools - for i, tool := range toolsResp.Tools { + 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: target.WorkloadID, + BackendID: backendID, } } - // Convert resources - for i, resource := range resourcesResp.Resources { + for i, resource := range resources { capabilities.Resources[i] = vmcp.Resource{ URI: resource.URI, Name: resource.Name, Description: resource.Description, MimeType: resource.MIMEType, - BackendID: target.WorkloadID, + BackendID: backendID, } } - // Convert resource templates (pass-through: no URI-template rewriting, like resources) - for i, template := range resourceTemplatesResp.ResourceTemplates { + // 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: target.WorkloadID, + BackendID: backendID, } } - // Convert prompts - for i, prompt := range promptsResp.Prompts { + for i, prompt := range prompts { args := make([]vmcp.PromptArgument, len(prompt.Arguments)) for j, arg := range prompt.Arguments { args[j] = vmcp.PromptArgument{ @@ -987,15 +1235,109 @@ func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.B Required: arg.Required, } } - capabilities.Prompts[i] = vmcp.Prompt{ Name: prompt.Name, Description: prompt.Description, Arguments: args, - BackendID: target.WorkloadID, + BackendID: backendID, + } + } + + return capabilities +} + +// ListCapabilities queries a backend for its MCP capabilities, selecting the +// Legacy or Modern path by revision. Returns tools, resources, and prompts +// exposed by the backend. +// +// Like the call verbs it routes through dispatch, so a mis-cached backend (e.g. a +// first-probe blip that pinned Legacy) self-corrects: the failing path triggers a +// re-probe and one retry under the corrected revision. +func (h *httpBackendClient) ListCapabilities(ctx context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + slog.Debug("querying capabilities from backend", "backend", target.WorkloadName, "url", target.BaseURL) + var out *vmcp.CapabilityList + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernListCapabilities(ctx, target) + } else { + out, err = h.legacyListCapabilities(ctx, target) + } + return err + }) + return out, err +} + +// modernListCapabilities resolves capabilities via Modern server/discover + +// enumeration. A -3202x protocol error is tolerated as nil caps (the backend is +// Modern but rejected our discover), yielding an empty list — consistent with +// probeRevision's classification. +func (h *httpBackendClient) modernListCapabilities( + ctx context.Context, target *vmcp.BackendTarget, +) (*vmcp.CapabilityList, error) { + caps, err := h.modernDiscover(ctx, target) + if err != nil && !errors.Is(err, errModernProtocolError) { + return nil, wrapBackendError(err, target.WorkloadID, "modern discover") + } + return h.modernEnumerate(ctx, target, caps) +} + +// legacyListCapabilities is the initialize + enumerate path (unchanged behavior). +func (h *httpBackendClient) legacyListCapabilities( + ctx context.Context, target *vmcp.BackendTarget, +) (*vmcp.CapabilityList, error) { + // Create a client for this backend (not yet initialized) + c, err := h.clientFactory(ctx, target, false) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer func() { + if err := c.Close(); err != nil { + slog.Debug("failed to close client", "error", err) } + }() + + // Initialize the client and get server capabilities + serverCaps, err := h.legacyInit(ctx, c, target.WorkloadID) + if err != nil { + return nil, err } + slog.Debug("backend capabilities", + "backend", target.WorkloadID, + "tools", serverCaps.Tools != nil, + "resources", serverCaps.Resources != nil, + "prompts", serverCaps.Prompts != nil) + + // Query each capability type based on server advertisement + // Check for nil BEFORE passing to functions to avoid interface{} nil pointer issues + toolsResp, err := queryTools(ctx, c, serverCaps.Tools != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list tools") + } + + resourcesResp, err := queryResources(ctx, c, serverCaps.Resources != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list resources") + } + + // Resource templates share the same server capability advertisement as resources. + resourceTemplatesResp, err := queryResourceTemplates(ctx, c, serverCaps.Resources != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list resource templates") + } + + promptsResp, err := queryPrompts(ctx, c, serverCaps.Prompts != nil, target.WorkloadID) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "list prompts") + } + + // Convert MCP types to vmcp types (shared with the Modern enumeration path). + capabilities := newCapabilityListFromMCP( + target.WorkloadID, + toolsResp.Tools, resourcesResp.Resources, resourceTemplatesResp.ResourceTemplates, promptsResp.Prompts, + ) + // TODO: Query server capabilities to detect logging/sampling support // This requires additional MCP protocol support for capabilities introspection @@ -1009,10 +1351,131 @@ 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. +// 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 with it. Every call verb AND ListCapabilities route through this +// seam so revision selection — and self-correction — lives in one place. // -//nolint:gocyclo // this function is complex because it handles tool calls with various content types and error handling. +// On a revision mismatch (isRevisionMismatch), dispatch always re-probes +// authoritatively and flips the cache so future calls use the corrected +// revision. It RETRIES fn only when the failure proves the backend did NOT +// execute the request — a protocol rejection (errWrongEra) or a Legacy +// initialize-step failure. A Legacy-shaped success body (errLegacyResponseBody) +// means a lenient backend MAY have executed a side-effecting request, so the +// cache is flipped but fn is NOT re-run. The retry is unconditionally single +// (no re-check), so it can never loop. +// +// When the revision was 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, +) 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 + } + + 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 +// 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, @@ -1021,7 +1484,53 @@ 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") + } + defer hc.CloseIdleConnections() + 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 { @@ -1034,9 +1543,9 @@ func (h *httpBackendClient) CallTool( }() // 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 @@ -1071,6 +1580,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) @@ -1090,10 +1606,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 } @@ -1107,12 +1623,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) } } @@ -1128,16 +1644,80 @@ func (h *httpBackendClient) CallTool( StructuredContent: structuredContent, IsError: result.IsError, Meta: responseMeta, - }, nil + } } -// ReadResource retrieves a resource from the backend MCP server. -// Returns the complete resource result including _meta field. +// ReadResource retrieves a resource from the backend MCP server, selecting the +// Legacy or Modern path by revision. Returns the complete resource result +// including _meta. func (h *httpBackendClient) ReadResource( ctx context.Context, target *vmcp.BackendTarget, uri string, ) (*vmcp.ResourceReadResult, error) { slog.Debug("reading resource from backend", "resource", uri, "backend", target.WorkloadName) + var out *vmcp.ResourceReadResult + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernReadResource(ctx, target, uri) + } else { + out, err = h.legacyReadResource(ctx, target, uri) + } + return err + }) + return out, err +} + +// modernReadResource invokes resources/read over the Modern shim. The URI is +// translated to the backend's capability name for both the body and the Mcp-Name +// header (server rejects a mismatch), and the result _meta is forwarded to core. +func (h *httpBackendClient) modernReadResource( + ctx context.Context, target *vmcp.BackendTarget, uri string, +) (*vmcp.ResourceReadResult, error) { + backendURI := target.GetBackendCapabilityName(uri) + if backendURI != uri { + slog.Debug("translating resource URI", "client_uri", uri, "backend_uri", backendURI) + } + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + // mcp.ResourceContents is an interface with no JSON unmarshaler, so it cannot + // decode directly. Decode the wire shape, rebuild the discriminated mcp types + // (blob takes precedence, symmetric with conversion.ToMCPResourceContents), + // then hand off to the SAME converter legacyReadResource uses so the + // content->vmcp mapping is shared, not duplicated. + var res struct { + Contents []struct { + URI string `json:"uri"` + MIMEType string `json:"mimeType"` + Text string `json:"text"` + Blob string `json:"blob"` + } `json:"contents"` + Meta map[string]any `json:"_meta"` + } + params := map[string]any{"uri": backendURI} + if err := modernCall(ctx, hc, target.BaseURL, "resources/read", params, backendURI, &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 { @@ -1150,8 +1730,8 @@ func (h *httpBackendClient) ReadResource( }() // Initialize the client - if _, err := initializeClient(ctx, c); err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + if _, err := h.legacyInit(ctx, c, target.WorkloadID); err != nil { + return nil, err } // Read the resource using the original URI from the backend's perspective. @@ -1184,8 +1764,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, @@ -1193,7 +1773,69 @@ func (h *httpBackendClient) GetPrompt( arguments map[string]any, ) (*vmcp.PromptGetResult, error) { slog.Debug("getting prompt from backend", "prompt", name, "backend", target.WorkloadName) + var out *vmcp.PromptGetResult + err := h.dispatch(ctx, target, func(ctx context.Context, rev mcpparser.Revision) error { + var err error + if rev == mcpparser.RevisionModern { + out, err = h.modernGetPrompt(ctx, target, name, arguments) + } else { + out, err = h.legacyGetPrompt(ctx, target, name, arguments) + } + return err + }) + return out, err +} +// modernGetPrompt invokes prompts/get over the Modern shim. The prompt name is +// translated to the backend's capability name for both the body and the Mcp-Name +// header (server rejects a mismatch); the result _meta is forwarded to core. +// Message content is decoded via mcp.UnmarshalContent so it goes through the same +// content conversion as the Legacy path. +func (h *httpBackendClient) modernGetPrompt( + ctx context.Context, target *vmcp.BackendTarget, name string, arguments map[string]any, +) (*vmcp.PromptGetResult, error) { + backendPromptName := target.GetBackendCapabilityName(name) + if backendPromptName != name { + slog.Debug("translating prompt name", "client_name", name, "backend_name", backendPromptName) + } + hc, err := h.buildModernHTTPClient(ctx, target) + if err != nil { + return nil, wrapBackendError(err, target.WorkloadID, "create client") + } + defer hc.CloseIdleConnections() + params := map[string]any{ + "name": backendPromptName, + "arguments": conversion.ConvertPromptArguments(arguments), + } + var res struct { + Description string `json:"description"` + Messages []struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` + } `json:"messages"` + Meta map[string]any `json:"_meta"` + } + if err := modernCall(ctx, hc, target.BaseURL, "prompts/get", params, backendPromptName, &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 { @@ -1206,8 +1848,8 @@ func (h *httpBackendClient) GetPrompt( }() // Initialize the client - if _, err := initializeClient(ctx, c); err != nil { - return nil, wrapBackendError(err, target.WorkloadID, "initialize client") + if _, err := h.legacyInit(ctx, c, target.WorkloadID); err != nil { + return nil, err } // Get the prompt using the original prompt name from the backend's perspective. @@ -1239,9 +1881,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, @@ -1249,9 +1892,91 @@ 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") + } + defer hc.CloseIdleConnections() + refMap, err := modernCompletionRef(target, ref) + if err != nil { + return nil, err + } + params := map[string]any{ + "ref": refMap, + "argument": map[string]any{"name": argName, "value": argValue}, + } + if len(contextArgs) > 0 { + params["context"] = map[string]any{"arguments": contextArgs} + } + var res struct { + Completion struct { + Values []string `json:"values"` + Total int `json:"total"` + HasMore bool `json:"hasMore"` + } `json:"completion"` + } + err = modernCall(ctx, hc, target.BaseURL, "completion/complete", params, "", &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 { @@ -1264,9 +1989,9 @@ func (h *httpBackendClient) Complete( }() // 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/client_test.go b/pkg/vmcp/client/client_test.go index edea258f7b..4969267d4c 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -43,6 +43,8 @@ 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/networking" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/auth" authmocks "github.com/stacklok/toolhive/pkg/vmcp/auth/mocks" @@ -73,6 +75,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) @@ -366,6 +373,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) @@ -396,6 +404,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) @@ -426,6 +435,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) @@ -1770,3 +1780,92 @@ type stubSamplingRequester struct{} func (*stubSamplingRequester) RequestSampling(context.Context, vmcp.SamplingRequest) (*vmcp.SamplingResult, error) { return nil, errors.New("stub: no downstream session") } + +// idleSpy is a RoundTripper that records CloseIdleConnections calls. +type idleSpy struct{ closed int } + +func (*idleSpy) RoundTrip(*http.Request) (*http.Response, error) { + return nil, errors.New("unused") +} +func (s *idleSpy) CloseIdleConnections() { s.closed++ } + +// TestModernChain_CloseIdleConnectionsForwards verifies the trace/identity/auth +// wrapper chain forwards CloseIdleConnections down to the concrete transport at +// the bottom (a plain hc.CloseIdleConnections would otherwise be a silent no-op). +func TestModernChain_CloseIdleConnectionsForwards(t *testing.T) { + t.Parallel() + + spy := &idleSpy{} + chain := &tracePropagatingRoundTripper{ + base: &identityPropagatingRoundTripper{ + base: &authRoundTripper{base: spy}, + }, + } + + // Reached the way http.Client.CloseIdleConnections reaches its transport. + chain.CloseIdleConnections() + assert.Equal(t, 1, spy.closed, "CloseIdleConnections must reach the bottom transport") +} + +// TestNewBackendHTTPClient_RedirectPolicy verifies the shared choke point installs +// SameHostRedirectPolicy: a cross-host 30x is refused (so the RoundTripper-injected +// credential never reaches the attacker host), while a same-host redirect is +// followed. +func TestNewBackendHTTPClient_RedirectPolicy(t *testing.T) { + t.Parallel() + + // inject mimics the auth/header-forward chain: it sets a credential on EVERY hop. + inject := roundTripperFunc(func(req *http.Request) (*http.Response, error) { + req.Header.Set("Authorization", "secret") + return http.DefaultTransport.RoundTrip(req) + }) + + t.Run("cross-host redirect refused, no credential leak", func(t *testing.T) { + t.Parallel() + + var attackerSawAuth atomic.Bool + attacker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + attackerSawAuth.Store(true) + } + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(attacker.Close) + + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Redirect(w, &http.Request{}, attacker.URL, http.StatusFound) + })) + t.Cleanup(origin.Close) + + hc := newBackendHTTPClient(inject, 5*time.Second) + resp, err := hc.Get(origin.URL) //nolint:noctx // test + if resp != nil { + _ = resp.Body.Close() + } + require.Error(t, err) + assert.ErrorIs(t, err, networking.ErrRedirectRefused) + assert.False(t, attackerSawAuth.Load(), "credential must not reach the cross-host redirect target") + }) + + t.Run("same-host redirect followed", func(t *testing.T) { + t.Parallel() + + var landed atomic.Bool + mux := http.NewServeMux() + mux.HandleFunc("/start", func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "/end", http.StatusFound) + }) + mux.HandleFunc("/end", func(w http.ResponseWriter, _ *http.Request) { + landed.Store(true) + w.WriteHeader(http.StatusOK) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + hc := newBackendHTTPClient(inject, 5*time.Second) + resp, err := hc.Get(srv.URL + "/start") //nolint:noctx // test + require.NoError(t, err) + _ = resp.Body.Close() + assert.True(t, landed.Load(), "same-host redirect must be followed") + }) +} diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go new file mode 100644 index 0000000000..28a071637d --- /dev/null +++ b/pkg/vmcp/client/modern.go @@ -0,0 +1,363 @@ +// 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") + +// errLegacyResponseBody is returned when a Modern request gets a well-formed 200 +// JSON-RPC SUCCESS result that is Legacy-shaped (no resultType). Unlike +// errWrongEra (a transport/protocol rejection that proves the backend did NOT +// process the request), a success body means a lenient Legacy backend MAY have +// executed the request — so the caller MUST NOT auto-retry it (double-execution +// of a side-effecting tool). The cache may still be reclassified. +var errLegacyResponseBody = errors.New("backend returned a Legacy-shaped body (no resultType); it may have executed") + +// errModernInputRequired is 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)") + +// 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") + +// 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. +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: 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, + 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) { + // 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. + + resp, err := hc.Do(req) + if err != nil { + // A transport/network failure (connection refused, timeout, ctx cancel) is + // transient — not a not-Modern signal — so it must not poison the revision cache. + return fmt.Errorf("%w: sending %s request: %w", errModernTransient, method, err) + } + defer func() { + // Drain so the connection can be reused (go-style rule); the readers below + // may stop early (SSE match) or the body may be a bare error page. Bounded + // by maxResponseSize so a hostile backend can't stall us on an unbounded + // drain now that this path is live in production (probeRevision). + _, _ = io.CopyN(io.Discard, resp.Body, maxResponseSize) + _ = resp.Body.Close() + }() + + result, rpcErr, err := readModernEnvelope(resp, id) + if err != nil { + return err + } + return interpretModernResult(result, rpcErr, method, out) +} + +// interpretModernResult maps a decoded Modern JSON-RPC response to an error or +// decodes it into out. Split from modernCall to keep each within the cyclomatic +// limit; see the modernCall doc for the error taxonomy. +func interpretModernResult(result json.RawMessage, rpcErr *modernRPCError, method string, out any) error { + if rpcErr != nil { + if rpcErr.Code == jsonRPCCodeMethodNotFound { + return fmt.Errorf("%w: %s", mcp.ErrMethodNotFound, rpcErr.Message) + } + if isModernProtocolCode(rpcErr.Code) { + return fmt.Errorf("%w: %s (rpc code %d)", errModernProtocolError, rpcErr.Message, rpcErr.Code) + } + return fmt.Errorf("modern %s: rpc error %d: %s", method, rpcErr.Code, rpcErr.Message) + } + + // The Modern result is an envelope keyed by resultType (modern_envelope.go). + var envelope struct { + ResultType string `json:"resultType"` + } + if json.Unmarshal(result, &envelope) != nil { + return errWrongEra + } + switch envelope.ResultType { + case modernResultTypeComplete: + // proceed to decode + case "": + // A JSON-RPC success result with no resultType is a Legacy-shaped body: a + // lenient Legacy backend that ignored our Modern headers and executed the + // request. Distinct from errWrongEra so the caller does not auto-retry. + return errLegacyResponseBody + default: + 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. +// +// 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") { + return readModernSSE(body, wantID) + } + + data, err := io.ReadAll(body) + if err != nil { + return nil, nil, fmt.Errorf("%w: reading response body: %w", errModernTransient, err) + } + if len(bytes.TrimSpace(data)) == 0 { + return nil, nil, errWrongEra + } + var env modernRPCEnvelope + if json.Unmarshal(data, &env) != nil { + return nil, nil, errWrongEra + } + if env.Error == nil && len(env.Result) == 0 { + return nil, nil, errWrongEra + } + return env.Result, env.Error, nil +} + +// readModernSSE scans an SSE body for the response whose id matches wantID, +// 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) + // Cap the token at maxResponseSize (the doc-promised bound) so a valid single + // data: event up to that size decodes; the outer io.LimitReader already bounds + // the total, so this cannot over-allocate. + sc.Buffer(make([]byte, 0, 64*1024), maxResponseSize) + for sc.Scan() { + data, ok := strings.CutPrefix(sc.Text(), "data:") + if !ok { + continue + } + var env modernRPCEnvelope + if json.Unmarshal([]byte(strings.TrimSpace(data)), &env) != nil { + continue + } + if env.Method != "" { + 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("%w: reading SSE stream: %w", errModernTransient, err) + } + return nil, nil, errWrongEra +} + +// modernIDMatches reports whether the raw JSON id equals wantID. +func modernIDMatches(raw json.RawMessage, wantID int64) bool { + if len(raw) == 0 { + return false + } + var n int64 + return json.Unmarshal(raw, &n) == nil && n == wantID +} + +// isModernProtocolCode reports whether code is one of the 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/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) +} diff --git a/pkg/vmcp/client/modern_enumerate_test.go b/pkg/vmcp/client/modern_enumerate_test.go new file mode 100644 index 0000000000..14c25c1da6 --- /dev/null +++ b/pkg/vmcp/client/modern_enumerate_test.go @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// modernReq decodes a Modern JSON-RPC request, reading the body exactly once +// (callers must not read r.Body separately) and returning the id and cursor. +func modernReq(t *testing.T, r *http.Request) (id any, cursor string) { + t.Helper() + body, err := readAll(t, r) + require.NoError(t, err) + var req struct { + ID any `json:"id"` + Params struct { + Cursor string `json:"cursor"` + } `json:"params"` + } + require.NoError(t, json.Unmarshal(body, &req)) + return req.ID, req.Params.Cursor +} + +// writeModernResult writes a Modern JSON-RPC success envelope for the given +// request id, wrapping result under a "complete" resultType. +func writeModernResult(t *testing.T, w http.ResponseWriter, id any, result map[string]any) { + t.Helper() + if result["resultType"] == nil { + result["resultType"] = "complete" + } + out, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + require.NoError(t, err) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(out) +} + +// TestModernEnumerate_Pagination verifies tools/list follows nextCursor across +// pages and aggregates every tool (no page-1 truncation, #5851). +func TestModernEnumerate_Pagination(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "tools/list", r.Header.Get("Mcp-Method")) + id, cursor := modernReq(t, r) + if cursor == "" { + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "t1", "inputSchema": map[string]any{"type": "object"}}}, + "nextCursor": "page2", + }) + return + } + require.Equal(t, "page2", cursor) + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "t2", "inputSchema": map[string]any{"type": "object"}}}, + }) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + caps := &mcpmcp.ServerCapabilities{Tools: &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{}} + + list, err := h.modernEnumerate(context.Background(), target, caps) + require.NoError(t, err) + require.Len(t, list.Tools, 2, "both pages must aggregate") + assert.Equal(t, "t1", list.Tools[0].Name) + assert.Equal(t, "t2", list.Tools[1].Name) + assert.Equal(t, "b", list.Tools[0].BackendID) +} + +// TestModernEnumerate_GatesOnFlags verifies only advertised capabilities are +// enumerated: a backend advertising only Tools is never asked for resources or +// prompts. +func TestModernEnumerate_GatesOnFlags(t *testing.T) { + t.Parallel() + + var toolsList, otherList atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Mcp-Method") == "tools/list" { + toolsList.Add(1) + } else { + otherList.Add(1) + } + id, _ := modernReq(t, r) + writeModernResult(t, w, id, map[string]any{"tools": []any{}}) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + caps := &mcpmcp.ServerCapabilities{Tools: &struct { + ListChanged bool `json:"listChanged,omitempty"` + }{}} + + _, err := h.modernEnumerate(context.Background(), target, caps) + require.NoError(t, err) + + assert.Equal(t, int32(1), toolsList.Load(), "tools advertised: listed once") + assert.Zero(t, otherList.Load(), "resources/prompts not advertised: must not be listed") +} + +// TestModernEnumerate_NilCapsEmpty verifies nil caps (a -3202x backend) yields +// an empty list with no list calls at all. +func TestModernEnumerate_NilCapsEmpty(t *testing.T) { + t.Parallel() + + var called atomic.Bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + called.Store(true) + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + + list, err := h.modernEnumerate(context.Background(), target, nil) + require.NoError(t, err) + assert.Empty(t, list.Tools) + assert.Empty(t, list.Resources) + assert.Empty(t, list.Prompts) + assert.False(t, called.Load(), "nil caps must not issue any list request") +} diff --git a/pkg/vmcp/client/modern_integration_test.go b/pkg/vmcp/client/modern_integration_test.go new file mode 100644 index 0000000000..e4371467de --- /dev/null +++ b/pkg/vmcp/client/modern_integration_test.go @@ -0,0 +1,197 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// startEchoBackend stands up a real in-process MCP server over streamable-HTTP +// exposing a single "echo" tool, returning the URL of its /mcp endpoint. Mirrors +// the server package's startRealMCPBackend test utility (not importable across +// the package boundary). +func startEchoBackend(t *testing.T) string { + t.Helper() + + mcpSrv := mcpserver.NewMCPServer("real-backend", "1.0.0") + mcpSrv.AddTool( + mcpmcp.NewTool("echo", + mcpmcp.WithDescription("Echoes the input back"), + mcpmcp.WithString("input", mcpmcp.Required()), + ), + func(_ context.Context, req mcpmcp.CallToolRequest) (*mcpmcp.CallToolResult, error) { + args, _ := req.Params.Arguments.(map[string]any) + input, _ := args["input"].(string) + return &mcpmcp.CallToolResult{Content: []mcpmcp.Content{mcpmcp.NewTextContent(input)}}, nil + }, + ) + + mux := http.NewServeMux() + mux.Handle("/mcp", mcpserver.NewStreamableHTTPServer(mcpSrv)) + ts := httptest.NewServer(mux) + t.Cleanup(ts.Close) + return ts.URL + "/mcp" +} + +// newModernVMCPServer stands up a real vMCP server with Modern (2026-07-28) +// dispatch enabled, aggregating the echo backend at backendURL. A well-formed +// Modern request routes through classifyingHandler -> dispatchModern (the Phase-2 +// server side), so this is the real target modernCall must satisfy end-to-end. +func newModernVMCPServer(t *testing.T, backendURL string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: backendURL, + TransportType: "streamable-http", + } + mockRegistry := mocks.NewMockBackendRegistry(ctrl) + mockRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + authReg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, authReg.RegisterStrategy( + authtypes.StrategyTypeUnauthenticated, strategies.NewUnauthenticatedStrategy(), + )) + + backendClient, err := NewHTTPBackendClient(authReg) + require.NoError(t, err) + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + srv, err := server.New( + context.Background(), + &server.Config{ + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: vmcpsession.NewSessionFactory(authReg), + Aggregator: agg, + // 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") +} + +// TestIntegration_ModernListCapabilities proves the full Modern enumeration path +// end-to-end: ListCapabilities probes the Phase-2 dispatchModern server as +// Modern, then enumerates its real capabilities (tools/list) via the shim — no +// initialize handshake, no Mcp-Session-Id — and returns the echo backend's tool. +func TestIntegration_ModernListCapabilities(t *testing.T) { + t.Parallel() + + backendURL := startEchoBackend(t) + vmcpSrv := newModernVMCPServer(t, backendURL) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{ + WorkloadID: "vmcp-modern", + WorkloadName: "vMCP Modern", + BaseURL: vmcpSrv.URL + "/mcp", + TransportType: "streamable-http", + } + + caps, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + + // Classified Modern (not Legacy initialize) and the echo tool enumerated. + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev) + + require.Len(t, caps.Tools, 1, "the echo backend's tool must enumerate through Modern") + assert.Equal(t, "echo", caps.Tools[0].Name) + assert.Equal(t, target.WorkloadID, caps.Tools[0].BackendID) + assert.Empty(t, caps.Resources, "echo backend exposes no resources") + assert.Empty(t, caps.Prompts, "echo backend exposes no prompts") +} diff --git a/pkg/vmcp/client/modern_test.go b/pkg/vmcp/client/modern_test.go new file mode 100644 index 0000000000..bd367efd6e --- /dev/null +++ b/pkg/vmcp/client/modern_test.go @@ -0,0 +1,358 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// completeEnvelope is a minimal valid Modern success body for a given method's +// result payload merged with the resultType envelope key. +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 a legacy-response-body", + status: http.StatusOK, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}`, + wantErr: errLegacyResponseBody, + }, + { + name: "other JSON-RPC error surfaces as a call error, not wrong-era", + status: http.StatusOK, + contentType: "application/json", + body: `{"jsonrpc":"2.0","id":1,"error":{"code":-32603,"message":"boom"}}`, + wantMsg: "boom", + }, + { + name: "401 is auth, not wrong-era", + status: http.StatusUnauthorized, + contentType: "text/plain", + body: "", + wantErr: errModernAuth, + }, + { + name: "403 is auth, not wrong-era", + status: http.StatusForbidden, + contentType: "application/json", + body: "", + wantErr: errModernAuth, + }, + { + name: "407 proxy-auth is auth, not wrong-era", + status: http.StatusProxyAuthRequired, + contentType: "text/plain", + body: "", + wantErr: errModernAuth, + }, + { + name: "503 is transient, not wrong-era", + status: http.StatusServiceUnavailable, + contentType: "text/plain", + body: "", + wantErr: errModernTransient, + }, + { + name: "429 is transient, not wrong-era", + status: http.StatusTooManyRequests, + contentType: "application/json", + body: "", + wantErr: errModernTransient, + }, + { + name: "401 tagged text/event-stream is still auth, not wrong-era", + status: http.StatusUnauthorized, + contentType: "text/event-stream", + body: "", + wantErr: errModernAuth, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", tt.contentType) + w.WriteHeader(tt.status) + _, _ = w.Write([]byte(tt.body)) + })) + t.Cleanup(srv.Close) + + err := modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil) + require.Error(t, err) + if tt.wantErr != nil { + assert.ErrorIs(t, err, tt.wantErr) + // Auth/transient statuses must never masquerade as the not-Modern + // signal (that is what would poison a cached-Modern backend). + if tt.wantErr == errModernAuth || tt.wantErr == errModernTransient { + assert.NotErrorIs(t, err, errWrongEra) + assert.NotErrorIs(t, err, errLegacyResponseBody) + } + return + } + // A valid JSON-RPC error body means the backend IS Modern: the error + // must surface (not be reclassified as wrong-era or method-not-found). + assert.NotErrorIs(t, err, errWrongEra) + assert.NotErrorIs(t, err, mcp.ErrMethodNotFound) + assert.Contains(t, err.Error(), tt.wantMsg) + }) + } +} + +// TestModernCall_LargeSSEEvent verifies a single SSE data: event larger than the +// old 4 MiB scanner cap (but under maxResponseSize) decodes successfully. +func TestModernCall_LargeSSEEvent(t *testing.T) { + t.Parallel() + + big := strings.Repeat("x", 5*1024*1024) // 5 MiB > old 4 MiB token cap + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + w.Header().Set("Content-Type", "text/event-stream") + out, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, + "result": map[string]any{"resultType": "complete", "blob": big}}) + require.NoError(t, err) + _, _ = w.Write([]byte("data: " + string(out) + "\n\n")) + })) + t.Cleanup(srv.Close) + + var got struct { + Blob string `json:"blob"` + } + require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", &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() + return io.ReadAll(r.Body) +} diff --git a/pkg/vmcp/client/reclassify_test.go b/pkg/vmcp/client/reclassify_test.go new file mode 100644 index 0000000000..9e2105752e --- /dev/null +++ b/pkg/vmcp/client/reclassify_test.go @@ -0,0 +1,304 @@ +// 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") +} + +// TestDispatch_TransientDoesNotReclassify verifies an auth blip or transient +// outage on a cached-Modern backend neither retries nor flips the cache to Legacy +// (F1: status-blind errWrongEra used to poison the cache here). +func TestDispatch_TransientDoesNotReclassify(t *testing.T) { + t.Parallel() + + for _, blip := range []error{errModernAuth, errModernTransient} { + blip := blip + t.Run(blip.Error(), func(t *testing.T) { + t.Parallel() + + // This server WOULD classify Modern if re-probed — proving no re-probe fired. + srv := modernDiscoverServer(t) + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + h.setRevision(target.WorkloadID, mcpparser.RevisionModern) + + attempts := 0 + err := h.dispatch(context.Background(), target, func(_ context.Context, _ mcpparser.Revision) error { + attempts++ + return fmt.Errorf("%w: HTTP blip", blip) + }) + require.Error(t, err) + assert.Equal(t, 1, attempts, "a transient/auth blip must not retry") + + rev, ok := h.cachedRevision(target.WorkloadID) + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev, "a blip must not flip a Modern backend to Legacy") + }) + } +} + +// TestReclassify_WarnsOnlyOnActualChange captures slog to confirm the WARN (which +// gates the reclassification counter in the same branch) fires only when the +// revision actually changes. +// +//nolint:paralleltest // swaps the global slog default; must not run concurrently with other tests +func TestReclassify_WarnsOnlyOnActualChange(t *testing.T) { + // Not parallel: swaps the global slog default. + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + srv := modernDiscoverServer(t) // re-probe classifies Modern + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b", BaseURL: srv.URL, TransportType: "streamable-http"} + + // prev=Legacy -> re-probe Modern: change, must WARN. + got := h.reclassify(context.Background(), target, mcpparser.RevisionLegacy) + assert.Equal(t, mcpparser.RevisionModern, got) + assert.Contains(t, buf.String(), "reclassified", "a real change must WARN") + + // prev=Modern -> re-probe Modern: no change, must be silent. + buf.Reset() + got = h.reclassify(context.Background(), target, mcpparser.RevisionModern) + assert.Equal(t, mcpparser.RevisionModern, got) + assert.Empty(t, buf.String(), "a no-op re-probe must not WARN") +} + +// TestListCapabilities_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/client/revision_test.go b/pkg/vmcp/client/revision_test.go new file mode 100644 index 0000000000..260e6c92b3 --- /dev/null +++ b/pkg/vmcp/client/revision_test.go @@ -0,0 +1,224 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" + "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" + authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" +) + +// newProbeClient builds a real httpBackendClient with an unauthenticated +// registry so probeRevision's buildBackendRoundTripper succeeds. +func newProbeClient(t *testing.T) *httpBackendClient { + t.Helper() + reg := vmcpauth.NewDefaultOutgoingAuthRegistry() + require.NoError(t, reg.RegisterStrategy(authtypes.StrategyTypeUnauthenticated, &strategies.UnauthenticatedStrategy{})) + c, err := NewHTTPBackendClient(reg) + require.NoError(t, err) + return c.(*httpBackendClient) +} + +// discoverEnvelope is a valid Modern server/discover success body echoing the +// request id. +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, + }, + { + // A valid Modern envelope with a non-"complete" resultType proves the peer + // is Modern (it decoded), so it must NOT fall back to Legacy. + name: "input_required envelope -> Modern", + handler: func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + writeModernResult(t, w, id, map[string]any{"resultType": "input_required"}) + }, + wantRev: mcpparser.RevisionModern, + }, + { + name: "discover -32601 (method not found) -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"method not found"}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "400 session required (-32600) -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"session required"}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "405 method not allowed -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "empty body -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + }, + wantRev: mcpparser.RevisionLegacy, + }, + { + name: "200 with Legacy-shaped result (no resultType) -> Legacy", + handler: func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}`)) + }, + wantRev: mcpparser.RevisionLegacy, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(tt.handler) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "b1", BaseURL: srv.URL, TransportType: "streamable-http"} + + rev, err := h.probeRevision(context.Background(), target) + require.NoError(t, err) + assert.Equal(t, tt.wantRev, rev) + + // The result is cached under the workload id. + cached, ok := h.cachedRevision("b1") + require.True(t, ok) + assert.Equal(t, tt.wantRev, cached) + }) + } +} + +// TestProbeRevision_TransientLeavesUnprobed verifies a dead backend (connection +// refused) is INCONCLUSIVE: probeRevision returns the error uncached rather than +// caching Legacy, so a transient outage cannot poison the revision cache and the +// next call re-probes. +func TestProbeRevision_TransientLeavesUnprobed(t *testing.T) { + t.Parallel() + + // A server we immediately close: connections are refused. + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + url := srv.URL + srv.Close() + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "dead", BaseURL: url, TransportType: "streamable-http"} + + _, err := h.probeRevision(context.Background(), target) + require.Error(t, err) + + _, ok := h.cachedRevision("dead") + assert.False(t, ok, "a transient probe failure must leave the backend unprobed") +} + +// TestListCapabilities_ModernServedFromCache verifies the cache: a Modern +// backend is probed once, and a second ListCapabilities is served from the +// cached revision (discover + enumerate, never a Legacy initialize handshake). +func TestListCapabilities_ModernServedFromCache(t *testing.T) { + t.Parallel() + + var initializeCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id, _ := modernReq(t, r) + switch r.Header.Get("Mcp-Method") { + case "server/discover": + writeModernResult(t, w, id, map[string]any{"capabilities": map[string]any{"tools": map[string]any{}}}) + case "tools/list": + writeModernResult(t, w, id, map[string]any{ + "tools": []any{map[string]any{"name": "echo", "inputSchema": map[string]any{"type": "object"}}}, + }) + default: + // Any non-Modern method (e.g. initialize) is a regression. + initializeCalls.Add(1) + w.WriteHeader(http.StatusNotFound) + } + })) + t.Cleanup(srv.Close) + + h := newProbeClient(t) + target := &vmcp.BackendTarget{WorkloadID: "modern", BaseURL: srv.URL, TransportType: "streamable-http"} + + caps1, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.Len(t, caps1.Tools, 1) + + rev, ok := h.cachedRevision("modern") + require.True(t, ok) + assert.Equal(t, mcpparser.RevisionModern, rev) + + // Second call is served via the cached Modern revision (discover+enumerate), + // not a re-probe ladder, and still returns the enumerated tool. + caps2, err := h.ListCapabilities(context.Background(), target) + require.NoError(t, err) + require.Len(t, caps2.Tools, 1) + assert.Equal(t, "echo", caps2.Tools[0].Name) + + assert.Zero(t, initializeCalls.Load(), "a Modern backend must never receive a Legacy initialize") +} diff --git a/pkg/vmcp/headerforward/transport.go b/pkg/vmcp/headerforward/transport.go index cf2c9cfcb3..e87710f197 100644 --- a/pkg/vmcp/headerforward/transport.go +++ b/pkg/vmcp/headerforward/transport.go @@ -64,6 +64,15 @@ func (h *headerForwardRoundTripper) RoundTrip(req *http.Request) (*http.Response return h.base.RoundTrip(reqCopy) } +// CloseIdleConnections forwards to the wrapped RoundTripper so http.Client's +// CloseIdleConnections reaches the concrete *http.Transport at the bottom of the +// chain (this wrapper would otherwise silently swallow the call). +func (h *headerForwardRoundTripper) CloseIdleConnections() { + if c, ok := h.base.(interface{ CloseIdleConnections() }); ok { + c.CloseIdleConnections() + } +} + // BuildHeaderForwardTripper constructs a headerForwardRoundTripper for the // backend's pre-resolved HeaderForwardConfig. Returns base unchanged when no // header injection is configured or the effective header set is empty. diff --git a/pkg/vmcp/headerforward/transport_test.go b/pkg/vmcp/headerforward/transport_test.go index c8f2b42766..464026709c 100644 --- a/pkg/vmcp/headerforward/transport_test.go +++ b/pkg/vmcp/headerforward/transport_test.go @@ -398,3 +398,22 @@ func TestMergeForwardedHeaders_RestrictedHeadersList(t *testing.T) { } } } + +// closeIdleSpy is a RoundTripper that records CloseIdleConnections calls. +type closeIdleSpy struct{ closed int } + +func (*closeIdleSpy) RoundTrip(*http.Request) (*http.Response, error) { + return nil, errors.New("unused") +} +func (s *closeIdleSpy) CloseIdleConnections() { s.closed++ } + +// TestHeaderForwardRoundTripper_CloseIdleConnections verifies the wrapper forwards +// CloseIdleConnections to its base rather than swallowing it. +func TestHeaderForwardRoundTripper_CloseIdleConnections(t *testing.T) { + t.Parallel() + + spy := &closeIdleSpy{} + rt := &headerForwardRoundTripper{base: spy} + rt.CloseIdleConnections() + assert.Equal(t, 1, spy.closed) +} diff --git a/pkg/vmcp/health/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 dae74b5f67..7a683386ee 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry.go @@ -13,23 +13,61 @@ package backendtelemetry import ( "context" "fmt" + "sync" "time" + "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/telemetry" transporttypes "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/pkg/vmcp" ) +// 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" ) +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( @@ -103,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 { @@ -153,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 0c96135357..5a43f7605e 100644 --- a/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go +++ b/pkg/vmcp/internal/backendtelemetry/backendtelemetry_test.go @@ -4,9 +4,60 @@ package backendtelemetry import ( + "context" "testing" + + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" ) +// fakeRevClient embeds vmcp.BackendClient (nil — its methods are never called +// here) and adds the optional CachedRevision accessor. +type fakeRevClient struct { + vmcp.BackendClient + rev mcpparser.Revision + ok bool +} + +func (f fakeRevClient) CachedRevision(string) (mcpparser.Revision, bool) { return f.rev, f.ok } + +// fakeNoRevClient embeds vmcp.BackendClient but does NOT implement revisionReporter. +type fakeNoRevClient struct{ vmcp.BackendClient } + +// TestTelemetryBackendClient_CachedRevisionForwarding verifies the decorator +// forwards CachedRevision to a client that reports it, and reports nothing for a +// client that doesn't. +func TestTelemetryBackendClient_CachedRevisionForwarding(t *testing.T) { + t.Parallel() + + d := telemetryBackendClient{backendClient: fakeRevClient{rev: mcpparser.RevisionModern, ok: true}} + rev, ok := d.CachedRevision("b") + if !ok || rev != mcpparser.RevisionModern { + t.Fatalf("CachedRevision = (%v, %v), want (Modern, true)", rev, ok) + } + if got := d.revisionLabel("b"); got != "2026-07-28" { + t.Errorf("revisionLabel = %q, want 2026-07-28", got) + } + + dn := telemetryBackendClient{backendClient: fakeNoRevClient{}} + if _, ok := dn.CachedRevision("b"); ok { + t.Error("CachedRevision should report false for a client without the accessor") + } + if got := dn.revisionLabel("b"); got != "" { + t.Errorf("revisionLabel = %q, want empty for unprobed/unsupported", got) + } +} + +// TestRecordRevisionReclassification is a smoke test: the counter lazily binds to +// the global meter provider and increments without panicking (the noop provider +// makes the value unobservable here — the WARN in the same reclassify branch is +// asserted in the client package's reclassify test). +func TestRecordRevisionReclassification(t *testing.T) { + t.Parallel() + RecordRevisionReclassification(context.Background()) + RecordRevisionReclassification(context.Background()) +} + func TestMapActionToMCPMethod(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/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") +} 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.