Allow adding or removing attachment when editing a message

Author: brichet

packages/jupyter-chat/src/components/chat-input.tsx

@@ -28,15 +28,17 @@ const INPUT_BOX_CLASS = 'jp-chat-input-container';
 
 export function ChatInput(props: ChatInput.IProps): JSX.Element {
   const { documentManager, model } = props;
-  const [input, setInput] = useState<string>(model.value || '');
+  const [input, setInput] = useState<string>(model.value);
   const inputRef = useRef<HTMLInputElement>();
 
   const chatCommands = useChatCommands(model, props.chatCommandRegistry);
 
   const [sendWithShiftEnter, setSendWithShiftEnter] = useState<boolean>(
     model.config.sendWithShiftEnter ?? false
   );
-  const [attachments, setAttachments] = useState<IAttachment[]>([]);
+  const [attachments, setAttachments] = useState<IAttachment[]>(
+    model.attachments
+  );
 
   // Display the include selection menu if it is not explicitly hidden, and if at least
   // one of the tool to check for text or cell selection is enabled.

packages/jupyter-chat/src/components/chat-messages.tsx

@@ -4,6 +4,7 @@
  */
 
 import { Button } from '@jupyter/react-components';
+import { IDocumentManager } from '@jupyterlab/docmanager';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import {
   LabIcon,
@@ -21,7 +22,7 @@ import { ChatInput } from './chat-input';
 import { MarkdownRenderer } from './markdown-renderer';
 import { ScrollContainer } from './scroll-container';
 import { IChatCommandRegistry } from '../chat-commands';
-import { InputModel } from '../input-model';
+import { IInputModel, InputModel } from '../input-model';
 import { IChatModel } from '../model';
 import { IChatMessage, IUser } from '../types';
 
@@ -43,6 +44,7 @@ type BaseMessageProps = {
   rmRegistry: IRenderMimeRegistry;
   model: IChatModel;
   chatCommandRegistry?: IChatCommandRegistry;
+  documentManager?: IDocumentManager;
 };
 
 /**
@@ -341,6 +343,7 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
     const [deleted, setDeleted] = useState<boolean>(false);
     const [canEdit, setCanEdit] = useState<boolean>(false);
     const [canDelete, setCanDelete] = useState<boolean>(false);
+    const [inputModel, setInputModel] = useState<IInputModel | null>(null);
 
     // Look if the message can be deleted or edited.
     useEffect(() => {
@@ -356,6 +359,25 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
       }
     }, [model, message]);
 
+    // Create an input model only if the message is edited.
+    useEffect(() => {
+      if (edit && canEdit) {
+        setInputModel(
+          new InputModel({
+            value: message.body,
+            activeCellManager: model.activeCellManager,
+            selectionWatcher: model.selectionWatcher,
+            config: {
+              sendWithShiftEnter: model.config.sendWithShiftEnter
+            },
+            attachments: message.attachments
+          })
+        );
+      } else {
+        setInputModel(null);
+      }
+    }, [edit]);
+
     // Cancel the current edition of the message.
     const cancelEdition = (): void => {
       setEdit(false);
@@ -369,6 +391,7 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
       // Update the message
       const updatedMessage = { ...message };
       updatedMessage.body = input;
+      updatedMessage.attachments = inputModel?.attachments;
       model.updateMessage!(id, updatedMessage);
       setEdit(false);
     };
@@ -386,22 +409,14 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
       <div ref={ref} data-index={props.index}></div>
     ) : (
       <div ref={ref} data-index={props.index}>
-        {edit && canEdit ? (
+        {edit && canEdit && inputModel ? (
           <ChatInput
             onSend={(input: string) => updateMessage(message.id, input)}
             onCancel={() => cancelEdition()}
-            model={
-              new InputModel({
-                value: message.body,
-                activeCellManager: model.activeCellManager,
-                selectionWatcher: model.selectionWatcher,
-                config: {
-                  sendWithShiftEnter: model.config.sendWithShiftEnter
-                }
-              })
-            }
+            model={inputModel}
             hideIncludeSelection={true}
             chatCommandRegistry={props.chatCommandRegistry}
+            documentManager={props.documentManager}
           />
         ) : (
           <MarkdownRenderer
@@ -413,7 +428,9 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
             rendered={props.renderedPromise}
           />
         )}
-        {message.attachments && (
+        {message.attachments && !edit && (
+          // Display the attachments only if message is not edited, otherwise the
+          // input component display them.
           <AttachmentPreviewList attachments={message.attachments} />
         )}
       </div>

packages/jupyter-chat/src/components/chat.tsx

@@ -36,6 +36,7 @@ export function ChatBody(props: Chat.IChatBodyProps): JSX.Element {
         rmRegistry={props.rmRegistry}
         model={model}
         chatCommandRegistry={props.chatCommandRegistry}
+        documentManager={props.documentManager}
       />
       <ChatInput
         onSend={onSend}

packages/jupyter-chat/src/input-model.ts

@@ -113,6 +113,7 @@ export interface IInputModel extends IDisposable {
 export class InputModel implements IInputModel {
   constructor(options: InputModel.IOptions) {
     this._value = options.value || '';
+    this._attachments = options.attachments || [];
     this.cursorIndex = options.cursorIndex || this.value.length;
     this._activeCellManager = options.activeCellManager ?? null;
     this._selectionWatcher = options.selectionWatcher ?? null;
@@ -316,7 +317,7 @@ export class InputModel implements IInputModel {
   private _value: string;
   private _cursorIndex: number | null = null;
   private _currentWord: string | null = null;
-  private _attachments: IAttachment[] = [];
+  private _attachments: IAttachment[];
   private _activeCellManager: IActiveCellManager | null;
   private _selectionWatcher: ISelectionWatcher | null;
   private _config: InputModel.IConfig;
@@ -335,6 +336,12 @@ export namespace InputModel {
      * The initial value of the input.
      */
     value?: string;
+
+    /**
+     * The initial attachments.
+     */
+    attachments?: IAttachment[];
+
     /**
      * The current cursor index.
      * This refers to the index of the character in front of the cursor.

packages/jupyterlab-chat/src/model.ts

@@ -194,6 +194,12 @@ export class LabChatModel extends ChatModel implements DocumentRegistry.IModel {
         edited: true
       };
     }
+    const attachmentIds = updatedMessage.attachments?.map(attachment =>
+      this.sharedModel.setAttachment(attachment)
+    );
+    if (attachmentIds) {
+      message.attachments = attachmentIds;
+    }
     this.sharedModel.updateMessage(index, message as IYmessage);
   }