@marimo-team/frontend 0.23.16-dev35 → 0.23.16-dev36

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

@@ -1,28 +1,81 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-/* oxlint-disable typescript/no-explicit-any */
-
-import { Transport } from "@open-rpc/client-js/build/transports/Transport";
+import type {
+  JSONRPCMessage,
+  JSONRPCNotification,
+  JSONRPCRequest,
+  Transport,
+} from "@marimo-team/codemirror-languageserver";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { store } from "@/core/state/jotai";
 import { CopilotLanguageServerClient } from "../language-server";
+import { copilotStatusState } from "../state";
+
+const copilotState = vi.hoisted(() => ({ isEnabled: true }));
+
+vi.mock("../state", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("../state")>();
+  return {
+    ...actual,
+    isCopilotEnabled: () => copilotState.isEnabled,
+  };
+});
+
+class MockTransport implements Transport {
+  private readonly handlers = new Set<(message: JSONRPCMessage) => void>();
+  readonly sent: JSONRPCMessage[] = [];
+
+  connect = vi.fn().mockResolvedValue(undefined);
+  close = vi.fn();
+  send = vi.fn((message: JSONRPCMessage) => {
+    this.sent.push(message);
+    if (isRequest(message) && message.method === "initialize") {
+      queueMicrotask(() => {
+        this.emit({
+          jsonrpc: "2.0",
+          id: message.id,
+          result: { capabilities: {} },
+        });
+      });
+    }
+  });
+
+  onMessage(handler: (message: JSONRPCMessage) => void): () => void {
+    this.handlers.add(handler);
+    return () => {
+      this.handlers.delete(handler);
+    };
+  }
+
+  emit(message: JSONRPCMessage) {
+    for (const handler of this.handlers) {
+      handler(message);
+    }
+  }
+}
 
-// Mock transport
-class MockTransport extends Transport {
-  override connect = vi.fn().mockResolvedValue(undefined);
-  override close = vi.fn();
-  override sendData = vi.fn().mockResolvedValue({});
-  override subscribe = vi.fn();
-  override unsubscribe = vi.fn();
-  override parseData = vi.fn();
+function isRequest(message: JSONRPCMessage): message is JSONRPCRequest {
+  return "method" in message && "id" in message;
+}
+
+function notifications(
+  transport: MockTransport,
+  method: string,
+): JSONRPCNotification[] {
+  return transport.sent.filter(
+    (message): message is JSONRPCNotification =>
+      "method" in message && !("id" in message) && message.method === method,
+  );
 }
 
 describe("CopilotLanguageServerClient", () => {
   let mockTransport: MockTransport;
 
   beforeEach(() => {
+    copilotState.isEnabled = true;
     mockTransport = new MockTransport();
   });
 
-  it("should initialize without copilot settings", () => {
+  it("initializes without copilot settings", () => {
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
@@ -32,178 +85,421 @@ describe("CopilotLanguageServerClient", () => {
     expect(client).toBeDefined();
   });
 
-  it("should initialize with empty copilot settings", () => {
+  it("initializes with copilot settings", () => {
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
-      copilotSettings: {},
+      copilotSettings: {
+        http: {
+          proxy: "http://proxy.example.com:8888",
+          proxyStrictSSL: true,
+        },
+        telemetry: {
+          telemetryLevel: "off",
+        },
+      },
     });
 
     expect(client).toBeDefined();
   });
 
-  it("should initialize with copilot settings", () => {
+  it("sends configuration after initialization", async () => {
     const copilotSettings = {
       http: {
         proxy: "http://proxy.example.com:8888",
         proxyStrictSSL: true,
       },
-      telemetry: {
-        telemetryLevel: "off",
-      },
     };
-
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
       copilotSettings,
     });
-
-    expect(client).toBeDefined();
+    await client.initializePromise;
+    await vi.waitFor(() => {
+      expect(
+        notifications(mockTransport, "workspace/didChangeConfiguration"),
+      ).toEqual([
+        {
+          jsonrpc: "2.0",
+          method: "workspace/didChangeConfiguration",
+          params: { settings: copilotSettings },
+        },
+      ]);
+    });
   });
 
-  it("should send configuration after initialization when settings are provided", async () => {
+  it.each([
+    ["empty", {}],
+    ["omitted", undefined],
+  ])(
+    "does not send configuration when settings are %s",
+    async (_, copilotSettings) => {
+      const client = new CopilotLanguageServerClient({
+        rootUri: "file:///test",
+        workspaceFolders: null,
+        transport: mockTransport,
+        copilotSettings,
+      });
+      await client.initializePromise;
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      expect(
+        notifications(mockTransport, "workspace/didChangeConfiguration"),
+      ).toEqual([]);
+    },
+  );
+
+  it("accepts enterprise, proxy, and telemetry settings together", async () => {
     const copilotSettings = {
       http: {
         proxy: "http://proxy.example.com:8888",
         proxyStrictSSL: true,
+        proxyKerberosServicePrincipal: "HTTP/proxy.example.com",
+      },
+      telemetry: {
+        telemetryLevel: "all",
+      },
+      "github-enterprise": {
+        uri: "https://github.enterprise.com",
       },
     };
-
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
       copilotSettings,
     });
-
-    // Spy on the notify method
-    const notifySpy = vi.spyOn(client as any, "notify");
-
-    // Wait for initialization
     await client.initializePromise;
-
-    // Wait for next tick to let the configuration be sent
-    await new Promise((resolve) => setTimeout(resolve, 0));
-
-    // Verify that notify was called with the correct parameters
-    expect(notifySpy).toHaveBeenCalledWith("workspace/didChangeConfiguration", {
-      settings: copilotSettings,
+    await vi.waitFor(() => {
+      expect(
+        notifications(mockTransport, "workspace/didChangeConfiguration"),
+      ).toEqual([
+        {
+          jsonrpc: "2.0",
+          method: "workspace/didChangeConfiguration",
+          params: { settings: copilotSettings },
+        },
+      ]);
     });
   });
 
-  it("should not send configuration after initialization when settings are empty", async () => {
+  it("reopens the latest Copilot document after re-initialization", async () => {
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
-      copilotSettings: {},
     });
+    await client.initializePromise;
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: "file:///test.py",
+        languageId: "python",
+        version: 0,
+        text: "initial",
+      },
+    });
+    await client.textDocumentDidChange({
+      textDocument: { uri: "file:///test.py", version: 1 },
+      contentChanges: [{ text: "updated" }],
+    });
+    mockTransport.sent.length = 0;
+
+    await client.reInitialize();
+
+    expect(notifications(mockTransport, "textDocument/didOpen")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didOpen",
+        params: {
+          textDocument: {
+            uri: "file:///test.py",
+            languageId: "python",
+            version: 1,
+            text: "updated",
+          },
+        },
+      },
+    ]);
+  });
 
-    // Spy on the notify method
-    const notifySpy = vi.spyOn(client as any, "notify");
+  it("reopens the document with the languageId it was opened with", async () => {
+    const client = new CopilotLanguageServerClient({
+      rootUri: "file:///test",
+      workspaceFolders: null,
+      transport: mockTransport,
+    });
+    await client.initializePromise;
+    // The Copilot plugin opens as "copilot", not "python"
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: "file:///test.py",
+        languageId: "copilot",
+        version: 0,
+        text: "initial",
+      },
+    });
+    await client.textDocumentDidChange({
+      textDocument: { uri: "file:///test.py", version: 1 },
+      contentChanges: [{ text: "updated" }],
+    });
+    mockTransport.sent.length = 0;
+
+    await client.reInitialize();
+
+    expect(notifications(mockTransport, "textDocument/didOpen")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didOpen",
+        params: {
+          textDocument: {
+            uri: "file:///test.py",
+            languageId: "copilot",
+            version: 1,
+            text: "updated",
+          },
+        },
+      },
+    ]);
+  });
 
-    // Wait for initialization
+  it("ignores an empty didChange without corrupting document state", async () => {
+    const client = new CopilotLanguageServerClient({
+      rootUri: "file:///test",
+      workspaceFolders: null,
+      transport: mockTransport,
+    });
     await client.initializePromise;
+    mockTransport.sent.length = 0;
 
-    // Wait for next tick
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    await expect(
+      client.textDocumentDidChange({
+        textDocument: { uri: "file:///test.py", version: 1 },
+        contentChanges: [],
+      }),
+    ).resolves.toBeUndefined();
 
-    // Verify that notify was NOT called with didChangeConfiguration
-    expect(notifySpy).not.toHaveBeenCalledWith(
-      "workspace/didChangeConfiguration",
-      expect.anything(),
-    );
+    expect(notifications(mockTransport, "textDocument/didOpen")).toEqual([]);
+    expect(notifications(mockTransport, "textDocument/didChange")).toEqual([]);
   });
 
-  it("should not send configuration after initialization when settings are not provided", async () => {
+  it("balances document references when Copilot is disabled before close", async () => {
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
     });
-
-    // Spy on the notify method
-    const notifySpy = vi.spyOn(client as any, "notify");
-
-    // Wait for initialization
     await client.initializePromise;
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: "file:///test.py",
+        languageId: "python",
+        version: 7,
+        text: "value = 1",
+      },
+    });
 
-    // Wait for next tick
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    copilotState.isEnabled = false;
+    await client.textDocumentDidClose({
+      textDocument: { uri: "file:///test.py" },
+    });
+    copilotState.isEnabled = true;
+    mockTransport.sent.length = 0;
+    await client.reInitialize();
 
-    // Verify that notify was NOT called with didChangeConfiguration
-    expect(notifySpy).not.toHaveBeenCalledWith(
-      "workspace/didChangeConfiguration",
-      expect.anything(),
-    );
+    expect(notifications(mockTransport, "textDocument/didOpen")).toEqual([]);
   });
 
-  it("should accept github-enterprise settings", async () => {
-    const copilotSettings = {
-      "github-enterprise": {
-        uri: "https://github.enterprise.com",
-      },
-    };
-
+  it("does not close the shared document until the last mounted editor closes it (relies on base client's own refcounting)", async () => {
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
-      copilotSettings,
     });
-
-    // Spy on the notify method
-    const notifySpy = vi.spyOn(client as any, "notify");
-
-    // Wait for initialization
     await client.initializePromise;
 
-    // Wait for next tick to let the configuration be sent
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    // Two editors (e.g. two visible cells) share this client and both open
+    // the same document.
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: "file:///test.py",
+        languageId: "python",
+        version: 1,
+        text: "value = 1",
+      },
+    });
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: "file:///test.py",
+        languageId: "python",
+        version: 1,
+        text: "value = 1",
+      },
+    });
+    mockTransport.sent.length = 0;
 
-    // Verify that notify was called with the enterprise settings
-    expect(notifySpy).toHaveBeenCalledWith("workspace/didChangeConfiguration", {
-      settings: copilotSettings,
+    // Closing the first editor must not close the shared document while the
+    // second editor is still mounted.
+    await client.textDocumentDidClose({
+      textDocument: { uri: "file:///test.py" },
     });
-  });
+    expect(notifications(mockTransport, "textDocument/didClose")).toEqual([]);
 
-  it("should accept all supported settings at once", async () => {
-    const copilotSettings = {
-      http: {
-        proxy: "http://proxy.example.com:8888",
-        proxyStrictSSL: true,
-        proxyKerberosServicePrincipal: "HTTP/proxy.example.com",
+    // Closing the last editor closes the document for real.
+    await client.textDocumentDidClose({
+      textDocument: { uri: "file:///test.py" },
+    });
+    expect(notifications(mockTransport, "textDocument/didClose")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didClose",
+        params: { textDocument: { uri: "file:///test.py" } },
       },
-      telemetry: {
-        telemetryLevel: "all",
+    ]);
+  });
+
+  it("keeps document versions monotonic after opening a versioned document", async () => {
+    const client = new CopilotLanguageServerClient({
+      rootUri: "file:///test",
+      workspaceFolders: null,
+      transport: mockTransport,
+    });
+    await client.initializePromise;
+    await client.textDocumentDidOpen({
+      textDocument: {
+        uri: "file:///test.py",
+        languageId: "python",
+        version: 7,
+        text: "value = 1",
       },
-      "github-enterprise": {
-        uri: "https://github.enterprise.com",
+    });
+    mockTransport.sent.length = 0;
+
+    await client.textDocumentDidChange({
+      textDocument: { uri: "file:///test.py", version: 8 },
+      contentChanges: [{ text: "value = 2" }],
+    });
+
+    expect(notifications(mockTransport, "textDocument/didChange")).toEqual([
+      {
+        jsonrpc: "2.0",
+        method: "textDocument/didChange",
+        params: {
+          textDocument: { uri: "file:///test.py", version: 8 },
+          contentChanges: [{ text: "value = 2" }],
+        },
       },
-    };
+    ]);
+  });
 
+  it("ignores malformed custom notifications", async () => {
     const client = new CopilotLanguageServerClient({
       rootUri: "file:///test",
       workspaceFolders: null,
       transport: mockTransport,
-      copilotSettings,
     });
+    await client.initializePromise;
 
-    // Spy on the notify method
-    const notifySpy = vi.spyOn(client as any, "notify");
+    expect(() => {
+      mockTransport.emit({
+        jsonrpc: "2.0",
+        method: "statusNotification",
+        params: { busy: "not-a-boolean" },
+      });
+      mockTransport.emit({
+        jsonrpc: "2.0",
+        method: "window/logMessage",
+        params: null,
+      });
+    }).not.toThrow();
+  });
 
-    // Wait for initialization
-    await client.initializePromise;
+  describe("status notifications", () => {
+    beforeEach(() => {
+      store.set(copilotStatusState, {
+        busy: false,
+        kind: null,
+        message: null,
+      });
+    });
 
-    // Wait for next tick to let the configuration be sent
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    it.each([
+      [
+        "omitted message and unknown kind",
+        { kind: "Inactive", busy: false },
+        { busy: false, kind: "Inactive", message: null, status: undefined },
+      ],
+      [
+        "omitted busy",
+        { kind: "Error", message: "Not signed in" },
+        {
+          busy: false,
+          kind: "Error",
+          message: "Not signed in",
+          status: undefined,
+        },
+      ],
+      [
+        "a fully populated payload",
+        {
+          busy: true,
+          kind: "Normal",
+          message: "Working",
+          status: "InProgress",
+        },
+        {
+          busy: true,
+          kind: "Normal",
+          message: "Working",
+          status: "InProgress",
+        },
+      ],
+      [
+        "unusable field types",
+        { busy: 1, kind: 2, message: 3, status: 4 },
+        { busy: false, kind: null, message: null, status: undefined },
+      ],
+    ])(
+      "normalizes rather than discards a status update with %s",
+      async (_label, params, expected) => {
+        const client = new CopilotLanguageServerClient({
+          rootUri: "file:///test",
+          workspaceFolders: null,
+          transport: mockTransport,
+        });
+        await client.initializePromise;
+
+        mockTransport.emit({
+          jsonrpc: "2.0",
+          method: "didChangeStatus",
+          params,
+        });
+
+        expect(store.get(copilotStatusState)).toEqual(expected);
+      },
+    );
 
-    // Verify that notify was called with all settings
-    expect(notifySpy).toHaveBeenCalledWith("workspace/didChangeConfiguration", {
-      settings: copilotSettings,
+    it("ignores a status update with non-object params", async () => {
+      const client = new CopilotLanguageServerClient({
+        rootUri: "file:///test",
+        workspaceFolders: null,
+        transport: mockTransport,
+      });
+      await client.initializePromise;
+
+      mockTransport.emit({
+        jsonrpc: "2.0",
+        method: "statusNotification",
+        params: "not-an-object",
+      });
+
+      expect(store.get(copilotStatusState)).toEqual({
+        busy: false,
+        kind: null,
+        message: null,
+      });
     });
   });
 });