Implement message queue component

package.json

@@ -89,7 +89,7 @@
         "@mui/icons-material": "^7",
         "@mui/material": "^7",
         "ai": "^6.0.116",
-        "jupyter-chat-components": "^0.2.0",
+        "jupyter-chat-components": "^0.4.0",
         "jupyter-secrets-manager": "^0.5.0",
         "yaml": "^2.8.1",
         "zod": "^4.3.6"

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 ]
 
 dependencies = [
-    "jupyter-chat-components >=0.2.0,<0.3",
+    "jupyter-chat-components >=0.4.0,<0.5",
     "jupyter-secrets-manager >=0.5,<0.6",
     "jupyterlab-ai-commands >=0.3.1,<0.4",
 ]

src/chat-model.ts

@@ -232,7 +232,8 @@ export class AIChatModel extends AbstractChatModel {
       stopStreaming: () => this.stopStreaming(),
       clearMessages: () => this.clearMessages(),
       agentManager: this._agentManager,
-      addSystemMessage: (body: string) => this.addSystemMessage(body)
+      addSystemMessage: (body: string) => this.addSystemMessage(body),
+      removeQueuedMessage: (id: string) => this.removeQueuedMessage(id)
     };
   }
 
@@ -267,8 +268,19 @@ export class AIChatModel extends AbstractChatModel {
     this.messageAdded(message);
   }
 
+  /**
+   * Overridden to ensure the message queue display stays at the bottom.
+   */
+  public messageAdded(message: IMessageContent): void {
+    super.messageAdded(message);
+    if (this._queueMessageId && message.id !== this._queueMessageId) {
+      this._updateQueueUI();
+    }
+  }
+
   /**
    * Sends a message to the AI and generates a response.
+   * If agent is busy, the message is queued.
    * @param message The user message to send
    */
   async sendMessage(message: INewMessage): Promise<void> {
@@ -278,18 +290,6 @@ export class AIChatModel extends AbstractChatModel {
       return;
     }
 
-    // Add user message to chat
-    const userMessage: IMessageContent = {
-      body: message.body,
-      sender: this.user || { username: 'user', display_name: 'User' },
-      id: UUID.uuid4(),
-      time: Date.now() / 1000,
-      type: 'msg',
-      raw_time: false,
-      attachments: [...this.input.attachments]
-    };
-    this.messageAdded(userMessage);
-
     // Check if we have valid configuration
     if (!this._agentManager.hasValidConfig()) {
       const errorMessage: IMessageContent = {
@@ -304,7 +304,34 @@ export class AIChatModel extends AbstractChatModel {
       return;
     }
 
+    // Create user message content
+    const userMessage: IMessageContent = {
+      body: message.body,
+      sender: this.user || { username: 'user', display_name: 'User' },
+      id: UUID.uuid4(),
+      time: Date.now() / 1000,
+      type: 'msg',
+      raw_time: false,
+      attachments: [...this.input.attachments]
+    };
+
+    if (this._isBusy) {
+      this._messageQueue.push({
+        id: UUID.uuid4(),
+        body: message.body,
+        _originalMsg: userMessage
+      });
+      this.input.clearAttachments();
+      this._updateQueueUI();
+      return;
+    }
+
+    this._isBusy = true;
+    this.messageAdded(userMessage);
+
     try {
+      this.updateWriters([{ user: this._getAIUser() }]);
+
       // Process attachments and add their content to the message
       let enhancedMessage: UserContent = message.body;
       if (this.input.attachments.length > 0) {
@@ -327,8 +354,6 @@ export class AIChatModel extends AbstractChatModel {
         }
       }
 
-      this.updateWriters([{ user: this._getAIUser() }]);
-
       await this._agentManager.generateResponse(enhancedMessage);
     } catch (error) {
       const errorMessage: IMessageContent = {
@@ -341,10 +366,19 @@ export class AIChatModel extends AbstractChatModel {
       };
       this.messageAdded(errorMessage);
     } finally {
-      this.updateWriters([]);
+      this._drainQueue();
     }
   }
 
+  /**
+   * Removes a queued message by its ID.
+   * @param messageId The ID of the queued message to remove
+   */
+  removeQueuedMessage(messageId: string): void {
+    this._messageQueue = this._messageQueue.filter(msg => msg.id !== messageId);
+    this._updateQueueUI();
+  }
+
   /**
    * Save the chat as json file.
    */
@@ -873,6 +907,110 @@ export class AIChatModel extends AbstractChatModel {
     });
   }
 
+  /**
+   * Processes the next message in the queue, or marks the agent as idle.
+   */
+  private async _drainQueue(): Promise<void> {
+    if (this._messageQueue.length === 0) {
+      this._isBusy = false;
+      this.updateWriters([]);
+      // Remove the queue display message
+      if (this._queueMessageId) {
+        const queueMsg = this.messages.find(
+          msg => msg.id === this._queueMessageId
+        );
+        if (queueMsg) {
+          const idx = this.messages.indexOf(queueMsg);
+          if (idx !== -1) {
+            this.messagesDeleted(idx, 1);
+          }
+        }
+        this._queueMessageId = null;
+      }
+      return;
+    }
+
+    // Dequeue and push to chat
+    const next = this._messageQueue.shift()!;
+    next._originalMsg.time = Date.now() / 1000;
+    this.messageAdded(next._originalMsg!);
+    this._updateQueueUI();
+
+    try {
+      // Process attachments now.
+      let body = next.body;
+      if (next._originalMsg?.attachments?.length) {
+        const { textContents } = await Private.processAttachments(
+          next._originalMsg.attachments,
+          this.input.documentManager
+        );
+        if (textContents.length > 0) {
+          body += '\n\n--- Attached Files ---\n' + textContents.join('\n\n');
+        }
+      }
+
+      await this._agentManager.generateResponse(body);
+    } catch (error) {
+      const errorMessage: IMessageContent = {
+        body: `Error generating AI response: ${(error as Error).message}`,
+        sender: this._getAIUser(),
+        id: UUID.uuid4(),
+        time: Date.now() / 1000,
+        type: 'msg',
+        raw_time: false
+      };
+      this.messageAdded(errorMessage);
+    } finally {
+      this._drainQueue();
+    }
+  }
+
+  /**
+   * Creates or updates the message-queue chat component.
+   */
+  private _updateQueueUI(): void {
+    // Delete existing queue message if it exists to
+    // ensure the new one goes to the bottom.
+    if (this._queueMessageId) {
+      const existingMsg = this.messages.find(
+        msg => msg.id === this._queueMessageId
+      );
+      if (existingMsg) {
+        const idx = this.messages.indexOf(existingMsg);
+        if (idx !== -1) {
+          this.messagesDeleted(idx, 1);
+        }
+      }
+      this._queueMessageId = null;
+    }
+
+    if (this._messageQueue.length === 0) {
+      return;
+    }
+
+    this._queueMessageId = UUID.uuid4();
+    const queueBody = {
+      data: {
+        'application/vnd.jupyter.chat.components': 'message-queue'
+      },
+      metadata: {
+        messages: this._messageQueue.map(m => ({ id: m.id, body: m.body })),
+        targetId: this.name
+      }
+    };
+
+    const queueMessage: IMessageContent = {
+      body: '',
+      mime_model: queueBody,
+      sender: this._getAIUser(),
+      id: this._queueMessageId,
+      time: Date.now() / 1000,
+      type: 'msg',
+      raw_time: false
+    };
+    this.messageAdded(queueMessage);
+  }
+
   // Private fields
   private _settingsModel: IAISettingsModel;
   private _user: IUser;
@@ -884,19 +1022,30 @@ export class AIChatModel extends AbstractChatModel {
   private _autosave: boolean = false;
   private _autosaveChanged = new Signal<AIChatModel, boolean>(this);
   private _autosaveDebouncer: Debouncer;
+  private _messageQueue: Private.IQueuedItem[] = [];
+  private _isBusy = false;
+  private _queueMessageId: string | null = null;
 }
 
 namespace Private {
-  type IDisplayOutput =
+  export interface IQueuedItem {
+    id: string;
+    body: string;
+    _originalMsg: IMessageContent;
+  }
+
+  export type IDisplayOutput =
     | nbformat.IDisplayData
     | nbformat.IDisplayUpdate
     | nbformat.IExecuteResult;
 
-  const isPlainObject = (value: unknown): value is Record<string, unknown> => {
+  export const isPlainObject = (
+    value: unknown
+  ): value is Record<string, unknown> => {
     return typeof value === 'object' && value !== null && !Array.isArray(value);
   };
 
-  const isDisplayOutput = (value: unknown): value is IDisplayOutput => {
+  export const isDisplayOutput = (value: unknown): value is IDisplayOutput => {
     if (!isPlainObject(value)) {
       return false;
     }
@@ -909,7 +1058,7 @@ namespace Private {
     );
   };
 
-  const toMimeBundle = (
+  export const toMimeBundle = (
     value: IDisplayOutput,
     trustedMimeTypes: ReadonlySet<string>
   ): IMimeModelBody | null => {
@@ -937,7 +1086,7 @@ namespace Private {
    * Tool outputs are not guaranteed to be raw Jupyter IOPub messages; they are
    * often wrapped objects (for example `{ success, result: { outputs: [...] } }`).
    */
-  const toDisplayOutputs = (value: unknown): IDisplayOutput[] => {
+  export const toDisplayOutputs = (value: unknown): IDisplayOutput[] => {
     if (isDisplayOutput(value)) {
       return [value];
     }
@@ -1438,6 +1587,10 @@ export namespace AIChatModel {
      * The agent manager of the chat.
      */
     agentManager: IAgentManager;
+    /**
+     * Removes a queued message by its ID.
+     */
+    removeQueuedMessage: (id: string) => void;
   }
 
   /**

src/index.ts

@@ -565,7 +565,7 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
         );
 
         if (aiWriting) {
-          widget.inputToolbarRegistry?.hide('send');
+          widget.inputToolbarRegistry?.show('send');
           widget.inputToolbarRegistry?.show('stop');
         } else {
           widget.inputToolbarRegistry?.hide('stop');
@@ -659,6 +659,20 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
 
     if (chatComponentsFactory) {
       chatComponentsFactory.toolCallApproval = toolCallApproval;
+
+      /**
+       * The callback to remove a queued message.
+       */
+      chatComponentsFactory.removeQueuedMessage = (
+        targetId: string,
+        messageId: string
+      ) => {
+        const model = tracker.find(chat => chat.model.name === targetId)?.model;
+        if (!model) {
+          return;
+        }
+        (model as AIChatModel).removeQueuedMessage(messageId);
+      };
     }
 
     return tracker;

src/widgets/main-area-chat.ts

@@ -99,7 +99,7 @@ export class MainAreaChat extends MainAreaWidget<ChatWidget> {
     );
 
     if (aiWriting) {
-      this.content.inputToolbarRegistry?.hide('send');
+      this.content.inputToolbarRegistry?.show('send');
       this.content.inputToolbarRegistry?.show('stop');
     } else {
       this.content.inputToolbarRegistry?.hide('stop');