diff --git a/server/streamable_http.go b/server/streamable_http.go index 2ec34f2b1..f205bbddd 100644 --- a/server/streamable_http.go +++ b/server/streamable_http.go @@ -959,7 +959,17 @@ func (s *StreamableHTTPServer) handleSamplingResponse(w HTTPResponseWriter, r *H if err := json.Unmarshal(responseMessage.Error, &jsonrpcError); err != nil { response.err = fmt.Errorf("failed to parse error: %v", err) } else { - response.err = fmt.Errorf("sampling error %d: %s", jsonrpcError.Code, jsonrpcError.Message) + method := mcp.MethodSamplingCreateMessage + if sessionInterface, ok := s.activeSessions.Load(sessionID); ok { + if session, ok := sessionInterface.(*streamableHttpSession); ok { + if pendingInterface, exists := session.samplingRequests.Load(requestID); exists { + if pending, ok := pendingInterface.(pendingClientRequest); ok { + method = pending.method + } + } + } + } + response.err = fmt.Errorf("%s error %d: %s", method, jsonrpcError.Code, jsonrpcError.Message) } } else if responseMessage.Result != nil { // Store the result to be unmarshaled later @@ -997,18 +1007,20 @@ func (s *StreamableHTTPServer) deliverSamplingResponse(w HTTPResponseWriter, ses } // Look up the dedicated response channel for this specific request - responseChannelInterface, exists := session.samplingRequests.Load(response.requestID) + pendingInterface, exists := session.samplingRequests.Load(response.requestID) if !exists { writeHTTPError(w, "No pending sampling request found for the given request ID", http.StatusBadRequest) return fmt.Errorf("no pending request found for session %s, request %d", sessionID, response.requestID) } - responseChan, ok := responseChannelInterface.(chan samplingResponseItem) + pending, ok := pendingInterface.(pendingClientRequest) if !ok { writeHTTPError(w, "Failed to deliver response", http.StatusInternalServerError) - return fmt.Errorf("invalid response channel type for session %s, request %d", sessionID, response.requestID) + return fmt.Errorf("invalid pending request type for session %s, request %d", sessionID, response.requestID) } + responseChan := pending.response + // Attempt to deliver the response with timeout to prevent indefinite blocking select { case responseChan <- response: @@ -1270,6 +1282,11 @@ type samplingRequestItem struct { response chan samplingResponseItem } +type pendingClientRequest struct { + response chan samplingResponseItem + method mcp.MCPMethod +} + type samplingResponseItem struct { requestID int64 result json.RawMessage @@ -1410,7 +1427,10 @@ func (s *streamableHttpSession) RequestSampling(ctx context.Context, request mcp } // Store the pending request - s.samplingRequests.Store(requestID, responseChan) + s.samplingRequests.Store(requestID, pendingClientRequest{ + response: responseChan, + method: mcp.MethodSamplingCreateMessage, + }) defer s.samplingRequests.Delete(requestID) // Send the sampling request via the channel (non-blocking) @@ -1467,7 +1487,10 @@ func (s *streamableHttpSession) ListRoots(ctx context.Context, request mcp.ListR } // Store the pending request - s.samplingRequests.Store(requestID, responseChan) + s.samplingRequests.Store(requestID, pendingClientRequest{ + response: responseChan, + method: mcp.MethodListRoots, + }) defer s.samplingRequests.Delete(requestID) // Send the list roots request via the channel (non-blocking) @@ -1512,7 +1535,10 @@ func (s *streamableHttpSession) RequestElicitation(ctx context.Context, request } // Store the pending request - s.samplingRequests.Store(requestID, responseChan) + s.samplingRequests.Store(requestID, pendingClientRequest{ + response: responseChan, + method: mcp.MethodElicitationCreate, + }) defer s.samplingRequests.Delete(requestID) // Send the sampling request via the channel (non-blocking) diff --git a/server/streamable_http_sampling_test.go b/server/streamable_http_sampling_test.go index 77a6e4155..005a1749a 100644 --- a/server/streamable_http_sampling_test.go +++ b/server/streamable_http_sampling_test.go @@ -214,3 +214,68 @@ func TestStreamableHTTPServer_SamplingQueueFull(t *testing.T) { t.Errorf("Expected queue full error, got: %v", err) } } + +// TestStreamableHTTPServer_ClientRequestErrorIncludesMethod verifies JSON-RPC client +// errors include the originating MCP method (e.g. elicitation/create vs sampling/createMessage). +func TestStreamableHTTPServer_ClientRequestErrorIncludesMethod(t *testing.T) { + mcpServer := NewMCPServer("test-server", "1.0.0") + mcpServer.EnableSampling() + + httpServer := NewStreamableHTTPServer(mcpServer, WithStateLess(true)) + testServer := httptest.NewServer(httpServer) + defer testServer.Close() + + sessionID := "test-session-elicitation-error" + session := newStreamableHttpSession(sessionID, httpServer.sessionTools, httpServer.sessionResources, httpServer.sessionResourceTemplates, httpServer.sessionLogLevels) + httpServer.activeSessions.Store(sessionID, session) + + requestID := int64(99) + responseChan := make(chan samplingResponseItem, 1) + session.samplingRequests.Store(requestID, pendingClientRequest{ + response: responseChan, + method: mcp.MethodElicitationCreate, + }) + + body := map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "error": map[string]any{ + "code": -32601, + "message": "Method not found", + }, + } + payload, err := json.Marshal(body) + if err != nil { + t.Fatalf("Failed to marshal body: %v", err) + } + + req, err := http.NewRequest("POST", testServer.URL, bytes.NewReader(payload)) + if err != nil { + t.Fatalf("Failed to create request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", sessionID) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to send request: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("Expected status %d, got %d", http.StatusAccepted, resp.StatusCode) + } + + select { + case response := <-responseChan: + if response.err == nil { + t.Fatal("Expected error response, got nil") + } + expected := "elicitation/create error -32601: Method not found" + if response.err.Error() != expected { + t.Errorf("Expected error %q, got %q", expected, response.err.Error()) + } + case <-time.After(2 * time.Second): + t.Fatal("Timed out waiting for elicitation error response") + } +}