diff --git a/pkg/mcp/parser.go b/pkg/mcp/parser.go index f797f857b..d6344390d 100644 --- a/pkg/mcp/parser.go +++ b/pkg/mcp/parser.go @@ -174,25 +174,30 @@ 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, - "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, + "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, + "tasks/list": handleListMethod, + "tasks/get": handleTaskIDMethod, + "tasks/cancel": handleTaskIDMethod, + "tasks/result": handleTaskIDMethod, + "notifications/tasks/status": handleTaskStatusNotificationMethod, } // staticResourceIDs maps methods to their static resource IDs @@ -418,6 +423,34 @@ func handleCancelledNotificationMethod(paramsMap map[string]interface{}) (string return "", paramsMap } +// handleTaskIDMethod extracts resource ID for task operations (tasks/get, tasks/cancel, tasks/result). +// Returns the taskId parameter as the resource identifier, or empty string if not present. +// Handles both string and numeric taskId values. +func handleTaskIDMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) { + if taskId, ok := paramsMap["taskId"].(string); ok { + return taskId, nil + } + // Handle numeric task IDs + if taskId, ok := paramsMap["taskId"].(float64); ok { + return strconv.FormatFloat(taskId, 'f', 0, 64), nil + } + return "", nil +} + +// handleTaskStatusNotificationMethod extracts resource ID for task status notifications. +// Returns the taskId parameter as the resource identifier while preserving all notification parameters. +// Handles both string and numeric taskId values. +func handleTaskStatusNotificationMethod(paramsMap map[string]interface{}) (string, map[string]interface{}) { + if taskId, ok := paramsMap["taskId"].(string); ok { + return taskId, paramsMap + } + // Handle numeric task IDs + if taskId, ok := paramsMap["taskId"].(float64); ok { + return strconv.FormatFloat(taskId, '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 { diff --git a/pkg/mcp/parser_test.go b/pkg/mcp/parser_test.go index f7858ced3..9228b3936 100644 --- a/pkg/mcp/parser_test.go +++ b/pkg/mcp/parser_test.go @@ -404,6 +404,85 @@ func TestExtractResourceAndArguments(t *testing.T) { "requestId": float64(456), }, }, + { + name: "tasks/get with taskId", + method: "tasks/get", + params: `{"taskId":"786512e2-9e0d-44bd-8f29-789f320fe840"}`, + expectedResourceID: "786512e2-9e0d-44bd-8f29-789f320fe840", + expectedArguments: nil, + }, + { + name: "tasks/cancel with taskId", + method: "tasks/cancel", + params: `{"taskId":"abc-123-def-456"}`, + expectedResourceID: "abc-123-def-456", + expectedArguments: nil, + }, + { + name: "tasks/result with taskId", + method: "tasks/result", + params: `{"taskId":"task-result-id-789"}`, + expectedResourceID: "task-result-id-789", + expectedArguments: nil, + }, + { + name: "tasks/get with numeric taskId", + method: "tasks/get", + params: `{"taskId":12345}`, + expectedResourceID: "12345", + expectedArguments: nil, + }, + { + name: "tasks/cancel with numeric taskId", + method: "tasks/cancel", + params: `{"taskId":67890}`, + expectedResourceID: "67890", + expectedArguments: nil, + }, + { + name: "tasks/result with numeric taskId", + method: "tasks/result", + params: `{"taskId":11111}`, + expectedResourceID: "11111", + expectedArguments: nil, + }, + { + name: "tasks/list with cursor", + method: "tasks/list", + params: `{"cursor":"next-page-cursor"}`, + expectedResourceID: "next-page-cursor", + expectedArguments: nil, + }, + { + name: "tasks/list without cursor", + method: "tasks/list", + params: `{}`, + expectedResourceID: "", + expectedArguments: nil, + }, + { + name: "notifications/tasks/status with taskId", + method: "notifications/tasks/status", + params: `{"taskId":"status-notification-task-id","status":"completed","createdAt":"2025-11-25T10:30:00Z","ttl":60000}`, + expectedResourceID: "status-notification-task-id", + expectedArguments: map[string]interface{}{ + "taskId": "status-notification-task-id", + "status": "completed", + "createdAt": "2025-11-25T10:30:00Z", + "ttl": float64(60000), + }, + }, + { + name: "notifications/tasks/status with numeric taskId", + method: "notifications/tasks/status", + params: `{"taskId":99999,"status":"running","createdAt":"2025-11-25T10:35:00Z"}`, + expectedResourceID: "99999", + expectedArguments: map[string]interface{}{ + "taskId": float64(99999), + "status": "running", + "createdAt": "2025-11-25T10:35:00Z", + }, + }, { name: "completion/complete with PromptReference", method: "completion/complete", @@ -623,6 +702,36 @@ func TestExtractResourceAndArguments(t *testing.T) { "reason": "User cancelled", }, }, + { + name: "tasks/get with missing taskId", + method: "tasks/get", + params: `{}`, + expectedResourceID: "", + expectedArguments: nil, + }, + { + name: "tasks/cancel with missing taskId", + method: "tasks/cancel", + params: `{}`, + expectedResourceID: "", + expectedArguments: nil, + }, + { + name: "tasks/result with missing taskId", + method: "tasks/result", + params: `{}`, + expectedResourceID: "", + expectedArguments: nil, + }, + { + name: "notifications/tasks/status with missing taskId", + method: "notifications/tasks/status", + params: `{"status":"completed"}`, + expectedResourceID: "", + expectedArguments: map[string]interface{}{ + "status": "completed", + }, + }, { name: "tools/list with empty cursor", method: "tools/list",