UNPKG

@mastra/core

Version:

Mastra is a framework for building AI-powered applications and agents with a modern TypeScript stack.

400 lines (395 loc) • 12.2 kB
'use strict'; var chunkVEYVZLLD_cjs = require('../chunk-VEYVZLLD.cjs'); var chunkCVF4W47C_cjs = require('../chunk-CVF4W47C.cjs'); var crypto = require('crypto'); var promises = require('fs/promises'); var net = require('net'); var path = require('path'); function _interopDefault (e) { return e && e.__esModule ? e : { default: e }; } var net__default = /*#__PURE__*/_interopDefault(net); function writeFrame(socket, frame) { return new Promise((resolve, reject) => { const onError = (error) => { socket.off("drain", onDrain); reject(error); }; const onDrain = () => { socket.off("error", onError); resolve(); }; socket.once("error", onError); const drained = socket.write(`${JSON.stringify(frame)} `, () => { if (drained) { socket.off("error", onError); resolve(); } }); if (!drained) { socket.once("drain", onDrain); } }); } function nextTick() { return new Promise((resolve) => setImmediate(resolve)); } function readFrames(socket, onFrame) { let buffer = ""; socket.setEncoding("utf8"); socket.on("data", (chunk) => { buffer += chunk; while (true) { const newlineIndex = buffer.indexOf("\n"); if (newlineIndex === -1) break; const line = buffer.slice(0, newlineIndex); buffer = buffer.slice(newlineIndex + 1); if (!line.trim()) continue; try { onFrame(JSON.parse(line)); } catch { } } }); } var UnixSocketPubSub = class extends chunkCVF4W47C_cjs.PubSub { socketPath; #server; #clientSocket; #isBroker = false; #closed = false; #starting; #callbacks = /* @__PURE__ */ new Map(); #subscribeWaiters = /* @__PURE__ */ new Map(); #brokerClients = /* @__PURE__ */ new Map(); #pendingWrites = []; constructor(socketPath) { super(); this.socketPath = socketPath; } get supportedModes() { return ["push"]; } get isBroker() { return this.#isBroker; } async publish(topic, event) { await this.#ensureStarted(); if (this.#isBroker) { await this.#publishFromBroker(topic, event); return; } const socket = this.#clientSocket; if (!socket || socket.destroyed) { await this.#ensureStarted(true); } await this.#sendToBroker({ type: "publish", topic, event }); } async subscribe(topic, cb, options) { if (options?.group) { throw new Error("UnixSocketPubSub does not support grouped subscriptions yet"); } const callbacks = this.#callbacks.get(topic) ?? /* @__PURE__ */ new Set(); const hadCallback = callbacks.has(cb); const wasConnected = Boolean(this.#clientSocket && !this.#clientSocket.destroyed); callbacks.add(cb); this.#callbacks.set(topic, callbacks); try { await this.#ensureStarted(); if (!this.#isBroker && !hadCallback && wasConnected) { await this.#sendSubscribeToBroker(topic); } } catch (error) { if (!hadCallback) { callbacks.delete(cb); if (callbacks.size === 0) { this.#callbacks.delete(topic); } } throw error; } } async unsubscribe(topic, cb) { const callbacks = this.#callbacks.get(topic); callbacks?.delete(cb); if (callbacks?.size === 0) { this.#callbacks.delete(topic); if (!this.#isBroker && this.#clientSocket && !this.#clientSocket.destroyed) { await this.#sendToBroker({ type: "unsubscribe", topic }); await nextTick(); } } } async flush() { await Promise.allSettled(this.#pendingWrites); this.#pendingWrites = []; } async close() { this.#closed = true; this.#callbacks.clear(); this.#clientSocket?.destroy(); this.#clientSocket = void 0; this.#rejectSubscribeWaiters(new Error("UnixSocketPubSub is closed")); for (const client of this.#brokerClients.values()) { client.socket.destroy(); } this.#brokerClients.clear(); if (this.#server) { await new Promise((resolve) => this.#server?.close(() => resolve())); this.#server = void 0; } if (this.#isBroker) { await promises.unlink(this.socketPath).catch(() => { }); } this.#isBroker = false; } async #ensureStarted(forceReconnect = false) { if (this.#closed) { throw new Error("UnixSocketPubSub is closed"); } if (!forceReconnect && (this.#isBroker || this.#clientSocket && !this.#clientSocket.destroyed)) { return; } if (this.#starting) { return this.#starting; } this.#starting = this.#start(forceReconnect).finally(() => { this.#starting = void 0; }); return this.#starting; } async #start(forceReconnect) { if (forceReconnect) { this.#clientSocket?.destroy(); this.#clientSocket = void 0; this.#isBroker = false; } this.#throwIfClosed(); await promises.mkdir(path.dirname(this.socketPath), { recursive: true }); this.#throwIfClosed(); try { await this.#listen(); this.#throwIfClosed(); this.#isBroker = true; return; } catch (error) { if (this.#closed) { await this.close(); throw new Error("UnixSocketPubSub is closed"); } const code = error.code; if (code !== "EADDRINUSE") throw error; } try { await this.#connectClient(); this.#throwIfClosed(); } catch (error) { if (this.#closed) { await this.close(); throw new Error("UnixSocketPubSub is closed"); } const code = error.code; if (code === "ECONNREFUSED" || code === "ENOENT" || code === "ENOTSOCK") { await promises.unlink(this.socketPath).catch(() => { }); this.#throwIfClosed(); await this.#listen(); this.#throwIfClosed(); this.#isBroker = true; return; } throw error; } } #throwIfClosed() { if (this.#closed) { throw new Error("UnixSocketPubSub is closed"); } } #listen() { return new Promise((resolve, reject) => { const server = net__default.default.createServer((socket) => this.#handleBrokerClient(socket)); const onError = (error) => { server.off("listening", onListening); reject(error); }; const onListening = () => { server.off("error", onError); this.#server = server; resolve(); }; server.once("error", onError); server.once("listening", onListening); server.listen(this.socketPath); }); } #connectClient() { return new Promise((resolve, reject) => { const socket = net__default.default.createConnection(this.socketPath); const onError = (error) => { socket.off("connect", onConnect); reject(error); }; const onConnect = () => { socket.off("error", onError); this.#clientSocket = socket; this.#isBroker = false; readFrames(socket, (frame) => this.#handleServerFrame(frame)); socket.on("close", () => { if (this.#clientSocket === socket) this.#clientSocket = void 0; this.#rejectSubscribeWaiters(new Error("UnixSocketPubSub broker connection closed")); }); socket.on("error", (error) => { if (this.#clientSocket === socket) this.#clientSocket = void 0; this.#rejectSubscribeWaiters(error); }); void this.#resubscribeClient().then(resolve, reject); }; socket.once("error", onError); socket.once("connect", onConnect); }); } async #resubscribeClient() { for (const topic of this.#callbacks.keys()) { await this.#sendSubscribeToBroker(topic); } } async #sendSubscribeToBroker(topic) { let waiter; const subscribed = new Promise((resolve, reject) => { waiter = { resolve, reject }; const waiters = this.#subscribeWaiters.get(topic) ?? []; waiters.push(waiter); this.#subscribeWaiters.set(topic, waiters); }); try { await this.#sendToBroker({ type: "subscribe", topic }); } catch (error) { this.#removeSubscribeWaiter(topic, waiter); throw error; } await subscribed; } #removeSubscribeWaiter(topic, waiter) { if (!waiter) return; const waiters = this.#subscribeWaiters.get(topic); if (!waiters) return; const nextWaiters = waiters.filter((item) => item !== waiter); if (nextWaiters.length === 0) { this.#subscribeWaiters.delete(topic); return; } this.#subscribeWaiters.set(topic, nextWaiters); } #settleSubscribeWaiters(topic, error) { const waiters = this.#subscribeWaiters.get(topic); this.#subscribeWaiters.delete(topic); if (error) { waiters?.forEach((waiter) => waiter.reject(error)); return; } waiters?.forEach((waiter) => waiter.resolve()); } #rejectSubscribeWaiters(error) { for (const topic of this.#subscribeWaiters.keys()) { this.#settleSubscribeWaiters(topic, error); } } #handleBrokerClient(socket) { const client = { socket, subscriptions: /* @__PURE__ */ new Set() }; this.#brokerClients.set(socket, client); readFrames(socket, (frame) => { const clientFrame = frame; if (clientFrame.type === "subscribe") { client.subscriptions.add(clientFrame.topic); void writeFrame(socket, { type: "subscribed", topic: clientFrame.topic }).catch(() => { }); } else if (clientFrame.type === "unsubscribe") { client.subscriptions.delete(clientFrame.topic); } else if (clientFrame.type === "publish") { void this.#publishFromBroker(clientFrame.topic, clientFrame.event); } }); socket.on("close", () => this.#brokerClients.delete(socket)); socket.on("error", () => this.#brokerClients.delete(socket)); } #handleServerFrame(frame) { if (frame.type === "subscribed") { this.#settleSubscribeWaiters(frame.topic); return; } if (frame.type !== "event") return; const event = { ...frame.event, createdAt: new Date(frame.event.createdAt) }; this.#deliverLocal(frame.topic, event); } async #publishFromBroker(topic, event) { const brokerEvent = { ...event, id: crypto.randomUUID(), createdAt: /* @__PURE__ */ new Date(), deliveryAttempt: 1 }; this.#deliverLocal(topic, brokerEvent); const frame = { type: "event", topic, event: brokerEvent }; for (const client of this.#brokerClients.values()) { if (!client.subscriptions.has(topic) || client.socket.destroyed) continue; const write = writeFrame(client.socket, frame).catch(() => { }); this.#pendingWrites.push(write); } await this.flush(); } #deliverLocal(topic, event) { const callbacks = this.#callbacks.get(topic); if (!callbacks) return; for (const cb of callbacks) { try { const result = cb( event, async () => { }, async () => { } ); if (result && typeof result.catch === "function") { void result.catch(() => { }); } } catch { } } } async #sendToBroker(frame) { const socket = this.#clientSocket; if (!socket || socket.destroyed) { await this.#ensureStarted(true); } const activeSocket = this.#clientSocket; if (!activeSocket || activeSocket.destroyed) { throw new Error("UnixSocketPubSub is not connected to a broker"); } await writeFrame(activeSocket, frame); } }; Object.defineProperty(exports, "CachingPubSub", { enumerable: true, get: function () { return chunkVEYVZLLD_cjs.CachingPubSub; } }); Object.defineProperty(exports, "withCaching", { enumerable: true, get: function () { return chunkVEYVZLLD_cjs.withCaching; } }); Object.defineProperty(exports, "EventEmitterPubSub", { enumerable: true, get: function () { return chunkCVF4W47C_cjs.EventEmitterPubSub; } }); Object.defineProperty(exports, "PubSub", { enumerable: true, get: function () { return chunkCVF4W47C_cjs.PubSub; } }); exports.UnixSocketPubSub = UnixSocketPubSub; //# sourceMappingURL=index.cjs.map //# sourceMappingURL=index.cjs.map