feat(vscode): support attach file context to inline edit (#3766)

* feat(vscode): support attach file context to inline edit

* refact: use separate file pick widget

clients/tabby-agent/src/chat/inlineEdit.ts

@@ -1,5 +1,5 @@
-import type { Connection, CancellationToken } from "vscode-languageserver";
-import type { TextDocument } from "vscode-languageserver-textdocument";
+import type { Connection, CancellationToken, Range, URI } from "vscode-languageserver";
+import { TextDocument } from "vscode-languageserver-textdocument";
 import type { TextDocuments } from "../lsp/textDocuments";
 import type { Feature } from "../feature";
 import type { Configurations } from "../config";
@@ -17,14 +17,22 @@ import {
   ChatEditMutexError,
   ServerCapabilities,
   ChatEditResolveParams,
+  ClientCapabilities,
+  ReadFileParams,
+  ReadFileRequest,
 } from "../protocol";
 import cryptoRandomString from "crypto-random-string";
 import { isEmptyRange } from "../utils/range";
-import { readResponseStream, Edit, applyWorkspaceEdit } from "./utils";
+import { readResponseStream, Edit, applyWorkspaceEdit, truncateFileContent } from "./utils";
 import { initMutexAbortController, mutexAbortController, resetMutexAbortController } from "./global";
+import { readFile } from "fs-extra";
+import { getLogger } from "../logger";
+import { isBrowser } from "../env";
 
 export class ChatEditProvider implements Feature {
+  private logger = getLogger("ChatEditProvider");
   private lspConnection: Connection | undefined = undefined;
+  private clientCapabilities: ClientCapabilities | undefined = undefined;
   private currentEdit: Edit | undefined = undefined;
 
   constructor(
@@ -33,8 +41,9 @@ export class ChatEditProvider implements Feature {
     private readonly documents: TextDocuments<TextDocument>,
   ) {}
 
-  initialize(connection: Connection): ServerCapabilities {
+  initialize(connection: Connection, clientCapabilities: ClientCapabilities): ServerCapabilities {
     this.lspConnection = connection;
+    this.clientCapabilities = clientCapabilities;
     connection.onRequest(ChatEditCommandRequest.type, async (params) => {
       return this.provideEditCommands(params);
     });
@@ -92,6 +101,44 @@ export class ChatEditProvider implements Feature {
     return result;
   }
 
+  async fetchFileContent(uri: URI, range?: Range, token?: CancellationToken) {
+    this.logger.trace("Prepare to fetch text content...");
+    let text: string | undefined = undefined;
+    const targetDocument = this.documents.get(uri);
+    if (targetDocument) {
+      this.logger.trace("Fetching text content from synced text document.", {
+        uri: targetDocument.uri,
+        range: range,
+      });
+      text = targetDocument.getText(range);
+      this.logger.trace("Fetched text content from synced text document.", { text });
+    } else if (this.clientCapabilities?.tabby?.workspaceFileSystem) {
+      const params: ReadFileParams = {
+        uri: uri,
+        format: "text",
+        range: range
+          ? {
+              start: { line: range.start.line, character: 0 },
+              end: { line: range.end.line, character: range.end.character },
+            }
+          : undefined,
+      };
+      this.logger.trace("Fetching text content from ReadFileRequest.", { params });
+      const result = await this.lspConnection?.sendRequest(ReadFileRequest.type, params, token);
+      this.logger.trace("Fetched text content from ReadFileRequest.", { result });
+      text = result?.text;
+    } else if (!isBrowser) {
+      try {
+        const content = await readFile(uri, "utf-8");
+        const textDocument = TextDocument.create(uri, "text", 0, content);
+        text = textDocument.getText(range);
+      } catch (error) {
+        this.logger.trace("Failed to fetch text content from file system.", { error });
+      }
+    }
+    return text;
+  }
+
   async provideEdit(params: ChatEditParams, token: CancellationToken): Promise<ChatEditToken | null> {
     if (params.format !== "previewChanges") {
       return null;
@@ -130,7 +177,7 @@ export class ChatEditProvider implements Feature {
     if (mutexAbortController && !mutexAbortController.signal.aborted) {
       throw {
         name: "ChatEditMutexError",
-        message: "Another smart edit is already in progress",
+        message: "Another chat edit is already in progress",
       } as ChatEditMutexError;
     }
 
@@ -170,6 +217,24 @@ export class ChatEditProvider implements Feature {
       }
     }
 
+    const fileContext =
+      (
+        await Promise.all(
+          (params.context ?? []).slice(0, config.chat.edit.fileContext.maxFiles).map(async (item) => {
+            const content = await this.fetchFileContent(item.uri, item.range, token);
+            if (!content) {
+              return undefined;
+            }
+            const fileContent = truncateFileContent(content, config.chat.edit.fileContext.maxCharsPerFile);
+            return `File "${item.uri}" referer to "@${item.referer}" in command, content:\n ${fileContent} \n`;
+          }),
+        )
+      )
+        .filter((item): item is string => item !== undefined)
+        .join("\n") ?? "";
+
+    this.logger.debug(`fileContext: ${fileContext}`);
+
     const messages: { role: "user"; content: string }[] = [
       {
         role: "user",
@@ -189,13 +254,16 @@ export class ChatEditProvider implements Feature {
                 return userCommand;
               case "{{languageId}}":
                 return document.languageId;
+              case "{{fileContext}}":
+                return fileContext;
               default:
                 return "";
             }
           },
         ),
       },
     ];
+    this.logger.debug(`messages: ${JSON.stringify(messages)}`);
     const readableStream = await this.tabbyApiClient.fetchChatStream(
       {
         messages,

clients/tabby-agent/src/chat/utils.ts

@@ -362,3 +362,17 @@ export function getCommentPrefix(languageId: string) {
   }
   return "";
 }
+
+export function truncateFileContent(content: string, maxLength: number): string {
+  if (content.length <= maxLength) {
+    return content;
+  }
+
+  content = content.slice(0, maxLength);
+  const lastNewLine = content.lastIndexOf("\n");
+  if (lastNewLine > 0) {
+    content = content.slice(0, lastNewLine + 1);
+  }
+
+  return content;
+}

clients/tabby-agent/src/config/default.ts

@@ -68,6 +68,10 @@ export const defaultConfigData: ConfigData = {
     edit: {
       documentMaxChars: 3000,
       commandMaxChars: 200,
+      fileContext: {
+        maxFiles: 20,
+        maxCharsPerFile: 3000,
+      },
       responseDocumentTag: ["<GENERATEDCODE>", "</GENERATEDCODE>"],
       responseCommentTag: undefined,
       promptTemplate: {

clients/tabby-agent/src/config/type.d.ts

@@ -78,6 +78,10 @@ export type ConfigData = {
     edit: {
       documentMaxChars: number;
       commandMaxChars: number;
+      fileContext: {
+        maxFiles: number;
+        maxCharsPerFile: number;
+      };
       responseDocumentTag: string[];
       responseCommentTag: string[] | undefined;
       promptTemplate: {

clients/tabby-agent/src/protocol.ts

@@ -442,8 +442,34 @@ export type ChatEditParams = {
    *    use {@link ChatEditResolveRequest} to resolve it later.
    */
   format: "previewChanges";
+
+  /**
+   * list of file contexts.
+   */
+  context?: ChatEditFileContext[];
 };
 
+/**
+ * Represents a file context use in {@link ChatEditParams}.
+ */
+export interface ChatEditFileContext {
+  /**
+   * The symbol in the user command that refer to this file context.
+   */
+  referer: string;
+
+  /**
+   * The uri of the file.
+   */
+  uri: URI;
+
+  /**
+   * The context range in the file.
+   * If the range is not provided, the whole file is considered.
+   */
+  range?: Range;
+}
+
 export type ChatEditToken = string;
 
 export type ChatFeatureNotAvailableError = {

clients/vscode/.mocha.env.js

@@ -0,0 +1,2 @@
+process.env.NODE_ENV = "test";
+process.env.IS_TEST = 1;

clients/vscode/.mocharc.js

@@ -0,0 +1,4 @@
+module.exports = {
+  spec: ["src/**/*.test.ts"],
+  require: ["ts-node/register", "./.mocha.env.js"],
+};

clients/vscode/package.json

@@ -454,10 +454,12 @@
     "ovsx:publish-prerelease": "ovsx publish --no-dependencies --pre-release --pat $OVSX_PAT",
     "package": "pnpm run vscode:package",
     "publish": "pnpm run vscode:publish && pnpm run ovsx:publish",
-    "publish-prerelease": "pnpm run vscode:publish-prerelease && pnpm run ovsx:publish-prerelease"
+    "publish-prerelease": "pnpm run vscode:publish-prerelease && pnpm run ovsx:publish-prerelease",
+    "test": "mocha"
   },
   "devDependencies": {
     "@types/dedent": "^0.7.2",
+    "@types/chai": "^4.3.5",
     "@types/deep-equal": "^1.0.4",
     "@types/diff": "^5.2.1",
     "@types/fs-extra": "^11.0.1",
@@ -474,6 +476,7 @@
     "@vscode/vsce": "^2.15.0",
     "assert": "^2.0.0",
     "debounce": "^2.2.0",
+    "chai": "^4.3.11",
     "dedent": "^0.7.0",
     "deep-equal": "^2.2.1",
     "diff": "^5.2.0",
@@ -483,6 +486,7 @@
     "eslint-config-prettier": "^9.0.0",
     "fs-extra": "^11.1.1",
     "get-installed-path": "^4.0.8",
+    "mocha": "^10.2.0",
     "object-hash": "^3.0.0",
     "ovsx": "^0.9.5",
     "prettier": "^3.0.0",

clients/vscode/src/chat/utils.ts

@@ -172,11 +172,6 @@ export function localUriToListFileItem(uri: Uri, gitProvider: GitProvider): List
   };
 }
 
-export function escapeGlobPattern(query: string): string {
-  // escape special glob characters: * ? [ ] { } ( ) ! @
-  return query.replace(/[*?[\]{}()!@]/g, "\\$&");
-}
-
 // Notebook cell uri conversion
 
 function isJupyterNotebookFilepath(uri: Uri): boolean {

clients/vscode/src/chat/webview.ts

@@ -54,10 +54,9 @@ import {
   chatPanelLocationToVSCodeRange,
   isValidForSyncActiveEditorSelection,
   localUriToListFileItem,
-  escapeGlobPattern,
   vscodeRangeToChatPanelLineRange,
 } from "./utils";
-import { findFiles } from "../findFiles";
+import { caseInsensitivePattern, findFiles } from "../findFiles";
 import { wrapCancelableFunction } from "../cancelableFunction";
 import mainHtml from "./html/main.html";
 import errorHtml from "./html/error.html";
@@ -526,17 +525,7 @@ export class ChatWebview {
         }
 
         try {
-          const caseInsensitivePattern = searchQuery
-            .split("")
-            .map((char) => {
-              if (char.toLowerCase() !== char.toUpperCase()) {
-                return `{${char.toLowerCase()},${char.toUpperCase()}}`;
-              }
-              return escapeGlobPattern(char);
-            })
-            .join("");
-
-          const globPattern = `**/${caseInsensitivePattern}*`;
+          const globPattern = caseInsensitivePattern(searchQuery);
           this.logger.info(`Searching files with pattern: ${globPattern}, limit: ${maxResults}`);
           const files = await this.findFiles(globPattern, { maxResults });
           this.logger.info(`Found ${files.length} files.`);

clients/vscode/src/findFiles.ts

@@ -235,3 +235,22 @@ export async function findFiles(
     });
   }
 }
+
+export function escapeGlobPattern(query: string): string {
+  // escape special glob characters: * ? [ ] { } ( ) ! @
+  return query.replace(/[*?[\]{}()!@]/g, "\\$&");
+}
+
+export function caseInsensitivePattern(query: string) {
+  const caseInsensitivePattern = query
+    .split("")
+    .map((char) => {
+      if (char.toLowerCase() !== char.toUpperCase()) {
+        return `{${char.toLowerCase()},${char.toUpperCase()}}`;
+      }
+      return escapeGlobPattern(char);
+    })
+    .join("");
+
+  return `**/${caseInsensitivePattern}*`;
+}