fix: :bug: ssh compatibility by reading from vscode.workspace.fs

9 files changed, 151 insertions, 196 deletions

diff --git a/continuedev/src/continuedev/plugins/context_providers/file.py b/continuedev/src/continuedev/plugins/context_providers/file.py
index a748379e..9846dd3e 100644
--- a/continuedev/src/continuedev/plugins/context_providers/file.py
+++ b/continuedev/src/continuedev/plugins/context_providers/file.py
@@ -1,25 +1,20 @@
 import asyncio
 import os
-from fnmatch import fnmatch
 from typing import List
 
 from ...core.context import ContextProvider
 from ...core.main import ContextItem, ContextItemDescription, ContextItemId
+from ...core.sdk import ContinueSDK
 from .util import remove_meilisearch_disallowed_chars
 
-MAX_SIZE_IN_BYTES = 1024 * 1024 * 1
+MAX_SIZE_IN_CHARS = 25_000
 
 
-def get_file_contents(filepath: str) -> str:
+async def get_file_contents(filepath: str, sdk: ContinueSDK) -> str:
     try:
-        filesize = os.path.getsize(filepath)
-        if filesize > MAX_SIZE_IN_BYTES:
-            return None
-
-        with open(filepath, "r") as f:
-            return f.read()
-    except Exception:
-        # Some files cannot be read, e.g. binary files
+        return (await sdk.ide.readFile(filepath))[:MAX_SIZE_IN_CHARS]
+    except Exception as e:
+        print(f"Failed to read file {filepath}: {e}")
         return None
 
 
@@ -105,7 +100,7 @@ class FileContextProvider(ContextProvider):
     async def get_context_item_for_filepath(
         self, absolute_filepath: str
     ) -> ContextItem:
-        content = get_file_contents(absolute_filepath)
+        content = await get_file_contents(absolute_filepath, self.sdk)
         if content is None:
             return None
 
@@ -128,26 +123,35 @@ class FileContextProvider(ContextProvider):
         )
 
     async def provide_context_items(self, workspace_dir: str) -> List[ContextItem]:
+        contents = await self.sdk.ide.listDirectoryContents(workspace_dir)
+        if contents is None:
+            return []
+
         absolute_filepaths: List[str] = []
-        for root, dir_names, file_names in os.walk(workspace_dir):
-            dir_names[:] = [
-                d
-                for d in dir_names
-                if not any(fnmatch(d, pattern) for pattern in self.ignore_patterns)
+        for filepath in contents[:1000]:
+            absolute_filepaths.append(filepath)
+
+        # for root, dir_names, file_names in os.walk(workspace_dir):
+        #     dir_names[:] = [
+        #         d
+        #         for d in dir_names
+        #         if not any(fnmatch(d, pattern) for pattern in self.ignore_patterns)
+        #     ]
+        #     for file_name in file_names:
+        #         absolute_filepaths.append(os.path.join(root, file_name))
+
+        #         if len(absolute_filepaths) > 1000:
+        #             break
+
+        #     if len(absolute_filepaths) > 1000:
+        #         break
+
+        items = await asyncio.gather(
+            *[
+                self.get_context_item_for_filepath(filepath)
+                for filepath in absolute_filepaths
             ]
-            for file_name in file_names:
-                absolute_filepaths.append(os.path.join(root, file_name))
-
-                if len(absolute_filepaths) > 1000:
-                    break
-
-            if len(absolute_filepaths) > 1000:
-                break
-
-        items = []
-        for absolute_filepath in absolute_filepaths:
-            item = await self.get_context_item_for_filepath(absolute_filepath)
-            if item is not None:
-                items.append(item)
+        )
+        items = list(filter(lambda item: item is not None, items))
 
         return items
diff --git a/continuedev/src/continuedev/server/ide.py b/continuedev/src/continuedev/server/ide.py
index 610a1a48..871724db 100644
--- a/continuedev/src/continuedev/server/ide.py
+++ b/continuedev/src/continuedev/server/ide.py
@@ -120,6 +120,10 @@ class TerminalContentsResponse(BaseModel):
     contents: str
 
 
+class ListDirectoryContentsResponse(BaseModel):
+    contents: List[str]
+
+
 T = TypeVar("T", bound=BaseModel)
 
 
@@ -241,6 +245,7 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
             "getUserSecret",
             "runCommand",
             "getTerminalContents",
+            "listDirectoryContents",
         ]:
             self.sub_queue.post(message_type, data)
         elif message_type == "workspaceDirectory":
@@ -477,6 +482,15 @@ class IdeProtocolServer(AbstractIdeProtocolServer):
         )
         return resp.fileEdit
 
+    async def listDirectoryContents(self, directory: str) -> List[str]:
+        """List the contents of a directory"""
+        resp = await self._send_and_receive_json(
+            {"directory": directory},
+            ListDirectoryContentsResponse,
+            "listDirectoryContents",
+        )
+        return resp.contents
+
     async def applyFileSystemEdit(self, edit: FileSystemEdit) -> EditDiff:
         """Apply a file edit"""
         backward = None
diff --git a/continuedev/src/continuedev/server/ide_protocol.py b/continuedev/src/continuedev/server/ide_protocol.py
index 435c82e2..f37c1737 100644
--- a/continuedev/src/continuedev/server/ide_protocol.py
+++ b/continuedev/src/continuedev/server/ide_protocol.py
@@ -147,5 +147,9 @@ class AbstractIdeProtocolServer(ABC):
     def onFileSaved(self, filepath: str, contents: str):
         """Called when a file is saved"""
 
+    @abstractmethod
+    async def listDirectoryContents(self, directory: str) -> List[str]:
+        """List directory contents"""
+
     workspace_directory: str
     unique_id: str
diff --git a/docs/docs/customization.md b/docs/docs/customization.md
index 3240c185..a1a9111e 100644
--- a/docs/docs/customization.md
+++ b/docs/docs/customization.md
@@ -156,7 +156,7 @@ config = ContinueConfig(
         default=OpenAI(api_key="my-api-key", model="gpt-3.5-turbo", openai_server_info=OpenAIServerInfo(
             api_base="https://my-azure-openai-instance.openai.azure.com/",
             engine="my-azure-openai-deployment",
-            api_version="2023-03-15-preview",
+            api_version="2023-03-15-preview",  # NOTE: It is recommended not to change api_version. Newer versions may not work correctly.
             api_type="azure"
         ))
     )
diff --git a/extension/react-app/src/pages/settings.tsx b/extension/react-app/src/pages/settings.tsx
index 8fd91ff5..8bb9b5d2 100644
--- a/extension/react-app/src/pages/settings.tsx
+++ b/extension/react-app/src/pages/settings.tsx
@@ -94,7 +94,7 @@ function Settings() {
     const temperature = formMethods.watch("temperature");
     // const models = formMethods.watch("models");
 
-    if (systemMessage) client.setSystemMessage(systemMessage);
+    client.setSystemMessage(systemMessage || "");
     if (temperature) client.setTemperature(temperature);
 
     // if (models) {
diff --git a/extension/src/continueIdeClient.ts b/extension/src/continueIdeClient.ts
index 430bb9dd..19575b13 100644
--- a/extension/src/continueIdeClient.ts
+++ b/extension/src/continueIdeClient.ts
@@ -12,10 +12,10 @@ import {
   rejectSuggestionCommand,
 } from "./suggestions";
 import { FileEditWithFullContents } from "../schema/FileEditWithFullContents";
-import * as fs from "fs";
 import { WebsocketMessenger } from "./util/messenger";
 import { diffManager } from "./diffs";
 const os = require("os");
+const path = require("path");
 
 const continueVirtualDocumentScheme = "continue";
 
@@ -253,7 +253,7 @@ class IdeProtocolClient {
         break;
       case "readFile":
         messenger.send("readFile", {
-          contents: this.readFile(data.filepath),
+          contents: await this.readFile(data.filepath),
         });
         break;
       case "getTerminalContents":
@@ -261,6 +261,44 @@ class IdeProtocolClient {
           contents: await this.getTerminalContents(),
         });
         break;
+      case "listDirectoryContents":
+        messenger.send("listDirectoryContents", {
+          contents: (
+            await vscode.workspace.fs.readDirectory(
+              vscode.Uri.file(data.directory)
+            )
+          )
+            .map(([name, type]) => name)
+            .filter((name) => {
+              const DEFAULT_IGNORE_DIRS = [
+                ".git",
+                ".vscode",
+                ".idea",
+                ".vs",
+                ".venv",
+                "env",
+                ".env",
+                "node_modules",
+                "dist",
+                "build",
+                "target",
+                "out",
+                "bin",
+                ".pytest_cache",
+                ".vscode-test",
+                ".continue",
+                "__pycache__",
+              ];
+              if (
+                !DEFAULT_IGNORE_DIRS.some((dir) =>
+                  name.split(path.sep).includes(dir)
+                )
+              ) {
+                return name;
+              }
+            }),
+        });
+        break;
       case "editFile":
         const fileEdit = await this.editFile(data.edit);
         messenger.send("editFile", {
@@ -306,7 +344,7 @@ class IdeProtocolClient {
         this.showSuggestion(data.edit);
         break;
       case "showDiff":
-        this.showDiff(data.filepath, data.replacement, data.step_index);
+        await this.showDiff(data.filepath, data.replacement, data.step_index);
         break;
       case "getSessionId":
       case "connected":
@@ -385,8 +423,8 @@ class IdeProtocolClient {
     );
   }
 
-  showDiff(filepath: string, replacement: string, step_index: number) {
-    diffManager.writeDiff(filepath, replacement, step_index);
+  async showDiff(filepath: string, replacement: string, step_index: number) {
+    await diffManager.writeDiff(filepath, replacement, step_index);
   }
 
   openFile(filepath: string) {
@@ -506,19 +544,14 @@ class IdeProtocolClient {
       });
   }
 
-  readFile(filepath: string): string {
+  async readFile(filepath: string): Promise<string> {
     let contents: string | undefined;
-    vscode.window.visibleTextEditors
-      .filter((editor) => this.editorIsCode(editor))
-      .forEach((editor) => {
-        if (editor.document.uri.fsPath === filepath) {
-          contents = editor.document.getText();
-        }
-      });
     if (typeof contents === "undefined") {
-      if (fs.existsSync(filepath)) {
-        contents = fs.readFileSync(filepath, "utf-8");
-      } else {
+      try {
+        contents = await vscode.workspace.fs
+          .readFile(vscode.Uri.file(filepath))
+          .then((bytes) => new TextDecoder().decode(bytes));
+      } catch {
         contents = "";
       }
     }
diff --git a/extension/src/diffs.ts b/extension/src/diffs.ts
index efaf7626..d2d1dae3 100644
--- a/extension/src/diffs.ts
+++ b/extension/src/diffs.ts
@@ -14,6 +14,19 @@ interface DiffInfo {
   range: vscode.Range;
 }
 
+async function readFile(path: string): Promise<string> {
+  return await vscode.workspace.fs
+    .readFile(vscode.Uri.file(path))
+    .then((bytes) => new TextDecoder().decode(bytes));
+}
+
+async function writeFile(path: string, contents: string) {
+  await vscode.workspace.fs.writeFile(
+    vscode.Uri.file(path),
+    new TextEncoder().encode(contents)
+  );
+}
+
 export const DIFF_DIRECTORY = path
   .join(os.homedir(), ".continue", "diffs")
   .replace(/^C:/, "c:");
@@ -27,13 +40,9 @@ class DiffManager {
     return this.diffs.get(newFilepath);
   }
 
-  private setupDirectory() {
+  private async setupDirectory() {
     // Make sure the diff directory exists
-    if (!fs.existsSync(DIFF_DIRECTORY)) {
-      fs.mkdirSync(DIFF_DIRECTORY, {
-        recursive: true,
-      });
-    }
+    await vscode.workspace.fs.createDirectory(vscode.Uri.file(DIFF_DIRECTORY));
   }
 
   constructor() {
@@ -57,15 +66,17 @@ class DiffManager {
     return path.join(DIFF_DIRECTORY, this.escapeFilepath(originalFilepath));
   }
 
-  private openDiffEditor(
+  private async openDiffEditor(
     originalFilepath: string,
     newFilepath: string
-  ): vscode.TextEditor | undefined {
-    // If the file doesn't yet exist or the basename is a single digit number (git hash object or something), don't open the diff editor
-    if (
-      !fs.existsSync(newFilepath) ||
-      path.basename(originalFilepath).match(/^\d$/)
-    ) {
+  ): Promise<vscode.TextEditor | undefined> {
+    // If the file doesn't yet exist or the basename is a single digit number (vscode terminal), don't open the diff editor
+    try {
+      await vscode.workspace.fs.stat(vscode.Uri.file(newFilepath));
+    } catch {
+      return undefined;
+    }
+    if (path.basename(originalFilepath).match(/^\d$/)) {
       return undefined;
     }
 
@@ -128,21 +139,21 @@ class DiffManager {
     return 0;
   }
 
-  writeDiff(
+  async writeDiff(
     originalFilepath: string,
     newContent: string,
     step_index: number
-  ): string {
-    this.setupDirectory();
+  ): Promise<string> {
+    await this.setupDirectory();
 
     // Create or update existing diff
     const newFilepath = this.getNewFilepath(originalFilepath);
-    fs.writeFileSync(newFilepath, newContent);
+    await writeFile(newFilepath, newContent);
 
     // Open the diff editor if this is a new diff
     if (!this.diffs.has(newFilepath)) {
       // Figure out the first line that is different
-      const oldContent = ideProtocolClient.readFile(originalFilepath);
+      const oldContent = await ideProtocolClient.readFile(originalFilepath);
       const line = this._findFirstDifferentLine(oldContent, newContent);
 
       const diffInfo: DiffInfo = {
@@ -157,7 +168,10 @@ class DiffManager {
     // Open the editor if it hasn't been opened yet
     const diffInfo = this.diffs.get(newFilepath);
     if (diffInfo && !diffInfo?.editor) {
-      diffInfo.editor = this.openDiffEditor(originalFilepath, newFilepath);
+      diffInfo.editor = await this.openDiffEditor(
+        originalFilepath,
+        newFilepath
+      );
       this.diffs.set(newFilepath, diffInfo);
     }
 
@@ -207,7 +221,7 @@ class DiffManager {
     return undefined;
   }
 
-  acceptDiff(newFilepath?: string) {
+  async acceptDiff(newFilepath?: string) {
     // When coming from a keyboard shortcut, we have to infer the newFilepath from visible text editors
     if (!newFilepath) {
       newFilepath = this.inferNewFilepath();
@@ -227,18 +241,18 @@ class DiffManager {
     vscode.workspace.textDocuments
       .find((doc) => doc.uri.fsPath === newFilepath)
       ?.save()
-      .then(() => {
-        fs.writeFileSync(
+      .then(async () => {
+        await writeFile(
           diffInfo.originalFilepath,
-          fs.readFileSync(diffInfo.newFilepath)
+          await readFile(diffInfo.newFilepath)
         );
         this.cleanUpDiff(diffInfo);
       });
 
-    recordAcceptReject(true, diffInfo);
+    await recordAcceptReject(true, diffInfo);
   }
 
-  rejectDiff(newFilepath?: string) {
+  async rejectDiff(newFilepath?: string) {
     // If no newFilepath is provided and there is only one in the dictionary, use that
     if (!newFilepath) {
       newFilepath = this.inferNewFilepath();
@@ -266,13 +280,13 @@ class DiffManager {
         this.cleanUpDiff(diffInfo);
       });
 
-    recordAcceptReject(false, diffInfo);
+    await recordAcceptReject(false, diffInfo);
   }
 }
 
 export const diffManager = new DiffManager();
 
-function recordAcceptReject(accepted: boolean, diffInfo: DiffInfo) {
+async function recordAcceptReject(accepted: boolean, diffInfo: DiffInfo) {
   const devDataDir = devDataPath();
   const suggestionsPath = path.join(devDataDir, "suggestions.json");
 
@@ -280,10 +294,10 @@ function recordAcceptReject(accepted: boolean, diffInfo: DiffInfo) {
   let suggestions = [];
 
   // Check if suggestions.json exists
-  if (fs.existsSync(suggestionsPath)) {
-    const rawData = fs.readFileSync(suggestionsPath, "utf-8");
+  try {
+    const rawData = await readFile(suggestionsPath);
     suggestions = JSON.parse(rawData);
-  }
+  } catch {}
 
   // Add the new suggestion to the list
   suggestions.push({
@@ -296,19 +310,15 @@ function recordAcceptReject(accepted: boolean, diffInfo: DiffInfo) {
   // ideProtocolClient.sendAcceptRejectSuggestion(accepted);
 
   // Write the updated suggestions back to the file
-  fs.writeFileSync(
-    suggestionsPath,
-    JSON.stringify(suggestions, null, 4),
-    "utf-8"
-  );
+  await writeFile(suggestionsPath, JSON.stringify(suggestions, null, 4));
 }
 
 export async function acceptDiffCommand(newFilepath?: string) {
-  diffManager.acceptDiff(newFilepath);
+  await diffManager.acceptDiff(newFilepath);
   ideProtocolClient.sendAcceptRejectDiff(true);
 }
 
 export async function rejectDiffCommand(newFilepath?: string) {
-  diffManager.rejectDiff(newFilepath);
+  await diffManager.rejectDiff(newFilepath);
   ideProtocolClient.sendAcceptRejectDiff(false);
 }
diff --git a/extension/src/util/util.ts b/extension/src/util/util.ts
index 15b34267..38c955e7 100644
--- a/extension/src/util/util.ts
+++ b/extension/src/util/util.ts
@@ -1,5 +1,3 @@
-import { RangeInFile } from "../../schema/RangeInFile";
-import * as fs from "fs";
 const os = require("os");
 
 function charIsEscapedAtIndex(index: number, str: string): boolean {
@@ -52,49 +50,6 @@ export function convertSingleToDoubleQuoteJSON(json: string): string {
   return newJson;
 }
 
-export async function readRangeInFile(
-  rangeInFile: RangeInFile
-): Promise<string> {
-  const range = rangeInFile.range;
-  return new Promise((resolve, reject) => {
-    fs.readFile(rangeInFile.filepath, (err, data) => {
-      if (err) {
-        reject(err);
-      } else {
-        let lines = data.toString().split("\n");
-        if (range.start.line === range.end.line) {
-          resolve(
-            lines[rangeInFile.range.start.line].slice(
-              rangeInFile.range.start.character,
-              rangeInFile.range.end.character
-            )
-          );
-        } else {
-          let firstLine = lines[range.start.line].slice(range.start.character);
-          let lastLine = lines[range.end.line].slice(0, range.end.character);
-          let middleLines = lines.slice(range.start.line + 1, range.end.line);
-          resolve([firstLine, ...middleLines, lastLine].join("\n"));
-        }
-      }
-    });
-  });
-}
-
-export function codeSelectionsToVirtualFileSystem(
-  codeSelections: RangeInFile[]
-): {
-  [filepath: string]: string;
-} {
-  let virtualFileSystem: { [filepath: string]: string } = {};
-  for (let cs of codeSelections) {
-    if (!cs.filepath) continue;
-    if (cs.filepath in virtualFileSystem) continue;
-    let content = fs.readFileSync(cs.filepath, "utf8");
-    virtualFileSystem[cs.filepath] = content;
-  }
-  return virtualFileSystem;
-}
-
 export function debounced(delay: number, fn: Function) {
   let timerId: NodeJS.Timeout | null;
   return function (...args: any[]) {
diff --git a/extension/src/util/vscode.ts b/extension/src/util/vscode.ts
index 3110d589..bf0fa1e5 100644
--- a/extension/src/util/vscode.ts
+++ b/extension/src/util/vscode.ts
@@ -1,6 +1,4 @@
 import * as vscode from "vscode";
-import * as path from "path";
-import * as fs from "fs";
 
 export function translate(range: vscode.Range, lines: number): vscode.Range {
   return new vscode.Range(
@@ -21,39 +19,6 @@ export function getNonce() {
   return text;
 }
 
-export function getTestFile(
-  filename: string,
-  createFile: boolean = false
-): string {
-  let basename = path.basename(filename).split(".")[0];
-  switch (path.extname(filename)) {
-    case ".py":
-      basename += "_test";
-      break;
-    case ".js":
-    case ".jsx":
-    case ".ts":
-    case ".tsx":
-      basename += ".test";
-      break;
-    default:
-      basename += "_test";
-  }
-
-  const directory = path.join(path.dirname(filename), "tests");
-  const testFilename = path.join(directory, basename + path.extname(filename));
-
-  // Optionally, create the file if it doesn't exist
-  if (createFile && !fs.existsSync(testFilename)) {
-    if (!fs.existsSync(directory)) {
-      fs.mkdirSync(directory);
-    }
-    fs.writeFileSync(testFilename, "");
-  }
-
-  return testFilename;
-}
-
 export function getExtensionUri(): vscode.Uri {
   return vscode.extensions.getExtension("Continue.continue")!.extensionUri;
 }
@@ -100,36 +65,6 @@ export function getRightViewColumn(): vscode.ViewColumn {
   return column;
 }
 
-export async function readFileAtRange(
-  range: vscode.Range,
-  filepath: string
-): Promise<string> {
-  return new Promise((resolve, reject) => {
-    fs.readFile(filepath, (err, data) => {
-      if (err) {
-        reject(err);
-      } else {
-        let lines = data.toString().split("\n");
-        if (range.isSingleLine) {
-          resolve(
-            lines[range.start.line].slice(
-              range.start.character,
-              range.end.character
-            )
-          );
-        } else {
-          const firstLine = lines[range.start.line].slice(
-            range.start.character
-          );
-          const lastLine = lines[range.end.line].slice(0, range.end.character);
-          const middleLines = lines.slice(range.start.line + 1, range.end.line);
-          resolve([firstLine, ...middleLines, lastLine].join("\n"));
-        }
-      }
-    });
-  });
-}
-
 let showTextDocumentInProcess = false;
 
 export function openEditorAndRevealRange(