@tanstack/ai-persistence/dist/esm/subagent-runs.js

Composable state persistence for TanStack AI messages, runs, interrupts, metadata, and locks.

import { mergeStoredMessages } from "./merge-stored.js";
import { StreamProcessor, convertMessagesToModelMessages, modelMessagesToUIMessages, subagentHostMessageId, subagentInputOf, wireSubagentInfo } from "@tanstack/ai";
//#region src/subagent-runs.ts
function withSubagentInfo(metadata, info) {
	const source = metadata != null && typeof metadata === "object" ? metadata : {};
	const tanstack = source.tanstack != null && typeof source.tanstack === "object" ? source.tanstack : {};
	return {
		...source,
		tanstack: {
			...tanstack,
			subagent: info
		}
	};
}
function childStoreId(subagentRunId) {
	return `subagent:${subagentRunId}`;
}
/**
* Head metadata key for what a later call of the child needs:
* - `base`: how many stored messages came before the current call. A resume
*   of that call keeps the number.
* - `parentThreadId`: the thread that started the child. Only that thread can
*   continue it.
* - `prompt`: the task of the first call, the first user message of the child.
*/
var CHILD_KEY = "tanstack:subagentChild";
function storedMarker(messages) {
	const marker = messages[0]?.metadata?.[CHILD_KEY];
	const read = (key) => marker !== null && typeof marker === "object" ? Reflect.get(marker, key) : void 0;
	const base = read("base");
	const parentThreadId = read("parentThreadId");
	const prompt = read("prompt");
	return {
		...typeof base === "number" && { base },
		...typeof parentThreadId === "string" && { parentThreadId },
		...typeof prompt === "string" && { prompt }
	};
}
/** The thread a chat writes its children under: its own child thread, if any. */
function ownerThread(caller) {
	return caller.subagentRunId === void 0 ? caller.threadId : childStoreId(caller.subagentRunId);
}
/** The `prompt` of an agent tool call, if any. */
function promptOf(call) {
	if (call === void 0) return void 0;
	let args;
	try {
		args = JSON.parse(call.args);
	} catch {
		return;
	}
	const prompt = args !== null && typeof args === "object" ? Reflect.get(args, "prompt") : void 0;
	return typeof prompt === "string" ? prompt : void 0;
}
function readSubagentRunId(chunk) {
	if (!("subagentRunId" in chunk)) return;
	const id = chunk.subagentRunId;
	return typeof id === "string" && id !== "" ? id : void 0;
}
function assistantId(runId) {
	return subagentHostMessageId(runId);
}
function messageText(message) {
	return typeof message.content === "string" ? message.content : "";
}
function readModelRunId(message) {
	const metadata = message.metadata;
	if (metadata == null || typeof metadata !== "object") return;
	if (!("tanstack" in metadata)) return;
	const tanstack = metadata.tanstack;
	if (tanstack == null || typeof tanstack !== "object") return;
	if (!("runId" in tanstack)) return;
	const runId = tanstack.runId;
	return typeof runId === "string" && runId !== "" ? runId : void 0;
}
function withRunId(message, runId) {
	const metadata = message.metadata;
	const tanstack = metadata != null && typeof metadata === "object" && "tanstack" in metadata && metadata.tanstack != null && typeof metadata.tanstack === "object" ? metadata.tanstack : {};
	return {
		...message,
		metadata: {
			...metadata,
			tanstack: {
				...tanstack,
				runId
			}
		}
	};
}
function sameChildText(stored, incoming) {
	if (stored === "" || incoming === "") return false;
	if (stored === incoming) return true;
	const blocks = stored.split("\n\n").filter((block) => block.includes(":\n"));
	if (blocks.length === 0) return false;
	return blocks.every((block) => incoming.includes(block));
}
function keepSubagentRunIds(stored, merged) {
	const dropped = stored.filter((message) => {
		if (readModelRunId(message) === void 0) return false;
		return !merged.some((item) => item.id !== void 0 && item.id === message.id);
	});
	const turnOf = (list, index) => list.slice(0, index).findLast((message) => message.role === "user")?.id;
	const used = /* @__PURE__ */ new Set();
	for (const previous of dropped) {
		const runId = readModelRunId(previous);
		if (runId === void 0) continue;
		if (merged.some((message) => readModelRunId(message) === runId)) continue;
		const previousText = messageText(previous);
		const previousTurn = turnOf(stored, stored.indexOf(previous));
		const index = merged.findIndex((message, messageIndex) => {
			if (used.has(messageIndex)) return false;
			if (message.role !== "assistant") return false;
			if (readModelRunId(message) !== void 0) return false;
			if (previousText === "") return previousTurn !== void 0 && turnOf(merged, messageIndex) === previousTurn && merged[messageIndex - 1]?.role === "user";
			return sameChildText(previousText, messageText(message));
		});
		if (index === -1) continue;
		const host = merged[index];
		if (!host) continue;
		used.add(index);
		merged[index] = withRunId(host, runId);
	}
	return merged;
}
function childText(messages) {
	return messages.flatMap((message) => message.role === "assistant" ? message.parts.flatMap((part) => part.type === "text" && part.content.trim() !== "" ? [part.content.trim()] : []) : []).join("\n\n");
}
function withoutCards(messages) {
	return messages.map((message) => ({
		...message,
		parts: message.parts.filter((part) => part.type !== "subagent")
	}));
}
/**
* Card data for a stored child transcript. It rides on the first message, in
* `metadata.tanstack.subagent`, the same shape the wire uses.
*/
function storedSubagentInfo(messages) {
	return wireSubagentInfo(messages[0]);
}
/** A stored child transcript without its card-only placeholder. */
function withoutPlaceholder(messages) {
	return messages.filter((message) => storedSubagentInfo([message])?.placeholder !== true);
}
function createSubagentRunRecorder(stores) {
	const children = /* @__PURE__ */ new Map();
	const intervalMs = stores.intervalMs ?? 1e3;
	const answered = /* @__PURE__ */ new Map();
	const parentSavedAt = /* @__PURE__ */ new Map();
	const calls = /* @__PURE__ */ new Map();
	/** Drop the open tool calls of a run that ended. */
	function dropCalls(runId) {
		for (const [id, call] of calls) if (call.runId === runId) calls.delete(id);
	}
	async function loadMessages(threadId) {
		return stores.messages.loadThread(threadId);
	}
	function lineage(subagentRunId) {
		const notes = [];
		let id = subagentRunId;
		for (let depth = 0; id !== void 0 && depth < 64; depth++) {
			const note = children.get(id);
			if (!note) break;
			notes.push(note);
			id = note.parentSubagentRunId;
		}
		return notes;
	}
	async function saveChild(subagentRunId) {
		const note = children.get(subagentRunId);
		if (!note) return;
		note.savedAt = Date.now();
		const info = {
			name: note.name,
			status: note.status,
			...note.parentSubagentRunId !== void 0 && { parentSubagentRunId: note.parentSubagentRunId },
			...note.parentToolCallId !== void 0 && { parentToolCallId: note.parentToolCallId },
			...note.interruptIds !== void 0 && { interruptIds: note.interruptIds },
			...note.error !== void 0 && { error: note.error },
			...note.metadata !== void 0 && { metadata: note.metadata }
		};
		const transcript = convertMessagesToModelMessages(withoutCards(note.processor.getMessages()));
		const marker = { [CHILD_KEY]: {
			base: note.base,
			parentThreadId: note.parentThreadId,
			...note.prompt !== void 0 && { prompt: note.prompt }
		} };
		const [first, ...rest] = transcript;
		const head = first ? {
			...first,
			metadata: {
				...withSubagentInfo(first.metadata, info),
				...marker
			}
		} : {
			id: `child:${subagentRunId}`,
			role: "assistant",
			content: "",
			metadata: {
				...withSubagentInfo(void 0, {
					...info,
					placeholder: true
				}),
				...marker
			}
		};
		await stores.messages.saveThread(childStoreId(subagentRunId), [head, ...rest]);
	}
	async function saveParent(threadId, runId) {
		const notes = [...children.values()].filter((note) => note.parentRunId === runId && note.parentSubagentRunId === void 0 && note.parentToolCallId === void 0);
		if (notes.length === 0) return;
		const content = notes.map((note) => {
			const text = childText(note.processor.getMessages());
			return text === "" ? "" : `${note.name}:\n${text}`;
		}).filter((block) => block !== "").join("\n\n");
		const stored = await loadMessages(threadId);
		const index = stored.findIndex((message) => readModelRunId(message) === runId);
		const host = stored[index];
		if (host) {
			const next = [...stored];
			next[index] = {
				...host,
				content
			};
			await stores.messages.saveThread(threadId, next);
			return;
		}
		const assistant = {
			id: assistantId(runId),
			role: "assistant",
			content,
			metadata: { tanstack: { runId } }
		};
		await stores.messages.saveThread(threadId, [...stored, assistant]);
	}
	async function startChild(input, chunk) {
		const id = chunk.subagentRunId;
		const call = chunk.parentToolCallId === void 0 ? void 0 : calls.get(chunk.parentToolCallId);
		if (chunk.parentToolCallId !== void 0) calls.delete(chunk.parentToolCallId);
		const record = await stores.runs?.createOrResume({
			runId: id,
			threadId: childStoreId(id),
			startedAt: Date.now(),
			parentRunId: chunk.parentSubagentRunId ?? input.runId,
			subagentRunId: id,
			name: chunk.name
		});
		const existing = children.get(id);
		if (existing) {
			existing.threadId = input.threadId;
			existing.runId = input.runId;
			existing.status = "running";
			delete existing.interruptIds;
			if (chunk.metadata !== void 0) existing.metadata = chunk.metadata;
		} else {
			const raw = await loadMessages(childStoreId(id));
			const stored = withoutPlaceholder(raw);
			const marker = storedMarker(raw);
			const resumes = storedSubagentInfo(raw)?.status === "suspended";
			const checked = subagentInputOf(chunk.metadata);
			const prompt = raw.length > 0 ? marker.prompt : checked === void 0 ? promptOf(call) : JSON.stringify(checked);
			const followUp = raw.length > 0 && !resumes && checked === void 0 ? promptOf(call) : void 0;
			const initial = followUp === void 0 ? stored : [...stored, {
				id: `${childStoreId(id)}:prompt:${stored.length}`,
				role: "user",
				content: followUp
			}];
			children.set(id, {
				name: chunk.name,
				threadId: input.threadId,
				runId: input.runId,
				savedAt: 0,
				parentRunId: record?.parentRunId ?? chunk.parentSubagentRunId ?? input.runId,
				...chunk.parentSubagentRunId !== void 0 && { parentSubagentRunId: chunk.parentSubagentRunId },
				...chunk.parentToolCallId !== void 0 && { parentToolCallId: chunk.parentToolCallId },
				...chunk.metadata !== void 0 && { metadata: chunk.metadata },
				processor: new StreamProcessor({
					subagentRunId: id,
					initialMessages: modelMessagesToUIMessages(initial)
				}),
				status: "running",
				base: resumes && marker.base !== void 0 ? marker.base : stored.length,
				parentThreadId: ownerThread({
					threadId: input.threadId,
					subagentRunId: chunk.parentSubagentRunId ?? input.subagentRunId
				}),
				...prompt !== void 0 && { prompt }
			});
		}
		if (record && record.status !== "running") await stores.runs?.update(id, { status: "running" });
		if (chunk.parentSubagentRunId !== void 0) children.get(chunk.parentSubagentRunId)?.processor.processChunk(chunk);
		await saveChild(id);
	}
	/**
	* The stored transcript of a child, for a call that continues it by
	* `sessionId`. It starts with the task of the first call as a user message.
	* `undefined` when `caller` did not start a child with this id. While the
	* child waits on an interrupt, it is the transcript from before that call:
	* the resume adds the newer messages of the child itself.
	*/
	async function loadChild(subagentRunId, caller) {
		const stored = await loadMessages(childStoreId(subagentRunId));
		const marker = storedMarker(stored);
		if (marker.parentThreadId !== ownerThread(caller)) return void 0;
		const messages = withoutPlaceholder(stored);
		const info = storedSubagentInfo(stored);
		const transcript = info?.status === "suspended" && marker.base !== void 0 ? messages.slice(0, marker.base) : messages;
		return {
			messages: [...marker.prompt === void 0 ? [] : [{
				id: `${childStoreId(subagentRunId)}:prompt`,
				role: "user",
				content: marker.prompt
			}], ...transcript],
			...info !== void 0 && { agent: info.name }
		};
	}
	async function settleChild(chunk) {
		const id = chunk.subagentRunId;
		const note = children.get(id);
		if (!note) return;
		try {
			await writeSettled(note, id, chunk);
		} finally {
			if (note.parentSubagentRunId !== void 0 || note.parentToolCallId !== void 0) children.delete(id);
		}
	}
	async function writeSettled(note, id, chunk) {
		note.processor.finalizeStream();
		if (note.parentSubagentRunId !== void 0) children.get(note.parentSubagentRunId)?.processor.processChunk(chunk);
		if (chunk.type === "SUBAGENT_ERROR") {
			const stopped = chunk.message === "Stopped";
			note.status = "error";
			note.error = {
				message: chunk.message,
				...chunk.code !== void 0 && { code: chunk.code }
			};
			await saveChild(id);
			await stores.runs?.update(id, {
				status: stopped ? "aborted" : "failed",
				finishedAt: Date.now(),
				...!stopped ? { error: { message: chunk.message } } : {}
			});
			return;
		}
		if (chunk.outcome?.type === "suspended") {
			note.status = "suspended";
			note.interruptIds = chunk.outcome.interruptIds ?? [];
			await saveChild(id);
			await stores.runs?.update(id, { status: "interrupted" });
			return;
		}
		note.status = "finished";
		await saveChild(id);
		await stores.runs?.update(id, {
			status: "completed",
			finishedAt: Date.now()
		});
	}
	async function commitAnswers(runId) {
		const entries = answered.get(runId) ?? [];
		answered.delete(runId);
		for (const entry of entries) if (entry.status === "cancelled") await stores.interrupts?.cancel(entry.interruptId);
		else await stores.interrupts?.resolve(entry.interruptId, entry.payload);
	}
	/** The notes a run fed. */
	function notesOf(runId) {
		return [...children].filter(([, note]) => note.runId === runId);
	}
	function forget(runId) {
		for (const [id] of notesOf(runId)) children.delete(id);
		parentSavedAt.delete(runId);
	}
	async function settleOpenChildren(runId, status, error) {
		for (const [subagentRunId] of notesOf(runId)) {
			await saveChild(subagentRunId);
			const current = await stores.runs?.get(subagentRunId);
			if (current && current.status !== "running") continue;
			await stores.runs?.update(subagentRunId, {
				status,
				finishedAt: Date.now(),
				...error ? { error } : {}
			});
		}
	}
	/** Write the parent messages of this thread's routed children. */
	async function saveParents(threadId) {
		const runIds = new Set([...children.values()].filter((note) => note.threadId === threadId).map((note) => note.parentRunId));
		for (const runId of runIds) await saveParent(threadId, runId);
	}
	return {
		loadChild,
		async start(input) {
			await stores.runs?.createOrResume({
				runId: input.runId,
				threadId: input.threadId,
				startedAt: Date.now()
			});
			if (input.resume?.length) answered.set(input.runId, [...input.resume]);
			const incoming = convertMessagesToModelMessages([...input.messages]);
			const stored = await loadMessages(input.threadId);
			const merged = mergeStoredMessages(stored, incoming);
			await stores.messages.saveThread(input.threadId, keepSubagentRunIds(stored, merged));
		},
		async chunk(input) {
			const chunk = input.chunk;
			if (chunk.type === "TOOL_CALL_START") calls.set(chunk.toolCallId, {
				name: chunk.toolCallName,
				args: "",
				runId: input.runId
			});
			if (chunk.type === "TOOL_CALL_ARGS") {
				const call = calls.get(chunk.toolCallId);
				if (call) call.args += chunk.delta;
			}
			if (chunk.type === "TOOL_CALL_RESULT") calls.delete(chunk.toolCallId);
			if (chunk.type === "SUBAGENT_STARTED") {
				await startChild(input, chunk);
				return;
			}
			if (chunk.type === "SUBAGENT_FINISHED" || chunk.type === "SUBAGENT_ERROR") {
				await settleChild(chunk);
				await saveParents(input.threadId);
				return;
			}
			const subagentRunId = readSubagentRunId(chunk);
			if (!subagentRunId) return;
			const notes = lineage(subagentRunId);
			if (notes.length === 0) return;
			for (const note of notes) note.processor.processChunk(chunk);
			const now = Date.now();
			const note = notes[0];
			if (note && now - note.savedAt >= intervalMs) await saveChild(subagentRunId);
			if (chunk.type === "TEXT_MESSAGE_CONTENT" && now - (parentSavedAt.get(input.runId) ?? 0) >= intervalMs) {
				parentSavedAt.set(input.runId, now);
				await saveParents(input.threadId);
			}
		},
		async suspend(input) {
			await commitAnswers(input.runId);
			await saveParents(input.threadId);
			for (const [subagentRunId] of notesOf(input.runId)) await saveChild(subagentRunId);
			for (const interrupt of input.interrupts) await stores.interrupts?.create({
				interruptId: interrupt.id,
				runId: input.runId,
				threadId: input.threadId,
				requestedAt: Date.now(),
				payload: { ...interrupt }
			});
			await stores.runs?.update(input.runId, { status: "interrupted" });
			forget(input.runId);
		},
		async finish(input) {
			dropCalls(input.runId);
			await commitAnswers(input.runId);
			await saveParents(input.threadId);
			await settleOpenChildren(input.runId, "completed");
			await stores.runs?.update(input.runId, {
				status: "completed",
				finishedAt: Date.now()
			});
			forget(input.runId);
		},
		async abort(input) {
			dropCalls(input.runId);
			const aborted = input.error instanceof Error && input.error.name === "AbortError";
			const message = input.error instanceof Error ? input.error.message : "Run failed";
			await settleOpenChildren(input.runId, aborted ? "aborted" : "failed", aborted ? void 0 : { message });
			await saveParents(input.threadId);
			await stores.runs?.update(input.runId, {
				status: aborted ? "aborted" : "failed",
				finishedAt: Date.now(),
				...!aborted ? { error: { message } } : {}
			});
			answered.delete(input.runId);
			forget(input.runId);
		}
	};
}
//#endregion
export { createSubagentRunRecorder, storedSubagentInfo };

//# sourceMappingURL=subagent-runs.js.map