Add support for custom remote function mapping via Annotation

ballerina/Ballerina.toml

@@ -1,7 +1,7 @@
 [package]
 org = "ballerina"
 name = "websocket"
-version = "2.14.1"
+version = "2.15.0"
 authors = ["Ballerina"]
 keywords = ["ws", "network", "bi-directional", "streaming", "service", "client"]
 repository = "https://github.com/ballerina-platform/module-ballerina-websocket"
@@ -15,8 +15,8 @@ graalvmCompatible = true
 [[platform.java21.dependency]]
 groupId = "io.ballerina.stdlib"
 artifactId = "websocket-native"
-version = "2.14.1"
-path = "../native/build/libs/websocket-native-2.14.1-SNAPSHOT.jar"
+version = "2.15.0"
+path = "../native/build/libs/websocket-native-2.15.0-SNAPSHOT.jar"
 
 [[platform.java21.dependency]]
 groupId = "io.ballerina.stdlib"
@@ -85,5 +85,5 @@ version = "4.1.118.Final"
 path = "./lib/netty-handler-proxy-4.1.118.Final.jar"
 
 [[platform.java21.dependency]]
-path = "../test-utils/build/libs/websocket-test-utils-2.14.1-SNAPSHOT.jar"
+path = "../test-utils/build/libs/websocket-test-utils-2.15.0-SNAPSHOT.jar"
 scope = "testOnly"

ballerina/CompilerPlugin.toml

@@ -3,4 +3,10 @@ id = "websocket-compiler-plugin"
 class = "io.ballerina.stdlib.websocket.plugin.WebSocketCompilerPlugin"
 
 [[dependency]]
-path = "../compiler-plugin/build/libs/websocket-compiler-plugin-2.14.1-SNAPSHOT.jar"
+path = "../compiler-plugin/build/libs/websocket-compiler-plugin-2.15.0-SNAPSHOT.jar"
+
+[[dependency]]
+path = "../native/build/libs/websocket-native-2.15.0-SNAPSHOT.jar"
+
+[[dependency]]
+path = "./lib/http-native-2.14.0.jar"

ballerina/Dependencies.toml

@@ -342,7 +342,7 @@ dependencies = [
 [[package]]
 org = "ballerina"
 name = "websocket"
-version = "2.14.1"
+version = "2.15.0"
 dependencies = [
 	{org = "ballerina", name = "auth"},
 	{org = "ballerina", name = "constraint"},

ballerina/annotation.bal

@@ -14,10 +14,6 @@
 // specific language governing permissions and limitations
 // under the License.
 
-///////////////////////////
-/// Service Annotations ///
-///////////////////////////
-
 # Configurations for a WebSocket service.
 #
 # + subProtocols - Negotiable sub protocol by the service
@@ -42,3 +38,13 @@ public type WSServiceConfig record {|
 
 # The annotation which is used to configure a WebSocket service.
 public annotation WSServiceConfig ServiceConfig on service;
+
+# Configurations used to define dispatching rules for remote functions.
+#
+# + value - The value which is going to be used for dispatching to custom remote functions.
+public type WsDispatcherMapping record {|
+    string value;
+|};
+
+# The annotation which is used to configure the dispatching rules for WebSocket remote functions.
+public const annotation WsDispatcherMapping DispatcherMapping on function;

ballerina/build.gradle

@@ -122,6 +122,7 @@ task updateTomlFiles {
         ballerinaTomlFile.text = newConfig
 
         def newPluginConfig = compilerPluginTomlFilePlaceHolder.text.replace("@project.version@", project.version)
+        newPluginConfig = newPluginConfig.replace("@stdlib.httpnative.version@", stdlibDependentHttpNativeVersion)
         compilerPluginTomlFile.text = newPluginConfig
     }
 }

ballerina/tests/test_dispatcher_mapping_annotation.bal

@@ -0,0 +1,66 @@
+//  Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+//
+//  WSO2 LLC. licenses this file to you under the Apache License,
+//  Version 2.0 (the "License"); you may not use this file except
+//  in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing,
+//  software distributed under the License is distributed on an
+//  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+//  KIND, either express or implied. See the License for the
+//  specific language governing permissions and limitations
+//  under the License.
+
+import ballerina/test;
+
+type Subscribe record {|
+    string event = "subscribe";
+    string data;
+|};
+
+@ServiceConfig {
+    dispatcherKey: "event"
+}
+service / on new Listener(22103) {
+    resource function get .() returns Service|UpgradeError {
+        return new WsService22103();
+    }
+}
+
+service class WsService22103 {
+    *Service;
+
+    @DispatcherMapping {
+        value: "subscribe"
+    }
+    remote function onSubscribeMessage(Subscribe message) returns string {
+        return "onSubscribeMessage";
+    }
+
+    remote function onSubscribeMessageError(Caller caller, error message) returns error? {
+        check caller->writeMessage("onSubscribeMessageError");
+    }
+}
+
+@test:Config {
+    groups: ["dispatcherMappingAnnotation"]
+}
+public function testDispatcherMappingAnnotation() returns error? {
+    Client wsClient = check new ("ws://localhost:22103/");
+    check wsClient->writeMessage({event: "subscribe", data: "test"});
+    string res = check wsClient->readMessage();
+    test:assertEquals(res, "onSubscribeMessage");
+}
+
+@test:Config {
+    groups: ["dispatcherMappingAnnotation"]
+}
+public function testDispatcherMappingAnnotationWithCustomOnError() returns error? {
+    Client wsClient = check new ("ws://localhost:22103/");
+    check wsClient->writeMessage({event: "subscribe", invalidField: "test"});
+    string res = check wsClient->readMessage();
+    test:assertEquals(res, "onSubscribeMessageError");
+}

build-config/resources/CompilerPlugin.toml

@@ -4,3 +4,9 @@ class = "io.ballerina.stdlib.websocket.plugin.WebSocketCompilerPlugin"
 
 [[dependency]]
 path = "../compiler-plugin/build/libs/websocket-compiler-plugin-@project.version@.jar"
+
+[[dependency]]
+path = "../native/build/libs/websocket-native-@project.version@.jar"
+
+[[dependency]]
+path = "./lib/http-native-@stdlib.httpnative.version@.jar"

changelog.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Added
 
+- [Support Custom Remote Function Mapping via Annotation](https://github.com/ballerina-platform/ballerina-library/issues/7733)
 - [Implement websocket close frame support](https://github.com/ballerina-platform/ballerina-library/issues/7578)
 
 ### Fixed

compiler-plugin-tests/src/test/java/io/ballerina/stdlib/websocket/compiler/WebSocketServiceValidationTest.java

@@ -572,6 +572,20 @@ public void testRemoteFunctionWithStreamAndCloseFrameReturnTypes() {
         Assert.assertEquals(diagnosticResult.errorCount(), 0);
     }
 
+    @Test
+    public void testDispatcherMappingAnnotation() {
+        Package currentPackage = loadPackage("sample_package_63");
+        PackageCompilation compilation = currentPackage.getCompilation();
+        DiagnosticResult diagnosticResult = compilation.diagnosticResult();
+        Assert.assertEquals(diagnosticResult.errorCount(), 3);
+        Diagnostic firstDiagnostic = (Diagnostic) diagnosticResult.errors().toArray()[0];
+        assertDiagnostic(firstDiagnostic, PluginConstants.CompilationErrors.RE_DECLARED_REMOTE_FUNCTIONS);
+        Diagnostic secondDiagnostic = (Diagnostic) diagnosticResult.errors().toArray()[1];
+        assertDiagnostic(secondDiagnostic, PluginConstants.CompilationErrors.DUPLICATED_DISPATCHER_MAPPING_VALUE);
+        Diagnostic thirdDiagnostic = (Diagnostic) diagnosticResult.errors().toArray()[2];
+        assertDiagnostic(thirdDiagnostic, PluginConstants.CompilationErrors.INVALID_FUNCTION_ANNOTATION);
+    }
+
     private void assertDiagnostic(Diagnostic diagnostic, PluginConstants.CompilationErrors error) {
         Assert.assertEquals(diagnostic.diagnosticInfo().code(), error.getErrorCode());
         Assert.assertEquals(diagnostic.diagnosticInfo().messageFormat(),

compiler-plugin-tests/src/test/resources/ballerina_sources/sample_package_63/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "websocket_test"
+name = "sample_63"
+version = "0.1.0"

compiler-plugin-tests/src/test/resources/ballerina_sources/sample_package_63/server.bal

@@ -0,0 +1,60 @@
+//  Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+//
+//  WSO2 LLC. licenses this file to you under the Apache License,
+//  Version 2.0 (the "License"); you may not use this file except
+//  in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing,
+//  software distributed under the License is distributed on an
+//  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+//  KIND, either express or implied. See the License for the
+//  specific language governing permissions and limitations
+//  under the License.
+
+import ballerina/websocket as ws;
+
+type Subscribe record {|
+    string event = "subscribe";
+    string data;
+|};
+
+@ws:ServiceConfig {
+    dispatcherKey: "event"
+}
+service / on new ws:Listener(9090) {
+    resource function get .() returns ws:Service|ws:UpgradeError {
+        return new WsService();
+    }
+}
+
+service class WsService {
+    *ws:Service;
+
+    remote function onSubscribe(Subscribe message) returns string {
+        return "onSubscribe";
+    }
+
+    @ws:DispatcherMapping {
+        value: "subscribe"
+    }
+    remote function onSubscribeMessage(Subscribe message) returns string {
+        return "onSubscribeMessage";
+    }
+
+    @ws:DispatcherMapping {
+        value: "subscribe"
+    }
+    remote function onSubscribeText(Subscribe message) returns string {
+        return "onSubscribeText";
+    }
+
+    @ws:DispatcherMapping {
+        value: "ping"
+    }
+    remote function onPing(Subscribe message) returns string {
+        return "onPing";
+    }
+}

compiler-plugin/src/main/java/io/ballerina/stdlib/websocket/plugin/PluginConstants.java

@@ -79,6 +79,11 @@ public enum CompilationErrors {
                 "WEBSOCKET_215"),
         INVALID_REMOTE_FUNCTIONS("Cannot have `{0}` with `onMessage` remote function",
                 "WEBSOCKET_216"),
+        RE_DECLARED_REMOTE_FUNCTIONS("Cannot have `{0}` because the message type `{1}` is already " +
+                "associated with `{2}` remote function", "WEBSOCKET_217"),
+        DUPLICATED_DISPATCHER_MAPPING_VALUE("DispatcherMapping annotation value `{0}` is already " +
+                "exists", "WEBSOCKET_218"),
+        INVALID_FUNCTION_ANNOTATION("Invalid annotation provided for `{0}` remote function", "WEBSOCKET_219"),
         INVALID_RESOURCE_ERROR("There should be only one `get` resource for the service",
                 "WEBSOCKET_101"),
         MORE_THAN_ONE_RESOURCE_PARAM_ERROR("There should be only http:Request as a parameter",