Update MCP parser to support latest specification

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 {