@tanstack/ai
Version:
Type-safe TypeScript AI SDK for streaming chat, tool calling, agents, structured outputs, and multimodal generation.
490 lines (489 loc) • 17.4 kB
JavaScript
import { EventType } from "../../../types.js";
import { addTokenUsage, isTanstackUsage, rebuildTokenUsage, toSpecTokenUsage } from "../../../utilities/ag-ui-usage.js";
import { tanstackMetadata, withTanstackMetadata } from "../../../utilities/merge-metadata.js";
import { INTERRUPT_BINDING_METADATA_KEY } from "../../../interrupt-resume.js";
import { EMIT_STREAM_CHUNK, SUBAGENT_TOOL } from "../tools/tool-calls.js";
//#region src/activities/chat/agents/spawn.ts
var SUBAGENT_STARTED = EventType.SUBAGENT_STARTED;
var SUBAGENT_FINISHED = EventType.SUBAGENT_FINISHED;
var SUBAGENT_ERROR = EventType.SUBAGENT_ERROR;
function createSubagentSink() {
return {
interrupts: [],
usage: []
};
}
function createSubagentId() {
return `subagent-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`;
}
function childRunId(parentRunId, subagentRunId) {
return `${parentRunId}:${subagentRunId}`;
}
/**
* Bind child interrupts to the parent run. The client resumes the parent run,
* so each binding must name that run. The resumed child then validates with
* the parent's interrupted run id.
*/
function rebindInterrupts(interrupts, runId) {
return interrupts.map((interrupt) => {
const binding = interrupt.metadata?.[INTERRUPT_BINDING_METADATA_KEY];
if (typeof binding !== "object" || binding === null) return interrupt;
return {
...interrupt,
metadata: {
...interrupt.metadata,
[INTERRUPT_BINDING_METADATA_KEY]: {
...binding,
interruptedRunId: runId,
generation: 0
}
}
};
});
}
function createAbortError() {
const error = /* @__PURE__ */ new Error("Aborted");
error.name = "AbortError";
return error;
}
function isAbortError(error, signal) {
if (signal?.aborted) return true;
return error instanceof Error && (error.name === "AbortError" || error.message === "Aborted");
}
function stoppedEvent(subagentRunId) {
return {
type: SUBAGENT_ERROR,
subagentRunId,
message: "Stopped",
timestamp: Date.now()
};
}
function childThreadId(sandbox, parentThreadId, name) {
return sandbox === "inherit" ? parentThreadId : `${parentThreadId}:${name}`;
}
function linkAbort(parent) {
const controller = new AbortController();
if (!parent) return {
controller,
dispose: () => {}
};
if (parent.aborted) {
controller.abort();
return {
controller,
dispose: () => {}
};
}
const onAbort = () => controller.abort();
parent.addEventListener("abort", onAbort, { once: true });
return {
controller,
dispose: () => parent.removeEventListener("abort", onAbort)
};
}
function orAbort(promise, signal) {
if (!signal) return promise;
if (signal.aborted) return Promise.reject(createAbortError());
return new Promise((resolve, reject) => {
const onAbort = () => reject(createAbortError());
signal.addEventListener("abort", onAbort, { once: true });
promise.then((value) => {
signal.removeEventListener("abort", onAbort);
resolve(value);
}, (error) => {
signal.removeEventListener("abort", onAbort);
reject(error);
});
});
}
function agentByName(agents, name) {
const agent = agents.find((entry) => entry.name === name);
if (!agent) throw new Error(`Unknown subagent: ${name}`);
return agent;
}
function openAgentStream(entry, bag, ctx, sink, parentToolCallId) {
const agent = agentByName(bag.agents, entry.name);
const resume = entry.resume;
const subagentRunId = resume?.subagentRunId ?? createSubagentId();
if (resume !== void 0 && ctx.interruptedRunId === void 0) throw new Error(`Subagent "${entry.name}" has interrupt answers, but the run has no parentRunId. Pass the interrupted run id as parentRunId.`);
const resumed = resume !== void 0 ? {
messages: [...ctx.messages, ...resume.messages],
parentRunId: ctx.interruptedRunId,
resume: resume.entries
} : void 0;
return spawnAgentStream(agent, {
messages: resumed?.messages ?? ctx.messages,
...ctx.abortSignal ? { abortSignal: ctx.abortSignal } : {},
threadId: childThreadId(bag.sandbox, ctx.threadId, entry.name),
runId: childRunId(ctx.parentRunId, subagentRunId),
parentRunId: resumed?.parentRunId ?? ctx.parentRunId,
subagentRunId,
...resumed ? { resume: resumed.resume } : {}
}, sink, parentToolCallId);
}
var ROUTER_PICK_ERROR = "subagents.router must return main, a name, a list of names, { names, order }, or { steps }.";
function assertOrder(order) {
if (order !== void 0 && order !== "parallel" && order !== "sequence") throw new Error("subagents.router order must be parallel or sequence.");
}
function normalizeNames(names, agents) {
if (names.length === 0) throw new Error(ROUTER_PICK_ERROR);
const hasMain = names.includes("main");
if (hasMain && names.length > 1) throw new Error("Do not mix main into a subagent list.");
if (hasMain) return ["main"];
for (const name of names) agentByName(agents, name);
return [...names];
}
function isStringList(pick) {
return Array.isArray(pick);
}
function normalizeRouterPick(pick, agents) {
if (pick === "main" || typeof pick === "string") return { steps: [{ names: normalizeNames([pick], agents) }] };
if (isStringList(pick)) return { steps: [{ names: normalizeNames(pick, agents) }] };
if ("steps" in pick) {
if (pick.steps.length === 0) throw new Error(ROUTER_PICK_ERROR);
const steps = pick.steps.map((step) => {
assertOrder(step.order);
const names = normalizeNames(step.names, agents);
return step.order === void 0 ? { names } : {
names,
order: step.order
};
});
const flat = steps.flatMap((step) => step.names);
if (flat.includes("main") && flat.length > 1) throw new Error("Do not mix main into a subagent list.");
return { steps };
}
assertOrder(pick.order);
const names = normalizeNames(pick.names, agents);
return { steps: [pick.order === void 0 ? { names } : {
names,
order: pick.order
}] };
}
/**
* Tag a child chunk with its subagent. A chunk that a nested child already
* tagged keeps its own id, and a nested child's start names this child as
* its parent.
*/
function attributeChunk(chunk, subagentRunId) {
if (chunk.type === SUBAGENT_STARTED) return chunk.parentSubagentRunId !== void 0 ? chunk : {
...chunk,
parentSubagentRunId: subagentRunId
};
if (chunk.type === SUBAGENT_FINISHED || chunk.type === SUBAGENT_ERROR) return chunk;
if ("subagentRunId" in chunk && typeof chunk.subagentRunId === "string") return chunk;
return {
...chunk,
subagentRunId
};
}
function runUsage(chunk) {
if (chunk?.type !== EventType.RUN_FINISHED) return [];
if (Array.isArray(chunk.usage)) return chunk.usage;
return isTanstackUsage(chunk.usage) ? toSpecTokenUsage(chunk.usage).usage : [];
}
/** The full usage of a run: token counts plus cost and the other fields. */
function fullUsage(chunk) {
if (chunk?.type !== EventType.RUN_FINISHED) return void 0;
return rebuildTokenUsage(chunk.usage, tanstackMetadata(chunk)?.usage);
}
/** Add a finished child run's usage to the sink. */
function collectUsage(sink, finished) {
sink.usage.push(...runUsage(finished));
const full = fullUsage(finished);
if (full) sink.total = sink.total ? addTokenUsage(sink.total, full) : full;
}
/**
* Put the children's usage on a parent terminal. `usage[]` keeps one entry per
* model call. `metadata.tanstack.usage` holds the summed cost and the other
* TanStack fields, so `fromSpecTokenUsage` reads the full total. Empties the
* sink, so the next parent terminal does not count it again.
*
* `RUN_ERROR` is accepted too: a turn that failed still spent whatever its
* children spent. Such a chunk carries no usage of its own, so `runUsage` and
* `fullUsage` return empty for it and the children's total stands alone.
*/
function withChildUsage(chunk, sink) {
if (sink.usage.length === 0 && !sink.total) return chunk;
const own = fullUsage(chunk);
const total = own && sink.total ? addTokenUsage(own, sink.total) : own ?? sink.total;
const usage = [...runUsage(chunk), ...sink.usage.splice(0)];
sink.total = void 0;
const leftover = total ? toSpecTokenUsage(total).leftover : void 0;
const next = {
...chunk,
usage
};
if (!leftover) return next;
return withTanstackMetadata(next, { usage: leftover });
}
async function* spawnAgentStream(agent, ctx, sink, parentToolCallId) {
const id = ctx.subagentRunId;
yield {
type: SUBAGENT_STARTED,
subagentRunId: id,
name: agent.name,
description: agent.description,
...ctx.parentSubagentRunId !== void 0 ? { parentSubagentRunId: ctx.parentSubagentRunId } : {},
...parentToolCallId !== void 0 ? { parentToolCallId } : {},
timestamp: Date.now()
};
let iterator;
let finished;
try {
if (ctx.abortSignal?.aborted) {
yield stoppedEvent(id);
return;
}
iterator = (await orAbort(Promise.resolve(agent.run(ctx)), ctx.abortSignal))[Symbol.asyncIterator]();
while (true) {
if (ctx.abortSignal?.aborted) {
yield stoppedEvent(id);
return;
}
const result = await orAbort(iterator.next(), ctx.abortSignal);
if (result.done) break;
const chunk = result.value;
if (chunk.type === EventType.RUN_STARTED || chunk.type === EventType.MESSAGES_SNAPSHOT) continue;
if (chunk.type === EventType.RUN_FINISHED) {
if (sink) collectUsage(sink, chunk);
finished = chunk;
continue;
}
if (chunk.type === EventType.RUN_ERROR) {
yield {
type: SUBAGENT_ERROR,
subagentRunId: id,
message: chunk.message || "Subagent failed",
...chunk.code ? { code: chunk.code } : {},
timestamp: Date.now()
};
return;
}
yield attributeChunk(chunk, id);
}
if (ctx.abortSignal?.aborted) {
yield stoppedEvent(id);
return;
}
const outcome = finished?.type === EventType.RUN_FINISHED ? finished.outcome : void 0;
if (outcome?.type === "cancelled") {
yield stoppedEvent(id);
return;
}
if (outcome?.type === "interrupt") {
const interrupts = outcome.interrupts.map((interrupt) => interrupt.subagentRunId ? interrupt : {
...interrupt,
subagentRunId: id
});
sink?.interrupts.push(...interrupts);
yield {
type: SUBAGENT_FINISHED,
subagentRunId: id,
outcome: {
type: "suspended",
interruptIds: interrupts.filter((interrupt) => interrupt.subagentRunId === id).map((interrupt) => interrupt.id)
},
timestamp: Date.now()
};
return;
}
const result = finished?.type === EventType.RUN_FINISHED ? finished.result : void 0;
yield {
type: SUBAGENT_FINISHED,
subagentRunId: id,
...result !== void 0 ? { result } : {},
timestamp: Date.now()
};
} catch (error) {
yield {
type: SUBAGENT_ERROR,
subagentRunId: id,
message: isAbortError(error, ctx.abortSignal) ? "Stopped" : error instanceof Error ? error.message : String(error),
timestamp: Date.now()
};
} finally {
try {
await iterator?.return?.();
} catch {}
}
}
async function* mergeAgentStreams(streams) {
const readers = streams.map((stream) => {
const iterator = stream[Symbol.asyncIterator]();
return {
iterator,
next: iterator.next()
};
});
try {
while (readers.length > 0) {
const indexed = readers.map((reader, index) => reader.next.then((result) => ({
index,
result,
reader
})));
const winner = await Promise.race(indexed);
if (winner.result.done) {
readers.splice(winner.index, 1);
continue;
}
yield winner.result.value;
winner.reader.next = winner.reader.iterator.next();
}
} finally {
for (const reader of readers) reader.iterator.return?.().catch(() => {});
}
}
/** True when a child in these chunks failed or stopped for outside input. */
function stopsSequence(chunks, id) {
return chunks.some((chunk) => chunk.type === SUBAGENT_ERROR && chunk.subagentRunId === id || chunk.type === SUBAGENT_FINISHED && chunk.subagentRunId === id && chunk.outcome?.type === "suspended");
}
async function* spawnNamedAgents(entries, bag, ctx, sink) {
if (bag.sandbox === "inherit" && entries.length > 1) throw new Error("subagents.sandbox 'inherit' cannot start two children in one turn");
const group = linkAbort(ctx.abortSignal);
const groupCtx = {
...ctx,
abortSignal: group.controller.signal
};
try {
if (bag.order === "sequence") {
let messages = ctx.messages;
for (const entry of entries) {
const chunks = [];
let id;
for await (const chunk of openAgentStream(entry, bag, {
...groupCtx,
messages
}, sink)) {
if (chunk.type === SUBAGENT_STARTED && id === void 0) id = chunk.subagentRunId;
chunks.push(chunk);
yield chunk;
}
if (id !== void 0 && stopsSequence(chunks, id)) return;
const text = [entry.resume?.text, collectNamedText(chunks, [entry.name])].filter((part) => part !== void 0 && part !== "").join("");
if (text) messages = [...messages, {
role: "assistant",
content: text
}];
}
return;
}
const streams = entries.map((entry) => openAgentStream(entry, bag, groupCtx, sink));
const onlyStream = streams.length === 1 ? streams[0] : void 0;
if (onlyStream) {
yield* onlyStream;
return;
}
yield* mergeAgentStreams(streams);
} finally {
group.controller.abort();
group.dispose();
}
}
/**
* Text of the named direct children, in `names` order. Text from nested
* children stays out: their chunks carry their own id.
*/
function collectNamedText(chunks, names) {
const nameByRunId = /* @__PURE__ */ new Map();
const textByName = /* @__PURE__ */ new Map();
for (const chunk of chunks) {
if (chunk.type === SUBAGENT_STARTED) {
if (chunk.parentSubagentRunId === void 0) nameByRunId.set(chunk.subagentRunId, chunk.name);
continue;
}
if (chunk.type !== EventType.TEXT_MESSAGE_CONTENT) continue;
if (!("subagentRunId" in chunk) || typeof chunk.subagentRunId !== "string") continue;
const name = nameByRunId.get(chunk.subagentRunId);
if (!name) continue;
textByName.set(name, `${textByName.get(name) ?? ""}${chunk.delta}`);
}
return names.map((name) => textByName.get(name)?.trim() ?? "").filter((text) => text.length > 0).join("\n\n");
}
/**
* The parent conversation up to the message that carries this tool call, with
* that message's tool calls removed. Its string text stays; array content is
* dropped.
*/
function messagesBeforeCall(messages, toolCallId) {
const index = messages.findIndex((message) => message.toolCalls?.some((call) => call.id === toolCallId));
if (index === -1) return [...messages];
const host = messages[index];
const kept = messages.slice(0, index);
if (host && typeof host.content === "string" && host.content !== "") {
const { toolCalls: _calls, ...text } = host;
kept.push(text);
}
return kept;
}
/**
* Record the parent messages when the model calls a subagent tool, so the
* child reads the conversation as it is at that call.
*/
function subagentCallMessages(names) {
const byCall = /* @__PURE__ */ new Map();
return {
middleware: {
name: "subagent-call-messages",
onBeforeToolCall(ctx, hook) {
if (!names.has(hook.toolName)) return void 0;
byCall.set(hook.toolCallId, messagesBeforeCall(ctx.messages, hook.toolCallId));
}
},
messagesFor: (toolCallId) => toolCallId === void 0 ? void 0 : byCall.get(toolCallId)
};
}
function createSyntheticSubagentTools(bag, parent) {
return bag.agents.map((agent) => ({
name: agent.name,
description: agent.description,
[SUBAGENT_TOOL]: true,
execute: async (_input, context) => {
const toolContext = context;
const toolCallId = toolContext?.toolCallId;
const suspended = parent.turn?.children.find((child) => child.status === "suspended" && child.parentToolCallId !== void 0 && child.parentToolCallId === toolCallId);
const entry = suspended ? {
name: agent.name,
resume: {
subagentRunId: suspended.subagentRunId,
messages: suspended.messages,
entries: suspended.resume,
text: suspended.text
}
} : { name: agent.name };
const sink = createSubagentSink();
const link = linkAbort(parent.abortSignal);
let subagentRunId = suspended?.subagentRunId ?? "";
let text = suspended?.text ?? "";
let error;
try {
for await (const chunk of openAgentStream(entry, bag, {
messages: parent.messagesFor?.(toolCallId) ?? parent.messages,
abortSignal: link.controller.signal,
threadId: parent.threadId,
parentRunId: parent.runId,
...parent.interruptedRunId !== void 0 ? { interruptedRunId: parent.interruptedRunId } : {}
}, sink, toolCallId)) {
if (chunk.type === SUBAGENT_STARTED && subagentRunId === "") subagentRunId = chunk.subagentRunId;
if (chunk.type === EventType.TEXT_MESSAGE_CONTENT && "subagentRunId" in chunk && chunk.subagentRunId === subagentRunId) text += chunk.delta;
if (chunk.type === SUBAGENT_ERROR && chunk.subagentRunId === subagentRunId) error = chunk.message;
toolContext?.[EMIT_STREAM_CHUNK]?.(chunk);
}
} finally {
link.dispose();
}
parent.sink.usage.push(...sink.usage);
if (sink.total) parent.sink.total = parent.sink.total ? addTokenUsage(parent.sink.total, sink.total) : sink.total;
return {
subagentRunId,
text,
...error !== void 0 ? { error } : {},
...sink.interrupts.length > 0 ? { interrupts: sink.interrupts } : {}
};
}
}));
}
//#endregion
export { SUBAGENT_ERROR, SUBAGENT_FINISHED, SUBAGENT_STARTED, collectNamedText, collectUsage, createSubagentId, createSubagentSink, createSyntheticSubagentTools, normalizeRouterPick, rebindInterrupts, spawnAgentStream, spawnNamedAgents, subagentCallMessages, withChildUsage };
//# sourceMappingURL=spawn.js.map