Update MCP parser to support latest specification (#1993)

* Update MCP parser to support latest specification

Add support for new MCP methods from the 2025-06-18 specification:
- elicitation/create for user input requests
- sampling/createMessage for message creation
- resources/subscribe and resources/unsubscribe for resource subscriptions
- resources/templates/list for template listing
- roots/list for listing roots
- notifications/progress for progress notifications
- notifications/cancelled for cancellation notifications
- Additional notification methods for list changes

Update completion/complete handler to properly handle both
PromptReference and ResourceTemplateReference types as specified
in the latest schema.

Add comprehensive test coverage for all new methods to ensure
proper extraction of resource IDs and arguments.

* test: add comprehensive test coverage for MCP parser

- Added tests for all new MCP methods from 2025-06-18 specification
- Added edge case tests for error handling and malformed inputs
- Added tests for JSON-RPC notifications (messages without ID)
- Improved test coverage from ~66% to 91.7% for pkg/mcp
- All parser functions now have 83-100% coverage

* Address PR feedback

Signed-off-by: Juan Antonio Osorio <[email protected]>

---------

Signed-off-by: Juan Antonio Osorio <[email protected]>

Author: JAORMX

pkg/mcp/parser.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"io"
 	"net/http"
+	"strconv"
 	"strings"
 
 	"golang.org/x/exp/jsonrpc2"
@@ -135,18 +136,25 @@ func parseMCPRequest(bodyBytes []byte) *ParsedMCPRequest {
 		return nil
 	}
 
-	// Handle only request messages
+	// Handle only request messages (both calls with ID and notifications without ID)
 	req, ok := msg.(*jsonrpc2.Request)
 	if !ok {
+		// Response or error messages are not parsed here
 		return nil
 	}
 
 	// Extract resource ID and arguments based on the method
 	resourceID, arguments := extractResourceAndArguments(req.Method, req.Params)
 
+	// Determine the ID - will be nil for notifications
+	var id interface{}
+	if req.ID.IsValid() {
+		id = req.ID.Raw()
+	}
+
 	return &ParsedMCPRequest{
 		Method:     req.Method,
-		ID:         req.ID.Raw(),
+		ID:         id,
 		Params:     req.Params,
 		ResourceID: resourceID,
 		Arguments:  arguments,
@@ -162,24 +170,36 @@ type methodHandler func(map[string]interface{}) (string, map[string]interface{})
 
 // methodHandlers maps MCP methods to their respective handlers
 var methodHandlers = map[string]methodHandler{
-	"initialize":            handleInitializeMethod,
-	"tools/call":            handleNamedResourceMethod,
-	"prompts/get":           handleNamedResourceMethod,
-	"resources/read":        handleResourceReadMethod,
-	"resources/list":        handleListMethod,
-	"tools/list":            handleListMethod,
-	"prompts/list":          handleListMethod,
-	"progress/update":       handleProgressMethod,
-	"notifications/message": handleNotificationMethod,
-	"logging/setLevel":      handleLoggingMethod,
-	"completion/complete":   handleCompletionMethod,
+	"initialize":               handleInitializeMethod,
+	"tools/call":               handleNamedResourceMethod,
+	"prompts/get":              handleNamedResourceMethod,
+	"resources/read":           handleResourceReadMethod,
+	"resources/list":           handleListMethod,
+	"tools/list":               handleListMethod,
+	"prompts/list":             handleListMethod,
+	"progress/update":          handleProgressMethod,
+	"notifications/message":    handleNotificationMethod,
+	"logging/setLevel":         handleLoggingMethod,
+	"completion/complete":      handleCompletionMethod,
+	"elicitation/create":       handleElicitationMethod,
+	"sampling/createMessage":   handleSamplingMethod,
+	"resources/subscribe":      handleResourceSubscribeMethod,
+	"resources/unsubscribe":    handleResourceUnsubscribeMethod,
+	"resources/templates/list": handleListMethod,
+	"roots/list":               handleListMethod,
+	"notifications/progress":   handleProgressNotificationMethod,
+	"notifications/cancelled":  handleCancelledNotificationMethod,
 }
 
 // staticResourceIDs maps methods to their static resource IDs
 var staticResourceIDs = map[string]string{
-	"ping":                             "ping",
-	"notifications/roots/list_changed": "roots",
-	"notifications/initialized":        "initialized",
+	"ping":                                 "ping",
+	"notifications/roots/list_changed":     "roots",
+	"notifications/initialized":            "initialized",
+	"notifications/prompts/list_changed":   "prompts",
+	"notifications/resources/list_changed": "resources",
+	"notifications/resources/updated":      "resources",
+	"notifications/tools/list_changed":     "tools",
 }
 
 func extractResourceAndArguments(method string, params json.RawMessage) (string, map[string]interface{}) {
@@ -277,14 +297,114 @@ func handleLoggingMethod(paramsMap map[string]interface{}) (string, map[string]i
 	return "", nil
 }
 
-// handleCompletionMethod extracts resource ID for completion requests
+// handleCompletionMethod extracts resource ID for completion requests.
+// For PromptReference: extracts the prompt name
+// For ResourceTemplateReference: extracts the template URI
+// For legacy string ref: returns the string value
+// Always returns paramsMap as arguments since completion requests need the full context
+// including the argument being completed and any context from previous completions.
 func handleCompletionMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	// Check if ref is a map (PromptReference or ResourceTemplateReference)
+	if ref, ok := paramsMap["ref"].(map[string]interface{}); ok {
+		// Try to extract name for PromptReference
+		if name, ok := ref["name"].(string); ok {
+			return name, paramsMap
+		}
+		// Try to extract uri for ResourceTemplateReference
+		if uri, ok := ref["uri"].(string); ok {
+			return uri, paramsMap
+		}
+	}
+	// Fallback to string ref (legacy support)
 	if ref, ok := paramsMap["ref"].(string); ok {
-		return ref, nil
+		return ref, paramsMap
+	}
+	return "", paramsMap
+}
+
+// handleElicitationMethod extracts resource ID for elicitation requests
+func handleElicitationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	// The message field could be used as a resource identifier
+	if message, ok := paramsMap["message"].(string); ok {
+		return message, paramsMap
+	}
+	return "", paramsMap
+}
+
+// handleSamplingMethod extracts resource ID for sampling/createMessage requests.
+// Returns the model name from modelPreferences if available, otherwise returns a
+// truncated version of the systemPrompt. The 50-character truncation provides a
+// reasonable balance between uniqueness and readability for authorization and audit logs.
+func handleSamplingMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	// Use model preferences or system prompt as identifier if available
+	if modelPrefs, ok := paramsMap["modelPreferences"].(map[string]interface{}); ok && modelPrefs != nil {
+		// Try direct name field first (simplified structure)
+		if name, ok := modelPrefs["name"].(string); ok && name != "" {
+			return name, paramsMap
+		}
+		// Try to get model name from hints array (full spec structure)
+		if hints, ok := modelPrefs["hints"].([]interface{}); ok && len(hints) > 0 {
+			if hint, ok := hints[0].(map[string]interface{}); ok {
+				if name, ok := hint["name"].(string); ok && name != "" {
+					return name, paramsMap
+				}
+			}
+		}
+	}
+	if systemPrompt, ok := paramsMap["systemPrompt"].(string); ok && systemPrompt != "" {
+		// Use first 50 chars of system prompt as identifier
+		// This provides a reasonable balance between uniqueness and readability
+		if len(systemPrompt) > 50 {
+			return systemPrompt[:50], paramsMap
+		}
+		return systemPrompt, paramsMap
+	}
+	return "", paramsMap
+}
+
+// handleResourceSubscribeMethod extracts resource ID for resource subscribe operations
+func handleResourceSubscribeMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	if uri, ok := paramsMap["uri"].(string); ok {
+		return uri, nil
 	}
 	return "", nil
 }
 
+// handleResourceUnsubscribeMethod extracts resource ID for resource unsubscribe operations
+func handleResourceUnsubscribeMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	if uri, ok := paramsMap["uri"].(string); ok {
+		return uri, nil
+	}
+	return "", nil
+}
+
+// handleProgressNotificationMethod extracts resource ID for progress notifications.
+// Extracts the progressToken which can be either a string or numeric value.
+func handleProgressNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	if token, ok := paramsMap["progressToken"].(string); ok {
+		return token, paramsMap
+	}
+	// Also handle numeric progress tokens
+	if token, ok := paramsMap["progressToken"].(float64); ok {
+		return strconv.FormatFloat(token, 'f', 0, 64), paramsMap
+	}
+	return "", paramsMap
+}
+
+// handleCancelledNotificationMethod extracts resource ID for cancelled notifications.
+// Extracts the requestId which can be either a string or numeric value.
+func handleCancelledNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) {
+	// Extract request ID as the resource identifier
+	if requestId, ok := paramsMap["requestId"].(string); ok {
+		return requestId, paramsMap
+	}
+	// Handle numeric request IDs
+	if requestId, ok := paramsMap["requestId"].(float64); ok {
+		return strconv.FormatFloat(requestId, 'f', 0, 64), paramsMap
+	}
+	return "", paramsMap
+}
+
 // GetMCPMethod is a convenience function to get the MCP method from the context.
 func GetMCPMethod(ctx context.Context) string {
 	if parsed := GetParsedMCPRequest(ctx); parsed != nil {