Add chat toolbar

jupyter_ai_tutor/handlers.py

@@ -2,6 +2,7 @@
 
 import tornado
 from jupyter_server.base.handlers import APIHandler
+from tornado.iostream import StreamClosedError
 
 
 class ExplainHandler(APIHandler):
@@ -65,13 +66,19 @@ async def post(self):
                     self.write(f"data: {json.dumps({'text': text})}\n\n")
                     self.flush()
 
+        except StreamClosedError:
+            return  # Client disconnected; stop streaming
         except Exception as e:
             self.log.exception("Error during tutor LLM call")
-            self.write(f"data: {json.dumps({'error': str(e)})}\n\n")
-            self.flush()
-
-        self.write("data: [DONE]\n\n")
-        self.flush()
-        self.finish()
-
+            try:
+                self.write(f"data: {json.dumps({'error': str(e)})}\n\n")
+                self.flush()
+            except StreamClosedError:
+                return
 
+        try:
+            self.write("data: [DONE]\n\n")
+            self.flush()
+            self.finish()
+        except StreamClosedError:
+            pass

src/api.ts

@@ -9,7 +9,8 @@ import { ServerConnection } from '@jupyterlab/services';
  */
 export async function* streamExplanation(
   body: string,
-  description?: string
+  description?: string,
+  signal?: AbortSignal
 ): AsyncGenerator<string, void, undefined> {
   const settings = ServerConnection.makeSettings();
   const url = URLExt.join(settings.baseUrl, 'api/jupyter-ai-tutor/explain');
@@ -19,7 +20,8 @@ export async function* streamExplanation(
     {
       method: 'POST',
       body: JSON.stringify({ body, description }),
-      headers: { 'Content-Type': 'application/json' }
+      headers: { 'Content-Type': 'application/json' },
+      signal
     },
     settings
   );

src/components/clear-button.tsx

@@ -0,0 +1,67 @@
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+
+import type { TranslationBundle } from '@jupyterlab/translation';
+
+import { clearIcon } from '@jupyterlab/ui-components';
+
+import React from 'react';
+
+import { ITutorChatContext } from '../model';
+
+/**
+ * Properties of the clear button.
+ */
+export interface IClearButtonProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The function to clear messages.
+   */
+  clearMessages: () => Promise<void>;
+  /**
+   * The application language translator.
+   */
+  translator: TranslationBundle;
+}
+
+/**
+ * The clear button component.
+ */
+export function ClearButton(props: IClearButtonProps): JSX.Element {
+  const { translator: trans } = props;
+  const tooltip = trans.__('Clear chat');
+  return (
+    <TooltippedButton
+      onClick={props.clearMessages}
+      tooltip={tooltip}
+      buttonProps={{
+        title: tooltip,
+        variant: 'outlined'
+      }}
+    >
+      <clearIcon.react />
+    </TooltippedButton>
+  );
+}
+
+/**
+ * Factory returning the clear button toolbar item.
+ */
+export function clearItem(
+  translator: TranslationBundle
+): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const { model } = props;
+      const clearMessages = async () =>
+        await (model.chatContext as ITutorChatContext).clearMessages();
+      const clearProps: IClearButtonProps = {
+        ...props,
+        clearMessages,
+        translator
+      };
+      return ClearButton(clearProps);
+    },
+    position: 0,
+    hidden: true
+  };
+}

src/components/index.ts

@@ -0,0 +1,2 @@
+export * from './clear-button';
+export * from './stop-button';

src/components/stop-button.tsx

@@ -0,0 +1,69 @@
+import { InputToolbarRegistry, TooltippedIconButton } from '@jupyter/chat';
+
+import type { TranslationBundle } from '@jupyterlab/translation';
+
+import StopIcon from '@mui/icons-material/Stop';
+
+import React from 'react';
+
+import { ITutorChatContext } from '../model';
+
+/**
+ * Properties of the stop button.
+ */
+export interface IStopButtonProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The function to stop streaming.
+   */
+  stopStreaming: () => void;
+  /**
+   * The application language translator.
+   */
+  translator: TranslationBundle;
+}
+
+/**
+ * The stop button component.
+ */
+export function StopButton(props: IStopButtonProps): JSX.Element {
+  const { translator: trans } = props;
+  const tooltip = trans.__('Stop streaming');
+  return (
+    <TooltippedIconButton
+      onClick={props.stopStreaming}
+      tooltip={tooltip}
+      buttonProps={{
+        title: tooltip
+      }}
+      sx={{
+        backgroundColor: 'var(--mui-palette-error-main, #d32f2f);'
+      }}
+    >
+      <StopIcon />
+    </TooltippedIconButton>
+  );
+}
+
+/**
+ * Factory returning the stop button toolbar item.
+ */
+export function stopItem(
+  translator: TranslationBundle
+): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const { model } = props;
+      const stopStreaming = () =>
+        (model.chatContext as ITutorChatContext).stopStreaming();
+      const stopProps: IStopButtonProps = {
+        ...props,
+        stopStreaming,
+        translator
+      };
+      return StopButton(stopProps);
+    },
+    position: 50,
+    hidden: true // Hidden by default, shown when streaming
+  };
+}

src/index.ts

@@ -2,7 +2,9 @@ import {
   AttachmentOpenerRegistry,
   ChatWidget,
   IAttachment,
-  INotebookAttachment
+  IChatModel,
+  INotebookAttachment,
+  InputToolbarRegistry
 } from '@jupyter/chat';
 import {
   JupyterFrontEnd,
@@ -15,7 +17,8 @@ import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { ITranslator, nullTranslator } from '@jupyterlab/translation';
 import { infoIcon } from '@jupyterlab/ui-components';
 
-import { TutorChatModel } from './model';
+import { clearItem, stopItem } from './components';
+import { TUTOR_USER, TutorChatModel } from './model';
 import { isContinuous } from './utils';
 
 const INFO_ICON_BASE_64 = btoa(infoIcon.svgstr);
@@ -53,6 +56,12 @@ const plugin: JupyterFrontEndPlugin<void> = {
     const { commands } = app;
     const trans = (translator ?? nullTranslator).load('jupyterlab');
 
+    // The input toolbar registry
+    const inputToolbarRegistry = InputToolbarRegistry.defaultToolbarRegistry();
+    inputToolbarRegistry.hide('send');
+    inputToolbarRegistry.addItem('stop', stopItem(trans));
+    inputToolbarRegistry.addItem('clear', clearItem(trans));
+
     // The attachment opener registry.
     const attachmentOpenerRegistry = new AttachmentOpenerRegistry();
 
@@ -111,7 +120,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
       welcomeMessage: trans.__(
         `## Select a code cell and click **Explain Code** <img src="data:image/svg+xml;base64,${INFO_ICON_BASE_64}" /> to get started.`
       ),
-      attachmentOpenerRegistry
+      attachmentOpenerRegistry,
+      inputToolbarRegistry
     });
     chatWidget.id = 'jupyter-ai-tutor-panel';
     chatWidget.title.label = trans.__('Tutor');
@@ -124,6 +134,33 @@ const plugin: JupyterFrontEndPlugin<void> = {
       commands.notifyCommandChanged(CommandIDs.explainCode);
     });
 
+    // Listen for writers change to display the stop button.
+    function writersChanged(_: IChatModel, writers: IChatModel.IWriter[]) {
+      // Check if AI is currently writing (streaming)
+      const aiWriting = writers.some(
+        writer => writer.user.username === TUTOR_USER.username
+      );
+
+      if (aiWriting) {
+        inputToolbarRegistry?.show('stop');
+      } else {
+        inputToolbarRegistry?.hide('stop');
+      }
+    }
+
+    tutorModel.writersChanged?.connect(writersChanged);
+
+    function messagesChanged(model: IChatModel) {
+      if (model.messages.length) {
+        inputToolbarRegistry?.show('clear');
+      } else {
+        inputToolbarRegistry?.hide('clear');
+      }
+    }
+
+    tutorModel.messagesUpdated.connect(messagesChanged);
+
+    // the command to ask for explanation.
     commands.addCommand(CommandIDs.explainCode, {
       label: trans.__('Explain Code'),
       caption: trans.__('Send cell content to AI tutor for explanation'),

src/model.ts

@@ -1,5 +1,4 @@
 import {
-  AbstractChatContext,
   AbstractChatModel,
   IAttachment,
   IChatContext,
@@ -18,29 +17,32 @@ interface ITutorNewMessage extends INewMessage {
   description?: string;
 }
 
-const TUTOR_USER: IUser = {
+export const TUTOR_USER: IUser = {
   username: 'tutor',
   display_name: 'Tutor',
   initials: 'T',
   bot: true,
   avatar_url: AI_AVATAR
 };
 
-class TutorChatContext extends AbstractChatContext {
-  get users(): IUser[] {
-    const users = new Map<string, IUser>();
-    for (const msg of this._model.messages) {
-      users.set(msg.sender.username, msg.sender);
-    }
-    return Array.from(users.values());
-  }
+export interface ITutorChatContext extends IChatContext {
+  /**
+   * The stop streaming callback.
+   */
+  stopStreaming: () => void;
+  /**
+   * The clear messages callback.
+   */
+  clearMessages: () => Promise<void>;
 }
 
 /**
  * Chat model for the AI tutor panel.
  * Routes messages to the tutor backend and streams responses into the chat.
  */
 export class TutorChatModel extends AbstractChatModel {
+  private _abortController: AbortController | null = null;
+
   constructor(options?: IChatModel.IOptions) {
     super(options);
     this.name = 'Tutor';
@@ -81,28 +83,51 @@ export class TutorChatModel extends AbstractChatModel {
     this.messageAdded(tutorMsgContent);
 
     const streamingMsg = this.messages[this.messages.length - 1];
+    this._abortController = new AbortController();
 
     try {
       let accumulated = '';
       for await (const chunk of streamExplanation(
         message.body,
-        message.description
+        message.description,
+        this._abortController.signal
       )) {
         accumulated += chunk;
         streamingMsg.update({ body: accumulated });
       }
     } catch (err) {
+      if (err instanceof Error && err.name === 'AbortError') return;
       const errorText = err instanceof Error ? err.message : String(err);
       streamingMsg.update({
         body: `Sorry, an error occurred: ${errorText}`
       });
       console.error('Tutor explanation failed:', err);
     } finally {
+      this._abortController = null;
       this.updateWriters([]);
     }
   }
 
-  createChatContext(): IChatContext {
-    return new TutorChatContext({ model: this });
+  createChatContext(): ITutorChatContext {
+    return {
+      name: this.name,
+      user: this.user,
+      users: [],
+      messages: this.messages,
+      stopStreaming: () => this.stopStreaming(),
+      clearMessages: () => this.clearMessages()
+    };
   }
+
+  stopStreaming = (): void => {
+    this._abortController?.abort();
+  };
+
+  /**
+   * Clears all messages from the chat and resets conversation state.
+   */
+  clearMessages = async (): Promise<void> => {
+    this.stopStreaming();
+    this.messagesDeleted(0, this.messages.length);
+  };
 }

style/base.css

@@ -4,6 +4,18 @@
     https://jupyterlab.readthedocs.io/en/stable/developer/css.html
 */
 
-#jupyter-ai-tutor-panel .jp-chat-input-container {
+/* Hide the input area */
+/* stylelint-disable-next-line selector-class-pattern */
+#jupyter-ai-tutor-panel .jp-chat-input-container .MuiAutocomplete-root {
   display: none;
 }
+
+/* Hide the whole input container if toolbar is empty */
+#jupyter-ai-tutor-panel
+  .jp-chat-input-container:has(.jp-chat-input-toolbar:empty) {
+  display: none;
+}
+
+#jupyter-ai-tutor-panel .jp-chat-input-container .jp-chat-input-toolbar {
+  border-top: none;
+}