Repository: jupyterlab/jupyter-ai (language: python)

Issue:
Remove "Include selection" checkbox


### Problem

The current UI for including a selection with your prompt is confusing. It requires a user to make a text selection in the main panel, after which the "Include selection" checkbox will appear. If a user checks this box, the checkbox remains checked until the user unchecks it manually while still having a text selection.

### Proposed Solution

Instead, we should implement this capability as a new feature of the "Send message" button, specifically by making it a dropdown button like the Send button in Slack's UI. Within the dropdown menu that is opened on click, there should be two additional options: "Send with text selection" and "Send with cell selection".


### Additional context

Sample image from Slack's UI:

![image](https://github.com/jupyterlab/jupyter-ai/assets/44106031/599faba2-cb83-47dc-b5d1-26c45facae4c)



Discussion / hints:

Fix patch:
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py
index f8c9f6f6b..318f9a5dd 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py
@@ -89,7 +89,7 @@ async def process_message(self, message: HumanChatMessage):
         selection: CellWithErrorSelection = message.selection
 
         # parse additional instructions specified after `/fix`
-        extra_instructions = message.body[4:].strip() or "None."
+        extra_instructions = message.prompt[4:].strip() or "None."
 
         self.get_llm_chain()
         with self.pending("Analyzing error"):
diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py
index 9f3fcee71..a08a87df9 100644
--- a/packages/jupyter-ai/jupyter_ai/handlers.py
+++ b/packages/jupyter-ai/jupyter_ai/handlers.py
@@ -247,12 +247,17 @@ async def on_message(self, message):
             self.log.error(e)
             return
 
+        message_body = chat_request.prompt
+        if chat_request.selection:
+            message_body += f"\n\n```\n{chat_request.selection.source}\n```\n"
+
         # message broadcast to chat clients
         chat_message_id = str(uuid.uuid4())
         chat_message = HumanChatMessage(
             id=chat_message_id,
             time=time.time(),
-            body=chat_request.prompt,
+            body=message_body,
+            prompt=chat_request.prompt,
             selection=chat_request.selection,
             client=self.chat_client,
         )
diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py
index 331f4cd94..f2fa098bb 100644
--- a/packages/jupyter-ai/jupyter_ai/models.py
+++ b/packages/jupyter-ai/jupyter_ai/models.py
@@ -14,22 +14,28 @@ class CellError(BaseModel):
     traceback: List[str]
 
 
+class TextSelection(BaseModel):
+    type: Literal["text"] = "text"
+    source: str
+
+
+class CellSelection(BaseModel):
+    type: Literal["cell"] = "cell"
+    source: str
+
+
 class CellWithErrorSelection(BaseModel):
     type: Literal["cell-with-error"] = "cell-with-error"
     source: str
     error: CellError
 
 
-Selection = Union[CellWithErrorSelection]
+Selection = Union[TextSelection, CellSelection, CellWithErrorSelection]
 
 
 # the type of message used to chat with the agent
 class ChatRequest(BaseModel):
     prompt: str
-    # TODO: This currently is only used when a user runs the /fix slash command.
-    # In the future, the frontend should set the text selection on this field in
-    # the `HumanChatMessage` it sends to JAI, instead of appending the text
-    # selection to `body` in the frontend.
     selection: Optional[Selection]
 
 
@@ -88,8 +94,13 @@ class HumanChatMessage(BaseModel):
     id: str
     time: float
     body: str
-    client: ChatClient
+    """The formatted body of the message to be rendered in the UI. Includes both
+    `prompt` and `selection`."""
+    prompt: str
+    """The prompt typed into the chat input by the user."""
     selection: Optional[Selection]
+    """The selection included with the prompt, if any."""
+    client: ChatClient
 
 
 class ClearMessage(BaseModel):
diff --git a/packages/jupyter-ai/src/components/chat-input.tsx b/packages/jupyter-ai/src/components/chat-input.tsx
index 64ed7447c..4ab9bcd36 100644
--- a/packages/jupyter-ai/src/components/chat-input.tsx
+++ b/packages/jupyter-ai/src/components/chat-input.tsx
@@ -6,9 +6,6 @@ import {
   SxProps,
   TextField,
   Theme,
-  FormGroup,
-  FormControlLabel,
-  Checkbox,
   InputAdornment,
   Typography
 } from '@mui/material';
@@ -27,15 +24,11 @@ import { ISignal } from '@lumino/signaling';
 import { AiService } from '../handler';
 import { SendButton, SendButtonProps } from './chat-input/send-button';
 import { useActiveCellContext } from '../contexts/active-cell-context';
+import { ChatHandler } from '../chat_handler';
 
 type ChatInputProps = {
-  value: string;
-  onChange: (newValue: string) => unknown;
-  onSend: (selection?: AiService.Selection) => unknown;
-  hasSelection: boolean;
-  includeSelection: boolean;
+  chatHandler: ChatHandler;
   focusInputSignal: ISignal<unknown, void>;
-  toggleIncludeSelection: () => unknown;
   sendWithShiftEnter: boolean;
   sx?: SxProps<Theme>;
   /**
@@ -105,6 +98,7 @@ function renderSlashCommandOption(
 }
 
 export function ChatInput(props: ChatInputProps): JSX.Element {
+  const [input, setInput] = useState('');
   const [slashCommandOptions, setSlashCommandOptions] = useState<
     SlashCommandOption[]
   >([]);
@@ -159,24 +153,24 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
    * chat input. Close the autocomplete when the user clears the chat input.
    */
   useEffect(() => {
-    if (props.value === '/') {
+    if (input === '/') {
       setOpen(true);
       return;
     }
 
-    if (props.value === '') {
+    if (input === '') {
       setOpen(false);
       return;
     }
-  }, [props.value]);
+  }, [input]);
 
   /**
    * Effect: Set current slash command
    */
   useEffect(() => {
-    const matchedSlashCommand = props.value.match(/^\s*\/\w+/);
+    const matchedSlashCommand = input.match(/^\s*\/\w+/);
     setCurrSlashCommand(matchedSlashCommand && matchedSlashCommand[0]);
-  }, [props.value]);
+  }, [input]);
 
   /**
    * Effect: ensure that the `highlighted` is never `true` when `open` is
@@ -190,25 +184,27 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
     }
   }, [open, highlighted]);
 
-  // TODO: unify the `onSend` implementation in `chat.tsx` and here once text
-  // selection is refactored.
-  function onSend() {
-    // case: /fix
+  function onSend(selection?: AiService.Selection) {
+    const prompt = input;
+    setInput('');
+
+    // if the current slash command is `/fix`, we always include a code cell
+    // with error output in the selection.
     if (currSlashCommand === '/fix') {
       const cellWithError = activeCell.manager.getContent(true);
       if (!cellWithError) {
         return;
       }
 
-      props.onSend({
-        ...cellWithError,
-        type: 'cell-with-error'
+      props.chatHandler.sendMessage({
+        prompt,
+        selection: { ...cellWithError, type: 'cell-with-error' }
       });
       return;
     }
 
-    // default case
-    props.onSend();
+    // otherwise, send a ChatRequest with the prompt and selection
+    props.chatHandler.sendMessage({ prompt, selection });
   }
 
   function handleKeyDown(event: React.KeyboardEvent<HTMLDivElement>) {
@@ -244,7 +240,7 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
     </span>
   );
 
-  const inputExists = !!props.value.trim();
+  const inputExists = !!input.trim();
   const sendButtonProps: SendButtonProps = {
     onSend,
     sendWithShiftEnter: props.sendWithShiftEnter,
@@ -258,9 +254,9 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
       <Autocomplete
         autoHighlight
         freeSolo
-        inputValue={props.value}
+        inputValue={input}
         onInputChange={(_, newValue: string) => {
-          props.onChange(newValue);
+          setInput(newValue);
         }}
         onHighlightChange={
           /**
@@ -322,23 +318,10 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
             FormHelperTextProps={{
               sx: { marginLeft: 'auto', marginRight: 0 }
             }}
-            helperText={props.value.length > 2 ? helperText : ' '}
+            helperText={input.length > 2 ? helperText : ' '}
           />
         )}
       />
-      {props.hasSelection && (
-        <FormGroup sx={{ display: 'flex', flexDirection: 'row' }}>
-          <FormControlLabel
-            control={
-              <Checkbox
-                checked={props.includeSelection}
-                onChange={props.toggleIncludeSelection}
-              />
-            }
-            label="Include selection"
-          />
-        </FormGroup>
-      )}
     </Box>
   );
 }
diff --git a/packages/jupyter-ai/src/components/chat-input/send-button.tsx b/packages/jupyter-ai/src/components/chat-input/send-button.tsx
index 6ab79a355..69ad3efc6 100644
--- a/packages/jupyter-ai/src/components/chat-input/send-button.tsx
+++ b/packages/jupyter-ai/src/components/chat-input/send-button.tsx
@@ -1,10 +1,18 @@
-import React from 'react';
+import React, { useCallback, useState } from 'react';
+import { Box, Menu, MenuItem, Typography } from '@mui/material';
+import KeyboardArrowDown from '@mui/icons-material/KeyboardArrowDown';
 import SendIcon from '@mui/icons-material/Send';
 
-import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button';
+import { TooltippedButton } from '../mui-extras/tooltipped-button';
+import { includeSelectionIcon } from '../../icons';
+import { useActiveCellContext } from '../../contexts/active-cell-context';
+import { useSelectionContext } from '../../contexts/selection-context';
+import { AiService } from '../../handler';
+
+const FIX_TOOLTIP = '/fix requires an active code cell with an error';
 
 export type SendButtonProps = {
-  onSend: () => unknown;
+  onSend: (selection?: AiService.Selection) => unknown;
   sendWithShiftEnter: boolean;
   currSlashCommand: string | null;
   inputExists: boolean;
@@ -12,34 +20,163 @@ export type SendButtonProps = {
 };
 
 export function SendButton(props: SendButtonProps): JSX.Element {
+  const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLElement | null>(null);
+  const [menuOpen, setMenuOpen] = useState(false);
+  const [textSelection] = useSelectionContext();
+  const activeCell = useActiveCellContext();
+
+  const openMenu = useCallback((el: HTMLElement | null) => {
+    setMenuAnchorEl(el);
+    setMenuOpen(true);
+  }, []);
+
+  const closeMenu = useCallback(() => {
+    setMenuOpen(false);
+  }, []);
+
   const disabled =
     props.currSlashCommand === '/fix'
       ? !props.inputExists || !props.activeCellHasError
       : !props.inputExists;
 
+  const includeSelectionDisabled = !(activeCell.exists || textSelection);
+
+  const includeSelectionTooltip =
+    props.currSlashCommand === '/fix'
+      ? FIX_TOOLTIP
+      : textSelection
+      ? `${textSelection.text.split('\n').length} lines selected`
+      : activeCell.exists
+      ? 'Code from 1 active cell'
+      : 'No selection or active cell';
+
   const defaultTooltip = props.sendWithShiftEnter
     ? 'Send message (SHIFT+ENTER)'
     : 'Send message (ENTER)';
 
   const tooltip =
     props.currSlashCommand === '/fix' && !props.activeCellHasError
-      ? '/fix requires a code cell with an error output selected'
+      ? FIX_TOOLTIP
       : !props.inputExists
       ? 'Message must not be empty'
       : defaultTooltip;
 
+  function sendWithSelection() {
+    // if the current slash command is `/fix`, `props.onSend()` should always
+    // include the code cell with error output, so the `selection` argument does
+    // not need to be defined.
+    if (props.currSlashCommand === '/fix') {
+      props.onSend();
+      closeMenu();
+      return;
+    }
+
+    // otherwise, parse the text selection or active cell, with the text
+    // selection taking precedence.
+    if (textSelection?.text) {
+      props.onSend({
+        type: 'text',
+        source: textSelection.text
+      });
+      closeMenu();
+      return;
+    }
+
+    if (activeCell.exists) {
+      props.onSend({
+        type: 'cell',
+        source: activeCell.manager.getContent(false).source
+      });
+      closeMenu();
+      return;
+    }
+  }
+
   return (
-    <TooltippedIconButton
-      onClick={() => props.onSend()}
-      disabled={disabled}
-      tooltip={tooltip}
-      iconButtonProps={{
-        size: 'small',
-        color: 'primary',
-        title: defaultTooltip
-      }}
-    >
-      <SendIcon />
-    </TooltippedIconButton>
+    <Box sx={{ display: 'flex', flexWrap: 'nowrap' }}>
+      <TooltippedButton
+        onClick={() => props.onSend()}
+        disabled={disabled}
+        tooltip={tooltip}
+        buttonProps={{
+          size: 'small',
+          title: defaultTooltip,
+          variant: 'contained'
+        }}
+        sx={{
+          minWidth: 'unset',
+          borderRadius: '2px 0px 0px 2px'
+        }}
+      >
+        <SendIcon />
+      </TooltippedButton>
+      <TooltippedButton
+        onClick={e => {
+          openMenu(e.currentTarget);
+        }}
+        disabled={disabled}
+        tooltip=""
+        buttonProps={{
+          variant: 'contained',
+          onKeyDown: e => {
+            if (e.key !== 'Enter' && e.key !== ' ') {
+              return;
+            }
+            openMenu(e.currentTarget);
+            // stopping propagation of this event prevents the prompt from being
+            // sent when the dropdown button is selected and clicked via 'Enter'.
+            e.stopPropagation();
+          }
+        }}
+        sx={{
+          minWidth: 'unset',
+          padding: '4px 0px',
+          borderRadius: '0px 2px 2px 0px',
+          borderLeft: '1px solid white'
+        }}
+      >
+        <KeyboardArrowDown />
+      </TooltippedButton>
+      <Menu
+        open={menuOpen}
+        onClose={closeMenu}
+        anchorEl={menuAnchorEl}
+        anchorOrigin={{
+          vertical: 'top',
+          horizontal: 'right'
+        }}
+        transformOrigin={{
+          vertical: 'bottom',
+          horizontal: 'right'
+        }}
+        sx={{
+          '& .MuiMenuItem-root': {
+            display: 'flex',
+            alignItems: 'center',
+            gap: '8px'
+          },
+          '& svg': {
+            lineHeight: 0
+          }
+        }}
+      >
+        <MenuItem
+          onClick={e => {
+            sendWithSelection();
+            // prevent sending second message with no selection
+            e.stopPropagation();
+          }}
+          disabled={includeSelectionDisabled}
+        >
+          <includeSelectionIcon.react />
+          <Box>
+            <Typography display="block">Send message with selection</Typography>
+            <Typography display="block" sx={{ opacity: 0.618 }}>
+              {includeSelectionTooltip}
+            </Typography>
+          </Box>
+        </MenuItem>
+      </Menu>
+    </Box>
   );
 }
diff --git a/packages/jupyter-ai/src/components/chat-messages.tsx b/packages/jupyter-ai/src/components/chat-messages.tsx
index 37d1f02ac..ec2e0cf1a 100644
--- a/packages/jupyter-ai/src/components/chat-messages.tsx
+++ b/packages/jupyter-ai/src/components/chat-messages.tsx
@@ -201,11 +201,6 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
       }}
     >
       {sortedMessages.map(message => {
-        // render selection in HumanChatMessage, if any
-        const markdownStr =
-          message.type === 'human' && message.selection
-            ? message.body + '\n\n```\n' + message.selection.source + '\n```\n'
-            : message.body;
         return (
           <Box key={message.id} sx={{ padding: 4 }}>
             <ChatMessageHeader
@@ -215,7 +210,7 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
             />
             <RendermimeMarkdown
               rmRegistry={props.rmRegistry}
-              markdownStr={markdownStr}
+              markdownStr={message.body}
               complete={
                 message.type === 'agent-stream' ? !!message.complete : true
               }
diff --git a/packages/jupyter-ai/src/components/chat.tsx b/packages/jupyter-ai/src/components/chat.tsx
index 2189fae44..09edfef5a 100644
--- a/packages/jupyter-ai/src/components/chat.tsx
+++ b/packages/jupyter-ai/src/components/chat.tsx
@@ -14,10 +14,7 @@ import { PendingMessages } from './pending-messages';
 import { ChatInput } from './chat-input';
 import { ChatSettings } from './chat-settings';
 import { AiService } from '../handler';
-import {
-  SelectionContextProvider,
-  useSelectionContext
-} from '../contexts/selection-context';
+import { SelectionContextProvider } from '../contexts/selection-context';
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
 import { CollaboratorsContextProvider } from '../contexts/collaborators-context';
@@ -66,9 +63,6 @@ function ChatBody({
     getPersonaName(messages)
   );
   const [showWelcomeMessage, setShowWelcomeMessage] = useState<boolean>(false);
-  const [includeSelection, setIncludeSelection] = useState(true);
-  const [input, setInput] = useState('');
-  const [textSelection] = useSelectionContext();
   const [sendWithShiftEnter, setSendWithShiftEnter] = useState(true);
 
   /**
@@ -107,25 +101,6 @@ function ChatBody({
     };
   }, [chatHandler]);
 
-  // no need to append to messageGroups imperatively here. all of that is
-  // handled by the listeners registered in the effect hooks above.
-  // TODO: unify how text selection & cell selection are handled
-  const onSend = async (selection?: AiService.Selection) => {
-    setInput('');
-
-    const prompt =
-      input +
-      (includeSelection && textSelection?.text
-        ? '\n\n```\n' + textSelection.text + '\n```'
-        : '');
-
-    // send message to backend
-    const messageId = await chatHandler.sendMessage({ prompt, selection });
-
-    // await reply from agent
-    await chatHandler.replyFor(messageId);
-  };
-
   const openSettingsView = () => {
     setShowWelcomeMessage(false);
     chatViewHandler(ChatView.Settings);
@@ -168,15 +143,8 @@ function ChatBody({
         <PendingMessages messages={pendingMessages} />
       </ScrollContainer>
       <ChatInput
-        value={input}
-        onChange={setInput}
-        onSend={onSend}
-        hasSelection={!!textSelection?.text}
-        includeSelection={includeSelection}
+        chatHandler={chatHandler}
         focusInputSignal={focusInputSignal}
-        toggleIncludeSelection={() =>
-          setIncludeSelection(includeSelection => !includeSelection)
-        }
         sx={{
           paddingLeft: 4,
           paddingRight: 4,
diff --git a/packages/jupyter-ai/src/components/mui-extras/tooltipped-button.tsx b/packages/jupyter-ai/src/components/mui-extras/tooltipped-button.tsx
new file mode 100644
index 000000000..363a4e832
--- /dev/null
+++ b/packages/jupyter-ai/src/components/mui-extras/tooltipped-button.tsx
@@ -0,0 +1,87 @@
+import React from 'react';
+import { Button, ButtonProps, SxProps, TooltipProps } from '@mui/material';
+
+import { ContrastingTooltip } from './contrasting-tooltip';
+
+export type TooltippedButtonProps = {
+  onClick: React.MouseEventHandler<HTMLButtonElement>;
+  tooltip: string;
+  children: JSX.Element;
+  disabled?: boolean;
+  placement?: TooltipProps['placement'];
+  /**
+   * The offset of the tooltip popup.
+   *
+   * The expected syntax is defined by the Popper library:
+   * https://popper.js.org/docs/v2/modifiers/offset/
+   */
+  offset?: [number, number];
+  'aria-label'?: string;
+  /**
+   * Props passed directly to the MUI `Button` component.
+   */
+  buttonProps?: ButtonProps;
+  /**
+   * Styles applied to the MUI `Button` component.
+   */
+  sx?: SxProps;
+};
+
+/**
+ * A component that renders an MUI `Button` with a high-contrast tooltip
+ * provided by `ContrastingTooltip`. This component differs from the MUI
+ * defaults in the following ways:
+ *
+ * - Shows the tooltip on hover even if disabled.
+ * - Renders the tooltip above the button by default.
+ * - Renders the tooltip closer to the button by default.
+ * - Lowers the opacity of the Button when disabled.
+ * - Renders the Button with `line-height: 0` to avoid showing extra
+ * vertical space in SVG icons.
+ *
+ * NOTE TO DEVS: Please keep this component's features synchronized with
+ * features available to `TooltippedIconButton`.
+ */
+export function TooltippedButton(props: TooltippedButtonProps): JSX.Element {
+  return (
+    <ContrastingTooltip
+      title={props.tooltip}
+      placement={props.placement ?? 'top'}
+      slotProps={{
+        popper: {
+          modifiers: [
+            {
+              name: 'offset',
+              options: {
+                offset: [0, -8]
+              }
+            }
+          ]
+        }
+      }}
+    >
+      {/*
+        By default, tooltips never appear when the Button is disabled. The
+        official way to support this feature in MUI is to wrap the child Button
+        element in a `span` element.
+
+        See: https://mui.com/material-ui/react-tooltip/#disabled-elements
+      */}
+      <span style={{ cursor: 'default' }}>
+        <Button
+          {...props.buttonProps}
+          onClick={props.onClick}
+          disabled={props.disabled}
+          sx={{
+            lineHeight: 0,
+            ...(props.disabled && { opacity: 0.5 }),
+            ...props.sx
+          }}
+          aria-label={props['aria-label']}
+        >
+          {props.children}
+        </Button>
+      </span>
+    </ContrastingTooltip>
+  );
+}
diff --git a/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx b/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx
index 6f97c7462..31279b4d9 100644
--- a/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx
+++ b/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx
@@ -1,10 +1,15 @@
 import React from 'react';
-import { IconButton, IconButtonProps, TooltipProps } from '@mui/material';
+import {
+  IconButton,
+  IconButtonProps,
+  SxProps,
+  TooltipProps
+} from '@mui/material';
 
 import { ContrastingTooltip } from './contrasting-tooltip';
 
 export type TooltippedIconButtonProps = {
-  onClick: () => unknown;
+  onClick: React.MouseEventHandler<HTMLButtonElement>;
   tooltip: string;
   children: JSX.Element;
   disabled?: boolean;
@@ -21,6 +26,10 @@ export type TooltippedIconButtonProps = {
    * Props passed directly to the MUI `IconButton` component.
    */
   iconButtonProps?: IconButtonProps;
+  /**
+   * Styles applied to the MUI `IconButton` component.
+   */
+  sx?: SxProps;
 };
 
 /**
@@ -34,6 +43,9 @@ export type TooltippedIconButtonProps = {
  * - Lowers the opacity of the IconButton when disabled.
  * - Renders the IconButton with `line-height: 0` to avoid showing extra
  * vertical space in SVG icons.
+ *
+ * NOTE TO DEVS: Please keep this component's features synchronized with
+ * features available to `TooltippedButton`.
  */
 export function TooltippedIconButton(
   props: TooltippedIconButtonProps
@@ -62,12 +74,16 @@ export function TooltippedIconButton(
 
         See: https://mui.com/material-ui/react-tooltip/#disabled-elements
       */}
-      <span>
+      <span style={{ cursor: 'default' }}>
         <IconButton
           {...props.iconButtonProps}
           onClick={props.onClick}
           disabled={props.disabled}
-          sx={{ lineHeight: 0, ...(props.disabled && { opacity: 0.5 }) }}
+          sx={{
+            lineHeight: 0,
+            ...(props.disabled && { opacity: 0.5 }),
+            ...props.sx
+          }}
           aria-label={props['aria-label']}
         >
           {props.children}
diff --git a/packages/jupyter-ai/src/contexts/active-cell-context.tsx b/packages/jupyter-ai/src/contexts/active-cell-context.tsx
index 72e93a8ca..a8a6ebcaf 100644
--- a/packages/jupyter-ai/src/contexts/active-cell-context.tsx
+++ b/packages/jupyter-ai/src/contexts/active-cell-context.tsx
@@ -83,7 +83,7 @@ export class ActiveCellManager {
    * `ActiveCellContentWithError` object that describes both the active cell and
    * the error output.
    */
-  getContent(withError: false): CellContent | null;
+  getContent(withError: false): CellContent;
   getContent(withError: true): CellWithErrorContent | null;
   getContent(withError = false): CellContent | CellWithErrorContent | null {
     const sharedModel = this._activeCell?.model.sharedModel;
diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts
index b93a4571a..b69c6ecef 100644
--- a/packages/jupyter-ai/src/handler.ts
+++ b/packages/jupyter-ai/src/handler.ts
@@ -57,13 +57,26 @@ export namespace AiService {
     traceback: string[];
   };
 
+  export type TextSelection = {
+    type: 'text';
+    source: string;
+  };
+
+  export type CellSelection = {
+    type: 'cell';
+    source: string;
+  };
+
   export type CellWithErrorSelection = {
     type: 'cell-with-error';
     source: string;
     error: CellError;
   };
 
-  export type Selection = CellWithErrorSelection;
+  export type Selection =
+    | TextSelection
+    | CellSelection
+    | CellWithErrorSelection;
 
   export type ChatRequest = {
     prompt: string;
@@ -101,9 +114,20 @@ export namespace AiService {
     type: 'human';
     id: string;
     time: number;
+    /**
+     * The formatted body of the message to be rendered in the UI. Includes both
+     * `prompt` and `selection`.
+     */
     body: string;
-    client: ChatClient;
+    /**
+     * The prompt typed into the chat input by the user.
+     */
+    prompt: string;
+    /**
+     * The selection included with the prompt, if any.
+     */
     selection?: Selection;
+    client: ChatClient;
   };
 
   export type ConnectionMessage = {
diff --git a/packages/jupyter-ai/src/icons.ts b/packages/jupyter-ai/src/icons.ts
index 12d194260..29aee13ea 100644
--- a/packages/jupyter-ai/src/icons.ts
+++ b/packages/jupyter-ai/src/icons.ts
@@ -5,6 +5,7 @@ import { LabIcon } from '@jupyterlab/ui-components';
 import chatSvgStr from '../style/icons/chat.svg';
 import jupyternautSvg from '../style/icons/jupyternaut.svg';
 import replaceCellSvg from '../style/icons/replace-cell.svg';
+import includeSelectionSvg from '../style/icons/include-selection.svg';
 
 export const chatIcon = new LabIcon({
   name: 'jupyter-ai::chat',
@@ -21,6 +22,11 @@ export const replaceCellIcon = new LabIcon({
   svgstr: replaceCellSvg
 });
 
+export const includeSelectionIcon = new LabIcon({
+  name: 'jupyter-ai::include-selection',
+  svgstr: includeSelectionSvg
+});
+
 // this icon is only used in the status bar.
 // to configure the icon shown on agent replies in the chat UI, please specify a
 // custom `Persona`.
diff --git a/packages/jupyter-ai/style/icons/include-selection.svg b/packages/jupyter-ai/style/icons/include-selection.svg
new file mode 100644
index 000000000..665ece235
--- /dev/null
+++ b/packages/jupyter-ai/style/icons/include-selection.svg
@@ -0,0 +1,5 @@
+<svg width="18" height="9" viewBox="0 0 18 9" fill="none" xmlns="http://www.w3.org/2000/svg">
+    <path fill-rule="evenodd" clip-rule="evenodd"
+        d="M-0.0018417 1.33824L0.837379 0.499023L1.35717 1.01871C1.35741 1.01862 1.35764 1.01853 1.35787 1.01845L3.411 3.07158H3.41049L4.83909 4.49987L3.41011 5.92872H3.41075L1.35769 7.98178C1.35749 7.98171 1.35729 7.98164 1.35709 7.98156L0.837023 8.50158L-0.00219727 7.66236L3.1547 4.49987L-0.0018417 1.33824ZM2.6821 8.07158H16.143C16.932 8.07158 17.5716 7.43201 17.5716 6.64301V2.35729C17.5716 1.56832 16.932 0.928721 16.143 0.928721H2.68236L4.82522 3.07158H15.4287V5.92872H4.82496L2.6821 8.07158ZM1.74238 4.50002L0.428719 5.81655V3.18349L1.74238 4.50002Z"
+        fill="#757575" />
+</svg>


Test patch:
diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py
index 6fab65763..d2e73ce6c 100644
--- a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py
+++ b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py
@@ -89,7 +89,13 @@ def chat_client():
 
 @pytest.fixture
 def human_chat_message(chat_client):
-    return HumanChatMessage(id="test", time=0, body="test message", client=chat_client)
+    return HumanChatMessage(
+        id="test",
+        time=0,
+        body="test message",
+        prompt="test message",
+        client=chat_client,
+    )
 
 
 def test_learn_index_permissions(tmp_path):