Add window.postMessage support.

backend/chainlit/__init__.py

@@ -169,6 +169,33 @@ async def with_parent_id(message: Message):
     return func
 
 
+@trace
+def on_window_message(func: Callable[[str], Any]) -> Callable:
+    """
+    Hook to react to javascript postMessage events coming from the UI.
+
+    Args:
+        func (Callable[[str], Any]): The function to be called when a window message is received.
+                                     Takes the message content as a string parameter.
+
+    Returns:
+        Callable[[str], Any]: The decorated on_window_message function.
+    """
+    config.code.on_window_message = wrap_user_function(func)
+    return func
+
+
+@trace
+def send_window_message(data: Any):
+    """
+    Send custom data to the host window via a window.postMessage event.
+
+    Args:
+        data (Any): The data to send with the event.
+    """
+    asyncio.create_task(context.emitter.send_window_message(data))
+
+
 @trace
 def on_chat_start(func: Callable) -> Callable:
     """
@@ -414,6 +441,8 @@ def acall(self):
     "CompletionGeneration",
     "GenerationMessage",
     "on_logout",
+    "on_window_message",
+    "send_window_message",
     "on_chat_start",
     "on_chat_end",
     "on_chat_resume",

backend/chainlit/config.py

@@ -281,9 +281,9 @@ class CodeSettings:
     on_chat_end: Optional[Callable[[], Any]] = None
     on_chat_resume: Optional[Callable[["ThreadDict"], Any]] = None
     on_message: Optional[Callable[["Message"], Any]] = None
+    on_window_message: Optional[Callable[[str], Any]] = None
     on_audio_chunk: Optional[Callable[["AudioChunk"], Any]] = None
     on_audio_end: Optional[Callable[[List["ElementBased"]], Any]] = None
-
     author_rename: Optional[Callable[[str], str]] = None
     on_settings_update: Optional[Callable[[Dict[str, Any]], Any]] = None
     set_chat_profiles: Optional[Callable[[Optional["User"]], List["ChatProfile"]]] = (

backend/chainlit/emitter.py

@@ -367,3 +367,7 @@ def send_action_response(
         return self.emit(
             "action_response", {"id": id, "status": status, "response": response}
         )
+
+    def send_window_message(self, data: Any):
+        """Send custom data to the host window."""
+        return self.emit("window_message", data)

backend/chainlit/socket.py

@@ -314,6 +314,13 @@ async def message(sid, payload: MessagePayload):
     session.current_task = task
 
 
+@sio.on("window_message")
+async def window_message(sid, data):
+    """Handle a message send by the host window."""
+    if config.code.on_window_message:
+        await config.code.on_window_message(data)
+
+
 @sio.on("audio_chunk")
 async def audio_chunk(sid, payload: AudioChunkPayload):
     """Handle an audio chunk sent by the user."""

frontend/src/AppWrapper.tsx

@@ -3,12 +3,18 @@ import { useEffect } from 'react';
 import { useTranslation } from 'react-i18next';
 import getRouterBasename from 'utils/router';
 
-import { useApi, useAuth, useConfig } from '@chainlit/react-client';
+import {
+  useApi,
+  useAuth,
+  useChatInteract,
+  useConfig
+} from '@chainlit/react-client';
 
 export default function AppWrapper() {
   const { isAuthenticated, isReady } = useAuth();
   const { language: languageInUse } = useConfig();
   const { i18n } = useTranslation();
+  const { windowMessage } = useChatInteract();
 
   function handleChangeLanguage(languageBundle: any): void {
     i18n.addResourceBundle(languageInUse, 'translation', languageBundle);
@@ -33,6 +39,14 @@ export default function AppWrapper() {
     handleChangeLanguage(translations.translation);
   }, [translations]);
 
+  useEffect(() => {
+    const handleWindowMessage = (event: MessageEvent) => {
+      windowMessage(event.data);
+    }
+    window.addEventListener('message', handleWindowMessage);
+    return () => window.removeEventListener('message', handleWindowMessage);
+  }, [windowMessage]);
+
   if (!isReady) {
     return null;
   }

libs/react-client/src/useChatInteract.ts

@@ -90,6 +90,13 @@ const useChatInteract = () => {
     [session?.socket]
   );
 
+  const windowMessage = useCallback(
+    (data: any) => {
+      session?.socket.emit('window_message', data);
+    },
+    [session?.socket]
+  );
+
   const sendAudioChunk = useCallback(
     (isStart: boolean, mimeType: string, elapsedTime: number, data: Blob) => {
       session?.socket.emit('audio_chunk', {
@@ -179,6 +186,7 @@ const useChatInteract = () => {
     replyMessage,
     sendMessage,
     editMessage,
+    windowMessage,
     sendAudioChunk,
     endAudioStream,
     stopTask,

libs/react-client/src/useChatSession.ts

@@ -268,6 +268,12 @@ const useChatSession = () => {
       socket.on('token_usage', (count: number) => {
         setTokenCount((old) => old + count);
       });
+
+      socket.on('window_message', (data: any) => {
+        if (window.parent && window.parent !== window) {
+          window.parent.postMessage(data, '*');
+        }
+      });
     },
     [setSession, sessionId, chatProfile]
   );