From 26fdec4e57b640ab0a39aa5a046f9730edb03fe5 Mon Sep 17 00:00:00 2001
From: Stefano Furi <stefano.furi@jetbrains.com>
Date: Thu, 4 Sep 2025 14:27:10 +0200
Subject: [PATCH] feat: add LSP, WebSocket based autocompletion logic

- Related to KTL-2807: Enhance autocompletion
- LSP-based completion both in REST and WS APIs.
---
 src/config.js                               |   4 +
 src/utils/websockets/webSocketConnection.js | 155 ++++++++++++++++++++
 src/webdemo-api.js                          |  39 +++++
 3 files changed, 198 insertions(+)
 create mode 100644 src/utils/websockets/webSocketConnection.js

diff --git a/src/config.js b/src/config.js
index 86a5b5b6..222b437f 100644
--- a/src/config.js
+++ b/src/config.js
@@ -60,6 +60,10 @@ export const API_URLS = {
   COMPLETE(version) {
     return `${this.server}/api/${version}/compiler/complete`;
   },
+  // TODO(KTL-3773): support multiple Kotlin versions with LSP-based completions
+  COMPLETE_WEBSOCKET(/* version */) {
+    return `ws://${this.server}/api/complete`;
+  },
   get VERSIONS() {
     return `${this.server}/versions`;
   },
diff --git a/src/utils/websockets/webSocketConnection.js b/src/utils/websockets/webSocketConnection.js
new file mode 100644
index 00000000..1ca757df
--- /dev/null
+++ b/src/utils/websockets/webSocketConnection.js
@@ -0,0 +1,155 @@
+import { API_URLS } from '../../config';
+
+const WS_COMPLETIONS_TIMEOUT_DURATION = 10_000;
+const MAX_RECONNECT_RETRIES = 5;
+
+const ConnectionState = {
+  DISCONNECTED: 'disconnected',
+  CONNECTED: 'connected',
+  CONNECTING: 'connecting',
+};
+
+const state = {
+  connection: null,
+  connectionState: ConnectionState.DISCONNECTED,
+  retryAttempt: 0,
+};
+
+const pendingRequest = new Map();
+let nextId = 1;
+let queuedRequest = null;
+
+export function completionCallback(callback, project, line, ch) {
+  const id = String(nextId++);
+
+  const timeout = setTimeout(() => {
+    pendingRequest.delete(id);
+    callback([]);
+  }, WS_COMPLETIONS_TIMEOUT_DURATION);
+
+  pendingRequest.set(id, { callback, timeout });
+
+  const files = project['files'];
+
+  const message = {
+    requestId: id,
+    completionRequest: { files },
+    line,
+    ch,
+  };
+
+  if (isConnectionReady()) {
+    safeSend(state.connection, message);
+    try {
+      state.connection.send(JSON.stringify(message));
+    } catch (error) {
+      queuedRequest = message;
+      handleDisconnect();
+    }
+  } else {
+    queuedRequest = message;
+    connect();
+  }
+}
+
+function connect() {
+  if (
+    isConnectionReady() ||
+    state.connectionState === ConnectionState.CONNECTING
+  ) {
+    return;
+  }
+
+  state.connectionState = ConnectionState.CONNECTING;
+
+  try {
+    const ws = new WebSocket(API_URLS.COMPLETE_WEBSOCKET());
+    state.connection = ws;
+
+    ws.onopen = () => {
+      state.connectionState = ConnectionState.CONNECTED;
+      state.retryAttempt = 0;
+
+      if (queuedRequest) safeSend(ws, queuedRequest);
+    };
+
+    ws.onmessage = handleMessage;
+
+    ws.onerror = () => {
+      handleDisconnect(true);
+    };
+
+    ws.onclose = (event) => {
+      const tryReconnect = event.code !== 1000;
+      handleDisconnect(tryReconnect);
+    };
+  } catch {
+    handleDisconnect(true);
+  }
+  return true;
+}
+
+function handleMessage(msg) {
+  let payload;
+  try {
+    payload = JSON.parse(msg.data);
+  } catch {
+    return;
+  }
+
+  if (payload.requestId) {
+    const entry = pendingRequest.get(payload.requestId);
+    if (!entry) return;
+
+    clearTimeout(entry.timeout);
+    pendingRequest.delete(payload.requestId);
+
+    if (payload.completions) {
+      entry.callback(payload.completions);
+    } else {
+      console.error(payload.message);
+      entry.callback([]);
+    }
+  }
+}
+
+function handleDisconnect(tryReconnect) {
+  if (isConnectionReady()) return;
+  state.connectionState = ConnectionState.DISCONNECTED;
+  pendingRequest.forEach((_, entry) => {
+    entry.callback([]);
+    clearTimeout(entry.timeout);
+  });
+
+  if (tryReconnect && state.retryAttempt <= MAX_RECONNECT_RETRIES) {
+    state.retryAttempt++;
+    connect();
+  } else {
+    try {
+      state.connection.close();
+    } finally {
+      state.connection = null;
+      state.retryAttempt = 0;
+    }
+  }
+}
+
+function safeSend(ws, message) {
+  try {
+    ws.send(JSON.stringify(message));
+    if (queuedRequest && queuedRequest.requestId === message.requestId)
+      queuedRequest = null;
+  } catch (err) {
+    console.warn(`Failed to send ${message}: `, err);
+    queuedRequest = message;
+    handleDisconnect(true);
+  }
+}
+
+function isConnectionReady() {
+  return (
+    state.connection &&
+    state.connectionState === ConnectionState.CONNECTED &&
+    state.connection.readyState === WebSocket.OPEN
+  );
+}
diff --git a/src/webdemo-api.js b/src/webdemo-api.js
index 64bade9d..9283a666 100644
--- a/src/webdemo-api.js
+++ b/src/webdemo-api.js
@@ -9,6 +9,7 @@ import {
   processJUnitResults,
   processJVMOutput,
 } from './view/output-view';
+import { completionCallback } from './utils/websockets/webSocketConnection';
 
 /**
  * @typedef {Object} KotlinVersion
@@ -225,6 +226,44 @@ export default class WebDemoApi {
     });
   }
 
+  /**
+   * Request for getting list of different completion proposals through WebSockets
+   *
+   * Note: currently `compilerVersion` is not used as it is not supported by lsp-based
+   * completions. Please refer to KTL-3773.
+   *
+   * @param code - string code
+   * @param cursor - cursor position in code
+   * @param compilerVersion - string kotlin compiler
+   * @param hiddenDependencies   - read only additional files
+   * @param platform - kotlin platform {@see TargetPlatform}
+   * @param callback
+   */
+  static getWSAutoCompletion(
+    code,
+    cursor,
+    compilerVersion,
+    platform,
+    hiddenDependencies,
+    callback,
+  ) {
+    const { line, ch, ...options } = cursor;
+    const files = [buildFileObject(code, DEFAULT_FILE_NAME)].concat(
+      hiddenDependencies.map((file, index) =>
+        buildFileObject(file, `hiddenDependency${index}.kt`),
+      ),
+    );
+
+    const project = {
+      args: '',
+      files,
+      confType: platform.id,
+      ...(options || {}),
+    };
+
+    completionCallback(callback, project, line, ch);
+  }
+
   /**
    * Request for getting errors of current file
    *