@memo-code/memo 0.8.0 → 0.8.5

package/dist/web/server/stream/stream.service.js

@@ -0,0 +1,166 @@
+"use strict";
+var __decorate = (this && this.__decorate) || function (decorators, target, key, desc) {
+    var c = arguments.length, r = c < 3 ? target : desc === null ? desc = Object.getOwnPropertyDescriptor(target, key) : desc, d;
+    if (typeof Reflect === "object" && typeof Reflect.decorate === "function") r = Reflect.decorate(decorators, target, key, desc);
+    else for (var i = decorators.length - 1; i >= 0; i--) if (d = decorators[i]) r = (c < 3 ? d(r) : c > 3 ? d(target, key, r) : d(target, key)) || r;
+    return c > 3 && r && Object.defineProperty(target, key, r), r;
+};
+var StreamService_1;
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.StreamService = void 0;
+const common_1 = require("@nestjs/common");
+const ws_1 = require("ws");
+function matchSessionStreamPath(pathname) {
+    const match = pathname.match(/^\/api\/chat\/sessions\/([^/]+)\/stream$/);
+    if (!match)
+        return null;
+    const sessionId = match[1]?.trim();
+    return sessionId && sessionId.length > 0
+        ? decodeURIComponent(sessionId)
+        : null;
+}
+function readBearerTokenFromHeader(authorization) {
+    if (!authorization)
+        return null;
+    const raw = Array.isArray(authorization) ? authorization[0] : authorization;
+    if (!raw)
+        return null;
+    const [scheme, token] = raw.split(' ');
+    if (scheme !== 'Bearer' || !token)
+        return null;
+    const trimmed = token.trim();
+    return trimmed.length > 0 ? trimmed : null;
+}
+let StreamService = StreamService_1 = class StreamService {
+    logger = new common_1.Logger(StreamService_1.name);
+    socketsBySession = new Map();
+    listenersBySession = new Map();
+    globalListeners = new Set();
+    wsServer = null;
+    attached = false;
+    attach(options) {
+        if (this.attached)
+            return;
+        this.attached = true;
+        this.wsServer = new ws_1.WebSocketServer({ noServer: true });
+        options.httpServer.on('upgrade', (request, socket, head) => {
+            const requestUrl = new URL(request.url ?? '/', 'http://127.0.0.1');
+            const sessionId = matchSessionStreamPath(requestUrl.pathname);
+            if (!sessionId || !this.wsServer)
+                return;
+            const queryToken = requestUrl.searchParams.get('accessToken')?.trim() ?? null;
+            const headerToken = readBearerTokenFromHeader(request.headers.authorization);
+            const accessToken = headerToken ?? queryToken;
+            if (!accessToken) {
+                socket.write('HTTP/1.1 401 Unauthorized\\r\\n\\r\\n');
+                socket.destroy();
+                return;
+            }
+            void options
+                .verifyAccessToken(accessToken)
+                .then(() => {
+                this.wsServer?.handleUpgrade(request, socket, head, (ws) => {
+                    this.registerSocket(sessionId, ws);
+                });
+            })
+                .catch(() => {
+                socket.write('HTTP/1.1 401 Unauthorized\\r\\n\\r\\n');
+                socket.destroy();
+            });
+        });
+    }
+    broadcast(sessionId, payload) {
+        if (this.globalListeners.size > 0) {
+            for (const listener of this.globalListeners) {
+                listener(sessionId, payload);
+            }
+        }
+        const listeners = this.listenersBySession.get(sessionId);
+        if (listeners && listeners.size > 0) {
+            for (const listener of listeners) {
+                listener(payload);
+            }
+        }
+        const sockets = this.socketsBySession.get(sessionId);
+        if (!sockets || sockets.size === 0)
+            return;
+        const message = JSON.stringify(payload);
+        for (const socket of sockets) {
+            if (socket.readyState === socket.OPEN) {
+                socket.send(message);
+            }
+        }
+    }
+    disconnectSession(sessionId) {
+        const sockets = this.socketsBySession.get(sessionId);
+        if (!sockets)
+            return;
+        for (const socket of sockets) {
+            socket.close(1000, 'session closed');
+        }
+        this.socketsBySession.delete(sessionId);
+    }
+    subscribe(sessionId, listener) {
+        let listeners = this.listenersBySession.get(sessionId);
+        if (!listeners) {
+            listeners = new Set();
+            this.listenersBySession.set(sessionId, listeners);
+        }
+        listeners.add(listener);
+        return () => {
+            const target = this.listenersBySession.get(sessionId);
+            if (!target)
+                return;
+            target.delete(listener);
+            if (target.size === 0) {
+                this.listenersBySession.delete(sessionId);
+            }
+        };
+    }
+    subscribeAll(listener) {
+        this.globalListeners.add(listener);
+        return () => {
+            this.globalListeners.delete(listener);
+        };
+    }
+    async onModuleDestroy() {
+        if (this.wsServer) {
+            for (const [, sockets] of this.socketsBySession) {
+                for (const socket of sockets) {
+                    socket.close(1001, 'server shutdown');
+                }
+            }
+            this.socketsBySession.clear();
+            this.listenersBySession.clear();
+            this.globalListeners.clear();
+            await new Promise((resolve) => {
+                this.wsServer?.close(() => resolve());
+            });
+            this.wsServer = null;
+        }
+    }
+    registerSocket(sessionId, socket) {
+        let sockets = this.socketsBySession.get(sessionId);
+        if (!sockets) {
+            sockets = new Set();
+            this.socketsBySession.set(sessionId, sockets);
+        }
+        sockets.add(socket);
+        socket.on('close', () => {
+            const target = this.socketsBySession.get(sessionId);
+            if (!target)
+                return;
+            target.delete(socket);
+            if (target.size === 0) {
+                this.socketsBySession.delete(sessionId);
+            }
+        });
+        socket.on('error', (error) => {
+            this.logger.warn(`socket error for session=${sessionId}: ${error.message}`);
+        });
+    }
+};
+exports.StreamService = StreamService;
+exports.StreamService = StreamService = StreamService_1 = __decorate([
+    (0, common_1.Injectable)()
+], StreamService);