UNPKG

@mastra/core

Version:
1,299 lines 83.8 kB
const require_error = require("./error-B-e62x-A.cjs"); const require_event_emitter = require("./event-emitter-80LrFIBS.cjs"); require("./request-context-ByoZMp-j.cjs"); const require_utils = require("./utils-Bw7FoAI3.cjs"); const require_signals = require("./signals-D2CulJo3.cjs"); const require_types = require("./types-ftn4_F2h.cjs"); const require_base = require("./base-CfojwyC3.cjs"); let crypto = require("crypto"); //#region src/agent/thread-stream-runtime.ts const AGENT_THREAD_KEY_SEPARATOR = "\0"; const AGENT_THREAD_STREAM_TOPIC_PREFIX = "agent.thread-stream"; /** * Lease TTL for the cross-process thread lease acquired in the idle-wake * path. Kept short so a crashed owner process frees the thread quickly; a * background timer renews it while the run is still running. Overridable via * `MASTRA_AGENT_THREAD_LEASE_TTL_MS` (production keeps the 15s default). */ const AGENT_THREAD_LEASE_TTL_MS = require_utils.readPositiveIntEnv("MASTRA_AGENT_THREAD_LEASE_TTL_MS", 15e3); /** * Interval at which the owner process renews its lease. Defaults to TTL/3, * leaving room for two missed renewals (network blip, GC pause) before the * lease expires. Overridable via `MASTRA_AGENT_THREAD_LEASE_RENEW_INTERVAL_MS`. */ const AGENT_THREAD_LEASE_RENEW_INTERVAL_MS = require_utils.readPositiveIntEnv("MASTRA_AGENT_THREAD_LEASE_RENEW_INTERVAL_MS", Math.floor(AGENT_THREAD_LEASE_TTL_MS / 3)); /** * TTL for a suspended run's warm in-memory state — the parked thread-run record * (swept by #sweepStaleSuspendedRecords). The Mastra internal-workflow registry * reads the same `MASTRA_SUSPENDED_RUN_TTL_MS` so both expire on one bound. A * suspended run is kept warm so a same-instance resume can reattach and the thread * stays blocked; once it lapses the state is evicted and resume falls back to the * durable snapshot. Multi-instance deployments (resume rarely lands on the origin) * can shed it sooner; 30 minute default. */ const AGENT_SUSPENDED_RUN_TTL_MS = require_utils.readPositiveIntEnv("MASTRA_SUSPENDED_RUN_TTL_MS", 1800 * 1e3); let defaultAgentThreadPubSub = new require_event_emitter.EventEmitterPubSub(); function withThreadMemory(memory, resourceId, threadId) { return { ...memory && typeof memory === "object" ? memory : {}, resource: memory?.resource ?? resourceId, thread: memory?.thread ?? threadId }; } function createRuntimeState() { return { threadRunsById: /* @__PURE__ */ new Map(), threadRunsByStreamId: /* @__PURE__ */ new Map(), threadKeysByRunId: /* @__PURE__ */ new Map(), remoteThreadKeysByRunId: /* @__PURE__ */ new Map(), activeThreadRunIds: /* @__PURE__ */ new Map(), activeThreadStreamIds: /* @__PURE__ */ new Map(), streamSeqByRunId: /* @__PURE__ */ new Map(), approvalSuspendedRunIds: /* @__PURE__ */ new Set(), suspendedRunIds: /* @__PURE__ */ new Set(), suspensionMetadataByRunId: /* @__PURE__ */ new Map(), pendingSignalsByThread: /* @__PURE__ */ new Map(), preRunSignalsByThread: /* @__PURE__ */ new Map(), pendingIdleSignalsByThread: /* @__PURE__ */ new Map(), pendingContinuationsByThread: /* @__PURE__ */ new Map(), watchedThreadStreamIds: /* @__PURE__ */ new Set(), preparedRunsById: /* @__PURE__ */ new Map(), abortedRunIds: /* @__PURE__ */ new Set(), leaseRenewalTimers: /* @__PURE__ */ new Map() }; } var AgentThreadStreamRuntime = class { #id; #statesByPubSub = /* @__PURE__ */ new WeakMap(); #getPubSub(pubsub) { return pubsub ?? defaultAgentThreadPubSub; } /** * Resolve the {@link LeaseProvider} for the configured pubsub. Leasing is * a separate capability from event delivery: a backend only implements it * when it can genuinely coordinate a distributed lock (Redis via SET-NX, * in-memory for single-process). We feature-detect once here so all lease * call sites can use the resolved provider unconditionally. * * `CachingPubSub` exposes its inner's lease provider via `getLeaseProvider` * (caching is transparent to leasing). Otherwise we duck-type the pubsub * directly. Backends that cannot lease fall back to {@link NoopLeaseProvider} * (always-win / no-op), preserving single-process behavior. */ #getLeaseProvider(pubsub) { const resolved = this.#getPubSub(pubsub); const unwrap = resolved.getLeaseProvider; if (typeof unwrap === "function") return unwrap.call(resolved) ?? require_event_emitter.NoopLeaseProvider; return require_event_emitter.isLeaseProvider(resolved) ? resolved : require_event_emitter.NoopLeaseProvider; } #resolveLeaseProvider(pubsub) { const provider = this.#getLeaseProvider(pubsub); return { provider, isFallback: provider === require_event_emitter.NoopLeaseProvider }; } async #hasLiveThreadLease(pubsub, key, runId) { const { provider, isFallback } = this.#resolveLeaseProvider(pubsub); if (isFallback) return true; return provider.getLeaseOwner(key).then((owner) => owner === runId).catch(() => false); } #getSourceId() { this.#id ??= (0, crypto.randomUUID)(); return this.#id; } /** * Fire-and-forget release of the cross-process thread lease held by * this owner. Safe to call when no lease was ever acquired — the * pubsub's `releaseLease` is a no-op for non-owners (Lua-guarded * GET+DEL on Redis), and the default in-memory implementation is * identical. Also stops the renewal timer if one is running for * this run. */ #releaseThreadLease(pubsub, key, runId) { const resolved = this.#getPubSub(pubsub); this.#stopLeaseRenewal(resolved, runId); this.#getLeaseProvider(resolved).releaseLease(key, runId).catch(() => {}); } /** * Start a background timer that renews the cross-process lease at * TTL/3 intervals while the run is still going. If the lease is lost * (e.g. expired due to clock skew or pubsub outage) the renewal * stops itself — there's nothing useful we can do from the runner * side beyond log; the original owner will keep running until the run * itself errors or completes. */ #startLeaseRenewal(pubsub, key, runId) { const state = this.#getState(pubsub); if (state.leaseRenewalTimers.has(runId)) return; const leaseProvider = this.#getLeaseProvider(pubsub); const timer = setInterval(() => { leaseProvider.renewLease(key, runId, AGENT_THREAD_LEASE_TTL_MS).then((renewed) => { if (!renewed) this.#stopLeaseRenewal(pubsub, runId); }).catch(() => {}); }, AGENT_THREAD_LEASE_RENEW_INTERVAL_MS); if (typeof timer === "object" && timer && typeof timer.unref === "function") timer.unref(); state.leaseRenewalTimers.set(runId, timer); } #stopLeaseRenewal(pubsub, runId) { const state = this.#getState(pubsub); const timer = state.leaseRenewalTimers.get(runId); if (!timer) return; clearInterval(timer); state.leaseRenewalTimers.delete(runId); } /** * Hand the cross-process thread lease from a finishing run (`fromRunId`) * to the run that will drain queued follow-up work next (`toRunId`), * without the lease key ever going empty. * * The previous owner releases its renewal timer and the new owner starts * its own; the lease key is re-stamped by `transferLease` (with a full fresh * TTL). On atomic backends (Redis, in-memory) a racing process cannot win a * freed key between a release and a re-acquire. Backends that can't transfer * atomically implement `transferLease` as release+acquire internally and own * that race cost. Returns `true` if the new owner now holds the lease. */ async #transferThreadLease(pubsub, key, fromRunId, toRunId) { const resolved = this.#getPubSub(pubsub); const held = await this.#getLeaseProvider(resolved).transferLease(key, fromRunId, toRunId, AGENT_THREAD_LEASE_TTL_MS).catch(() => false); this.#stopLeaseRenewal(resolved, fromRunId); if (held) this.#startLeaseRenewal(resolved, key, toRunId); return held; } /** * Ensure this process owns the cross-process lease for `toRunId` before it * starts a run, regardless of whether it already held the lease. * * - When `fromRunId` is provided (draining after a run this process owned), * atomically transfer the held lease to `toRunId` — gap-free, no empty key. * - When `fromRunId` is absent, or the transfer reports the old owner no * longer holds the lease, fall back to a fresh `acquireLease`. This covers * a *different* process that observed the owner finish via pub/sub and now * wants to wake the thread: it never held the lease, so it must win one. * * On success the renewal timer is started for `toRunId`. On failure the * returned `owner` is the current holder so the caller can forward work to it. */ async #acquireOrTransferThreadLease(pubsub, key, toRunId, fromRunId) { const resolved = this.#getPubSub(pubsub); if (fromRunId) { if (await this.#transferThreadLease(pubsub, key, fromRunId, toRunId)) return { acquired: true, owner: toRunId }; } const result = await this.#getLeaseProvider(resolved).acquireLease(key, toRunId, AGENT_THREAD_LEASE_TTL_MS).catch(() => ({ acquired: false, owner: void 0 })); if (result.acquired) { this.#startLeaseRenewal(resolved, key, toRunId); return { acquired: true, owner: toRunId }; } return { acquired: false, owner: result.owner }; } /** * Whether the thread has any queued follow-up work that a finishing run's * completion handler would drain next: pending follow-up signals (including * any pre-run leftover that will be folded in), queued continuations, or * queued idle signals. */ #hasPendingThreadWork(state, key) { return (state.pendingSignalsByThread.get(key)?.length ?? 0) > 0 || (state.preRunSignalsByThread.get(key)?.length ?? 0) > 0 || (state.pendingContinuationsByThread.get(key)?.length ?? 0) > 0 || (state.pendingIdleSignalsByThread.get(key)?.length ?? 0) > 0; } #getState(pubsub) { const resolvedPubSub = this.#getPubSub(pubsub); let state = this.#statesByPubSub.get(resolvedPubSub); if (!state) { state = createRuntimeState(); this.#statesByPubSub.set(resolvedPubSub, state); } return state; } #threadKey(resourceId, threadId) { return [resourceId ?? "", threadId].join(AGENT_THREAD_KEY_SEPARATOR); } #threadTopic(key) { return `${AGENT_THREAD_STREAM_TOPIC_PREFIX}.${encodeURIComponent(key)}`; } #isApprovalSuspendedRun(state, runId) { return state.approvalSuspendedRunIds.has(runId); } #isSuspendedRun(state, runId) { return state.suspendedRunIds.has(runId) || this.#isApprovalSuspendedRun(state, runId); } #isThreadBlockingRun(state, record) { return record.output.status === "running" || record.output.status === "suspended" || record.lifecycle === "suspending" || record.lifecycle === "suspended" || !!record.suspension || this.#isSuspendedRun(state, record.runId); } #serializeSignal(signal) { return signal; } #nextStreamIdentity(state, runId) { const streamSeq = (state.streamSeqByRunId.get(runId) ?? 0) + 1; state.streamSeqByRunId.set(runId, streamSeq); return { streamId: (0, crypto.randomUUID)(), streamSeq }; } #markRunSuspending(state, runId, streamId, suspension) { state.suspendedRunIds.add(runId); state.suspensionMetadataByRunId.set(runId, suspension); const record = state.threadRunsByStreamId.get(streamId) ?? state.threadRunsById.get(runId); if (record) { record.lifecycle = "suspending"; record.suspension = suspension; } if (suspension.kind === "approval") state.approvalSuspendedRunIds.add(runId); } #clearSuspendedRun(state, runId) { state.suspendedRunIds.delete(runId); state.suspensionMetadataByRunId.delete(runId); state.approvalSuspendedRunIds.delete(runId); } #generateSignalMessageId(agent, target) { return agent.getMastraInstance?.()?.generateId({ idType: "message", source: "agent", entityId: agent.id, threadId: target.threadId, resourceId: target.resourceId }) ?? (0, crypto.randomUUID)(); } #createMessageSignalInput(message) { return { ...typeof message === "string" || Array.isArray(message) ? { contents: message } : message, type: "user", tagName: "user" }; } getThreadState(options, pubsub) { const state = this.#getState(pubsub); const key = this.#threadKey(options.resourceId, options.threadId); const activeRunId = state.activeThreadRunIds.get(key); if (!activeRunId) return "idle"; const activeRecord = state.threadRunsById.get(activeRunId); if (activeRecord && !this.#isThreadBlockingRun(state, activeRecord)) { state.activeThreadRunIds.delete(key); return "idle"; } return "active"; } #publish(pubsub, key, event) { this.#publishAndWait(pubsub, key, event).catch(() => {}); } async #publishAndWait(pubsub, key, event) { await this.#getPubSub(pubsub).publish(this.#threadTopic(key), { type: event.type, runId: event.runId, data: event }); } #withBroadcastStream(output, pubsub, key, streamId) { const runtime = this; const parts = []; const waiters = /* @__PURE__ */ new Set(); let started = false; let done = false; let error; const wake = () => { const pending = [...waiters]; waiters.clear(); for (const waiter of pending) waiter(); }; const emitPart = async (part) => { if (part && typeof part === "object" && "type" in part) { const typedPart = part; if (typedPart.type === "tool-call-approval" || typedPart.type === "tool-call-suspended") runtime.#markRunSuspending(runtime.#getState(pubsub), output.runId, streamId, { toolCallId: typedPart.payload?.toolCallId, toolName: typedPart.payload?.toolName, kind: typedPart.type === "tool-call-approval" ? "approval" : "generic-tool" }); } parts.push(part); await runtime.#publishAndWait(pubsub, key, { type: "stream-part", runId: output.runId, streamId, part, sourceId: runtime.#getSourceId() }); wake(); }; const start = () => { if (started) return; started = true; (async () => { try { const source = output.fullStream; if (!source) return; if (typeof source.getReader === "function") { const reader = source.getReader(); try { while (true) { const { value: part, done: streamDone } = await reader.read(); if (streamDone) break; await emitPart(part); } } finally { reader.releaseLock(); } } else for await (const part of source) await emitPart(part); } catch (caught) { error = caught; } finally { done = true; wake(); } })(); }; const createStream = () => { let index = 0; let closed = false; let waiter; return new ReadableStream({ async pull(controller) { start(); while (!closed) { if (index < parts.length) { controller.enqueue(parts[index++]); return; } if (error) { controller.error(error); return; } if (done) { controller.close(); return; } await new Promise((resolve) => { waiter = resolve; waiters.add(resolve); }); if (waiter) { waiters.delete(waiter); waiter = void 0; } } }, cancel() { closed = true; if (waiter) { waiters.delete(waiter); waiter(); waiter = void 0; } } }); }; return { output, createSubscriberStream: createStream, startBroadcast: start }; } #getThreadTarget(options) { const thread = options?.memory?.thread; return { threadId: options?.requestContext?.get("mastra__threadId") || (typeof thread === "string" ? thread : thread?.id), resourceId: options?.requestContext?.get("mastra__resourceId") || options?.memory?.resource }; } prepareRunOptions(options, pubsub) { const { threadId } = this.#getThreadTarget(options); if (!threadId || !options.runId) return options; const state = this.#getState(pubsub); const abortController = new AbortController(); const upstreamAbortSignal = options.abortSignal; const abort = () => abortController.abort(); if (upstreamAbortSignal?.aborted) abort(); else upstreamAbortSignal?.addEventListener("abort", abort, { once: true }); state.preparedRunsById.set(options.runId, { abortController, cleanup: () => upstreamAbortSignal?.removeEventListener("abort", abort) }); if (state.abortedRunIds.has(options.runId)) abort(); return { ...options, abortSignal: abortController.signal }; } abortRun(runId, pubsub) { const state = this.#getState(pubsub); const preparedRun = state.preparedRunsById.get(runId); if (!preparedRun) { state.abortedRunIds.add(runId); return false; } preparedRun.abortController.abort(); state.abortedRunIds.add(runId); const key = state.threadKeysByRunId.get(runId); if (key) { const streamId = state.activeThreadRunIds.get(key) === runId ? state.activeThreadStreamIds.get(key) : void 0; this.#releaseThreadLease(pubsub, key, runId); this.#publish(pubsub, key, { type: "run-aborted", runId, streamId }); } return true; } getActiveThreadRunId(options, pubsub) { const state = this.#getState(pubsub); const key = this.#threadKey(options.resourceId, options.threadId); const activeRunId = state.activeThreadRunIds.get(key); if (!activeRunId) return void 0; const record = state.threadRunsById.get(activeRunId); if (record && !this.#isThreadBlockingRun(state, record)) return void 0; return activeRunId; } getResumableThreadRun(options, pubsub) { const state = this.#getState(pubsub); const key = this.#threadKey(options.resourceId, options.threadId); const record = state.threadRunsById.get(options.runId); const isSuspended = this.#isSuspendedRun(state, options.runId); if (!record || state.threadKeysByRunId.get(options.runId) !== key || !isSuspended) return; const suspension = record.suspension ?? state.suspensionMetadataByRunId.get(options.runId); if (options.toolCallId && suspension?.toolCallId && suspension.toolCallId !== options.toolCallId) return; return { runId: options.runId, toolCallId: options.toolCallId ?? suspension?.toolCallId }; } abortThread(options, pubsub) { const resolvedPubSub = this.#getPubSub(pubsub); const state = this.#getState(resolvedPubSub); const key = this.#threadKey(options.resourceId, options.threadId); const runId = this.getActiveThreadRunId(options, resolvedPubSub); if (!runId) return false; if (state.preparedRunsById.has(runId)) return this.abortRun(runId, resolvedPubSub); if (state.threadKeysByRunId.get(runId) === key) { this.abortRun(runId, resolvedPubSub); return true; } if (state.remoteThreadKeysByRunId.get(runId) !== key) return false; const streamId = state.activeThreadStreamIds.get(key); if (!streamId) return false; this.#publish(resolvedPubSub, key, { type: "run-abort-requested", runId, streamId }); return true; } /** @internal */ resetForTests() { for (const pubsub of [defaultAgentThreadPubSub]) { this.#resetState(pubsub); pubsub.close?.(); } defaultAgentThreadPubSub = new require_event_emitter.EventEmitterPubSub(); } #resetState(pubsub) { const state = this.#statesByPubSub.get(pubsub); if (!state) return; state.preparedRunsById.forEach((preparedRun) => { preparedRun.abortController.abort(); preparedRun.cleanup(); }); state.leaseRenewalTimers.forEach((timer) => clearInterval(timer)); state.leaseRenewalTimers.clear(); state.threadRunsById.clear(); state.threadRunsByStreamId.clear(); state.threadKeysByRunId.clear(); state.remoteThreadKeysByRunId.clear(); state.activeThreadRunIds.clear(); state.approvalSuspendedRunIds.clear(); state.suspendedRunIds.clear(); state.suspensionMetadataByRunId.clear(); state.pendingSignalsByThread.clear(); state.preRunSignalsByThread.clear(); state.pendingIdleSignalsByThread.clear(); state.pendingContinuationsByThread.clear(); state.activeThreadStreamIds.clear(); state.streamSeqByRunId.clear(); state.watchedThreadStreamIds.clear(); state.preparedRunsById.clear(); state.abortedRunIds.clear(); } #cleanupPreparedRun(state, runId) { state.preparedRunsById.get(runId)?.cleanup(); state.preparedRunsById.delete(runId); state.abortedRunIds.delete(runId); } async #persistSignal(agent, signal, resourceId, threadId, requestContext) { if (signal.transient) return; const memory = await agent.getMemory({ requestContext }); if (!memory) return; await memory.saveMessages({ messages: [signal.toDBMessage({ resourceId, threadId })] }); } #broadcastPersistedSignal(state, pubsub, key, runId, signal, resourceId, threadId) { let finish; const finished = new Promise((resolve) => { finish = resolve; }); const parts = [ { type: "start", runId }, { ...signal.toDataPart(), runId }, { type: "finish", runId, payload: { stepResult: { reason: "stop" }, output: { usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 } } } } ]; const output = { runId, status: "running", fullStream: new ReadableStream({ start(controller) { for (const part of parts) controller.enqueue(part); controller.close(); finish(); } }), _waitUntilFinished: () => finished }; const { streamId, streamSeq } = this.#nextStreamIdentity(state, runId); const { output: outputForSubscribers, createSubscriberStream, startBroadcast } = this.#withBroadcastStream(output, pubsub, key, streamId); const record = { agent: { id: `persisted-signal:${signal.id}` }, output: outputForSubscribers, runId, streamId, streamSeq, lifecycle: "running", threadId, resourceId, streamOptions: {}, createSubscriberStream }; state.threadRunsById.set(runId, record); state.threadRunsByStreamId.set(streamId, record); state.threadKeysByRunId.set(runId, key); state.activeThreadStreamIds.set(key, streamId); this.#publishAndWait(pubsub, key, { type: "run-registered", runId, streamId, streamSeq }).then(startBroadcast, startBroadcast); outputForSubscribers._waitUntilFinished().finally(() => { setTimeout(() => { state.threadRunsByStreamId.delete(streamId); if (state.threadRunsById.get(runId) === record) { state.threadRunsById.delete(runId); state.threadKeysByRunId.delete(runId); } if (state.activeThreadRunIds.get(key) === runId && state.activeThreadStreamIds.get(key) === streamId) { state.activeThreadRunIds.delete(key); state.activeThreadStreamIds.delete(key); } this.#releaseThreadLease(pubsub, key, runId); this.#publish(pubsub, key, { type: "run-completed", runId, streamId }); }, 0); }); } async #persistAndBroadcastIdleSignal(state, pubsub, key, runId, agent, signal, resourceId, threadId, requestContext) { if (signal.transient) return; await this.#persistSignal(agent, signal, resourceId, threadId, requestContext); this.#broadcastPersistedSignal(state, pubsub, key, runId, signal, resourceId, threadId); } /** * Evict SUSPENDED records parked longer than {@link AGENT_SUSPENDED_RUN_TTL_MS}. * Called lazily on each registration so cleanup is proportional to activity and * zero-cost when idle — mirrors the internal-workflow registry sweep. Bounds the * records left behind by abandoned suspends and by resumes that land on a * different instance (which never clean the origin instance's record). * * When the expiring record is still the run's current record — an abandoned * suspend, not one superseded by a same-instance resume — the teardown mirrors * #watchThreadRunCompletion's terminal path: it clears run-level state, releases * the cross-process lease, and publishes `run-completed` so remote subscribers * stop treating the thread as blocked and drain any queued follow-up work. A * superseded older stream just has its stream entry dropped; the resumed run * keeps its lease, suspended marker, and active slot. */ #sweepStaleSuspendedRecords(state, pubsub) { const now = Date.now(); for (const [streamId, record] of state.threadRunsByStreamId) { if (record.lifecycle !== "suspended" || record.suspendedAt === void 0) continue; if (now - record.suspendedAt <= AGENT_SUSPENDED_RUN_TTL_MS) continue; state.threadRunsByStreamId.delete(streamId); state.watchedThreadStreamIds.delete(streamId); if (state.threadRunsById.get(record.runId) !== record) continue; const staleKey = this.#threadKey(record.resourceId, record.threadId); state.threadRunsById.delete(record.runId); state.threadKeysByRunId.delete(record.runId); this.#clearSuspendedRun(state, record.runId); this.#releaseThreadLease(pubsub, staleKey, record.runId); if (state.activeThreadRunIds.get(staleKey) === record.runId && state.activeThreadStreamIds.get(staleKey) === streamId) { state.activeThreadRunIds.delete(staleKey); state.activeThreadStreamIds.delete(staleKey); } this.#publish(pubsub, staleKey, { type: "run-completed", runId: record.runId, streamId }); } } registerRun(agent, output, streamOptions, pubsub) { const { threadId, resourceId } = this.#getThreadTarget(streamOptions); if (!threadId) return; const state = this.#getState(pubsub); this.#sweepStaleSuspendedRecords(state, pubsub); const key = this.#threadKey(resourceId, threadId); const { streamId, streamSeq } = this.#nextStreamIdentity(state, output.runId); const { output: outputForSubscribers, createSubscriberStream, startBroadcast } = this.#withBroadcastStream(output, pubsub, key, streamId); const record = { agent, output: outputForSubscribers, runId: output.runId, streamId, streamSeq, lifecycle: "running", threadId, resourceId, streamOptions, createSubscriberStream }; this.#clearSuspendedRun(state, output.runId); state.threadRunsById.set(output.runId, record); state.threadRunsByStreamId.set(streamId, record); state.threadKeysByRunId.set(output.runId, key); state.activeThreadRunIds.set(key, output.runId); state.activeThreadStreamIds.set(key, streamId); const resolvedPubSub = this.#getPubSub(pubsub); const registered = (async () => { if ((await this.#getLeaseProvider(resolvedPubSub).acquireLease(key, output.runId, AGENT_THREAD_LEASE_TTL_MS).catch(() => ({ acquired: true }))).acquired) this.#startLeaseRenewal(resolvedPubSub, key, output.runId); await this.#publishAndWait(pubsub, key, { type: "run-registered", runId: output.runId, streamId, streamSeq }); })(); registered.then(startBroadcast, startBroadcast); this.#watchThreadRunCompletion(state, pubsub, key, record); return registered; } #watchThreadRunCompletion(state, pubsub, key, record) { if (state.watchedThreadStreamIds.has(record.streamId)) return; state.watchedThreadStreamIds.add(record.streamId); record.output._waitUntilFinished().finally(() => { state.watchedThreadStreamIds.delete(record.streamId); this.#cleanupPreparedRun(state, record.runId); if (record.output.status === "suspended" && this.#isSuspendedRun(state, record.runId)) { record.lifecycle = "suspended"; record.suspendedAt = Date.now(); this.#publish(pubsub, key, { type: "run-suspended", runId: record.runId, streamId: record.streamId }); return; } record.lifecycle = "completed"; this.#clearSuspendedRun(state, record.runId); state.threadRunsByStreamId.delete(record.streamId); if (state.threadRunsById.get(record.runId) === record) { state.threadRunsById.delete(record.runId); state.threadKeysByRunId.delete(record.runId); } if (state.activeThreadRunIds.get(key) === record.runId && state.activeThreadStreamIds.get(key) === record.streamId) { state.activeThreadRunIds.delete(key); state.activeThreadStreamIds.delete(key); } this.#publish(pubsub, key, { type: "run-completed", runId: record.runId, streamId: record.streamId }); if (this.#hasPendingThreadWork(state, key)) this.#drainPendingSignals(state, pubsub, key, record); else this.#releaseThreadLease(pubsub, key, record.runId); }); } async #drainPendingSignals(state, pubsub, key, previousRun) { if (state.activeThreadRunIds.has(key)) return; const preRunLeftover = state.preRunSignalsByThread.get(key); if (preRunLeftover?.length) { state.preRunSignalsByThread.delete(key); state.pendingSignalsByThread.set(key, [...preRunLeftover, ...state.pendingSignalsByThread.get(key) ?? []]); } const queue = state.pendingSignalsByThread.get(key); const signal = queue?.shift(); if (signal && queue) { if (queue.length === 0) state.pendingSignalsByThread.delete(key); const nextRunId = (0, crypto.randomUUID)(); state.activeThreadRunIds.set(key, nextRunId); state.threadKeysByRunId.set(nextRunId, key); const owns = await this.#acquireOrTransferThreadLease(pubsub, key, nextRunId, previousRun.runId); if (!owns.acquired) { if (state.activeThreadRunIds.get(key) === nextRunId) state.activeThreadRunIds.delete(key); state.threadKeysByRunId.delete(nextRunId); state.preRunSignalsByThread.delete(key); const restored = state.pendingSignalsByThread.get(key) ?? []; state.pendingSignalsByThread.set(key, [signal, ...restored]); if (owns.owner) { await this.#publishAndWait(pubsub, key, { type: "signal-enqueued", runId: owns.owner, signal: this.#serializeSignal(signal), sourceId: this.#getSourceId() }).catch(() => {}); state.pendingSignalsByThread.get(key)?.shift(); if ((state.pendingSignalsByThread.get(key)?.length ?? 0) === 0) state.pendingSignalsByThread.delete(key); } return; } const output = await previousRun.agent.stream(signal, { ...previousRun.streamOptions, runId: nextRunId, memory: withThreadMemory(previousRun.streamOptions.memory, previousRun.resourceId ?? "", previousRun.threadId ?? "") }); if (queue.length > 0) { const nextRecord = state.threadRunsById.get(output.runId); if (nextRecord) this.#watchThreadRunCompletion(state, pubsub, key, nextRecord); } return; } if (await this.#drainPendingContinuations(state, pubsub, key, previousRun.runId)) return; if (await this.#drainPendingIdleSignals(state, pubsub, key, previousRun.runId)) return; this.#releaseThreadLease(pubsub, key, previousRun.runId); } async #drainPendingContinuations(state, pubsub, key, fromRunId) { if (state.activeThreadRunIds.has(key)) return false; const queue = state.pendingContinuationsByThread.get(key); const pending = queue?.shift(); if (!pending || !queue) return false; if (queue.length === 0) state.pendingContinuationsByThread.delete(key); if (fromRunId) { state.activeThreadRunIds.set(key, pending.runId); state.threadKeysByRunId.set(pending.runId, key); if (!(await this.#acquireOrTransferThreadLease(pubsub, key, pending.runId, fromRunId)).acquired) { if (state.activeThreadRunIds.get(key) === pending.runId) state.activeThreadRunIds.delete(key); state.threadKeysByRunId.delete(pending.runId); state.preRunSignalsByThread.delete(key); const restored = state.pendingContinuationsByThread.get(key) ?? []; state.pendingContinuationsByThread.set(key, [pending, ...restored]); return false; } } this.#startContinuation(state, pubsub, key, pending); return true; } #startContinuation(state, pubsub, key, pending) { state.activeThreadRunIds.set(key, pending.runId); state.threadKeysByRunId.set(pending.runId, key); pending.agent.stream(pending.messages, { ...pending.streamOptions, runId: pending.runId, memory: withThreadMemory(pending.streamOptions?.memory, pending.resourceId, pending.threadId) }).then((output) => { if ((state.pendingContinuationsByThread.get(key)?.length ?? 0) > 0) { const nextRecord = state.threadRunsById.get(output.runId); if (nextRecord) this.#watchThreadRunCompletion(state, pubsub, key, nextRecord); } }).catch((err) => { state.threadKeysByRunId.delete(pending.runId); this.#cleanupPreparedRun(state, pending.runId); if (state.activeThreadRunIds.get(key) === pending.runId) state.activeThreadRunIds.delete(key); this.#publish(pubsub, key, { type: "run-failed", runId: pending.runId, error: require_error.getErrorFromUnknown(err).message }); this.#drainPendingContinuations(state, pubsub, key, pending.runId).then(async (started) => { if (started) return; if (await this.#drainPendingIdleSignals(state, pubsub, key, pending.runId)) return; this.#releaseThreadLease(pubsub, key, pending.runId); }); }); } continueWithMessages(agent, messages, target, pubsub) { const state = this.#getState(pubsub); const key = this.#threadKey(target.resourceId, target.threadId); const runId = target.runId ?? (0, crypto.randomUUID)(); const pending = { agent, messages, runId, resourceId: target.resourceId, threadId: target.threadId, streamOptions: target.streamOptions }; const activeRunId = state.activeThreadRunIds.get(key); const activeRecord = activeRunId ? state.threadRunsById.get(activeRunId) : void 0; if (state.activeThreadRunIds.has(key)) { const queue = state.pendingContinuationsByThread.get(key) ?? []; queue.push(pending); state.pendingContinuationsByThread.set(key, queue); if (activeRecord) this.#watchThreadRunCompletion(state, pubsub, key, activeRecord); return { accepted: true, runId }; } this.#startContinuation(state, pubsub, key, pending); return { accepted: true, runId }; } async #drainPendingIdleSignals(state, pubsub, key, fromRunId) { if (state.activeThreadRunIds.has(key)) return false; const idleQueue = state.pendingIdleSignalsByThread.get(key); const pendingIdle = idleQueue?.shift(); if (!pendingIdle || !idleQueue) return false; if (idleQueue.length === 0) state.pendingIdleSignalsByThread.delete(key); state.activeThreadRunIds.set(key, pendingIdle.runId); state.threadKeysByRunId.set(pendingIdle.runId, key); const owns = await this.#acquireOrTransferThreadLease(pubsub, key, pendingIdle.runId, fromRunId); if (!owns.acquired) { if (state.activeThreadRunIds.get(key) === pendingIdle.runId) state.activeThreadRunIds.delete(key); state.threadKeysByRunId.delete(pendingIdle.runId); state.preRunSignalsByThread.delete(key); if (owns.owner) await this.#publishAndWait(pubsub, key, { type: "signal-enqueued", runId: owns.owner, signal: this.#serializeSignal(pendingIdle.signal), sourceId: this.#getSourceId() }).catch(() => {}); await this.#drainPendingIdleSignals(state, pubsub, key, fromRunId); return true; } try { const output = await pendingIdle.agent.stream(pendingIdle.signal, { ...pendingIdle.streamOptions, runId: pendingIdle.runId, memory: withThreadMemory(pendingIdle.streamOptions?.memory, pendingIdle.resourceId, pendingIdle.threadId) }); if ((idleQueue?.length ?? 0) > 0) { const nextRecord = state.threadRunsById.get(output.runId); if (nextRecord) this.#watchThreadRunCompletion(state, pubsub, key, nextRecord); } } catch (err) { state.threadKeysByRunId.delete(pendingIdle.runId); this.#cleanupPreparedRun(state, pendingIdle.runId); if (state.activeThreadRunIds.get(key) === pendingIdle.runId) state.activeThreadRunIds.delete(key); this.#publish(pubsub, key, { type: "run-failed", runId: pendingIdle.runId, error: require_error.getErrorFromUnknown(err).message }); if (!await this.#drainPendingIdleSignals(state, pubsub, key, pendingIdle.runId)) this.#releaseThreadLease(pubsub, key, pendingIdle.runId); } return true; } /** * Drains queued signals for a run. * * - `scope: 'pending'` (default) returns active-run follow-up signals — each * becomes its own model turn via `signalDrainStep`. * - `scope: 'pre-run'` returns signals queued before the run's first model * request — the first LLM step folds these into that request. */ drainPendingSignals(runId, pubsub, scope = "pending") { const state = this.#getState(pubsub); const record = state.threadRunsById.get(runId); const key = record ? this.#threadKey(record.resourceId, record.threadId) : state.threadKeysByRunId.get(runId); if (!key) return []; const signalsByThread = scope === "pre-run" ? state.preRunSignalsByThread : state.pendingSignalsByThread; const queue = signalsByThread.get(key); if (!queue || queue.length === 0) return []; signalsByThread.delete(key); return queue; } async waitForCrossAgentThreadRun(agent, options, pubsub) { const { threadId, resourceId } = this.#getThreadTarget(options); if (!threadId) return; const state = this.#getState(pubsub); const key = this.#threadKey(resourceId, threadId); while (true) { const activeRunId = state.activeThreadRunIds.get(key); if (!activeRunId) return; const activeRecord = state.threadRunsById.get(activeRunId); if (activeRecord) { if (activeRecord.agent.id === agent.id || !this.#isThreadBlockingRun(state, activeRecord)) return; await activeRecord.output._waitUntilFinished().catch(() => {}); continue; } if (state.threadKeysByRunId.get(activeRunId) === key) return; await this.#waitForRemoteRunToFinish(pubsub, key, activeRunId); } } async #waitForRemoteRunToFinish(pubsub, key, runId) { const resolvedPubSub = this.#getPubSub(pubsub); const { provider, isFallback } = this.#resolveLeaseProvider(resolvedPubSub); const topic = this.#threadTopic(key); let timer; let subscribed = false; let settled = false; let resolveWait; const wait = new Promise((resolve) => { resolveWait = resolve; }); const clearRemoteActive = (streamId) => { const state = this.#getState(resolvedPubSub); if (state.activeThreadRunIds.get(key) !== runId || streamId && state.activeThreadStreamIds.get(key) !== streamId) return; state.activeThreadRunIds.delete(key); state.activeThreadStreamIds.delete(key); if (state.remoteThreadKeysByRunId.get(runId) === key) state.remoteThreadKeysByRunId.delete(runId); }; const finish = () => { if (settled) return; settled = true; if (timer) clearTimeout(timer); resolveWait(); }; const checkLease = async () => { if (settled) return; if (isFallback) return; const owner = await provider.getLeaseOwner(key).catch(() => void 0); if (settled) return; if (owner !== runId) { clearRemoteActive(); finish(); return; } timer = setTimeout(() => void checkLease(), AGENT_THREAD_LEASE_TTL_MS); }; const onEvent = (event) => { const data = event.data; if ((data?.type === "run-completed" || data?.type === "run-aborted" || data?.type === "run-failed") && data.runId === runId) { clearRemoteActive(data.streamId); finish(); } }; try { await resolvedPubSub.subscribe(topic, onEvent); subscribed = true; if (!isFallback) timer = setTimeout(() => void checkLease(), AGENT_THREAD_LEASE_TTL_MS); await wait; } catch { finish(); await wait; } finally { if (timer) clearTimeout(timer); if (subscribed) await resolvedPubSub.unsubscribe(topic, onEvent).catch(() => {}); } } async subscribeToThread(agent, options, pubsub) { const resolvedPubSub = this.#getPubSub(pubsub); const state = this.#getState(resolvedPubSub); const key = this.#threadKey(options.resourceId, options.threadId); const topic = this.#threadTopic(key); const seenStreamIds = /* @__PURE__ */ new Set(); const pendingRuns = []; const waiters = []; const remoteRuns = /* @__PURE__ */ new Map(); let done = false; const wake = () => { while (waiters.length) waiters.shift()?.(); }; const activeRunId = () => { const runId = state.activeThreadRunIds.get(key); if (!runId) return null; const record = state.threadRunsById.get(runId); if (!record) return runId; return this.#isThreadBlockingRun(state, record) ? runId : null; }; const enqueueRun = (record) => { if (done || seenStreamIds.has(record.streamId)) return; seenStreamIds.add(record.streamId); pendingRuns.push(record); wake(); }; const createRemoteRun = (runId, streamId, streamSeq) => { const remoteRun = { parts: [], waiters: [], finishWaiters: [], done: false, stream: void 0, closed: false }; remoteRun.stream = new ReadableStream({ pull(controller) { const drain = () => { if (remoteRun.closed) return; while (remoteRun.parts.length > 0) controller.enqueue(remoteRun.parts.shift()); if (remoteRun.done) { remoteRun.closed = true; controller.close(); } }; drain(); if (!remoteRun.done && !remoteRun.closed) remoteRun.waiters.push(drain); }, cancel() { remoteRun.done = true; remoteRun.closed = true; remoteRun.waiters.length = 0; while (remoteRun.finishWaiters.length) remoteRun.finishWaiters.shift()?.(); } }); remoteRuns.set(streamId, remoteRun); return { agent, output: { runId, status: "running", fullStream: remoteRun.stream, _waitUntilFinished: async () => { if (remoteRun.done) return; await new Promise((resolve) => remoteRun.finishWaiters.push(resolve)); } }, runId, streamId, streamSeq, lifecycle: "running", threadId: options.threadId, resourceId: options.resourceId, streamOptions: {} }; }; const localStreamIds = /* @__PURE__ */ new Set(); const replayedStreamIds = /* @__PURE__ */ new Set(); let currentReader = null; let activeReaderRunId = null; let cancelledByAbort = false; const markActiveIfLive = async (runId, streamId, local) => { if (!local && !await this.#hasLiveThreadLease(resolvedPubSub, key, runId)) return; state.activeThreadRunIds.set(key, runId); state.activeThreadStreamIds.set(key, streamId); if (!local) state.remoteThreadKeysByRunId.set(runId, key); }; const clearActiveIfCurrent = (runId, streamId) => { if (state.activeThreadRunIds.get(key) !== runId || streamId && state.activeThreadStreamIds.get(key) !== streamId) return; state.activeThreadRunIds.delete(key); state.activeThreadStreamIds.delete(key); if (state.remoteThreadKeysByRunId.get(runId) === key) state.remoteThreadKeysByRunId.delete(runId); }; const handleEvent = async (event) => { if (done) return; const data = event.data; if (!data) return; if (data.type === "run-registered") { const localRecord = state.threadRunsByStreamId.get(data.streamId); if (localRecord) localStreamIds.add(data.streamId); else replayedStreamIds.add(data.streamId); await markActiveIfLive(data.runId, data.streamId, Boolean(localRecord)); const record = localRecord ?? createRemoteRun(data.runId, data.streamId, data.streamSeq); enqueueRun(record); wake(); return; } if (data.type === "stream-part") { if (data.sourceId === this.#id && (localStreamIds.has(data.streamId) || !replayedStreamIds.has(data.streamId))) return; if (state.activeThreadRunIds.get(key) !== data.runId || state.activeThreadStreamIds.get(key) !== data.streamId) await markActiveIfLive(data.runId, data.streamId, false); let remoteRun = remoteRuns.get(data.streamId); if (!remoteRun) { enqueueRun(createRemoteRun(data.runId, data.streamId, state.streamSeqByRunId.get(data.runId) ?? 1)); remoteRun = remoteRuns.get(data.streamId); if (!remoteRun) return; } remoteRun.parts.push(data.part); while (remoteRun.waiters.length) remoteRun.waiters.shift()?.(); return; } if (data.type === "signal-enqueued") { if (data.sourceId === this.#id) return; const signalsByThread = data.preRun ? state.preRunSignalsByThread : state.pendingSignalsByThread; const queue = signalsByThread.get(key) ?? []; queue.push(require_signals.createSignal(data.signal)); signalsByThread.set(key, queue); return; } if (data.type === "run-abort-requested") { if (state.preparedRunsById.has(data.runId) && state.threadKeysByRunId.get(data.runId) === key && state.activeThreadRunIds.get(key) === data.runId && state.activeThreadStreamIds.get(key) === data.streamId && await this.#hasLiveThreadLease(resolvedPubSub, key, data.runId)) this.abortRun(data.runId, resolvedPubSub); return; } if (data.type === "run-failed") { const eventStreamId = data.streamId ?? data.runId; clearActiveIfCurrent(data.runId, data.streamId); let errorRun; let remoteRun = remoteRuns.get(eventStreamId); if (!remoteRun) { errorRun = createRemoteRun(data.runId, eventStreamId, state.streamSeqByRunId.get(data.runId) ?? 1); remoteRun = remoteRuns.get(eventStreamId); } if (remoteRun) { remoteRun.parts.push({ type: "error", payload: { error: new Error(data.error) } }); remoteRun.done = true; while (remoteRun.waiters.length) remoteRun.waiters.shift()?.(); while (remoteRun.finishWaiters.length) remoteRun.finishWaiters.shift()?.(); remoteRuns.delete(eventStreamId); seenStreamIds.delete(eventStreamId); } if (errorRun) enqueueRun(errorRun); await this.#drainPendingIdleSignals(state, resolvedPubSub, key, data.runId); wake(); return; } if (data.type === "run-completed" || data.type === "run-aborted" || data.type === "run-suspended") { const eventStreamId = data.streamId ?? data.runId; if (data.type === "run-suspended") { state.suspendedRunIds.add(data.runId); const record = state.threadRunsByStreamId.get(eventStreamId) ?? state.threadRunsById.get(data.runId); if (record) record.lifecycle = "suspended"; } else clearActiveIfCurrent(data.runId, data.streamId); if (data.type !== "run-suspended") this.#clearSuspendedRun(state, data.runId); const remoteRun = remoteRuns.get(eventStreamId); if (remoteRun) { remoteRun.done = true; while (remoteRun.waiters.length) remoteRun.waiters.shift()?.(); while (remoteRun.finishWaiters.length) remoteRun.finishWaiters.shift()?.(); remoteRuns.delete(eventStreamId); seenStreamIds.delete(eventStreamId); } if (data.type === "run-aborted" && activeReaderRunId === data.runId && currentReader) { cancelledByAbort = true; try { currentReader.cancel(); } catch {} } if (data.type !== "run-suspended") await this.#drainPendingIdleSignals(state, resolvedPubSub, key, data.runId); wake(); } }; let eventTail = Promise.resolve(); const onEvent = (event) => { eventTail = eventTail.then(() => handleEvent(event)).catch(() => {}); }; await resolvedPubSub.subscribe(topic, onEvent); const currentRunId = activeRunId(); const currentRecord = currentRunId ? state.threadRunsById.get(currentRunId) : void 0; if (currentRecord) { localStreamIds.add(currentRecord.streamId); enqueueRun(currentRecord); } const unsubscribe = () => { if (done) return; done = true; resolvedPubSub.unsubscribe(topic, onEvent).catch(() => {}); if (currentReader) try { currentReader.cancel(); } catch {} wake(); }; return { activeRunId, abort: () => this.abortThread(options, resolvedPubSub), unsubscribe, stream: (async function* () { try { while (!done || pendingRuns.length > 0) { if (pendingRuns.length === 0) { await new Promise((resolve) => waiters.push(resolve)); continue; } const run = pendingRuns.shift(); const reader = (run.createSubscriberStream?.() ?? run.output.fullStream).getReader(); currentReader = reader; activeReaderRunId = run.runId; let readerReleased = false; try { while (true) { const { value: part, done: streamDone } = await reader.read(); if (streamDone) break; const typedPart = part; yield typedPart && typeof typedPart === "object" && !("runId" in typedPart) ? { ...typedPart, runId: run.runId } : typedPart; if (done) break; const finishReason = typedPart.finishReason ?? typedPart.payload?.finishReason; if (typedPart.type === "error" || typedPart.type === "abort" || typedPart.type === "finish" && finishReason !== "tool-calls") { readerReleased = true; (async () => { try { while (true) { const { done: d } = await reader.read(); if (d) break; } } catch {} reader.releaseLock(); })(); break; } } if (!readerReleased && !done && cancelledByAbort) { yield { type: "abort", runId: run.runId }; cancelledByAbort = false; } } finally { currentReader = null; activeReaderRunId = null; if (!readerReleased) reader.releaseLock(); } } } finally { unsubscribe(); } })() }; } sendMessage(agent, message, target, pubsub) { return this.sendSignal(agent, this.#createMessageSignalInput(message), target, pubsub); } queueMessage(agent, message, target, pubsub) { const state = this.#getState(pubsub); const acceptedAt = /* @__PURE__ */ new Date(); let key; let runId = target.runId; let activeRecord; if (target.resourceId && target.threadId) { key = this.#threadKey(target.resourceId, target.threadId); const activeRunId = state.activeThreadRunIds.get(key); activeRecord = activeRunId ? state.threadRunsById.get(activeRunId) : void 0; if (activeRecord && !this.#isThreadBlockingRun(state, activeRecord)) { state.activeThreadRunIds.delete(key); activeRecord = void 0; } runId ??= activeRunId; } if (runId) { activeRecord ??= state.threadRunsById.get(runId); if (activeRecord) key ??= this.#threadKey(activeRecord.resourceId, activeRecord.threadId); } const resourceId = target.resourceId ?? activeRecord?.resourceId; const threadId = target.threadId ?? activeRecord?.threadId; if (!resourceId || !threadId) throw new Error("resourceId and threadId are required to queue a message"); key ??= this.#threadKey(resourceId, threadId); const signal = require_signals.createMessageSignal(message, { id: this.#generateSignalMessag