diff --git a/src/elli_websocket_handler.erl b/src/elli_websocket_handler.erl
index a51c060..6b4836c 100644
--- a/src/elli_websocket_handler.erl
+++ b/src/elli_websocket_handler.erl
@@ -8,9 +8,9 @@
 %% Licensed 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.
@@ -24,16 +24,16 @@
 -callback websocket_init(Req :: elli:req(), Args :: any()) ->
     {ok, Headers :: elli:headers(), State :: any()} |
     {ok, Headers :: elli:headers(), hibernate, State :: any()} |
-    {ok, Headers :: elli:headers(), Timeout :: non_neg_integer(), State :: any()} | 
-    {ok, Headers :: elli:headers(), hibernate, Timeout :: non_neg_integer(), State :: any()} | 
+    {ok, Headers :: elli:headers(), Timeout :: non_neg_integer(), State :: any()} |
+    {ok, Headers :: elli:headers(), hibernate, Timeout :: non_neg_integer(), State :: any()} |
     {shutdown, Headers :: elli:headers()}.
 
 
 -callback websocket_handle(Req :: elli:req(), Message :: elli_websocket:message(), State :: any()) ->
-    {reply, ReplyMessage :: elli_websocket:message(), State}.
+    {reply, ReplyMessage :: elli_websocket:message(), State :: any()}.
 
 -callback websocket_info(Req :: elli:req(), Message :: any(), State :: any()) ->
     {ok, ReplyState :: any()}.
 
--callback websocket_handle_event(Event :: elli_websocket:event(), Args :: list(), State :: any()) -> 
+-callback websocket_handle_event(Event :: elli_websocket:event(), Args :: list(), State :: any()) ->
     ok.