GH-45 Support WebSockets in annotated routing plugin (Resolves #45)

dzikoysk · dzikoysk · commit a0b505dbfaa4 · 2026-02-15T20:55:53.000+01:00
diff --git a/README.md b/README.md
@@ -124,6 +124,28 @@ static final class ExampleEndpoints {
         ctx.status(500).result("Something went wrong: " + e.getClass());
     }
 
+    // register websocket endpoints with @Ws annotation
+    // method must return WsHandler, keeping all WS event listeners grouped together
+    @Ws("/events")
+    WsHandler websocketEvents() {
+        return new WsHandler() {
+            @Override
+            public void onConnect(WsConnectContext ctx) {
+                System.out.println("WebSocket connected: " + ctx.sessionId());
+            }
+
+            @Override
+            public void onMessage(WsMessageContext ctx) {
+                ctx.send("Echo: " + ctx.message());
+            }
+
+            @Override
+            public void onClose(WsCloseContext ctx) {
+                System.out.println("WebSocket closed: " + ctx.sessionId());
+            }
+        };
+    }
+
 }
 
 public static void main(String[] args) {
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -114,7 +114,7 @@ subprojects {
     apply(plugin = "org.jetbrains.kotlin.kapt")
 
     dependencies {
-        val javalin = "7.0.0-beta.1"
+        val javalin = "7.0.0-beta.3"
         compileOnly("io.javalin:javalin:$javalin")
         testImplementation("io.javalin:javalin:$javalin")
         testImplementation("io.javalin:javalin-testtools:$javalin")
diff --git a/routing-annotations/routing-annotated-specification/src/main/kotlin/io/javalin/community/routing/annotations/RoutingAnnotations.kt b/routing-annotations/routing-annotated-specification/src/main/kotlin/io/javalin/community/routing/annotations/RoutingAnnotations.kt
@@ -2,6 +2,11 @@ package io.javalin.community.routing.annotations
 
 import io.javalin.event.JavalinLifecycleEvent
 import io.javalin.http.HttpStatus
+import io.javalin.websocket.WsBinaryMessageContext
+import io.javalin.websocket.WsCloseContext
+import io.javalin.websocket.WsConnectContext
+import io.javalin.websocket.WsErrorContext
+import io.javalin.websocket.WsMessageContext
 import kotlin.annotation.AnnotationRetention.RUNTIME
 import kotlin.annotation.AnnotationTarget.CLASS
 import kotlin.annotation.AnnotationTarget.FUNCTION
@@ -97,4 +102,16 @@ annotation class Form(val value: String = "")
 
 @Retention(RUNTIME)
 @Target(VALUE_PARAMETER)
-annotation class Cookie(val value: String = "")
+annotation class Cookie(val value: String = "")
+
+@Retention(RUNTIME)
+@Target(FUNCTION)
+annotation class Ws(val value: String = "")
+
+interface WsHandler {
+    fun onConnect(ctx: WsConnectContext) {}
+    fun onError(ctx: WsErrorContext) {}
+    fun onClose(ctx: WsCloseContext) {}
+    fun onMessage(ctx: WsMessageContext) {}
+    fun onBinaryMessage(ctx: WsBinaryMessageContext) {}
+}
diff --git a/routing-annotations/routing-annotated/src/main/kotlin/io/javalin/community/routing/annotations/AnnotatedRouting.kt b/routing-annotations/routing-annotated/src/main/kotlin/io/javalin/community/routing/annotations/AnnotatedRouting.kt
@@ -20,6 +20,7 @@ import io.javalin.event.JavalinLifecycleEvent.SERVER_STOP_FAILED
 import io.javalin.http.BadRequestResponse
 import io.javalin.http.Context
 import io.javalin.http.Handler
+import io.javalin.websocket.WsHandlerType
 
 fun interface HandlerResultConsumer<T> {
     fun handle(ctx: Context, value: T)
@@ -73,6 +74,7 @@ object AnnotatedRouting : RoutingApiInitializer<AnnotatedRoutingConfig> {
         val registeredEventListeners = mutableMapOf<JavalinLifecycleEvent, AnnotatedEvent>()
         val registeredRoutes = mutableListOf<AnnotatedRoute>()
         val registeredExceptionHandlers = mutableListOf<AnnotatedException>()
+        val registeredWsRoutes = mutableListOf<AnnotatedWsRoute>()
 
         configuration.registeredRoutes.forEach {
             val detectedEventListeners = loader.loadEventHandlers(it)
@@ -83,6 +85,9 @@ object AnnotatedRouting : RoutingApiInitializer<AnnotatedRoutingConfig> {
 
             val detectedExceptionHandlers = loader.loadExceptionHandlers(it)
             registeredExceptionHandlers.addAll(detectedExceptionHandlers)
+
+            val detectedWsRoutes = loader.loadWsHandlers(it)
+            registeredWsRoutes.addAll(detectedWsRoutes)
         }
 
         registeredEventListeners.forEach { (key, event) ->
@@ -118,6 +123,10 @@ object AnnotatedRouting : RoutingApiInitializer<AnnotatedRoutingConfig> {
                 annotatedException.handler.invoke(ctx, exception)
             }
         }
+
+        registeredWsRoutes.forEach { wsRoute ->
+            state.internalRouter.addWsHandler(WsHandlerType.WEBSOCKET, wsRoute.path, wsRoute.wsConfig)
+        }
     }
 
     private fun createVersionedRoute(apiVersionHeader: String, id: RouteIdentifier, routes: List<DslRoute<Context, Unit>>): Handler {
diff --git a/routing-annotations/routing-annotated/src/main/kotlin/io/javalin/community/routing/annotations/ReflectiveEndpointLoader.kt b/routing-annotations/routing-annotated/src/main/kotlin/io/javalin/community/routing/annotations/ReflectiveEndpointLoader.kt
@@ -12,18 +12,22 @@ import io.javalin.router.InternalRouter
 import io.javalin.util.Util.firstOrNull
 import io.javalin.validation.Validation
 import io.javalin.validation.Validator
+import io.javalin.websocket.WsConfig
 import java.lang.reflect.AnnotatedElement
 import java.lang.reflect.Method
 import java.lang.reflect.Parameter
 import java.lang.reflect.ParameterizedType
 import java.lang.reflect.Type
 import java.util.*
+import java.util.function.Consumer
 import kotlin.reflect.KClass
 
 typealias AnnotatedRoute = DefaultDslRoute<Context, Unit>
 typealias AnnotatedException = DefaultDslException<Context, Exception, Unit>
 typealias AnnotatedEvent = () -> Unit
 
+data class AnnotatedWsRoute(val path: String, val wsConfig: Consumer<WsConfig>)
+
 internal class ReflectiveEndpointLoader(
     private val internalRouter: InternalRouter,
     private val resultHandlers: Map<Class<*>, HandlerResultConsumer<*>>
@@ -164,6 +168,44 @@ internal class ReflectiveEndpointLoader(
         return dslEvents
     }
 
+    fun loadWsHandlers(endpoint: Any): List<AnnotatedWsRoute> {
+        val endpointClass = endpoint::class.java
+
+        val endpointPath = endpointClass.getAnnotation<Endpoints>()
+            ?.value
+            ?: ""
+
+        val wsRoutes = mutableListOf<AnnotatedWsRoute>()
+
+        getAllDeclaredMethods(endpointClass).forEach { method ->
+            val wsAnnotation = method.getAnnotation<Ws>() ?: return@forEach
+
+            require(method.trySetAccessible()) {
+                "Unable to access method $method in class $endpointClass"
+            }
+
+            require(WsHandler::class.java.isAssignableFrom(method.returnType)) {
+                "Method $method annotated with @Ws must return WsHandler"
+            }
+
+            val path = "/$endpointPath/${wsAnnotation.value}".replace(repeatedPathSeparatorRegex, "/").dropLastWhile { it == '/' }
+
+            wsRoutes.add(AnnotatedWsRoute(
+                path = path,
+                wsConfig = { wsConfig ->
+                    val handler = method.invoke(endpoint) as WsHandler
+                    wsConfig.onConnect { handler.onConnect(it) }
+                    wsConfig.onError { handler.onError(it) }
+                    wsConfig.onClose { handler.onClose(it) }
+                    wsConfig.onMessage { handler.onMessage(it) }
+                    wsConfig.onBinaryMessage { handler.onBinaryMessage(it) }
+                }
+            ))
+        }
+
+        return wsRoutes
+    }
+
     @Suppress("UNCHECKED_CAST")
     private fun findResultHandler(method: Method): HandlerResultConsumer<Any?> =
         resultHandlers.asSequence()
diff --git a/routing-annotations/routing-annotated/src/test/java/io/javalin/community/routing/annotations/AnnotatedRoutingTest.kt b/routing-annotations/routing-annotated/src/test/java/io/javalin/community/routing/annotations/AnnotatedRoutingTest.kt
@@ -11,6 +11,8 @@ import io.javalin.http.HttpStatus
 import io.javalin.router.Endpoint
 import io.javalin.testtools.HttpClient
 import io.javalin.testtools.JavalinTest
+import io.javalin.websocket.WsConnectContext
+import io.javalin.websocket.WsMessageContext
 import kong.unirest.Unirest
 import kong.unirest.Unirest.request
 import org.assertj.core.api.Assertions.assertThat
@@ -19,7 +21,14 @@ import org.junit.jupiter.api.Nested
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.assertDoesNotThrow
 import java.io.Closeable
+import java.net.URI
+import java.net.http.HttpClient as JdkHttpClient
+import java.net.http.WebSocket as JdkWebSocket
 import java.util.*
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionStage
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.TimeUnit
 
 class AnnotatedRoutingTest {
 
@@ -554,4 +563,79 @@ class AnnotatedRoutingTest {
 
     }
 
+    @Nested
+    inner class WebSockets {
+
+        @Test
+        fun `should register and handle websocket endpoint`() {
+            val log = ConcurrentLinkedQueue<String>()
+
+            val app = Javalin.create { cfg ->
+                cfg.routes(Annotated) {
+                    registerEndpoints(
+                        @Endpoints("/ws")
+                        object {
+                            @Ws("/echo")
+                            fun echo(): WsHandler = object : WsHandler {
+                                override fun onConnect(ctx: WsConnectContext) {
+                                    log.add("connected")
+                                }
+                                override fun onMessage(ctx: WsMessageContext) {
+                                    log.add("received: ${ctx.message()}")
+                                    ctx.send("echo: ${ctx.message()}")
+                                }
+                            }
+                        }
+                    )
+                }
+            }.start(0)
+
+            try {
+                val received = CompletableFuture<String>()
+
+                val wsClient = JdkHttpClient.newHttpClient().newWebSocketBuilder()
+                    .buildAsync(
+                        URI.create("ws://localhost:${app.port()}/ws/echo"),
+                        object : JdkWebSocket.Listener {
+                            override fun onText(webSocket: JdkWebSocket, data: CharSequence, last: Boolean): CompletionStage<*> {
+                                received.complete(data.toString())
+                                return CompletableFuture.completedFuture(null)
+                            }
+                        }
+                    )
+                    .get(2, TimeUnit.SECONDS)
+
+                wsClient.sendText("hello", true).get(2, TimeUnit.SECONDS)
+                val response = received.get(2, TimeUnit.SECONDS)
+
+                assertThat(response).isEqualTo("echo: hello")
+                assertThat(log).contains("connected", "received: hello")
+
+                wsClient.sendClose(JdkWebSocket.NORMAL_CLOSURE, "done").get(2, TimeUnit.SECONDS)
+            } finally {
+                app.stop()
+            }
+        }
+
+        @Test
+        fun `should throw if @Ws method does not return WsHandler`() {
+            assertThatThrownBy {
+                Javalin.create { cfg ->
+                    cfg.routes(Annotated) {
+                        registerEndpoints(
+                            @Endpoints
+                            object {
+                                @Ws("/bad")
+                                fun bad(): String = "not a WsHandler"
+                            }
+                        )
+                    }
+                }
+            }
+            .isInstanceOf(IllegalArgumentException::class.java)
+            .hasMessageContaining("must return WsHandler")
+        }
+
+    }
+
 }
diff --git a/routing-annotations/routing-annotated/src/test/java/io/javalin/community/routing/annotations/example/AnnotatedRoutingExample.java b/routing-annotations/routing-annotated/src/test/java/io/javalin/community/routing/annotations/example/AnnotatedRoutingExample.java
@@ -8,15 +8,16 @@
 import io.javalin.openapi.OpenApiParam;
 import io.javalin.openapi.OpenApiResponse;
 import io.javalin.router.Endpoint;
+import io.javalin.websocket.WsConnectContext;
+import io.javalin.websocket.WsCloseContext;
+import io.javalin.websocket.WsMessageContext;
 import kong.unirest.HttpResponse;
 import kong.unirest.Unirest;
 import org.jetbrains.annotations.Nullable;
 import java.io.Serializable;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Objects;
 
-import static io.javalin.community.routing.annotations.AnnotatedRouting.Annotated;
 import static io.javalin.http.Header.AUTHORIZATION;
 import static io.javalin.openapi.HttpMethod.GET;
 
@@ -102,6 +103,28 @@ void defaultExceptionHandler(Exception e, Context ctx) {
             ctx.status(500).result("Something went wrong: " + e.getClass());
         }
 
+        // register websocket endpoints with @Ws annotation
+        // method must return WsHandler, keeping all WS event listeners grouped together
+        @Ws("/events")
+        WsHandler websocketEvents() {
+            return new WsHandler() {
+                @Override
+                public void onConnect(WsConnectContext ctx) {
+                    System.out.println("WebSocket connected: " + ctx.sessionId());
+                }
+
+                @Override
+                public void onMessage(WsMessageContext ctx) {
+                    ctx.send("Echo: " + ctx.message());
+                }
+
+                @Override
+                public void onClose(WsCloseContext ctx) {
+                    System.out.println("WebSocket closed: " + ctx.sessionId());
+                }
+            };
+        }
+
     }
 
     public static void main(String[] args) {
diff --git a/routing-core/src/main/java/io/javalin/community/routing/JavalinConfigUnsafe.java b/routing-core/src/main/java/io/javalin/community/routing/JavalinConfigUnsafe.java
diff --git a/routing-core/src/main/kotlin/io/javalin/community/routing/JavalinRoutingExtensions.kt b/routing-core/src/main/kotlin/io/javalin/community/routing/JavalinRoutingExtensions.kt
@@ -98,7 +98,7 @@ fun <SETUP> JavalinConfig.routes(
     initializer: RoutingApiInitializer<SETUP>,
     setup: Consumer<SETUP>,
 ): JavalinConfig = also {
-    initializer.initialize(state = JavalinConfigUnsafe.getState(this)) {
+    initializer.initialize(state = this.unsafe) {
         setup.accept(this)
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ fun <SETUP> JavalinConfig.routes(`
`98`	`98`	`initializer: RoutingApiInitializer<SETUP>,`
`99`	`99`	`setup: Consumer<SETUP>,`
`100`	`100`	`): JavalinConfig = also {`
`101`		`- initializer.initialize(state = JavalinConfigUnsafe.getState(this)) {`
	`101`	`+ initializer.initialize(state = this.unsafe) {`
`102`	`102`	`setup.accept(this)`
`103`	`103`	`}`
`104`	`104`	`}`