diff --git a/jupyter_collaboration/app.py b/jupyter_collaboration/app.py
index 801d016e..8d7cac3d 100644
--- a/jupyter_collaboration/app.py
+++ b/jupyter_collaboration/app.py
@@ -1,6 +1,8 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+from typing import Any, Dict
+
 from jupyter_server.extension.application import ExtensionApp
 from traitlets import Float, Int, Type
 from ypy_websocket.ystore import BaseYStore
@@ -55,7 +57,7 @@ def initialize_settings(self):
         )
 
     def initialize_handlers(self):
-        rooms = {}
+        rooms: Dict[str, Any] = {}
         self.handlers.extend(
             [
                 (r"/api/yjs/roomid/(.*)", YDocRoomIdHandler, {"rooms": rooms}),
diff --git a/jupyter_collaboration/handlers.py b/jupyter_collaboration/handlers.py
index a8131946..a3f63a81 100644
--- a/jupyter_collaboration/handlers.py
+++ b/jupyter_collaboration/handlers.py
@@ -147,7 +147,7 @@ async def __anext__(self):
             raise StopAsyncIteration()
         return message
 
-    def initialize(self, rooms: dict):
+    def initialize(self, rooms: Dict[str, Any]) -> None:
         self._rooms: Dict[str, Any] = rooms
 
         # Config
@@ -302,7 +302,7 @@ async def recv(self):
     def on_message(self, message):
         assert self.websocket_server is not None
         message_type = message[0]
-        if message_type == YMessageType.AWARENESS:
+        if self.room is not None and message_type == YMessageType.AWARENESS:
             # awareness
             skip = False
             changes = self.room.awareness.get_changes(message[1:])
@@ -405,7 +405,7 @@ def check_origin(self, origin):
 class YDocRoomIdHandler(APIHandler):
     auth_resource = "contents"
 
-    def initialize(self, rooms: dict):
+    def initialize(self, rooms: Dict[str, Any]) -> None:
         self._rooms: Dict[str, Any] = rooms
 
     @web.authenticated