diff --git a/include/main.nut b/include/main.nut
index 38e3ab8..4c23b70 100644
--- a/include/main.nut
+++ b/include/main.nut
@@ -86,7 +86,7 @@ addEventHandler("onWebsocketConnect", function(socket, url)
         print("[PyG2O] Successfully connected to " + url);
 });
 
-addEventHandler("onWebsocketDisconnect", function(socket, url)
+addEventHandler("onWebsocketClose", function(socket, url, message)
 {
     if (socket != _pyg2o_server_connection || _reconnect)
         return;