@jupyter/chat 0.23.0-alpha.2 → 0.23.0-alpha.3

package/lib/widgets/multichat-panel.d.ts

@@ -1,3 +1,5 @@
+import { ToolbarRegistry } from '@jupyterlab/apputils';
+import { IObservableList } from '@jupyterlab/observables';
 import { TranslationBundle } from '@jupyterlab/translation';
 import { PanelWithToolbar, ReactiveToolbar } from '@jupyterlab/ui-components';
 import { Message } from '@lumino/messaging';
@@ -6,7 +8,8 @@ import { Panel, Widget } from '@lumino/widgets';
 import { ChatWidget } from './chat-widget';
 import { Chat, IInputToolbarRegistryFactory } from '../components';
 import { IChatModel } from '../model';
-import { IChatPlaceholderFactory } from '../tokens';
+import { IChatPanel, IChatPlaceholderFactory } from '../tokens';
+import { ChatArea } from '../types';
 /**
  * A panel widget with a reactive toolbar.
  */
@@ -28,9 +31,9 @@ export declare class MultiChatPanel extends PanelWithToolbar {
      */
     get current(): SidePanelWidget | undefined;
     /**
-     * A signal emitting when a chat widget is opened in the panel.
+     * A signal emitting when a chat panel is opened in the sidepanel.
      */
-    get chatOpened(): ISignal<MultiChatPanel, ChatWidget>;
+    get chatOpened(): ISignal<MultiChatPanel, IChatPanel>;
     /**
      * A signal emitting when the panel visibility changed.
      */
@@ -106,10 +109,10 @@ export declare class MultiChatPanel extends PanelWithToolbar {
     private _chatNamesChanged;
     private _chatOptions;
     private _inputToolbarFactory?;
+    private _chatToolbarFactory?;
     private _updateChatListDebouncer;
     private _createModel?;
     private _getChatNames?;
-    private _openInMain?;
     private _renameChat?;
     private _placeholderFactory?;
     private _openChatWidget?;
@@ -132,6 +135,10 @@ export declare namespace MultiChatPanel {
          * The input toolbar factory;
          */
         inputToolbarFactory?: IInputToolbarRegistryFactory;
+        /**
+         * An optional toolbar factory for each opened chat.
+         */
+        chatToolbarFactory?: (panel: IChatPanel) => IObservableList<ToolbarRegistry.IToolbarItem>;
         /**
          * An optional callback to create a chat model.
          *
@@ -147,12 +154,6 @@ export declare namespace MultiChatPanel {
         getChatNames?: () => Promise<{
             [name: string]: string;
         }>;
-        /**
-         * An optional callback to open the chat in the main area.
-         *
-         * @param name - the name of the chat to move.
-         */
-        openInMain?: (name: string) => Promise<boolean>;
         /**
          * An optional callback to rename a chat.
          *
@@ -188,10 +189,14 @@ export declare namespace MultiChatPanel {
 /**
  * A widget containing the chat and its toolbar.
  */
-declare class SidePanelWidget extends ReactivePanelWithToolbar {
+declare class SidePanelWidget extends ReactivePanelWithToolbar implements IChatPanel {
     constructor(options: SidePanelWidget.IOptions);
     protected onAfterAttach(msg: Message): void;
     protected onResize(msg: Widget.ResizeMessage): void;
+    /**
+     * The area of the widget.
+     */
+    get area(): ChatArea;
     /**
      * The chat widget embedded in the sidepanel widget.
      */
@@ -220,16 +225,11 @@ declare class SidePanelWidget extends ReactivePanelWithToolbar {
      * Update the title based on the chat name.
      */
     private _updateTitle;
-    /**
-     * Enable/disable unread icon.
-     */
-    private _unreadChanged;
     /**
      * Trigger reactive toolbar overflow computation from rendered toolbar size.
      */
     private _updateReactiveToolbar;
     private _chatWidget;
-    private _markAsRead;
     private _displayName;
     private _titleWidget;
     private _nameChanged;
@@ -254,14 +254,14 @@ declare namespace SidePanelWidget {
          * The displayed name of the chat.
          */
         displayName?: string;
-        /**
-         * The callback to open the chat in main area.
-         */
-        openInMain?: (name: string) => Promise<boolean>;
         /**
          * The callback to rename the chat.
          */
         renameChat?: boolean | ((oldName: string) => Promise<string | null>);
+        /**
+         * An optional toolbar factory.
+         */
+        toolbarFactory?: (panel: IChatPanel) => IObservableList<ToolbarRegistry.IToolbarItem>;
         /**
          * The translation bundle.
          */

package/lib/widgets/multichat-panel.js

@@ -8,7 +8,7 @@
  */
 import { InputDialog } from '@jupyterlab/apputils';
 import { nullTranslator } from '@jupyterlab/translation';
-import { addIcon, closeIcon, launchIcon, PanelWithToolbar, ReactiveToolbar, ReactWidget, Spinner, Toolbar, ToolbarButton } from '@jupyterlab/ui-components';
+import { addIcon, closeIcon, PanelWithToolbar, ReactiveToolbar, ReactWidget, Spinner, Toolbar, ToolbarButton } from '@jupyterlab/ui-components';
 import { ArrayExt } from '@lumino/algorithm';
 import { MessageLoop } from '@lumino/messaging';
 import { Debouncer } from '@lumino/polling';
@@ -19,7 +19,7 @@ import { ChatSelectorPopup } from './chat-selector-popup';
 import { ChatWidget } from './chat-widget';
 import { defaultPlaceholder } from './placeholder';
 import { TRANSLATION_DOMAIN } from '../context';
-import { chatIcon, readIcon } from '../icons';
+import { chatIcon } from '../icons';
 const SIDEPANEL_CLASS = 'jp-chat-sidepanel';
 const ADD_BUTTON_CLASS = 'jp-chat-add';
 const OPEN_SELECT_CLASS = 'jp-chat-open';
@@ -120,9 +120,9 @@ export class MultiChatPanel extends PanelWithToolbar {
         this.addClass(SIDEPANEL_CLASS);
         this._chatOptions = options;
         this._inputToolbarFactory = options.inputToolbarFactory;
+        this._chatToolbarFactory = options.chatToolbarFactory;
         this._getChatNames = options.getChatNames;
         this._createModel = options.createModel;
-        this._openInMain = options.openInMain;
         this._renameChat = options.renameChat;
         this._placeholderFactory = options.placeholderFactory;
         if (this._createModel) {
@@ -168,7 +168,7 @@ export class MultiChatPanel extends PanelWithToolbar {
             : undefined;
     }
     /**
-     * A signal emitting when a chat widget is opened in the panel.
+     * A signal emitting when a chat panel is opened in the sidepanel.
      */
     get chatOpened() {
         return this._chatOpened;
@@ -298,8 +298,8 @@ export class MultiChatPanel extends PanelWithToolbar {
         const widget = new SidePanelWidget({
             widget: chatWidget,
             displayName: name,
-            openInMain: this._openInMain,
             renameChat: this._renameChat,
+            toolbarFactory: this._chatToolbarFactory,
             onClose: (name, disposeModel = true) => {
                 this.unsetLoadedModel(name, disposeModel);
             },
@@ -314,7 +314,7 @@ export class MultiChatPanel extends PanelWithToolbar {
         if (this._chatSelectorPopup) {
             this._chatSelectorPopup.setCurrentChat(name);
         }
-        this._chatOpened.emit(chatWidget);
+        this._chatOpened.emit(widget);
         return chatWidget;
     }
     /**
@@ -372,14 +372,8 @@ export class MultiChatPanel extends PanelWithToolbar {
  */
 class SidePanelWidget extends ReactivePanelWithToolbar {
     constructor(options) {
-        var _a, _b, _c, _d;
+        var _a;
         super();
-        /**
-         * Enable/disable unread icon.
-         */
-        this._unreadChanged = (_, unread) => {
-            this._markAsRead.enabled = unread.length > 0;
-        };
         this._nameChanged = new Signal(this);
         this._chatWidget = options.widget;
         this._displayName = (_a = options.displayName) !== null && _a !== void 0 ? _a : options.widget.model.name;
@@ -399,18 +393,6 @@ class SidePanelWidget extends ReactivePanelWithToolbar {
         });
         // Add the chat widget
         this.addWidget(this._chatWidget);
-        // Add toolbar buttons
-        this._markAsRead = new ToolbarButton({
-            icon: readIcon,
-            iconLabel: trans.__('Mark chat as read'),
-            className: 'jp-mod-styled',
-            onClick: () => {
-                if (this.model) {
-                    this.model.unreadMessages = [];
-                }
-            }
-        });
-        this.toolbar.addItem('markRead', this._markAsRead);
         if (options.renameChat) {
             const renameButton = new ToolbarButton({
                 iconClass: 'jp-EditIcon',
@@ -447,21 +429,6 @@ class SidePanelWidget extends ReactivePanelWithToolbar {
             });
             this.toolbar.addItem('rename', renameButton);
         }
-        if (options.openInMain) {
-            const moveToMain = new ToolbarButton({
-                icon: launchIcon,
-                iconLabel: trans.__('Move the chat to the main area'),
-                className: 'jp-mod-styled',
-                onClick: async () => {
-                    var _a;
-                    const name = this.model.name;
-                    if (await ((_a = options.openInMain) === null || _a === void 0 ? void 0 : _a.call(options, name))) {
-                        options.onClose(this._displayName, false);
-                    }
-                }
-            });
-            this.toolbar.addItem('moveMain', moveToMain);
-        }
         const closeButton = new ToolbarButton({
             icon: closeIcon,
             iconLabel: trans.__('Close the chat'),
@@ -471,9 +438,25 @@ class SidePanelWidget extends ReactivePanelWithToolbar {
             }
         });
         this.toolbar.addItem('close', closeButton);
-        // Update mark as read button state
-        (_b = this.model.unreadChanged) === null || _b === void 0 ? void 0 : _b.connect(this._unreadChanged);
-        this._markAsRead.enabled = ((_d = (_c = this.model) === null || _c === void 0 ? void 0 : _c.unreadMessages.length) !== null && _d !== void 0 ? _d : 0) > 0;
+        if (options.toolbarFactory) {
+            const items = options.toolbarFactory(this);
+            for (let i = 0; i < items.length; i++) {
+                const { name, widget } = items.get(i);
+                this.toolbar.insertBefore('close', name, widget);
+            }
+            items.changed.connect((_, change) => {
+                if (change.type === 'add') {
+                    for (const { name, widget } of change.newValues) {
+                        this.toolbar.insertBefore('close', name, widget);
+                    }
+                }
+                else if (change.type === 'remove') {
+                    for (const { widget } of change.oldValues) {
+                        widget.dispose();
+                    }
+                }
+            });
+        }
     }
     onAfterAttach(msg) {
         super.onAfterAttach(msg);
@@ -483,6 +466,12 @@ class SidePanelWidget extends ReactivePanelWithToolbar {
         super.onResize(msg);
         this._updateReactiveToolbar();
     }
+    /**
+     * The area of the widget.
+     */
+    get area() {
+        return 'sidebar';
+    }
     /**
      * The chat widget embedded in the sidepanel widget.
      */
@@ -520,11 +509,6 @@ class SidePanelWidget extends ReactivePanelWithToolbar {
      * Dispose of the resources held by the widget.
      */
     dispose() {
-        var _a;
-        const model = this.model;
-        if (model) {
-            (_a = model.unreadChanged) === null || _a === void 0 ? void 0 : _a.disconnect(this._unreadChanged);
-        }
         super.dispose();
     }
     /**

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jupyter/chat",
-  "version": "0.23.0-alpha.2",
+  "version": "0.23.0-alpha.3",
   "description": "A package that provides UI components that can be used to create a chat in a Jupyterlab extension.",
   "keywords": [
     "jupyter",
@@ -55,6 +55,7 @@
     "@jupyterlab/filebrowser": "^4.6.0",
     "@jupyterlab/fileeditor": "^4.6.0",
     "@jupyterlab/notebook": "^4.6.0",
+    "@jupyterlab/observables": "^5.6.0",
     "@jupyterlab/rendermime": "^4.6.0",
     "@jupyterlab/translation": "^4.6.0",
     "@jupyterlab/ui-components": "^4.6.0",

package/src/__tests__/input-model.spec.ts

@@ -0,0 +1,129 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import { InputModel } from '../input-model';
+import { INewMessage } from '../types';
+
+describe('test input model', () => {
+  describe('metadata', () => {
+    it('should start with empty metadata', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      expect(model.getMetadata()).toEqual({});
+    });
+
+    it('should seed metadata from options', () => {
+      const model = new InputModel({
+        onSend: jest.fn(),
+        metadata: { persona: 'kiro' }
+      });
+      expect(model.getMetadata()).toEqual({ persona: 'kiro' });
+    });
+
+    it('should merge patches with updateMetadata', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      model.updateMetadata({ persona: 'kiro' });
+      model.updateMetadata({ model: { id: 'claude-opus-48' } });
+      expect(model.getMetadata()).toEqual({
+        persona: 'kiro',
+        model: { id: 'claude-opus-48' }
+      });
+    });
+
+    it('should overwrite existing keys on update', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      model.updateMetadata({ persona: 'kiro' });
+      model.updateMetadata({ persona: 'jupyternaut' });
+      expect(model.getMetadata()).toEqual({ persona: 'jupyternaut' });
+    });
+
+    it('should shallow-merge: a top-level key replaces the whole value', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      model.updateMetadata({ model: { id: 'a' } });
+      // Passing `model` again replaces it wholesale (no recursive merge).
+      model.updateMetadata({ model: { id: 'b' } });
+      expect(model.getMetadata()).toEqual({ model: { id: 'b' } });
+    });
+
+    it('should not be mutated by later changes to a patch', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      const patch = { model: { id: 'a' } };
+      model.updateMetadata(patch);
+      // Mutating the patch after the fact must not reach into stored metadata.
+      patch.model.id = 'tampered';
+      expect(model.getMetadata()).toEqual({ model: { id: 'a' } });
+    });
+
+    it('should clear metadata', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      model.updateMetadata({ persona: 'kiro' });
+      model.clearMetadata();
+      expect(model.getMetadata()).toEqual({});
+    });
+
+    it('should emit metadataChanged on update and clear', () => {
+      const model = new InputModel({ onSend: jest.fn() });
+      const emitted: any[] = [];
+      model.metadataChanged?.connect((_, metadata) => {
+        emitted.push({ ...metadata });
+      });
+      model.updateMetadata({ persona: 'kiro' });
+      model.clearMetadata();
+      expect(emitted).toEqual([{ persona: 'kiro' }, {}]);
+    });
+  });
+
+  describe('send', () => {
+    it('should attach metadata to the message when non-empty', () => {
+      const onSend = jest.fn();
+      const model = new InputModel({ onSend });
+      model.updateMetadata({ persona: 'kiro' });
+      model.send('hello');
+
+      const message: INewMessage = onSend.mock.calls[0][0];
+      expect(message.body).toBe('hello');
+      expect(message.metadata).toEqual({ persona: 'kiro' });
+    });
+
+    it('should omit metadata from the message when empty', () => {
+      const onSend = jest.fn();
+      const model = new InputModel({ onSend });
+      model.send('hello');
+
+      const message: INewMessage = onSend.mock.calls[0][0];
+      expect(message.metadata).toBeUndefined();
+    });
+
+    it('should send a copy of the metadata', () => {
+      const onSend = jest.fn();
+      const model = new InputModel({ onSend });
+      model.updateMetadata({ persona: 'kiro' });
+      model.send('hello');
+
+      const message: INewMessage = onSend.mock.calls[0][0];
+      model.updateMetadata({ persona: 'jupyternaut' });
+      expect(message.metadata).toEqual({ persona: 'kiro' });
+    });
+
+    it('should keep metadata after sending (sticky selection)', () => {
+      // Unlike attachments/mentions, metadata carries the picker's
+      // persona/model/settings selection, which is sticky across messages.
+      const model = new InputModel({ onSend: jest.fn() });
+      model.updateMetadata({ persona: 'kiro' });
+      model.send('hello');
+      expect(model.getMetadata()).toEqual({ persona: 'kiro' });
+    });
+  });
+});
+
+// `IMessageMetadata` is intentionally empty in the source; consumers augment it
+// with their own fields via module augmentation. We do the same here purely so
+// the tests can exercise `updateMetadata` with representative fields — this
+// stays in the test file and no consumer-specific fields leak into the source.
+declare module '../types' {
+  interface IMessageMetadata {
+    persona?: string;
+    model?: { id: string };
+  }
+}

package/src/__tests__/multichat-panel.spec.ts

@@ -3,13 +3,16 @@
  * Distributed under the terms of the Modified BSD License.
  */
 
-import { ToolbarButton } from '@jupyterlab/ui-components';
+import { ToolbarRegistry } from '@jupyterlab/apputils';
 import { RenderMimeRegistry } from '@jupyterlab/rendermime';
+import { ObservableList } from '@jupyterlab/observables';
+import { ToolbarButton, PanelWithToolbar } from '@jupyterlab/ui-components';
 import { Widget } from '@lumino/widgets';
 
 import { IChatPlaceholderFactory } from '../tokens';
 import { MultiChatPanel } from '../widgets/multichat-panel';
 import { defaultPlaceholder, Placeholder } from '../widgets/placeholder';
+import { MockChatModel } from './mocks';
 
 describe('MultiChatPanel', () => {
   let rmRegistry: RenderMimeRegistry;
@@ -79,40 +82,6 @@ describe('MultiChatPanel', () => {
     });
   });
 
-  describe('openInMain', () => {
-    it('should not dispose the model when moving to main area', async () => {
-      const { MockChatModel } = await import('./mocks');
-      const model = new MockChatModel();
-      const openInMain = jest.fn().mockResolvedValue(true);
-
-      const panel = new MultiChatPanel({ rmRegistry, openInMain });
-      panel.open({ model, displayName: 'test-chat' });
-
-      // Find the 'moveMain' toolbar button by name and invoke its click handler directly.
-      // ToolbarButton stores the onClick handler independently of DOM/React rendering,
-      // so this works without attaching the widget to the document.
-      const toolbar = panel.current!.toolbar;
-      const names = Array.from(toolbar.names());
-      const moveMainIdx = names.indexOf('moveMain');
-      expect(moveMainIdx).toBeGreaterThan(-1);
-
-      const moveButton = (toolbar.layout as any).widgets[
-        moveMainIdx
-      ] as ToolbarButton;
-
-      moveButton.onClick();
-      // Wait for the async onClick handler: openInMain resolves, then onClose is called.
-      await Promise.resolve();
-
-      expect(openInMain).toHaveBeenCalledWith(model.name);
-      expect(model.isDisposed).toBe(false);
-      expect(panel.getLoadedModel('test-chat')).toBeUndefined();
-
-      model.dispose();
-      panel.dispose();
-    });
-  });
-
   describe('placeholderFactory', () => {
     it('should use defaultPlaceholder when no factory is provided', () => {
       const panel = new MultiChatPanel({ rmRegistry });
@@ -207,7 +176,7 @@ describe('MultiChatPanel', () => {
       panel.dispose();
     });
 
-    it('should call factory.create again when the placeholder is re-added after closing a chat', async () => {
+    it('should call factory.create again when the placeholder is re-added after closing a chat', () => {
       const factory: IChatPlaceholderFactory = {
         create: jest.fn().mockReturnValue(new Widget())
       };
@@ -219,7 +188,6 @@ describe('MultiChatPanel', () => {
       });
 
       // Simulate opening and closing a chat to trigger _addPlaceholder again.
-      const { MockChatModel } = await import('./mocks');
       const model = new MockChatModel();
       panel.open({ model, displayName: 'test-chat' });
       panel.unsetLoadedModel('test-chat');
@@ -240,7 +208,95 @@ describe('MultiChatPanel', () => {
         '.jp-Toolbar-responsive-opener'
       );
       expect(opener).not.toBeNull();
+      panel.dispose();
+    });
+  });
+
+  describe('chatToolbarFactory', () => {
+    const makeFactory = (items: ToolbarRegistry.IToolbarItem[]) =>
+      jest.fn().mockReturnValue(new ObservableList({ values: items }));
+
+    it('should call factory with the panel when a chat is opened', () => {
+      const factory = makeFactory([{ name: 'myButton', widget: new Widget() }]);
+      const panel = new MultiChatPanel({
+        rmRegistry,
+        chatToolbarFactory: factory
+      });
+      const model = new MockChatModel();
+      const chatWidget = panel.open({ model, displayName: 'test' });
+      expect(factory).toHaveBeenCalledTimes(1);
+      expect(factory.mock.calls[0][0].widget).toBe(chatWidget);
+      panel.dispose();
+    });
 
+    it('should add the created widget to the chat toolbar', () => {
+      const panel = new MultiChatPanel({
+        rmRegistry,
+        chatToolbarFactory: makeFactory([
+          { name: 'myButton', widget: new Widget() }
+        ])
+      });
+      const model = new MockChatModel();
+      panel.open({ model, displayName: 'test' });
+      const toolbar = (panel.current as unknown as PanelWithToolbar).toolbar;
+      expect(Array.from(toolbar.names())).toContain('myButton');
+      panel.dispose();
+    });
+
+    it('should insert the item before "close"', () => {
+      const panel = new MultiChatPanel({
+        rmRegistry,
+        chatToolbarFactory: makeFactory([
+          { name: 'myButton', widget: new Widget() }
+        ])
+      });
+      const model = new MockChatModel();
+      panel.open({ model, displayName: 'test' });
+      const toolbar = (panel.current as unknown as PanelWithToolbar).toolbar;
+      const names = Array.from(toolbar.names());
+      expect(names.indexOf('myButton')).toBeLessThan(names.indexOf('close'));
+      panel.dispose();
+    });
+
+    it('should add all items when factory returns multiple items', () => {
+      const panel = new MultiChatPanel({
+        rmRegistry,
+        chatToolbarFactory: makeFactory([
+          { name: 'item1', widget: new Widget() },
+          { name: 'item2', widget: new Widget() }
+        ])
+      });
+      const model = new MockChatModel();
+      panel.open({ model, displayName: 'test' });
+      const toolbar = (panel.current as unknown as PanelWithToolbar).toolbar;
+      const names = Array.from(toolbar.names());
+      expect(names).toContain('item1');
+      expect(names).toContain('item2');
+      panel.dispose();
+    });
+
+    it('should call factory again for each newly opened chat', () => {
+      const factory = jest.fn().mockImplementation(
+        () =>
+          new ObservableList<ToolbarRegistry.IToolbarItem>({
+            values: [{ name: 'myButton', widget: new Widget() }]
+          })
+      );
+      const panel = new MultiChatPanel({
+        rmRegistry,
+        chatToolbarFactory: factory
+      });
+      const model1 = new MockChatModel();
+      panel.open({ model: model1, displayName: 'chat1' });
+      panel.unsetLoadedModel('chat1');
+
+      const model2 = new MockChatModel();
+      panel.open({ model: model2, displayName: 'chat2' });
+
+      expect(factory).toHaveBeenCalledTimes(2);
+      expect(factory.mock.results[0].value).not.toBe(
+        factory.mock.results[1].value
+      );
       panel.dispose();
     });
   });

package/src/components/messages/message.tsx

@@ -10,8 +10,8 @@ import { MessageRenderer } from './message-renderer';
 import { AttachmentPreviewList } from '../attachments';
 import { ChatInput } from '../input';
 import { useChatContext } from '../../context';
-import { IInputModel, InputModel } from '../../input-model';
-import { IMessageContent, IMessage } from '../../types';
+import { InputModel } from '../../input-model';
+import { IMessageContent, IMessage, INewMessage } from '../../types';
 import { replaceSpanToMention } from '../../utils';
 
 export const MESSAGE_CONTAINER_CLASS = 'jp-chat-message-container';
@@ -96,8 +96,8 @@ export function ChatMessageBase(props: ChatMessageProps): JSX.Element {
     });
     const inputModel = new InputModel({
       chatContext: model.createChatContext(),
-      onSend: (input: string, model?: IInputModel) =>
-        updateMessage(message.id, input, model),
+      onSend: (newMessage: INewMessage) =>
+        updateMessage(message.id, newMessage),
       onCancel: () => cancelEdition(),
       value: body,
       activeCellManager: model.activeCellManager,
@@ -121,17 +121,27 @@ export function ChatMessageBase(props: ChatMessageProps): JSX.Element {
 
   // Update the content of the message.
   const updateMessage = useCallback(
-    (id: string, input: string, inputModel?: IInputModel): void => {
-      if (!canEdit || !inputModel) {
+    (id: string, newMessage: INewMessage): void => {
+      const inputModel = model.getEditionModel(id);
+      if (!canEdit || !inputModel || !model.updateMessage) {
+        setEdit(false);
         return;
       }
+
       // Update the message
       const updatedMessage = { ...message };
-      updatedMessage.body = input;
-      updatedMessage.attachments = inputModel.attachments;
-      updatedMessage.mentions = inputModel.mentions;
-      model.updateMessage!(id, updatedMessage);
-      model.getEditionModel(message.id)?.dispose();
+      if (newMessage.body) {
+        updatedMessage.body = newMessage.body;
+      }
+      if (newMessage.attachments) {
+        updatedMessage.attachments = newMessage.attachments;
+      }
+      if (newMessage.mentions) {
+        updatedMessage.mentions = newMessage.mentions;
+      }
+
+      model.updateMessage(id, updatedMessage);
+      inputModel.dispose();
       setEdit(false);
     },
     [model, message, canEdit]

package/src/index.ts

@@ -13,4 +13,5 @@ export * from './registers';
 export * from './selection-watcher';
 export * from './tokens';
 export * from './types';
+export * from './utils';
 export * from './widgets';