@marimo-team/frontend 0.25.2-dev97 → 0.25.2

package/src/core/codemirror/copilot/__tests__/rpc-cancellation.test.ts

@@ -0,0 +1,201 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  JSONRPCClient,
+  type JSONRPCMessage,
+  type Transport,
+} from "@marimo-team/codemirror-languageserver";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+describe("LSP request cancellation", () => {
+  let client: JSONRPCClient;
+  let emit: (message: JSONRPCMessage) => void;
+  let disconnect: (error: Error) => void;
+  let connected: ReturnType<typeof Promise.withResolvers<undefined>>;
+  let transport: Transport & {
+    send: ReturnType<typeof vi.fn<(message: JSONRPCMessage) => void>>;
+  };
+
+  beforeEach(() => {
+    vi.useFakeTimers();
+    connected = Promise.withResolvers<undefined>();
+    transport = {
+      connect: () => connected.promise,
+      close: vi.fn(),
+      send: vi.fn<(message: JSONRPCMessage) => void>(),
+      onMessage: (handler) => {
+        emit = handler;
+        return vi.fn();
+      },
+      onClose: (handler) => {
+        disconnect = handler;
+        return vi.fn();
+      },
+    };
+    client = new JSONRPCClient(transport);
+  });
+
+  afterEach(() => {
+    client.close();
+    vi.useRealTimers();
+  });
+
+  async function connect() {
+    connected.resolve(undefined);
+    await vi.advanceTimersByTimeAsync(0);
+  }
+
+  it("does not send an already-aborted request", async () => {
+    await connect();
+    const controller = new AbortController();
+    controller.abort();
+    await expect(
+      client.request("complete", {}, 1000, controller.signal),
+    ).rejects.toMatchObject({ name: "AbortError" });
+    expect(transport.send).not.toHaveBeenCalled();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("cancels while connecting without sending either request or cancellation", async () => {
+    const controller = new AbortController();
+    const response = client.request("complete", {}, 1000, controller.signal);
+    controller.abort();
+    await expect(response).rejects.toMatchObject({ name: "AbortError" });
+    await connect();
+    expect(transport.send).not.toHaveBeenCalled();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("rejects a failed connection and removes the waiting request's abort listener", async () => {
+    const controller = new AbortController();
+    const remove = vi.spyOn(controller.signal, "removeEventListener");
+    const response = client.request("complete", {}, 1000, controller.signal);
+    const settled = response.catch((error: unknown) => error);
+    connected.reject(new Error("connection refused"));
+    await expect(settled).resolves.toMatchObject({
+      message: "connection refused",
+    });
+    expect(remove).toHaveBeenCalledExactlyOnceWith(
+      "abort",
+      expect.any(Function),
+    );
+    controller.abort();
+    expect(transport.send).not.toHaveBeenCalled();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("cancels exactly the matching request, ignores late replies, and leaves others active", async () => {
+    await connect();
+    const controller = new AbortController();
+    const first = client.request(
+      "complete",
+      { position: 1 },
+      1000,
+      controller.signal,
+    );
+    const second = client.request("complete", { position: 2 }, 1000);
+    controller.abort();
+    controller.abort();
+    await expect(first).rejects.toMatchObject({ name: "AbortError" });
+    emit({ jsonrpc: "2.0", id: 0, result: "stale" });
+    emit({ jsonrpc: "2.0", id: 1, result: "latest" });
+    await expect(second).resolves.toBe("latest");
+    expect(transport.send.mock.calls).toEqual([
+      [{ jsonrpc: "2.0", id: 0, method: "complete", params: { position: 1 } }],
+      [{ jsonrpc: "2.0", id: 1, method: "complete", params: { position: 2 } }],
+      [{ jsonrpc: "2.0", method: "$/cancelRequest", params: { id: 0 } }],
+    ]);
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it.each([
+    "success",
+    "error",
+    "timeout",
+    "disconnect",
+    "close",
+    "send failure",
+  ])("removes the abort listener and timer after %s", async (outcome) => {
+    await connect();
+    const controller = new AbortController();
+    const remove = vi.spyOn(controller.signal, "removeEventListener");
+    if (outcome === "send failure") {
+      transport.send.mockImplementationOnce(() => {
+        throw new Error("send failed");
+      });
+    }
+    const response = client.request("complete", {}, 1000, controller.signal);
+    // Attach the rejection handler before advancing timers or dropping the connection.
+    const settled = response.then(
+      (value) => ({ status: "fulfilled", value }),
+      (error: unknown) => ({ status: "rejected", error }),
+    );
+    if (outcome === "success") {
+      emit({ jsonrpc: "2.0", id: 0, result: "ok" });
+    }
+    if (outcome === "error") {
+      emit({
+        jsonrpc: "2.0",
+        id: 0,
+        error: { code: -32603, message: "failed" },
+      });
+    }
+    if (outcome === "timeout") {
+      await vi.advanceTimersByTimeAsync(1000);
+    }
+    if (outcome === "disconnect") {
+      disconnect(new Error("disconnected"));
+    }
+    if (outcome === "close") {
+      client.close();
+    }
+    const expectedError = expect.any(Error);
+    expect(await settled).toEqual(
+      outcome === "success"
+        ? { status: "fulfilled", value: "ok" }
+        : { status: "rejected", error: expectedError },
+    );
+    expect(remove).toHaveBeenCalledExactlyOnceWith(
+      "abort",
+      expect.any(Function),
+    );
+    expect(vi.getTimerCount()).toBe(0);
+    const calls = transport.send.mock.calls.length;
+    controller.abort();
+    expect(transport.send).toHaveBeenCalledTimes(calls);
+  });
+
+  it("sends exactly one cancellation at the timeout boundary and ignores a late reply", async () => {
+    await connect();
+    const controller = new AbortController();
+    const settled = client
+      .request("complete", {}, 1000, controller.signal)
+      .catch((error: unknown) => error);
+    await vi.advanceTimersByTimeAsync(999);
+    expect(transport.send.mock.calls).toEqual([
+      [{ jsonrpc: "2.0", id: 0, method: "complete", params: {} }],
+    ]);
+    await vi.advanceTimersByTimeAsync(1);
+    expect(transport.send.mock.calls).toEqual([
+      [{ jsonrpc: "2.0", id: 0, method: "complete", params: {} }],
+      [{ jsonrpc: "2.0", method: "$/cancelRequest", params: { id: 0 } }],
+    ]);
+    emit({ jsonrpc: "2.0", id: 0, result: "late" });
+    controller.abort();
+    await expect(settled).resolves.toMatchObject({ code: -32001 });
+    expect(transport.send).toHaveBeenCalledTimes(2);
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("settles cancellation even if sending the cancellation notification fails", async () => {
+    await connect();
+    const controller = new AbortController();
+    const response = client.request("complete", {}, 1000, controller.signal);
+    transport.send.mockImplementationOnce(() => {
+      throw new Error("connection gone");
+    });
+    controller.abort();
+    await expect(response).rejects.toMatchObject({ name: "AbortError" });
+    expect(vi.getTimerCount()).toBe(0);
+  });
+});

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

@@ -39,7 +39,7 @@ const copilotCompartment = new Compartment();
 const logger = Logger.get("@github/copilot-language-server");
 
 const commonInlineCompletionConfig = {
-  delay: 300, // default is 500ms
+  delay: 150,
   includeKeymap: true,
   events: {
     // Only show suggestions when the editor is focused
@@ -79,7 +79,7 @@ export const copilotBundle = (config: CompletionConfig): Extension => {
     extensions.push(
       inlineCompletion({
         ...commonInlineCompletionConfig,
-        fetchFn: async (state, _abortSignal, view) => {
+        fetchFn: async (state, signal, view) => {
           const currentPosition = state.selection.main.head;
           if (!isCopilotEnabled()) {
             return "";
@@ -94,16 +94,22 @@ export const copilotBundle = (config: CompletionConfig): Extension => {
             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 client = getCopilotClient();
+          const allCode = await client.flushDocumentChanges(
+            state.doc.toString(),
+          );
+          if (signal.aborted || state.doc !== view.state.doc) {
+            return "";
+          }
           const request = getCopilotRequest(state, allCode);
-          const response = await getCopilotClient().getCompletion(request);
+          const response = await client.getCompletion(request, signal);
 
           // If we are at a new position, ignore the response
-          if (currentPosition !== view.state.selection.main.head) {
+          if (
+            signal.aborted ||
+            state.doc !== view.state.doc ||
+            currentPosition !== view.state.selection.main.head
+          ) {
             Logger.debug("ignoring response because of new position");
             return "";
           }
@@ -119,7 +125,7 @@ export const copilotBundle = (config: CompletionConfig): Extension => {
     extensions.push(
       inlineCompletion({
         ...commonInlineCompletionConfig,
-        fetchFn: async (state, _abortSignal, view) => {
+        fetchFn: async (state, signal, view) => {
           const currentPosition = state.selection.main.head;
 
           if (state.doc.length === 0) {
@@ -146,10 +152,15 @@ export const copilotBundle = (config: CompletionConfig): Extension => {
           let res = await API.post<AiInlineCompletionRequest, string>(
             "/ai/inline_completion",
             { prefix, suffix, language },
+            { signal },
           );
 
           // If we are at a new position, ignore the response
-          if (currentPosition !== view.state.selection.main.head) {
+          if (
+            signal.aborted ||
+            state.doc !== view.state.doc ||
+            currentPosition !== view.state.selection.main.head
+          ) {
             Logger.debug("ignoring response because of new position");
             return "";
           }

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

@@ -1,7 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { LanguageServerClient } from "@marimo-team/codemirror-languageserver";
-import { throttle } from "lodash-es";
 import type {
   CompletionItem,
   CompletionList,
@@ -58,6 +57,7 @@ interface UntypedLanguageServerMethods {
     method: string,
     params: unknown,
     timeout: number,
+    signal?: AbortSignal,
   ) => Promise<unknown>;
   notify: (method: string, params: unknown) => Promise<void>;
 }
@@ -113,6 +113,8 @@ function isLogMessageParams(
  */
 export class CopilotLanguageServerClient extends LanguageServerClient {
   private documentVersion = 0;
+  private pendingChange: DidChangeTextDocumentParams | undefined;
+  private completionController: AbortController | undefined;
   private openDocumentCount = 0;
   private lastOpenedDocument: DidOpenTextDocumentParams | undefined;
   private copilotSettings: Record<string, unknown> = {};
@@ -142,9 +144,11 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
    */
   public async reInitialize(): Promise<void> {
     logger.log("#reInitialize: Re-initializing LSP connection");
+    this.completionController?.abort();
     this.initializePromise = this.initialize();
     await this.initializePromise;
     await this.sendConfiguration();
+    this.applyPendingChange();
     if (this.openDocumentCount > 0 && this.lastOpenedDocument) {
       await this.sendNotification(
         "textDocument/didOpen",
@@ -168,12 +172,14 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
   private async _request<Method extends keyof LSPRequestMap>(
     method: Method,
     params: LSPRequestMap[Method][0],
+    signal?: AbortSignal,
   ): Promise<LSPRequestMap[Method][1]> {
     assertHasLanguageServerRpc(this);
     return (await this.request(
       method,
       params,
       REQUEST_TIMEOUT_MS,
+      signal,
     )) as LSPRequestMap[Method][1];
   }
 
@@ -230,6 +236,10 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
       return;
     }
     this.openDocumentCount--;
+    if (this.openDocumentCount === 0) {
+      this.pendingChange = undefined;
+      this.completionController?.abort();
+    }
     return super.textDocumentDidClose(params);
   }
 
@@ -281,27 +291,72 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
       );
     }
 
-    const text = getCodes(change.text);
-    const version = ++this.documentVersion;
-    // `reInitialize` replays this as a didOpen, so keep the original languageId
+    this.completionController?.abort();
+    this.pendingChange = {
+      ...params,
+      textDocument: VersionedTextDocumentIdentifier.create(
+        params.textDocument.uri,
+        ++this.documentVersion,
+      ),
+    };
+  }
+
+  private applyPendingChange(
+    assembledText?: string,
+  ): DidChangeTextDocumentParams | undefined {
+    const pending = this.pendingChange;
+    if (!pending) {
+      return;
+    }
+    this.pendingChange = undefined;
+    const text = assembledText ?? getCodes(pending.contentChanges[0].text);
     this.lastOpenedDocument = {
       textDocument: {
-        uri: params.textDocument.uri,
+        ...pending.textDocument,
         languageId:
           this.lastOpenedDocument?.textDocument.languageId ??
           DEFAULT_LANGUAGE_ID,
-        version,
         text,
       },
     };
-    return super.textDocumentDidChange({
-      ...params,
-      contentChanges: [{ text: text }],
-      textDocument: VersionedTextDocumentIdentifier.create(
-        params.textDocument.uri,
-        version,
-      ),
-    });
+    return { ...pending, contentChanges: [{ text }] };
+  }
+
+  /** Assemble and send the latest notebook once, just before requesting a suggestion. */
+  public async flushDocumentChanges(currentCode?: string): Promise<string> {
+    // Another cell may have changed since the focused editor last typed.
+    // The focused cell must be last in the shared document for cursor mapping.
+    const assembledText =
+      currentCode !== undefined && this.lastOpenedDocument
+        ? getCodes(currentCode)
+        : undefined;
+    if (
+      currentCode !== undefined &&
+      this.lastOpenedDocument &&
+      (this.pendingChange
+        ? this.pendingChange.contentChanges[0].text !== currentCode
+        : assembledText !== this.lastOpenedDocument.textDocument.text)
+    ) {
+      this.pendingChange = {
+        textDocument: VersionedTextDocumentIdentifier.create(
+          this.lastOpenedDocument.textDocument.uri,
+          ++this.documentVersion,
+        ),
+        contentChanges: [{ text: currentCode }],
+      };
+    }
+    const pending = this.pendingChange;
+    const change = this.applyPendingChange(assembledText);
+    if (change) {
+      try {
+        await super.textDocumentDidChange(change);
+      } catch (error) {
+        // Preserve newer edits if the send failed while another change arrived.
+        this.pendingChange ??= pending;
+        throw error;
+      }
+    }
+    return this.lastOpenedDocument?.textDocument.text ?? "";
   }
 
   public override textDocumentHover(params: HoverParams): Promise<Hover> {
@@ -351,49 +406,32 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     }
   }
 
-  private getCompletionInternal = async (
-    params: InlineCompletionParams,
-    version: number,
-  ): Promise<InlineCompletionList | InlineCompletionItem[] | null> => {
-    return await this._request("textDocument/inlineCompletion", {
-      ...params,
-      textDocument: {
-        ...VersionedTextDocumentIdentifier.create(
-          params.textDocument.uri,
-          version,
-        ),
-      },
-    });
-  };
-
-  // 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,
-  );
-
   public async getCompletion(
     params: InlineCompletionParams,
+    signal?: AbortSignal,
   ): Promise<InlineCompletionList | InlineCompletionItem[] | null> {
-    if (this.isDisabled()) {
+    if (this.isDisabled() || signal?.aborted || this.openDocumentCount === 0) {
       return null;
     }
 
+    this.completionController?.abort();
+    const controller = new AbortController();
+    this.completionController = controller;
+    const abort = () => controller.abort();
+    signal?.addEventListener("abort", abort, { once: true });
     const requestVersion = this.documentVersion;
-
-    // If version is 0, it means the document hasn't been opened yet
-    if (requestVersion === 0) {
-      return null;
-    }
-
-    // Start a loading indicator
     setGitHubCopilotLoadingVersion(requestVersion);
-    let response: InlineCompletionList | InlineCompletionItem[] | null;
     try {
-      response = await this.throttledGetCompletionInternal(
+      await this.flushDocumentChanges();
+      if (
+        controller.signal.aborted ||
+        this.isDisabled() ||
+        requestVersion === 0
+      ) {
+        return null;
+      }
+      const response = await this._request(
+        "textDocument/inlineCompletion",
         {
           ...params,
           textDocument: VersionedTextDocumentIdentifier.create(
@@ -401,24 +439,35 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
             requestVersion,
           ),
         },
-        requestVersion,
+        controller.signal,
       );
+      if (
+        controller.signal.aborted ||
+        this.isDisabled() ||
+        requestVersion !== this.documentVersion
+      ) {
+        return null;
+      }
+      return response;
     } catch (error) {
-      // A suggestion that times out should fail quietly rather than surface as
-      // an unhandled rejection out of the inline-completion fetcher.
-      logger.warn("#getCompletion: Failed to fetch completions", error);
+      if (!controller.signal.aborted) {
+        logger.warn("#getCompletion: Failed to fetch completions", error);
+      }
       return null;
     } finally {
-      // 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;
+      signal?.removeEventListener("abort", abort);
+      // A superseded request can have the same document version as its successor.
+      if (this.completionController === controller) {
+        this.completionController = undefined;
+        clearGitHubCopilotLoadingVersion(requestVersion);
+      }
     }
+  }
 
-    return response ?? null;
+  public override close(): void {
+    this.pendingChange = undefined;
+    this.completionController?.abort();
+    super.close();
   }
 
   /**

package/src/core/codemirror/language/__tests__/sql.test.ts

@@ -8,6 +8,7 @@ import { EditorState, type Extension } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { requestId } from "@/__tests__/branded";
 import type { CellId } from "@/core/cells/ids";
 import type {
   CompletionConfig,
@@ -1106,7 +1107,7 @@ describe("CustomSqlParser", () => {
     const request = vi.spyOn(ValidateSQL, "request").mockResolvedValue({
       error: null,
       parse_result: { success: false, errors: [error] },
-      request_id: "request-id",
+      request_id: requestId("request-id"),
       validate_result: null,
     });
     const state = EditorState.create({
@@ -1150,7 +1151,7 @@ describe("CustomSqlParser", () => {
     vi.spyOn(ValidateSQL, "request").mockResolvedValue({
       error: null,
       parse_result: { success: false, errors: [error] },
-      request_id: "request-id",
+      request_id: requestId("request-id"),
       validate_result: null,
     });
     const state = EditorState.create({
@@ -1194,7 +1195,7 @@ describe("CustomSqlParser", () => {
     const request = vi.spyOn(ValidateSQL, "request").mockResolvedValue({
       error: null,
       parse_result: { success: false, errors: [error] },
-      request_id: "request-id",
+      request_id: requestId("request-id"),
       validate_result: null,
     });
     const state = createNamedConnectionState({