diff --git a/pkg/vmcp/server/backend_enrichment.go b/pkg/vmcp/server/backend_enrichment.go deleted file mode 100644 index cd940a996f..0000000000 --- a/pkg/vmcp/server/backend_enrichment.go +++ /dev/null @@ -1,120 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "bytes" - "context" - "encoding/json" - "io" - "log/slog" - "net/http" - - "github.com/stacklok/toolhive/pkg/audit" - "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/discovery" -) - -// withBackendEnrichment wraps h with the legacy audit backend-enrichment middleware -// when audit is enabled on the legacy (server.New) path. On the Serve path -// (s.core != nil) the per-session tool/resource handlers write the backend label -// directly into the audit BackendInfo (serve_handlers.go), so the middleware — with -// its per-request body re-read and routing-table lookup — is not wired and h is -// returned unchanged (#5512 review). The legacy wiring is retired with the rest of the -// discovery path in #5445. -func (s *Server) withBackendEnrichment(h http.Handler) http.Handler { - if s.config.AuditConfig != nil && s.core == nil { - slog.Info("backend enrichment middleware enabled for audit events") - return backendEnrichmentMiddleware(h) - } - return h -} - -// backendEnrichmentMiddleware wraps an HTTP handler to add backend routing information -// to audit events by parsing MCP requests and resolving the target backend. -// -// This is the legacy (server.New) path's audit labelling: it reads the routing table -// the discovery middleware injected into the request context. It is only wired when -// s.core == nil (see (*Server).Handler) — on the Serve path the per-session handlers -// label the audit event directly (serve_handlers.go), so this middleware, with its -// per-request body re-read, is not in the chain. Retired with the discovery path in #5445. -// -// It is a free function (not a *Server method): with the Serve branch gone it reads no -// Server state — resolution comes entirely from the request context. -func backendEnrichmentMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Read and parse the request body to extract MCP method and parameters - var requestBody []byte - if r.Body != nil { - var err error - requestBody, err = io.ReadAll(r.Body) - // Always restore body for next handler, even on error - if err != nil { - // Log the error and restore an empty body to ensure consistent behavior - slog.Warn("failed to read request body in backend enrichment middleware", - "error", err) - r.Body = io.NopCloser(bytes.NewReader([]byte{})) - } else { - // Restore body with the read content - r.Body = io.NopCloser(bytes.NewReader(requestBody)) - } - } - - // Parse MCP request to extract tool/resource name - var mcpRequest struct { - Method string `json:"method"` - Params map[string]any `json:"params"` - } - - if len(requestBody) > 0 && json.Unmarshal(requestBody, &mcpRequest) == nil { - backendName := resolveBackendName(r.Context(), mcpRequest.Method, mcpRequest.Params) - - // Mutate the existing BackendInfo from audit middleware - if backendName != "" { - if backendInfo, ok := audit.BackendInfoFromContext(r.Context()); ok && backendInfo != nil { - backendInfo.BackendName = backendName - } - } - } - - // Call next handler - next.ServeHTTP(w, r) - }) -} - -// resolveBackendName resolves the backend handling an MCP request from the routing -// table the discovery middleware injected into the request context (legacy path). -// The middleware is only wired when s.core == nil, so there is no Serve-path branch. -func resolveBackendName(ctx context.Context, method string, params map[string]any) string { - caps, ok := discovery.DiscoveredCapabilitiesFromContext(ctx) - if !ok || caps == nil || caps.RoutingTable == nil { - return "" - } - return lookupBackendName(method, params, caps.RoutingTable) -} - -// lookupBackendName looks up which backend handles a given MCP request. -func lookupBackendName(method string, params map[string]any, routingTable *vmcp.RoutingTable) string { - switch method { - case "tools/call": - if toolName, ok := params["name"].(string); ok { - if target, exists := routingTable.Tools[toolName]; exists { - return target.WorkloadName - } - } - case "resources/read": - if uri, ok := params["uri"].(string); ok { - if target, exists := routingTable.Resources[uri]; exists { - return target.WorkloadName - } - } - case "prompts/get": - if promptName, ok := params["name"].(string); ok { - if target, exists := routingTable.Prompts[promptName]; exists { - return target.WorkloadName - } - } - } - return "" -} diff --git a/pkg/vmcp/server/backend_enrichment_test.go b/pkg/vmcp/server/backend_enrichment_test.go deleted file mode 100644 index 60b7a82d11..0000000000 --- a/pkg/vmcp/server/backend_enrichment_test.go +++ /dev/null @@ -1,749 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. -// SPDX-License-Identifier: Apache-2.0 - -package server - -import ( - "bytes" - "context" - "errors" - "io" - "net/http" - "net/http/httptest" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/stacklok/toolhive/pkg/audit" - "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/aggregator" - "github.com/stacklok/toolhive/pkg/vmcp/discovery" -) - -// Common test constants -const toolsCallRequest = `{"method":"tools/call","params":{"name":"test-tool"}}` - -// errorReader is a reader that always returns an error -type errorReader struct{} - -func (errorReader) Read([]byte) (int, error) { - return 0, errors.New("simulated read error") -} - -// createTestHandler creates a handler that tracks if it was called -func createTestHandler() (http.Handler, *bool) { - called := false - handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - called = true - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("OK")) - }) - return handler, &called -} - -func TestBackendEnrichmentMiddleware(t *testing.T) { - t.Parallel() - - t.Run("enriches backend name for tools/call request", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - // Create routing table with a tool - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - // Create aggregated capabilities with routing table - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - // Create backend info that will be mutated - backendInfo := &audit.BackendInfo{} - - // Create request with MCP tools/call body - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(toolsCallRequest))) - - // Add capabilities and backend info to context - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - // Create response recorder - rr := httptest.NewRecorder() - - // Create server instance and wrap handler with middleware - middleware := backendEnrichmentMiddleware(nextHandler) - - // Execute middleware - middleware.ServeHTTP(rr, req) - - // Verify handler was called - assert.True(t, *handlerCalled, "next handler should be called") - assert.Equal(t, http.StatusOK, rr.Code) - - // Verify backend name was enriched - assert.Equal(t, "backend-1", backendInfo.BackendName) - }) - - t.Run("enriches backend name for resources/read request", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - // Create routing table with a resource - routingTable := &vmcp.RoutingTable{ - Resources: map[string]*vmcp.BackendTarget{ - "file:///test/resource": { - WorkloadName: "backend-2", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - requestBody := `{"method":"resources/read","params":{"uri":"file:///test/resource"}}` - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(requestBody))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - assert.Equal(t, "backend-2", backendInfo.BackendName) - }) - - t.Run("enriches backend name for prompts/get request", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - // Create routing table with a prompt - routingTable := &vmcp.RoutingTable{ - Prompts: map[string]*vmcp.BackendTarget{ - "test-prompt": { - WorkloadName: "backend-3", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - requestBody := `{"method":"prompts/get","params":{"name":"test-prompt"}}` - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(requestBody))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - assert.Equal(t, "backend-3", backendInfo.BackendName) - }) - - t.Run("handles missing routing table gracefully", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - // Create capabilities without routing table - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: nil, - } - - backendInfo := &audit.BackendInfo{} - - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(toolsCallRequest))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - // Backend name should remain empty - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("handles missing capabilities in context", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - backendInfo := &audit.BackendInfo{} - - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(toolsCallRequest))) - - // Only add backend info, no capabilities - ctx := audit.WithBackendInfo(req.Context(), backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("handles missing backend info in context", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(toolsCallRequest))) - - // Only add capabilities, no backend info - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - // Should not panic, should proceed normally - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - }) - - t.Run("handles malformed JSON request", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - // Malformed JSON - requestBody := `{"method":"tools/call","params":{"name":"test-tool"` - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(requestBody))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - // Should not panic, should proceed with next handler - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("handles tool not found in routing table", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - // Routing table with different tool - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "other-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - requestBody := `{"method":"tools/call","params":{"name":"unknown-tool"}}` - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(requestBody))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - // Backend name should remain empty since tool not found - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("properly restores request body for next handler", func(t *testing.T) { - t.Parallel() - var bodyRead []byte - - // Create a handler that reads the body - bodyReadingHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var err error - bodyRead, err = io.ReadAll(r.Body) - require.NoError(t, err) - w.WriteHeader(http.StatusOK) - }) - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(toolsCallRequest))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(bodyReadingHandler) - middleware.ServeHTTP(rr, req) - - // Verify body was properly restored and readable by next handler - assert.Equal(t, toolsCallRequest, string(bodyRead)) - assert.Equal(t, "backend-1", backendInfo.BackendName) - }) - - t.Run("handles nil request body", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - req := httptest.NewRequest(http.MethodGet, "/mcp", nil) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - // Should not panic, should proceed normally - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("handles empty request body", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(""))) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - assert.True(t, *handlerCalled) - assert.Equal(t, http.StatusOK, rr.Code) - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("handles body read error gracefully", func(t *testing.T) { - t.Parallel() - - nextHandler, handlerCalled := createTestHandler() - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - // Create request with error reader that will fail on read - req := httptest.NewRequest(http.MethodPost, "/mcp", io.NopCloser(errorReader{})) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - // Should not panic, should proceed with next handler - assert.True(t, *handlerCalled, "next handler should still be called even on read error") - assert.Equal(t, http.StatusOK, rr.Code) - // Backend name should remain empty since we couldn't read the body - assert.Empty(t, backendInfo.BackendName) - }) - - t.Run("restores empty body after read error for downstream handlers", func(t *testing.T) { - t.Parallel() - var bodyRead []byte - var readErr error - - // Create a handler that tries to read the body - bodyReadingHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - bodyRead, readErr = io.ReadAll(r.Body) - w.WriteHeader(http.StatusOK) - }) - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": { - WorkloadName: "backend-1", - }, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - // Create request with error reader - req := httptest.NewRequest(http.MethodPost, "/mcp", io.NopCloser(errorReader{})) - - ctx := discovery.WithDiscoveredCapabilities(req.Context(), caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(bodyReadingHandler) - middleware.ServeHTTP(rr, req) - - // Verify body was restored (as empty) and readable by next handler without error - assert.NoError(t, readErr, "downstream handler should be able to read restored body without error") - assert.Empty(t, bodyRead, "restored body should be empty after read error") - }) -} - -func TestLookupBackendName(t *testing.T) { - t.Parallel() - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "tool-1": {WorkloadName: "backend-a"}, - "tool-2": {WorkloadName: "backend-b"}, - }, - Resources: map[string]*vmcp.BackendTarget{ - "file:///resource-1": {WorkloadName: "backend-c"}, - "file:///resource-2": {WorkloadName: "backend-d"}, - }, - Prompts: map[string]*vmcp.BackendTarget{ - "prompt-1": {WorkloadName: "backend-e"}, - "prompt-2": {WorkloadName: "backend-f"}, - }, - } - - t.Run("looks up tool by name", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": "tool-1", - } - - result := lookupBackendName("tools/call", params, routingTable) - assert.Equal(t, "backend-a", result) - }) - - t.Run("looks up resource by URI", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "uri": "file:///resource-1", - } - - result := lookupBackendName("resources/read", params, routingTable) - assert.Equal(t, "backend-c", result) - }) - - t.Run("looks up prompt by name", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": "prompt-1", - } - - result := lookupBackendName("prompts/get", params, routingTable) - assert.Equal(t, "backend-e", result) - }) - - t.Run("returns empty string for unknown tool", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": "unknown-tool", - } - - result := lookupBackendName("tools/call", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("returns empty string for unknown resource", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "uri": "file:///unknown-resource", - } - - result := lookupBackendName("resources/read", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("returns empty string for unknown prompt", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": "unknown-prompt", - } - - result := lookupBackendName("prompts/get", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("returns empty string for unknown method", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": "tool-1", - } - - result := lookupBackendName("unknown/method", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles missing parameter for tools/call", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "other": "value", - } - - result := lookupBackendName("tools/call", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles missing parameter for resources/read", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "other": "value", - } - - result := lookupBackendName("resources/read", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles missing parameter for prompts/get", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "other": "value", - } - - result := lookupBackendName("prompts/get", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles non-string parameter for tools/call", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": 123, // Integer instead of string - } - - result := lookupBackendName("tools/call", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles non-string parameter for resources/read", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "uri": []string{"not", "a", "string"}, - } - - result := lookupBackendName("resources/read", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles non-string parameter for prompts/get", func(t *testing.T) { - t.Parallel() - params := map[string]any{ - "name": map[string]string{"not": "a string"}, - } - - result := lookupBackendName("prompts/get", params, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles nil params map", func(t *testing.T) { - t.Parallel() - - result := lookupBackendName("tools/call", nil, routingTable) - assert.Empty(t, result) - }) - - t.Run("handles empty routing table", func(t *testing.T) { - t.Parallel() - emptyTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{}, - Resources: map[string]*vmcp.BackendTarget{}, - Prompts: map[string]*vmcp.BackendTarget{}, - } - - params := map[string]any{ - "name": "tool-1", - } - - result := lookupBackendName("tools/call", params, emptyTable) - assert.Empty(t, result) - }) -} - -func TestBackendEnrichmentMiddleware_ContextPropagation(t *testing.T) { - t.Parallel() - - t.Run("preserves context for downstream handlers", func(t *testing.T) { - t.Parallel() - - var receivedCtx context.Context - nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - receivedCtx = r.Context() - w.WriteHeader(http.StatusOK) - }) - - routingTable := &vmcp.RoutingTable{ - Tools: map[string]*vmcp.BackendTarget{ - "test-tool": {WorkloadName: "backend-1"}, - }, - } - - caps := &aggregator.AggregatedCapabilities{ - RoutingTable: routingTable, - } - - backendInfo := &audit.BackendInfo{} - - req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader([]byte(toolsCallRequest))) - - // Add some custom context value - type contextKey string - const testKey contextKey = "test-key" - ctx := context.WithValue(req.Context(), testKey, "test-value") - ctx = discovery.WithDiscoveredCapabilities(ctx, caps) - ctx = audit.WithBackendInfo(ctx, backendInfo) - req = req.WithContext(ctx) - - rr := httptest.NewRecorder() - - middleware := backendEnrichmentMiddleware(nextHandler) - middleware.ServeHTTP(rr, req) - - // Verify custom context value was preserved - require.NotNil(t, receivedCtx) - assert.Equal(t, "test-value", receivedCtx.Value(testKey)) - - // Verify capabilities and backend info are still accessible - receivedCaps, ok := discovery.DiscoveredCapabilitiesFromContext(receivedCtx) - assert.True(t, ok) - assert.NotNil(t, receivedCaps) - - receivedBackendInfo, ok := audit.BackendInfoFromContext(receivedCtx) - assert.True(t, ok) - assert.NotNil(t, receivedBackendInfo) - assert.Equal(t, "backend-1", receivedBackendInfo.BackendName) - }) -} diff --git a/pkg/vmcp/server/serve_session_test.go b/pkg/vmcp/server/serve_session_test.go index 53653f1870..85856e9d71 100644 --- a/pkg/vmcp/server/serve_session_test.go +++ b/pkg/vmcp/server/serve_session_test.go @@ -338,11 +338,11 @@ func TestServeLazyInjectsToolsForRehydratedSession(t *testing.T) { sessionID := initResp.Header.Get("Mcp-Session-Id") require.NotEmpty(t, sessionID) - // Wait until the session is fully registered with adapted tools in the manager. + // Wait until the session is fully registered in the manager. require.Eventually(t, func() bool { - tools, gErr := srv.vmcpSessionMgr.GetAdaptedTools(sessionID) - return gErr == nil && len(tools) > 0 - }, 2*time.Second, 10*time.Millisecond, "session should be registered with adapted tools") + _, ok := srv.vmcpSessionMgr.GetMultiSession(sessionID) + return ok + }, 2*time.Second, 10*time.Millisecond, "session should be registered") // Empty-store (cross-pod) branch: a fresh SDK session with no tools gets them injected. rehydrated := &fakeSDKSession{id: sessionID, tools: map[string]server.ServerTool{}} @@ -350,7 +350,7 @@ func TestServeLazyInjectsToolsForRehydratedSession(t *testing.T) { assert.Contains(t, rehydrated.tools, testTool.Name, "an empty per-session store should be re-injected from the vMCP session manager") - // No-op branch: a populated store is left untouched (early return before GetAdaptedTools). + // No-op branch: a populated store is left untouched (early return before re-derivation). populated := &fakeSDKSession{id: sessionID, tools: map[string]server.ServerTool{ "preexisting": {Tool: mcp.Tool{Name: "preexisting"}}, }} diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 75d324d188..84c3343bca 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -553,9 +553,6 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { // when auth middleware is nil. mcpHandler = mcpparser.ParsingMiddleware(mcpHandler) - // Apply backend enrichment middleware (legacy path only — see withBackendEnrichment). - mcpHandler = s.withBackendEnrichment(mcpHandler) - // Apply discovery middleware (runs after audit/auth middleware) — legacy path only. // Discovery middleware performs per-request capability aggregation with user context, // injecting the routing table into the request context (the discovery-into-context seam). @@ -1044,26 +1041,19 @@ func (s *Server) lazyInjectSessionTools(ctx context.Context) { sessionID := sess.SessionID() // Re-derive the tool set the same way registration did, so cross-pod re-injection - // matches the advertised set. On the Serve path (s.core != nil) that means a fresh - // core.ListTools for the request identity (the core is stateless, so re-deriving on - // pod B is the cache-miss equivalent of the once-per-session registration call), - // optimizer-wrapped into find_tool/call_tool when the optimizer is enabled — both - // via serveSessionTools, the same helper registration uses; on the legacy path it - // reads the factory-built session's tools via GetAdaptedTools. + // matches the advertised set: a fresh core.ListTools for the request identity (the core + // is stateless, so re-deriving on pod B is the cache-miss equivalent of the + // once-per-session registration call), optimizer-wrapped into find_tool/call_tool when + // the optimizer is enabled — via serveSessionTools, the same helper registration uses. // - // Note: on the Serve path this lists under the CURRENT request identity, not the - // session's bound identity. For the realistic same-principal load-balanced case they - // are equal, so the advertised set is identical across pods. A cross-identity re-hydration - // would advertise the requester's own filtered set, but the call-time binding check - // (enforceSessionBinding, run before core.CallTool/ReadResource) is the backstop — it - // rejects a mismatched caller, so no other principal's capabilities can be invoked. - adaptedTools, err := func() ([]server.ServerTool, error) { - if s.core != nil { - identity, _ := auth.IdentityFromContext(ctx) - return s.serveSessionTools(ctx, sessionID, identity) - } - return s.vmcpSessionMgr.GetAdaptedTools(sessionID) - }() + // Note: this lists under the CURRENT request identity, not the session's bound identity. + // For the realistic same-principal load-balanced case they are equal, so the advertised + // set is identical across pods. A cross-identity re-hydration would advertise the + // requester's own filtered set, but the call-time binding check (enforceSessionBinding, + // run before core.CallTool/ReadResource) is the backstop — it rejects a mismatched + // caller, so no other principal's capabilities can be invoked. + identity, _ := auth.IdentityFromContext(ctx) + adaptedTools, err := s.serveSessionTools(ctx, sessionID, identity) if err != nil || len(adaptedTools) == 0 { slog.Debug("lazyInjectSessionTools: no tools available for session", "session_id", sessionID) return @@ -1137,53 +1127,12 @@ func (s *Server) handleSessionRegistrationImpl(ctx context.Context, session serv return retErr } - // Serve path: the core is the single authoritative aggregation. Source the - // advertised tool/resource set from core.ListTools/ListResources (called once - // per session here) and install handlers that route through the core; the - // session factory's own aggregation is not used. CreateSession above still - // establishes the bound session record (identity binding, TTL, Validate). The - // returned error becomes retErr (named return), so the defer terminates the - // session on failure. - if s.core != nil { - return s.injectCoreSessionCapabilities(ctx, session) - } - - // Legacy server.New path: uniform registration — same code path regardless of - // which decorators are active. session.Tools() returns the final decorated tool list. - adaptedTools, retErr := s.vmcpSessionMgr.GetAdaptedTools(sessionID) - if retErr != nil { - slog.Error("failed to get session-scoped tools", - "session_id", sessionID, - "error", retErr) - return retErr - } - - adaptedResources, retErr := s.vmcpSessionMgr.GetAdaptedResources(sessionID) - if retErr != nil { - slog.Error("failed to get session-scoped resources", - "session_id", sessionID, - "error", retErr) - return retErr - } - - if len(adaptedResources) > 0 { - if err := setSessionResourcesDirect(session, adaptedResources); err != nil { - slog.Error("failed to add session resources", "session_id", sessionID, "error", err) - return err - } - } - - if len(adaptedTools) > 0 { - if err := setSessionToolsDirect(session, adaptedTools); err != nil { - slog.Error("failed to add session tools", "session_id", sessionID, "error", err) - return err - } - } - - slog.Info("session capabilities injected", - "session_id", sessionID, - "tool_count", len(adaptedTools)) - return nil + // The core is the single authoritative aggregation: source the advertised tool/resource + // set from core.ListTools/ListResources (called once per session here) and install + // handlers that route through the core. CreateSession above still establishes the bound + // session record (identity binding, TTL, Validate). The returned error becomes retErr + // (named return), so the defer terminates the session on failure. + return s.injectCoreSessionCapabilities(ctx, session) } // backendHealth returns the core-owned backend health reporter, or nil when health diff --git a/pkg/vmcp/server/session_manager_interface.go b/pkg/vmcp/server/session_manager_interface.go index d9b35c15ff..c7960a6e4b 100644 --- a/pkg/vmcp/server/session_manager_interface.go +++ b/pkg/vmcp/server/session_manager_interface.go @@ -27,15 +27,6 @@ type SessionManager interface { // connections and replaces the placeholder with a fully-formed MultiSession. CreateSession(ctx context.Context, sessionID string) (vmcpsession.MultiSession, error) - // GetAdaptedTools returns SDK-format tools for the given session with session-scoped - // handlers. This enables session-scoped routing: each tool call goes through the - // session's backend connections rather than the global router. - GetAdaptedTools(sessionID string) ([]mcpserver.ServerTool, error) - - // GetAdaptedResources returns SDK-format resources for the given session with - // session-scoped handlers, analogous to GetAdaptedTools for resources. - GetAdaptedResources(sessionID string) ([]mcpserver.ServerResource, error) - // GetMultiSession retrieves the fully-formed MultiSession for the given session ID. // Returns (nil, false) if the session does not exist or is still a placeholder. // Used to access session-scoped backend tool metadata (e.g. for conflict validation). diff --git a/pkg/vmcp/server/sessionmanager/session_manager.go b/pkg/vmcp/server/sessionmanager/session_manager.go index 54076d047e..c6175cff77 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager.go +++ b/pkg/vmcp/server/sessionmanager/session_manager.go @@ -14,7 +14,6 @@ package sessionmanager import ( "context" - "encoding/json" "errors" "fmt" "log/slog" @@ -22,14 +21,12 @@ import ( "time" "github.com/google/uuid" - "github.com/mark3labs/mcp-go/mcp" mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/cache" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" - "github.com/stacklok/toolhive/pkg/vmcp/conversion" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" @@ -825,219 +822,6 @@ func (sm *Manager) DecorateSession(sessionID string, fn func(sessiontypes.MultiS return nil } -// GetAdaptedTools returns SDK-format tools for the given session, with handlers -// that delegate tool invocations directly to the session's CallTool() method. -// -// When the session factory is configured with an aggregator (WithAggregator), -// tools are in their final resolved form — overrides and conflict resolution -// applied via ProcessPreQueriedCapabilities. Each handler passes the resolved -// tool name to CallTool, which translates it back to the original backend name -// via GetBackendCapabilityName. -// -// Without an aggregator, raw backend tool names are used as-is (no overrides -// or conflict resolution applied). -func (sm *Manager) GetAdaptedTools(sessionID string) ([]mcpserver.ServerTool, error) { - multiSess, ok := sm.GetMultiSession(sessionID) - if !ok { - return nil, fmt.Errorf("Manager.GetAdaptedTools: session %q not found or not a multi-session", sessionID) - } - - domainTools := multiSess.Tools() - sdkTools := make([]mcpserver.ServerTool, 0, len(domainTools)) - - for _, domainTool := range domainTools { - schemaJSON, err := json.Marshal(domainTool.InputSchema) - if err != nil { - return nil, fmt.Errorf("Manager.GetAdaptedTools: failed to marshal schema for tool %s: %w", domainTool.Name, err) - } - - tool := mcp.Tool{ - Name: domainTool.Name, - Description: domainTool.Description, - RawInputSchema: schemaJSON, - Annotations: conversion.ToMCPToolAnnotations(domainTool.Annotations), - } - if domainTool.OutputSchema != nil { - outputSchemaJSON, marshalErr := json.Marshal(domainTool.OutputSchema) - if marshalErr != nil { - slog.Warn("failed to marshal tool output schema", - "tool", domainTool.Name, "error", marshalErr) - } else { - tool.RawOutputSchema = outputSchemaJSON - } - } - - capturedSess := multiSess - capturedSessionID := sessionID - capturedToolName := domainTool.Name - handler := func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { - args, ok := req.Params.Arguments.(map[string]any) - if !ok { - wrappedErr := fmt.Errorf("%w: arguments must be object, got %T", vmcp.ErrInvalidInput, req.Params.Arguments) - slog.Warn("invalid arguments for tool", "tool", capturedToolName, "error", wrappedErr) - return mcp.NewToolResultError(wrappedErr.Error()), nil - } - - meta := conversion.FromMCPMeta(req.Params.Meta) - caller, _ := auth.IdentityFromContext(ctx) - - result, callErr := capturedSess.CallTool(ctx, caller, capturedToolName, args, meta) - if callErr != nil { - if errors.Is(callErr, sessiontypes.ErrUnauthorizedCaller) || errors.Is(callErr, sessiontypes.ErrNilCaller) { - slog.Warn("caller authorization failed, terminating session", - "session_id", capturedSessionID, "tool", capturedToolName, "error", callErr) - if _, termErr := sm.Terminate(capturedSessionID); termErr != nil { - slog.Error("failed to terminate session after auth failure", - "session_id", capturedSessionID, "error", termErr) - } - return mcp.NewToolResultError(fmt.Sprintf("Unauthorized: %v", callErr)), nil - } - return mcp.NewToolResultError(callErr.Error()), nil - } - - return &mcp.CallToolResult{ - Result: mcp.Result{ - Meta: conversion.ToMCPMeta(result.Meta), - }, - Content: conversion.ToMCPContents(result.Content), - StructuredContent: result.StructuredContent, - IsError: result.IsError, - }, nil - } - - sdkTools = append(sdkTools, mcpserver.ServerTool{ - Tool: tool, - Handler: handler, - }) - slog.Debug("Manager.GetAdaptedTools: adapted tool", "session_id", sessionID, "tool", domainTool.Name) - } - - return sdkTools, nil -} - -// GetAdaptedResources returns SDK-format resources for the given session, with handlers -// that delegate read requests directly to the session's ReadResource() method. -func (sm *Manager) GetAdaptedResources(sessionID string) ([]mcpserver.ServerResource, error) { - multiSess, ok := sm.GetMultiSession(sessionID) - if !ok { - return nil, fmt.Errorf("Manager.GetAdaptedResources: session %q not found or not a multi-session", sessionID) - } - - domainResources := multiSess.Resources() - sdkResources := make([]mcpserver.ServerResource, 0, len(domainResources)) - - for _, domainResource := range domainResources { - resource := mcp.Resource{ - Name: domainResource.Name, - URI: domainResource.URI, - Description: domainResource.Description, - MIMEType: domainResource.MimeType, - } - - capturedSess := multiSess - capturedSessionID := sessionID - capturedResourceURI := domainResource.URI - handler := func(ctx context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { - caller, _ := auth.IdentityFromContext(ctx) - - result, readErr := capturedSess.ReadResource(ctx, caller, capturedResourceURI) - if readErr != nil { - if errors.Is(readErr, sessiontypes.ErrUnauthorizedCaller) || errors.Is(readErr, sessiontypes.ErrNilCaller) { - slog.Warn("caller authorization failed, terminating session", - "session_id", capturedSessionID, "resource", capturedResourceURI, "error", readErr) - if _, termErr := sm.Terminate(capturedSessionID); termErr != nil { - slog.Error("failed to terminate session after auth failure", - "session_id", capturedSessionID, "error", termErr) - } - return nil, fmt.Errorf("unauthorized: %w", readErr) - } - return nil, readErr - } - - return conversion.ToMCPResourceContents(result.Contents), nil - } - - sdkResources = append(sdkResources, mcpserver.ServerResource{ - Resource: resource, - Handler: handler, - }) - slog.Debug("Manager.GetAdaptedResources: adapted resource", "session_id", sessionID, "uri", domainResource.URI) - } - - return sdkResources, nil -} - -// GetAdaptedPrompts returns SDK-format prompts for the given session, with handlers -// that delegate prompt requests directly to the session's GetPrompt() method. -func (sm *Manager) GetAdaptedPrompts(sessionID string) ([]mcpserver.ServerPrompt, error) { - multiSess, ok := sm.GetMultiSession(sessionID) - if !ok { - return nil, fmt.Errorf("Manager.GetAdaptedPrompts: session %q not found or not a multi-session", sessionID) - } - - domainPrompts := multiSess.Prompts() - sdkPrompts := make([]mcpserver.ServerPrompt, 0, len(domainPrompts)) - - for _, domainPrompt := range domainPrompts { - prompt := mcp.Prompt{ - Name: domainPrompt.Name, - Description: domainPrompt.Description, - } - for _, arg := range domainPrompt.Arguments { - prompt.Arguments = append(prompt.Arguments, mcp.PromptArgument{ - Name: arg.Name, - Description: arg.Description, - Required: arg.Required, - }) - } - - capturedSess := multiSess - capturedSessionID := sessionID - capturedPromptName := domainPrompt.Name - handler := func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { - caller, _ := auth.IdentityFromContext(ctx) - - args := make(map[string]any, len(req.Params.Arguments)) - for k, v := range req.Params.Arguments { - args[k] = v - } - result, getErr := capturedSess.GetPrompt(ctx, caller, capturedPromptName, args) - if getErr != nil { - if errors.Is(getErr, sessiontypes.ErrUnauthorizedCaller) || errors.Is(getErr, sessiontypes.ErrNilCaller) { - slog.Warn("caller authorization failed, terminating session", - "session_id", capturedSessionID, "prompt", capturedPromptName, "error", getErr) - if _, termErr := sm.Terminate(capturedSessionID); termErr != nil { - slog.Error("failed to terminate session after auth failure", - "session_id", capturedSessionID, "error", termErr) - } - return nil, fmt.Errorf("unauthorized: %w", getErr) - } - return nil, getErr - } - - mcpMessages := make([]mcp.PromptMessage, 0, len(result.Messages)) - for _, msg := range result.Messages { - mcpMessages = append(mcpMessages, mcp.PromptMessage{ - Role: mcp.Role(msg.Role), - Content: conversion.ToMCPContent(msg.Content), - }) - } - return &mcp.GetPromptResult{ - Description: result.Description, - Messages: mcpMessages, - }, nil - } - - sdkPrompts = append(sdkPrompts, mcpserver.ServerPrompt{ - Prompt: prompt, - Handler: handler, - }) - slog.Debug("Manager.GetAdaptedPrompts: adapted prompt", "session_id", sessionID, "prompt", domainPrompt.Name) - } - - return sdkPrompts, nil -} - // listAllBackends returns all backends from the registry as a pointer slice. func (sm *Manager) listAllBackends(ctx context.Context) []*vmcp.Backend { raw := sm.backendReg.List(ctx) diff --git a/pkg/vmcp/server/sessionmanager/session_manager_test.go b/pkg/vmcp/server/sessionmanager/session_manager_test.go index 8a8fb70cd8..69f20b5948 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager_test.go +++ b/pkg/vmcp/server/sessionmanager/session_manager_test.go @@ -11,7 +11,6 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" @@ -1104,903 +1103,6 @@ func TestSessionManager_GetMultiSession(t *testing.T) { }) } -// --------------------------------------------------------------------------- -// Tests: GetAdaptedTools -// --------------------------------------------------------------------------- - -func TestSessionManager_GetAdaptedTools(t *testing.T) { - t.Parallel() - - t.Run("returns error for unknown session", func(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - sess := newMockSession(t, ctrl, "", nil) - factory := newMockFactory(t, ctrl, sess) - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - _, err := sm.GetAdaptedTools("no-such-session") - require.Error(t, err) - assert.Contains(t, err.Error(), "not found or not a multi-session") - }) - - t.Run("returns tools with correct names and schemas", func(t *testing.T) { - t.Parallel() - - tools := []vmcp.Tool{ - { - Name: "alpha", - Description: "first tool", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "input": map[string]any{"type": "string"}, - }, - }, - }, - {Name: "beta", Description: "second tool"}, - } - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - return newMockSession(t, ctrl, id, tools), nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 2) - - byName := map[string]mcp.Tool{} - for _, st := range adaptedTools { - byName[st.Tool.Name] = st.Tool - } - - require.Contains(t, byName, "alpha") - require.Contains(t, byName, "beta") - - // InputSchema must be marshalled into RawInputSchema so clients - // receive the full parameter schema. - assert.NotEmpty(t, byName["alpha"].RawInputSchema) - assert.Contains(t, string(byName["alpha"].RawInputSchema), `"type"`) - }) - - t.Run("preserves annotations and output schema", func(t *testing.T) { - t.Parallel() - - boolPtr := func(b bool) *bool { return &b } - tools := []vmcp.Tool{ - { - Name: "annotated", - Description: "tool with annotations", - InputSchema: map[string]any{"type": "object"}, - OutputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "result": map[string]any{"type": "string"}, - }, - }, - Annotations: &vmcp.ToolAnnotations{ - Title: "Annotated Tool", - ReadOnlyHint: boolPtr(true), - DestructiveHint: boolPtr(false), - }, - }, - { - Name: "plain", - Description: "tool without annotations or output schema", - InputSchema: map[string]any{"type": "object"}, - }, - } - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - return newMockSession(t, ctrl, id, tools), nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 2) - - byName := map[string]mcp.Tool{} - for _, st := range adaptedTools { - byName[st.Tool.Name] = st.Tool - } - - // Verify annotations are preserved on the annotated tool. - annotated := byName["annotated"] - assert.Equal(t, "Annotated Tool", annotated.Annotations.Title) - require.NotNil(t, annotated.Annotations.ReadOnlyHint) - assert.True(t, *annotated.Annotations.ReadOnlyHint) - require.NotNil(t, annotated.Annotations.DestructiveHint) - assert.False(t, *annotated.Annotations.DestructiveHint) - assert.Nil(t, annotated.Annotations.IdempotentHint) - assert.Nil(t, annotated.Annotations.OpenWorldHint) - - // Verify output schema is preserved. - assert.NotNil(t, annotated.RawOutputSchema) - assert.Contains(t, string(annotated.RawOutputSchema), `"result"`) - - // Verify nil annotations produce zero-valued annotations and nil output schema. - plain := byName["plain"] - assert.Empty(t, plain.Annotations.Title) - assert.Nil(t, plain.Annotations.ReadOnlyHint) - assert.Nil(t, plain.RawOutputSchema) - }) - - t.Run("handlers delegate to session CallTool", func(t *testing.T) { - t.Parallel() - - tools := []vmcp.Tool{{Name: "greet", Description: "greets user"}} - ctrl := gomock.NewController(t) - - callToolResult := &vmcp.ToolCallResult{ - Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "Hello, world!"}}, - } - 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) { - sess := newMockSession(t, ctrl, id, tools) - sess.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - Return(callToolResult, nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 1) - - // Invoke the handler. - handler := adaptedTools[0].Handler - require.NotNil(t, handler) - - result, handlerErr := handler(context.Background(), newCallToolRequest("greet", nil)) - require.NoError(t, handlerErr) - require.NotNil(t, result) - require.Len(t, result.Content, 1) - // mcp.Content is an interface; assert the concrete TextContent type. - textContent, ok := result.Content[0].(mcp.TextContent) - require.True(t, ok, "expected TextContent") - assert.Equal(t, "Hello, world!", textContent.Text) - assert.False(t, result.IsError) - }) - - t.Run("handler returns tool error when CallTool fails", func(t *testing.T) { - t.Parallel() - - tools := []vmcp.Tool{{Name: "boom", Description: "always fails"}} - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := newMockSession(t, ctrl, id, tools) - sess.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - Return(nil, errors.New("backend exploded")).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 1) - - result, handlerErr := adaptedTools[0].Handler(context.Background(), newCallToolRequest("boom", nil)) - require.NoError(t, handlerErr, "handler should not return an error — it should wrap it in a tool result") - require.NotNil(t, result) - assert.True(t, result.IsError, "IsError should be set for failed tool calls") - }) - - t.Run("handler returns error result for non-object arguments", func(t *testing.T) { - t.Parallel() - - tools := []vmcp.Tool{{Name: "strict", Description: "requires object args"}} - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - return newMockSession(t, ctrl, id, tools), nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 1) - - // Pass a non-object argument (string instead of map). - req := mcp.CallToolRequest{} - req.Params.Name = "strict" - req.Params.Arguments = "not-an-object" - - result, handlerErr := adaptedTools[0].Handler(context.Background(), req) - require.NoError(t, handlerErr, "handler must not return a Go error") - require.NotNil(t, result) - assert.True(t, result.IsError, "non-object arguments should produce an error tool result") - }) - - t.Run("handler forwards request meta to CallTool", func(t *testing.T) { - t.Parallel() - - tools := []vmcp.Tool{{Name: "meta-tool", Description: "checks meta forwarding"}} - ctrl := gomock.NewController(t) - - var capturedMeta map[string]any - 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) { - sess := newMockSession(t, ctrl, id, tools) - sess.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, _ *auth.Identity, _ string, _ map[string]any, meta map[string]any) (*vmcp.ToolCallResult, error) { - capturedMeta = meta - return &vmcp.ToolCallResult{}, nil - }).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 1) - - // Build a request with a progress token in _meta. - req := mcp.CallToolRequest{} - req.Params.Name = "meta-tool" - req.Params.Arguments = map[string]any{} - req.Params.Meta = &mcp.Meta{ProgressToken: mcp.ProgressToken("tok-1")} - - _, handlerErr := adaptedTools[0].Handler(context.Background(), req) - require.NoError(t, handlerErr) - - // The meta must have been forwarded to CallTool. - require.NotNil(t, capturedMeta, "meta should be forwarded to CallTool") - assert.Equal(t, "tok-1", capturedMeta["progressToken"]) - }) - - t.Run("handler terminates session on authorization errors", func(t *testing.T) { - t.Parallel() - - // Test both ErrUnauthorizedCaller and ErrNilCaller - testCases := []struct { - name string - authError error - expectError string - }{ - { - name: "ErrUnauthorizedCaller", - authError: sessiontypes.ErrUnauthorizedCaller, - expectError: "Unauthorized", - }, - { - name: "ErrNilCaller", - authError: sessiontypes.ErrNilCaller, - expectError: "Unauthorized", - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - tools := []vmcp.Tool{{Name: "auth-tool", Description: "requires authorization"}} - ctrl := gomock.NewController(t) - authErr := tc.authError - 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) { - sess := newMockSession(t, ctrl, id, tools) - sess.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - Return(nil, authErr).Times(1) - // Close() is called when the session is terminated after auth failure - sess.EXPECT().Close().Return(nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedTools, err := sm.GetAdaptedTools(sessionID) - require.NoError(t, err) - require.Len(t, adaptedTools, 1) - - // Call the tool - should return an error result - req := newCallToolRequest("auth-tool", map[string]any{}) - result, handlerErr := adaptedTools[0].Handler(context.Background(), req) - require.NoError(t, handlerErr, "handler should not return Go error") - require.NotNil(t, result) - - // Verify error result contains "Unauthorized" - assert.True(t, result.IsError, "result should indicate error") - require.Len(t, result.Content, 1, "result should have content") - textContent, ok := result.Content[0].(mcp.TextContent) - require.True(t, ok, "expected TextContent") - assert.Contains(t, textContent.Text, tc.expectError) - - // Verify subsequent GetAdaptedTools fails (session no longer exists) - _, err = sm.GetAdaptedTools(sessionID) - assert.Error(t, err, "GetAdaptedTools should fail after session termination") - // gomock verifies Close() was called exactly once via Times(1) - }) - } - }) -} - -// --------------------------------------------------------------------------- -// Tests: GetAdaptedResources -// --------------------------------------------------------------------------- - -func TestSessionManager_GetAdaptedResources(t *testing.T) { - t.Parallel() - - t.Run("returns error for unknown session", func(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - sess := newMockSession(t, ctrl, "", nil) - factory := newMockFactory(t, ctrl, sess) - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - _, err := sm.GetAdaptedResources("no-such-session") - require.Error(t, err) - assert.Contains(t, err.Error(), "not found or not a multi-session") - }) - - t.Run("returns resources with correct fields", func(t *testing.T) { - t.Parallel() - - resources := []vmcp.Resource{ - { - Name: "config", - URI: "file:///etc/config.json", - Description: "Configuration file", - MimeType: "application/json", - }, - { - Name: "readme", - URI: "file:///README.md", - Description: "Readme", - MimeType: "text/markdown", - }, - } - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := newMockSession(t, ctrl, id, nil) - // Override default Resources() AnyTimes with a specific return - sess.EXPECT().Resources().Return(resources).AnyTimes() - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedResources, err := sm.GetAdaptedResources(sessionID) - require.NoError(t, err) - require.Len(t, adaptedResources, 2) - - byURI := map[string]mcp.Resource{} - for _, sr := range adaptedResources { - byURI[sr.Resource.URI] = sr.Resource - } - - require.Contains(t, byURI, "file:///etc/config.json") - require.Contains(t, byURI, "file:///README.md") - - assert.Equal(t, "config", byURI["file:///etc/config.json"].Name) - assert.Equal(t, "application/json", byURI["file:///etc/config.json"].MIMEType) - assert.Equal(t, "readme", byURI["file:///README.md"].Name) - assert.Equal(t, "text/markdown", byURI["file:///README.md"].MIMEType) - }) - - t.Run("handler delegates to session ReadResource", func(t *testing.T) { - t.Parallel() - - resources := []vmcp.Resource{ - { - Name: "data", - URI: "file:///data.txt", - MimeType: "text/plain", - }, - } - readResult := &vmcp.ResourceReadResult{ - Contents: []vmcp.ResourceContent{ - {URI: "file:///data.txt", MimeType: "text/plain", Text: "hello resource"}, - }, - } - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := newMockSession(t, ctrl, id, nil) - sess.EXPECT().Resources().Return(resources).AnyTimes() - sess.EXPECT().ReadResource(gomock.Any(), gomock.Any(), "file:///data.txt"). - Return(readResult, nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedResources, err := sm.GetAdaptedResources(sessionID) - require.NoError(t, err) - require.Len(t, adaptedResources, 1) - - req := mcp.ReadResourceRequest{} - req.Params.URI = "file:///data.txt" - contents, handlerErr := adaptedResources[0].Handler(context.Background(), req) - require.NoError(t, handlerErr) - require.Len(t, contents, 1) - - textContents, ok := contents[0].(mcp.TextResourceContents) - require.True(t, ok, "expected TextResourceContents") - assert.Equal(t, "file:///data.txt", textContents.URI) - assert.Equal(t, "text/plain", textContents.MIMEType) - assert.Equal(t, "hello resource", textContents.Text) - }) - - t.Run("handler returns error when ReadResource fails", func(t *testing.T) { - t.Parallel() - - resources := []vmcp.Resource{ - { - Name: "broken", - URI: "file:///broken.txt", - MimeType: "text/plain", - }, - } - readErr := errors.New("read failed") - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := newMockSession(t, ctrl, id, nil) - sess.EXPECT().Resources().Return(resources).AnyTimes() - sess.EXPECT().ReadResource(gomock.Any(), gomock.Any(), "file:///broken.txt"). - Return(nil, readErr).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedResources, err := sm.GetAdaptedResources(sessionID) - require.NoError(t, err) - require.Len(t, adaptedResources, 1) - - req := mcp.ReadResourceRequest{} - req.Params.URI = "file:///broken.txt" - contents, handlerErr := adaptedResources[0].Handler(context.Background(), req) - require.Error(t, handlerErr) - assert.Nil(t, contents) - assert.ErrorContains(t, handlerErr, "read failed") - }) - - t.Run("handler preserves empty MimeType from backend", func(t *testing.T) { - t.Parallel() - - resources := []vmcp.Resource{ - { - Name: "binary", - URI: "file:///binary.bin", - // MimeType intentionally empty - }, - } - readResult := &vmcp.ResourceReadResult{ - Contents: []vmcp.ResourceContent{ - {URI: "file:///binary.bin", MimeType: "", Text: "binary data"}, - }, - } - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := newMockSession(t, ctrl, id, nil) - sess.EXPECT().Resources().Return(resources).AnyTimes() - sess.EXPECT().ReadResource(gomock.Any(), gomock.Any(), "file:///binary.bin"). - Return(readResult, nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedResources, err := sm.GetAdaptedResources(sessionID) - require.NoError(t, err) - require.Len(t, adaptedResources, 1) - - req := mcp.ReadResourceRequest{} - req.Params.URI = "file:///binary.bin" - contents, handlerErr := adaptedResources[0].Handler(context.Background(), req) - require.NoError(t, handlerErr) - require.Len(t, contents, 1) - - textContents, ok := contents[0].(mcp.TextResourceContents) - require.True(t, ok, "expected TextResourceContents") - assert.Equal(t, "", textContents.MIMEType) - }) - - t.Run("handler terminates session on authorization errors", func(t *testing.T) { - t.Parallel() - - testCases := []struct { - name string - authError error - }{ - { - name: "ErrUnauthorizedCaller", - authError: sessiontypes.ErrUnauthorizedCaller, - }, - { - name: "ErrNilCaller", - authError: sessiontypes.ErrNilCaller, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - resources := []vmcp.Resource{ - { - Name: "protected", - URI: "file:///protected.txt", - }, - } - authErr := tc.authError - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := newMockSession(t, ctrl, id, nil) - sess.EXPECT().Resources().Return(resources).AnyTimes() - sess.EXPECT().ReadResource(gomock.Any(), gomock.Any(), "file:///protected.txt"). - Return(nil, authErr).Times(1) - // Close() is called when the session is terminated after auth failure - sess.EXPECT().Close().Return(nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedResources, err := sm.GetAdaptedResources(sessionID) - require.NoError(t, err) - require.Len(t, adaptedResources, 1) - - req := mcp.ReadResourceRequest{} - req.Params.URI = "file:///protected.txt" - contents, handlerErr := adaptedResources[0].Handler(context.Background(), req) - require.Error(t, handlerErr, "handler should return an error for auth failures") - assert.Nil(t, contents) - assert.ErrorContains(t, handlerErr, "unauthorized") - - // Verify subsequent GetAdaptedResources fails (session no longer exists) - _, err = sm.GetAdaptedResources(sessionID) - assert.Error(t, err, "GetAdaptedResources should fail after session termination") - // gomock verifies Close() was called exactly once via Times(1) - }) - } - }) -} - -// --------------------------------------------------------------------------- -// Tests: GetAdaptedPrompts -// --------------------------------------------------------------------------- - -func TestSessionManager_GetAdaptedPrompts(t *testing.T) { - t.Parallel() - - t.Run("returns error for unknown session", func(t *testing.T) { - t.Parallel() - - ctrl := gomock.NewController(t) - sess := newMockSession(t, ctrl, "", nil) - factory := newMockFactory(t, ctrl, sess) - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - _, err := sm.GetAdaptedPrompts("no-such-session") - require.Error(t, err) - assert.Contains(t, err.Error(), "not found or not a multi-session") - }) - - t.Run("returns prompts with correct fields and arguments", func(t *testing.T) { - t.Parallel() - - prompts := []vmcp.Prompt{ - { - Name: "greet", - Description: "Greet someone", - Arguments: []vmcp.PromptArgument{ - {Name: "name", Description: "Who to greet", Required: true}, - {Name: "language", Description: "Language to use", Required: false}, - }, - }, - { - Name: "summarize", - Description: "Summarize text", - }, - } - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - // Create mock directly (without newMockSession) so there is no - // pre-existing Prompts().Return(nil).AnyTimes() that would win - // the FIFO expectation race over our specific prompts list. - sess := sessionmocks.NewMockMultiSession(ctrl) - sess.EXPECT().ID().Return(id).AnyTimes() - sess.EXPECT().GetMetadata().Return(map[string]string{}).AnyTimes() - sess.EXPECT().Prompts().Return(prompts).AnyTimes() - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedPrompts, err := sm.GetAdaptedPrompts(sessionID) - require.NoError(t, err) - require.Len(t, adaptedPrompts, 2) - - byName := map[string]mcp.Prompt{} - for _, sp := range adaptedPrompts { - byName[sp.Prompt.Name] = sp.Prompt - } - - require.Contains(t, byName, "greet") - assert.Equal(t, "Greet someone", byName["greet"].Description) - require.Len(t, byName["greet"].Arguments, 2) - assert.Equal(t, "name", byName["greet"].Arguments[0].Name) - assert.True(t, byName["greet"].Arguments[0].Required) - assert.Equal(t, "language", byName["greet"].Arguments[1].Name) - assert.False(t, byName["greet"].Arguments[1].Required) - - require.Contains(t, byName, "summarize") - assert.Equal(t, "Summarize text", byName["summarize"].Description) - assert.Empty(t, byName["summarize"].Arguments) - }) - - t.Run("handler delegates to session GetPrompt", func(t *testing.T) { - t.Parallel() - - prompts := []vmcp.Prompt{ - { - Name: "hello", - Description: "Say hello", - Arguments: []vmcp.PromptArgument{{Name: "name", Required: true}}, - }, - } - getResult := &vmcp.PromptGetResult{ - Description: "A greeting", - Messages: []vmcp.PromptMessage{ - {Role: "assistant", Content: vmcp.Content{Type: vmcp.ContentTypeText, Text: "Hello, world!"}}, - }, - } - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := sessionmocks.NewMockMultiSession(ctrl) - sess.EXPECT().ID().Return(id).AnyTimes() - sess.EXPECT().GetMetadata().Return(map[string]string{}).AnyTimes() - sess.EXPECT().Prompts().Return(prompts).AnyTimes() - sess.EXPECT().GetPrompt(gomock.Any(), gomock.Any(), "hello", gomock.Any()). - Return(getResult, nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedPrompts, err := sm.GetAdaptedPrompts(sessionID) - require.NoError(t, err) - require.Len(t, adaptedPrompts, 1) - - req := mcp.GetPromptRequest{} - req.Params.Name = "hello" - req.Params.Arguments = map[string]string{"name": "Alice"} - result, handlerErr := adaptedPrompts[0].Handler(context.Background(), req) - require.NoError(t, handlerErr) - require.NotNil(t, result) - assert.Equal(t, "A greeting", result.Description) - require.Len(t, result.Messages, 1) - assert.Equal(t, mcp.RoleAssistant, result.Messages[0].Role) - }) - - t.Run("handler returns error when GetPrompt fails", func(t *testing.T) { - t.Parallel() - - prompts := []vmcp.Prompt{{Name: "broken"}} - getErr := errors.New("prompt backend error") - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := sessionmocks.NewMockMultiSession(ctrl) - sess.EXPECT().ID().Return(id).AnyTimes() - sess.EXPECT().GetMetadata().Return(map[string]string{}).AnyTimes() - sess.EXPECT().Prompts().Return(prompts).AnyTimes() - sess.EXPECT().GetPrompt(gomock.Any(), gomock.Any(), "broken", gomock.Any()). - Return(nil, getErr).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedPrompts, err := sm.GetAdaptedPrompts(sessionID) - require.NoError(t, err) - require.Len(t, adaptedPrompts, 1) - - req := mcp.GetPromptRequest{} - req.Params.Name = "broken" - result, handlerErr := adaptedPrompts[0].Handler(context.Background(), req) - require.Error(t, handlerErr) - assert.Nil(t, result) - assert.ErrorContains(t, handlerErr, "prompt backend error") - }) - - t.Run("handler terminates session on authorization errors", func(t *testing.T) { - t.Parallel() - - testCases := []struct { - name string - authError error - }{ - {name: "ErrUnauthorizedCaller", authError: sessiontypes.ErrUnauthorizedCaller}, - {name: "ErrNilCaller", authError: sessiontypes.ErrNilCaller}, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - prompts := []vmcp.Prompt{{Name: "secret"}} - authErr := tc.authError - - ctrl := gomock.NewController(t) - factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) - factory.EXPECT(). - MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { - sess := sessionmocks.NewMockMultiSession(ctrl) - sess.EXPECT().ID().Return(id).AnyTimes() - sess.EXPECT().GetMetadata().Return(map[string]string{}).AnyTimes() - sess.EXPECT().Prompts().Return(prompts).AnyTimes() - sess.EXPECT().GetPrompt(gomock.Any(), gomock.Any(), "secret", gomock.Any()). - Return(nil, authErr).Times(1) - // Close() is called when the session is terminated after auth failure. - sess.EXPECT().Close().Return(nil).Times(1) - return sess, nil - }).Times(1) - - registry := newFakeRegistry() - sm, _ := newTestSessionManager(t, factory, registry) - - sessionID := sm.Generate() - _, err := sm.CreateSession(context.Background(), sessionID) - require.NoError(t, err) - - adaptedPrompts, err := sm.GetAdaptedPrompts(sessionID) - require.NoError(t, err) - require.Len(t, adaptedPrompts, 1) - - req := mcp.GetPromptRequest{} - req.Params.Name = "secret" - result, handlerErr := adaptedPrompts[0].Handler(context.Background(), req) - require.Error(t, handlerErr, "handler should return an error for auth failures") - assert.Nil(t, result) - assert.ErrorContains(t, handlerErr, "unauthorized") - - // Verify subsequent GetAdaptedPrompts fails (session no longer exists). - _, err = sm.GetAdaptedPrompts(sessionID) - assert.Error(t, err, "GetAdaptedPrompts should fail after session termination") - // gomock verifies Close() was called exactly once via Times(1) - }) - } - }) -} - // --------------------------------------------------------------------------- // Tests: DecorateSession // --------------------------------------------------------------------------- @@ -2751,15 +1853,3 @@ func TestTerminate_LegacyFormatSession_TakesPlaceholderPath(t *testing.T) { assert.Equal(t, MetadataValTrue, metadata[MetadataKeyTerminated], "legacy session Terminate must set MetadataKeyTerminated rather than deleting") } - -// --------------------------------------------------------------------------- -// Helper -// --------------------------------------------------------------------------- - -// newCallToolRequest builds a minimal mcp.CallToolRequest for handler tests. -func newCallToolRequest(name string, args map[string]any) mcp.CallToolRequest { - req := mcp.CallToolRequest{} - req.Params.Name = name - req.Params.Arguments = args - return req -}