@marimo-team/frontend 0.23.16-dev5 → 0.23.16-dev50

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

@@ -7,6 +7,7 @@ import type {
   CompletionList,
   CompletionParams,
   DidChangeTextDocumentParams,
+  DidCloseTextDocumentParams,
   DidOpenTextDocumentParams,
   Hover,
   HoverParams,
@@ -16,7 +17,9 @@ import type {
 } from "vscode-languageserver-protocol";
 import { VersionedTextDocumentIdentifier } from "vscode-languageserver-protocol";
 import { store } from "@/core/state/jotai";
+import { invariant } from "@/utils/invariant";
 import { Logger } from "@/utils/Logger";
+import { hasFunctionProperty, isRecord } from "@/utils/records";
 import { getCodes } from "./getCodes";
 import {
   clearGitHubCopilotLoadingVersion,
@@ -30,42 +33,88 @@ import type {
   GitHubCopilotStatusNotificationParams,
   GitHubCopilotStatusResult,
 } from "./types";
+import type { LanguageAdapterType } from "../language/types";
 
 const logger = Logger.get("@github/copilot-language-server");
+const REQUEST_TIMEOUT_MS = 10_000;
+// Only used for the synthetic didOpen sent when a change arrives before any open
+const DEFAULT_LANGUAGE_ID: LanguageAdapterType = "python";
+type NoParams = Record<string, never>;
 
 // A map of request methods and their parameters and return types
 export interface LSPRequestMap {
-  checkStatus: [{}, GitHubCopilotStatusResult];
-  signIn: [{}, GitHubCopilotSignInInitiateResult];
+  checkStatus: [NoParams, GitHubCopilotStatusResult];
+  signIn: [NoParams, GitHubCopilotSignInInitiateResult];
   signInConfirm: [GitHubCopilotSignInConfirmParams, GitHubCopilotStatusResult];
-  signOut: [{}, GitHubCopilotStatusResult];
+  signOut: [NoParams, GitHubCopilotStatusResult];
   "textDocument/inlineCompletion": [
     InlineCompletionParams,
     InlineCompletionList | InlineCompletionItem[] | null,
   ];
 }
 
-export interface LSPEventMap {
-  statusNotification: GitHubCopilotStatusNotificationParams;
-  didChangeStatus: GitHubCopilotStatusNotificationParams;
-  "window/logMessage": { type: number; message: string };
+interface UntypedLanguageServerMethods {
+  request: (
+    method: string,
+    params: unknown,
+    timeout: number,
+  ) => Promise<unknown>;
+  notify: (method: string, params: unknown) => Promise<void>;
 }
 
-export type EnhancedNotification = {
-  [key in keyof LSPEventMap]: {
-    jsonrpc: "2.0";
-    id?: null | undefined;
-    method: key;
-    params: LSPEventMap[key];
+/**
+ * `LanguageServerClient#request`/`#notify` are protected and generic over the
+ * library's own `LSPRequestMap`/`LSPNotifyMap`
+ * This asserts the two inherited methods we need are present.
+ */
+function assertHasLanguageServerRpc(
+  value: unknown,
+): asserts value is UntypedLanguageServerMethods {
+  invariant(
+    isRecord(value) &&
+      hasFunctionProperty(value, "request") &&
+      hasFunctionProperty(value, "notify"),
+    "LanguageServerClient is missing request/notify",
+  );
+}
+
+/**
+ * Normalizes rather than validates: Copilot omits `message`/`busy` on most
+ * notifications and ships `kind` values we don't know about, so rejecting on an
+ * exact shape match would silently freeze the status indicator.
+ */
+function parseCopilotStatusNotificationParams(
+  value: unknown,
+): GitHubCopilotStatusNotificationParams | null {
+  if (!isRecord(value)) {
+    return null;
+  }
+  const { busy, kind, message, status } = value;
+  return {
+    busy: typeof busy === "boolean" ? busy : false,
+    kind: typeof kind === "string" ? kind : null,
+    message: typeof message === "string" ? message : null,
+    status: typeof status === "string" ? status : undefined,
   };
-}[keyof LSPEventMap];
+}
+
+function isLogMessageParams(
+  value: unknown,
+): value is { type: number; message: string } {
+  return (
+    isRecord(value) &&
+    typeof value.type === "number" &&
+    typeof value.message === "string"
+  );
+}
 
 /**
  * A client for the Copilot language server.
  */
 export class CopilotLanguageServerClient extends LanguageServerClient {
   private documentVersion = 0;
-  private hasOpenedDocument = false;
+  private openDocumentCount = 0;
+  private lastOpenedDocument: DidOpenTextDocumentParams | undefined;
   private copilotSettings: Record<string, unknown> = {};
 
   constructor(
@@ -80,10 +129,11 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
   }
 
   private attachInitializeListener() {
-    // Send configuration after initialization
-    this.initializePromise.then(() => {
-      this.sendConfiguration();
-    });
+    void this.initializePromise
+      .then(() => this.sendConfiguration())
+      .catch((error: unknown) => {
+        logger.warn("#initialize: Failed to send configuration", error);
+      });
   }
 
   /**
@@ -95,6 +145,12 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     this.initializePromise = this.initialize();
     await this.initializePromise;
     await this.sendConfiguration();
+    if (this.openDocumentCount > 0 && this.lastOpenedDocument) {
+      await this.sendNotification(
+        "textDocument/didOpen",
+        this.lastOpenedDocument,
+      );
+    }
   }
 
   private async sendConfiguration() {
@@ -103,7 +159,9 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     if (!settings || Object.keys(settings).length === 0) {
       return;
     }
-    await this.notify("workspace/didChangeConfiguration", { settings });
+    await this.sendNotification("workspace/didChangeConfiguration", {
+      settings,
+    });
     logger.debug("#sendConfiguration: Configuration sent", settings);
   }
 
@@ -111,20 +169,21 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     method: Method,
     params: LSPRequestMap[Method][0],
   ): Promise<LSPRequestMap[Method][1]> {
-    return await (
-      this as unknown as {
-        request: (
-          method: Method,
-          params: LSPRequestMap[Method][0],
-        ) => Promise<LSPRequestMap[Method][1]>;
-      }
-    ).request(method, params);
+    assertHasLanguageServerRpc(this);
+    return (await this.request(
+      method,
+      params,
+      REQUEST_TIMEOUT_MS,
+    )) as LSPRequestMap[Method][1];
   }
 
-  // oxlint-disable-next-line typescript/no-explicit-any
-  override async notify(method: any, params: any): Promise<any> {
+  private async sendNotification(
+    method: string,
+    params: unknown,
+  ): Promise<void> {
     logger.debug("#notify", method, params);
-    return super.notify(method, params);
+    assertHasLanguageServerRpc(this);
+    return this.notify(method, params);
   }
 
   override getInitializationOptions() {
@@ -151,14 +210,29 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
 
   override async textDocumentDidOpen(
     params: DidOpenTextDocumentParams,
-  ): Promise<DidOpenTextDocumentParams> {
+  ): Promise<boolean> {
     if (this.isDisabled()) {
-      return params;
+      return false;
     }
-    this.hasOpenedDocument = true;
+    this.openDocumentCount++;
+    this.lastOpenedDocument = params;
+    this.documentVersion = Math.max(
+      this.documentVersion,
+      params.textDocument.version,
+    );
     return super.textDocumentDidOpen(params);
   }
 
+  override async textDocumentDidClose(
+    params: DidCloseTextDocumentParams,
+  ): Promise<void> {
+    if (this.openDocumentCount === 0) {
+      return;
+    }
+    this.openDocumentCount--;
+    return super.textDocumentDidClose(params);
+  }
+
   override async textDocumentCompletion(
     params: CompletionParams,
   ): Promise<CompletionList | CompletionItem[]> {
@@ -168,18 +242,27 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
 
   override async textDocumentDidChange(
     params: DidChangeTextDocumentParams,
-  ): Promise<DidChangeTextDocumentParams> {
+  ): Promise<void> {
     if (this.isDisabled()) {
-      return params;
+      return;
+    }
+
+    const change = params.contentChanges[0];
+    if (!change) {
+      logger.warn(
+        "#textDocumentDidChange: Ignoring an update with no content changes.",
+        params,
+      );
+      return;
     }
 
-    if (!this.hasOpenedDocument) {
+    if (this.openDocumentCount === 0) {
       await this.textDocumentDidOpen({
         textDocument: {
           uri: params.textDocument.uri,
-          languageId: "python",
+          languageId: DEFAULT_LANGUAGE_ID,
           version: params.textDocument.version,
-          text: params.contentChanges[0].text,
+          text: change.text,
         },
       });
     }
@@ -191,7 +274,6 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
         changes,
       );
     }
-    const change = changes[0];
     if ("range" in change) {
       logger.warn(
         "#textDocumentDidChange: Copilot doesn't support range changes.",
@@ -200,12 +282,24 @@ 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.lastOpenedDocument = {
+      textDocument: {
+        uri: params.textDocument.uri,
+        languageId:
+          this.lastOpenedDocument?.textDocument.languageId ??
+          DEFAULT_LANGUAGE_ID,
+        version,
+        text,
+      },
+    };
     return super.textDocumentDidChange({
       ...params,
       contentChanges: [{ text: text }],
       textDocument: VersionedTextDocumentIdentifier.create(
         params.textDocument.uri,
-        ++this.documentVersion,
+        version,
       ),
     });
   }
@@ -264,9 +358,11 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     return await this._request("textDocument/inlineCompletion", {
       ...params,
       textDocument: {
-        ...params.textDocument,
-        version: version,
-      } as VersionedTextDocumentIdentifier,
+        ...VersionedTextDocumentIdentifier.create(
+          params.textDocument.uri,
+          version,
+        ),
+      },
     });
   };
 
@@ -287,8 +383,6 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     }
 
     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) {
@@ -297,12 +391,27 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
 
     // 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);
+    let response: InlineCompletionList | InlineCompletionItem[] | null;
+    try {
+      response = await this.throttledGetCompletionInternal(
+        {
+          ...params,
+          textDocument: VersionedTextDocumentIdentifier.create(
+            params.textDocument.uri,
+            requestVersion,
+          ),
+        },
+        requestVersion,
+      );
+    } 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);
+      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) {
@@ -316,29 +425,32 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
    * Handle notifications from the Copilot language server.
    * Uses onNotification to listen for statusNotification, didChangeStatus, and window/logMessage.
    */
-  private handleNotification: Parameters<
-    LanguageServerClient["onNotification"]
-  >[0] = (notif) => {
-    if (!notif.params) {
+  private handleNotification = (notification: unknown): void => {
+    if (!isRecord(notification) || typeof notification.method !== "string") {
       return;
     }
 
-    const notification = notif as unknown as EnhancedNotification;
-
-    // Handle statusNotification
-    if (notification.method === "statusNotification") {
-      store.set(copilotStatusState, notification.params);
-    }
-
-    // Handle didChangeStatus
-    if (notification.method === "didChangeStatus") {
-      store.set(copilotStatusState, notification.params);
+    if (
+      notification.method === "statusNotification" ||
+      notification.method === "didChangeStatus"
+    ) {
+      const status = parseCopilotStatusNotificationParams(notification.params);
+      if (status) {
+        store.set(copilotStatusState, status);
+      } else {
+        logger.warn(
+          "#statusNotification: Ignoring a status update with no params",
+          notification.params,
+        );
+      }
+      return;
     }
 
-    // Handle window/logMessage
-    if (notification.method === "window/logMessage") {
-      const params = notification.params as { type: number; message: string };
-      const { type, message } = params;
+    if (
+      notification.method === "window/logMessage" &&
+      isLogMessageParams(notification.params)
+    ) {
+      const { type, message } = notification.params;
       // Map LSP log types to console methods
       // type: 1 = Error, 2 = Warning, 3 = Info, 4 = Log
       switch (type) {

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

@@ -1,10 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-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 { prettyError } from "@/utils/errors";
-import { Logger } from "@/utils/Logger";
+import { ReconnectingWebSocketTransport } from "@/core/lsp/transport";
 
 export interface LazyWebsocketTransportOptions {
   /**
@@ -24,224 +21,35 @@ export interface LazyWebsocketTransportOptions {
   showError: (title: string, description: string | React.ReactNode) => void;
 
   /**
-   * Number of retry attempts for connection.
+   * Number of connection attempts in each retry cycle.
    * @default 3
    */
   retries?: number;
 
   /**
-   * Delay between retry attempts in milliseconds.
+   * Delay between connection attempts in milliseconds.
    * @default 1000
    */
   retryDelayMs?: number;
-
-  /**
-   * Default timeout for sendData operations in milliseconds.
-   * Used when the caller does not provide an explicit timeout.
-   * @default 5000
-   */
-  maxTimeoutMs?: number;
-}
-
-interface Subscription {
-  event: "pending" | "notification" | "response" | "error";
-  handler: Parameters<Transport["subscribe"]>[1];
 }
 
 /**
- * A WebSocket transport that lazily connects after waiting for prerequisites.
- *
- * This transport:
- * - Waits for copilot to be enabled before connecting
- * - Waits for the websocket to be available before connecting
- * - Forwards subscriptions/unsubscriptions to the delegate
- * - Handles reconnection automatically
+ * Copilot's lazy transport is the shared reconnecting LSP transport with
+ * Copilot-specific readiness checks, retries, and user-facing errors.
  */
-export class LazyWebsocketTransport extends Transport {
-  private delegate: WebSocketTransport | undefined;
-  private pendingSubscriptions: Subscription[] = [];
-  private readonly options: Required<LazyWebsocketTransportOptions>;
-  private needsReInitialization = false;
-
-  /**
-   * Callback invoked after the transport reconnects following a close or connection failure.
-   * Used by the LSP client to re-run the initialize handshake on the new connection.
-   */
-  onReconnect?: () => Promise<void>;
-
+export class LazyWebsocketTransport extends ReconnectingWebSocketTransport {
   constructor(options: LazyWebsocketTransportOptions) {
-    super();
-    this.delegate = undefined;
-    this.pendingSubscriptions = [];
-    this.options = {
+    super({
+      getWsUrl: options.getWsUrl,
+      waitForConnection: options.waitForReady,
       retries: options.retries ?? 3,
       retryDelayMs: options.retryDelayMs ?? 1000,
-      maxTimeoutMs: options.maxTimeoutMs ?? 5000,
-      getWsUrl: options.getWsUrl,
-      waitForReady: options.waitForReady,
-      showError: options.showError,
-    };
-  }
-
-  override subscribe(...args: Parameters<Transport["subscribe"]>): void {
-    // Register handler on parent Transport
-    super.subscribe(...args);
-
-    const [event, handler] = args;
-
-    // Track the subscription
-    this.pendingSubscriptions.push({ event, handler });
-
-    // Also register on delegate if it exists
-    if (this.delegate) {
-      this.delegate.subscribe(event, handler);
-    }
-  }
-
-  override unsubscribe(
-    ...args: Parameters<Transport["unsubscribe"]>
-  ): import("events").EventEmitter | undefined {
-    // Unsubscribe from parent Transport
-    const result = super.unsubscribe(...args);
-
-    const [event, handler] = args;
-    // Also unsubscribe from delegate if it exists
-    if (this.delegate) {
-      this.delegate.unsubscribe(event, handler);
-    }
-
-    // Remove from pending subscriptions if handler is provided
-    if (handler) {
-      if (event) {
-        // Remove specific handler for specific event
-        const index = this.pendingSubscriptions.findIndex(
-          (sub) => sub.event === event && sub.handler === handler,
+      onConnectionFailure: (error) => {
+        options.showError(
+          "GitHub Copilot Connection Error",
+          `Failed to connect to GitHub Copilot. Please check your settings and try again.\n\n${prettyError(error)}`,
         );
-        if (index > -1) {
-          this.pendingSubscriptions.splice(index, 1);
-        }
-      } else {
-        // Remove handler from all events
-        this.pendingSubscriptions = this.pendingSubscriptions.filter(
-          (sub) => sub.handler !== handler,
-        );
-      }
-    } else if (event) {
-      // Remove all handlers for specific event
-      this.pendingSubscriptions = this.pendingSubscriptions.filter(
-        (sub) => sub.event !== event,
-      );
-    } else {
-      // Remove all subscriptions
-      this.pendingSubscriptions = [];
-    }
-
-    return result;
-  }
-
-  private createDelegate(): WebSocketTransport {
-    const delegate = new WebSocketTransport(this.options.getWsUrl());
-    // Register all pending subscriptions on the delegate
-    for (const { event, handler } of this.pendingSubscriptions) {
-      delegate.subscribe(event, handler);
-    }
-    return delegate;
-  }
-
-  private async tryConnect(): Promise<void> {
-    for (let attempt = 1; attempt <= this.options.retries; attempt++) {
-      try {
-        // Create delegate, if it doesn't exist
-        if (!this.delegate) {
-          this.delegate = this.createDelegate();
-        }
-        await this.delegate.connect();
-        Logger.log("Copilot#connect: Connected successfully");
-        return;
-      } catch (error) {
-        Logger.warn(
-          `Copilot#connect: Connection attempt ${attempt}/${this.options.retries} failed`,
-          error,
-        );
-        if (attempt === this.options.retries) {
-          this.delegate = undefined;
-          this.needsReInitialization = true;
-          // Show error toast on final retry
-          this.options.showError(
-            "GitHub Copilot Connection Error",
-            `Failed to connect to GitHub Copilot. Please check your settings and try again.\n\n${prettyError(error)}`,
-          );
-          throw error;
-        }
-        await new Promise((resolve) =>
-          setTimeout(resolve, this.options.retryDelayMs),
-        );
-      }
-    }
-  }
-
-  override async connect(): Promise<void> {
-    // Wait for all prerequisites to be ready
-    await this.options.waitForReady();
-
-    // Try connecting with retries
-    return this.tryConnect();
-  }
-
-  override close(): void {
-    this.delegate?.close();
-    this.delegate = undefined;
-    this.needsReInitialization = true;
-  }
-
-  override async sendData(
-    data: JSONRPCRequestData,
-    timeout: number | null | undefined,
-  ): Promise<unknown> {
-    // If delegate is undefined, try to reconnect
-    if (!this.delegate) {
-      Logger.log("Copilot#sendData: Delegate not initialized, reconnecting...");
-      try {
-        // Ensure prerequisites are ready before attempting connection
-        await this.options.waitForReady();
-        await this.tryConnect();
-      } catch (error) {
-        Logger.error("Copilot#sendData: Failed to reconnect transport", error);
-        throw new Error(
-          "Unable to connect to GitHub Copilot. Please check your settings and try again.",
-          { cause: error },
-        );
-      }
-
-      // Re-run LSP initialization handshake after reconnecting
-      if (this.needsReInitialization && this.onReconnect) {
-        Logger.log(
-          "Copilot#sendData: Re-initializing LSP after reconnection...",
-        );
-        try {
-          await this.onReconnect();
-          this.needsReInitialization = false;
-        } catch (error) {
-          // Close the uninitialized connection so the next attempt starts fresh
-          this.close();
-          Logger.error(
-            "Copilot#sendData: LSP re-initialization after reconnection failed",
-            error,
-          );
-          throw error;
-        }
-      }
-    }
-
-    // After reconnection, delegate should be initialized
-    if (!this.delegate) {
-      throw new Error(
-        "Failed to initialize GitHub Copilot connection. Please try again.",
-      );
-    }
-
-    // Use maxTimeoutMs as default when no timeout is provided
-    timeout = timeout ?? this.options.maxTimeoutMs;
-    return this.delegate.sendData(data, timeout);
+      },
+    });
   }
 }

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

@@ -35,11 +35,18 @@ export type GitHubCopilotStatus =
  */
 export type GitHubCopilotRequestStatus = "InProgress" | "Warning" | "Normal";
 
+/**
+ * Payload of Copilot's `statusNotification`/`didChangeStatus`. `kind` and
+ * `status` are open strings, not unions: this is a vendor shape we don't
+ * control, and an unrecognized value must not discard the notification.
+ */
 export interface GitHubCopilotStatusNotificationParams {
   busy: boolean;
-  kind: "Warning" | "Normal" | "Error" | null;
+  /** Known values: "Warning" | "Normal" | "Error" | "Inactive" */
+  kind: string | null;
   message: string | null;
-  status?: "InProgress" | "Warning" | "Normal";
+  /** Known values: "InProgress" | "Warning" | "Normal" */
+  status?: string;
 }
 
 export interface GitHubCopilotStatusResult {