Format the message to retrieve the user avatar color

packages/collaboration-extension/src/chat.ts

@@ -12,12 +12,14 @@ import {
 } from '@jupyterlab/application';
 import { ReactWidget, IThemeManager } from '@jupyterlab/apputils';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
-
 import {
   buildChatSidebar,
   buildErrorWidget,
-  ChatHandler
+  ChatHandler,
+  ChatService
 } from '@jupyterlab/chat';
+import { IGlobalAwareness } from '@jupyter/collaboration';
+import { Awareness } from 'y-protocols/awareness';
 
 /**
  * Initialization of the @jupyterlab/chat extension.
@@ -27,17 +29,18 @@ export const chat: JupyterFrontEndPlugin<void> = {
   description: 'A chat extension for Jupyterlab',
   autoStart: true,
   optional: [ILayoutRestorer, IThemeManager],
-  requires: [IRenderMimeRegistry],
+  requires: [IGlobalAwareness, IRenderMimeRegistry],
   activate: async (
     app: JupyterFrontEnd,
+    awareness: Awareness,
     rmRegistry: IRenderMimeRegistry,
     restorer: ILayoutRestorer | null,
     themeManager: IThemeManager | null
   ) => {
     /**
      * Initialize chat handler, open WS connection
      */
-    const chatHandler = new ChatHandler();
+    const chatHandler = new CollaborativeChatHandler({ awareness });
 
     let chatWidget: ReactWidget | null = null;
     try {
@@ -58,3 +61,48 @@ export const chat: JupyterFrontEndPlugin<void> = {
     console.log('Collaboration chat initialized');
   }
 };
+
+/**
+ * The collaborative chat handler.
+ */
+class CollaborativeChatHandler extends ChatHandler {
+  /**
+   * Create a new collaborative chat handler.
+   */
+  constructor(options: Private.IOptions) {
+    super();
+    this._awareness = options.awareness;
+  }
+
+  /**
+   * A function called before transferring the message to the panel(s).
+   * Can be useful if some actions are required on the message.
+   *
+   * It is used in this case to retrieve the user avatar color, unknown on server side.
+   */
+  protected formatChatMessage(
+    message: ChatService.IChatMessage
+  ): ChatService.IChatMessage {
+    const sender = Array.from(this._awareness.states.values()).find(
+      awareness => awareness.user.username === message.sender.username
+    )?.user;
+    if (sender) {
+      message.sender.color = sender.color;
+    }
+    return message;
+  }
+
+  private _awareness: Awareness;
+}
+
+/**
+ * The private namespace
+ */
+namespace Private {
+  /**
+   * Options for the collaborative chat handler.
+   */
+  export interface IOptions extends ChatHandler.IOptions {
+    awareness: Awareness;
+  }
+}