diff --git a/src/kemal/dsl.cr b/src/kemal/dsl.cr
index 8b42b6a..6b5d589 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, HTTP::Context -> Void)
+def ws(path, &block : HTTP::WebSocket, HTTP::Server::Context -> Void)
   Kemal::WebSocketHandler.new path, &block
 end