fix typing issues raised by mypy

packages/jupyter-ai/jupyter_ai/chat_handlers/base.py

@@ -9,6 +9,7 @@
     Awaitable,
     ClassVar,
     Dict,
+    get_args as get_type_args,
     List,
     Literal,
     Optional,
@@ -265,8 +266,8 @@ async def _default_handle_exc(self, e: Exception, message: HumanChatMessage):
     def broadcast_message(self, message: Message):
         """
         Broadcasts a message to all WebSocket connections. If there are no
-        WebSocket connections, this method directly appends to
-        `self.chat_history`.
+        WebSocket connections and the message is a chat message, this method
+        directly appends to `self.chat_history`.
         """
         broadcast = False
         for websocket in self._root_chat_handlers.values():
@@ -278,7 +279,9 @@ def broadcast_message(self, message: Message):
             break
 
         if not broadcast:
-            self._chat_history.append(message)
+            if isinstance(message, get_type_args(ChatMessage)):
+                cast(ChatMessage, message)
+                self._chat_history.append(message)
 
     def reply(self, response: str, human_msg: Optional[HumanChatMessage] = None):
         """

packages/jupyter-ai/jupyter_ai/handlers.py

@@ -121,17 +121,17 @@ def loop(self) -> AbstractEventLoop:
     def pending_messages(self) -> List[PendingMessage]:
         return self.settings["pending_messages"]
 
+    @pending_messages.setter
+    def pending_messages(self, new_pending_messages):
+        self.settings["pending_messages"] = new_pending_messages
+
     @property
     def cleared_message_ids(self) -> Set[str]:
         """Set of `HumanChatMessage.id` that were cleared via `ClearRequest`."""
         if "cleared_message_ids" not in self.settings:
             self.settings["cleared_message_ids"] = set()
         return self.settings["cleared_message_ids"]
 
-    @pending_messages.setter
-    def pending_messages(self, new_pending_messages):
-        self.settings["pending_messages"] = new_pending_messages
-
     def initialize(self):
         self.log.debug("Initializing websocket connection %s", self.request.path)