diff --git a/src/kemal/dsl.cr b/src/kemal/dsl.cr
index c18dbeb..8b42b6a 100644
--- a/src/kemal/dsl.cr
+++ b/src/kemal/dsl.cr
@@ -6,6 +6,6 @@ HTTP_METHODS = %w(get post put patch delete options)
   end
 {% end %}
 
-def ws(path, &block : HTTP::WebSocket -> _)
+def ws(path, &block : HTTP::WebSocket, HTTP::Context -> Void)
   Kemal::WebSocketHandler.new path, &block
 end
diff --git a/src/kemal/websocket_handler.cr b/src/kemal/websocket_handler.cr
index d5ac49e..cf13f91 100644
--- a/src/kemal/websocket_handler.cr
+++ b/src/kemal/websocket_handler.cr
@@ -1,7 +1,7 @@
 # Kemal::WebSocketHandler is used for each define WebSocket route.
 # For each WebSocket route a new handler is created and registered to global handlers.
 class Kemal::WebSocketHandler < HTTP::WebSocketHandler
-  def initialize(@path, &@proc : HTTP::WebSocket ->)
+  def initialize(@path, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
     Kemal.config.add_ws_handler self
   end