@marimo-team/frontend 0.14.18-dev4 → 0.14.18-dev41

package/src/core/codemirror/copilot/__tests__/trim-utils.test.ts

@@ -0,0 +1,291 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { describe, expect, it } from "vitest";
+import { trimAutocompleteResponse } from "../trim-utils";
+
+const trim = trimAutocompleteResponse;
+
+describe("trimAutocompleteResponse", () => {
+  it("should return the response unchanged when no prefix or suffix", () => {
+    expect(
+      trim({
+        response: "console.log('hello');",
+        prefix: "",
+        suffix: "",
+      }),
+    ).toBe("console.log('hello');");
+  });
+
+  it("should handle empty response", () => {
+    expect(
+      trim({
+        response: "",
+        prefix: "prefix",
+        suffix: "suffix",
+      }),
+    ).toBe("");
+  });
+
+  it("should handle null/undefined inputs gracefully", () => {
+    expect(
+      trim({
+        response: "",
+        prefix: "",
+        suffix: "",
+      }),
+    ).toBe("");
+  });
+
+  describe("basic prefix trimming", () => {
+    it("should trim exact prefix match", () => {
+      expect(
+        trim({
+          response: "const x = 42;",
+          prefix: "const x = ",
+          suffix: "",
+        }),
+      ).toBe("42;");
+    });
+
+    it("should not trim when prefix doesn't match", () => {
+      expect(
+        trim({
+          response: "const x = 42;",
+          prefix: "let x = ",
+          suffix: "",
+        }),
+      ).toBe("const x = 42;");
+    });
+
+    it("should handle empty prefix", () => {
+      expect(
+        trim({
+          response: "const x = 42;",
+          prefix: "",
+          suffix: "",
+        }),
+      ).toBe("const x = 42;");
+    });
+  });
+
+  describe("basic suffix trimming", () => {
+    it("should trim exact suffix match", () => {
+      expect(
+        trim({
+          response: "42 + 1",
+          prefix: "",
+          suffix: " + 1",
+        }),
+      ).toBe("42");
+    });
+
+    it("should not trim when suffix doesn't match", () => {
+      expect(
+        trim({
+          response: "42 + 1",
+          prefix: "",
+          suffix: " + 2",
+        }),
+      ).toBe("42 + 1");
+    });
+
+    it("should handle empty suffix", () => {
+      expect(
+        trim({
+          response: "const x = 42;",
+          prefix: "",
+          suffix: "",
+        }),
+      ).toBe("const x = 42;");
+    });
+  });
+
+  describe("combined prefix and suffix trimming", () => {
+    it("should trim both prefix and suffix", () => {
+      expect(
+        trim({
+          response: "def func():    print('hello')    return",
+          prefix: "def func():",
+          suffix: "    return",
+        }),
+      ).toBe("    print('hello')");
+    });
+
+    it("should handle overlapping prefix and suffix", () => {
+      expect(
+        trim({
+          response: "abcdefabc",
+          prefix: "abc",
+          suffix: "abc",
+        }),
+      ).toBe("def");
+    });
+  });
+
+  describe("multiple occurrences", () => {
+    it("should trim multiple consecutive prefixes", () => {
+      expect(
+        trim({
+          response: ">>>>>>hello",
+          prefix: ">>>",
+          suffix: "",
+        }),
+      ).toBe(">>>hello");
+    });
+
+    it("should trim multiple consecutive suffixes", () => {
+      expect(
+        trim({
+          response: "hello<<<<<<",
+          prefix: "",
+          suffix: "<<<",
+        }),
+      ).toBe("hello<<<");
+    });
+  });
+
+  describe("real-world scenarios", () => {
+    it("should handle multiline Python function completion", () => {
+      expect(
+        trim({
+          response:
+            "def calculate_sum(a, b):\n    result = a + b\n    return result",
+          prefix: "def calculate_sum(a, b):",
+          suffix: "\n    return result",
+        }),
+      ).toBe("\n    result = a + b");
+    });
+
+    it("should handle Python list comprehension", () => {
+      expect(
+        trim({
+          response: "1, 2, 3, 4, 5];",
+          prefix: "const items = [",
+          suffix: "];",
+        }),
+      ).toBe("1, 2, 3, 4, 5");
+    });
+
+    it("should handle HTML tag completion", () => {
+      expect(
+        trim({
+          response: "  <p>Hello World</p></div>",
+          prefix: '<div class="container">',
+          suffix: "</div>",
+        }),
+      ).toBe("  <p>Hello World</p>");
+    });
+
+    it("should handle SQL query completion", () => {
+      expect(
+        trim({
+          response: " age > 18 AND status = 'active' ORDER BY id",
+          prefix: "SELECT * FROM users WHERE",
+          suffix: "ORDER BY id",
+        }),
+      ).toBe(" age > 18 AND status = 'active' ");
+    });
+
+    it("should handle markdown completion", () => {
+      expect(
+        trim({
+          response: "## Introduction\n\nThis is a sample",
+          prefix: "## ",
+          suffix: "",
+        }),
+      ).toBe("Introduction\n\nThis is a sample");
+    });
+  });
+
+  describe("edge cases", () => {
+    it("should handle response that is exactly the prefix", () => {
+      expect(
+        trim({
+          response: "hello",
+          prefix: "hello",
+          suffix: "",
+        }),
+      ).toBe("");
+    });
+
+    it("should handle response that is exactly the suffix", () => {
+      expect(
+        trim({
+          response: "world",
+          prefix: "",
+          suffix: "world",
+        }),
+      ).toBe("");
+    });
+
+    it("should handle response that is exactly prefix + suffix", () => {
+      expect(
+        trim({
+          response: "helloworld",
+          prefix: "hello",
+          suffix: "world",
+        }),
+      ).toBe("");
+    });
+
+    it("should handle very long prefix/suffix", () => {
+      expect(
+        trim({
+          response: `${"a".repeat(1000)}middle${"b".repeat(1000)}`,
+          prefix: "a".repeat(1000),
+          suffix: "b".repeat(1000),
+        }),
+      ).toBe("middle");
+    });
+
+    it("should handle special characters in prefix/suffix", () => {
+      expect(
+        trim({
+          response: "function test() { /* comment */ return 42; /* end */ }",
+          prefix: "function test() { /* comment */",
+          suffix: "/* end */ }",
+        }),
+      ).toBe(" return 42; ");
+    });
+
+    it("should handle unicode characters", () => {
+      expect(
+        trim({
+          response: "const 🚀 = '🎉 Hello World! 🌟';",
+          prefix: "const 🚀 = ",
+          suffix: ";",
+        }),
+      ).toBe("'🎉 Hello World! 🌟'");
+    });
+
+    it("should handle newlines in prefix/suffix", () => {
+      expect(
+        trim({
+          response: "if (condition) {\n  console.log('test');\n}",
+          prefix: "if (condition) {\n",
+          suffix: "\n}",
+        }),
+      ).toBe("  console.log('test');");
+    });
+
+    it("should handle cases where AI response doesn't include context", () => {
+      expect(
+        trim({
+          response: "42",
+          prefix: "const x = ",
+          suffix: ";",
+        }),
+      ).toBe("42");
+    });
+
+    it("should handle partial matches correctly (not trim them)", () => {
+      // Should not trim partial matches
+      expect(
+        trim({
+          response: "function test() hello world return x",
+          prefix: "function test() {",
+          suffix: "return x;",
+        }),
+      ).toBe("function test() hello world return x");
+    });
+  });
+});

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

@@ -0,0 +1,123 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { languageServerWithClient } from "@marimo-team/codemirror-languageserver";
+import { WebSocketTransport } from "@open-rpc/client-js";
+import type { JSONRPCRequestData } from "@open-rpc/client-js/build/Request";
+import { Transport } from "@open-rpc/client-js/build/transports/Transport";
+import { toast } from "@/components/ui/use-toast";
+import { waitForConnectionOpen } from "@/core/network/connection";
+import { getRuntimeManager } from "@/core/runtime/config";
+import { Logger } from "@/utils/Logger";
+import { once } from "@/utils/once";
+import { CopilotLanguageServerClient } from "./language-server";
+import { waitForEnabledCopilot } from "./state";
+
+// Dummy file for the copilot language server
+export const COPILOT_FILENAME = "/__marimo_copilot__.py";
+export const LANGUAGE_ID = "copilot";
+const FILE_URI = `file://${COPILOT_FILENAME}`;
+
+export const createWSTransport = once(() => {
+  return new LazyWebsocketTransport();
+});
+
+/**
+ * Custom WSTransport that:
+ *  - waits for copilot to be enabled
+ *  - wait for the websocket to be available
+ */
+class LazyWebsocketTransport extends Transport {
+  private delegate: WebSocketTransport | undefined;
+
+  constructor() {
+    super();
+    this.delegate = undefined;
+  }
+
+  private getWsUrl(): string {
+    const runtimeManager = getRuntimeManager();
+    return runtimeManager.getLSPURL("copilot").toString();
+  }
+
+  private async tryConnect(retries = 3, delayMs = 1000): Promise<void> {
+    for (let attempt = 1; attempt <= retries; attempt++) {
+      try {
+        // Create delegate, if it doesn't exist
+        if (!this.delegate) {
+          this.delegate = new WebSocketTransport(this.getWsUrl());
+        }
+        await this.delegate.connect();
+        Logger.log("Copilot#connect: Connected successfully");
+        return;
+      } catch (error) {
+        Logger.warn(
+          `Copilot#connect: Connection attempt ${attempt}/${retries} failed`,
+          error,
+        );
+        if (attempt === retries) {
+          this.delegate = undefined;
+          // Show error toast on final retry
+          toast({
+            variant: "danger",
+            title: "GitHub Copilot Connection Error",
+            description:
+              "Failed to connect to GitHub Copilot. Please check settings and try again.",
+          });
+          throw error;
+        }
+        await new Promise((resolve) => setTimeout(resolve, delayMs));
+      }
+    }
+  }
+
+  override async connect() {
+    // Wait for copilot to be enabled
+    await waitForEnabledCopilot();
+    // Wait for ws to be available with retries
+    await waitForConnectionOpen();
+
+    // Try connecting with retries
+    return this.tryConnect();
+  }
+
+  override close() {
+    this.delegate?.close();
+    this.delegate = undefined;
+  }
+
+  override async sendData(
+    data: JSONRPCRequestData,
+    timeout: number | null | undefined,
+  ) {
+    // Clamp timeout to 5 seconds
+    timeout = Math.min(timeout ?? 5000, 5000);
+    return this.delegate?.sendData(data, timeout);
+  }
+}
+
+export const getCopilotClient = once(
+  () =>
+    new CopilotLanguageServerClient({
+      rootUri: FILE_URI,
+      workspaceFolders: null,
+      transport: createWSTransport(),
+    }),
+);
+
+export function copilotServer() {
+  return languageServerWithClient({
+    documentUri: FILE_URI,
+    client: getCopilotClient(),
+    languageId: LANGUAGE_ID,
+    // Disable all basic LSP features
+    // we only need textDocument/didChange
+    hoverEnabled: false,
+    completionEnabled: false,
+    definitionEnabled: false,
+    renameEnabled: false,
+    codeActionsEnabled: false,
+    signatureHelpEnabled: false,
+    diagnosticsEnabled: false,
+    sendIncrementalChanges: false,
+  });
+}

package/src/core/codemirror/copilot/copilot-config.tsx

@@ -0,0 +1,285 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { useAtom } from "jotai";
+import { CheckIcon, CopyIcon, Loader2Icon, XIcon } from "lucide-react";
+import { memo, useState } from "react";
+import { Button, buttonVariants } from "@/components/ui/button";
+import { Label } from "@/components/ui/label";
+import { ExternalLink } from "@/components/ui/links";
+import { toast } from "@/components/ui/use-toast";
+import { copyToClipboard } from "@/utils/copy";
+import { Logger } from "@/utils/Logger";
+import { getCopilotClient } from "./client";
+import { copilotSignedInState, isGitHubCopilotSignedInState } from "./state";
+import type { GitHubCopilotStatus } from "./types";
+
+export const CopilotConfig = memo(() => {
+  const [copilotSignedIn, copilotChangeSignIn] = useAtom(
+    isGitHubCopilotSignedInState,
+  );
+  const [step, setStep] = useAtom(copilotSignedInState);
+  const [localData, setLocalData] = useState<{ url: string; code: string }>();
+  const [loading, setLoading] = useState(false);
+
+  const trySignIn = async (evt: React.MouseEvent) => {
+    evt.preventDefault();
+    setLoading(true);
+    try {
+      const { verificationUri, status, userCode } = await initiateSignIn();
+      if (isSignedIn(status)) {
+        copilotChangeSignIn(true);
+      } else {
+        setStep("signingIn");
+        setLocalData({ url: verificationUri, code: userCode });
+      }
+    } finally {
+      setLoading(false);
+    }
+  };
+
+  const tryFinishSignIn = async (evt: React.MouseEvent) => {
+    evt.preventDefault();
+    if (!localData) {
+      return;
+    }
+
+    setLoading(true);
+    try {
+      const result = await handleSignInConfirmation(localData.code);
+      if (result.success) {
+        copilotChangeSignIn(true);
+        setStep("signedIn");
+      } else if (result.error === "connection") {
+        setStep("connectionError");
+        toast({
+          title: "GitHub Copilot Connection Error",
+          description: "Lost connection during sign-in. Please try again.",
+          variant: "danger",
+          action: <Button onClick={trySignIn}>Retry</Button>,
+        });
+      } else {
+        setStep("signInFailed");
+      }
+    } finally {
+      setLoading(false);
+    }
+  };
+
+  const signOut = async (evt: React.MouseEvent) => {
+    evt.preventDefault();
+    await handleSignOut();
+    copilotChangeSignIn(false);
+    setStep("signedOut");
+  };
+
+  const renderBody = () => {
+    // If we don't have a step set, infer it from the current state
+    const resolvedStep = step ?? (copilotSignedIn ? "signedIn" : "connecting");
+
+    switch (resolvedStep) {
+      case "connecting":
+        return <Label className="font-normal flex">Connecting...</Label>;
+      case "signedOut":
+        return (
+          <Button onClick={trySignIn} size="xs" variant="link">
+            Sign in to GitHub Copilot
+          </Button>
+        );
+
+      case "signingIn":
+        return (
+          <ol className="ml-4 text-sm list-decimal [&>li]:mt-2">
+            <li>
+              <div className="flex items-center">
+                Copy this code:
+                <strong className="ml-2">{localData?.code}</strong>
+                <CopyIcon
+                  className="ml-2 cursor-pointer opacity-60 hover:opacity-100 h-3 w-3"
+                  onClick={async () => {
+                    if (!localData) {
+                      return;
+                    }
+                    await copyToClipboard(localData.code);
+                    toast({
+                      description: "Copied to clipboard",
+                    });
+                  }}
+                />
+              </div>
+            </li>
+            <li>
+              Enter the code at this link:
+              <a
+                href={localData?.url}
+                target="_blank"
+                rel="noreferrer"
+                className={buttonVariants({ variant: "link", size: "xs" })}
+              >
+                {localData?.url}
+              </a>
+            </li>
+            <li>
+              Click here when done:
+              <Button size="xs" onClick={tryFinishSignIn} className="ml-1">
+                {loading && (
+                  <Loader2Icon className="h-3 w-3 mr-1 animate-spin" />
+                )}
+                Done
+              </Button>
+            </li>
+          </ol>
+        );
+
+      case "signInFailed":
+        return (
+          <div className="flex flex-col gap-1">
+            <div className="text-destructive text-sm">
+              Sign in failed. Please try again.
+            </div>
+            {loading ? (
+              <Loader2Icon className="h-3 w-3 mr-1 animate-spin" />
+            ) : (
+              <Button onClick={trySignIn} size="xs" variant="link">
+                Connect to GitHub Copilot
+              </Button>
+            )}
+          </div>
+        );
+
+      case "signedIn":
+        return (
+          <div className="flex items-center gap-5">
+            <Label className="font-normal flex items-center">
+              <div className="inline-flex items-center justify-center bg-(--grass-7) rounded-full p-1 mr-2">
+                <CheckIcon className="h-3 w-3 text-white" />
+              </div>
+              Connected
+            </Label>
+            <Button onClick={signOut} size="xs" variant="text">
+              Disconnect
+            </Button>
+          </div>
+        );
+
+      case "connectionError":
+        return (
+          <div className="flex flex-col gap-1">
+            <Label className="font-normal flex">
+              <XIcon className="h-4 w-4 mr-1" />
+              Connection Error
+            </Label>
+            <div className="text-sm">Unable to connect to GitHub Copilot.</div>
+            <Button onClick={trySignIn} size="xs" variant="link">
+              Retry Connection
+            </Button>
+          </div>
+        );
+
+      case "notConnected":
+        return (
+          <div className="flex flex-col gap-1">
+            <Label className="font-normal flex">
+              <XIcon className="h-4 w-4 mr-1" />
+              Unable to connect
+            </Label>
+            <div className="text-sm">
+              For troubleshooting, see the{" "}
+              <ExternalLink href="https://docs.marimo.io/getting_started/index.html#github-copilot">
+                docs
+              </ExternalLink>
+              .
+            </div>
+          </div>
+        );
+    }
+  };
+
+  return renderBody();
+});
+
+CopilotConfig.displayName = "CopilotConfig";
+
+// Utility functions
+const MAX_RETRIES = 5;
+const RETRY_DELAY_MS = 3000;
+
+async function initiateSignIn() {
+  const client = getCopilotClient();
+  return await client.signInInitiate();
+}
+
+async function handleSignInConfirmation(userCode: string) {
+  const client = getCopilotClient();
+
+  try {
+    Logger.log("Copilot#tryFinishSignIn: Attempting to confirm sign-in");
+    const { status } = await client.signInConfirm({ userCode });
+
+    if (isSignedIn(status)) {
+      Logger.log("Copilot#tryFinishSignIn: Sign-in confirmed successfully");
+      return { success: true };
+    }
+
+    Logger.warn(
+      "Copilot#tryFinishSignIn: Sign-in confirmation returned unexpected status",
+      { status },
+    );
+    return { success: false };
+  } catch (error) {
+    if (isConnectionError(error)) {
+      Logger.error(
+        "Copilot#tryFinishSignIn: Connection error during sign-in",
+        error,
+      );
+      return { success: false, error: "connection" };
+    }
+
+    return await handleSignInRetry(client);
+  }
+}
+
+async function handleSignInRetry(client: ReturnType<typeof getCopilotClient>) {
+  for (let i = 0; i < MAX_RETRIES; i++) {
+    try {
+      await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY_MS));
+      const signedIn = await client.signedIn();
+      if (signedIn) {
+        Logger.log(
+          "Copilot#tryFinishSignIn: Successfully signed in after retry",
+        );
+        return { success: true };
+      }
+    } catch (retryError) {
+      Logger.warn("Copilot#tryFinishSignIn: Retry attempt failed", {
+        attempt: i + 1,
+        maxRetries: MAX_RETRIES,
+      });
+
+      if (isConnectionError(retryError)) {
+        return { success: false, error: "connection" };
+      }
+    }
+  }
+
+  Logger.error("Copilot#tryFinishSignIn: All sign-in attempts failed");
+  return { success: false };
+}
+
+async function handleSignOut() {
+  const client = getCopilotClient();
+  await client.signOut();
+}
+
+function isConnectionError(error: unknown): boolean {
+  return (
+    error instanceof Error &&
+    (error.message.includes("ECONNREFUSED") ||
+      error.message.includes("WebSocket") ||
+      error.message.includes("network"))
+  );
+}
+
+function isSignedIn(status: GitHubCopilotStatus): boolean {
+  return (
+    status === "SignedIn" || status === "AlreadySignedIn" || status === "OK"
+  );
+}