@marimo-team/frontend 0.14.18-dev9 → 0.15.0

package/src/core/codemirror/copilot/extension.ts

@@ -0,0 +1,285 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import {
+  Compartment,
+  type EditorState,
+  type Extension,
+  Prec,
+  type Text,
+} from "@codemirror/state";
+import type { EditorView } from "@codemirror/view";
+import { keymap } from "@codemirror/view";
+import {
+  inlineCompletion,
+  rejectInlineCompletion,
+} from "@marimo-team/codemirror-ai";
+import {
+  copilotPlugin as codeiumCopilotPlugin,
+  codeiumOtherDocumentsConfig,
+  Language,
+} from "@valtown/codemirror-codeium";
+import {
+  type InlineCompletionItem,
+  type InlineCompletionList,
+  type InlineCompletionParams,
+  InlineCompletionTriggerKind,
+} from "vscode-languageserver-protocol";
+import type { CompletionConfig } from "@/core/config/config-schema";
+import type { AiInlineCompletionRequest } from "@/core/kernel/messages";
+import { API } from "@/core/network/api";
+import { Logger } from "@/utils/Logger";
+import { languageAdapterState } from "../language/extension";
+import { isInVimMode } from "../utils";
+import { COPILOT_FILENAME, copilotServer, getCopilotClient } from "./client";
+import { getCodes } from "./getCodes";
+import { isCopilotEnabled } from "./state";
+import { trimAutocompleteResponse } from "./trim-utils";
+
+const copilotCompartment = new Compartment();
+
+const logger = Logger.get("@github/copilot-language-server");
+
+const commonInlineCompletionConfig = {
+  delay: 300, // default is 500ms
+  includeKeymap: true,
+  events: {
+    // Only show suggestions when the editor is focused
+    shouldShowSuggestion: (view: EditorView) => view.hasFocus,
+    beforeSuggestionFetch: (view: EditorView) => view.hasFocus,
+  },
+};
+
+export const copilotBundle = (config: CompletionConfig): Extension => {
+  if (process.env.NODE_ENV === "test") {
+    return [];
+  }
+
+  const extensions: Extension[] = [];
+
+  if (config.copilot === "codeium" && config.codeium_api_key) {
+    extensions.push(
+      codeiumCopilotPlugin({
+        apiKey: config.codeium_api_key,
+        language: Language.PYTHON,
+      }),
+      codeiumOtherDocumentsConfig.of({
+        override: async () => {
+          return [
+            {
+              text: getCodes(""),
+              language: Language.PYTHON,
+              editorLanguage: "python",
+            },
+          ];
+        },
+      }),
+    );
+  }
+
+  if (config.copilot === "github") {
+    extensions.push(
+      inlineCompletion({
+        ...commonInlineCompletionConfig,
+        fetchFn: async (state, _abortSignal, view) => {
+          const currentPosition = state.selection.main.head;
+          if (!isCopilotEnabled()) {
+            return "";
+          }
+
+          if (!view.hasFocus) {
+            return "";
+          }
+
+          // If no text, don't fetch
+          if (state.doc.length === 0) {
+            return "";
+          }
+
+          // wait 20ms so that the view is updated first
+          await new Promise((resolve) => setTimeout(resolve, 20));
+
+          const currentCode = state.doc.toString();
+          const allCode = getCodes(currentCode);
+          const request = getCopilotRequest(state, allCode);
+          const response = await getCopilotClient().getCompletion(request);
+
+          // If we are at a new position, ignore the response
+          if (currentPosition !== view.state.selection.main.head) {
+            Logger.debug("ignoring response because of new position");
+            return "";
+          }
+
+          const suggestion = getSuggestion(response, request.position, state);
+          return suggestion;
+        },
+      }),
+    );
+  }
+
+  if (config.copilot === "custom") {
+    extensions.push(
+      inlineCompletion({
+        ...commonInlineCompletionConfig,
+        fetchFn: async (state, _abortSignal, view) => {
+          const currentPosition = state.selection.main.head;
+
+          if (state.doc.length === 0) {
+            return "";
+          }
+
+          // If not focused, don't fetch
+          if (!view.hasFocus) {
+            return "";
+          }
+
+          const prefix = state.doc.sliceString(0, state.selection.main.head);
+          const suffix = state.doc.sliceString(
+            state.selection.main.head,
+            state.doc.length,
+          );
+
+          // If no prefix, don't fetch
+          if (prefix.length === 0) {
+            return "";
+          }
+
+          const language = state.field(languageAdapterState).type;
+          let res = await API.post<AiInlineCompletionRequest, string>(
+            "/ai/inline_completion",
+            { prefix, suffix, language },
+          );
+
+          // If we are at a new position, ignore the response
+          if (currentPosition !== view.state.selection.main.head) {
+            Logger.debug("ignoring response because of new position");
+            return "";
+          }
+
+          // Sometimes the prefix might get included in the response, so we need to trim it
+          res = trimAutocompleteResponse({ response: res, prefix, suffix });
+
+          return res;
+        },
+      }),
+    );
+  }
+
+  return [
+    ...extensions,
+    Prec.highest(
+      keymap.of([
+        {
+          key: "Escape",
+          run: (view: EditorView) => {
+            const status = rejectInlineCompletion(view);
+            // When in vim mode, we need to propagate escape to exit insert mode.
+            if (isInVimMode(view)) {
+              return false;
+            }
+            return status;
+          },
+        },
+      ]),
+    ),
+    // place in own compartment so it doesn't interfere with other LSP
+    copilotCompartment.of(copilotServer()),
+  ];
+};
+
+function getCopilotRequest(
+  state: EditorState,
+  allCode: string,
+): InlineCompletionParams {
+  // We need to update the position of the cursor because added newlines
+  // from appending the other code
+  const currentCode = state.doc.toString();
+  const numberOfNewLines =
+    allCode.split("\n").length - currentCode.split("\n").length;
+
+  const position = offsetToPos(state.doc, state.selection.main.head);
+  position.line += numberOfNewLines;
+  return {
+    textDocument: {
+      uri: `file://${COPILOT_FILENAME}`,
+      version: "replace_me" as unknown as number,
+    },
+    position: position,
+    context: {
+      triggerKind: InlineCompletionTriggerKind.Automatic,
+    },
+    formattingOptions: {
+      tabSize: state.tabSize,
+      insertSpaces: true,
+    },
+  } as InlineCompletionParams;
+}
+
+function getSuggestion(
+  response: InlineCompletionList | InlineCompletionItem[] | null,
+  userPosition: InlineCompletionParams["position"],
+  state: EditorState,
+): string {
+  if (!response) {
+    logger.debug("No response from copilot");
+    return "";
+  }
+
+  const first = Array.isArray(response) ? response[0] : response.items[0];
+  if (!first) {
+    logger.debug("No response from copilot");
+    return "";
+  }
+
+  const { insertText, range } = first;
+  const insertTextString = String(insertText);
+
+  if (!range) {
+    logger.error("No range from copilot");
+    return insertTextString;
+  }
+
+  // Calculate the start of the suggestion relative to the current position
+  const startOffset = range.start.character - userPosition.character;
+
+  // If startOffset is negative, we need to trim the beginning of displayText
+  const resultText =
+    startOffset < 0 ? insertTextString.slice(-startOffset) : insertTextString;
+
+  // If the end of the suggestion already exists next in the document, we should trim it,
+  // for example closing quotes, brackets, etc.
+  // e.g
+  // current content: print("hello|")
+  //                     position ^
+  // suggestion: world")
+  // we need to remove the closing ") since we already have it
+  //
+  // See https://github.com/marimo-team/marimo/issues/830
+
+  // Loop through from the whole word to the end of the suggestion
+  // if we find a match, we trim it off the end of the suggestion
+  const remainingText = state.doc.sliceString(
+    state.selection.main.head,
+    state.doc.length,
+  );
+
+  for (let i = 0; i < resultText.length; i++) {
+    const remainingResultText = resultText.slice(i);
+    if (remainingText.startsWith(remainingResultText)) {
+      return resultText.slice(0, i);
+    }
+  }
+
+  return resultText;
+}
+
+export const exportedForTesting = {
+  getCopilotRequest,
+  getSuggestion,
+};
+
+function offsetToPos(doc: Text, offset: number) {
+  const line = doc.lineAt(offset);
+  return {
+    line: line.number - 1,
+    character: offset - line.from,
+  };
+}

package/src/core/codemirror/copilot/getCodes.ts

@@ -0,0 +1,130 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import { getAllEditorViews, notebookAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { variablesAtom } from "@/core/variables/state";
+import type { Variables } from "@/core/variables/types";
+import { Objects } from "@/utils/objects";
+import { getEditorCodeAsPython } from "../language/utils";
+
+export function getCodes(otherCode: string) {
+  const codes = getOtherCellsCode(otherCode);
+
+  return [...codes, otherCode].join("\n");
+}
+
+export function getOtherCellsCode(otherCode: string) {
+  // Get all other cells' code
+  // Put `import` statements at the top, as it can help copilot give better suggestions
+  // TODO: we should sort this topologically
+  const codes = getAllEditorViews()
+    .map((editorView) => {
+      const code = getEditorCodeAsPython(editorView);
+      if (code === otherCode) {
+        return null;
+      }
+      return code;
+    })
+    .filter(Boolean)
+    .sort((a, b) => {
+      if (a.startsWith("import") && !b.startsWith("import")) {
+        return -1;
+      }
+      if (!a.startsWith("import") && b.startsWith("import")) {
+        return 1;
+      }
+      return 0;
+    });
+
+  return codes;
+}
+
+const notebookCellCodes = atom((get) => {
+  const notebook = get(notebookAtom);
+  const codes = Objects.fromEntries(
+    notebook.cellIds.inOrderIds.map((id) => {
+      const handle = notebook.cellHandles[id];
+      return [
+        id,
+        handle?.current ? getEditorCodeAsPython(handle.current.editorView) : "",
+      ];
+    }),
+  );
+  return codes;
+});
+const inOrderCellIdsAtom = atom((get) => {
+  const notebook = get(notebookAtom);
+  return notebook.cellIds.inOrderIds;
+});
+const topologicalCellIdsAtom = atom((get) => {
+  const cellIds = get(inOrderCellIdsAtom);
+  const variables = get(variablesAtom);
+  return getTopologicalCellIds(cellIds, variables);
+});
+
+export const topologicalCodesAtom = atom((get) => {
+  const sortedCellIds = get(topologicalCellIdsAtom);
+  const codes = get(notebookCellCodes);
+  return { cellIds: sortedCellIds, codes };
+});
+
+export function getTopologicalCellIds(cellIds: CellId[], variables: Variables) {
+  // Build adjacency list
+  const adjacency = new Map<CellId, CellId[]>();
+  cellIds.forEach((id) => adjacency.set(id, []));
+
+  // Start with all cells with no declaredBy or usedBy
+  const noDepCells = new Set<CellId>(cellIds); // Cells with no declaredBy
+  const noUsedByCells = new Set<CellId>(cellIds); // Cells with no usedBy
+
+  // Link "declaredBy -> usedBy"
+  for (const { declaredBy, usedBy } of Object.values(variables)) {
+    for (const declCell of declaredBy) {
+      noDepCells.delete(declCell);
+      for (const useCell of usedBy) {
+        noUsedByCells.delete(useCell);
+
+        if (useCell !== declCell) {
+          adjacency.get(declCell)?.push(useCell);
+        }
+      }
+    }
+  }
+
+  // Kahn's algorithm for topological sort
+  const inDegree = new Map<CellId, number>();
+  cellIds.forEach((id) => inDegree.set(id, 0));
+
+  adjacency.forEach((targets) => {
+    targets.forEach((t) => {
+      inDegree.set(t, (inDegree.get(t) || 0) + 1);
+    });
+  });
+
+  const queue: CellId[] = [];
+  inDegree.forEach((deg, id) => {
+    if (deg === 0) {
+      queue.push(id);
+    }
+  });
+
+  const sorted: CellId[] = [];
+  while (queue.length > 0) {
+    const current = queue.shift();
+    if (!current) {
+      break;
+    }
+    sorted.push(current);
+    adjacency.get(current)?.forEach((t) => {
+      inDegree.set(t, (inDegree.get(t) || 0) - 1);
+      if (inDegree.get(t) === 0) {
+        queue.push(t);
+      }
+    });
+  }
+
+  // Put noDepCells at the end
+  const filteredSorted = sorted.filter((id) => !noDepCells.has(id));
+  return [...filteredSorted, ...noDepCells];
+}

package/src/core/codemirror/copilot/language-server.ts

@@ -0,0 +1,255 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { LanguageServerClient } from "@marimo-team/codemirror-languageserver";
+import { throttle } from "lodash-es";
+import type {
+  CompletionItem,
+  CompletionList,
+  CompletionParams,
+  DidChangeTextDocumentParams,
+  DidOpenTextDocumentParams,
+  Hover,
+  HoverParams,
+  InlineCompletionItem,
+  InlineCompletionList,
+  InlineCompletionParams,
+} from "vscode-languageserver-protocol";
+import { VersionedTextDocumentIdentifier } from "vscode-languageserver-protocol";
+import { Logger } from "@/utils/Logger";
+import { getCodes } from "./getCodes";
+import {
+  clearGitHubCopilotLoadingVersion,
+  isCopilotEnabled,
+  setGitHubCopilotLoadingVersion,
+} from "./state";
+import type {
+  GitHubCopilotSignInConfirmParams,
+  GitHubCopilotSignInInitiateResult,
+  GitHubCopilotStatusResult,
+} from "./types";
+
+const logger = Logger.get("@github/copilot-language-server");
+
+// A map of request methods and their parameters and return types
+export interface LSPRequestMap {
+  checkStatus: [{}, GitHubCopilotStatusResult];
+  signIn: [{}, GitHubCopilotSignInInitiateResult];
+  signInConfirm: [GitHubCopilotSignInConfirmParams, GitHubCopilotStatusResult];
+  signOut: [{}, GitHubCopilotStatusResult];
+  "textDocument/inlineCompletion": [
+    InlineCompletionParams,
+    InlineCompletionList | InlineCompletionItem[] | null,
+  ];
+}
+
+/**
+ * A client for the Copilot language server.
+ */
+export class CopilotLanguageServerClient extends LanguageServerClient {
+  private documentVersion = 0;
+  private hasOpenedDocument = false;
+
+  private async _request<Method extends keyof LSPRequestMap>(
+    method: Method,
+    params: LSPRequestMap[Method][0],
+  ): Promise<LSPRequestMap[Method][1]> {
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      return await (this as any).request(method, params);
+    } catch (error) {
+      logger.error("#request: Error", error);
+      throw error;
+    }
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  override async notify(method: any, params: any): Promise<any> {
+    logger.debug("#notify", method, params);
+    return super.notify(method, params);
+  }
+
+  override getInitializationOptions() {
+    const info = {
+      name: "marimo",
+      version: "0.1.0",
+    };
+    return {
+      ...super.getInitializationOptions(),
+      workspaceFolders: [],
+      capabilities: {
+        workspace: { workspaceFolders: false },
+      },
+      initializationOptions: {
+        editorInfo: info,
+        editorPluginInfo: info,
+      },
+    };
+  }
+
+  private isDisabled() {
+    return !isCopilotEnabled();
+  }
+
+  override async textDocumentDidOpen(
+    params: DidOpenTextDocumentParams,
+  ): Promise<DidOpenTextDocumentParams> {
+    if (this.isDisabled()) {
+      return params;
+    }
+    this.hasOpenedDocument = true;
+    return super.textDocumentDidOpen(params);
+  }
+
+  override async textDocumentCompletion(
+    params: CompletionParams,
+  ): Promise<CompletionList | CompletionItem[]> {
+    // Not used in Copilot
+    return [];
+  }
+
+  override async textDocumentDidChange(
+    params: DidChangeTextDocumentParams,
+  ): Promise<DidChangeTextDocumentParams> {
+    if (this.isDisabled()) {
+      return params;
+    }
+
+    if (!this.hasOpenedDocument) {
+      await this.textDocumentDidOpen({
+        textDocument: {
+          uri: params.textDocument.uri,
+          languageId: "python",
+          version: params.textDocument.version,
+          text: params.contentChanges[0].text,
+        },
+      });
+    }
+
+    const changes = params.contentChanges;
+    if (changes.length !== 1) {
+      logger.warn(
+        "#textDocumentDidChange: Multiple changes detected. This is not supported.",
+        changes,
+      );
+    }
+    const change = changes[0];
+    if ("range" in change) {
+      logger.warn(
+        "#textDocumentDidChange: Copilot doesn't support range changes.",
+        change,
+      );
+    }
+
+    const text = getCodes(change.text);
+    return super.textDocumentDidChange({
+      ...params,
+      contentChanges: [{ text: text }],
+      textDocument: VersionedTextDocumentIdentifier.create(
+        params.textDocument.uri,
+        ++this.documentVersion,
+      ),
+    });
+  }
+
+  override textDocumentHover(params: HoverParams): Promise<Hover> {
+    // Not used in Copilot
+    return Promise.resolve({ contents: [] });
+  }
+
+  // AUTH
+  signOut() {
+    return this._request("signOut", {});
+  }
+
+  async signInInitiate() {
+    logger.log("#signInInitiate: Starting sign-in flow");
+    try {
+      const result = await this._request("signIn", {});
+      logger.log("#signInInitiate: Sign-in flow started successfully");
+      return result;
+    } catch (error) {
+      logger.warn("#signInInitiate: Failed to start sign-in flow", error);
+      throw error;
+    }
+  }
+
+  async signInConfirm(params: GitHubCopilotSignInConfirmParams) {
+    logger.log("#signInConfirm: Confirming sign-in");
+    try {
+      const result = await this._request("signInConfirm", params);
+      logger.log("#signInConfirm: Sign-in confirmed successfully");
+      return result;
+    } catch (error) {
+      logger.warn("#signInConfirm: Failed to confirm sign-in", error);
+      throw error;
+    }
+  }
+
+  async signedIn() {
+    try {
+      const { status } = await this._request("checkStatus", {});
+      logger.log("#checkStatus: Status check completed", { status });
+      return (
+        status === "SignedIn" || status === "AlreadySignedIn" || status === "OK"
+      );
+    } catch (error) {
+      logger.warn("#signedIn: Failed to check sign-in status", error);
+      throw error;
+    }
+  }
+
+  private getCompletionInternal = async (
+    params: InlineCompletionParams,
+    version: number,
+  ): Promise<InlineCompletionList | InlineCompletionItem[] | null> => {
+    return await this._request("textDocument/inlineCompletion", {
+      ...params,
+      textDocument: {
+        ...params.textDocument,
+        version: version,
+      } as VersionedTextDocumentIdentifier,
+    });
+  };
+
+  // Even though the copilot extension has a debounce,
+  // there are multiple requests sent at the same time
+  // when multiple Codemirror instances are mounted at the same time.
+  // So we throttle it to ignore multiple requests at the same time.
+  private throttledGetCompletionInternal = throttle(
+    this.getCompletionInternal,
+    200,
+  );
+
+  async getCompletion(
+    params: InlineCompletionParams,
+  ): Promise<InlineCompletionList | InlineCompletionItem[] | null> {
+    if (this.isDisabled()) {
+      return null;
+    }
+
+    const requestVersion = this.documentVersion;
+    (params.textDocument as VersionedTextDocumentIdentifier).version =
+      requestVersion;
+
+    // If version is 0, it means the document hasn't been opened yet
+    if (requestVersion === 0) {
+      return null;
+    }
+
+    // Start a loading indicator
+    setGitHubCopilotLoadingVersion(requestVersion);
+    const response = await this.throttledGetCompletionInternal(
+      params,
+      requestVersion,
+    );
+    // Stop the loading indicator (only if the version hasn't changed)
+    clearGitHubCopilotLoadingVersion(requestVersion);
+
+    // If the document version has changed since the request was made, return an empty response
+    if (requestVersion !== this.documentVersion) {
+      return null;
+    }
+
+    return response ?? null;
+  }
+}