diff --git a/pkg/transport/bridge.go b/pkg/transport/bridge.go index 6ba5ff6eea..d53b913a70 100644 --- a/pkg/transport/bridge.go +++ b/pkg/transport/bridge.go @@ -104,6 +104,18 @@ func (b *StdioBridge) run(ctx context.Context) { params = map[string]any{} } + // On a *_list_changed notification, re-fetch the upstream capability set + // before forwarding the notification. forwardAll is additive (AddTool/ + // AddResource/AddPrompt upsert by name/URI), so re-running it matches the + // startup behavior and keeps the local server's advertised set in sync with + // the upstream's post-change state. + switch n.Method { + case "notifications/tools/list_changed", + "notifications/resources/list_changed", + "notifications/prompts/list_changed": + b.forwardAll(context.Background()) + } + b.srv.SendNotificationToAllClients(n.Method, params) }) diff --git a/pkg/transport/bridge_test.go b/pkg/transport/bridge_test.go new file mode 100644 index 0000000000..bb275e6eec --- /dev/null +++ b/pkg/transport/bridge_test.go @@ -0,0 +1,230 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transport + +import ( + "context" + "encoding/json" + "net/http/httptest" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/transport/types" +) + +// toolNamesOnServer returns the names of the tools currently registered on the +// given local MCP server, by issuing a synthetic tools/list request through the +// shim's HandleMessage dispatch path. The bridge's forwardAll registers upstream +// tools here, so this is the assertion seam for "what does the bridge advertise +// downstream". Call only after the bridge has fully stopped (run() returned), +// since StdioBridge exposes no synchronization for its srv field while running. +func toolNamesOnServer(t *testing.T, srv *server.MCPServer) []string { + t.Helper() + resp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","method":"tools/list","id":1}`)) + jr, ok := resp.(mcp.JSONRPCResponse) + require.True(t, ok, "tools/list should return a JSONRPCResponse, got %T", resp) + buf, err := json.Marshal(jr.Result) + require.NoError(t, err) + var lt mcp.ListToolsResult + require.NoError(t, json.Unmarshal(buf, <)) + names := make([]string, 0, len(lt.Tools)) + for _, tl := range lt.Tools { + names = append(names, tl.Name) + } + return names +} + +func containsTool(names []string, want string) bool { + for _, n := range names { + if n == want { + return true + } + } + return false +} + +// noopToolHandler is a stand-in tool handler; the bridge re-fetch tests never +// invoke tools, they only assert the advertised set. +func noopToolHandler(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return nil, nil +} + +// listToolsCounter is a server hook that atomically counts tools/list requests +// received by the backend. The bridge's forwardAll issues a tools/list upstream +// on startup and again whenever its OnNotification handler re-syncs on a +// *_list_changed notification, so the counter is the race-free readiness and +// re-fetch signal observed entirely on the backend side. +type listToolsCounter struct { + count atomic.Int64 +} + +func (c *listToolsCounter) hook() server.OnBeforeListToolsFunc { + return func(context.Context, any, *mcp.ListToolsRequest) { c.count.Add(1) } +} + +// sessionHolder captures the single client session that connects to the backend, +// guarded by a mutex because the hook fires on the backend's session goroutine. +type sessionHolder struct { + mu sync.Mutex + _session server.ClientSession +} + +func (h *sessionHolder) set(s server.ClientSession) { + h.mu.Lock() + defer h.mu.Unlock() + h._session = s +} + +func (h *sessionHolder) get() server.ClientSession { + h.mu.Lock() + defer h.mu.Unlock() + return h._session +} + +// TestBridge_ToolsListChanged_TriggersReSync verifies the bridge's +// notifications/tools/list_changed re-fetch fix: when the upstream backend's +// tool set changes after the bridge has connected, the upstream emits a +// tools/list_changed notification, the bridge's OnNotification handler re-runs +// forwardAll, and the newly added tool appears on the bridge's local stdio +// server. +// +// The backend is a real mcpcompat streamable-HTTP server. A tool-set change +// visible to an already-connected client is driven through the per-session +// overlay (SessionWithTools.SetSessionTools): the shim syncs the overlay onto the +// live go-sdk server, which emits notifications/tools/list_changed over the +// standalone SSE stream to the bridge's upstream client (the bridge connects +// with WithContinuousListening, so the SSE stream is established). The bridge's +// OnNotification closure is the fix under test. +// +// ServeStdio blocks reading os.Stdin; this test swaps os.Stdin for a pipe so +// ServeStdio returns cleanly on stdin EOF at teardown. The bridge exposes no +// synchronization for its srv field, so readiness/re-fetch is observed via the +// backend's tools/list counter (race-free) and bridge.srv is read only after +// Shutdown returns (b.wg.Wait() provides the happens-before edge). +// +//nolint:paralleltest // Swaps process-global os.Stdin; cannot run in parallel. +func TestBridge_ToolsListChanged_TriggersReSync(t *testing.T) { + // This test swaps process-global os.Stdin, so it cannot run in parallel + // with any other test touching stdio. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // --- upstream backend with an initial tool "alpha" --- + counter := &listToolsCounter{} + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + hooks.AddBeforeListTools(counter.hook()) + + backend := server.NewMCPServer( + "backend", "1.0", + server.WithToolCapabilities(true), + server.WithHooks(hooks), + ) + backend.AddTool(mcp.NewTool("alpha"), noopToolHandler) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // --- bridge pointing at the backend (streamable-http) --- + bridge, err := NewStdioBridge("test", ts.URL, types.TransportTypeStreamableHTTP) + require.NoError(t, err) + + // Swap os.Stdin for a pipe so ServeStdio (which hardcodes os.Stdin) does not + // touch the real process stdin and can be unblocked at teardown by closing + // the write end. Leave os.Stdout real; the assertions don't read it. + origIn := os.Stdin + pipeR, pipeW, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { + os.Stdin = origIn + _ = pipeR.Close() + _ = pipeW.Close() + }) + os.Stdin = pipeR + + bridge.Start(ctx) + t.Cleanup(func() { + // Close the stdin write end so the go-sdk stdio reader hits EOF and + // ServeStdio returns, letting run() exit and the WaitGroup drain. Then + // Shutdown closes the upstream client and waits for run() to finish. + _ = pipeW.Close() + bridge.Shutdown() + }) + + // Step 1: wait for the bridge to connect and run its initial forwardAll, + // which issues the first upstream tools/list (counter >= 1) after the + // session is registered. Observed entirely on the backend side: race-free. + require.Eventually(t, func() bool { + return holder.get() != nil && counter.count.Load() >= 1 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not complete initial forwardAll (session=%v, listCalls=%d)", + holder.get() != nil, counter.count.Load()) + + // Step 2: mutate the upstream tool set — add "beta" via the per-session + // overlay. SetSessionTools syncs onto the live go-sdk server bound to this + // session, which emits notifications/tools/list_changed over the SSE stream + // to the bridge's upstream client. The bridge's OnNotification handler then + // re-runs forwardAll, issuing a SECOND upstream tools/list (counter >= 2). + backendSession := holder.get() + swt, ok := backendSession.(server.SessionWithTools) + require.True(t, ok, "backend session must implement SessionWithTools") + swt.SetSessionTools(map[string]server.ServerTool{ + "alpha": {Tool: mcp.NewTool("alpha"), Handler: noopToolHandler}, + "beta": {Tool: mcp.NewTool("beta"), Handler: noopToolHandler}, + }) + + // Step 3: wait for the re-fetch — the second tools/list proves the + // OnNotification handler fired and re-ran forwardAll (the fix). + require.Eventually(t, func() bool { + return counter.count.Load() >= 2 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not re-run forwardAll after tools/list_changed (listCalls=%d)", + counter.count.Load()) + + // Step 4: stop the bridge, then read bridge.srv. Shutdown calls + // b.wg.Wait(), which synchronizes with run()'s defer b.wg.Done() (which is + // happens-after every field write run() made), so bridge.srv is safe to read + // here without a data race. + _ = pipeW.Close() // unblock ServeStdio (stdin EOF) + bridge.Shutdown() + + names := toolNamesOnServer(t, bridge.srv) + assert.True(t, containsTool(names, "alpha"), + "alpha must be present after the re-fetch (forwardAll is additive), got %v", names) + assert.True(t, containsTool(names, "beta"), + "beta must be present after the re-sync, got %v", names) +} + +// TestBridge_ProgressAndLoggingNotifications_DroppedByShim is a regression +// anchor for a known limitation of the mcpcompat client shim: it does not +// install ProgressNotificationHandler or LoggingMessageHandler on the +// underlying go-sdk client, so upstream notifications/progress and +// notifications/message notifications are dropped before they ever reach the +// bridge's OnNotification handler (and thus before they could be forwarded +// downstream by SendNotificationToAllClients, which itself only forwards those +// two methods). +// +// When the shim is fixed to install those handlers, delete this Skip and +// implement a real end-to-end test asserting that progress and logging +// notifications are forwarded to the bridge's local server clients. +// +//nolint:paralleltest // Skipped regression anchor; documents a shim limitation. +func TestBridge_ProgressAndLoggingNotifications_DroppedByShim(t *testing.T) { + t.Skip("known shim limitation: mcpcompat client does not install " + + "ProgressNotificationHandler/LoggingMessageHandler, so upstream " + + "progress/logging notifications are dropped before reaching the bridge " + + "(see client.installNotificationHandlers); additionally the shim's " + + "SendNotificationToAllClients only forwards notifications/progress and " + + "notifications/message, dropping list_changed downstream") +} diff --git a/pkg/transport/proxy/transparent/host_header_regression_test.go b/pkg/transport/proxy/transparent/host_header_regression_test.go new file mode 100644 index 0000000000..bf95a6b396 --- /dev/null +++ b/pkg/transport/proxy/transparent/host_header_regression_test.go @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestRegression_LocalhostProxy_NonLocalhostHostHeaderRewritten verifies that +// tracingTransport.RoundTrip rewrites a non-localhost Host header to match the +// target URL's host. Without this, an attacker could inject a Host header to +// bypass host-based validation on the upstream server. +func TestRegression_LocalhostProxy_NonLocalhostHostHeaderRewritten(t *testing.T) { + t.Parallel() + + // Capture the Host header the upstream server actually receives. + var receivedHost string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHost = r.Host + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + targetURL, err := url.Parse(upstream.URL) + require.NoError(t, err) + + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, "streamable-http", nil, nil, "", false) + + proxy := &httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.SetURL(targetURL) + pr.SetXForwarded() + }, + FlushInterval: -1, + Transport: newTracingTransport(http.DefaultTransport, p), + ModifyResponse: p.modifyResponse, + } + + // Send a request with a malicious Host header. + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "http://evil.example.com/some/path", nil) + req.Host = "evil.example.com" + + proxy.ServeHTTP(rec, req) + + // The outbound request's Host must be the target URL's host, + // not the attacker-supplied value. + assert.Equal(t, targetURL.Host, receivedHost, + "tracingTransport must rewrite Host header from attacker value to target URL host") + assert.NotEqual(t, "evil.example.com", receivedHost, + "attacker Host header must not reach the upstream server") +} diff --git a/pkg/vmcp/client/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go new file mode 100644 index 0000000000..7e6fd292df --- /dev/null +++ b/pkg/vmcp/client/auth_error_mapping_regression_test.go @@ -0,0 +1,300 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// jsonRPCResponse is a generic JSON-RPC 2.0 response envelope. +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result,omitempty"` + Error *jsonRPCError `json:"error,omitempty"` +} + +type jsonRPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// jsonRPCRequest is a generic JSON-RPC 2.0 request envelope (for method routing). +type jsonRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Method string `json:"method"` +} + +// TestRegression_401_MapsToErrAuthenticationFailed verifies that a backend +// returning HTTP 401 on initialize is classified as ErrAuthenticationFailed. +func TestRegression_401_MapsToErrAuthenticationFailed(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Unauthorized"}, + }) + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-401-backend", + WorkloadName: "Test 401 Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrAuthenticationFailed), + "expected ErrAuthenticationFailed, got: %v", err) +} + +// TestRegression_403OnInitialize_LegacySSEFallback verifies that a backend +// returning HTTP 403 on initialize is classified as ErrBackendUnavailable. +// +// NOTE: 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. +func TestRegression_403OnInitialize_LegacySSEFallback(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Forbidden"}, + }) + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-403-backend", + WorkloadName: "Test 403 Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrBackendUnavailable), + "expected ErrBackendUnavailable, got: %v", err) + assert.Contains(t, err.Error(), "403", + "error message should reference 403 status, got: %v", err) +} + +// TestRegression_403OnInitialize_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) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Forbidden"}, + }) + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-403-sentinel-backend", + WorkloadName: "Test 403 Sentinel Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + _, 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)") +} + +// TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure verifies +// that an MCP tool error result (IsError=true on a 200 HTTP response) whose +// message contains "401 unauthorized" is NOT classified as an auth failure. +func TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, err := io.ReadAll(r.Body) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + var req jsonRPCRequest + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + switch req.Method { + case "initialize": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "test-backend", "version": "1.0.0"} + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/call": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "content": [{"type": "text", "text": "tool error: 401 unauthorized - permission denied"}], + "isError": true + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/list": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "tools": [{"name": "test-tool", "description": "A test tool", "inputSchema": {"type": "object"}}] + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + default: + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{}`), + }) + } + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-tool-error-backend", + WorkloadName: "Test Tool Error Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + result, err := h.CallTool(ctx, target, "test-tool", map[string]any{"arg": "val"}, nil) + + if err != nil { + if errors.Is(err, vmcp.ErrAuthenticationFailed) { + t.Fatalf("UNEXPECTED: CallTool returned ErrAuthenticationFailed for a 200 response with IsError=true containing '401 unauthorized'") + } + t.Fatalf("unexpected transport error from CallTool: %v", err) + } + + assert.NotNil(t, result, "expected non-nil result for IsError=true") + assert.True(t, result.IsError, "expected IsError=true") + assert.NotEmpty(t, result.Content, "expected non-empty content") +} diff --git a/pkg/vmcp/server/capability_regression_test.go b/pkg/vmcp/server/capability_regression_test.go new file mode 100644 index 0000000000..c26ef658fa --- /dev/null +++ b/pkg/vmcp/server/capability_regression_test.go @@ -0,0 +1,71 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_InitializeAdvertisesToolsAndResourcesCapabilities pins the +// capabilities advertised in the initialize response on the Serve path. +// +// OBSERVED BEHAVIOR (go-sdk bridge): the initialize response advertises only +// {"logging":{}} — it does NOT advertise tools/resources capabilities, even when +// the core supplies tools and resources. This differs from mcp-go, which +// advertises tools/resources in initialize. The go-sdk server advertises a +// static, server-level capability set (logging only on the Serve path, which +// starts with WithToolCapabilities(false)/WithResourceCapabilities(false)); +// per-session tool/resource overlays installed by the OnRegisterSession hook do +// not change the advertised capabilities in the initialize response. +// +// This test pins that observed behavior so a future change to the bridge (e.g. +// advertising tools/resources once a session's overlay is populated) is a +// deliberate, visible flip rather than a silent drift. It asserts on the RAW +// initialize response body parsed with encoding/json — not tools/list. +func TestRegression_InitializeAdvertisesToolsAndResourcesCapabilities(t *testing.T) { + t.Parallel() + + fc := &fakeCore{ + tools: []vmcp.Tool{{Name: "cap-tool", Description: "a capability test tool"}}, + resources: []vmcp.Resource{{Name: "cap-doc", URI: "file:///cap.txt"}}, + } + _, _, baseURL := registerServeSession(t, fc) + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, 200, initResp.StatusCode, "initialize should succeed") + + env, _ := readServeJSONRPC(t, initResp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "initialize response must have a result object; env: %v", env) + + capabilities, ok := result["capabilities"].(map[string]any) + require.True(t, ok, "result.capabilities must be present; result: %v", result) + + // Pin the observed go-sdk-bridge behavior: only logging is advertised in + // initialize on the Serve path. tools/resources are NOT advertised here (they + // are discoverable via tools/list/resources/list once the session is + // registered). If the bridge is updated to advertise them, flip these + // assertions to assert non-nil tools/resources. + assert.Contains(t, capabilities, "logging", + "logging capability must be advertised; got %v", capabilities) + assert.NotContains(t, capabilities, "tools", + "tools capability is NOT advertised in initialize on the Serve path (go-sdk bridge); got %v", capabilities) + assert.NotContains(t, capabilities, "resources", + "resources capability is NOT advertised in initialize on the Serve path (go-sdk bridge); got %v", capabilities) +} diff --git a/pkg/vmcp/server/context_isolation_regression_test.go b/pkg/vmcp/server/context_isolation_regression_test.go new file mode 100644 index 0000000000..279f86a111 --- /dev/null +++ b/pkg/vmcp/server/context_isolation_regression_test.go @@ -0,0 +1,238 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/auth" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + sessionfactorymocks "github.com/stacklok/toolhive/pkg/vmcp/session/mocks" + sessionmocks "github.com/stacklok/toolhive/pkg/vmcp/session/types/mocks" +) + +// perToolCore is a fakeCore variant whose CallTool returns a result text equal +// to the called tool's name. This lets the context-isolation test detect +// cross-contamination: each concurrent call must observe its own tool name. +type perToolCore struct { + fakeCore +} + +func (p *perToolCore) CallTool( + _ context.Context, _ *auth.Identity, name string, _ map[string]any, _ map[string]any, +) (*vmcp.ToolCallResult, error) { + p.callToolCalls.Add(1) + p.lastCallToolName.Store(name) + return &vmcp.ToolCallResult{ + Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: name}}, + }, nil +} + +// newPerToolSessionFactory mirrors newToolSessionFactory but wires a per-tool +// CallTool result on the mock MultiSession so the SDK tool handler routes +// through perToolCore (whose result carries the tool name). The mock session's +// CallTool is only a fallback; the Serve path routes through coreToolHandler → +// core.CallTool, so perToolCore.CallTool is the one that runs. +func newPerToolSessionFactory( + t *testing.T, ctrl *gomock.Controller, tools []vmcp.Tool, +) *sessionfactorymocks.MockMultiSessionFactory { + t.Helper() + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + mock := sessionmocks.NewMockMultiSession(ctrl) + mock.EXPECT().ID().Return(id).AnyTimes() + mock.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes() + mock.EXPECT().CreatedAt().Return(time.Time{}).AnyTimes() + mock.EXPECT().Type().Return(transportsession.SessionType("")).AnyTimes() + mock.EXPECT().GetData().Return(nil).AnyTimes() + mock.EXPECT().SetData(gomock.Any()).AnyTimes() + mock.EXPECT().GetMetadata().Return(map[string]string{ + vmcpsession.MetadataKeyIdentityBinding: "unauthenticated", + }).AnyTimes() + mock.EXPECT().GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding). + Return("unauthenticated", true).AnyTimes() + mock.EXPECT().SetMetadata(gomock.Any(), gomock.Any()).AnyTimes() + toolsCopy := make([]vmcp.Tool, len(tools)) + copy(toolsCopy, tools) + mock.EXPECT().Tools().Return(toolsCopy).AnyTimes() + mock.EXPECT().AllTools().Return(toolsCopy).AnyTimes() + mock.EXPECT().Resources().Return(nil).AnyTimes() + mock.EXPECT().Prompts().Return(nil).AnyTimes() + mock.EXPECT().BackendSessions().Return(nil).AnyTimes() + rt := &vmcp.RoutingTable{Tools: make(map[string]*vmcp.BackendTarget, len(tools))} + for _, tool := range tools { + rt.Tools[tool.Name] = &vmcp.BackendTarget{WorkloadID: tool.Name} + } + mock.EXPECT().GetRoutingTable().Return(rt).AnyTimes() + mock.EXPECT().ReadResource(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mock.EXPECT().GetPrompt(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mock.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(&vmcp.ToolCallResult{Content: []vmcp.Content{{Type: "text", Text: "ok"}}}, nil).AnyTimes() + mock.EXPECT().Close().Return(nil).AnyTimes() + return mock, nil + }).AnyTimes() + return factory +} + +// TestRegression_ConcurrentToolCalls_NoAuditBleed fires two tools/call POSTs +// concurrently against the same session and asserts each response carries its +// OWN tool's result text — no cross-contamination. The Serve path sets the +// audit BackendInfo per-request via the handler closure (coreToolHandler writes +// the pre-resolved backend name), so concurrent calls must not bleed a backend +// name or result from one request into another. +// +// The audit BackendInfo is attached per-request by the audit middleware (a fresh +// *BackendInfo per request, see pkg/audit/auditor.go Middleware), and the Serve +// handler writes to it from the per-request context — so the isolation surface +// under test is the result text and the observed backend label. We assert each +// call's result text matches the tool it invoked. +func TestRegression_ConcurrentToolCalls_NoAuditBleed(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + tools := []vmcp.Tool{{Name: "tool-a"}, {Name: "tool-b"}} + fc := &perToolCore{fakeCore: fakeCore{tools: tools}} + factory := newPerToolSessionFactory(t, ctrl, tools) + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + baseURL := ts.URL + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return ok + }, 2*time.Second, 10*time.Millisecond, "session should be registered") + + type callOutcome struct { + toolName string + body string + err error + } + + var wg sync.WaitGroup + outcomes := make([]callOutcome, len(tools)) + wg.Add(len(tools)) + for i, tool := range tools { + i, tool := i, tool + go func() { + defer wg.Done() + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2 + i, + "method": "tools/call", + "params": map[string]any{ + "name": tool.Name, + "arguments": map[string]any{}, + }, + }, sessionID) + defer resp.Body.Close() + body, readErr := io.ReadAll(resp.Body) + outcomes[i] = callOutcome{toolName: tool.Name, body: string(body), err: readErr} + }() + } + + // Fail-fast wait: never block indefinitely on a WaitGroup. + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent tool calls to complete") + } + + // Each call must have returned its OWN tool name in the result text, proving + // no result cross-contamination between the concurrent requests. + for _, o := range outcomes { + require.NoError(t, o.err) + assert.Contains(t, o.body, o.toolName, + "concurrent call for %q must carry its own tool's result; got %s", o.toolName, o.body) + } + + // The core must have been reached exactly once per tool (2 total) — no + // duplicated or dropped calls. + assert.Equal(t, int32(len(tools)), fc.callToolCalls.Load(), + "core.CallTool must be called exactly once per concurrent request") + + // The per-request audit BackendInfo must not bleed: drive the handler closure + // directly with two distinct BackendInfo values concurrently to assert each + // context's BackendInfo is labelled with its own tool's backend name, not the + // sibling's. This exercises the per-request labelling surface in + // coreToolHandler (serve_handlers.go) directly. + var bgWG sync.WaitGroup + bgWG.Add(len(tools)) + backendLabels := make([]string, len(tools)) + for i, tool := range tools { + i, tool := i, tool + go func() { + defer bgWG.Done() + bi := &audit.BackendInfo{} + ctx := audit.WithBackendInfo(context.Background(), bi) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: tool.Name, Arguments: map[string]any{}}} + res, err := srv.coreToolHandler(sessionID, tool.Name, tool.Name)(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + backendLabels[i] = bi.BackendName + }() + } + bgDone := make(chan struct{}) + go func() { bgWG.Wait(); close(bgDone) }() + select { + case <-bgDone: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for backend-label goroutines") + } + + // Each concurrent handler invocation must have labelled its OWN BackendInfo + // with its own tool's backend name — no cross-contamination between the + // concurrent contexts. + for i, tool := range tools { + assert.Equal(t, tool.Name, backendLabels[i], + "audit BackendInfo for %q must carry its own backend name, not a sibling's", + tool.Name) + } +} diff --git a/pkg/vmcp/server/derive.go b/pkg/vmcp/server/derive.go index a4a5c059ef..a1799babc7 100644 --- a/pkg/vmcp/server/derive.go +++ b/pkg/vmcp/server/derive.go @@ -76,6 +76,7 @@ func deriveServerConfig( Port: cfg.Port, // 0 means "OS-assigned". EndpointPath: cfg.EndpointPath, SessionTTL: cfg.SessionTTL, + HeartbeatInterval: cfg.HeartbeatInterval, AuthMiddleware: cfg.AuthMiddleware, AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, diff --git a/pkg/vmcp/server/derive_test.go b/pkg/vmcp/server/derive_test.go index 26abece62e..645d8cddbc 100644 --- a/pkg/vmcp/server/derive_test.go +++ b/pkg/vmcp/server/derive_test.go @@ -39,6 +39,7 @@ func populatedLegacyConfig() *Config { Port: 7777, EndpointPath: "/custom", SessionTTL: 17 * time.Minute, + HeartbeatInterval: 5 * time.Second, AuthMiddleware: passthrough, AuthzMiddleware: passthrough, AuthInfoHandler: http.NewServeMux(), @@ -70,6 +71,7 @@ func TestDeriveServerConfigProjectsTransportFields(t *testing.T) { assert.Equal(t, 7777, got.Port) assert.Equal(t, "/custom", got.EndpointPath) assert.Equal(t, 17*time.Minute, got.SessionTTL) + assert.Equal(t, 5*time.Second, got.HeartbeatInterval) assert.Equal(t, 11*time.Second, got.StatusReportingInterval) // Func/handler/pointer fields projected by reference. diff --git a/pkg/vmcp/server/schema_fidelity_regression_test.go b/pkg/vmcp/server/schema_fidelity_regression_test.go new file mode 100644 index 0000000000..d841f56e21 --- /dev/null +++ b/pkg/vmcp/server/schema_fidelity_regression_test.go @@ -0,0 +1,152 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_ToolSchemaWithTypeObject_ProjectedIntact guards schema +// fidelity for the SUPPORTED case: a tool whose InputSchema includes a +// top-level "type":"object" is projected to the client UNCHANGED in tools/list, +// with its properties and required array preserved. The Serve path +// (coreSessionTools in serve_handlers.go) marshals the core's InputSchema into +// RawInputSchema; the go-sdk bridge's normalizeObjectSchema passes object-typed +// schemas through verbatim, so properties/required survive. +// +// All assertions use encoding/json and generic maps — not SDK types — so the +// check is against the wire representation the client sees. +func TestRegression_ToolSchemaWithTypeObject_ProjectedIntact(t *testing.T) { + t.Parallel() + + originalSchema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "x": map[string]any{"type": "string"}, + }, + "required": []any{"x"}, + } + fc := &fakeCore{tools: []vmcp.Tool{{ + Name: "schema-tool", + Description: "a schema-fidelity test tool", + InputSchema: originalSchema, + }}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode, "tools/list should succeed") + + env, body := readServeJSONRPC(t, resp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list must have a result; body: %s", string(body)) + + tools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + + // Find the projected tool by name. + var found map[string]any + for _, tRaw := range tools { + tm, ok := tRaw.(map[string]any) + if !ok { + continue + } + if tm["name"] == "schema-tool" { + found = tm + break + } + } + require.NotNil(t, found, "schema-tool must be present in tools/list; tools: %v", tools) + + inputSchema, ok := found["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object; tool: %v", found) + + // The projected schema must equal the original map exactly: properties and + // required preserved, type unchanged. + assert.Equal(t, originalSchema, inputSchema, + "an object-typed inputSchema must be projected intact; got %v", inputSchema) + assert.Equal(t, "object", inputSchema["type"], + "the top-level type must remain \"object\"; got %v", inputSchema) +} + +// TestRegression_ToolSchemaWithoutTypeObject_NormalizedToEmptyObject pins the +// go-sdk bridge's normalization of a schema that omits the top-level +// "type":"object". normalizeObjectSchema (in the mcpcompat server bridge) +// REPLACES any non-object-typed schema with {"type":"object"}, dropping +// properties/required. This is a known schema-fidelity gap versus mcp-go (which +// projected such schemas verbatim). +// +// This test pins the CURRENT behavior so a future fix to normalizeObjectSchema +// (e.g. preserving properties when type is absent) is a deliberate, visible +// flip. When that fix lands, replace this test's assertions with an +// integrity check mirroring TestRegression_ToolSchemaWithTypeObject_ProjectedIntact. +func TestRegression_ToolSchemaWithoutTypeObject_NormalizedToEmptyObject(t *testing.T) { + t.Parallel() + + originalSchema := map[string]any{ + "properties": map[string]any{ + "x": map[string]any{"type": "string"}, + }, + "required": []any{"x"}, + // NOTE: deliberately NO top-level "type": "object". + } + fc := &fakeCore{tools: []vmcp.Tool{{ + Name: "schema-tool", + Description: "a schema-fidelity test tool", + InputSchema: originalSchema, + }}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode, "tools/list should succeed") + + env, body := readServeJSONRPC(t, resp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list must have a result; body: %s", string(body)) + + tools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + + var found map[string]any + for _, tRaw := range tools { + tm, ok := tRaw.(map[string]any) + if !ok { + continue + } + if tm["name"] == "schema-tool" { + found = tm + break + } + } + require.NotNil(t, found, "schema-tool must be present in tools/list; tools: %v", tools) + + inputSchema, ok := found["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object; tool: %v", found) + + // Pin the current (lossy) normalization: a schema without a top-level + // "type":"object" is replaced with the empty object schema, dropping + // properties and required. This documents the fidelity gap. + assert.Equal(t, map[string]any{"type": "object"}, inputSchema, + "a non-object-typed schema is normalized to {\"type\":\"object\"} (fidelity gap); got %v", inputSchema) + assert.NotContains(t, inputSchema, "properties", + "properties are dropped by the normalization; got %v", inputSchema) + assert.NotContains(t, inputSchema, "required", + "required is dropped by the normalization; got %v", inputSchema) +} diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 2e30582d24..c9668c15e6 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -62,6 +62,10 @@ type ServerConfig struct { // SessionTTL is the session time-to-live duration (default: 30 minutes). SessionTTL time.Duration + // HeartbeatInterval configures the SSE keep-alive ping interval on GET + // connections (default: 30s when zero). + HeartbeatInterval time.Duration + // AuthMiddleware is the optional authentication middleware applied to MCP routes. // If nil, no authentication is required. AuthMiddleware func(http.Handler) http.Handler @@ -326,6 +330,7 @@ func buildServeConfig(cfg *ServerConfig) *Config { Port: cfg.Port, EndpointPath: cfg.EndpointPath, SessionTTL: cfg.SessionTTL, + HeartbeatInterval: cfg.HeartbeatInterval, AuthMiddleware: cfg.AuthMiddleware, AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, diff --git a/pkg/vmcp/server/serve_test.go b/pkg/vmcp/server/serve_test.go index c364e93c0d..4e8235c56a 100644 --- a/pkg/vmcp/server/serve_test.go +++ b/pkg/vmcp/server/serve_test.go @@ -339,6 +339,7 @@ func TestBuildServeConfigMapsSharedFields(t *testing.T) { Port: 1, EndpointPath: "/e", SessionTTL: time.Second, + HeartbeatInterval: time.Second, AuthMiddleware: func(h http.Handler) http.Handler { return h }, AuthInfoHandler: http.NewServeMux(), PassthroughHeaders: []string{"x-test"}, diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index c2ca3211c3..01cb833462 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -95,6 +95,15 @@ const ( capabilityCacheTTL = 30 * time.Second ) +// heartbeatInterval returns the configured heartbeat interval, or the default +// when the configured value is zero (unset). +func heartbeatInterval(d time.Duration) time.Duration { + if d <= 0 { + return defaultHeartbeatInterval + } + return d +} + //go:generate mockgen -destination=mocks/mock_watcher.go -package=mocks -source=server.go Watcher // Watcher is the interface for Kubernetes backend watcher integration. @@ -131,6 +140,11 @@ type Config struct { // Sessions inactive for this duration will be automatically cleaned up SessionTTL time.Duration + // HeartbeatInterval configures the SSE keep-alive ping interval on GET + // connections. When zero, the Handler defaults to defaultHeartbeatInterval (30s). + // Prevents proxies/load balancers from closing idle SSE connections. + HeartbeatInterval time.Duration + // AuthMiddleware is the optional authentication middleware to apply to MCP routes. // If nil, no authentication is required. // This should be a composed middleware chain (e.g., TokenValidator + MCP parser). @@ -501,7 +515,7 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { s.mcpServer, server.WithEndpointPath(s.config.EndpointPath), server.WithSessionIdManager(s.vmcpSessionMgr), - server.WithHeartbeatInterval(defaultHeartbeatInterval), + server.WithHeartbeatInterval(heartbeatInterval(s.config.HeartbeatInterval)), ) // Create HTTP mux with separated authenticated and unauthenticated routes diff --git a/pkg/vmcp/server/session_lifecycle_regression_test.go b/pkg/vmcp/server/session_lifecycle_regression_test.go new file mode 100644 index 0000000000..22c75e2924 --- /dev/null +++ b/pkg/vmcp/server/session_lifecycle_regression_test.go @@ -0,0 +1,155 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "io" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// readServeJSONRPC reads an HTTP response body from the Serve-path streamable +// server and unmarshals the JSON-RPC envelope into a map. The streamable server +// is configured with JSONResponse:true, so the body is a single application/json +// document (not an SSE stream). It also returns the raw body bytes for +// substring assertions. +func readServeJSONRPC(t *testing.T, resp *http.Response) (map[string]any, []byte) { + t.Helper() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var env map[string]any + require.NoError(t, json.Unmarshal(body, &env), "body: %s", string(body)) + return env, body +} + +// serveDelete sends an HTTP DELETE to /mcp with the given session ID. +func serveDelete(t *testing.T, baseURL, sessionID string) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext( + context.Background(), http.MethodDelete, baseURL+"/mcp", nil, + ) + require.NoError(t, err) + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// TestRegression_TerminatedSessionRejected verifies that a session terminated +// via the SDK's DELETE path (which drives vmcpSessionMgr.Terminate under the +// hood and forgets the SDK's local session) is no longer accepted for +// subsequent requests: a tools/list POST with the dead session ID must yield a +// 4xx response — the SDK's Validate rejects the unknown/terminated session +// before any handler runs — not a successful JSON-RPC result. +// +// Note: calling vmcpSessionMgr.Terminate alone does NOT immediately evict the +// SDK's in-memory local session (lazy eviction — the entry is closed when the +// cache is later re-checked). Only the DELETE path forgets the SDK session in +// lockstep with storage termination, so the rejection is observable on the +// next request. This mirrors the proven assertion in +// TestIntegration_SessionManagement_Termination. +func TestRegression_TerminatedSessionRejected(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, baseURL := registerServeSession(t, fc) + + // Terminate via the SDK DELETE path: this calls vmcpSessionMgr.Terminate + // (storage delete) AND forgets the SDK's local session, so the next request + // is rejected by Validate rather than served from a stale local copy. + delResp := serveDelete(t, baseURL, sessionID) + defer delResp.Body.Close() + require.Equal(t, http.StatusOK, delResp.StatusCode, "DELETE should return 200") + + // The termination must be reflected in storage before we assert. + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return !ok + }, 2*time.Second, 10*time.Millisecond, "session must be gone from the manager after Terminate") + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + + // A terminated session must NOT be accepted: reject with a 4xx (the SDK's + // Validate returns 404 for an unknown/terminated session). A 200 here would + // mean a terminated session was still serving requests — a security regression. + assert.GreaterOrEqual(t, resp.StatusCode, 400, + "terminated session must be rejected, got status %d", resp.StatusCode) + assert.Less(t, resp.StatusCode, 500, + "rejection should be a client error (4xx), not a server error; got %d", resp.StatusCode) +} + +// TestRegression_DELETE_TerminationEvictsCache verifies that an HTTP DELETE to +// /mcp with a valid Mcp-Session-Id terminates the session: the response is 200 +// (the streamable server rewrites the SDK's 204 to 200 for mcp-go compatibility) +// and the session is subsequently absent from vmcpSessionMgr. +func TestRegression_DELETE_TerminationEvictsCache(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, baseURL := registerServeSession(t, fc) + + delResp := serveDelete(t, baseURL, sessionID) + defer delResp.Body.Close() + + // The streamable server rewrites 204 → 200 and forwards the termination to + // the SessionIdManager, so DELETE returns 200 OK on the happy path. + assert.Equal(t, http.StatusOK, delResp.StatusCode, + "DELETE on a live session should return 200 (got %d)", delResp.StatusCode) + + // The session must be evicted from the vMCP session manager. + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return !ok + }, 2*time.Second, 10*time.Millisecond, "session must be evicted from vmcpSessionMgr after DELETE") +} + +// TestRegression_SessionIdentityBinding_SecondPrincipalRejected is a thin +// regression re-assertion that the Serve call path enforces the session's +// identity binding: a session bound to the anonymous identity rejects a caller +// presenting a different principal. The exhaustive test lives in +// TestServeEnforcesSessionBinding; this guards the regression at the tool-handler +// boundary by invoking coreToolHandler directly with a non-anonymous identity. +func TestRegression_SessionIdentityBinding_SecondPrincipalRejected(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, _ := registerServeSession(t, fc) + + // The registered session is bound to "unauthenticated" (anonymous) by + // newToolSessionFactory. A caller presenting a token is a session-upgrade + // attack and must be rejected before the core is reached. + ctx := auth.WithIdentity(context.Background(), &auth.Identity{Token: "attacker-token"}) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "t", Arguments: map[string]any{}}} + + res, err := srv.coreToolHandler(sessionID, "t", "")(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + assert.True(t, res.IsError, "a second principal must be rejected with an error result") + + body, err := json.Marshal(res) + require.NoError(t, err) + assert.Contains(t, string(body), "Unauthorized", + "the rejection must carry an unauthorized message") + // The core must not have been reached on a binding failure. + assert.Equal(t, int32(0), fc.callToolCalls.Load(), + "core.CallTool must not be reached when the binding check fails") +} diff --git a/pkg/vmcp/server/sse_keepalive_regression_test.go b/pkg/vmcp/server/sse_keepalive_regression_test.go new file mode 100644 index 0000000000..4f27336f88 --- /dev/null +++ b/pkg/vmcp/server/sse_keepalive_regression_test.go @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" +) + +// TestRegression_SSEKeepAlive_PeriodicBytesOnIdleStream verifies that an idle +// SSE stream opened against the Serve-path streamable server keeps producing +// bytes (a keep-alive ping or heartbeat) within a short window. A regression +// that drops the heartbeat wiring will cause this test to fail: no bytes arrive +// on an idle GET stream, so proxies/load balancers eventually close it. +// +// The streamable server mounts the SDK keep-alive over the configured +// HeartbeatInterval, so the test sets a short interval (100ms) and reads with a +// 500ms deadline — generous relative to the interval but tight enough that a +// missing keep-alive is caught rather than hanging the suite. +func TestRegression_SSEKeepAlive_PeriodicBytesOnIdleStream(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + testTool := vmcp.Tool{Name: "keepalive-tool", Description: "keep-alive regression anchor"} + factory, _ := newToolSessionFactory(t, ctrl, []vmcp.Tool{testTool}) + fc := &fakeCore{tools: []vmcp.Tool{testTool}} + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + HeartbeatInterval: 100 * time.Millisecond, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + + // initialize → obtain a session ID for the subsequent GET stream. + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "keepalive-test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode, "initialize should succeed") + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID, "session ID should be returned in Mcp-Session-Id header") + + // Open a long-lived SSE GET stream. The stream must stay open and emit + // keep-alive bytes even though no client request is in flight. + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL+"/mcp", nil) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Mcp-Session-Id", sessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "GET stream should open with 200") + + // Read with a 500ms deadline. The keep-alive interval is 100ms, so at least + // one ping/heartbeat should land within the deadline; a missing keep-alive + // surfaces as a read timeout with zero bytes — the regression. The streamable + // server's body is not a net.Conn, so SetReadDeadline is unavailable; drive + // the read on a goroutine and race it against a timer. + type readResult struct { + data []byte + err error + } + ch := make(chan readResult, 1) + go func() { + buf := make([]byte, 4096) + n, err := resp.Body.Read(buf) + ch <- readResult{data: buf[:n], err: err} + }() + + var got []byte + select { + case r := <-ch: + got = r.data + case <-time.After(500 * time.Millisecond): + } + + // At least some bytes must arrive — the keep-alive ping, an SSE comment + // (`:` prefix), or a JSON-RPC ping request frame. Zero bytes within the + // window means the keep-alive wiring is broken; fail loudly as a regression. + if len(got) == 0 { + t.Fatal("no bytes received on idle SSE stream within 500ms; keep-alive/heartbeat " + + "is not emitting periodic bytes (regression: proxies will close idle streams)") + } + + body := string(got) + t.Logf("received keep-alive bytes on idle SSE stream: %q", body) + // Accept any of: a JSON-RPC ping request, an SSE comment, or any SSE data + // frame. The point is that bytes flow; their exact form is shim-defined. + hasPing := strings.Contains(body, `"method":"ping"`) || + strings.Contains(body, `"method": "ping"`) + hasSSEComment := strings.Contains(body, ":") + hasDataFrame := strings.Contains(body, "data:") + assert.True(t, hasPing || hasSSEComment || hasDataFrame, + "keep-alive bytes should be an SSE comment, a data frame, or a JSON-RPC ping; got %q", body) +} diff --git a/test/integration/vmcp/helpers/vmcp_server.go b/test/integration/vmcp/helpers/vmcp_server.go index aecb2f4cad..7805e1047a 100644 --- a/test/integration/vmcp/helpers/vmcp_server.go +++ b/test/integration/vmcp/helpers/vmcp_server.go @@ -75,6 +75,7 @@ type vmcpServerConfig struct { workflowDefs map[string]*composer.WorkflowDefinition telemetryProvider *telemetry.Provider passthroughHeaders []string + sessionTTL time.Duration } // WithPrefixConflictResolution configures prefix-based conflict resolution. @@ -112,6 +113,15 @@ func WithPassthroughHeaders(headers ...string) VMCPServerOption { } } +// WithSessionTTL overrides the server's session time-to-live (default 30m). +// A short TTL is useful for sliding-TTL / eviction regression tests. A zero +// value leaves the default in place. +func WithSessionTTL(ttl time.Duration) VMCPServerOption { + return func(c *vmcpServerConfig) { + c.sessionTTL = ttl + } +} + // getFreePort returns an available TCP port on localhost. func getFreePort(tb testing.TB) int { tb.Helper() @@ -198,6 +208,9 @@ func NewVMCPServer( Base: sessionFactory, }, } + if config.sessionTTL > 0 { + serverCfg.SessionTTL = config.sessionTTL + } vmcpServer, err := vmcpserver.Serve(ctx, coreVMCP, serverCfg) require.NoError(tb, err, "failed to create vMCP server") diff --git a/test/integration/vmcp/pagination_regression_test.go b/test/integration/vmcp/pagination_regression_test.go new file mode 100644 index 0000000000..9d0e971c40 --- /dev/null +++ b/test/integration/vmcp/pagination_regression_test.go @@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// TestRegression_Over1000Tools_CompleteSetReceived is a regression anchor for +// the vMCP pagination gap: a backend exposing more than the MCP page size +// (1000 tools) must surface its complete tool set across pagination cursors. +// +// Today the Serve-path tool query path issues a single ListTools with no +// cursor loop, so backends that paginate return only the first page. The test +// is intentionally skipped until the gap is closed (see gap-analysis V1 and the +// follow-up issue); the skeleton remains so the regression cannot regress +// silently — flipping the skip off re-runs the full assertion. +func TestRegression_Over1000Tools_CompleteSetReceived(t *testing.T) { + t.Parallel() + t.Skip("vMCP does not follow pagination cursors; see gap-analysis V1 / follow-up issue") + + ctx := context.Background() + + // Generate 1000+ tools on a single backend to exceed the MCP page size. + const toolCount = 1100 + tools := make([]helpers.BackendTool, 0, toolCount) + for i := 0; i < toolCount; i++ { + name := fmt.Sprintf("tool_%04d", i) + tools = append(tools, helpers.NewBackendTool( + name, + fmt.Sprintf("backend tool number %d", i), + func(_ context.Context, _ map[string]any) string { + return `{"ok": true}` + }, + )) + } + + backend := helpers.CreateBackendServer(t, tools, helpers.WithBackendName("many-tools")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("many-tools", + helpers.WithURL(backend.URL+"/mcp"), + helpers.WithMetadata("group", "pagination-test"), + ), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + ) + + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + client := helpers.NewMCPClient(ctx, t, vmcpURL) + defer client.Close() + + result := client.ListTools(ctx) + + // The complete tool set must be received; a single-page query returns at + // most the page size (1000), so fewer than toolCount indicates a dropped + // tail — the regression this test pins. + assert.GreaterOrEqual(t, len(result.Tools), toolCount, + "vMCP must surface the complete backend tool set across pagination cursors; "+ + "received %d of %d", len(result.Tools), toolCount) +} diff --git a/test/integration/vmcp/per_session_projection_regression_test.go b/test/integration/vmcp/per_session_projection_regression_test.go new file mode 100644 index 0000000000..6f4ac07141 --- /dev/null +++ b/test/integration/vmcp/per_session_projection_regression_test.go @@ -0,0 +1,424 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sort" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/env" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/authz/authorizers" + "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth/factory" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// principalHeader is the request header the test identity middleware reads to +// determine the authenticated principal. The Cedar authorizer resolves the +// principal from the identity's "sub" claim, so the middleware sets both +// Subject and Claims["sub"] to the header value. +const principalHeader = "X-Test-Principal" + +// newCedarVMCPServer builds a Cedar-authz-backed vMCP server (via server.New, +// mirroring pkg/vmcp/server/authz_integration_test.go) backed by a real MCP +// backend at backendURL. An identity middleware injects a principal derived +// from the X-Test-Principal request header so each session binds to its +// caller and the Cedar admission seam can resolve Client::"". +// +// A priority conflict resolver is used so tool names are NOT prefixed — the +// Cedar policies can name tools by their raw backend names (e.g. +// Tool::"secret-tool"). +func newCedarVMCPServer(t *testing.T, backendURL string, policies ...string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: backendURL, + TransportType: "streamable-http", + } + mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + authReg, err := vmcpauth.NewOutgoingAuthRegistry(context.Background(), &env.OSReader{}) + require.NoError(t, err) + factory := vmcpsession.NewSessionFactory(authReg) + + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + // Priority resolver keeps raw tool names so Cedar policies can name + // Tool::"secret-tool" rather than a prefixed variant. + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + // Identity middleware: derive the principal from the X-Test-Principal + // header so two sessions (alice, bob) bind to different identities. The + // Cedar authorizer resolves the principal from the "sub" claim, so both + // Subject and Claims["sub"] are set to the header value. + identityMiddleware := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + principal := r.Header.Get(principalHeader) + if principal == "" { + principal = "anonymous" + } + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: principal, + Name: principal, + Claims: map[string]any{"sub": principal, "name": principal}, + }} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) + } + + authzCfg, err := authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{Policies: policies, EntitiesJSON: "[]"}, + }) + require.NoError(t, err) + + srv, err := server.New( + context.Background(), + &server.Config{ + Name: "test-vmcp", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: factory, + Aggregator: agg, + AuthMiddleware: identityMiddleware, + Authz: authzCfg, + }, + router.NewSessionRouter(&vmcp.RoutingTable{}), + backendClient, + mockBackendRegistry, + nil, + ) + require.NoError(t, err) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// rawClient is a minimal JSON-RPC-over-HTTP client for the Cedar authz tests. +// It is needed because the helpers.MCPClient.CallTest helper uses require.NoError +// on the SDK error, which would abort the test on a policy-rejected call. This +// client surfaces the raw HTTP response so tests can inspect status codes and +// error bodies directly. +type rawClient struct { + baseURL string + sessionID string + nextID int +} + +func newRawClient(baseURL string) *rawClient { + return &rawClient{baseURL: baseURL, nextID: 1} +} + +func (c *rawClient) initialize(t *testing.T, principal string) { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + c.nextID++ + defer resp.Body.Close() + c.sessionID = resp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, c.sessionID, "initialize response missing Mcp-Session-Id header") + + // Send the initialized notification (no id, no response expected). + notif := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "method": "notifications/initialized", + }, c.sessionID) + notif.Body.Close() +} + +func (c *rawClient) listTools(t *testing.T, principal string) map[string]any { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "tools/list", + "params": map[string]any{}, + }, c.sessionID) + c.nextID++ + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "tools/list failed: %s", string(body)) + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + return parsed +} + +func (c *rawClient) callTool(t *testing.T, principal, name string, args map[string]any) (int, map[string]any) { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "tools/call", + "params": map[string]any{ + "name": name, + "arguments": args, + }, + }, c.sessionID) + c.nextID++ + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) // best-effort; body may be plain text on 404 + return resp.StatusCode, parsed +} + +func (c *rawClient) postMCP(t *testing.T, principal string, body map[string]any, sessionID string) *http.Response { + t.Helper() + rawBody, err := json.Marshal(body) + require.NoError(t, err) + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, c.baseURL+"/mcp", bytes.NewReader(rawBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(principalHeader, principal) + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// toolNamesFromResult extracts the tool names from a tools/list JSON-RPC result. +func toolNamesFromResult(t *testing.T, result map[string]any) []string { + t.Helper() + rawResult, ok := result["result"].(map[string]any) + require.True(t, ok, "response missing result object: %v", result) + toolsRaw, ok := rawResult["tools"].([]any) + require.True(t, ok, "result missing tools array: %v", rawResult) + names := make([]string, 0, len(toolsRaw)) + for _, t2 := range toolsRaw { + tool, ok := t2.(map[string]any) + require.True(t, ok) + name, ok := tool["name"].(string) + require.True(t, ok) + names = append(names, name) + } + sort.Strings(names) + return names +} + +// startAuthzBackend starts a real MCP backend exposing the named tools. Each +// tool echoes back its name and the input so a caller can confirm the backend +// was reached. +func startAuthzBackend(t *testing.T, toolNames ...string) string { + t.Helper() + tools := make([]helpers.BackendTool, 0, len(toolNames)) + for _, name := range toolNames { + tools = append(tools, helpers.NewBackendTool( + name, + "test tool "+name, + func(_ context.Context, args map[string]any) string { + return `{"tool":` + name + `,"args":` + toJSON(args) + `}` + }, + )) + } + backend := helpers.CreateBackendServer(t, tools, helpers.WithBackendName("real-backend")) + t.Cleanup(backend.Close) + return backend.URL + "/mcp" +} + +// toJSON is a tiny marshal helper that never fails the test (returns "null"). +func toJSON(v map[string]any) string { + b, err := json.Marshal(v) + if err != nil { + return "null" + } + return string(b) +} + +// TestRegression_TwoSessions_DifferentAuthz_DifferentToolsList proves the +// core admission seam projects a per-identity tool list: two sessions bound to +// different principals see different advertised tool sets under the same Cedar +// policy. Alice (permitted for secret-tool) sees both secret-tool and +// public-tool; Bob (only permitted for public-tool) sees only public-tool. +func TestRegression_TwoSessions_DifferentAuthz_DifferentToolsList(t *testing.T) { + t.Parallel() + + backendURL := startAuthzBackend(t, "secret-tool", "public-tool") + + ts := newCedarVMCPServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"secret-tool");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"public-tool");`, + ) + + // Alice's session. + alice := newRawClient(ts.URL) + alice.initialize(t, "alice") + aliceTools := toolNamesFromResult(t, alice.listTools(t, "alice")) + assert.Equal(t, []string{"public-tool", "secret-tool"}, aliceTools, + "alice (permitted for secret-tool) must see both tools") + + // Bob's session. + bob := newRawClient(ts.URL) + bob.initialize(t, "bob") + bobTools := toolNamesFromResult(t, bob.listTools(t, "bob")) + assert.Equal(t, []string{"public-tool"}, bobTools, + "bob (not permitted for secret-tool) must see only public-tool") +} + +// TestRegression_FilteredToolUncallableInSessionA_CallableInSessionB proves +// the list-filter/call-deny pairing: a tool filtered out of a session (bob's) +// is rejected as unknown when called, while the same tool is callable in a +// session where it was advertised (alice's). +func TestRegression_FilteredToolUncallableInSessionA_CallableInSessionB(t *testing.T) { + t.Parallel() + + backendURL := startAuthzBackend(t, "secret-tool", "public-tool") + + ts := newCedarVMCPServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"secret-tool");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"public-tool");`, + ) + + // Alice can call secret-tool (it is in her session's advertised set). + alice := newRawClient(ts.URL) + alice.initialize(t, "alice") + require.Contains(t, toolNamesFromResult(t, alice.listTools(t, "alice")), "secret-tool") + + status, result := alice.callTool(t, "alice", "secret-tool", map[string]any{"input": "hi"}) + require.Equal(t, http.StatusOK, status, "alice's permitted call must succeed: %v", result) + assert.False(t, isToolResultError(result), "alice's call to secret-tool must not be an error: %v", result) + + // Bob cannot call secret-tool: it was filtered out of his session, so the + // SDK rejects it as an unknown tool. Poll briefly because the session's + // tool injection runs in a hook that fires after initialize returns. + bob := newRawClient(ts.URL) + bob.initialize(t, "bob") + require.NotContains(t, toolNamesFromResult(t, bob.listTools(t, "bob")), "secret-tool") + + var lastResult map[string]any + require.Eventually(t, func() bool { + var s int + s, lastResult = bob.callTool(t, "bob", "secret-tool", map[string]any{"input": "hi"}) + if s != http.StatusOK { + return true // 404/400 also indicates rejection + } + return isToolResultError(lastResult) + }, 5*time.Second, 50*time.Millisecond, + "bob's call to secret-tool (filtered out) must be rejected") +} + +// TestRegression_SetSessionTools_FixedAtInitialize_NoMidSessionReconciliation +// asserts that a session's advertised tool set is fixed at initialize time +// and is not re-aggregated on each tools/list call. Calling tools/list twice +// on the same session must return the same set. +func TestRegression_SetSessionTools_FixedAtInitialize_NoMidSessionReconciliation(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + backend := helpers.CreateBackendServer(t, []helpers.BackendTool{ + helpers.NewBackendTool("tool-a", "tool A", func(_ context.Context, _ map[string]any) string { + return `{"tool":"a"}` + }), + helpers.NewBackendTool("tool-b", "tool B", func(_ context.Context, _ map[string]any) string { + return `{"tool":"b"}` + }), + }, helpers.WithBackendName("ab-backend")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("ab-backend", helpers.WithURL(backend.URL+"/mcp")), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + ) + + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + client := helpers.NewMCPClient(ctx, t, vmcpURL) + defer client.Close() + + first := helpers.GetToolNames(client.ListTools(ctx)) + second := helpers.GetToolNames(client.ListTools(ctx)) + + sort.Strings(first) + sort.Strings(second) + + assert.Equal(t, []string{"ab-backend_tool-a", "ab-backend_tool-b"}, first, + "initial tools/list must show both tools") + assert.Equal(t, first, second, + "the session's tool set is fixed at initialize; a second tools/list must return the same set") +} + +// isToolResultError returns true when a tools/call JSON-RPC result has +// IsError=true (the SDK rejected the call) or the response carries a +// JSON-RPC error object. +func isToolResultError(result map[string]any) bool { + if _, ok := result["error"]; ok { + return true + } + if r, ok := result["result"].(map[string]any); ok { + if isErr, _ := r["isError"].(bool); isErr { + return true + } + } + return false +} + +// mustParseJSONRPC reads resp.Body and unmarshals it as a JSON-RPC envelope. +// If the body is not valid JSON (e.g. a plain-text 404), it returns an empty +// map so callers can fall back to status-code checks. +func mustParseJSONRPC(t *testing.T, resp *http.Response) map[string]any { + t.Helper() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) // best-effort; body may be plain text on 4xx + return parsed +} + +// Compile-time check that the mcpcompat imports are used (not mark3labs/mcp-go). +var ( + _ = mcp.LATEST_PROTOCOL_VERSION + _ = mcpserver.NewMCPServer +) diff --git a/test/integration/vmcp/sliding_ttl_regression_test.go b/test/integration/vmcp/sliding_ttl_regression_test.go new file mode 100644 index 0000000000..0a70ae2737 --- /dev/null +++ b/test/integration/vmcp/sliding_ttl_regression_test.go @@ -0,0 +1,135 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// TestRegression_SlidingSessionTTL_TrafficKeepsSessionAlive proves the +// session TTL is sliding: a session that receives traffic at least once per +// TTL window stays alive indefinitely, while an idle session is evicted once +// its TTL elapses. +// +// The transport session storage (LocalSessionDataStorage.Load) refreshes the +// last-access timestamp on every read, and the SDK's SessionIdManager.Validate +// is called on every request — so active traffic keeps the session alive while +// a session that goes idle for longer than the TTL is rejected on its next +// request. +// +// This test is timing-sensitive and must NOT run in parallel. +// +//nolint:paralleltest // timing-sensitive: relies on real TTL expiry and background cleanup +func TestRegression_SlidingSessionTTL_TrafficKeepsSessionAlive(t *testing.T) { + const sessionTTL = 2 * time.Second + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + backend := helpers.CreateBackendServer(t, []helpers.BackendTool{ + helpers.NewBackendTool("ping", "ping tool", func(_ context.Context, _ map[string]any) string { + return `{"pong":true}` + }), + }, helpers.WithBackendName("ping-backend")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("ping-backend", helpers.WithURL(backend.URL+"/mcp")), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + helpers.WithSessionTTL(sessionTTL), + ) + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + + // ── Active session: traffic every 500ms for 4s (2x the TTL) ────────────── + // + // If the TTL were fixed (not sliding), the session would expire at the 2s + // mark and the next tools/list would fail. A sliding TTL refreshes + // last-access on each request, so all calls must succeed. + activeClient := helpers.NewMCPClient(ctx, t, vmcpURL) + defer activeClient.Close() + + const ( + tickInterval = 500 * time.Millisecond + totalWindow = 4 * time.Second + ) + ticks := int(totalWindow / tickInterval) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + ticker := time.NewTicker(tickInterval) + defer ticker.Stop() + for i := 0; i < ticks; i++ { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + // tools/list routes through the session; each call refreshes + // last-access in the transport session storage. + _ = activeClient.ListTools(ctx) + } + }() + wg.Wait() + + // One final call after the full window: the session must still be alive. + final := activeClient.ListTools(ctx) + assert.NotEmpty(t, final.Tools, + "active session must survive past its TTL while receiving traffic") + + // ── Idle session: no traffic for 3s (> TTL) must be evicted ────────────── + // + // A second session that goes idle for longer than the TTL must be rejected + // on its next tool call. tools/list is NOT a reliable eviction signal here: + // the go-sdk StreamableHTTPHandler keeps its per-session transport in an + // in-memory map that is not TTL-gated, so tools/list continues to return + // the cached tool set even after the vMCP session storage has evicted the + // session. The eviction IS observable on tools/call, whose handler runs + // enforceSessionBinding → GetMultiSession → checkSession → storage.Load, + // which returns ErrSessionNotFound once the TTL has elapsed and the + // background cleanup sweep has removed the entry. + idleClient := newRawClient(vmcpURL) + idleClient.initialize(t, "") + + // Let the session go idle well past the TTL. The cleanup goroutine runs on + // a ttl/2 ticker, so waiting 2x the TTL after the idle period guarantees the + // background sweep has evicted the session. + idleFor := 3 * sessionTTL + time.Sleep(idleFor) + + // A tool call on the idle session must be rejected: enforceSessionBinding + // fails because the session is gone from storage. The SDK surfaces this as + // either an HTTP error (4xx) or a JSON-RPC error result — assert both. + resp := idleClient.postMCP(t, "", map[string]any{ + "jsonrpc": "2.0", + "id": idleClient.nextID, + "method": "tools/call", + "params": map[string]any{ + "name": "ping-backend_ping", + "arguments": map[string]any{}, + }, + }, idleClient.sessionID) + idleClient.nextID++ + defer resp.Body.Close() + + rejected := resp.StatusCode >= 400 || + isToolResultError(mustParseJSONRPC(t, resp)) + assert.True(t, rejected, + "idle session must be evicted after exceeding the TTL (status %d)", resp.StatusCode) +} + +// Compile-time check that this file stays in the vmcp_test package. +var _ vmcp.Backend