Add code toolbar to Jupyter AI chat

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

@@ -5,6 +5,7 @@ import SettingsIcon from '@mui/icons-material/Settings';
 import ArrowBackIcon from '@mui/icons-material/ArrowBack';
 import type { Awareness } from 'y-protocols/awareness';
 import type { IThemeManager } from '@jupyterlab/apputils';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 
 import { JlThemeProvider } from './jl-theme-provider';
 import { ChatMessages } from './chat-messages';
@@ -19,7 +20,10 @@ import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
 import { CollaboratorsContextProvider } from '../contexts/collaborators-context';
 import { IJaiCompletionProvider } from '../tokens';
-import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+import {
+  ActiveCellContextProvider,
+  ActiveCellManager
+} from '../contexts/active-cell-context';
 import { ScrollContainer } from './scroll-container';
 
 type ChatBodyProps = {
@@ -188,6 +192,7 @@ export type ChatProps = {
   chatView?: ChatView;
   completionProvider: IJaiCompletionProvider | null;
   openInlineCompleterSettings: () => void;
+  activeCellManager: ActiveCellManager;
 };
 
 enum ChatView {
@@ -202,51 +207,57 @@ export function Chat(props: ChatProps): JSX.Element {
     <JlThemeProvider themeManager={props.themeManager}>
       <SelectionContextProvider selectionWatcher={props.selectionWatcher}>
         <CollaboratorsContextProvider globalAwareness={props.globalAwareness}>
-          <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'
-            }}
+          <ActiveCellContextProvider
+            activeCellManager={props.activeCellManager}
           >
-            {/* top bar */}
-            <Box sx={{ display: 'flex', justifyContent: 'space-between' }}>
-              {view !== ChatView.Chat ? (
-                <IconButton onClick={() => setView(ChatView.Chat)}>
-                  <ArrowBackIcon />
-                </IconButton>
-              ) : (
-                <Box />
+            <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 ? (
+                  <IconButton onClick={() => setView(ChatView.Settings)}>
+                    <SettingsIcon />
+                  </IconButton>
+                ) : (
+                  <Box />
+                )}
+              </Box>
+              {/* body */}
+              {view === ChatView.Chat && (
+                <ChatBody
+                  chatHandler={props.chatHandler}
+                  setChatView={setView}
+                  rmRegistry={props.rmRegistry}
+                />
               )}
-              {view === ChatView.Chat ? (
-                <IconButton onClick={() => setView(ChatView.Settings)}>
-                  <SettingsIcon />
-                </IconButton>
-              ) : (
-                <Box />
+              {view === ChatView.Settings && (
+                <ChatSettings
+                  rmRegistry={props.rmRegistry}
+                  completionProvider={props.completionProvider}
+                  openInlineCompleterSettings={
+                    props.openInlineCompleterSettings
+                  }
+                />
               )}
             </Box>
-            {/* body */}
-            {view === ChatView.Chat && (
-              <ChatBody
-                chatHandler={props.chatHandler}
-                setChatView={setView}
-                rmRegistry={props.rmRegistry}
-              />
-            )}
-            {view === ChatView.Settings && (
-              <ChatSettings
-                rmRegistry={props.rmRegistry}
-                completionProvider={props.completionProvider}
-                openInlineCompleterSettings={props.openInlineCompleterSettings}
-              />
-            )}
-          </Box>
+          </ActiveCellContextProvider>
         </CollaboratorsContextProvider>
       </SelectionContextProvider>
     </JlThemeProvider>

packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx

@@ -0,0 +1,101 @@
+import React from 'react';
+import { Box } from '@mui/material';
+import { addAboveIcon, addBelowIcon } from '@jupyterlab/ui-components';
+
+import { CopyButton } from './copy-button';
+import { replaceCellIcon } from '../../icons';
+
+import {
+  ActiveCellManager,
+  useActiveCellContext
+} from '../../contexts/active-cell-context';
+import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button';
+
+export type CodeToolbarProps = {
+  /**
+   * The content of the Markdown code block this component is attached to.
+   */
+  content: string;
+};
+
+export function CodeToolbar(props: CodeToolbarProps): JSX.Element {
+  const [activeCellExists, activeCellManager] = useActiveCellContext();
+  const sharedToolbarButtonProps = {
+    content: props.content,
+    activeCellManager,
+    activeCellExists
+  };
+
+  return (
+    <Box
+      sx={{
+        display: 'flex',
+        justifyContent: 'flex-end',
+        alignItems: 'center',
+        padding: '6px 2px',
+        marginBottom: '1em',
+        border: '1px solid var(--jp-cell-editor-border-color)',
+        borderTop: 'none'
+      }}
+    >
+      <InsertAboveButton {...sharedToolbarButtonProps} />
+      <InsertBelowButton {...sharedToolbarButtonProps} />
+      <ReplaceButton {...sharedToolbarButtonProps} />
+      <CopyButton value={props.content} />
+    </Box>
+  );
+}
+
+type ToolbarButtonProps = {
+  content: string;
+  activeCellExists: boolean;
+  activeCellManager: ActiveCellManager;
+};
+
+function InsertAboveButton(props: ToolbarButtonProps) {
+  const tooltip = props.activeCellExists
+    ? 'Insert above active cell'
+    : 'Insert above active cell (no active cell)';
+
+  return (
+    <TooltippedIconButton
+      tooltip={tooltip}
+      onClick={() => props.activeCellManager.insertAbove(props.content)}
+      disabled={!props.activeCellExists}
+    >
+      <addAboveIcon.react height="16px" width="16px" />
+    </TooltippedIconButton>
+  );
+}
+
+function InsertBelowButton(props: ToolbarButtonProps) {
+  const tooltip = props.activeCellExists
+    ? 'Insert below active cell'
+    : 'Insert below active cell (no active cell)';
+
+  return (
+    <TooltippedIconButton
+      tooltip={tooltip}
+      disabled={!props.activeCellExists}
+      onClick={() => props.activeCellManager.insertBelow(props.content)}
+    >
+      <addBelowIcon.react height="16px" width="16px" />
+    </TooltippedIconButton>
+  );
+}
+
+function ReplaceButton(props: ToolbarButtonProps) {
+  const tooltip = props.activeCellExists
+    ? 'Replace active cell'
+    : 'Replace active cell (no active cell)';
+
+  return (
+    <TooltippedIconButton
+      tooltip={tooltip}
+      disabled={!props.activeCellExists}
+      onClick={() => props.activeCellManager.replace(props.content)}
+    >
+      <replaceCellIcon.react height="16px" width="16px" />
+    </TooltippedIconButton>
+  );
+}

packages/jupyter-ai/src/components/code-blocks/copy-button.tsx

@@ -0,0 +1,58 @@
+import React, { useState, useCallback, useRef } from 'react';
+
+import { copyIcon } from '@jupyterlab/ui-components';
+
+import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button';
+
+enum CopyStatus {
+  None,
+  Copying,
+  Copied
+}
+
+const COPYBTN_TEXT_BY_STATUS: Record<CopyStatus, string> = {
+  [CopyStatus.None]: 'Copy to Clipboard',
+  [CopyStatus.Copying]: 'Copying...',
+  [CopyStatus.Copied]: 'Copied!'
+};
+
+type CopyButtonProps = {
+  value: string;
+};
+
+export function CopyButton(props: CopyButtonProps): JSX.Element {
+  const [copyStatus, setCopyStatus] = useState<CopyStatus>(CopyStatus.None);
+  const timeoutId = useRef<number | null>(null);
+
+  const copy = useCallback(async () => {
+    // ignore if we are already copying
+    if (copyStatus === CopyStatus.Copying) {
+      return;
+    }
+
+    try {
+      await navigator.clipboard.writeText(props.value);
+    } catch (err) {
+      console.error('Failed to copy text: ', err);
+      setCopyStatus(CopyStatus.None);
+      return;
+    }
+
+    setCopyStatus(CopyStatus.Copied);
+    if (timeoutId.current) {
+      clearTimeout(timeoutId.current);
+    }
+    timeoutId.current = setTimeout(() => setCopyStatus(CopyStatus.None), 1000);
+  }, [copyStatus, props.value]);
+
+  return (
+    <TooltippedIconButton
+      tooltip={COPYBTN_TEXT_BY_STATUS[copyStatus]}
+      placement="top"
+      onClick={copy}
+      aria-label="Copy to Clipboard"
+    >
+      <copyIcon.react height="16px" width="16px" />
+    </TooltippedIconButton>
+  );
+}