diff --git a/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go b/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go index 255955597f..2ed35e584e 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go +++ b/pkg/transport/proxy/streamable/streamable_proxy_spec_test.go @@ -147,8 +147,9 @@ func TestInitializeSetsSessionHeader(t *testing.T) { require.NotEmpty(t, sessID, "server should set Mcp-Session-Id header") } -// TestPOSTNotificationOnlyAccepted checks single notification POST returns 202 with no body. -func TestPOSTNotificationOnlyAccepted(t *testing.T) { +// TestPOSTNotificationAndClientResponseAccepted checks that one-way client +// messages return 202 with no body and are forwarded to the destination. +func TestPOSTNotificationAndClientResponseAccepted(t *testing.T) { t.Parallel() const port = 8104 @@ -162,18 +163,32 @@ func TestPOSTNotificationOnlyAccepted(t *testing.T) { time.Sleep(50 * time.Millisecond) url := "http://127.0.0.1:8104" + StreamableHTTPEndpoint - // Notification (no id) - notif := `{"jsonrpc":"2.0","method":"progress","params":{"pct":50}}` + assertAccepted := func(name, message string) { + t.Helper() + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, url, bytes.NewReader([]byte(message))) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusAccepted, resp.StatusCode, name) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Empty(t, body, "%s: 202 should have no body", name) - req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader([]byte(notif))) - req.Header.Set("Content-Type", "application/json") - resp, err := http.DefaultClient.Do(req) - require.NoError(t, err) - defer resp.Body.Close() + select { + case msg := <-proxy.GetMessageChannel(): + encoded, err := jsonrpc2.EncodeMessage(msg) + require.NoError(t, err) + assert.JSONEq(t, message, string(encoded), name) + case <-time.After(time.Second): + t.Fatalf("%s was not forwarded to the destination", name) + } + } - assert.Equal(t, http.StatusAccepted, resp.StatusCode) - body, _ := io.ReadAll(resp.Body) - assert.Equal(t, 0, len(body), "202 should have no body") + assertAccepted("notification", `{"jsonrpc":"2.0","method":"progress","params":{"pct":50}}`) + assertAccepted("client response", `{"jsonrpc":"2.0","id":"server-1","result":{}}`) } // TestBatchRequestsRejected verifies that the streamable proxy rejects every diff --git a/pkg/vmcp/server/client_response_integration_test.go b/pkg/vmcp/server/client_response_integration_test.go new file mode 100644 index 0000000000..1368cdf05a --- /dev/null +++ b/pkg/vmcp/server/client_response_integration_test.go @@ -0,0 +1,285 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/auth" + mcpparser "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" +) + +type elicitingCore struct { + *fakeCore + requester vmcp.ElicitationRequester +} + +func (c *elicitingCore) CallTool( + ctx context.Context, _ *auth.Identity, _ string, _ map[string]any, _ map[string]any, +) (*vmcp.ToolCallResult, error) { + c.callToolCalls.Add(1) + result, err := c.requester.RequestElicitation(ctx, vmcp.ElicitationRequest{ + Message: "What is your name?", + RequestedSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string"}, + }, + }, + }) + if err != nil { + return nil, err + } + content, ok := result.Content.(map[string]any) + if !ok { + return nil, fmt.Errorf("unexpected elicitation content type %T", result.Content) + } + name, _ := content["name"].(string) + return &vmcp.ToolCallResult{ + Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: "hello " + name}}, + }, nil +} + +type responsePostObservation struct { + statusCode int + body []byte + readErr error +} + +type responsePostRecordingTransport struct { + base http.RoundTripper + observed chan<- responsePostObservation +} + +func legacyRequestAuthorization(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + next.ServeHTTP(w, r) + return + } + if mcpparser.GetParsedMCPRequest(r.Context()) == nil { + http.Error(w, "Invalid or malformed MCP request", http.StatusBadRequest) + return + } + next.ServeHTTP(w, r) + }) +} + +func (t *responsePostRecordingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + isResponse := false + if req.Method == http.MethodPost && req.Body != nil { + requestBody, err := io.ReadAll(req.Body) + if err != nil { + return nil, err + } + _ = req.Body.Close() + req.Body = io.NopCloser(bytes.NewReader(requestBody)) + + var payload map[string]json.RawMessage + if json.Unmarshal(requestBody, &payload) == nil { + _, hasID := payload["id"] + _, hasMethod := payload["method"] + _, hasResult := payload["result"] + _, hasError := payload["error"] + isResponse = hasID && !hasMethod && (hasResult || hasError) + } + } + + base := t.base + if base == nil { + base = http.DefaultTransport + } + resp, err := base.RoundTrip(req) + if err != nil || !isResponse { + return resp, err + } + + responseBody, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(responseBody)) + observation := responsePostObservation{ + statusCode: resp.StatusCode, + body: responseBody, + readErr: readErr, + } + select { + case t.observed <- observation: + default: + // This test expects one response POST; discard any retry after that observation. + } + return resp, nil +} + +func cleanupWithin(t *testing.T, name string, cleanup func() error) { + t.Helper() + done := make(chan error, 1) + go func() { done <- cleanup() }() + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(5 * time.Second): + t.Errorf("timed out cleaning up %s", name) + } +} + +// TestHandler_AcceptsClientResponsePost guards two Serve-path regressions: a +// Legacy client response POST must return 202, and configuring the retired HTTP +// authorization middleware must not reinstall it in Handler's request chain. +func TestHandler_AcceptsClientResponsePost(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + t.Cleanup(cancel) + + ctrl := gomock.NewController(t) + testTool := vmcp.Tool{ + Name: "elicit", + Description: "elicits a name before responding", + InputSchema: map[string]any{"type": "object"}, + } + factory, _ := newToolSessionFactory(t, ctrl, []vmcp.Tool{testTool}) + // Optimizer is a real construction-time Modern blocker. Enabling it makes the + // Modern-first client negotiate the session-based Legacy path under test. + optimizerFactory := &recordingOptimizerFactory{} + core := &elicitingCore{fakeCore: &fakeCore{tools: []vmcp.Tool{testTool}}} + + srv, err := Serve(ctx, core, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{ + Base: factory, + OptimizerFactory: optimizerFactory.build, + AdvertiseFromCore: true, + }, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { + cleanupWithin(t, "vMCP server", func() error { return srv.Stop(context.Background()) }) + }) + core.requester = NewSDKElicitationAdapter(srv.MCPServer()) + + // Regression canary: Handler no longer reads config.AuthzMiddleware (the HTTP + // authz block was removed in #5556). Setting it here proves that reintroducing + // that block would reject response POSTs as malformed. + srv.config.AuthMiddleware = mcpparser.ParsingMiddleware + srv.config.AuthzMiddleware = legacyRequestAuthorization + + handler, err := srv.Handler(ctx) + require.NoError(t, err) + httpServer := httptest.NewServer(handler) + t.Cleanup(func() { + cleanupWithin(t, "HTTP server", func() error { + httpServer.CloseClientConnections() + httpServer.Close() + return nil + }) + }) + + observed := make(chan responsePostObservation, 1) + httpClient := &http.Client{Transport: &responsePostRecordingTransport{ + base: http.DefaultTransport, + observed: observed, + }} + client, err := mcpclient.NewStreamableHttpClientWithOpts( + httpServer.URL+"/mcp", + []mcptransport.StreamableHTTPCOption{ + mcptransport.WithContinuousListening(), + mcptransport.WithHTTPBasicClient(httpClient), + }, + []mcpclient.ClientOption{mcpclient.WithElicitationHandler(mcpclient.ElicitationHandlerFunc( + func(_ context.Context, _ mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + return &mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{ + Action: mcp.ElicitationResponseActionAccept, + Content: map[string]any{"name": "grace"}, + }, + }, nil + }, + ))}, + ) + require.NoError(t, err) + require.NoError(t, client.Start(ctx)) + t.Cleanup(func() { + cancel() + cleanupWithin(t, "MCP client", client.Close) + }) + + initializeResult, err := client.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "response-test", Version: "1.0.0"}, + }, + }) + require.NoError(t, err) + require.Equal(t, mcp.LATEST_PROTOCOL_VERSION, initializeResult.ProtocolVersion) + + type callOutcome struct { + result *mcp.CallToolResult + err error + } + callDone := make(chan callOutcome, 1) + go func() { + result, err := client.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: optimizerdec.CallToolName, + Arguments: map[string]any{ + "tool_name": testTool.Name, + "parameters": map[string]any{}, + }, + }, + }) + callDone <- callOutcome{result: result, err: err} + }() + + var responsePost responsePostObservation + select { + case responsePost = <-observed: + case outcome := <-callDone: + t.Fatalf( + "tool call finished before client response POST: result=%v error=%v", + outcome.result, + outcome.err, + ) + case <-ctx.Done(): + t.Fatalf("timed out waiting for client response POST: %v", ctx.Err()) + } + require.NoError(t, responsePost.readErr) + if responsePost.statusCode != http.StatusAccepted { + cancel() + } + assert.Equal(t, http.StatusAccepted, responsePost.statusCode) + assert.Empty(t, responsePost.body) + + var outcome callOutcome + select { + case outcome = <-callDone: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for tool call after client response POST") + } + require.NoError(t, outcome.err) + require.NotNil(t, outcome.result) + require.False(t, outcome.result.IsError) + require.Len(t, outcome.result.Content, 1) + text, ok := mcp.AsTextContent(outcome.result.Content[0]) + require.True(t, ok) + assert.Equal(t, "hello grace", text.Text) +}