Merge pull request #4 from orus-dev/bug-patches

Bug patches

src/client.ts

@@ -2,79 +2,189 @@ import { WebSocket } from "ws";
 import { LiveRunMove } from "./types";
 import { randomUUID } from "crypto";
 
+/* ──────────────────────────
+   State
+────────────────────────── */
+
 let socket: WebSocket | null = null;
-let cookies: string | undefined;
+let connecting: Promise<WebSocket> | null = null;
+let currentOrigin: string | null = null;
+
+let cookies: Record<string, string> | undefined;
 
 type PendingResolver = {
   resolve: (value: any) => void;
   reject: (err: any) => void;
+  timeout: NodeJS.Timeout;
 };
 
 const pending = new Map<string, PendingResolver>();
 
-export function getCookies(): string | undefined {
+/* ──────────────────────────
+   Cookies
+────────────────────────── */
+
+export function getCookies(): Record<string, string> | undefined {
   return cookies;
 }
 
-export function setCookies(c: string) {
+export function setCookies(c: Record<string, string>) {
   cookies = c;
 }
 
+/* ──────────────────────────
+   Utils
+────────────────────────── */
+
 function getWsOrigin(useLocalhost: boolean) {
   return useLocalhost
     ? "ws://localhost:3000/api/run"
     : "wss://dev-run.selimaj.dev/api/run";
 }
 
-function ensureSocket(useLocalhost: boolean): WebSocket {
-  if (socket && socket.readyState === WebSocket.OPEN) {
+function rejectAllPending(err: Error) {
+  for (const { reject, timeout } of pending.values()) {
+    clearTimeout(timeout);
+    reject(err);
+  }
+  pending.clear();
+}
+
+/* ──────────────────────────
+   Socket management
+────────────────────────── */
+
+async function ensureSocket(useLocalhost: boolean): Promise<WebSocket> {
+  const origin = getWsOrigin(useLocalhost);
+
+  if (
+    socket &&
+    socket.readyState === WebSocket.OPEN &&
+    currentOrigin === origin
+  ) {
     return socket;
   }
 
-  socket = new WebSocket(getWsOrigin(useLocalhost));
+  if (connecting) {
+    return connecting;
+  }
+
+  if (socket && currentOrigin !== origin) closeClient();
+
+  connecting = new Promise<WebSocket>((resolve, reject) => {
+    const ws = new WebSocket(origin);
+
+    ws.onopen = async () => {
+      try {
+        const authCookies = getCookies();
+        if (!authCookies) {
+          throw new Error("Cookies not set");
+        }
+
+        socket = ws;
+        currentOrigin = origin;
+        connecting = null;
+
+        await sendUnsafe({
+          type: "auth",
+          cookies: authCookies,
+        });
+
+        resolve(ws);
+      } catch (err) {
+        ws.close();
+        reject(err);
+      }
+    };
+
+    ws.onmessage = (e) => {
+      let msg: any;
+      try {
+        msg = JSON.parse(e.data.toString());
+      } catch {
+        console.warn("Invalid WS message:", e.data);
+        return;
+      }
 
-  socket.onmessage = (e) => {
-    const msg = JSON.parse(e.data.toString());
       const { requestId, ok, data, error } = msg;
-
       if (!requestId) return;
 
       const pendingReq = pending.get(requestId);
       if (!pendingReq) return;
 
       pending.delete(requestId);
+      clearTimeout(pendingReq.timeout);
 
       ok ? pendingReq.resolve(data) : pendingReq.reject(error);
     };
 
-  socket.onclose = () => {
-    socket = null;
+    ws.onerror = (err) => {
+      rejectAllPending(new Error("WebSocket error"));
+      connecting = null;
+      reject(err);
     };
 
-  return socket;
+    ws.onclose = () => {
+      rejectAllPending(new Error("WebSocket closed"));
+      socket = null;
+      currentOrigin = null;
+      connecting = null;
+    };
+
+    ws.on("ping", () => {
+      ws.pong();
+    });
+  });
+
+  return connecting;
 }
 
-function send<T>(
-  useLocalhost: boolean,
-  payload: Record<string, any>,
-): Promise<T> {
-  return new Promise((resolve, reject) => {
-    const ws = ensureSocket(useLocalhost);
+/* ──────────────────────────
+   RPC helpers
+────────────────────────── */
+
+async function sendUnsafe<T>(payload: Record<string, any>): Promise<T> {
+  if (!socket || socket.readyState !== WebSocket.OPEN) {
+    throw new Error("WebSocket is not open");
+  }
 
   const requestId = randomUUID();
 
-    pending.set(requestId, { resolve, reject });
+  return new Promise<T>((resolve, reject) => {
+    const timeout = setTimeout(() => {
+      pending.delete(requestId);
+      reject(new Error("WebSocket request timed out"));
+    }, 10_000);
 
-    ws.send(
-      JSON.stringify({
-        ...payload,
-        requestId,
-        cookies, // optional: only needed if you authenticate via cookies manually
-      }),
-    );
+    pending.set(requestId, { resolve, reject, timeout });
+
+    if (!socket || socket.readyState !== WebSocket.OPEN) {
+      throw new Error("WebSocket closed while sending");
+    }
+
+    socket.send(JSON.stringify({ ...payload, requestId }));
   });
 }
 
+async function send<T>(
+  useLocalhost: boolean,
+  payload: Record<string, any>,
+): Promise<T> {
+  await ensureSocket(useLocalhost);
+  return sendUnsafe(payload);
+}
+
+/* ──────────────────────────
+   Public API
+────────────────────────── */
+
+export async function closeClient() {
+  if (!socket) return;
+  socket.close();
+  socket = null;
+  connecting = null;
+}
+
 export async function addRun(
   useLocalhost: boolean,
   problem: string,
@@ -111,3 +221,13 @@ export async function addRunMoves(
     moves,
   });
 }
+
+export async function getText(
+  useLocalhost: boolean,
+  runId: string,
+): Promise<string> {
+  return await send(useLocalhost, {
+    type: "getText",
+    runId,
+  });
+}

src/extension.ts

@@ -4,6 +4,7 @@ import { startServer, stopServer } from "./server";
 import os from "os";
 import { exec } from "child_process";
 import net from "net";
+import { closeClient } from "./client";
 
 var globalStatusBarItem: vscode.StatusBarItem | undefined;
 
@@ -105,6 +106,7 @@ export function activate(context: vscode.ExtensionContext) {
 
 export function deactivate() {
   stopServer(globalStatusBarItem as any);
+  closeClient();
 }
 
 export function canConnect(

src/move.ts

@@ -1,7 +1,7 @@
 import * as vscode from "vscode";
 import { diff_match_patch, Diff } from "diff-match-patch";
 import { LiveRunMove } from "./types";
-import { addRunMoves } from "./client";
+import { addRunMoves, getText } from "./client";
 import path from "path";
 
 export function startMonitoring(useLocalhost: boolean, runId: string) {
@@ -11,6 +11,7 @@ export function startMonitoring(useLocalhost: boolean, runId: string) {
   let moves: LiveRunMove[] = [];
 
   let lastText = "";
+  let lastCursorOffset: number | undefined;
   let lastEventTime = Date.now();
   let idleTimeout: NodeJS.Timeout | undefined;
   let isSending = false;
@@ -18,7 +19,9 @@ export function startMonitoring(useLocalhost: boolean, runId: string) {
   const IDLE_MS = 350;
 
   if (vscode.window.activeTextEditor) {
-    lastText = vscode.window.activeTextEditor.document.getText();
+    const editor = vscode.window.activeTextEditor;
+    lastText = editor.document.getText();
+    lastCursorOffset = editor.document.offsetAt(editor.selection.active);
   }
 
   async function flush(file: string | null, language: string | null) {
@@ -42,13 +45,13 @@ export function startMonitoring(useLocalhost: boolean, runId: string) {
     return latency;
   }
 
+  // --- Existing text change interval ---
   const readInterval = setInterval(() => {
     const editor = vscode.window.activeTextEditor;
     if (!editor) return;
 
     const newText = editor.document.getText();
-    if (newText === lastText) return;
-
+    if (newText !== lastText) {
       let latency = getLatency();
 
       const diffs = dmp.diff_main(lastText, newText);
@@ -76,8 +79,68 @@ export function startMonitoring(useLocalhost: boolean, runId: string) {
           ),
         IDLE_MS,
       );
+    }
   }, 500);
 
+  // --- New cursor tracking interval ---
+  const cursorInterval = setInterval(() => {
+    const editor = vscode.window.activeTextEditor;
+    if (!editor) return;
+
+    const cursorOffset = editor.document.offsetAt(editor.selection.active);
+
+    if (cursorOffset !== lastCursorOffset) {
+      const latency = getLatency();
+
+      moves.push({
+        moveId: moveId++,
+        latency,
+        cursor: cursorOffset,
+      });
+
+      lastCursorOffset = cursorOffset;
+
+      if (idleTimeout) clearTimeout(idleTimeout);
+      idleTimeout = setTimeout(
+        () =>
+          flush(
+            path.basename(editor.document.fileName || "") || null,
+            editor.document.languageId || null,
+          ),
+        IDLE_MS,
+      );
+    }
+  }, 1000); // every second
+
+  // Sync text every 5 seconds to prevent desync
+  const textStabilityInterval = setInterval(async () => {
+    const editor = vscode.window.activeTextEditor;
+    if (!editor) return;
+
+    const serverText = await getText(useLocalhost, runId);
+
+    const editorText = editor.document.getText();
+    if (editorText !== serverText) {
+      clearTimeout(idleTimeout);
+      moves = [
+        {
+          latency: 0,
+          cursor: editor.document.offsetAt(editor.selection.active),
+          moveId: moveId++,
+          changes: {
+            from: 0,
+            to: serverText?.length || 0,
+            insert: editorText,
+          },
+        },
+      ];
+      flush(
+        path.basename(editor?.document.fileName || "") || null,
+        editor?.document.languageId || null,
+      );
+    }
+  }, 5000); // every 5 seconds
+
   vscode.window.onDidChangeActiveTextEditor((editor) => {
     let latency = getLatency();
 
@@ -98,7 +161,9 @@ export function startMonitoring(useLocalhost: boolean, runId: string) {
   });
 
   return () => {
+    clearInterval(textStabilityInterval);
     clearInterval(readInterval);
+    clearInterval(cursorInterval);
     if (idleTimeout) clearTimeout(idleTimeout);
     flush(
       path.basename(vscode.window.activeTextEditor?.document.fileName || "") ||