watch websocket errors

oml/websocket.py

@@ -49,7 +49,11 @@ class Handler(WebSocketHandler):
 
     #websocket calls
     def on_message(self, message):
+        try:
             action, data = json.loads(message)
+        except json.decoder.JSONDecodeError:
+            logger.debug('invalid websocket message: %s', message)
+            return
         if state.tasks:
             state.tasks.queue(action, data)
 
@@ -62,7 +66,17 @@ class Handler(WebSocketHandler):
         if self.ws_connection is None:
             self.on_close()
         else:
-            state.main.add_callback(lambda: self.write_message(message))
+            state.main.add_callback(lambda: self._write_message(message))
+
+    async def _write_message(self, message):
+        try:
+            task = self.write_message(message)
+            await task
+        except tornado.iostream.StreamClosedError as e:
+            self.on_close()
+        except tornado.websocket.WebSocketClosedError as e:
+            self.on_close()
+
 
 def trigger_event(event, data):
     #if len(state.websockets):