diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f76d709..708ae7e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -125,7 +125,7 @@ jobs: - name: Execute integration tests working-directory: packages/jupyterlab-${{ matrix.extension }}-chat/ui-tests run: | - jlpm playwright test + jlpm playwright test --retries=2 - name: Upload Playwright Test report if: always() diff --git a/packages/jupyter-chat/package.json b/packages/jupyter-chat/package.json index e7735d0..a1496cc 100644 --- a/packages/jupyter-chat/package.json +++ b/packages/jupyter-chat/package.json @@ -55,7 +55,9 @@ "@emotion/react": "^11.10.5", "@emotion/styled": "^11.10.5", "@jupyter/react-components": "^0.15.2", + "@jupyterlab/application": "^4.2.0", "@jupyterlab/apputils": "^4.3.0", + "@jupyterlab/notebook": "^4.2.0", "@jupyterlab/rendermime": "^4.2.0", "@jupyterlab/ui-components": "^4.2.0", "@lumino/commands": "^2.0.0", diff --git a/packages/jupyter-chat/src/active-cell-manager.ts b/packages/jupyter-chat/src/active-cell-manager.ts new file mode 100644 index 0000000..0aac237 --- /dev/null +++ b/packages/jupyter-chat/src/active-cell-manager.ts @@ -0,0 +1,318 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +import { JupyterFrontEnd, LabShell } from '@jupyterlab/application'; +import { Cell, ICellModel } from '@jupyterlab/cells'; +import { IChangedArgs } from '@jupyterlab/coreutils'; +import { INotebookTracker, NotebookActions } from '@jupyterlab/notebook'; +import { IError as CellError } from '@jupyterlab/nbformat'; +import { ISignal, Signal } from '@lumino/signaling'; + +type CellContent = { + type: string; + source: string; +}; + +type CellWithErrorContent = { + type: 'code'; + source: string; + error: { + name: string; + value: string; + traceback: string[]; + }; +}; + +export interface IActiveCellManager { + /** + * Whether the notebook is available and an active cell exists. + */ + readonly available: boolean; + /** + * The `CellError` output within the active cell, if any. + */ + readonly activeCellError: CellError | null; + /** + * A signal emitting when the active cell changed. + */ + readonly availabilityChanged: ISignal; + /** + * A signal emitting when the error state of the active cell changed. + */ + readonly activeCellErrorChanged: ISignal; + /** + * Returns an `ActiveCellContent` object that describes the current active + * cell. If no active cell exists, this method returns `null`. + * + * When called with `withError = true`, this method returns `null` if the + * active cell does not have an error output. Otherwise it returns an + * `ActiveCellContentWithError` object that describes both the active cell and + * the error output. + */ + getContent(withError: boolean): CellContent | CellWithErrorContent | null; + /** + * Inserts `content` in a new cell above the active cell. + */ + insertAbove(content: string): void; + /** + * Inserts `content` in a new cell below the active cell. + */ + insertBelow(content: string): void; + /** + * Replaces the contents of the active cell. + */ + replace(content: string): Promise; +} + +/** + * The active cell manager namespace. + */ +export namespace ActiveCellManager { + /** + * The constructor options. + */ + export interface IOptions { + /** + * The notebook tracker. + */ + tracker: INotebookTracker; + /** + * The current shell of the application. + */ + shell: JupyterFrontEnd.IShell; + } +} + +/** + * A manager that maintains a reference to the current active notebook cell in + * the main panel (if any), and provides methods for inserting or appending + * content to the active cell. + * + * The current active cell should be obtained by listening to the + * `activeCellChanged` signal. + */ +export class ActiveCellManager implements IActiveCellManager { + constructor(options: ActiveCellManager.IOptions) { + this._notebookTracker = options.tracker; + this._notebookTracker.activeCellChanged.connect(this._onActiveCellChanged); + options.shell.currentChanged?.connect(this._onMainAreaChanged); + if (options.shell instanceof LabShell) { + options.shell.layoutModified?.connect(this._onMainAreaChanged); + } + this._onMainAreaChanged(); + } + + /** + * Whether the notebook is available and an active cell exists. + */ + get available(): boolean { + return this._available; + } + + /** + * The `CellError` output within the active cell, if any. + */ + get activeCellError(): CellError | null { + return this._activeCellError; + } + + /** + * A signal emitting when the active cell changed. + */ + get availabilityChanged(): ISignal { + return this._availabilityChanged; + } + + /** + * A signal emitting when the error state of the active cell changed. + */ + get activeCellErrorChanged(): ISignal { + return this._activeCellErrorChanged; + } + + /** + * Returns an `ActiveCellContent` object that describes the current active + * cell. If no active cell exists, this method returns `null`. + * + * When called with `withError = true`, this method returns `null` if the + * active cell does not have an error output. Otherwise it returns an + * `ActiveCellContentWithError` object that describes both the active cell and + * the error output. + */ + getContent(withError: false): CellContent | null; + getContent(withError: true): CellWithErrorContent | null; + getContent(withError = false): CellContent | CellWithErrorContent | null { + const sharedModel = this._notebookTracker.activeCell?.model.sharedModel; + if (!sharedModel) { + return null; + } + + // case where withError = false + if (!withError) { + return { + type: sharedModel.cell_type, + source: sharedModel.getSource() + }; + } + + // case where withError = true + const error = this._activeCellError; + if (error) { + return { + type: 'code', + source: sharedModel.getSource(), + error: { + name: error.ename, + value: error.evalue, + traceback: error.traceback + } + }; + } + + return null; + } + + /** + * Inserts `content` in a new cell above the active cell. + */ + insertAbove(content: string): void { + const notebookPanel = this._notebookTracker.currentWidget; + if (!notebookPanel || !notebookPanel.isVisible) { + return; + } + + // create a new cell above the active cell and mark new cell as active + NotebookActions.insertAbove(notebookPanel.content); + // replace content of this new active cell + this.replace(content); + } + + /** + * Inserts `content` in a new cell below the active cell. + */ + insertBelow(content: string): void { + const notebookPanel = this._notebookTracker.currentWidget; + if (!notebookPanel || !notebookPanel.isVisible) { + return; + } + + // create a new cell below the active cell and mark new cell as active + NotebookActions.insertBelow(notebookPanel.content); + // replace content of this new active cell + this.replace(content); + } + + /** + * Replaces the contents of the active cell. + */ + async replace(content: string): Promise { + const notebookPanel = this._notebookTracker.currentWidget; + if (!notebookPanel || !notebookPanel.isVisible) { + return; + } + // get reference to active cell directly from Notebook API. this avoids the + // possibility of acting on an out-of-date reference. + const activeCell = this._notebookTracker.activeCell; + if (!activeCell) { + return; + } + + // wait for editor to be ready + await activeCell.ready; + + // replace the content of the active cell + /** + * NOTE: calling this method sometimes emits an error to the browser console: + * + * ``` + * Error: Calls to EditorView.update are not allowed while an update is in progress + * ``` + * + * However, there seems to be no impact on the behavior/stability of the + * JupyterLab application after this error is logged. Furthermore, this is + * the official API for setting the content of a cell in JupyterLab 4, + * meaning that this is likely unavoidable. + */ + activeCell.editor?.model.sharedModel.setSource(content); + } + + private _onMainAreaChanged = () => { + const value = this._notebookTracker.currentWidget?.isVisible ?? false; + if (value !== this._notebookVisible) { + this._notebookVisible = value; + this._available = !!this._activeCell && this._notebookVisible; + this._availabilityChanged.emit(this._available); + } + }; + + /** + * Handle the change of active notebook cell. + */ + private _onActiveCellChanged = ( + _: INotebookTracker, + activeCell: Cell | null + ): void => { + if (this._activeCell !== activeCell) { + this._activeCell?.model.stateChanged.disconnect(this._cellStateChange); + this._activeCell = activeCell; + + activeCell?.ready.then(() => { + this._activeCell?.model.stateChanged.connect(this._cellStateChange); + this._available = !!this._activeCell && this._notebookVisible; + this._availabilityChanged.emit(this._available); + this._activeCell?.disposed.connect(() => { + this._activeCell = null; + }); + }); + } + }; + + /** + * Handle the change of the active cell state. + */ + private _cellStateChange = ( + _: ICellModel, + change: IChangedArgs + ): void => { + if (change.name === 'executionCount') { + const currSharedModel = this._activeCell?.model.sharedModel; + const prevActiveCellError = this._activeCellError; + let currActiveCellError: CellError | null = null; + if (currSharedModel && 'outputs' in currSharedModel) { + currActiveCellError = + currSharedModel.outputs.find( + (output): output is CellError => output.output_type === 'error' + ) || null; + } + + // for some reason, the `CellError` object is not referentially stable, + // meaning that this condition always evaluates to `true` and the + // `activeCellErrorChanged` signal is emitted every 200ms, even when the + // error output is unchanged. this is why we have to rely on + // `execution_count` to track changes to the error output. + if (prevActiveCellError !== currActiveCellError) { + this._activeCellError = currActiveCellError; + this._activeCellErrorChanged.emit(this._activeCellError); + } + } + }; + + /** + * The notebook tracker. + */ + private _notebookTracker: INotebookTracker; + /** + * Whether the current notebook panel is visible or not. + */ + private _notebookVisible: boolean = false; + /** + * The active cell. + */ + private _activeCell: Cell | null = null; + private _available: boolean = false; + private _activeCellError: CellError | null = null; + private _availabilityChanged = new Signal(this); + private _activeCellErrorChanged = new Signal(this); +} diff --git a/packages/jupyter-chat/src/components/chat-messages.tsx b/packages/jupyter-chat/src/components/chat-messages.tsx index 09c674f..d96603b 100644 --- a/packages/jupyter-chat/src/components/chat-messages.tsx +++ b/packages/jupyter-chat/src/components/chat-messages.tsx @@ -398,6 +398,7 @@ export function ChatMessage(props: ChatMessageProps): JSX.Element { setEdit(true) : undefined} delete={canDelete ? () => deleteMessage(message.id) : undefined} /> diff --git a/packages/jupyter-chat/src/components/code-blocks/code-toolbar.tsx b/packages/jupyter-chat/src/components/code-blocks/code-toolbar.tsx new file mode 100644 index 0000000..0553e02 --- /dev/null +++ b/packages/jupyter-chat/src/components/code-blocks/code-toolbar.tsx @@ -0,0 +1,143 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +import { addAboveIcon, addBelowIcon } from '@jupyterlab/ui-components'; +import { Box } from '@mui/material'; +import React, { useEffect, useState } from 'react'; + +import { CopyButton } from './copy-button'; +import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button'; +import { IActiveCellManager } from '../../active-cell-manager'; +import { replaceCellIcon } from '../../icons'; +import { IChatModel } from '../../model'; + +const CODE_TOOLBAR_CLASS = 'jp-chat-code-toolbar'; +const CODE_TOOLBAR_ITEM_CLASS = 'jp-chat-code-toolbar-item'; + +export type CodeToolbarProps = { + /** + * The chat model. + */ + model: IChatModel; + /** + * The content of the Markdown code block this component is attached to. + */ + content: string; +}; + +export function CodeToolbar(props: CodeToolbarProps): JSX.Element { + const { content, model } = props; + const [toolbarEnable, setToolbarEnable] = useState( + model.config.enableCodeToolbar ?? true + ); + + const activeCellManager = model.activeCellManager; + + const [toolbarBtnProps, setToolbarBtnProps] = useState({ + content: content, + activeCellManager: activeCellManager, + activeCellAvailable: activeCellManager?.available ?? false + }); + + useEffect(() => { + activeCellManager?.availabilityChanged.connect(() => { + setToolbarBtnProps({ + content, + activeCellManager: activeCellManager, + activeCellAvailable: activeCellManager.available + }); + }); + + model.configChanged.connect((_, config) => { + setToolbarEnable(config.enableCodeToolbar ?? true); + }); + }, [model]); + + return activeCellManager === null || !toolbarEnable ? ( + <> + ) : ( + + + + + + + ); +} + +type ToolbarButtonProps = { + content: string; + activeCellAvailable?: boolean; + activeCellManager: IActiveCellManager | null; + className?: string; +}; + +function InsertAboveButton(props: ToolbarButtonProps) { + const tooltip = props.activeCellAvailable + ? 'Insert above active cell' + : 'Insert above active cell (no active cell)'; + + return ( + props.activeCellManager?.insertAbove(props.content)} + disabled={!props.activeCellAvailable} + > + + + ); +} + +function InsertBelowButton(props: ToolbarButtonProps) { + const tooltip = props.activeCellAvailable + ? 'Insert below active cell' + : 'Insert below active cell (no active cell)'; + + return ( + props.activeCellManager?.insertBelow(props.content)} + > + + + ); +} + +function ReplaceButton(props: ToolbarButtonProps) { + const tooltip = props.activeCellAvailable + ? 'Replace active cell' + : 'Replace active cell (no active cell)'; + + return ( + props.activeCellManager?.replace(props.content)} + > + + + ); +} diff --git a/packages/jupyter-chat/src/components/code-blocks/copy-button.tsx b/packages/jupyter-chat/src/components/code-blocks/copy-button.tsx new file mode 100644 index 0000000..266bd70 --- /dev/null +++ b/packages/jupyter-chat/src/components/code-blocks/copy-button.tsx @@ -0,0 +1,68 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +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.None]: 'Copy to clipboard', + [CopyStatus.Copying]: 'Copying…', + [CopyStatus.Copied]: 'Copied!' +}; + +type CopyButtonProps = { + value: string; + className?: string; +}; + +export function CopyButton(props: CopyButtonProps): JSX.Element { + const [copyStatus, setCopyStatus] = useState(CopyStatus.None); + const timeoutId = useRef(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 = window.setTimeout( + () => setCopyStatus(CopyStatus.None), + 1000 + ); + }, [copyStatus, props.value]); + + return ( + + + + ); +} diff --git a/packages/jupyter-chat/src/components/copy-button.tsx b/packages/jupyter-chat/src/components/copy-button.tsx deleted file mode 100644 index 68a0206..0000000 --- a/packages/jupyter-chat/src/components/copy-button.tsx +++ /dev/null @@ -1,55 +0,0 @@ -/* - * Copyright (c) Jupyter Development Team. - * Distributed under the terms of the Modified BSD License. - */ - -import React, { useState, useCallback } from 'react'; - -import { Box, Button } from '@mui/material'; - -enum CopyStatus { - None, - Copied -} - -const COPYBTN_TEXT_BY_STATUS: Record = { - [CopyStatus.None]: 'Copy to Clipboard', - [CopyStatus.Copied]: 'Copied!' -}; - -type CopyButtonProps = { - value: string; -}; - -export function CopyButton(props: CopyButtonProps): JSX.Element { - const [copyStatus, setCopyStatus] = useState(CopyStatus.None); - - const copy = useCallback(async () => { - try { - await navigator.clipboard.writeText(props.value); - } catch (err) { - console.error('Failed to copy text: ', err); - setCopyStatus(CopyStatus.None); - return; - } - - setCopyStatus(CopyStatus.Copied); - setTimeout(() => setCopyStatus(CopyStatus.None), 1000); - }, [props.value]); - - return ( - - - - ); -} diff --git a/packages/jupyter-chat/src/components/mui-extras/contrasting-tooltip.tsx b/packages/jupyter-chat/src/components/mui-extras/contrasting-tooltip.tsx new file mode 100644 index 0000000..6100709 --- /dev/null +++ b/packages/jupyter-chat/src/components/mui-extras/contrasting-tooltip.tsx @@ -0,0 +1,27 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +import React from 'react'; +import { styled, Tooltip, TooltipProps, tooltipClasses } from '@mui/material'; + +/** + * A restyled MUI tooltip component that is dark by default to improve contrast + * against JupyterLab's default light theme. TODO: support dark themes. + */ +export const ContrastingTooltip = styled( + ({ className, ...props }: TooltipProps) => ( + + ) +)(({ theme }) => ({ + [`& .${tooltipClasses.tooltip}`]: { + backgroundColor: theme.palette.common.black, + color: theme.palette.common.white, + boxShadow: theme.shadows[1], + fontSize: 11 + }, + [`& .${tooltipClasses.arrow}`]: { + color: theme.palette.common.black + } +})); diff --git a/packages/jupyter-chat/src/components/mui-extras/tooltipped-icon-button.tsx b/packages/jupyter-chat/src/components/mui-extras/tooltipped-icon-button.tsx new file mode 100644 index 0000000..6954f88 --- /dev/null +++ b/packages/jupyter-chat/src/components/mui-extras/tooltipped-icon-button.tsx @@ -0,0 +1,84 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +import React from 'react'; +import { IconButton, IconButtonProps, TooltipProps } from '@mui/material'; + +import { ContrastingTooltip } from './contrasting-tooltip'; + +export type TooltippedIconButtonProps = { + onClick: () => unknown; + tooltip: string; + children: JSX.Element; + className?: string; + 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 `IconButton` component. + */ + iconButtonProps?: IconButtonProps; +}; + +/** + * A component that renders an MUI `IconButton` 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 IconButton when disabled. + * - Renders the IconButton with `line-height: 0` to avoid showing extra + * vertical space in SVG icons. + */ +export function TooltippedIconButton( + props: TooltippedIconButtonProps +): JSX.Element { + return ( + + {/* + By default, tooltips never appear when the IconButton 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 + */} + + + {props.children} + + + + ); +} diff --git a/packages/jupyter-chat/src/components/rendermime-markdown.tsx b/packages/jupyter-chat/src/components/rendermime-markdown.tsx index eda95f9..66a89f4 100644 --- a/packages/jupyter-chat/src/components/rendermime-markdown.tsx +++ b/packages/jupyter-chat/src/components/rendermime-markdown.tsx @@ -4,11 +4,12 @@ */ import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; -import React, { useState, useEffect, useRef } from 'react'; -import ReactDOM from 'react-dom'; +import React, { useState, useEffect } from 'react'; +import { createPortal } from 'react-dom'; -import { CopyButton } from './copy-button'; +import { CodeToolbar, CodeToolbarProps } from './code-blocks/code-toolbar'; import { MessageToolbar } from './toolbar'; +import { IChatModel } from '../model'; const MD_MIME_TYPE = 'text/markdown'; const RENDERMIME_MD_CLASS = 'jp-chat-rendermime-markdown'; @@ -17,6 +18,7 @@ type RendermimeMarkdownProps = { markdownStr: string; rmRegistry: IRenderMimeRegistry; appendContent?: boolean; + model: IChatModel; edit?: () => void; delete?: () => void; }; @@ -37,7 +39,11 @@ function RendermimeMarkdownBase(props: RendermimeMarkdownProps): JSX.Element { const [renderedContent, setRenderedContent] = useState( null ); - const containerRef = useRef(null); + + // each element is a two-tuple with the structure [codeToolbarRoot, codeToolbarProps]. + const [codeToolbarDefns, setCodeToolbarDefns] = useState< + Array<[HTMLDivElement, CodeToolbarProps]> + >([]); useEffect(() => { const renderContent = async () => { @@ -49,23 +55,29 @@ function RendermimeMarkdownBase(props: RendermimeMarkdownProps): JSX.Element { const renderer = props.rmRegistry.createRenderer(MD_MIME_TYPE); await renderer.renderModel(model); props.rmRegistry.latexTypesetter?.typeset(renderer.node); - - // Attach CopyButton to each 
 block
-      if (containerRef.current && renderer.node) {
-        const preBlocks = renderer.node.querySelectorAll('pre');
-        preBlocks.forEach(preBlock => {
-          const copyButtonContainer = document.createElement('div');
-          preBlock.parentNode?.insertBefore(
-            copyButtonContainer,
-            preBlock.nextSibling
-          );
-          ReactDOM.render(
-            ,
-            copyButtonContainer
-          );
-        });
+      if (!renderer.node) {
+        throw new Error(
+          'Rendermime was unable to render Markdown content within a chat message. Please report this upstream to Jupyter chat on GitHub.'
+        );
       }
 
+      const newCodeToolbarDefns: [HTMLDivElement, CodeToolbarProps][] = [];
+
+      // Attach CodeToolbar root element to each 
 block
+      const preBlocks = renderer.node.querySelectorAll('pre');
+      preBlocks.forEach(preBlock => {
+        const codeToolbarRoot = document.createElement('div');
+        preBlock.parentNode?.insertBefore(
+          codeToolbarRoot,
+          preBlock.nextSibling
+        );
+        newCodeToolbarDefns.push([
+          codeToolbarRoot,
+          { model: props.model, content: preBlock.textContent || '' }
+        ]);
+      });
+
+      setCodeToolbarDefns(newCodeToolbarDefns);
       setRenderedContent(renderer.node);
     };
 
@@ -73,7 +85,7 @@ function RendermimeMarkdownBase(props: RendermimeMarkdownProps): JSX.Element {
   }, [props.markdownStr, props.rmRegistry]);
 
   return (
-    

+    

       {renderedContent &&
         (appendContent ? (
           
 node && node.appendChild(renderedContent)} />
@@ -81,6 +93,18 @@ function RendermimeMarkdownBase(props: RendermimeMarkdownProps): JSX.Element {
           
 node && node.replaceChildren(renderedContent)} />
         ))}
       
+      {
+        // Render a `CodeToolbar` element underneath each code block.
+        // We use ReactDOM.createPortal() so each `CodeToolbar` element is able
+        // to use the context in the main React tree.
+        codeToolbarDefns.map(codeToolbarDefn => {
+          const [codeToolbarRoot, codeToolbarProps] = codeToolbarDefn;
+          return createPortal(
+            ,
+            codeToolbarRoot
+          );
+        })
+      }
     

   );
 }
diff --git a/packages/jupyter-chat/src/icons.ts b/packages/jupyter-chat/src/icons.ts
index e974770..e0dc6f8 100644
--- a/packages/jupyter-chat/src/icons.ts
+++ b/packages/jupyter-chat/src/icons.ts
@@ -9,6 +9,7 @@ import { LabIcon } from '@jupyterlab/ui-components';
 
 import chatSvgStr from '../style/icons/chat.svg';
 import readSvgStr from '../style/icons/read.svg';
+import replaceCellSvg from '../style/icons/replace-cell.svg';
 
 export const chatIcon = new LabIcon({
   name: 'jupyter-chat::chat',
@@ -19,3 +20,8 @@ export const readIcon = new LabIcon({
   name: 'jupyter-chat::read',
   svgstr: readSvgStr
 });
+
+export const replaceCellIcon = new LabIcon({
+  name: 'jupyter-ai::replace-cell',
+  svgstr: replaceCellSvg
+});
diff --git a/packages/jupyter-chat/src/index.ts b/packages/jupyter-chat/src/index.ts
index fb73421..d741b88 100644
--- a/packages/jupyter-chat/src/index.ts
+++ b/packages/jupyter-chat/src/index.ts
@@ -7,6 +7,7 @@ export * from './icons';
 export * from './model';
 export * from './registry';
 export * from './types';
+export * from './active-cell-manager';
 export * from './widgets/chat-error';
 export * from './widgets/chat-sidebar';
 export * from './widgets/chat-widget';
diff --git a/packages/jupyter-chat/src/model.ts b/packages/jupyter-chat/src/model.ts
index 2d267c1..9e39ee7 100644
--- a/packages/jupyter-chat/src/model.ts
+++ b/packages/jupyter-chat/src/model.ts
@@ -14,6 +14,7 @@ import {
   IConfig,
   IUser
 } from './types';
+import { IActiveCellManager } from './active-cell-manager';
 
 /**
  * The chat model interface.
@@ -49,6 +50,11 @@ export interface IChatModel extends IDisposable {
    */
   readonly messages: IChatMessage[];
 
+  /**
+   * Get the active cell manager.
+   */
+  readonly activeCellManager: IActiveCellManager | null;
+
   /**
    * A signal emitting when the messages list is updated.
    */
@@ -151,6 +157,8 @@ export class ChatModel implements IChatModel {
     this._config = { stackMessages: true, ...config };
 
     this._commands = options.commands;
+
+    this._activeCellManager = options.activeCellManager ?? null;
   }
 
   /**
@@ -160,6 +168,9 @@ export class ChatModel implements IChatModel {
     return this._messages;
   }
 
+  get activeCellManager(): IActiveCellManager | null {
+    return this._activeCellManager;
+  }
   /**
    * The chat model id.
    */
@@ -481,6 +492,7 @@ export class ChatModel implements IChatModel {
   private _config: IConfig;
   private _isDisposed = false;
   private _commands?: CommandRegistry;
+  private _activeCellManager: IActiveCellManager | null;
   private _notificationId: string | null = null;
   private _messagesUpdated = new Signal(this);
   private _configChanged = new Signal(this);
@@ -505,5 +517,10 @@ export namespace ChatModel {
      * Commands registry.
      */
     commands?: CommandRegistry;
+
+    /**
+     * Active cell manager
+     */
+    activeCellManager?: IActiveCellManager | null;
   }
 }
diff --git a/packages/jupyter-chat/src/types.ts b/packages/jupyter-chat/src/types.ts
index f55d667..ff193c1 100644
--- a/packages/jupyter-chat/src/types.ts
+++ b/packages/jupyter-chat/src/types.ts
@@ -35,6 +35,10 @@ export interface IConfig {
    * Whether to enable or not the notifications on unread messages.
    */
   unreadNotifications?: boolean;
+  /**
+   * Whether to enable or not the code toolbar.
+   */
+  enableCodeToolbar?: boolean;
 }
 
 /**
diff --git a/packages/jupyter-chat/style/icons/replace-cell.svg b/packages/jupyter-chat/style/icons/replace-cell.svg
new file mode 100644
index 0000000..74c5e49
--- /dev/null
+++ b/packages/jupyter-chat/style/icons/replace-cell.svg
@@ -0,0 +1,8 @@
+
+    
+    
+
diff --git a/packages/jupyterlab-collaborative-chat/package.json b/packages/jupyterlab-collaborative-chat/package.json
index 6e66777..2e474b1 100644
--- a/packages/jupyterlab-collaborative-chat/package.json
+++ b/packages/jupyterlab-collaborative-chat/package.json
@@ -66,6 +66,7 @@
     "@jupyterlab/coreutils": "^6.2.0",
     "@jupyterlab/docregistry": "^4.2.0",
     "@jupyterlab/launcher": "^4.2.0",
+    "@jupyterlab/notebook": "^4.2.0",
     "@jupyterlab/rendermime": "^4.2.0",
     "@jupyterlab/services": "^7.2.0",
     "@jupyterlab/settingregistry": "^4.2.0",
diff --git a/packages/jupyterlab-collaborative-chat/schema/factory.json b/packages/jupyterlab-collaborative-chat/schema/factory.json
index 5608cc7..676c130 100644
--- a/packages/jupyterlab-collaborative-chat/schema/factory.json
+++ b/packages/jupyterlab-collaborative-chat/schema/factory.json
@@ -7,15 +7,6 @@
   },
   "jupyter.lab.transform": true,
   "properties": {
-    "toolbar": {
-      "title": "File browser toolbar items",
-      "description": "Note: To disable a toolbar item,\ncopy it to User Preferences and add the\n\"disabled\" key. The following example will disable the uploader button:\n{\n  \"toolbar\": [\n    {\n      \"name\": \"uploader\",\n      \"disabled\": true\n    }\n  ]\n}\n\nToolbar description:",
-      "items": {
-        "$ref": "#/definitions/toolbarItem"
-      },
-      "type": "array",
-      "default": []
-    },
     "sendWithShiftEnter": {
       "description": "Whether to send a message via Shift-Enter instead of Enter.",
       "type": "boolean",
@@ -33,6 +24,21 @@
       "type": "boolean",
       "default": true,
       "readOnly": false
+    },
+    "enableCodeToolbar": {
+      "description": "Whether to enable or not the code toolbar.",
+      "type": "boolean",
+      "default": true,
+      "readOnly": false
+    },
+    "toolbar": {
+      "title": "File browser toolbar items",
+      "description": "Note: To disable a toolbar item,\ncopy it to User Preferences and add the\n\"disabled\" key. The following example will disable the uploader button:\n{\n  \"toolbar\": [\n    {\n      \"name\": \"uploader\",\n      \"disabled\": true\n    }\n  ]\n}\n\nToolbar description:",
+      "items": {
+        "$ref": "#/definitions/toolbarItem"
+      },
+      "type": "array",
+      "default": []
     }
   },
   "additionalProperties": false,
diff --git a/packages/jupyterlab-collaborative-chat/src/factory.ts b/packages/jupyterlab-collaborative-chat/src/factory.ts
index de75608..02d0820 100644
--- a/packages/jupyterlab-collaborative-chat/src/factory.ts
+++ b/packages/jupyterlab-collaborative-chat/src/factory.ts
@@ -3,7 +3,12 @@
  * Distributed under the terms of the Modified BSD License.
  */
 
-import { ChatWidget, IAutocompletionRegistry, IConfig } from '@jupyter/chat';
+import {
+  ChatWidget,
+  IActiveCellManager,
+  IAutocompletionRegistry,
+  IConfig
+} from '@jupyter/chat';
 import { IThemeManager } from '@jupyterlab/apputils';
 import { ABCWidgetFactory, DocumentRegistry } from '@jupyterlab/docregistry';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
@@ -101,6 +106,7 @@ export class CollaborativeChatModelFactory
     this._user = options.user;
     this._widgetConfig = options.widgetConfig;
     this._commands = options.commands;
+    this._activeCellManager = options.activeCellManager ?? null;
   }
 
   collaborative = true;
@@ -173,7 +179,8 @@ export class CollaborativeChatModelFactory
       ...options,
       user: this._user,
       widgetConfig: this._widgetConfig,
-      commands: this._commands
+      commands: this._commands,
+      activeCellManager: this._activeCellManager
     });
   }
 
@@ -181,4 +188,5 @@ export class CollaborativeChatModelFactory
   private _user: User.IIdentity | null;
   private _widgetConfig: IWidgetConfig;
   private _commands?: CommandRegistry;
+  private _activeCellManager: IActiveCellManager | null;
 }
diff --git a/packages/jupyterlab-collaborative-chat/src/index.ts b/packages/jupyterlab-collaborative-chat/src/index.ts
index 5229f48..d47a7da 100644
--- a/packages/jupyterlab-collaborative-chat/src/index.ts
+++ b/packages/jupyterlab-collaborative-chat/src/index.ts
@@ -4,7 +4,9 @@
  */
 
 import {
+  ActiveCellManager,
   AutocompletionRegistry,
+  IActiveCellManager,
   IAutocompletionRegistry,
   chatIcon,
   readIcon
@@ -30,6 +32,7 @@ import {
 import { PathExt } from '@jupyterlab/coreutils';
 import { DocumentRegistry } from '@jupyterlab/docregistry';
 import { ILauncher } from '@jupyterlab/launcher';
+import { INotebookTracker } from '@jupyterlab/notebook';
 import { IObservableList } from '@jupyterlab/observables';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { Contents } from '@jupyterlab/services';
@@ -43,13 +46,20 @@ import {
   CollaborativeChatModelFactory
 } from './factory';
 import { CollaborativeChatModel } from './model';
-import { chatFileType, CommandIDs, IChatPanel, IChatFactory } from './token';
+import {
+  chatFileType,
+  CommandIDs,
+  IChatPanel,
+  IChatFactory,
+  IActiveCellManagerToken
+} from './token';
 import { ChatPanel, CollaborativeChatPanel } from './widget';
 import { YChat } from './ychat';
 
 const FACTORY = 'Chat';
 
 const pluginIds = {
+  activeCellManager: 'jupyterlab-collaborative-chat:activeCellManager',
   autocompletionRegistry:
     'jupyterlab-collaborative-chat:autocompletionRegistry',
   chatCommands: 'jupyterlab-collaborative-chat:commands',
@@ -79,6 +89,7 @@ const docFactories: JupyterFrontEndPlugin = {
   autoStart: true,
   requires: [IRenderMimeRegistry],
   optional: [
+    IActiveCellManagerToken,
     IAutocompletionRegistry,
     ICollaborativeDrive,
     ILayoutRestorer,
@@ -91,6 +102,7 @@ const docFactories: JupyterFrontEndPlugin = {
   activate: (
     app: JupyterFrontEnd,
     rmRegistry: IRenderMimeRegistry,
+    activeCellManager: IActiveCellManager | null,
     autocompletionRegistry: IAutocompletionRegistry,
     drive: ICollaborativeDrive | null,
     restorer: ILayoutRestorer | null,
@@ -116,6 +128,7 @@ const docFactories: JupyterFrontEndPlugin = {
     let sendWithShiftEnter = false;
     let stackMessages = true;
     let unreadNotifications = true;
+    let enableCodeToolbar = true;
     function loadSetting(setting: ISettingRegistry.ISettings): void {
       // Read the settings and convert to the correct type
       sendWithShiftEnter = setting.get('sendWithShiftEnter')
@@ -123,10 +136,12 @@ const docFactories: JupyterFrontEndPlugin = {
       stackMessages = setting.get('stackMessages').composite as boolean;
       unreadNotifications = setting.get('unreadNotifications')
         .composite as boolean;
+      enableCodeToolbar = setting.get('enableCodeToolbar').composite as boolean;
       widgetConfig.configChanged.emit({
         sendWithShiftEnter,
         stackMessages,
-        unreadNotifications
+        unreadNotifications,
+        enableCodeToolbar
       });
     }
 
@@ -165,7 +180,8 @@ const docFactories: JupyterFrontEndPlugin = {
     const widgetConfig = new WidgetConfig({
       sendWithShiftEnter,
       stackMessages,
-      unreadNotifications
+      unreadNotifications,
+      enableCodeToolbar
     });
 
     // Namespace for the tracker
@@ -190,7 +206,8 @@ const docFactories: JupyterFrontEndPlugin = {
         const modelFactory = new CollaborativeChatModelFactory({
           user,
           widgetConfig,
-          commands: app.commands
+          commands: app.commands,
+          activeCellManager
         });
         app.docRegistry.addModelFactory(modelFactory);
       })
@@ -255,11 +272,12 @@ const chatCommands: JupyterFrontEndPlugin = {
   description: 'The commands to create or open a chat',
   autoStart: true,
   requires: [ICollaborativeDrive, IChatFactory],
-  optional: [IChatPanel, ICommandPalette, ILauncher],
+  optional: [IActiveCellManagerToken, IChatPanel, ICommandPalette, ILauncher],
   activate: (
     app: JupyterFrontEnd,
     drive: ICollaborativeDrive,
     factory: IChatFactory,
+    activeCellManager: IActiveCellManager | null,
     chatPanel: ChatPanel | null,
     commandPalette: ICommandPalette | null,
     launcher: ILauncher | null
@@ -463,7 +481,8 @@ const chatCommands: JupyterFrontEndPlugin = {
                 user,
                 sharedModel,
                 widgetConfig,
-                commands: app.commands
+                commands: app.commands,
+                activeCellManager
               });
 
               /**
@@ -506,13 +525,19 @@ const chatPanel: JupyterFrontEndPlugin = {
   autoStart: true,
   provides: IChatPanel,
   requires: [ICollaborativeDrive, IRenderMimeRegistry],
-  optional: [IAutocompletionRegistry, ILayoutRestorer, IThemeManager],
+  optional: [
+    IAutocompletionRegistry,
+    ILayoutRestorer,
+    INotebookTracker,
+    IThemeManager
+  ],
   activate: (
     app: JupyterFrontEnd,
     drive: ICollaborativeDrive,
     rmRegistry: IRenderMimeRegistry,
     autocompletionRegistry: IAutocompletionRegistry,
     restorer: ILayoutRestorer | null,
+    notebookTracker: INotebookTracker,
     themeManager: IThemeManager | null
   ): ChatPanel => {
     const { commands } = app;
@@ -588,4 +613,30 @@ const chatPanel: JupyterFrontEndPlugin = {
   }
 };
 
-export default [autocompletionPlugin, chatCommands, docFactories, chatPanel];
+/**
+ * Extension providing the active cell manager.
+ */
+const activeCellManager: JupyterFrontEndPlugin = {
+  id: pluginIds.activeCellManager,
+  description: 'the active cell manager plugin',
+  autoStart: true,
+  requires: [INotebookTracker],
+  provides: IActiveCellManagerToken,
+  activate: (
+    app: JupyterFrontEnd,
+    notebookTracker: INotebookTracker
+  ): IActiveCellManager => {
+    return new ActiveCellManager({
+      tracker: notebookTracker,
+      shell: app.shell
+    });
+  }
+};
+
+export default [
+  activeCellManager,
+  autocompletionPlugin,
+  chatCommands,
+  docFactories,
+  chatPanel
+];
diff --git a/packages/jupyterlab-collaborative-chat/src/token.ts b/packages/jupyterlab-collaborative-chat/src/token.ts
index 9d236c9..e83d8d1 100644
--- a/packages/jupyterlab-collaborative-chat/src/token.ts
+++ b/packages/jupyterlab-collaborative-chat/src/token.ts
@@ -93,3 +93,10 @@ export const CommandIDs = {
 export const IChatPanel = new Token(
   'jupyter-collaborative-chat:IChatPanel'
 );
+
+/**
+ * The active cell manager plugin.
+ */
+export const IActiveCellManagerToken = new Token(
+  'jupyter-collaborative-chat:IActiveCellManager'
+);
diff --git a/packages/jupyterlab-collaborative-chat/ui-tests/tests/autocompletion.spec.ts b/packages/jupyterlab-collaborative-chat/ui-tests/tests/autocompletion.spec.ts
index e1fb3ba..d8769cc 100644
--- a/packages/jupyterlab-collaborative-chat/ui-tests/tests/autocompletion.spec.ts
+++ b/packages/jupyterlab-collaborative-chat/ui-tests/tests/autocompletion.spec.ts
@@ -3,10 +3,11 @@
  * Distributed under the terms of the Modified BSD License.
  */
 
-import { expect, IJupyterLabPageFixture, test } from '@jupyterlab/galata';
-import { Locator } from '@playwright/test';
+import { expect, test } from '@jupyterlab/galata';
 
-const FILENAME = 'my-chat.chat';
+import { openChat } from './test-utils';
+
+const FILENAME = 'autocompletion.chat';
 const opener = '?';
 const commands = ['?test', '?other-test', '?last-test'];
 
@@ -46,26 +47,6 @@ const getPlugin = (pluginId: string): Promise => {
   });
 };
 
-const openChat = async (
-  page: IJupyterLabPageFixture,
-  filename: string
-): Promise => {
-  const panel = await page.activity.getPanelLocator(filename);
-  if (panel !== null && (await panel.count())) {
-    return panel;
-  }
-
-  await page.evaluate(async filepath => {
-    await window.jupyterapp.commands.execute('collaborative-chat:open', {
-      filepath
-    });
-  }, filename);
-  await page.waitForCondition(
-    async () => await page.activity.isTabActive(filename)
-  );
-  return (await page.activity.getPanelLocator(filename)) as Locator;
-};
-
 test.beforeEach(async ({ page }) => {
   // Expose a function to get a plugin.
   await page.evaluate(exposeDepsJs({ getPlugin }));
diff --git a/packages/jupyterlab-collaborative-chat/ui-tests/tests/code-toolbar.spec.ts b/packages/jupyterlab-collaborative-chat/ui-tests/tests/code-toolbar.spec.ts
new file mode 100644
index 0000000..90f381d
--- /dev/null
+++ b/packages/jupyterlab-collaborative-chat/ui-tests/tests/code-toolbar.spec.ts
@@ -0,0 +1,207 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import {
+  expect,
+  galata,
+  IJupyterLabPageFixture,
+  test
+} from '@jupyterlab/galata';
+
+import { openChat, sendMessage, USER } from './test-utils';
+
+test.use({
+  mockUser: USER,
+  mockSettings: {
+    ...galata.DEFAULT_SETTINGS,
+    'jupyterlab-collaborative-chat:factory': {
+      sendWithShiftEnter: true
+    }
+  }
+});
+
+const FILENAME = 'toolbar.chat';
+const CONTENT = 'print("This is a code cell")';
+const MESSAGE = `\`\`\`\n${CONTENT}\n\`\`\``;
+
+async function splitMainArea(page: IJupyterLabPageFixture, name: string) {
+  // Emulate drag and drop
+  const viewerHandle = page.activity.getTabLocator(name);
+  const viewerBBox = await viewerHandle.boundingBox();
+
+  await page.mouse.move(
+    viewerBBox!.x + 0.5 * viewerBBox!.width,
+    viewerBBox!.y + 0.5 * viewerBBox!.height
+  );
+  await page.mouse.down();
+  await page.mouse.move(viewerBBox!.x + 0.5 * viewerBBox!.width, 600);
+  await page.mouse.up();
+}
+
+test.describe('#codeToolbar', () => {
+  test.beforeEach(async ({ page }) => {
+    // Create a chat file
+    await page.filebrowser.contents.uploadContent('{}', 'text', FILENAME);
+  });
+
+  test.afterEach(async ({ page }) => {
+    if (await page.filebrowser.contents.fileExists(FILENAME)) {
+      await page.filebrowser.contents.deleteFile(FILENAME);
+    }
+  });
+
+  test('should have a code toolbar', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbar = message.locator('.jp-chat-code-toolbar');
+    const toolbarItems = message.locator('.jp-chat-code-toolbar-item');
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    await expect(message).toBeAttached();
+    await expect(toolbar).toBeAttached();
+    await expect(toolbarItems).toHaveCount(4);
+  });
+
+  test('should not have a code toolbar', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbar = message.locator('.jp-chat-code-toolbar');
+
+    await sendMessage(page, FILENAME, 'Simple message');
+    await expect(message).toBeAttached();
+    await expect(toolbar).not.toBeAttached();
+  });
+
+  test('buttons should be disabled without notebook', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    await expect(toolbarButtons).toHaveCount(4);
+    for (let i = 0; i < 3; i++) {
+      await expect(toolbarButtons.nth(i)).toBeDisabled();
+    }
+    await expect(toolbarButtons.nth(3)).toBeEnabled();
+  });
+
+  test('buttons should be disabled with a non visible notebook', async ({
+    page
+  }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    await page.notebook.createNew();
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    for (let i = 0; i < 3; i++) {
+      await expect(toolbarButtons.nth(i)).toBeDisabled();
+    }
+  });
+
+  test('buttons should be enabled with a visible notebook', async ({
+    page
+  }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    const notebook = await page.notebook.createNew();
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    await splitMainArea(page, notebook!);
+    for (let i = 0; i < 3; i++) {
+      await expect(toolbarButtons.nth(i)).toBeEnabled();
+    }
+  });
+
+  test('insert code above', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    const notebook = await page.notebook.createNew();
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    await splitMainArea(page, notebook!);
+    // activate the first cell.
+    await page.notebook.selectCells(0);
+
+    await toolbarButtons.first().click();
+
+    await page.activity.activateTab(notebook!);
+    await page.waitForCondition(
+      async () => (await page.notebook.getCellCount()) === 2
+    );
+    expect(await page.notebook.getCellTextInput(0)).toBe(`${CONTENT}\n`);
+  });
+
+  test('insert code below', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    const notebook = await page.notebook.createNew();
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    await splitMainArea(page, notebook!);
+
+    // activate the first cell.
+    await page.notebook.selectCells(0);
+
+    await toolbarButtons.nth(1).click();
+
+    await page.activity.activateTab(notebook!);
+    await page.waitForCondition(
+      async () => (await page.notebook.getCellCount()) === 2
+    );
+    expect(await page.notebook.getCellTextInput(1)).toBe(`${CONTENT}\n`);
+  });
+
+  test('replace active cell', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    const notebook = await page.notebook.createNew();
+
+    await sendMessage(page, FILENAME, MESSAGE);
+    await splitMainArea(page, notebook!);
+
+    // write content in the first cell.
+    const cell = await page.notebook.getCellLocator(0);
+    await cell?.getByRole('textbox').pressSequentially('initial content');
+    await toolbarButtons.nth(2).click();
+
+    await page.activity.activateTab(notebook!);
+    await page.waitForCondition(
+      async () =>
+        (await page.notebook.getCellTextInput(0)) !== 'initial content'
+    );
+    expect(await page.notebook.getCellTextInput(0)).toBe(`${CONTENT}\n`);
+  });
+
+  test('should copy code content', async ({ page }) => {
+    const chatPanel = await openChat(page, FILENAME);
+    const message = chatPanel.locator('.jp-chat-message');
+    const toolbarButtons = message.locator('.jp-chat-code-toolbar-item button');
+
+    const notebook = await page.notebook.createNew();
+
+    await sendMessage(page, FILENAME, MESSAGE);
+
+    // Copy the message code content to clipboard.
+    await toolbarButtons.last().click();
+
+    await page.activity.activateTab(notebook!);
+    const cell = await page.notebook.getCellLocator(0);
+    await cell?.getByRole('textbox').press('Control+V');
+    await page.waitForCondition(
+      async () => (await page.notebook.getCellTextInput(0)) !== ''
+    );
+    expect(await page.notebook.getCellTextInput(0)).toBe(`${CONTENT}\n`);
+  });
+});
diff --git a/packages/jupyterlab-collaborative-chat/ui-tests/tests/jupyterlab_collaborative_chat.spec.ts b/packages/jupyterlab-collaborative-chat/ui-tests/tests/jupyterlab_collaborative_chat.spec.ts
index d4251fa..9ac9e7a 100644
--- a/packages/jupyterlab-collaborative-chat/ui-tests/tests/jupyterlab_collaborative_chat.spec.ts
+++ b/packages/jupyterlab-collaborative-chat/ui-tests/tests/jupyterlab_collaborative_chat.spec.ts
@@ -13,19 +13,11 @@ import { Contents, User } from '@jupyterlab/services';
 import { ReadonlyJSONObject, UUID } from '@lumino/coreutils';
 import { Locator } from '@playwright/test';
 
+import { openChat, sendMessage, USER } from './test-utils';
+
 const FILENAME = 'my-chat.chat';
 const MSG_CONTENT = 'Hello World!';
-const USERNAME = UUID.uuid4();
-const USER: User.IUser = {
-  identity: {
-    username: USERNAME,
-    name: 'jovyan',
-    display_name: 'jovyan',
-    initials: 'JP',
-    color: 'var(--jp-collaborator-color1)'
-  },
-  permissions: {}
-};
+const USERNAME = USER.identity.username;
 
 test.use({
   mockUser: USER,
@@ -51,26 +43,6 @@ const readFileContent = async (
   }, filename);
 };
 
-const openChat = async (
-  page: IJupyterLabPageFixture,
-  filename: string
-): Promise => {
-  const panel = await page.activity.getPanelLocator(filename);
-  if (panel !== null && (await panel.count())) {
-    return panel;
-  }
-
-  await page.evaluate(async filepath => {
-    await window.jupyterapp.commands.execute('collaborative-chat:open', {
-      filepath
-    });
-  }, filename);
-  await page.waitForCondition(
-    async () => await page.activity.isTabActive(filename)
-  );
-  return (await page.activity.getPanelLocator(filename)) as Locator;
-};
-
 const openChatToSide = async (
   page: IJupyterLabPageFixture,
   filename: string
@@ -120,22 +92,6 @@ const openSidePanel = async (
   return panel.first();
 };
 
-const sendMessage = async (
-  page: IJupyterLabPageFixture,
-  filename: string = FILENAME,
-  content: string = MSG_CONTENT
-) => {
-  const chatPanel = await openChat(page, filename);
-  const input = chatPanel
-    .locator('.jp-chat-input-container')
-    .getByRole('combobox');
-  const sendButton = chatPanel
-    .locator('.jp-chat-input-container')
-    .getByRole('button');
-  await input.pressSequentially(content);
-  await sendButton.click();
-};
-
 test.describe('#commandPalette', () => {
   const name = FILENAME.replace('.chat', '');
 
@@ -590,7 +546,7 @@ test.describe('#messagesNavigation', () => {
       await expect(navigationBottom).toBeAttached();
       expect(navigationBottom).not.toHaveClass(/jp-chat-navigation-unread/);
 
-      await sendMessage(guestPage);
+      await sendMessage(guestPage, FILENAME, MSG_CONTENT);
 
       await expect(navigationBottom).toHaveClass(/jp-chat-navigation-unread/);
       expect(await navigationBottom.screenshot()).toMatchSnapshot(
@@ -604,7 +560,7 @@ test.describe('#messagesNavigation', () => {
       const navigationBottom = chatPanel.locator('.jp-chat-navigation-bottom');
       await messages.first().scrollIntoViewIfNeeded();
 
-      await sendMessage(guestPage);
+      await sendMessage(guestPage, FILENAME, MSG_CONTENT);
 
       await expect(navigationBottom).toHaveClass(/jp-chat-navigation-unread/);
       await navigationBottom.click();
@@ -686,7 +642,7 @@ test.describe('#notifications', () => {
     const messages = chatPanel.locator('.jp-chat-message');
     await messages.first().scrollIntoViewIfNeeded();
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     await page.waitForCondition(
       async () => (await page.notifications).length > 0
     );
@@ -707,7 +663,7 @@ test.describe('#notifications', () => {
     const messages = chatPanel.locator('.jp-chat-message');
     await messages.first().scrollIntoViewIfNeeded();
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     await page.waitForCondition(
       async () => (await page.notifications).length > 0
     );
@@ -727,7 +683,7 @@ test.describe('#notifications', () => {
     const messages = chatPanel.locator('.jp-chat-message');
     await messages.first().scrollIntoViewIfNeeded();
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     await page.waitForCondition(
       async () => (await page.notifications).length > 0
     );
@@ -738,7 +694,7 @@ test.describe('#notifications', () => {
       '1 incoming message(s) in my-chat.chat'
     );
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     notifications = await page.notifications;
     expect(notifications[0].message).toBe(
       '2 incoming message(s) in my-chat.chat'
@@ -750,7 +706,7 @@ test.describe('#notifications', () => {
     const messages = chatPanel.locator('.jp-chat-message');
     await messages.first().scrollIntoViewIfNeeded();
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     await page.waitForCondition(
       async () => (await page.notifications).length > 0
     );
@@ -780,7 +736,7 @@ test.describe('#notifications', () => {
       async () => (await page.notifications).length === 0
     );
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     await expect(messages).toHaveCount(messagesCount + 2);
 
     notifications = await page.notifications;
@@ -796,7 +752,7 @@ test.describe('#notifications', () => {
     const tabLabel = tab.locator('.lm-TabBar-tabLabel');
     await expect(tabLabel).toHaveText(FILENAME);
 
-    await sendMessage(guestPage);
+    await sendMessage(guestPage, FILENAME, MSG_CONTENT);
     const beforePseudo = tabLabel.evaluate(elem => {
       return window.getComputedStyle(elem, ':before');
     });
@@ -965,7 +921,7 @@ test.describe('#raw_time', () => {
     );
 
     // Send a new message
-    await sendMessage(page);
+    await sendMessage(page, FILENAME, MSG_CONTENT);
 
     expect(messages).toHaveCount(3);
     await expect(
diff --git a/packages/jupyterlab-collaborative-chat/ui-tests/tests/test-utils.ts b/packages/jupyterlab-collaborative-chat/ui-tests/tests/test-utils.ts
new file mode 100644
index 0000000..675548e
--- /dev/null
+++ b/packages/jupyterlab-collaborative-chat/ui-tests/tests/test-utils.ts
@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import { IJupyterLabPageFixture } from '@jupyterlab/galata';
+import { User } from '@jupyterlab/services';
+import { UUID } from '@lumino/coreutils';
+import { Locator } from '@playwright/test';
+
+export const USER: User.IUser = {
+  identity: {
+    username: UUID.uuid4(),
+    name: 'jovyan',
+    display_name: 'jovyan',
+    initials: 'JP',
+    color: 'var(--jp-collaborator-color1)'
+  },
+  permissions: {}
+};
+
+export const openChat = async (
+  page: IJupyterLabPageFixture,
+  filename: string
+): Promise => {
+  const panel = await page.activity.getPanelLocator(filename);
+  if (panel !== null && (await panel.count())) {
+    return panel;
+  }
+
+  await page.evaluate(async filepath => {
+    await window.jupyterapp.commands.execute('collaborative-chat:open', {
+      filepath
+    });
+  }, filename);
+  await page.waitForCondition(
+    async () => await page.activity.isTabActive(filename)
+  );
+  return (await page.activity.getPanelLocator(filename)) as Locator;
+};
+
+export const sendMessage = async (
+  page: IJupyterLabPageFixture,
+  filename: string,
+  content: string
+) => {
+  const chatPanel = await openChat(page, filename);
+  const input = chatPanel
+    .locator('.jp-chat-input-container')
+    .getByRole('combobox');
+  const sendButton = chatPanel
+    .locator('.jp-chat-input-container')
+    .getByRole('button');
+  await input.pressSequentially(content);
+  await sendButton.click();
+};
diff --git a/packages/jupyterlab-ws-chat/package.json b/packages/jupyterlab-ws-chat/package.json
index 20b7856..0b78f85 100644
--- a/packages/jupyterlab-ws-chat/package.json
+++ b/packages/jupyterlab-ws-chat/package.json
@@ -57,6 +57,7 @@
     "@jupyter/chat": "^0.2.0",
     "@jupyterlab/apputils": "^4.3.0",
     "@jupyterlab/coreutils": "^6.2.0",
+    "@jupyterlab/notebook": "^4.2.0",
     "@jupyterlab/rendermime": "^4.2.0",
     "@jupyterlab/services": "^7.2.0",
     "@jupyterlab/settingregistry": "^4.2.0",
diff --git a/packages/jupyterlab-ws-chat/schema/chat.json b/packages/jupyterlab-ws-chat/schema/chat.json
index e52eece..5b2863d 100644
--- a/packages/jupyterlab-ws-chat/schema/chat.json
+++ b/packages/jupyterlab-ws-chat/schema/chat.json
@@ -20,6 +20,12 @@
       "type": "boolean",
       "default": true,
       "readOnly": false
+    },
+    "enableCodeToolbar": {
+      "description": "Whether to enable or not the code toolbar.",
+      "type": "boolean",
+      "default": true,
+      "readOnly": false
     }
   },
   "additionalProperties": false
diff --git a/packages/jupyterlab-ws-chat/src/index.ts b/packages/jupyterlab-ws-chat/src/index.ts
index 33bd87e..1e0526a 100644
--- a/packages/jupyterlab-ws-chat/src/index.ts
+++ b/packages/jupyterlab-ws-chat/src/index.ts
@@ -4,6 +4,7 @@
  */
 
 import {
+  ActiveCellManager,
   AutocompletionRegistry,
   IAutocompletionRegistry,
   buildChatSidebar,
@@ -15,6 +16,7 @@ import {
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
 import { ReactWidget, IThemeManager } from '@jupyterlab/apputils';
+import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 
@@ -49,6 +51,7 @@ const chat: JupyterFrontEndPlugin = {
   optional: [
     IAutocompletionRegistry,
     ILayoutRestorer,
+    INotebookTracker,
     ISettingRegistry,
     IThemeManager
   ],
@@ -57,13 +60,23 @@ const chat: JupyterFrontEndPlugin = {
     rmRegistry: IRenderMimeRegistry,
     autocompletionRegistry: IAutocompletionRegistry,
     restorer: ILayoutRestorer | null,
+    notebookTracker: INotebookTracker,
     settingsRegistry: ISettingRegistry | null,
     themeManager: IThemeManager | null
   ) => {
+    // Create an active cell manager for code toolbar.
+    const activeCellManager = new ActiveCellManager({
+      tracker: notebookTracker,
+      shell: app.shell
+    });
+
     /**
      * Initialize chat handler, open WS connection
      */
-    const chatHandler = new WebSocketHandler({ commands: app.commands });
+    const chatHandler = new WebSocketHandler({
+      commands: app.commands,
+      activeCellManager
+    });
 
     /**
      * Load the settings.
@@ -71,6 +84,7 @@ const chat: JupyterFrontEndPlugin = {
     let sendWithShiftEnter = false;
     let stackMessages = true;
     let unreadNotifications = true;
+    let enableCodeToolbar = true;
     function loadSetting(setting: ISettingRegistry.ISettings): void {
       // Read the settings and convert to the correct type
       sendWithShiftEnter = setting.get('sendWithShiftEnter')
@@ -78,10 +92,12 @@ const chat: JupyterFrontEndPlugin = {
       stackMessages = setting.get('stackMessages').composite as boolean;
       unreadNotifications = setting.get('unreadNotifications')
         .composite as boolean;
+      enableCodeToolbar = setting.get('enableCodeToolbar').composite as boolean;
       chatHandler.config = {
         sendWithShiftEnter,
         stackMessages,
-        unreadNotifications
+        unreadNotifications,
+        enableCodeToolbar
       };
     }
 
diff --git a/yarn.lock b/yarn.lock
index 3f56fdd..4a0e2cb 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2281,7 +2281,9 @@ __metadata:
     "@emotion/react": ^11.10.5
     "@emotion/styled": ^11.10.5
     "@jupyter/react-components": ^0.15.2
+    "@jupyterlab/application": ^4.2.0
     "@jupyterlab/apputils": ^4.3.0
+    "@jupyterlab/notebook": ^4.2.0
     "@jupyterlab/rendermime": ^4.2.0
     "@jupyterlab/ui-components": ^4.2.0
     "@lumino/commands": ^2.0.0
@@ -2783,7 +2785,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@jupyterlab/notebook@npm:^4.2.3":
+"@jupyterlab/notebook@npm:^4.2.0, @jupyterlab/notebook@npm:^4.2.3":
   version: 4.2.3
   resolution: "@jupyterlab/notebook@npm:4.2.3"
   dependencies:
@@ -9856,6 +9858,7 @@ __metadata:
     "@jupyterlab/coreutils": ^6.2.0
     "@jupyterlab/docregistry": ^4.2.0
     "@jupyterlab/launcher": ^4.2.0
+    "@jupyterlab/notebook": ^4.2.0
     "@jupyterlab/rendermime": ^4.2.0
     "@jupyterlab/services": ^7.2.0
     "@jupyterlab/settingregistry": ^4.2.0
@@ -9903,6 +9906,7 @@ __metadata:
     "@jupyterlab/apputils": ^4.3.0
     "@jupyterlab/builder": ^4.2.0
     "@jupyterlab/coreutils": ^6.2.0
+    "@jupyterlab/notebook": ^4.2.0
     "@jupyterlab/rendermime": ^4.2.0
     "@jupyterlab/services": ^7.2.0
     "@jupyterlab/settingregistry": ^4.2.0