chore(monorepo/server): move server-side source code

apps/server/src/services/ws.ts

@@ -0,0 +1,316 @@
+import { WebSocketServer as WebSocketServer, WebSocket } from "ws";
+import { isDev, isElectron, randomString } from "./utils.js";
+import log from "./log.js";
+import sql from "./sql.js";
+import cls from "./cls.js";
+import config from "./config.js";
+import syncMutexService from "./sync_mutex.js";
+import protectedSessionService from "./protected_session.js";
+import becca from "../becca/becca.js";
+import AbstractBeccaEntity from "../becca/entities/abstract_becca_entity.js";
+
+import type { IncomingMessage, Server as HttpServer } from "http";
+import type { EntityChange } from "./entity_changes_interface.js";
+
+if (isDev) {
+    const chokidar = (await import("chokidar")).default;
+    const debounce = (await import("debounce")).default;
+    const debouncedReloadFrontend = debounce(() => reloadFrontend("source code change"), 200);
+    chokidar
+        .watch("src/public", {
+            ignored: "src/public/app/doc_notes/en/User Guide"
+        })
+        .on("add", debouncedReloadFrontend)
+        .on("change", debouncedReloadFrontend)
+        .on("unlink", debouncedReloadFrontend);
+}
+
+let webSocketServer!: WebSocketServer;
+let lastSyncedPush: number | null = null;
+
+interface Message {
+    type: string;
+    data?: {
+        lastSyncedPush?: number | null;
+        entityChanges?: any[];
+        shrinkImages?: boolean;
+    } | null;
+    lastSyncedPush?: number | null;
+
+    progressCount?: number;
+    taskId?: string;
+    taskType?: string | null;
+    message?: string;
+    reason?: string;
+    result?: string | Record<string, string | undefined>;
+
+    script?: string;
+    params?: any[];
+    noteId?: string;
+    messages?: string[];
+    startNoteId?: string;
+    currentNoteId?: string;
+    entityType?: string;
+    entityId?: string;
+    originEntityName?: "notes";
+    originEntityId?: string | null;
+    lastModifiedMs?: number;
+    filePath?: string;
+
+    // LLM streaming specific fields
+    chatNoteId?: string;
+    content?: string;
+    thinking?: string;
+    toolExecution?: {
+        action?: string;
+        tool?: string;
+        toolCallId?: string;
+        result?: string | Record<string, any>;
+        error?: string;
+        args?: Record<string, unknown>;
+    };
+    done?: boolean;
+    error?: string;
+    raw?: unknown;
+}
+
+type SessionParser = (req: IncomingMessage, params: {}, cb: () => void) => void;
+function init(httpServer: HttpServer, sessionParser: SessionParser) {
+    webSocketServer = new WebSocketServer({
+        verifyClient: (info, done) => {
+            sessionParser(info.req, {}, () => {
+                const allowed = isElectron || (info.req as any).session.loggedIn || (config.General && config.General.noAuthentication);
+
+                if (!allowed) {
+                    log.error("WebSocket connection not allowed because session is neither electron nor logged in.");
+                }
+
+                done(allowed);
+            });
+        },
+        server: httpServer
+    });
+
+    webSocketServer.on("connection", (ws, req) => {
+        (ws as any).id = randomString(10);
+
+        console.log(`websocket client connected`);
+
+        ws.on("message", async (messageJson) => {
+            const message = JSON.parse(messageJson as any);
+
+            if (message.type === "log-error") {
+                log.info(`JS Error: ${message.error}\r
+Stack: ${message.stack}`);
+            } else if (message.type === "log-info") {
+                log.info(`JS Info: ${message.info}`);
+            } else if (message.type === "ping") {
+                await syncMutexService.doExclusively(() => sendPing(ws));
+            } else {
+                log.error("Unrecognized message: ");
+                log.error(message);
+            }
+        });
+    });
+
+    webSocketServer.on("error", (error) => {
+        // https://github.com/zadam/trilium/issues/3374#issuecomment-1341053765
+        console.log(error);
+    });
+}
+
+function sendMessage(client: WebSocket, message: Message) {
+    const jsonStr = JSON.stringify(message);
+
+    if (client.readyState === WebSocket.OPEN) {
+        client.send(jsonStr);
+    }
+}
+
+function sendMessageToAllClients(message: Message) {
+    const jsonStr = JSON.stringify(message);
+
+    if (webSocketServer) {
+        // Special logging for LLM streaming messages
+        if (message.type === "llm-stream") {
+            log.info(`[WS-SERVER] Sending LLM stream message: chatNoteId=${message.chatNoteId}, content=${!!message.content}, thinking=${!!message.thinking}, toolExecution=${!!message.toolExecution}, done=${!!message.done}`);
+        } else if (message.type !== "sync-failed" && message.type !== "api-log-messages") {
+            log.info(`Sending message to all clients: ${jsonStr}`);
+        }
+
+        let clientCount = 0;
+        webSocketServer.clients.forEach(function each(client) {
+            if (client.readyState === WebSocket.OPEN) {
+                client.send(jsonStr);
+                clientCount++;
+            }
+        });
+
+        // Log WebSocket client count for debugging
+        if (message.type === "llm-stream") {
+            log.info(`[WS-SERVER] Sent LLM stream message to ${clientCount} clients`);
+        }
+    }
+}
+
+function fillInAdditionalProperties(entityChange: EntityChange) {
+    if (entityChange.isErased) {
+        return;
+    }
+
+    // fill in some extra data needed by the frontend
+    // first try to use becca, which works for non-deleted entities
+    // only when that fails, try to load from the database
+    if (entityChange.entityName === "attributes") {
+        entityChange.entity = becca.getAttribute(entityChange.entityId);
+
+        if (!entityChange.entity) {
+            entityChange.entity = sql.getRow(/*sql*/`SELECT * FROM attributes WHERE attributeId = ?`, [entityChange.entityId]);
+        }
+    } else if (entityChange.entityName === "branches") {
+        entityChange.entity = becca.getBranch(entityChange.entityId);
+
+        if (!entityChange.entity) {
+            entityChange.entity = sql.getRow(/*sql*/`SELECT * FROM branches WHERE branchId = ?`, [entityChange.entityId]);
+        }
+    } else if (entityChange.entityName === "notes") {
+        entityChange.entity = becca.getNote(entityChange.entityId);
+
+        if (!entityChange.entity) {
+            entityChange.entity = sql.getRow(/*sql*/`SELECT * FROM notes WHERE noteId = ?`, [entityChange.entityId]);
+
+            if (entityChange.entity?.isProtected) {
+                entityChange.entity.title = protectedSessionService.decryptString(entityChange.entity.title || "");
+            }
+        }
+    } else if (entityChange.entityName === "revisions") {
+        entityChange.noteId = sql.getValue<string>(
+            /*sql*/`SELECT noteId
+                                                    FROM revisions
+                                                    WHERE revisionId = ?`,
+            [entityChange.entityId]
+        );
+    } else if (entityChange.entityName === "note_reordering") {
+        entityChange.positions = {};
+
+        const parentNote = becca.getNote(entityChange.entityId);
+
+        if (parentNote) {
+            for (const childBranch of parentNote.getChildBranches()) {
+                if (childBranch?.branchId) {
+                    entityChange.positions[childBranch.branchId] = childBranch.notePosition;
+                }
+            }
+        }
+    } else if (entityChange.entityName === "options") {
+        entityChange.entity = becca.getOption(entityChange.entityId);
+
+        if (!entityChange.entity) {
+            entityChange.entity = sql.getRow(/*sql*/`SELECT * FROM options WHERE name = ?`, [entityChange.entityId]);
+        }
+    } else if (entityChange.entityName === "attachments") {
+        entityChange.entity = sql.getRow(
+            /*sql*/`SELECT attachments.*, LENGTH(blobs.content) AS contentLength
+                                                FROM attachments
+                                                JOIN blobs USING (blobId)
+                                                WHERE attachmentId = ?`,
+            [entityChange.entityId]
+        );
+    }
+
+    if (entityChange.entity instanceof AbstractBeccaEntity) {
+        entityChange.entity = entityChange.entity.getPojo();
+    }
+}
+
+// entities with higher number can reference the entities with lower number
+const ORDERING: Record<string, number> = {
+    etapi_tokens: 0,
+    attributes: 2,
+    branches: 2,
+    blobs: 0,
+    note_reordering: 2,
+    revisions: 2,
+    attachments: 3,
+    notes: 1,
+    options: 0,
+    note_embeddings: 3
+};
+
+function sendPing(client: WebSocket, entityChangeIds = []) {
+    if (entityChangeIds.length === 0) {
+        sendMessage(client, { type: "ping" });
+
+        return;
+    }
+
+    const entityChanges = sql.getManyRows<EntityChange>(/*sql*/`SELECT * FROM entity_changes WHERE id IN (???)`, entityChangeIds);
+    if (!entityChanges) {
+        return;
+    }
+
+    // sort entity changes since froca expects "referential order", i.e. referenced entities should already exist
+    // in froca.
+    // Froca needs this since it is an incomplete copy, it can't create "skeletons" like becca.
+    entityChanges.sort((a, b) => ORDERING[a.entityName] - ORDERING[b.entityName]);
+
+    for (const entityChange of entityChanges) {
+        try {
+            fillInAdditionalProperties(entityChange);
+        } catch (e: any) {
+            log.error(`Could not fill additional properties for entity change ${JSON.stringify(entityChange)} because of error: ${e.message}: ${e.stack}`);
+        }
+    }
+
+    sendMessage(client, {
+        type: "frontend-update",
+        data: {
+            lastSyncedPush,
+            entityChanges
+        }
+    });
+}
+
+function sendTransactionEntityChangesToAllClients() {
+    if (webSocketServer) {
+        const entityChangeIds = cls.getAndClearEntityChangeIds();
+
+        webSocketServer.clients.forEach((client) => sendPing(client, entityChangeIds));
+    }
+}
+
+function syncPullInProgress() {
+    sendMessageToAllClients({ type: "sync-pull-in-progress", lastSyncedPush });
+}
+
+function syncPushInProgress() {
+    sendMessageToAllClients({ type: "sync-push-in-progress", lastSyncedPush });
+}
+
+function syncFinished() {
+    sendMessageToAllClients({ type: "sync-finished", lastSyncedPush });
+}
+
+function syncFailed() {
+    sendMessageToAllClients({ type: "sync-failed", lastSyncedPush });
+}
+
+function reloadFrontend(reason: string) {
+    sendMessageToAllClients({ type: "reload-frontend", reason });
+}
+
+function setLastSyncedPush(entityChangeId: number) {
+    lastSyncedPush = entityChangeId;
+}
+
+export default {
+    init,
+    sendMessageToAllClients,
+    syncPushInProgress,
+    syncPullInProgress,
+    syncFinished,
+    syncFailed,
+    sendTransactionEntityChangesToAllClients,
+    setLastSyncedPush,
+    reloadFrontend
+};