Only allow users to stop streaming response to >their< last message

jupyterlab · Oct 11, 2024 · 8a94b19 · 8a94b19
1 parent 01d1c8c
commit 8a94b19
Show file tree

Hide file tree

Showing 9 changed files with 150 additions and 79 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py
@@ -155,10 +155,15 @@ def get_chat_user(self) -> ChatUser:
         # (`serverapp` is a property on all `JupyterHandler` subclasses)
         assert self.serverapp
         extensions = self.serverapp.extension_manager.extensions
-        collaborative = (
+        collaborative_legacy = (
             "jupyter_collaboration" in extensions
             and extensions["jupyter_collaboration"].enabled
         )
+        collaborative_v3 = (
+            "jupyter_server_ydoc" in extensions
+            and extensions["jupyter_server_ydoc"].enabled
+        )
+        collaborative = collaborative_legacy or collaborative_v3
 
         if collaborative:
             names = self.current_user.name.split(" ", maxsplit=2)
@@ -180,7 +185,7 @@ def get_chat_user(self) -> ChatUser:
         login = getpass.getuser()
         initials = login[0].capitalize()
         return ChatUser(
-            username=login,
+            username=self.current_user.username,
             initials=initials,
             name=login,
             display_name=login,
@@ -310,7 +315,8 @@ async def on_message(self, message):
         if isinstance(request, StopRequest):
             for history_message in self.chat_history[::-1]:
                 if (
-                    history_message.type == "agent-stream"
+                    history_message.id == request.target
+                    and history_message.type == "agent-stream"
                     and not history_message.complete
                 ):
                     try:
@@ -319,7 +325,7 @@ async def on_message(self, message):
                         # do nothing if the message was already interrupted
                         # or stream got completed (thread-safe way!)
                         pass
-            self.broadcast_message(StopMessage(target=history_message.id))
+                    self.broadcast_message(StopMessage(target=request.target))
             return
 
         chat_request = request

diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py
@@ -43,6 +43,10 @@ class StopRequest(BaseModel):
     """Request sent by human asking to stop streaming/generating the response"""
 
     type: Literal["stop"]
+    target: str
+    """
+    Message ID of the agent chat message to stop streaming.
+    """
 
 
 class ClearRequest(BaseModel):

diff --git a/packages/jupyter-ai/src/components/chat-input.tsx b/packages/jupyter-ai/src/components/chat-input.tsx
@@ -36,7 +36,10 @@ type ChatInputProps = {
    * `'Jupyternaut'`, but can differ for custom providers.
    */
   personaName: string;
-  streaming: boolean;
+  /**
+   * List of streaming messages that are owned by the current user.
+   */
+  streaming: AiService.AgentStreamMessage[];
 };
 
 /**
@@ -273,8 +276,11 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
 
   const sendButtonProps: SendButtonProps = {
     onSend,
-    onStop: () => {
-      props.chatHandler.sendMessage({ type: 'stop' });
+    onStop: (message: AiService.AgentStreamMessage) => {
+      props.chatHandler.sendMessage({
+        type: 'stop',
+        target: message.id
+      });
     },
     sendWithShiftEnter: props.sendWithShiftEnter,
     streaming: props.streaming,

diff --git a/packages/jupyter-ai/src/components/chat-input/send-button.tsx b/packages/jupyter-ai/src/components/chat-input/send-button.tsx
@@ -14,9 +14,9 @@ const FIX_TOOLTIP = '/fix requires an active code cell with an error';
 
 export type SendButtonProps = {
   onSend: (selection?: AiService.Selection) => unknown;
-  onStop: () => void;
+  onStop: (message: AiService.AgentStreamMessage) => void;
   sendWithShiftEnter: boolean;
-  streaming: boolean;
+  streaming: AiService.AgentStreamMessage[];
   currSlashCommand: string | null;
   inputExists: boolean;
   activeCellHasError: boolean;
@@ -37,7 +37,8 @@ export function SendButton(props: SendButtonProps): JSX.Element {
     setMenuOpen(false);
   }, []);
 
-  let disabled = props.streaming ? false : !props.inputExists;
+  const isStreaming = props.streaming.length !== 0;
+  let disabled = isStreaming ? false : !props.inputExists;
   if (props.currSlashCommand === '/fix' && !props.activeCellHasError) {
     disabled = true;
   }
@@ -60,7 +61,7 @@ export function SendButton(props: SendButtonProps): JSX.Element {
   const tooltip =
     props.currSlashCommand === '/fix' && !props.activeCellHasError
       ? FIX_TOOLTIP
-      : props.streaming
+      : isStreaming
       ? 'Stop streaming'
       : !props.inputExists
       ? 'Message must not be empty'
@@ -101,7 +102,7 @@ export function SendButton(props: SendButtonProps): JSX.Element {
   return (
     <Box sx={{ display: 'flex', flexWrap: 'nowrap' }}>
       <TooltippedButton
-        onClick={() => (props.streaming ? props.onStop() : props.onSend())}
+        onClick={() => (isStreaming ? props.onStop(props.streaming[props.streaming.length - 1]) : props.onSend())}
         disabled={disabled}
         tooltip={tooltip}
         buttonProps={{
@@ -114,7 +115,7 @@ export function SendButton(props: SendButtonProps): JSX.Element {
           borderRadius: '2px 0px 0px 2px'
         }}
       >
-        {props.streaming ? <StopIcon /> : <SendIcon />}
+        {isStreaming ? <StopIcon /> : <SendIcon />}
       </TooltippedButton>
       <TooltippedButton
         onClick={e => {

diff --git a/packages/jupyter-ai/src/components/chat.tsx b/packages/jupyter-ai/src/components/chat.tsx
@@ -7,6 +7,7 @@ import AddIcon from '@mui/icons-material/Add';
 import type { Awareness } from 'y-protocols/awareness';
 import type { IThemeManager } from '@jupyterlab/apputils';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+import type { User } from '@jupyterlab/services';
 import { ISignal } from '@lumino/signaling';
 
 import { JlThemeProvider } from './jl-theme-provider';
@@ -28,6 +29,7 @@ import {
   ActiveCellContextProvider,
   ActiveCellManager
 } from '../contexts/active-cell-context';
+import { UserContextProvider, useUserContext } from '../contexts/user-context';
 import { ScrollContainer } from './scroll-container';
 import { TooltippedIconButton } from './mui-extras/tooltipped-icon-button';
 import { TelemetryContextProvider } from '../contexts/telemetry-context';
@@ -76,6 +78,7 @@ function ChatBody({
     getPersonaName(messages)
   );
   const [sendWithShiftEnter, setSendWithShiftEnter] = useState(true);
+  const user = useUserContext();
 
   /**
    * Effect: fetch config on initial render
@@ -143,6 +146,13 @@ function ChatBody({
     );
   }
 
+  const myHumanMessageIds = new Set(messages.filter(m => m.type === 'human' && m.client.username === user?.identity.username).map(m => m.id));
+
+  const myStreamedMessages = messages.filter(m =>
+    m.type === 'agent-stream' && !m.complete
+    && myHumanMessageIds.has(m.reply_to)
+  ) as AiService.AgentStreamMessage[];
+
   return (
     <>
       <ScrollContainer sx={{ flexGrow: 1 }}>
@@ -157,10 +167,7 @@ function ChatBody({
       <ChatInput
         chatHandler={chatHandler}
         focusInputSignal={focusInputSignal}
-        streaming={
-          messages.filter(m => m.type === 'agent-stream' && !m.complete)
-            .length !== 0
-        }
+        streaming={myStreamedMessages}
         sx={{
           paddingLeft: 4,
           paddingRight: 4,
@@ -188,6 +195,7 @@ export type ChatProps = {
   focusInputSignal: ISignal<unknown, void>;
   messageFooter: IJaiMessageFooter | null;
   telemetryHandler: IJaiTelemetryHandler | null;
+  userManager: User.IManager
 };
 
 enum ChatView {
@@ -212,69 +220,71 @@ export function Chat(props: ChatProps): JSX.Element {
             activeCellManager={props.activeCellManager}
           >
             <TelemetryContextProvider telemetryHandler={props.telemetryHandler}>
-              <Box
-                // root box should not include padding as it offsets the vertical
-                // scrollbar to the left
-                sx={{
-                  width: '100%',
-                  height: '100%',
-                  boxSizing: 'border-box',
-                  background: 'var(--jp-layout-color0)',
-                  display: 'flex',
-                  flexDirection: 'column'
-                }}
-              >
-                {/* top bar */}
-                <Box sx={{ display: 'flex', justifyContent: 'space-between' }}>
-                  {view !== ChatView.Chat ? (
-                    <IconButton onClick={() => setView(ChatView.Chat)}>
-                      <ArrowBackIcon />
-                    </IconButton>
-                  ) : (
-                    <Box />
-                  )}
-                  {view === ChatView.Chat ? (
-                    <Box sx={{ display: 'flex' }}>
-                      {!showWelcomeMessage && (
-                        <TooltippedIconButton
-                          onClick={() =>
-                            props.chatHandler.sendMessage({ type: 'clear' })
-                          }
-                          tooltip="New chat"
-                        >
-                          <AddIcon />
-                        </TooltippedIconButton>
-                      )}
-                      <IconButton onClick={() => openSettingsView()}>
-                        <SettingsIcon />
+              <UserContextProvider userManager={props.userManager} >
+                <Box
+                  // root box should not include padding as it offsets the vertical
+                  // scrollbar to the left
+                  sx={{
+                    width: '100%',
+                    height: '100%',
+                    boxSizing: 'border-box',
+                    background: 'var(--jp-layout-color0)',
+                    display: 'flex',
+                    flexDirection: 'column'
+                  }}
+                >
+                  {/* top bar */}
+                  <Box sx={{ display: 'flex', justifyContent: 'space-between' }}>
+                    {view !== ChatView.Chat ? (
+                      <IconButton onClick={() => setView(ChatView.Chat)}>
+                        <ArrowBackIcon />
                       </IconButton>
-                    </Box>
-                  ) : (
-                    <Box />
+                    ) : (
+                      <Box />
+                    )}
+                    {view === ChatView.Chat ? (
+                      <Box sx={{ display: 'flex' }}>
+                        {!showWelcomeMessage && (
+                          <TooltippedIconButton
+                            onClick={() =>
+                              props.chatHandler.sendMessage({ type: 'clear' })
+                            }
+                            tooltip="New chat"
+                          >
+                            <AddIcon />
+                          </TooltippedIconButton>
+                        )}
+                        <IconButton onClick={() => openSettingsView()}>
+                          <SettingsIcon />
+                        </IconButton>
+                      </Box>
+                    ) : (
+                      <Box />
+                    )}
+                  </Box>
+                  {/* body */}
+                  {view === ChatView.Chat && (
+                    <ChatBody
+                      chatHandler={props.chatHandler}
+                      openSettingsView={openSettingsView}
+                      showWelcomeMessage={showWelcomeMessage}
+                      setShowWelcomeMessage={setShowWelcomeMessage}
+                      rmRegistry={props.rmRegistry}
+                      focusInputSignal={props.focusInputSignal}
+                      messageFooter={props.messageFooter}
+                    />
+                  )}
+                  {view === ChatView.Settings && (
+                    <ChatSettings
+                      rmRegistry={props.rmRegistry}
+                      completionProvider={props.completionProvider}
+                      openInlineCompleterSettings={
+                        props.openInlineCompleterSettings
+                      }
+                    />
                   )}
                 </Box>
-                {/* body */}
-                {view === ChatView.Chat && (
-                  <ChatBody
-                    chatHandler={props.chatHandler}
-                    openSettingsView={openSettingsView}
-                    showWelcomeMessage={showWelcomeMessage}
-                    setShowWelcomeMessage={setShowWelcomeMessage}
-                    rmRegistry={props.rmRegistry}
-                    focusInputSignal={props.focusInputSignal}
-                    messageFooter={props.messageFooter}
-                  />
-                )}
-                {view === ChatView.Settings && (
-                  <ChatSettings
-                    rmRegistry={props.rmRegistry}
-                    completionProvider={props.completionProvider}
-                    openInlineCompleterSettings={
-                      props.openInlineCompleterSettings
-                    }
-                  />
-                )}
-              </Box>
+              </UserContextProvider>
             </TelemetryContextProvider>
           </ActiveCellContextProvider>
         </CollaboratorsContextProvider>

diff --git a/packages/jupyter-ai/src/contexts/user-context.tsx b/packages/jupyter-ai/src/contexts/user-context.tsx
@@ -0,0 +1,39 @@
+import React, { useContext, useEffect, useState } from 'react';
+import type { User } from '@jupyterlab/services';
+import { PartialJSONObject } from '@lumino/coreutils';
+
+const UserContext = React.createContext<
+  User.IUser | null
+>(null);
+
+
+export function useUserContext(): User.IUser | null {
+  return useContext(UserContext);
+}
+
+type UserContextProviderProps = {
+  userManager: User.IManager;
+  children: React.ReactNode;
+};
+
+export function UserContextProvider({
+  userManager,
+  children
+}: UserContextProviderProps): JSX.Element {
+  const [user, setUser] = useState<User.IUser | null>(null);
+
+  useEffect(() => {
+    userManager.ready.then(() => {
+      setUser({identity: userManager.identity!, permissions: userManager.permissions as PartialJSONObject});
+    });
+    userManager.userChanged.connect((sender, newUser) => {
+      setUser(newUser);
+    });
+  }, []);
+
+  return (
+    <UserContext.Provider value={user}>
+      {children}
+    </UserContext.Provider>
+  );
+}
diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts
@@ -91,6 +91,7 @@ export namespace AiService {
 
   export type StopRequest = {
     type: 'stop';
+    target: string;
   };
 
   export type Collaborator = {

diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts
@@ -101,7 +101,8 @@ const plugin: JupyterFrontEndPlugin<IJaiCore> = {
         activeCellManager,
         focusInputSignal,
         messageFooter,
-        telemetryHandler
+        telemetryHandler,
+        app.serviceManager.user
       );
     } catch (e) {
       chatWidget = buildErrorWidget(themeManager);