feat(chat): support auto-sync active selection in current Notebook file

TabbyML · Dec 20, 2024 · eda2e0a · eda2e0a
1 parent 4184a0e
commit eda2e0a
Show file tree

Hide file tree

Showing 8 changed files with 174 additions and 23 deletions.
diff --git a/clients/tabby-chat-panel/src/index.ts b/clients/tabby-chat-panel/src/index.ts
@@ -44,12 +44,20 @@ export interface LineRange {
    */
   end: number
 }
-
+/**
+ * Represents a range of lines in a notebook file.
+ */
+interface NotebookCellRange extends LineRange {
+  /**
+   * 0-based cell index
+   */
+  cellIndex: number
+}
 /**
  * Represents a location in a file.
  * It could be a 1-based line number, a line range, a position or a position range.
  */
-export type Location = number | LineRange | Position | PositionRange
+export type Location = number | LineRange | Position | PositionRange | NotebookCellRange
 
 export interface FileContext {
   kind: 'file'

diff --git a/clients/vscode/src/chat/WebviewHelper.ts b/clients/vscode/src/chat/WebviewHelper.ts
@@ -266,7 +266,7 @@ export class WebviewHelper {
   }
 
   public isSupportedSchemeForActiveSelection(scheme: string) {
-    const supportedSchemes = ["file", "untitled"];
+    const supportedSchemes = ["file", "untitled", "vscode-notebook-cell"];
     return supportedSchemes.includes(scheme);
   }
 

diff --git a/clients/vscode/src/chat/fileContext.ts b/clients/vscode/src/chat/fileContext.ts
@@ -1,9 +1,22 @@
 import type { TextEditor, TextDocument } from "vscode";
 import type { FileContext } from "tabby-chat-panel";
 import type { GitProvider } from "../git/GitProvider";
-import { workspace, window, Position, Range, Selection, TextEditorRevealType, Uri, ViewColumn, commands } from "vscode";
+import {
+  workspace,
+  window,
+  Position,
+  Range,
+  Selection,
+  TextEditorRevealType,
+  Uri,
+  ViewColumn,
+  commands,
+  NotebookRange,
+  NotebookEditorRevealType,
+} from "vscode";
 import path from "path";
 import { getLogger } from "../logger";
+import { parseVscodeNotebookCellURI } from "./utils";
 
 const logger = getLogger("FileContext");
 
@@ -57,6 +70,15 @@ export async function showFileContext(fileContext: FileContext, gitProvider: Git
     return;
   }
 
+  if (fileContext.filepath.startsWith("vscode-notebook-cell")) {
+    const uri = Uri.parse(fileContext.filepath);
+    const cellUri = parseVscodeNotebookCellURI(uri);
+    if (cellUri?.scheme === "untitled") {
+      showUntitledNotebookCellContext(uri);
+      return;
+    }
+  }
+
   const document = await openTextDocument(
     {
       filePath: fileContext.filepath,
@@ -67,7 +89,6 @@ export async function showFileContext(fileContext: FileContext, gitProvider: Git
   if (!document) {
     throw new Error(`File not found: ${fileContext.filepath}`);
   }
-
   const editor = await window.showTextDocument(document, {
     viewColumn: ViewColumn.Active,
     preview: false,
@@ -81,6 +102,27 @@ export async function showFileContext(fileContext: FileContext, gitProvider: Git
   editor.revealRange(new Range(start, end), TextEditorRevealType.InCenter);
 }
 
+async function showUntitledNotebookCellContext(uri: Uri) {
+  const notebookDocument = workspace.notebookDocuments.find((notebook) => {
+    return notebook.getCells().some((cell) => cell.document.uri.toString() === uri.toString());
+  });
+
+  if (notebookDocument) {
+    const notebookEditor = await window.showNotebookDocument(notebookDocument);
+    const targetCell = notebookDocument.getCells().find((cell) => cell.document.uri.toString() === uri.toString());
+    if (notebookEditor && targetCell) {
+      const cellIndex = targetCell.index
+      // FIXME(@jueliang) set selection
+      notebookEditor.revealRange(new NotebookRange(cellIndex, cellIndex), NotebookEditorRevealType.InCenter)
+
+    } else {
+      throw new Error(`Cell not found in notebook: ${uri.toString()}`);
+    }
+  } else {
+    throw new Error(`Notebook not found for URI: ${uri.toString()}`);
+  }
+}
+
 export async function buildFilePathParams(uri: Uri, gitProvider: GitProvider): Promise<FilePathParams> {
   const workspaceFolder =
     workspace.getWorkspaceFolder(uri) ?? (uri.scheme === "untitled" ? workspace.workspaceFolders?.[0] : undefined);
@@ -105,18 +147,39 @@ export async function buildFilePathParams(uri: Uri, gitProvider: GitProvider): P
   };
 }
 
+async function openTextDocumentByAbsoluteFilepath(absoluteFilepath: Uri) {
+  if (!absoluteFilepath.scheme) {
+    return null;
+  }
+
+  if (absoluteFilepath.scheme.startsWith("vscode-notebook-cell")) {
+    const cell = parseVscodeNotebookCellURI(absoluteFilepath);
+    if (cell?.scheme === "untitled") {
+      const notebookDocument = workspace.notebookDocuments.find((notebook) => {
+        return (
+          notebook.isUntitled &&
+          notebook.getCells().some((cell) => cell.document.uri.toString() === absoluteFilepath.toString())
+        );
+      });
+      if (notebookDocument && typeof cell?.handle === "number") {
+        const notebookCell = notebookDocument.cellAt(cell.handle);
+        return notebookCell.document;
+      }
+    }
+  }
+
+  return await workspace.openTextDocument(absoluteFilepath);
+}
+
 export async function openTextDocument(
   filePathParams: FilePathParams,
   gitProvider: GitProvider,
 ): Promise<TextDocument | null> {
   const { filePath, gitRemoteUrl } = filePathParams;
-
   // Try parse as absolute path
   try {
     const absoluteFilepath = Uri.parse(filePath, true);
-    if (absoluteFilepath.scheme) {
-      return await workspace.openTextDocument(absoluteFilepath);
-    }
+    return openTextDocumentByAbsoluteFilepath(absoluteFilepath);
   } catch (err) {
     // ignore
   }
@@ -127,7 +190,7 @@ export async function openTextDocument(
     if (localGitRoot) {
       try {
         const absoluteFilepath = Uri.joinPath(localGitRoot, filePath);
-        return await workspace.openTextDocument(absoluteFilepath);
+        return await openTextDocumentByAbsoluteFilepath(absoluteFilepath);
       } catch (err) {
         // ignore
       }
@@ -138,7 +201,7 @@ export async function openTextDocument(
     // Try find file in workspace folder
     const absoluteFilepath = Uri.joinPath(root.uri, filePath);
     try {
-      return await workspace.openTextDocument(absoluteFilepath);
+      return await openTextDocumentByAbsoluteFilepath(absoluteFilepath);
     } catch (err) {
       // ignore
     }
@@ -148,7 +211,7 @@ export async function openTextDocument(
     if (localGitRoot) {
       try {
         const absoluteFilepath = Uri.joinPath(localGitRoot, filePath);
-        return await workspace.openTextDocument(absoluteFilepath);
+        return await openTextDocumentByAbsoluteFilepath(absoluteFilepath);
       } catch (err) {
         // ignore
       }
@@ -160,7 +223,7 @@ export async function openTextDocument(
   const files = await workspace.findFiles(filePath, undefined, 1);
   if (files[0]) {
     try {
-      return await workspace.openTextDocument(files[0]);
+      return await openTextDocumentByAbsoluteFilepath(files[0]);
     } catch (err) {
       // ignore
     }

diff --git a/clients/vscode/src/chat/utils.ts b/clients/vscode/src/chat/utils.ts
@@ -83,6 +83,11 @@ export function chatPanelLineRangeToVSCodeRange(lineRange: LineRange): VSCodeRan
   return new VSCodeRange(Math.max(0, lineRange.start - 1), 0, lineRange.end, 0);
 }
 
+export function chatPanelNotebookCellRangeToVSCodeRange(lineRange: LineRange): VSCodeRange {
+  // Do not minus 1 from end line number, as we want to include the last line.
+  return new VSCodeRange(Math.max(0, lineRange.start - 1), 0, lineRange.end, 0);
+}
+
 export function chatPanelLocationToVSCodeRange(location: Location): VSCodeRange | null {
   if (typeof location === "number") {
     const position = new VSCodePosition(Math.max(0, location - 1), 0);
@@ -100,3 +105,27 @@ export function chatPanelLocationToVSCodeRange(location: Location): VSCodeRange
   logger.warn(`Invalid location params.`, location);
   return null;
 }
+
+export function parseVscodeNotebookCellURI(uri: Uri) {
+  if (!uri.scheme) return undefined;
+  if (!uri.scheme.startsWith("vscode-notebook-cell")) return undefined;
+
+  const _lengths = ["W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f"];
+  const _padRegexp = new RegExp(`^[${_lengths.join("")}]+`);
+  const _radix = 7;
+  const fragment = uri.fragment.split("#").pop() || "";
+  const idx = fragment.indexOf("s");
+  if (idx < 0) {
+    return undefined;
+  }
+  const handle = parseInt(fragment.substring(0, idx).replace(_padRegexp, ""), _radix);
+  const scheme = Buffer.from(fragment.substring(idx + 1), "base64").toString("utf-8");
+
+  if (isNaN(handle)) {
+    return undefined;
+  }
+  return {
+    handle,
+    scheme,
+  };
+}
diff --git a/ee/tabby-ui/components/chat/chat-panel.tsx b/ee/tabby-ui/components/chat/chat-panel.tsx
@@ -13,7 +13,7 @@ import { updateEnableActiveSelection } from '@/lib/stores/chat-actions'
 import { useChatStore } from '@/lib/stores/chat-store'
 import { useMutation } from '@/lib/tabby/gql'
 import { setThreadPersistedMutation } from '@/lib/tabby/query'
-import { cn, getTitleFromMessages } from '@/lib/utils'
+import { cn, getTitleFromMessages, resolveFileNameForDisplay } from '@/lib/utils'
 import { Badge } from '@/components/ui/badge'
 import { Button } from '@/components/ui/button'
 import {
@@ -343,15 +343,14 @@ function ContextLabel({
   context: Context
   className?: string
 }) {
-  const [fileName] = context.filepath.split('/').slice(-1)
   const line =
     context.range.start === context.range.end
       ? `:${context.range.start}`
       : `:${context.range.start}-${context.range.end}`
 
   return (
     <span className={cn('truncate', className)}>
-      {fileName}
+      {resolveFileNameForDisplay(context.filepath)}
       <span className="text-muted-foreground">{line}</span>
     </span>
   )

diff --git a/ee/tabby-ui/components/chat/code-references.tsx b/ee/tabby-ui/components/chat/code-references.tsx
@@ -2,7 +2,7 @@ import React, { forwardRef, useEffect, useState } from 'react'
 import { isNil } from 'lodash-es'
 
 import { RelevantCodeContext } from '@/lib/types'
-import { cn } from '@/lib/utils'
+import { cn, resolveFileNameForDisplay } from '@/lib/utils'
 import {
   Tooltip,
   TooltipContent,
@@ -16,6 +16,7 @@ import {
   AccordionTrigger
 } from '../ui/accordion'
 import { IconExternalLink, IconFile, IconFileSearch2 } from '../ui/icons'
+import { VSCODE_NOTEBOOK_CELL_SCHEME } from '@/lib/constants'
 
 interface ContextReferencesProps {
   isInEditor?: boolean
@@ -85,9 +86,8 @@ export const CodeReferences = forwardRef<
           <AccordionTrigger
             className={cn('my-0 py-2 font-semibold', triggerClassname)}
           >
-            <span className="mr-2">{`Read ${totalContextLength} file${
-              isMultipleReferences ? 's' : ''
-            }`}</span>
+            <span className="mr-2">{`Read ${totalContextLength} file${isMultipleReferences ? 's' : ''
+              }`}</span>
           </AccordionTrigger>
           <AccordionContent className="space-y-2">
             {clientContexts?.map((item, index) => {
@@ -150,8 +150,9 @@ function ContextItem({
     !isNil(context.range?.end) &&
     context.range.start < context.range.end
   const pathSegments = context.filepath.split('/')
-  const fileName = pathSegments[pathSegments.length - 1]
   const path = pathSegments.slice(0, pathSegments.length - 1).join('/')
+  const isVscodeNotebookCell = path.startsWith(VSCODE_NOTEBOOK_CELL_SCHEME)
+  const showPath = !!path && !isVscodeNotebookCell
   const scores = context?.extra?.scores
   const onTooltipOpenChange = (v: boolean) => {
     if (!enableTooltip || !scores) return
@@ -177,7 +178,7 @@ function ContextItem({
           <div className="flex items-center gap-1 overflow-hidden">
             <IconFile className="shrink-0" />
             <div className="flex-1 truncate" title={context.filepath}>
-              <span>{fileName}</span>
+            <span>{resolveFileNameForDisplay(context.filepath)}</span>
               {context.range?.start && (
                 <span className="text-muted-foreground">
                   :{context.range.start}
@@ -188,7 +189,11 @@ function ContextItem({
                   -{context.range.end}
                 </span>
               )}
-              <span className="ml-2 text-xs text-muted-foreground">{path}</span>
+              {showPath && (
+                <span className="ml-2 text-xs text-muted-foreground">
+                  {path}
+                </span>
+              )}
             </div>
             {showClientCodeIcon && (
               <IconFileSearch2 className="shrink-0 text-muted-foreground" />

diff --git a/ee/tabby-ui/lib/constants/index.ts b/ee/tabby-ui/lib/constants/index.ts
@@ -15,3 +15,5 @@ export * as regex from './regex'
 export const ERROR_CODE_NOT_FOUND = 'NOT_FOUND'
 
 export const NEWLINE_CHARACTER = '\n'
+
+export const VSCODE_NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'
diff --git a/ee/tabby-ui/lib/utils/chat.ts b/ee/tabby-ui/lib/utils/chat.ts
@@ -8,6 +8,7 @@ import {
 import { MentionAttributes } from '@/lib/types'
 
 import { MARKDOWN_SOURCE_REGEX } from '../constants/regex'
+import { VSCODE_NOTEBOOK_CELL_SCHEME } from '../constants'
 
 export const isCodeSourceContext = (kind: ContextSourceKind) => {
   return [
@@ -105,3 +106,47 @@ export function checkSourcesAvailability(
 
   return { hasCodebaseSource, hasDocumentSource }
 }
+
+function parseVscodeNotebookCellURI(uri: string) {
+  if (!uri.startsWith(VSCODE_NOTEBOOK_CELL_SCHEME)) return undefined
+
+  const _lengths = ['W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f']
+  const _padRegexp = new RegExp(`^[${_lengths.join('')}]+`)
+  const _radix = 7
+  const fragment = uri.split('#').pop() || ''
+  const idx = fragment.indexOf('s')
+  if (idx < 0) {
+    return undefined
+  }
+  const handle = parseInt(
+    fragment.substring(0, idx).replace(_padRegexp, ''),
+    _radix
+  )
+  const scheme = Buffer.from(fragment.substring(idx + 1), 'base64').toString(
+    'utf-8'
+  )
+
+  if (isNaN(handle)) {
+    return undefined
+  }
+  return {
+    handle,
+    scheme
+  }
+}
+
+export function resolveFileNameForDisplay(uri: string) {
+  const regexPattern = `(?:${VSCODE_NOTEBOOK_CELL_SCHEME}:)?(.*?)(\\?|#|$)`
+  const regex = new RegExp(regexPattern)
+
+  const pathSegments = uri.split('/')
+  const fileName = pathSegments[pathSegments.length - 1]
+  const match = fileName.match(regex)
+  const displayName = match ? match[1] : null
+  const notebook = parseVscodeNotebookCellURI(uri)
+
+  if (displayName && notebook) {
+    return `${displayName} · Cell ${(notebook.handle || 0) + 1}`
+  }
+  return displayName
+}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -15,3 +15,5 @@ export * as regex from './regex'
		export const ERROR_CODE_NOT_FOUND = 'NOT_FOUND'

		export const NEWLINE_CHARACTER = '\n'

		export const VSCODE_NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'