@marimo-team/frontend 0.25.2-dev114 → 0.25.2-dev115

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

@@ -5,10 +5,12 @@ import type {
   JSONRPCRequest,
   Transport,
 } from "@marimo-team/codemirror-languageserver";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { InlineCompletionParams } from "vscode-languageserver-protocol";
 import { store } from "@/core/state/jotai";
+import * as notebookCodes from "../getCodes";
 import { CopilotLanguageServerClient } from "../language-server";
-import { copilotStatusState } from "../state";
+import { copilotStatusState, githubCopilotLoadingVersion } from "../state";
 
 const copilotState = vi.hoisted(() => ({ isEnabled: true }));
 
@@ -382,6 +384,7 @@ describe("CopilotLanguageServerClient", () => {
       contentChanges: [{ text: "value = 2" }],
     });
 
+    await client.flushDocumentChanges();
     expect(notifications(mockTransport, "textDocument/didChange")).toEqual([
       {
         jsonrpc: "2.0",
@@ -503,3 +506,509 @@ describe("CopilotLanguageServerClient", () => {
     });
   });
 });
+
+describe("Copilot completion requests", () => {
+  let transport: MockTransport;
+  let client: CopilotLanguageServerClient;
+  const params: InlineCompletionParams = {
+    textDocument: { uri: "file:///test.py" },
+    position: { line: 0, character: 1 },
+    context: { triggerKind: 2 },
+  };
+  const suggestion = [{ insertText: "result" }];
+
+  function requests() {
+    return transport.sent.filter(
+      (message): message is JSONRPCRequest =>
+        isRequest(message) &&
+        message.method === "textDocument/inlineCompletion",
+    );
+  }
+
+  async function change(text: string) {
+    await client.textDocumentDidChange({
+      textDocument: { uri: params.textDocument.uri, version: 1 },
+      contentChanges: [{ text }],
+    });
+  }
+
+  async function tick() {
+    await vi.advanceTimersByTimeAsync(0);
+  }
+
+  beforeEach(async () => {
+    vi.useFakeTimers();
+    copilotState.isEnabled = true;
+    store.set(githubCopilotLoadingVersion, null);
+    transport = new MockTransport();
+    client = new CopilotLanguageServerClient({
+      rootUri: "file:///test",
+      workspaceFolders: null,
+      transport,
+    });
+    await client.initializePromise;
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: params.textDocument.uri,
+        languageId: "python",
+        version: 1,
+        text: "x",
+      },
+    });
+    transport.sent.length = 0;
+  });
+
+  afterEach(() => {
+    client.close();
+    const shutdown = transport.sent.find(
+      (message) => isRequest(message) && message.method === "shutdown",
+    );
+    if (shutdown && isRequest(shutdown)) {
+      transport.emit({ jsonrpc: "2.0", id: shutdown.id, result: null });
+    }
+    vi.useRealTimers();
+    vi.restoreAllMocks();
+  });
+
+  it("coalesces edits and sends the latest document before the completion", async () => {
+    for (const text of ["p", "pr", "pri", "prin", "print"]) {
+      await change(text);
+    }
+    expect(notifications(transport, "textDocument/didChange")).toEqual([]);
+    const response = client.getCompletion(params);
+    await tick();
+    expect(transport.sent).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 6 },
+          contentChanges: [{ text: "print" }],
+        },
+      },
+      {
+        jsonrpc: "2.0",
+        id: expect.any(Number),
+        method: "textDocument/inlineCompletion",
+        params: {
+          ...params,
+          textDocument: { uri: params.textDocument.uri, version: 6 },
+        },
+      },
+    ]);
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[0].id,
+      result: suggestion,
+    });
+    await expect(response).resolves.toEqual(suggestion);
+    expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+    expect(await client.flushDocumentChanges()).toBe("print");
+    expect(notifications(transport, "textDocument/didChange")).toHaveLength(1);
+  });
+
+  it.each(["x", ""])(
+    "does not upload an unchanged focused cell %j again",
+    async (text) => {
+      const assemble = vi.spyOn(notebookCodes, "getCodes");
+      await change(text);
+      expect(await client.flushDocumentChanges(text)).toBe(text);
+      expect(await client.flushDocumentChanges(text)).toBe(text);
+      expect(await client.flushDocumentChanges(text)).toBe(text);
+      expect(assemble).toHaveBeenCalledTimes(3);
+      expect(notifications(transport, "textDocument/didChange")).toEqual([
+        {
+          jsonrpc: "2.0",
+          method: "textDocument/didChange",
+          params: {
+            textDocument: { uri: params.textDocument.uri, version: 2 },
+            contentChanges: [{ text }],
+          },
+        },
+      ]);
+
+      await change("next");
+      await client.flushDocumentChanges("next");
+      expect(
+        notifications(transport, "textDocument/didChange")[1].params,
+      ).toEqual({
+        textDocument: { uri: params.textDocument.uri, version: 3 },
+        contentChanges: [{ text: "next" }],
+      });
+    },
+  );
+
+  it("does not upload a document that already matches didOpen", async () => {
+    expect(await client.flushDocumentChanges("x")).toBe("x");
+    expect(notifications(transport, "textDocument/didChange")).toEqual([]);
+  });
+
+  it("compares the assembled notebook and uploads changes in another cell", async () => {
+    let otherCell = "import math";
+    const assemble = vi
+      .spyOn(notebookCodes, "getCodes")
+      .mockImplementation((code) => `${otherCell}\n${code}`);
+    await change("x");
+    expect(await client.flushDocumentChanges("x")).toBe("import math\nx");
+    expect(await client.flushDocumentChanges("x")).toBe("import math\nx");
+    otherCell = "import random";
+    expect(await client.flushDocumentChanges("x")).toBe("import random\nx");
+    expect(await client.flushDocumentChanges("x")).toBe("import random\nx");
+    expect(assemble).toHaveBeenCalledTimes(4);
+    expect(notifications(transport, "textDocument/didChange")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 2 },
+          contentChanges: [{ text: "import math\nx" }],
+        },
+      },
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 3 },
+          contentChanges: [{ text: "import random\nx" }],
+        },
+      },
+    ]);
+  });
+
+  it("assembles once with the focused cell last when another cell has a pending edit", async () => {
+    const assemble = vi
+      .spyOn(notebookCodes, "getCodes")
+      .mockImplementation((code) => `other\n${code}`);
+    await change("other");
+    expect(await client.flushDocumentChanges("x")).toBe("other\nx");
+    expect(assemble).toHaveBeenCalledExactlyOnceWith("x");
+    expect(notifications(transport, "textDocument/didChange")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 3 },
+          contentChanges: [{ text: "other\nx" }],
+        },
+      },
+    ]);
+  });
+
+  it("retries a failed upload even when its text matches the last attempted document", async () => {
+    const assemble = vi.spyOn(notebookCodes, "getCodes");
+    transport.send.mockImplementationOnce(() => {
+      throw new Error("socket closed");
+    });
+    await expect(client.flushDocumentChanges("latest")).rejects.toThrow(
+      "socket closed",
+    );
+    expect(await client.flushDocumentChanges("latest")).toBe("latest");
+    expect(await client.flushDocumentChanges("latest")).toBe("latest");
+    expect(assemble).toHaveBeenCalledTimes(3);
+    expect(notifications(transport, "textDocument/didChange")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 2 },
+          contentChanges: [{ text: "latest" }],
+        },
+      },
+    ]);
+  });
+
+  it("cancels a rapid request instead of giving the next caller its promise", async () => {
+    const first = client.getCompletion(params);
+    await tick();
+    const second = client.getCompletion({
+      ...params,
+      position: { line: 0, character: 0 },
+    });
+    await tick();
+    expect(requests()).toHaveLength(2);
+    expect(notifications(transport, "$/cancelRequest")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "$/cancelRequest",
+        params: { id: requests()[0].id },
+      },
+    ]);
+    await expect(first).resolves.toBeNull();
+    // Both requests use the same version; old cleanup must not hide the new spinner.
+    expect(store.get(githubCopilotLoadingVersion)).toBe(1);
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[1].id,
+      result: suggestion,
+    });
+    await expect(second).resolves.toEqual(suggestion);
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[0].id,
+      result: [{ insertText: "stale" }],
+    });
+    expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+  });
+
+  it("suppresses a request superseded before it reaches the wire", async () => {
+    const first = client.getCompletion(params);
+    const second = client.getCompletion(params);
+    await tick();
+    expect(requests()).toHaveLength(1);
+    await expect(first).resolves.toBeNull();
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[0].id,
+      result: suggestion,
+    });
+    await expect(second).resolves.toEqual(suggestion);
+  });
+
+  it("aborts on an edit before any replacement completion is requested", async () => {
+    const response = client.getCompletion(params);
+    await tick();
+    await change("changed");
+    await expect(response).resolves.toBeNull();
+    expect(notifications(transport, "$/cancelRequest")).toHaveLength(1);
+    expect(notifications(transport, "textDocument/didChange")).toEqual([]);
+    expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+  });
+
+  it.each([false, true])(
+    "handles external cancellation (already sent: %s)",
+    async (sent) => {
+      const controller = new AbortController();
+      if (!sent) {
+        controller.abort();
+      }
+      const response = client.getCompletion(params, controller.signal);
+      if (sent) {
+        await tick();
+        controller.abort();
+      }
+      await expect(response).resolves.toBeNull();
+      expect(requests()).toHaveLength(sent ? 1 : 0);
+      expect(notifications(transport, "$/cancelRequest")).toHaveLength(
+        sent ? 1 : 0,
+      );
+      expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+    },
+  );
+
+  it.each([null, [], { items: [] }])(
+    "handles empty response %j",
+    async (result) => {
+      const response = client.getCompletion(params);
+      await tick();
+      transport.emit({ jsonrpc: "2.0", id: requests()[0].id, result });
+      await expect(response).resolves.toEqual(result);
+      expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+    },
+  );
+
+  it("fails quietly on a server error and allows a subsequent request", async () => {
+    const response = client.getCompletion(params);
+    await tick();
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[0].id,
+      error: { code: -32603, message: "unavailable" },
+    });
+    await expect(response).resolves.toBeNull();
+    const retry = client.getCompletion(params);
+    await tick();
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[1].id,
+      result: suggestion,
+    });
+    await expect(retry).resolves.toEqual(suggestion);
+  });
+
+  it("cancels a timed-out request and clears its loading state", async () => {
+    const response = client.getCompletion(params);
+    await vi.advanceTimersByTimeAsync(10_000);
+    await expect(response).resolves.toBeNull();
+    expect(notifications(transport, "$/cancelRequest")).toHaveLength(1);
+    expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+  });
+
+  it("does not send requests while disabled or display a response after disabling", async () => {
+    copilotState.isEnabled = false;
+    await expect(client.getCompletion(params)).resolves.toBeNull();
+    expect(requests()).toEqual([]);
+    copilotState.isEnabled = true;
+    const response = client.getCompletion(params);
+    await tick();
+    copilotState.isEnabled = false;
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[0].id,
+      result: suggestion,
+    });
+    await expect(response).resolves.toBeNull();
+  });
+
+  it("preserves pending edits after a failed document upload and retries them", async () => {
+    await change("latest");
+    transport.send.mockImplementationOnce(() => {
+      throw new Error("socket closed");
+    });
+    await expect(client.getCompletion(params)).resolves.toBeNull();
+    expect(requests()).toEqual([]);
+    const response = client.getCompletion(params);
+    await tick();
+    expect(notifications(transport, "textDocument/didChange")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 2 },
+          contentChanges: [{ text: "latest" }],
+        },
+      },
+    ]);
+    transport.emit({
+      jsonrpc: "2.0",
+      id: requests()[0].id,
+      result: suggestion,
+    });
+    await expect(response).resolves.toEqual(suggestion);
+  });
+
+  it.each(["success", "failure"])(
+    "cancels a request during a stalled reconnect before handshake %s",
+    async (outcome) => {
+      transport.send.mockImplementationOnce((message) => {
+        transport.sent.push(message);
+      });
+      const reconnect = client.reInitialize().then(
+        () => "ready",
+        (error: unknown) => error,
+      );
+      const initialize = transport.sent.find(
+        (message) => isRequest(message) && message.method === "initialize",
+      );
+      if (!initialize || !isRequest(initialize)) {
+        throw new Error("Expected initialize request");
+      }
+      const controller = new AbortController();
+      const remove = vi.spyOn(controller.signal, "removeEventListener");
+      const settled = vi.fn();
+      const response = client
+        .getCompletion(params, controller.signal)
+        .then(settled);
+      await tick();
+      expect(settled).not.toHaveBeenCalled();
+      expect(requests()).toEqual([]);
+      expect(store.get(githubCopilotLoadingVersion)).toBe(1);
+      controller.abort();
+      await tick();
+      // Assert before answering initialize so the handshake cannot hide a stuck abort.
+      expect(settled).toHaveBeenCalledExactlyOnceWith(null);
+      expect(remove).toHaveBeenCalledExactlyOnceWith(
+        "abort",
+        expect.any(Function),
+      );
+      expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+      expect(notifications(transport, "$/cancelRequest")).toEqual([]);
+      transport.emit(
+        outcome === "success"
+          ? {
+              jsonrpc: "2.0",
+              id: initialize.id,
+              result: { capabilities: {} },
+            }
+          : {
+              jsonrpc: "2.0",
+              id: initialize.id,
+              error: { code: -32603, message: "handshake failed" },
+            },
+      );
+      const handshakeError = expect.objectContaining({
+        message: "handshake failed",
+      });
+      await expect(reconnect).resolves.toEqual(
+        outcome === "success" ? "ready" : handshakeError,
+      );
+      expect(requests()).toEqual([]);
+      let retryResult: unknown;
+      if (outcome === "success") {
+        const retry = client.getCompletion(params);
+        await tick();
+        transport.emit({
+          jsonrpc: "2.0",
+          id: requests()[0].id,
+          result: suggestion,
+        });
+        retryResult = await retry;
+      }
+      expect(requests()).toHaveLength(outcome === "success" ? 1 : 0);
+      expect(retryResult).toEqual(
+        outcome === "success" ? suggestion : undefined,
+      );
+      await response;
+      expect(settled).toHaveBeenCalledTimes(1);
+    },
+  );
+
+  it("replays the most recent unsent edit on reconnect", async () => {
+    await change("first");
+    await change("latest");
+    await client.reInitialize();
+    expect(notifications(transport, "textDocument/didOpen")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didOpen",
+        params: {
+          textDocument: {
+            uri: params.textDocument.uri,
+            languageId: "python",
+            version: 3,
+            text: "latest",
+          },
+        },
+      },
+    ]);
+    expect(notifications(transport, "textDocument/didChange")).toEqual([]);
+  });
+
+  it("cancels requests and drops pending changes when the final editor closes", async () => {
+    const response = client.getCompletion(params);
+    await tick();
+    await client.textDocumentDidClose({ textDocument: params.textDocument });
+    await expect(response).resolves.toBeNull();
+    await change("unsent");
+    await client.textDocumentDidClose({ textDocument: params.textDocument });
+    await client.flushDocumentChanges();
+    expect(notifications(transport, "textDocument/didChange")).toEqual([]);
+  });
+  it("does not request completions for a closed document", async () => {
+    await client.textDocumentDidClose({ textDocument: params.textDocument });
+    await expect(client.getCompletion(params)).resolves.toBeNull();
+    expect(requests()).toEqual([]);
+    expect(store.get(githubCopilotLoadingVersion)).toBeNull();
+  });
+
+  it("keeps a newer edit when an earlier upload fails", async () => {
+    await change("old");
+    transport.send.mockImplementationOnce(() => {
+      void change("new");
+      throw new Error("send failed");
+    });
+    await expect(client.flushDocumentChanges("old")).rejects.toThrow(
+      "send failed",
+    );
+    expect(await client.flushDocumentChanges("new")).toBe("new");
+    expect(notifications(transport, "textDocument/didChange")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: params.textDocument.uri, version: 3 },
+          contentChanges: [{ text: "new" }],
+        },
+      },
+    ]);
+  });
+});