fix(ui): socket yjs manager & refactor

ui/src/lib/i18n/locales/en.json

@@ -225,6 +225,7 @@
   "Member has been successfully added to the workspace.": "",
   "Member Removed": "",
   "Member has been successfully removed from the workspace.": "",
+  "Connection Error": "",
   "Empty workflow detected": "",
   "You cannot create a deployment without a workflow.": "",
   "Reload": ""

ui/src/lib/i18n/locales/es.json

@@ -225,6 +225,7 @@
   "Member has been successfully added to the workspace.": "El miembro ha sido agregado exitosamente al espacio de trabajo.",
   "Member Removed": "Miembro eliminado",
   "Member has been successfully removed from the workspace.": "El miembro ha sido eliminado exitosamente del espacio de trabajo.",
+  "Connection Error": "",
   "Empty workflow detected": "Se detectó un flujo de trabajo vacío",
   "You cannot create a deployment without a workflow.": "No puedes crear un despliegue sin un flujo de trabajo.",
   "Reload": "Recargar"

ui/src/lib/i18n/locales/fr.json

@@ -225,6 +225,7 @@
   "Member has been successfully added to the workspace.": "Le membre a été ajouté avec succès à l'espace de travail.",
   "Member Removed": "Membre supprimé",
   "Member has been successfully removed from the workspace.": "Le membre a été supprimé avec succès de l'espace de travail.",
+  "Connection Error": "",
   "Empty workflow detected": "Flux de travail vide détecté",
   "You cannot create a deployment without a workflow.": "Vous ne pouvez pas créer de déploiement sans flux de travail.",
   "Reload": "Recharger"

ui/src/lib/i18n/locales/ja.json

@@ -225,6 +225,7 @@
   "Member has been successfully added to the workspace.": "メンバーがワークスペースに正常に追加されました。",
   "Member Removed": "メンバー削除",
   "Member has been successfully removed from the workspace.": "メンバーがワークスペースから正常に削除されました。",
+  "Connection Error": "",
   "Empty workflow detected": "空のワークフローが検出されました",
   "You cannot create a deployment without a workflow.": "ワークフローがないとデプロイメントを作成できません。",
   "Reload": "リロード"

ui/src/lib/i18n/locales/zh.json

@@ -225,6 +225,7 @@
   "Member has been successfully added to the workspace.": "成员已成功添加到工作区。",
   "Member Removed": "成员已移除",
   "Member has been successfully removed from the workspace.": "成员已成功从工作区中移除。",
+  "Connection Error": "",
   "Empty workflow detected": "检测到空工作流",
   "You cannot create a deployment without a workflow.": "没有工作流，您无法创建部署。",
   "Reload": "重新加载"

ui/src/lib/utils.ts

@@ -5,3 +5,7 @@ import { twMerge } from "tailwind-merge";
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
 }
+
+export function sleep(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}

ui/src/lib/yjs/socketYjsManager.ts

@@ -0,0 +1,303 @@
+import * as Y from "yjs";
+
+import { sleep } from "../utils";
+
+import type { FlowMessage } from "./types";
+
+export type AccessTokenProvider = () => Promise<string> | string;
+
+export class SocketYjsManager {
+  protected ws!: WebSocket;
+  protected doc: Y.Doc;
+  protected socketReady = false;
+  protected firstSyncComplete = false;
+  protected accessTokenProvider: AccessTokenProvider | undefined;
+  protected projectId: string | undefined;
+  protected onUpdateHandlers: ((update: Uint8Array) => void)[] = [];
+  protected reconnectAttempts = 0;
+  protected maxReconnectAttempts = 5;
+  protected reconnectDelay = 1000;
+  protected reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+
+  constructor(doc: Y.Doc) {
+    this.doc = doc;
+
+    // Bind methods
+    this.onConnectionEstablished = this.onConnectionEstablished.bind(this);
+    this.onConnectionDisconnect = this.onConnectionDisconnect.bind(this);
+    this.onConnectionError = this.onConnectionError.bind(this);
+    this.onAuthenticateRequest = this.onAuthenticateRequest.bind(this);
+    this.onDocUpdate = this.onDocUpdate.bind(this);
+    this.onReady = this.onReady.bind(this);
+    this.onPeerUpdate = this.onPeerUpdate.bind(this);
+    this.handleMessage = this.handleMessage.bind(this);
+    this.reconnect = this.reconnect.bind(this);
+  }
+
+  public getDoc(): Y.Doc {
+    return this.doc;
+  }
+
+  // Replace the setupSocket method
+  async setupSocket(data: {
+    url: string;
+    roomId: string;
+    projectId: string;
+    accessTokenProvider: AccessTokenProvider;
+  }) {
+    this.accessTokenProvider = data.accessTokenProvider;
+    this.projectId = data.projectId;
+
+    try {
+      const token = await this.accessTokenProvider();
+      const wsUrl = new URL(data.url);
+      wsUrl.protocol = wsUrl.protocol.replace("http", "ws");
+      wsUrl.pathname = `/${data.roomId}`;
+
+      // Add query parameters for authentication
+      wsUrl.searchParams.set("user_id", this.doc.clientID.toString());
+      wsUrl.searchParams.set("project_id", data.projectId);
+      wsUrl.searchParams.set("token", token); // Pass token as query param since we can't set headers
+
+      this.ws = new WebSocket(wsUrl.href);
+      this.ws.binaryType = "arraybuffer";
+
+      this.setupWebSocketListeners();
+      this.setupDocListeners();
+
+      console.log("Attempting WebSocket connection to:", wsUrl.origin + wsUrl.pathname);
+    } catch (error) {
+      console.error("Failed to setup WebSocket:", error);
+      throw error;
+    }
+  }
+
+  // Update the reconnect method
+  private reconnect() {
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+    }
+
+    if (this.reconnectAttempts < this.maxReconnectAttempts) {
+      this.reconnectAttempts++;
+      const delay =
+        this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1);
+
+      this.reconnectTimer = setTimeout(async () => {
+        try {
+          if (this.ws) {
+            const originalUrl = new URL(this.ws.url);
+            const baseUrl = `${originalUrl.protocol}//${originalUrl.host}`;
+            const roomId = originalUrl.pathname.slice(1);
+            await this.setupSocket({
+              url: baseUrl,
+              roomId,
+              projectId: this.projectId || "",
+              accessTokenProvider: this.accessTokenProvider || (() => ""),
+            });
+          }
+        } catch (error) {
+          console.error("Reconnection failed:", error);
+        }
+      }, delay);
+    }
+  }
+
+  private setupWebSocketListeners() {
+    this.ws.addEventListener("open", this.onConnectionEstablished);
+    this.ws.addEventListener("close", this.onConnectionDisconnect);
+    this.ws.addEventListener("error", this.onConnectionError);
+    this.ws.addEventListener("message", this.handleMessage);
+  }
+
+  private setupDocListeners() {
+    this.doc.on("update", this.onDocUpdate);
+  }
+
+  protected onConnectionEstablished() {
+    this.reconnectAttempts = 0;
+    this.socketReady = true;
+    this.initializeRoom().catch(console.error);
+  }
+
+  protected onConnectionDisconnect() {
+    this.socketReady = false;
+    this.firstSyncComplete = false;
+    this.reconnect();
+  }
+
+  protected onConnectionError(error: Event) {
+    console.error("WebSocket error:", error);
+    this.reconnect();
+  }
+
+  protected async onAuthenticateRequest() {
+    const token = await this.accessTokenProvider?.();
+    if (token && this.ws.readyState === WebSocket.OPEN) {
+      this.ws.send(JSON.stringify({ type: "authenticate", token }));
+    }
+  }
+
+  protected async handleMessage(event: MessageEvent) {
+    try {
+      if (event.data instanceof ArrayBuffer) {
+        // Handle binary message (Yjs update)
+        const update = new Uint8Array(event.data);
+        await this.onPeerUpdate({ update });
+      } else if (typeof event.data === "string") {
+        // Handle text message
+        const data = JSON.parse(event.data);
+        if (data.type === "authenticate") {
+          await this.onAuthenticateRequest();
+        } else if (data.type === "ready") {
+          this.onReady();
+        }
+      }
+    } catch (error) {
+      console.error("Error handling message:", error);
+    }
+  }
+
+  protected async initializeRoom() {
+    try {
+      await this.sendFlowMessage({
+        event: {
+          tag: "Create",
+          content: { room_id: this.doc.clientID.toString() },
+        },
+      });
+
+      await this.sendFlowMessage({
+        event: {
+          tag: "Join",
+          content: { room_id: this.doc.clientID.toString() },
+        },
+      });
+
+      await this.sendFlowMessage({
+        event: {
+          tag: "Emit",
+          content: { data: "" },
+        },
+        session_command: {
+          tag: "Start",
+          content: {
+            project_id: this.projectId || "",
+            user: {
+              id: this.doc.clientID.toString(),
+              tenant_id: this.projectId,
+              name: "defaultName",
+              email: "[email protected]",
+            },
+          },
+        },
+      });
+
+      await this.syncData();
+    } catch (error) {
+      console.error("Failed to initialize room:", error);
+    }
+  }
+
+  async isReady(): Promise<boolean> {
+    if (this.socketReady) return true;
+    await sleep(100);
+    return await this.isReady();
+  }
+
+  protected onReady() {
+    this.socketReady = true;
+  }
+
+  protected async onPeerUpdate(data: { update: ArrayBuffer | Uint8Array }) {
+    const update =
+      data.update instanceof ArrayBuffer
+        ? new Uint8Array(data.update)
+        : data.update;
+    Y.applyUpdate(this.doc, update, "peer");
+    this.onUpdateHandlers.forEach((handler) => handler(update));
+  }
+
+  async syncData() {
+    await this.isReady();
+
+    const stateVector = Y.encodeStateVector(this.doc);
+    if (this.ws.readyState === WebSocket.OPEN) {
+      this.ws.send(stateVector);
+    }
+
+    if (!this.firstSyncComplete) {
+      this.firstSyncComplete = true;
+      queueMicrotask(() => this.syncData());
+    }
+  }
+
+  private async sendFlowMessage(message: FlowMessage): Promise<void> {
+    return new Promise((resolve, reject) => {
+      if (this.ws.readyState !== WebSocket.OPEN) {
+        reject(new Error("WebSocket is not connected"));
+        return;
+      }
+
+      try {
+        this.ws.send(JSON.stringify(message));
+        resolve();
+      } catch (error) {
+        reject(error);
+      }
+    });
+  }
+
+  protected onDocUpdate(update: Uint8Array, origin: unknown) {
+    if (origin === this.doc.clientID && this.ws.readyState === WebSocket.OPEN) {
+      this.sendFlowMessage({
+        event: {
+          tag: "Emit",
+          content: { data: "" },
+        },
+        session_command: {
+          tag: "MergeUpdates",
+          content: {
+            project_id: this.projectId || "",
+            data: new Uint8Array(update),
+            updated_by: this.doc.clientID.toString(),
+          },
+        },
+      });
+    }
+  }
+
+  onUpdate(handler: (update: Uint8Array) => void) {
+    this.onUpdateHandlers.push(handler);
+  }
+
+  destroy() {
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+    }
+
+    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
+      this.sendFlowMessage({
+        event: {
+          tag: "Emit",
+          content: { data: "" },
+        },
+        session_command: {
+          tag: "End",
+          content: {
+            project_id: this.projectId || "",
+            user: {
+              id: this.doc.clientID.toString(),
+              tenant_id: this.projectId,
+              name: "defaultName",
+              email: "[email protected]",
+            },
+          },
+        },
+      }).finally(() => {
+        this.ws.close();
+      });
+    }
+    this.doc.destroy();
+  }
+}

ui/src/lib/yjs/types.ts

@@ -1,3 +1,5 @@
+import { User } from "@flow/types/user";
+
 export type YJsonPrimitive = string | number | boolean | null | Uint8Array;
 
 export type YJsonValue =
@@ -6,3 +8,32 @@ export type YJsonValue =
   | {
       [key: string]: YJsonValue;
     };
+
+export type FlowMessage = {
+  event: {
+    tag: "Create" | "Join" | "Leave" | "Emit";
+    content: {
+      room_id?: string;
+      data?: string;
+    };
+  };
+  session_command?: SessionCommand;
+};
+
+export type SessionCommand = {
+  tag:
+    | "Start"
+    | "End"
+    | "Complete"
+    | "CheckStatus"
+    | "AddTask"
+    | "RemoveTask"
+    | "ListAllSnapshotsVersions"
+    | "MergeUpdates";
+  content: {
+    project_id: string;
+    user?: User;
+    data?: Uint8Array;
+    updated_by?: string;
+  };
+};