UNPKG

@mastra/core

Version:
847 lines (840 loc) • 26.6 kB
export { CachingPubSub, withCaching } from '../chunk-BOHPE4XX.js'; import { DefaultStepResult } from '../chunk-S2RGNJL3.js'; import { PubSub } from '../chunk-2S76KRB5.js'; export { EventEmitterPubSub, NoopLeaseProvider, PubSub, isLeaseProvider } from '../chunk-2S76KRB5.js'; import { DefaultGeneratedFile, DefaultGeneratedFileWithType } from '../chunk-HHYICF3X.js'; import { randomUUID } from 'crypto'; import { unlink, mkdir, open, stat } from 'fs/promises'; import net from 'net'; import { dirname } from 'path'; // src/events/codec/error.ts var MAX_CAUSE_DEPTH = 5; function serializeError(err, depth = 0) { const json = { name: err.name || "Error", message: err.message }; if (err.stack !== void 0) json.stack = err.stack; if (err.cause !== void 0) { if (err.cause instanceof Error && depth < MAX_CAUSE_DEPTH) { json.cause = serializeError(err.cause, depth + 1); } else { json.cause = err.cause; } } for (const key in err) { if (!Object.prototype.hasOwnProperty.call(err, key)) continue; if (key === "message" || key === "name" || key === "stack" || key === "cause") continue; json[key] = err[key]; } return json; } function rehydrateError(s) { const cause = s.cause !== void 0 ? s.cause && typeof s.cause === "object" && "message" in s.cause && "name" in s.cause ? rehydrateError(s.cause) : s.cause : void 0; const err = cause !== void 0 ? new Error(s.message, { cause }) : new Error(s.message); if (s.name) err.name = s.name; if (s.stack !== void 0) err.stack = s.stack; for (const key in s) { if (!Object.prototype.hasOwnProperty.call(s, key)) continue; if (key === "message" || key === "name" || key === "stack" || key === "cause") continue; err[key] = s[key]; } return err; } // src/events/codec/registry.ts var classRegistry = /* @__PURE__ */ new Map(); function registerClass(name, codec) { classRegistry.set(name, codec); } function getClassCodec(name) { return classRegistry.get(name); } // src/events/codec/registrations.ts (() => { registerClass("DefaultGeneratedFile", { toData: (f) => ({ data: f.base64, mediaType: f.mediaType }), fromData: (d) => new DefaultGeneratedFile({ data: d.data, mediaType: d.mediaType }) }); registerClass("DefaultGeneratedFileWithType", { toData: (f) => ({ data: f.base64, mediaType: f.mediaType }), fromData: (d) => new DefaultGeneratedFileWithType({ data: d.data, mediaType: d.mediaType }) }); registerClass("DefaultStepResult", { toData: (s) => ({ content: s.content, finishReason: s.finishReason, usage: s.usage, warnings: s.warnings, request: s.request, response: s.response, providerMetadata: s.providerMetadata, tripwire: s.tripwire }), fromData: (d) => new DefaultStepResult(d) }); return true; })(); // src/events/codec/tags.ts var CODEC_TAG = "__m_codec__"; var MAX_REGEXP_SOURCE_LENGTH = 1024; function isEnvelope(value) { const tag = value[CODEC_TAG]; if (typeof tag !== "string") return false; switch (tag) { case "Undefined": return true; case "Date": case "BigInt": case "URL": return typeof value.v === "string"; case "RegExp": { const v = value.v; if (!v || typeof v.source !== "string" || typeof v.flags !== "string") return false; if (v.source.length > MAX_REGEXP_SOURCE_LENGTH) return false; return /^[dgimsuvy]*$/.test(v.flags) && new Set(v.flags).size === v.flags.length; } case "Map": case "Set": return Array.isArray(value.v); case "Error": return typeof value.v === "object" && value.v !== null; case "Class": return typeof value.n === "string"; default: return false; } } // src/events/codec/codec.ts function encode(value) { return walk(value, /* @__PURE__ */ new WeakSet()); } function walk(v, seen) { if (v === void 0) return { [CODEC_TAG]: "Undefined" }; if (v === null) return null; const t = typeof v; if (t === "string" || t === "boolean") return v; if (t === "number") return Number.isFinite(v) ? v : null; if (t === "bigint") return { [CODEC_TAG]: "BigInt", v: v.toString() }; if (t === "function" || t === "symbol") return void 0; if (v instanceof Date) { return { [CODEC_TAG]: "Date", v: v.toISOString() }; } if (v instanceof RegExp) { return { [CODEC_TAG]: "RegExp", v: { source: v.source, flags: v.flags } }; } if (v instanceof URL) { return { [CODEC_TAG]: "URL", v: v.toString() }; } if (v instanceof Error) { return { [CODEC_TAG]: "Error", v: serializeError(v) }; } if (v instanceof Map) { if (seen.has(v)) return null; seen.add(v); const entries = []; for (const [k, val] of v.entries()) { entries.push([walk(k, seen), walk(val, seen)]); } return { [CODEC_TAG]: "Map", v: entries }; } if (v instanceof Set) { if (seen.has(v)) return null; seen.add(v); const values = []; for (const x of v) values.push(walk(x, seen)); return { [CODEC_TAG]: "Set", v: values }; } if (Array.isArray(v)) { if (seen.has(v)) return null; seen.add(v); return v.map((x) => walk(x, seen)); } if (t === "object") { if (seen.has(v)) return null; seen.add(v); const maybeToJSON = v.toJSON; if (typeof maybeToJSON === "function") { return walk(maybeToJSON.call(v), seen); } const ctor = v.constructor; const ctorName = ctor?.name; if (ctorName && ctorName !== "Object") { const reg = getClassCodec(ctorName); if (reg) { return { [CODEC_TAG]: "Class", n: ctorName, v: walk(reg.toData(v), seen) }; } } const out = {}; for (const k of Object.keys(v)) { if (k === "__proto__") continue; const raw = v[k]; if (raw === void 0) { out[k] = { [CODEC_TAG]: "Undefined" }; continue; } const encoded = walk(raw, seen); if (encoded === void 0) continue; out[k] = encoded; } return out; } return v; } function decodeRegExpEnvelope(v) { if (!v || typeof v !== "object") return /(?:)/; const candidate = v; if (typeof candidate.source !== "string") return /(?:)/; if (typeof candidate.flags !== "string") return /(?:)/; if (candidate.source.length > MAX_REGEXP_SOURCE_LENGTH) return /(?:)/; const flags = candidate.flags; if (!/^[dgimsuvy]*$/.test(flags)) return /(?:)/; if (new Set(flags).size !== flags.length) return /(?:)/; try { return new RegExp(candidate.source, flags); } catch { return /(?:)/; } } function decode(value) { if (value === null) return null; if (typeof value !== "object") return value; if (Array.isArray(value)) return value.map(decode); if (CODEC_TAG in value && isEnvelope(value)) { const env = value; switch (env[CODEC_TAG]) { case "Undefined": return void 0; case "Date": return new Date(env.v); case "BigInt": return BigInt(env.v); case "RegExp": return decodeRegExpEnvelope(env.v); case "URL": return new URL(env.v); case "Map": return new Map(env.v.map(([k, val]) => [decode(k), decode(val)])); case "Set": return new Set(env.v.map(decode)); case "Error": return rehydrateError(decodeSerializedError(env.v)); case "Class": { const reg = getClassCodec(env.n); const data = decode(env.v); return reg ? reg.fromData(data) : data; } } } const out = {}; for (const k of Object.keys(value)) { if (k === "__proto__") continue; out[k] = decode(value[k]); } return out; } function decodeSerializedError(s) { const decoded = decode(s); return decoded; } // src/events/unix-socket-pubsub.ts var DEFAULT_MAX_REMOTE_CLIENT_QUEUED_BYTES = 64 * 1024 * 1024; var MAX_LOCAL_REDELIVERIES = 6; var REDELIVERY_DELAY_MS = 100; function serializeFrame(frame) { return `${JSON.stringify(encode(frame))} `; } function writeSerializedFrame(socket, serializedFrame) { return new Promise((resolve, reject) => { let writeCompleted = false; let drainCompleted = true; let settled = false; const cleanup = () => { socket.off("error", onError); socket.off("close", onClose); socket.off("drain", onDrain); }; const settle = (error) => { if (settled) return; settled = true; cleanup(); if (error) { reject(error); return; } resolve(); }; const maybeResolve = () => { if (writeCompleted && drainCompleted) { settle(); } }; const onError = (error) => settle(error); const onClose = () => settle(new Error("UnixSocketPubSub socket closed before write completed")); const onDrain = () => { drainCompleted = true; maybeResolve(); }; socket.once("error", onError); socket.once("close", onClose); let drained; try { drained = socket.write(serializedFrame, (error) => { if (error) { settle(error); return; } writeCompleted = true; maybeResolve(); }); } catch (error) { settle(error); return; } if (!drained) { drainCompleted = false; socket.once("drain", onDrain); } }); } function writeFrame(socket, frame) { return writeSerializedFrame(socket, serializeFrame(frame)); } 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(decode(JSON.parse(line))); } catch { } } }); } var UnixSocketPubSub = class extends PubSub { socketPath; #server; #clientSocket; #isBroker = false; #closed = false; #starting; #callbacks = /* @__PURE__ */ new Map(); #subscribeWaiters = /* @__PURE__ */ new Map(); #brokerClients = /* @__PURE__ */ new Map(); #pendingWrites = /* @__PURE__ */ new Set(); #recovering; #maxRemoteClientQueuedBytes; constructor(socketPath, options = {}) { super(); this.socketPath = socketPath; this.#maxRemoteClientQueuedBytes = options.maxRemoteClientQueuedBytes ?? DEFAULT_MAX_REMOTE_CLIENT_QUEUED_BYTES; } get supportedModes() { return ["push"]; } get isBroker() { return this.#isBroker; } /** Number of remote clients currently connected to this broker. Always 0 for non-broker instances. */ get remoteClientCount() { return this.#isBroker ? this.#brokerClients.size : 0; } async publish(topic, event, options) { await this.#ensureStarted(); if (options?.localOnly) { const localEvent = { ...event, id: randomUUID(), createdAt: /* @__PURE__ */ new Date(), deliveryAttempt: 1 }; this.#deliverLocal(topic, localEvent); return; } if (this.#isBroker) { await this.#publishFromBroker(topic, event, void 0, options?.localOnly); return; } const socket = this.#clientSocket; if (!socket || socket.destroyed) { await this.#ensureStarted(true); } await this.#sendToBroker({ type: "publish", topic, event, localOnly: options?.localOnly }); } 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]); } 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()]) { this.#removeBrokerClient(client); } if (this.#server) { await new Promise((resolve) => this.#server?.close(() => resolve())); this.#server = void 0; } if (this.#isBroker) { await 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 mkdir(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" && code !== "EEXIST") 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") { this.#throwIfClosed(); await this.#electBroker(); return; } throw error; } } #throwIfClosed() { if (this.#closed) { throw new Error("UnixSocketPubSub is closed"); } } #listen() { return new Promise((resolve, reject) => { const server = net.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.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", () => this.#handleClientDisconnect(socket, new Error("UnixSocketPubSub broker connection closed")) ); socket.on("error", (error) => this.#handleClientDisconnect(socket, 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); } } #handleClientDisconnect(socket, error) { if (this.#clientSocket !== socket) return; this.#clientSocket = void 0; this.#rejectSubscribeWaiters(error); if (!this.#closed) { void this.#recoverClientConnection(); } } async #recoverClientConnection() { if (this.#recovering) return this.#recovering; this.#recovering = this.#recoverClientConnectionLoop().finally(() => { this.#recovering = void 0; }); return this.#recovering; } async #recoverClientConnectionLoop() { while (!this.#closed && !this.#isBroker && !(this.#clientSocket && !this.#clientSocket.destroyed)) { try { await this.#ensureStarted(true); return; } catch { if (this.#closed) return; await new Promise((resolve) => setTimeout(resolve, 10)); } } } /** * Serializes broker election across processes using an exclusive lock file. * Only the lock winner unlinks the stale socket and listens; losers wait * then connect as clients to the newly elected broker. */ async #electBroker() { const lockPath = this.socketPath + ".elect"; let lockFd; try { lockFd = await open(lockPath, "wx"); } catch (e) { if (e.code === "EEXIST") { if (await this.#isElectionLockStale(lockPath)) { await unlink(lockPath).catch(() => { }); throw new Error("Stale broker election lock removed"); } await new Promise((resolve) => setTimeout(resolve, 150)); try { await this.#connectClient(); this.#throwIfClosed(); return; } catch { throw new Error("Broker election in progress by another process"); } } throw e; } try { try { await this.#connectClient(); this.#throwIfClosed(); return; } catch { } await unlink(this.socketPath).catch(() => { }); this.#throwIfClosed(); await this.#listen(); this.#throwIfClosed(); this.#isBroker = true; } finally { await lockFd.close().catch(() => { }); await unlink(lockPath).catch(() => { }); } } async #isElectionLockStale(lockPath) { try { const lockStat = await stat(lockPath); return Date.now() - lockStat.mtimeMs > 2e3; } catch { return true; } } 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(), writeChain: Promise.resolve(), queuedBytes: 0 }; this.#brokerClients.set(socket, client); readFrames(socket, (frame) => { const clientFrame = frame; if (clientFrame.type === "subscribe") { client.subscriptions.add(clientFrame.topic); this.#enqueueBrokerClientWrite(client, { type: "subscribed", topic: clientFrame.topic }); } else if (clientFrame.type === "unsubscribe") { client.subscriptions.delete(clientFrame.topic); } else if (clientFrame.type === "publish") { void this.#publishFromBroker(clientFrame.topic, clientFrame.event, client, clientFrame.localOnly); } }); socket.on("close", () => this.#removeBrokerClient(client)); socket.on("error", () => this.#removeBrokerClient(client)); } #enqueueBrokerClientWrite(client, frame) { if (this.#brokerClients.get(client.socket) !== client || client.socket.destroyed) return; const serializedFrame = serializeFrame(frame); const queuedBytes = Buffer.byteLength(serializedFrame); if (client.queuedBytes + queuedBytes > this.#maxRemoteClientQueuedBytes) { this.#removeBrokerClient(client); return; } client.queuedBytes += queuedBytes; const write = client.writeChain.catch(() => { }).then(async () => { if (this.#brokerClients.get(client.socket) !== client || client.socket.destroyed) return; await writeSerializedFrame(client.socket, serializedFrame); }).catch(() => { this.#removeBrokerClient(client); }).finally(() => { client.queuedBytes = Math.max(0, client.queuedBytes - queuedBytes); }); client.writeChain = write; this.#pendingWrites.add(write); void write.finally(() => this.#pendingWrites.delete(write)); } #removeBrokerClient(client) { if (this.#brokerClients.get(client.socket) !== client) return; this.#brokerClients.delete(client.socket); client.subscriptions.clear(); client.queuedBytes = 0; client.writeChain = Promise.resolve(); if (!client.socket.destroyed) { client.socket.destroy(); } } #handleServerFrame(frame) { if (frame.type === "subscribed") { this.#settleSubscribeWaiters(frame.topic); return; } if (frame.type !== "event") return; this.#deliverLocal(frame.topic, frame.event); } async #publishFromBroker(topic, event, sourceClient, localOnly) { const brokerEvent = { ...event, id: randomUUID(), createdAt: /* @__PURE__ */ new Date(), deliveryAttempt: 1 }; this.#deliverLocal(topic, brokerEvent); if (this.#brokerClients.size === 0) return; if (localOnly) { if (sourceClient && sourceClient.subscriptions.has(topic) && !sourceClient.socket.destroyed) { this.#enqueueBrokerClientWrite(sourceClient, { type: "event", topic, event: brokerEvent }); } return; } let frame; for (const client of this.#brokerClients.values()) { if (!client.subscriptions.has(topic) || client.socket.destroyed) continue; frame ??= { type: "event", topic, event: brokerEvent }; this.#enqueueBrokerClientWrite(client, frame); } } #deliverLocal(topic, event) { const callbacks = this.#callbacks.get(topic); if (!callbacks) return; for (const cb of callbacks) { this.#invokeLocalCallback(topic, event, cb, 0); } } #invokeLocalCallback(topic, event, cb, attempt) { let nacked = false; const nack = async () => { if (nacked || this.#closed) return; nacked = true; if (attempt >= MAX_LOCAL_REDELIVERIES) return; const stillSubscribed = this.#callbacks.get(topic)?.has(cb); if (!stillSubscribed) return; const timer = setTimeout( () => { if (this.#closed) return; if (!this.#callbacks.get(topic)?.has(cb)) return; const redeliveredEvent = { ...event, deliveryAttempt: (event.deliveryAttempt ?? 1) + 1 }; this.#invokeLocalCallback(topic, redeliveredEvent, cb, attempt + 1); }, REDELIVERY_DELAY_MS * (attempt + 1) ); timer.unref?.(); }; try { const result = cb( event, async () => { }, nack ); if (result && typeof result.catch === "function") { void result.catch(() => { }); } } catch { } } async #sendToBroker(frame) { const maxRetries = 3; let lastError; for (let attempt = 0; attempt <= maxRetries; attempt++) { try { if (attempt === 0) { await this.#sendToActiveBroker(frame); } else { if (this.#closed) throw lastError; const failedSocket = this.#clientSocket; this.#clientSocket = void 0; failedSocket?.destroy(); await this.#ensureStarted(true); await this.#sendToActiveBroker(frame); } return; } catch (error) { lastError = error; if (this.#closed) throw error; const code = error?.code; const transient = code === "EPIPE" || code === "ECONNRESET" || code === "ENOTCONN" || error?.message?.includes("socket closed before write completed") || error?.message?.includes("broker connection closed") || error?.message?.includes("not connected to a broker"); if (!transient || attempt === maxRetries) throw error; await new Promise((resolve) => setTimeout(resolve, 10 * (attempt + 1))); } } } async #sendToActiveBroker(frame) { const socket = this.#clientSocket; if (!socket || socket.destroyed) { await this.#ensureStarted(true); } if (this.#isBroker) { await this.#handlePromotedBrokerFrame(frame); return; } const activeSocket = this.#clientSocket; if (!activeSocket || activeSocket.destroyed) { throw new Error("UnixSocketPubSub is not connected to a broker"); } await writeFrame(activeSocket, frame); } async #handlePromotedBrokerFrame(frame) { if (frame.type === "subscribe") { this.#settleSubscribeWaiters(frame.topic); } else if (frame.type === "publish") { await this.#publishFromBroker(frame.topic, frame.event); } } }; export { MAX_LOCAL_REDELIVERIES, UnixSocketPubSub }; //# sourceMappingURL=index.js.map //# sourceMappingURL=index.js.map