Bring to head of internal repo (#66)

* Project import generated by Copybara. GitOrigin-RevId: f6fda22160d88ef816ad7349aa5d3b5a1075ce42 * update Document proto --------- Co-authored-by: Copybara Bot <[email protected]>
Exafunction · Oct 3, 2024 · 4ee5013 · 4ee5013
1 parent 2f9c8bc
commit 4ee5013
Show file tree

Hide file tree

Showing 15 changed files with 198 additions and 67 deletions.
diff --git a/exa/language_server_pb/language_server.proto b/exa/language_server_pb/language_server.proto
@@ -63,9 +63,13 @@ message DocumentPosition {
 
 // Next ID: 9, Previous field: cursor_position.
 message Document {
-  string absolute_path = 1;
-  // Path relative to the root of the workspace.
-  string relative_path = 2;
+  // OS specific separators.
+  string absolute_path_migrate_me_to_uri = 1 [deprecated = true];
+  string absolute_uri = 12;
+  // Path relative to the root of the workspace. Slash separated.
+  // Leave empty if the document is not in the workspace.
+  string relative_path_migrate_me_to_workspace_uri = 2 [deprecated = true];
+  string workspace_uri = 13;
   string text = 3;
   // Language ID provided by the editor.
   string editor_language = 4 [(validate.rules).string.min_len = 1];
@@ -83,6 +87,13 @@ message Document {
       "\r\n"
     ]
   }];
+
+  // These fields are not used by the chrome extension.
+  Range visible_range = 9;
+  bool is_cutoff_start = 10;
+  bool is_cutoff_end = 11;
+  int32 lines_cutoff_start = 14;
+  int32 lines_cutoff_end = 15;
 }
 
 message ExperimentConfig {

diff --git a/generate.js b/generate.js
@@ -3,7 +3,7 @@ const { execSync } = require('child_process');
 require('dotenv').config();
 if (process.env.CODEIUM_ENV === 'monorepo') {
   execSync(
-    'pnpm buf generate ../../.. --path ../../language_server_pb/language_server.proto --path ../../seat_management_pb/seat_management.proto --include-imports'
+    'pnpm buf generate ../../.. --path ../../language_server_pb/language_server.proto --path ../../seat_management_pb/seat_management.proto --path ../../opensearch_clients_pb/opensearch_clients.proto --include-imports'
   );
 } else {
   execSync('pnpm buf generate');

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "codeium-chrome",
-  "version": "1.12.7",
+  "version": "1.20.4",
   "description": "",
   "license": "MIT",
   "scripts": {

diff --git a/src/codemirror.ts b/src/codemirror.ts
@@ -136,7 +136,7 @@ export class CodeMirrorManager {
         lineEnding: '\n',
         // We could use the regular path which could have a drive: prefix, but
         // this is probably unusual.
-        relativePath: relativePath,
+        absoluteUri: `file:///${relativePath}`,
       },
       editorOptions,
     });
@@ -330,6 +330,12 @@ export class CodeMirrorManager {
     if (event.key === 'Tab' && event.shiftKey) {
       return { consumeEvent: false, forceTriggerCompletion };
     }
+
+    // TODO(kevin): clean up autoHandle logic
+    // Currently we have:
+    // Jupyter Notebook:     tab = true,  escape = false
+    // Code Mirror Websites: tab = true,  escape = true
+    // Jupyter Lab:          tab = false, escape = false
     if (!event.metaKey && !event.ctrlKey && !event.altKey && !event.shiftKey) {
       if (alsoHandle.tab && event.key === tabKey && this.acceptCompletion()) {
         return { consumeEvent: true, forceTriggerCompletion };
@@ -340,7 +346,7 @@ export class CodeMirrorManager {
       // Special case if we are in jupyter notebooks and the tab key has been rebinded.
       // We do not want to consume the default keybinding, because it triggers the default
       // jupyter completion.
-      if (alsoHandle.tab && tabKey != 'Tab' && event.key === 'Tab') {
+      if (alsoHandle.tab && !alsoHandle.escape && tabKey !== 'Tab' && event.key === 'Tab') {
         return { consumeEvent: false, forceTriggerCompletion };
       }
     }

diff --git a/src/codemirrorInject.ts b/src/codemirrorInject.ts
@@ -6,15 +6,22 @@ declare type CodeMirror = typeof import('codemirror');
 export class CodeMirrorState {
   codeMirrorManager: CodeMirrorManager;
   docs: CodeMirror.Doc[] = [];
+  debounceMs: number = 0;
   hookedEditors = new WeakSet<CodeMirror.Editor>();
-  constructor(extensionId: string, cm: CodeMirror | undefined, readonly multiplayer: boolean) {
+  constructor(
+    extensionId: string,
+    cm: CodeMirror | undefined,
+    readonly multiplayer: boolean,
+    debounceMs?: number
+  ) {
     this.codeMirrorManager = new CodeMirrorManager(extensionId, {
       ideName: 'codemirror',
       ideVersion: `${cm?.version ?? 'unknown'}-${window.location.hostname}`,
     });
     if (cm !== undefined) {
       cm.defineInitHook(this.editorHook());
     }
+    this.debounceMs = debounceMs ?? 0;
   }
 
   editorHook(): (editor: CodeMirror.Editor) => void {
@@ -78,7 +85,7 @@ export class CodeMirrorState {
           undefined,
           undefined
         );
-      });
+      }, this.debounceMs);
     });
   }
 }
diff --git a/src/common.ts b/src/common.ts
@@ -13,9 +13,10 @@ import {
 } from '../proto/exa/language_server_pb/language_server_pb';
 
 const EXTENSION_NAME = 'chrome';
-const EXTENSION_VERSION = '1.12.7';
+const EXTENSION_VERSION = '1.20.4';
 
 export const CODEIUM_DEBUG = false;
+export const DEFAULT_PATH = 'unknown_url';
 
 export interface ClientSettings {
   apiKey?: string;

diff --git a/src/component/Options.tsx b/src/component/Options.tsx
@@ -155,6 +155,8 @@ const Options = () => {
   const jupyterNotebookKeybindingAcceptRef = createRef<HTMLInputElement>();
   const [jupyterNotebookKeybindingAcceptText, setJupyterNotebookKeybindingAcceptText] =
     useState('');
+  const [jupyterDebounceMs, setJupyterDebounceMs] = useState(0);
+  const jupyterDebounceMsRef = createRef<HTMLInputElement>();
 
   useEffect(() => {
     (async () => {
@@ -188,6 +190,11 @@ const Options = () => {
     })().catch((e) => {
       console.error(e);
     });
+    (async () => {
+      setJupyterDebounceMs((await getStorageItem('jupyterDebounceMs')) ?? 0);
+    })().catch((e) => {
+      console.error(e);
+    });
   }, []);
   // TODO(prem): Deduplicate with serviceWorker.ts/storage.ts.
   const resolvedPortalUrl = useMemo(() => {
@@ -405,6 +412,31 @@ const Options = () => {
           </Button>
         </Box>
       </Box>
+      <Box sx={{ my: 2, mx: 2 }}>
+        <Typography variant="h6"> Jupyter debounce time </Typography>
+        <TextField
+          id="jupyterDebounceMs"
+          label="Debounce time (ms)"
+          variant="standard"
+          fullWidth
+          type="number"
+          inputRef={jupyterDebounceMsRef}
+          value={jupyterDebounceMs}
+          onChange={(e) => setJupyterDebounceMs(Number(e.target.value))}
+        />
+        <Box sx={{ display: 'flex', justifyContent: 'flex-end' }}>
+          <Button
+            variant="text"
+            onClick={async () => {
+              const debounceTime = Number(jupyterDebounceMsRef.current?.value);
+              await setStorageItem('jupyterDebounceMs', debounceTime);
+            }}
+            sx={{ textTransform: 'none' }}
+          >
+            Save <LoginIcon />
+          </Button>
+        </Box>
+      </Box>
     </Box>
   );
 };

diff --git a/src/jupyterInject.ts b/src/jupyterInject.ts
@@ -1,9 +1,11 @@
 import type CodeMirror from 'codemirror';
 
 import { CodeMirrorManager } from './codemirror';
-import { JupyterNotebookKeyBindings } from './common';
+import { DEFAULT_PATH, JupyterNotebookKeyBindings } from './common';
 import { EditorOptions } from '../proto/exa/codeium_common_pb/codeium_common_pb';
 
+// Note: this file only deals with jupyter notebook 6 (not Jupyter Lab)
+
 declare class Cell {
   code_mirror: CodeMirror.Editor;
   notebook: Notebook;
@@ -63,7 +65,7 @@ class JupyterState {
     this.keybindings = keybindings;
   }
 
-  patchCellKeyEvent() {
+  patchCellKeyEvent(debounceMs?: number) {
     const beforeMainHandler = (doc: CodeMirror.Doc, event: KeyboardEvent) => {
       return this.codeMirrorManager.beforeMainKeyHandler(
         doc,
@@ -96,7 +98,7 @@ class JupyterState {
               return;
             }
           }
-          const textModels = [];
+          const textModels: CodeMirror.Doc[] = [];
 
           const editableCells = [...this.notebook.get_cells()];
           let currentModelWithOutput;
@@ -143,9 +145,10 @@ class JupyterState {
 
           const url = window.location.href;
           // URLs are usually of the form, http://localhost:XXXX/notebooks/path/to/notebook.ipynb
-          // We only want the path to the notebook.
+          // We only want path/to/notebook.ipynb
+          // If the URL is not of this form, we just use "unknown_url"
           const path = new URL(url).pathname;
-          const relativePath = path.endsWith('.ipynb') ? path : undefined;
+          const relativePath = path.endsWith('.ipynb') ? path.replace(/^\//, '') : undefined;
 
           await codeMirrorManager.triggerCompletion(
             true, // isNotebook
@@ -156,10 +159,10 @@ class JupyterState {
               tabSize: BigInt(editor.getOption('tabSize') ?? 4),
               insertSpaces: !(editor.getOption('indentWithTabs') ?? false),
             }),
-            relativePath,
-            undefined
+            relativePath ?? DEFAULT_PATH,
+            undefined // create_disposables
           );
-        });
+        }, debounceMs ?? 0);
       };
     };
     this.jupyter.CodeCell.prototype.handle_codemirror_keyevent = replaceOriginalHandler(
@@ -189,10 +192,11 @@ class JupyterState {
 export function inject(
   extensionId: string,
   jupyter: Jupyter,
-  keybindings: JupyterNotebookKeyBindings
+  keybindings: JupyterNotebookKeyBindings,
+  debounceMs?: number
 ): JupyterState {
   const jupyterState = new JupyterState(extensionId, jupyter, keybindings);
-  jupyterState.patchCellKeyEvent();
+  jupyterState.patchCellKeyEvent(debounceMs);
   jupyterState.patchShortcutManagerHandler();
   return jupyterState;
 }
diff --git a/src/jupyterlabPlugin.ts b/src/jupyterlabPlugin.ts
@@ -55,17 +55,21 @@ class CodeiumPlugin {
   codeMirrorManager: CodeMirrorManager;
   keybindings: Promise<JupyterLabKeyBindings>;
 
+  debounceMs: number;
+
   constructor(
     readonly extensionId: string,
     app: JupyterFrontEnd,
     notebookTracker: INotebookTracker,
     editorTracker: IEditorTracker,
-    documentManager: IDocumentManager
+    documentManager: IDocumentManager,
+    debounceMs: number
   ) {
     this.app = app;
     this.notebookTracker = notebookTracker;
     this.editorTracker = editorTracker;
     this.documentManager = documentManager;
+    this.debounceMs = debounceMs;
     this.codeMirrorManager = new CodeMirrorManager(extensionId, {
       ideName: 'jupyterlab',
       ideVersion: `${app.name.toLowerCase()} ${app.version}`,
@@ -224,15 +228,16 @@ class CodeiumPlugin {
           return keybindingDisposables;
         }
       );
-    });
+    }, this.debounceMs);
     void chrome.runtime.sendMessage(this.extensionId, { type: 'success' });
     return false;
   }
 }
 
 export function getPlugin(
   extensionId: string,
-  jupyterapp: JupyterFrontEnd
+  jupyterapp: JupyterFrontEnd,
+  debounceMs: number
 ): JupyterFrontEndPlugin<void> {
   return {
     id: 'codeium:plugin',
@@ -244,7 +249,14 @@ export function getPlugin(
       documentManager: IDocumentManager
     ) => {
       // This indirection is necessary to get us a `this` to store state in.
-      new CodeiumPlugin(extensionId, app, notebookTracker, editorTracker, documentManager);
+      new CodeiumPlugin(
+        extensionId,
+        app,
+        notebookTracker,
+        editorTracker,
+        documentManager,
+        debounceMs
+      );
     },
     requires: [
       // eslint-disable-next-line @typescript-eslint/ban-ts-comment

diff --git a/src/monacoCompletionProvider.ts b/src/monacoCompletionProvider.ts
@@ -4,6 +4,7 @@ import { IdeInfo, LanguageServerClient } from './common';
 import { getLanguage } from './monacoLanguages';
 import { TextAndOffsets, computeTextAndOffsets } from './notebook';
 import { numUtf8BytesToNumCodeUnits } from './utf';
+import { sleep } from './utils';
 import { Language } from '../proto/exa/codeium_common_pb/codeium_common_pb';
 import {
   CompletionItem,
@@ -197,9 +198,11 @@ function colabRelativePath(): string | undefined {
     return undefined;
   }
   if (fileId.source === 'drive') {
-    return `${fileId.fileId}.ipynb`;
+    let fileIdString = fileId.fileId;
+    fileIdString = fileIdString.replace(/^\//, '');
+    return `${fileIdString}.ipynb`;
   }
-  return fileId.fileId;
+  return fileId.fileId.replace(/^\//, '');
 }
 
 function deepnoteAndDatabricksRelativePath(url: string): string | undefined {
@@ -219,9 +222,11 @@ function deepnoteAndDatabricksRelativePath(url: string): string | undefined {
 export class MonacoCompletionProvider implements monaco.languages.InlineCompletionsProvider {
   modelUriToEditor = new Map<string, monaco.editor.ICodeEditor>();
   client: LanguageServerClient;
+  debounceMs: number;
 
-  constructor(readonly extensionId: string, readonly monacoSite: MonacoSite) {
+  constructor(readonly extensionId: string, readonly monacoSite: MonacoSite, debounceMs: number) {
     this.client = new LanguageServerClient(extensionId);
+    this.debounceMs = debounceMs;
   }
 
   getIdeInfo(): IdeInfo {
@@ -284,7 +289,11 @@ export class MonacoCompletionProvider implements monaco.languages.InlineCompleti
 
   private absolutePath(model: monaco.editor.ITextModel): string | undefined {
     // Given we are using path, note the docs on fsPath: https://microsoft.github.io/monaco-editor/api/classes/monaco.Uri.html#fsPath
-    return model.uri.path;
+    if (this.monacoSite === OMonacoSite.COLAB) {
+      // The colab absolute path is something like the cell number (i.e. /3)
+      return colabRelativePath();
+    }
+    return model.uri.path.replace(/^\//, '');
     // TODO(prem): Adopt some site-specific convention.
   }
 
@@ -434,14 +443,15 @@ export class MonacoCompletionProvider implements monaco.languages.InlineCompleti
         language: getLanguage(getEditorLanguage(model)),
         cursorOffset: BigInt(numUtf8Bytes),
         lineEnding: '\n',
-        relativePath: this.relativePath(),
-        absolutePath: this.absolutePath(model),
+        relativePathMigrateMeToWorkspaceUri: this.relativePath(),
+        absoluteUri: 'file:///' + this.absolutePath(model),
       },
       editorOptions: {
         tabSize: BigInt(model.getOptions().tabSize),
         insertSpaces: model.getOptions().insertSpaces,
       },
     });
+    await sleep(this.debounceMs ?? 0);
     const response = await this.client.getCompletions(request);
     if (response === undefined) {
       return;