assistant-cloud/dist/CloudMessagePersistence.js

Cloud integration for assistant-ui

import { CloudAPIError } from "./AssistantCloudAPI.js";
// src/CloudMessagePersistence.ts
const CLOUD_MESSAGE_PAGE_SIZE = 200;
const isOpen = (run) => typeof run === "object";
/** How long a running lease may go without a heartbeat before its message reads as stale. */
const DEFAULT_RUN_STALE_AFTER_MS = 3e4;
/** 409 codes after which the lease is gone for good: nothing more is written for its message. */
const LOST_LEASE_CODES = /* @__PURE__ */ new Set(["run_ended", "lease_mismatch"]);
const lostLeaseCode = (error) => error instanceof CloudAPIError && error.status === 409 && error.code !== void 0 && LOST_LEASE_CODES.has(error.code) ? error.code : void 0;
const isRunUnsupported = (error) => error instanceof CloudAPIError && error.status === 400;
/**
* Appends, updates and loads cloud messages while mapping local ids to cloud
* ids and chaining parent_id. A parent that is still being created is awaited,
* so concurrent appends land under the right parent. A parent whose create
* failed is created again before its child, and a child whose parent still
* cannot be created is stored as a root.
*
* A message appended as running holds a run lease, renewed by heartbeats on
* the server's interval until a write with an end reason ends it. Heartbeats
* keep running while the page is in the background. The server never
* expires a lease; readers show a run as stale once it has gone `staleAfter`
* without a beat, and a late run keeps writing its own message. A lease that
* is lost (ended elsewhere or refused) is final: the message gets no further
* writes from this persistence.
*/
var CloudMessagePersistence = class {
	idMapping = /* @__PURE__ */ new Map();
	failedAppends = /* @__PURE__ */ new Map();
	/** Every message this persistence opened a lease for; a repeated append of one writes its row instead of storing another. */
	runs = /* @__PURE__ */ new Map();
	runWrites = /* @__PURE__ */ new Map();
	/** Writes waiting on an in-flight create, chained so they apply in append order. */
	pendingWrites = /* @__PURE__ */ new Map();
	/** The deletes each write waiting on an in-flight create saw start; each resolves to whether it deleted the message. */
	waiters = /* @__PURE__ */ new Map();
	getCloud;
	staleAfter;
	constructor(cloud, options = {}) {
		this.getCloud = typeof cloud === "function" ? cloud : () => cloud;
		this.staleAfter = options.staleAfter ?? 3e4;
	}
	/**
	* Reads a listed message's run against `staleAfter`; `sinceReadMs` is the
	* time since the list was read. Undefined for no run or an ended one.
	*/
	runStateOf(run, sinceReadMs = 0) {
		if (run?.status !== "running") return void 0;
		const quietMs = run.last_heartbeat_ms_ago + sinceReadMs;
		return quietMs >= this.staleAfter ? { status: "stale" } : {
			status: "in-progress",
			staleInMs: this.staleAfter - quietMs
		};
	}
	/**
	* Persist a message to the cloud.
	*
	* @param threadId - Remote thread ID
	* @param messageId - Local message ID (used for tracking)
	* @param parentId - Local parent message ID (or null for first message)
	* @param format - Message format (e.g., "aui/v0", "ai-sdk/v6")
	* @param content - Message content (format-specific)
	* @param run - `running` stores the message with an open run lease that is
	*   kept alive by heartbeats until a write with an end reason closes it.
	*   Once the lease is lost (it was ended elsewhere or the server refused
	*   it), later writes for the message are dropped.
	*/
	async append(threadId, messageId, parentId, format, content, run) {
		if (this.runs.get(messageId) === "lost") return;
		const cloud = this.getCloud();
		const existing = this.idMapping.get(messageId);
		if (existing instanceof Promise) {
			const mapping = this.idMapping;
			const deletes = [];
			const waiters = this.waiters.get(messageId) ?? /* @__PURE__ */ new Set();
			waiters.add(deletes);
			this.waiters.set(messageId, waiters);
			const write = (this.pendingWrites.get(messageId) ?? existing).catch(() => void 0).then(async () => {
				for (let i = 0; i < deletes.length; i++) if (await deletes[i]) return;
				if (this.idMapping !== mapping) return;
				if (this.idMapping.has(messageId)) await this.update(threadId, messageId, format, content, run);
				else await this.append(threadId, messageId, parentId, format, content, run);
			});
			this.pendingWrites.set(messageId, write);
			try {
				await write;
			} finally {
				if (this.pendingWrites.get(messageId) === write) this.pendingWrites.delete(messageId);
				waiters.delete(deletes);
				if (waiters.size === 0 && this.waiters.get(messageId) === waiters) this.waiters.delete(messageId);
			}
			return;
		}
		if (existing !== void 0 && this.runs.has(messageId)) {
			await this.update(threadId, messageId, format, content, run);
			return;
		}
		this.failedAppends.delete(messageId);
		const task = (async () => {
			const body = {
				parent_id: parentId ? await this.resolveParentId(parentId) : null,
				...messageId.length > 0 && messageId.length <= 255 ? { external_id: messageId } : void 0,
				format,
				content
			};
			if (run !== "running") return (await cloud.threads.messages.create(threadId, body)).message_id;
			let created;
			try {
				created = await cloud.threads.messages.create(threadId, {
					...body,
					open_run: true
				});
			} catch (error) {
				if (!isRunUnsupported(error)) throw error;
				created = await cloud.threads.messages.create(threadId, body);
			}
			if (created.run) this.openRun(threadId, messageId, created.run);
			return created.message_id;
		})();
		this.idMapping.set(messageId, task);
		try {
			const remoteId = await task;
			if (this.idMapping.get(messageId) === task) this.idMapping.set(messageId, remoteId);
		} catch (err) {
			if (this.idMapping.get(messageId) === task) {
				this.idMapping.delete(messageId);
				this.failedAppends.set(messageId, {
					threadId,
					parentId,
					format,
					content
				});
			}
			throw err;
		}
	}
	async resolveParentId(parentId) {
		await Promise.resolve(this.idMapping.get(parentId)).catch(() => void 0);
		const failedParent = this.failedAppends.get(parentId);
		try {
			if (failedParent && !this.idMapping.has(parentId)) await this.append(failedParent.threadId, parentId, failedParent.parentId, failedParent.format, failedParent.content);
			return await this.idMapping.get(parentId) ?? parentId;
		} catch {
			return null;
		}
	}
	/**
	* Update an already-persisted message in the cloud.
	*/
	async update(threadId, messageId, _format, content, run) {
		if (this.runs.get(messageId) === "lost") return;
		const remoteId = await this.getRemoteId(messageId);
		if (this.runs.get(messageId) === "lost") return;
		if (!remoteId) {
			console.warn(`Skipping update for message ${messageId}: no remote id is mapped.`);
			return;
		}
		if (!this.runs.has(messageId)) {
			await this.getCloud().threads.messages.update(threadId, remoteId, { content });
			return;
		}
		const task = (this.runWrites.get(messageId) ?? Promise.resolve()).catch(() => void 0).then(() => this.writeRun(threadId, messageId, remoteId, content, run ?? "done"));
		this.runWrites.set(messageId, task);
		try {
			await task;
		} finally {
			if (this.runWrites.get(messageId) === task) this.runWrites.delete(messageId);
		}
	}
	async writeRun(threadId, messageId, remoteId, content, phase) {
		const open = this.runs.get(messageId);
		if (open === "lost") return;
		const cloud = this.getCloud();
		if (!isOpen(open)) {
			if (phase === "running") return;
			await cloud.threads.messages.update(threadId, remoteId, { content });
			return;
		}
		try {
			await cloud.threads.messages.update(open.threadId, remoteId, {
				content,
				lease_id: open.leaseId
			});
		} catch (error) {
			const code = lostLeaseCode(error);
			if (!code) throw error;
			this.loseRun(messageId, code);
			return;
		}
		if (phase === "running") {
			this.scheduleHeartbeat(messageId);
			return;
		}
		this.closeRun(messageId, "ended");
		await cloud.threads.runs.end(open.threadId, open.leaseId, { reason: phase }).catch((error) => {
			if (!(error instanceof CloudAPIError) || error.status >= 500) throw error;
		});
	}
	/**
	* Moves a stored message's id mapping and run lease to a new local id, for
	* runtimes that swap a placeholder's id once the backend names the message.
	*/
	async rename(from, to) {
		const entry = this.idMapping.get(from);
		if (entry === void 0 || this.idMapping.has(to)) return;
		await Promise.resolve(entry).catch(() => void 0);
		const remoteId = this.idMapping.get(from);
		if (typeof remoteId !== "string" || this.idMapping.has(to)) return;
		this.idMapping.delete(from);
		this.idMapping.set(to, remoteId);
		const run = this.runs.get(from);
		if (run === void 0) return;
		this.runs.delete(from);
		this.runs.set(to, run);
		if (isOpen(run)) this.scheduleHeartbeat(to);
	}
	/** Whether a run lease this persistence opened for the message is still held. */
	hasOpenRun(messageId) {
		return isOpen(this.runs.get(messageId));
	}
	openRun(threadId, messageId, lease) {
		const previous = this.runs.get(messageId);
		if (isOpen(previous) && previous.timer !== void 0) clearTimeout(previous.timer);
		this.runs.set(messageId, {
			threadId,
			leaseId: lease.lease_id,
			intervalMs: lease.heartbeat_interval_ms,
			timer: void 0
		});
		this.scheduleHeartbeat(messageId);
	}
	scheduleHeartbeat(messageId) {
		const open = this.runs.get(messageId);
		if (!isOpen(open)) return;
		if (open.timer !== void 0) clearTimeout(open.timer);
		open.timer = setTimeout(() => {
			open.timer = void 0;
			this.heartbeat(messageId, open);
		}, open.intervalMs);
	}
	async heartbeat(messageId, open) {
		if (this.runs.get(messageId) !== open) return;
		try {
			open.intervalMs = (await this.getCloud().threads.runs.heartbeat(open.threadId, open.leaseId)).heartbeat_interval_ms;
		} catch (error) {
			if (this.runs.get(messageId) !== open) return;
			if (error instanceof CloudAPIError && (error.status === 409 || error.status === 404)) {
				this.loseRun(messageId, error.code ?? String(error.status));
				return;
			}
		}
		if (this.runs.get(messageId) === open) this.scheduleHeartbeat(messageId);
	}
	/** The lease is gone for good: stop heartbeating and drop every later write for the message. */
	loseRun(messageId, reason) {
		if (this.runs.get(messageId) === "lost") return;
		this.closeRun(messageId, "lost");
		console.warn(`[assistant-cloud] Lost the run lease of message ${messageId} (${reason}); later writes for it are dropped.`);
	}
	closeRun(messageId, as) {
		const open = this.runs.get(messageId);
		if (isOpen(open) && open.timer !== void 0) clearTimeout(open.timer);
		this.runs.set(messageId, as);
	}
	/**
	* Delete mapped messages from the cloud and clear their local ID mappings.
	*/
	async delete(threadId, messageIds) {
		let settle;
		const deleting = new Promise((resolve) => {
			settle = resolve;
		});
		for (const localId of messageIds) for (const deletes of this.waiters.get(localId) ?? []) deletes.push(deleting);
		let deleted = false;
		try {
			await this.deleteMessages(threadId, messageIds);
			deleted = true;
		} finally {
			settle(deleted);
		}
	}
	async deleteMessages(threadId, messageIds) {
		const cloud = this.getCloud();
		const mapping = this.idMapping;
		const resolved = [];
		for (const localId of messageIds) {
			const entry = mapping.get(localId);
			if (!entry) continue;
			try {
				resolved.push({
					localId,
					entry,
					remoteId: await entry
				});
			} catch {
				continue;
			}
		}
		for (const localId of messageIds) this.failedAppends.delete(localId);
		if (resolved.length === 0) return;
		try {
			await cloud.threads.messages.delete(threadId, resolved.map(({ remoteId }) => remoteId));
		} catch (error) {
			if (!(error instanceof CloudAPIError && error.status === 404)) throw error;
		}
		for (const { localId, entry, remoteId } of resolved) {
			const current = mapping.get(localId);
			if (current === entry || current === remoteId) mapping.delete(localId);
		}
	}
	/**
	* Check if a message has been persisted (or is currently being persisted).
	*/
	isPersisted(messageId) {
		return this.idMapping.has(messageId);
	}
	/**
	* Get the remote ID for a local message ID (resolved).
	* Returns undefined if not persisted.
	*/
	async getRemoteId(messageId) {
		const entry = this.idMapping.get(messageId);
		if (!entry) return void 0;
		return entry;
	}
	getResolvedRemoteId(messageId) {
		const entry = this.idMapping.get(messageId);
		return typeof entry === "string" ? entry : void 0;
	}
	record(localId, remoteId) {
		this.idMapping.set(localId, remoteId);
	}
	/**
	* Load messages from the cloud and populate the ID mapping.
	*
	* The list endpoint caps a response at 200 rows, so pages are followed by
	* message ID cursor until a short page and concatenated in server order.
	*
	* The ID mapping is populated so that `isPersisted()` returns true for
	* loaded messages, preventing re-persistence of already-stored messages.
	*
	* A loaded ID that an append already maps keeps the remote ID from that append, and falls back to the loaded ID if the append fails.
	*
	* @param threadId - Remote thread ID
	* @param format - Optional format filter
	* @returns Array of cloud messages
	*/
	async load(threadId, format) {
		const idMapping = this.idMapping;
		const cloud = this.getCloud();
		const messages = [];
		const seen = /* @__PURE__ */ new Set();
		let after;
		while (true) {
			const page = await cloud.threads.messages.list(threadId, {
				...format ? { format } : void 0,
				limit: CLOUD_MESSAGE_PAGE_SIZE,
				...after ? { after } : void 0
			});
			const last = page.messages.at(-1);
			if (!last) break;
			const fresh = page.messages.filter((m) => !seen.has(m.id));
			if (fresh.length === 0) break;
			for (const m of fresh) seen.add(m.id);
			messages.push(...fresh);
			if (page.messages.length < CLOUD_MESSAGE_PAGE_SIZE) break;
			after = last.id;
		}
		if (this.idMapping === idMapping) for (const m of messages) {
			const entry = idMapping.get(m.id);
			if (entry === void 0) idMapping.set(m.id, m.id);
			else if (entry instanceof Promise) entry.catch(() => {
				const current = idMapping.get(m.id);
				if (current === void 0 || current === entry) idMapping.set(m.id, m.id);
			});
		}
		return messages;
	}
	/**
	* Reset the ID mapping (call when switching threads).
	*
	* Pending `load()` and `append()` calls are not cancelled and still settle
	* normally, but their results no longer populate the ID mapping.
	*/
	reset() {
		this.idMapping = /* @__PURE__ */ new Map();
		this.failedAppends = /* @__PURE__ */ new Map();
		this.pendingWrites = /* @__PURE__ */ new Map();
		this.waiters = /* @__PURE__ */ new Map();
		for (const run of this.runs.values()) if (isOpen(run) && run.timer !== void 0) clearTimeout(run.timer);
		this.runs = /* @__PURE__ */ new Map();
	}
};
export { CloudMessagePersistence, DEFAULT_RUN_STALE_AFTER_MS };

//# sourceMappingURL=CloudMessagePersistence.js.map