@dbos-inc/vercel-ai 0.2.5 → 0.4.4

package/src/agent-tool.ts

@@ -0,0 +1,122 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { DBOS } from '@dbos-inc/dbos-sdk';
+import type { FlexibleSchema, ModelMessage, Tool } from 'ai' with { 'resolution-mode': 'import' };
+import { writeDurableStream, writeToolRecord } from './durable-stream';
+import { isInWorkflowFunction } from './internal';
+
+/** Marks a tool built by agentTool: durableTools leaves it unwrapped (it is a child workflow, not a step) and binds its durable stream. */
+export const AGENT_TOOL: unique symbol = Symbol.for('@dbos-inc/vercel-ai/agentTool');
+
+// Captured at module load, outside any workflow: a cancellation triggered by an abort must not claim a function id in the parent's log.
+const outsideWorkflow = AsyncLocalStorage.snapshot();
+
+// The child's row appears shortly after the call starts and a cancel of a missing row is a no-op, so wait for it, giving up once the call has settled.
+async function cancelChild(childID: string, settled: () => boolean): Promise<void> {
+  for (let i = 0; i < 100 && !settled(); i++) {
+    if (await DBOS.getWorkflowStatus(childID)) {
+      // The child's own sub-agents are its children; without the cascade they would run on with nobody awaiting them.
+      await DBOS.cancelWorkflow(childID, { cancelChildren: true });
+      return;
+    }
+    await new Promise((resolve) => setTimeout(resolve, 50));
+  }
+}
+
+// What agentTool needs from an agent: the AI SDK's Agent interface, structurally.
+type StreamingAgent = {
+  stream(options: { prompt: string } | { messages: ModelMessage[] }): PromiseLike<{ consumeStream(): PromiseLike<void>; readonly text: PromiseLike<string> }>;
+};
+
+export interface AgentToolOptions<INPUT, AGENT extends StreamingAgent, OUTPUT> {
+  /** Name of the child workflow; must be unique. */
+  name: string;
+  description: string;
+  inputSchema: FlexibleSchema<INPUT>;
+  agent: AGENT;
+  /** Turns the tool's input into the sub-agent's prompt. */
+  prompt: (input: INPUT) => string | ModelMessage[];
+  /** Turns the sub-agent's result into the tool's output (default: its final text); must be serializable. */
+  output?: (result: Awaited<ReturnType<AGENT['stream']>>) => OUTPUT | Promise<OUTPUT>;
+  /** Record the call in this durable stream: a `data-dbos-subagent` part naming the child workflow, then its output. */
+  durableStream?: string;
+  /** Name of a DBOS queue to run each child on, e.g. to bound how many sub-agents run at once. */
+  queue?: string;
+  /** Workflow timeout for each child. */
+  timeoutMS?: number;
+}
+
+export type AgentTool<INPUT, OUTPUT> = Tool<INPUT, OUTPUT> & {
+  /** The registered child workflow; call it directly to run the sub-agent without a model in the loop. */
+  workflow: (input: INPUT) => Promise<OUTPUT>;
+};
+
+/**
+ * Wraps an agent as a tool whose every call runs as a child workflow: durable at model-call granularity, with its own
+ * concurrency guard, safe to call in parallel, and visible as a child in the parent's step list. Call it at module load,
+ * before `DBOS.launch()`, since it registers the child workflow.
+ */
+export function agentTool<INPUT, AGENT extends StreamingAgent, OUTPUT = string>(
+  options: AgentToolOptions<INPUT, AGENT, OUTPUT>,
+): AgentTool<INPUT, OUTPUT> {
+  const { name, agent, prompt, output } = options;
+  const run = async (input: INPUT): Promise<OUTPUT> => {
+    const request = prompt(input);
+    // stream, not generate: only streamed calls write to a durable stream.
+    const result = await agent.stream(typeof request === 'string' ? { prompt: request } : { messages: request });
+    await result.consumeStream();
+    return output ? await output(result as Awaited<ReturnType<AGENT['stream']>>) : ((await result.text) as OUTPUT);
+  };
+  const registered = DBOS.registerWorkflow(run, { name });
+  // Unbound on purpose: DBOS's registered invoker reads `this`, and as a method `this` would be the tool object.
+  const workflow = (input: INPUT): Promise<OUTPUT> => registered(input);
+  return build(options, registered, workflow, options.durableStream);
+}
+
+function build<INPUT, AGENT extends StreamingAgent, OUTPUT>(
+  options: AgentToolOptions<INPUT, AGENT, OUTPUT>,
+  registered: (input: INPUT) => Promise<OUTPUT>,
+  workflow: (input: INPUT) => Promise<OUTPUT>,
+  durableStream: string | undefined,
+): AgentTool<INPUT, OUTPUT> {
+  const { name, description, inputSchema, queue, timeoutMS } = options;
+  const execute = async (input: INPUT, execOptions: { toolCallId: string; abortSignal?: AbortSignal }): Promise<OUTPUT> => {
+    if (!isInWorkflowFunction()) return workflow(input);
+    const { toolCallId } = execOptions;
+    // The tool call id comes from the checkpointed model output, so the child id is the same on replay and known for cancellation.
+    const childID = `${DBOS.workflowID}-${toolCallId}`;
+    // Start and getResult each reserve their function id synchronously here, so parallel calls replay in order.
+    const started = DBOS.startWorkflow(registered, { workflowID: childID, queueName: queue, timeoutMS })(input);
+    const pending = DBOS.getResult<OUTPUT>(childID);
+    started.catch(() => {});
+    pending.catch(() => {});
+    let settled = false;
+    const cancel = () => void outsideWorkflow(() => cancelChild(childID, () => settled)).catch(() => {});
+    execOptions.abortSignal?.addEventListener('abort', cancel, { once: true });
+    if (execOptions.abortSignal?.aborted) cancel();
+    try {
+      if (durableStream) {
+        await writeDurableStream(durableStream, [
+          { type: 'data-dbos-subagent', id: toolCallId, data: { toolCallId, workflowID: childID, name } },
+        ]);
+      }
+      await started;
+      const result = (await pending) as OUTPUT;
+      // A tool record, not a raw chunk, so the reader masks a sub-agent's error text like any other tool's.
+      if (durableStream) await writeToolRecord(durableStream, toolCallId, { output: result });
+      return result;
+    } catch (error) {
+      if (durableStream) await writeToolRecord(durableStream, toolCallId, { errorText: error instanceof Error ? error.message : String(error) });
+      throw error;
+    } finally {
+      settled = true;
+      execOptions.abortSignal?.removeEventListener('abort', cancel);
+    }
+  };
+  return {
+    description,
+    inputSchema,
+    execute,
+    workflow,
+    [AGENT_TOOL]: (key: string) => build(options, registered, workflow, key),
+  } as unknown as AgentTool<INPUT, OUTPUT>;
+}

package/src/durable-stream.ts

@@ -0,0 +1,395 @@
+import { randomUUID } from 'node:crypto';
+import { DBOS, Error as DBOSErrors, StatusString } from '@dbos-inc/dbos-sdk';
+import type { UIMessageChunk } from 'ai' with { 'resolution-mode': 'import' };
+import type { LanguageModelV4FinishReason, LanguageModelV4StreamPart } from '@ai-sdk/provider' with { 'resolution-mode': 'import' };
+
+/** Durable stream config: the DBOS stream key, or the key plus batching limits for the model step's writes. */
+export type DurableStreamOptions = string | { key: string; maxBatchParts?: number; maxBatchDelayMs?: number };
+
+/** One DBOS stream value; the reader turns these into AI SDK UI message chunks. */
+export type DurableStreamRecord =
+  | { kind: 'model'; step: number; attempt: string; parts: LanguageModelV4StreamPart[] }
+  | { kind: 'model-end'; step: number; attempt: string; finishReason?: LanguageModelV4FinishReason; aborted?: true }
+  | { kind: 'tool'; step: number; attempt: number; toolCallId: string; output?: unknown; errorText?: string }
+  | { kind: 'ui'; step?: number; attempt?: number; chunks: UIMessageChunk[] }
+  | { kind: 'end'; finishReason: string };
+
+interface ResolvedDurableStream {
+  key: string;
+  maxBatchParts: number;
+  maxBatchDelayMs: number;
+}
+
+export function resolveDurableStream(options: DurableStreamOptions | undefined): ResolvedDurableStream | undefined {
+  if (options === undefined) return undefined;
+  const config = typeof options === 'string' ? { key: options } : options;
+  return { key: config.key, maxBatchParts: config.maxBatchParts ?? 20, maxBatchDelayMs: config.maxBatchDelayMs ?? 25 };
+}
+
+function stepInfo(): { step: number; attempt: number } {
+  return { step: DBOS.stepID ?? -1, attempt: DBOS.stepStatus?.currentAttempt ?? 1 };
+}
+
+// Only the parts toUIChunk renders; anything else would be written and never read.
+function isContentPart(part: LanguageModelV4StreamPart): boolean {
+  switch (part.type) {
+    case 'text-start':
+    case 'text-delta':
+    case 'text-end':
+    case 'reasoning-start':
+    case 'reasoning-delta':
+    case 'reasoning-end':
+    case 'tool-input-start':
+    case 'tool-input-delta':
+    case 'tool-call':
+    case 'tool-result':
+    case 'source':
+    case 'file':
+      return true;
+    default:
+      return false;
+  }
+}
+
+// Bytes become base64 so a file part stays compact in Postgres.
+function encodePart(part: LanguageModelV4StreamPart): LanguageModelV4StreamPart {
+  if (part.type === 'file' && part.data.type === 'data' && part.data.data instanceof Uint8Array) {
+    return { ...part, data: { type: 'data', data: Buffer.from(part.data.data).toString('base64') } };
+  }
+  return part;
+}
+
+// A transient write error (the SDK already retries offset conflicts) gets a few attempts before it fails the model call.
+async function writeWithRetry(key: string, record: DurableStreamRecord): Promise<void> {
+  for (let attempt = 1; ; attempt++) {
+    try {
+      return await DBOS.writeStream(key, record);
+    } catch (error) {
+      if (attempt === 3) throw error;
+      await new Promise((resolve) => setTimeout(resolve, 100 * attempt));
+    }
+  }
+}
+
+/** Batches a live model step's parts into step-scope stream writes; nothing is written on replay because the step body does not run. */
+export class ModelStreamWriter {
+  private pending: LanguageModelV4StreamPart[] = [];
+  private chain: Promise<void> = Promise.resolve();
+  private failure: unknown;
+  private timer: ReturnType<typeof setTimeout> | undefined;
+  private readonly step = DBOS.stepID ?? -1;
+  // Unique per execution of the step, so a recovered run's re-execution is distinguishable from the crashed one.
+  private readonly attempt = randomUUID();
+
+  constructor(private readonly config: ResolvedDurableStream) {}
+
+  push(part: LanguageModelV4StreamPart): void {
+    if (!isContentPart(part)) return;
+    this.pending.push(encodePart(part));
+    if (this.pending.length >= this.config.maxBatchParts) this.flush();
+    else this.timer ??= setTimeout(() => this.flush(), this.config.maxBatchDelayMs);
+  }
+
+  /** Flushes, records how the call ended, and resolves once every write is durable; a write that failed after retries fails the call here. */
+  async end(outcome: { finishReason: LanguageModelV4FinishReason } | { aborted: true }): Promise<void> {
+    this.flush();
+    this.write({ kind: 'model-end', step: this.step, attempt: this.attempt, ...outcome });
+    await this.chain;
+    if (this.failure !== undefined) throw this.failure;
+  }
+
+  /** After a failure: flush what streamed so the record matches what the consumer saw; the stream's end then comes from the workflow's status. */
+  async abandon(): Promise<void> {
+    this.flush();
+    await this.chain;
+  }
+
+  private flush(): void {
+    if (this.timer !== undefined) clearTimeout(this.timer);
+    this.timer = undefined;
+    if (this.pending.length === 0) return;
+    const parts = this.pending;
+    this.pending = [];
+    this.write({ kind: 'model', step: this.step, attempt: this.attempt, parts });
+  }
+
+  // Every link has a handler, so a rejection can never sit unobserved; after one failure later writes are skipped.
+  private write(record: DurableStreamRecord): void {
+    this.chain = this.chain
+      .then(() => (this.failure === undefined ? writeWithRetry(this.config.key, record) : undefined))
+      .catch((error: unknown) => {
+        this.failure ??= error;
+      });
+  }
+}
+
+/** Records a tool call's outcome from inside its step. */
+export function writeToolRecord(key: string, toolCallId: string, outcome: { output: unknown } | { errorText: string }): Promise<void> {
+  const record: DurableStreamRecord = { kind: 'tool', ...stepInfo(), toolCallId, ...outcome };
+  return DBOS.writeStream(key, record);
+}
+
+/**
+ * Appends UI message chunks to a durable stream. From a step the write is cheap and at-least-once, so give data parts
+ * stable ids; from workflow code it is a checkpointed step, so the number of calls must be deterministic.
+ */
+export function writeDurableStream(key: string, chunks: UIMessageChunk[]): Promise<void> {
+  const status = DBOS.stepStatus;
+  const record: DurableStreamRecord = { kind: 'ui', step: status?.stepID, attempt: status?.currentAttempt, chunks };
+  return DBOS.writeStream(key, record);
+}
+
+/** Marks the end of the turn explicitly and closes the stream; without it the reader infers the end from the last model call or the workflow's status. */
+export async function closeDurableStream(key: string, finishReason = 'stop'): Promise<void> {
+  const record: DurableStreamRecord = { kind: 'end', finishReason };
+  await DBOS.writeStream(key, record);
+  await DBOS.closeStream(key);
+}
+
+/** What the reader needs from DBOS: the `DBOS` class in a launched process, or a `DBOSClient` anywhere else. */
+export interface DurableStreamSource {
+  readStream<T>(workflowID: string, key: string, options?: { offset?: number }): AsyncGenerator<T, void, unknown>;
+  readStreamOffset<T>(workflowID: string, key: string, offset: number, options?: { timeoutSeconds?: number }): Promise<T>;
+  retrieveWorkflow(workflowID: string): { getStatus(): Promise<{ status: string; error?: unknown } | null> };
+}
+
+export interface ReadDurableStreamOptions {
+  workflowID: string;
+  key: string;
+  /** Id for the `start` chunk; omitted on a resume (`offset` > 0). */
+  messageId?: string;
+  /** Number of records already consumed, from the last `data-dbos-offset` chunk. */
+  offset?: number;
+  /** Defaults to `DBOS`; pass a `DBOSClient` to read from a process that has not launched DBOS. */
+  client?: DurableStreamSource;
+  /** Emit reasoning parts (default true, as in the AI SDK). */
+  sendReasoning?: boolean;
+  /** Emit source parts (default false, as in the AI SDK). */
+  sendSources?: boolean;
+  /** Text sent to clients for a workflow or tool error; defaults to a generic message, as in the AI SDK, so server details stay private. */
+  onError?: (error: unknown) => string;
+}
+
+/** Reads a durable stream as AI SDK UI message chunks, live or after the fact, resuming from `offset`. */
+export function readDurableStream(options: ReadDurableStreamOptions): ReadableStream<UIMessageChunk> {
+  const iterator = uiChunks(options);
+  return new ReadableStream<UIMessageChunk>({
+    async pull(controller) {
+      const { value, done } = await iterator.next();
+      if (done) controller.close();
+      else controller.enqueue(value);
+    },
+    async cancel() {
+      await iterator.return(undefined);
+    },
+  });
+}
+
+async function* uiChunks(options: ReadDurableStreamOptions): AsyncGenerator<UIMessageChunk> {
+  const { workflowID, key, sendReasoning = true, sendSources = false, onError = () => 'An error occurred.' } = options;
+  const client: DurableStreamSource = options.client ?? DBOS;
+  const state: ReaderState = {
+    offset: options.offset ?? 0,
+    resumed: (options.offset ?? 0) > 0,
+    openParts: new Map(),
+    ended: false,
+  };
+  if (state.offset === 0) yield { type: 'start', messageId: options.messageId };
+  const emit = (record: DurableStreamRecord) => emitRecord(state, record, { sendReasoning, sendSources, onError });
+
+  // Phase 1: everything already stored, one value per query until an offset is empty; a superseded attempt is skipped whole.
+  const history: DurableStreamRecord[] = [];
+  for (;;) {
+    try {
+      history.push(await client.readStreamOffset<DurableStreamRecord>(workflowID, key, state.offset + history.length, { timeoutSeconds: 0 }));
+    } catch (error) {
+      if (!DBOSErrors.isStreamTimeoutError(error)) throw error;
+      break;
+    }
+  }
+  const finalAttempt = new Map<number, string>();
+  for (const record of history) {
+    if (record.kind === 'model' || record.kind === 'model-end') finalAttempt.set(record.step, record.attempt);
+  }
+  for (const record of history) {
+    const stale = (record.kind === 'model' || record.kind === 'model-end') && finalAttempt.get(record.step) !== record.attempt;
+    yield* stale ? skipRecord(state) : emit(record);
+    if (state.ended) return;
+  }
+
+  // Phase 2: live; a re-executed step shows up as a new attempt and is handed off in place.
+  for await (const record of client.readStream<DurableStreamRecord>(workflowID, key, { offset: state.offset })) {
+    yield* emit(record);
+    if (state.ended) return;
+  }
+
+  // No end record: the workflow's status decides how the turn ended (a stream closed while it still runs counts as finished).
+  const status = (await client.retrieveWorkflow(workflowID).getStatus())?.status;
+  yield* closeStep(state);
+  if (status === StatusString.CANCELLED) {
+    yield { type: 'abort' };
+  } else if (status === undefined || status === StatusString.SUCCESS || status === StatusString.PENDING || status === StatusString.ENQUEUED) {
+    // The AI SDK's finish schema has no 'unknown'; a turn with no model call ends with the reason omitted.
+    yield state.finishReason === undefined ? { type: 'finish' } : { type: 'finish', finishReason: state.finishReason as UIFinishReason };
+  } else {
+    const error = (await client.retrieveWorkflow(workflowID).getStatus())?.error;
+    yield { type: 'error', errorText: onError(error ?? new Error('The workflow ended before the response completed.')) };
+    yield { type: 'finish', finishReason: 'error' };
+  }
+}
+
+interface ReaderState {
+  offset: number;
+  resumed: boolean;
+  openStep?: number;
+  openAttempt?: string;
+  // Text/reasoning parts of the open attempt that have started but not ended, by UI part id.
+  openParts: Map<string, 'text' | 'reasoning'>;
+  finishReason?: string;
+  ended: boolean;
+}
+
+function offsetChunk(state: ReaderState): UIMessageChunk {
+  return { type: 'data-dbos-offset', data: { offset: state.offset }, transient: true } as UIMessageChunk;
+}
+
+function* skipRecord(state: ReaderState): Generator<UIMessageChunk> {
+  state.offset += 1;
+  yield offsetChunk(state);
+}
+
+function* closeStep(state: ReaderState): Generator<UIMessageChunk> {
+  if (state.openStep !== undefined) yield { type: 'finish-step' };
+  state.openStep = undefined;
+  state.openAttempt = undefined;
+  state.openParts.clear();
+}
+
+// A live re-execution of the open step: end the stale attempt's parts and tell the client which ones to discard.
+function* supersede(state: ReaderState, attempt: string): Generator<UIMessageChunk> {
+  for (const [id, kind] of state.openParts) yield { type: kind === 'text' ? 'text-end' : 'reasoning-end', id };
+  yield {
+    type: 'data-dbos-superseded',
+    data: { attempt: state.openAttempt, parts: [...state.openParts.keys()] },
+    transient: true,
+  } as UIMessageChunk;
+  state.openParts.clear();
+  state.openAttempt = attempt;
+}
+
+function* emitRecord(
+  state: ReaderState,
+  record: DurableStreamRecord,
+  filter: { sendReasoning: boolean; sendSources: boolean; onError: (error: unknown) => string },
+): Generator<UIMessageChunk> {
+  state.offset += 1;
+  switch (record.kind) {
+    case 'model': {
+      if (state.openStep !== record.step) {
+        yield* closeStep(state);
+        if (!state.resumed) yield { type: 'start-step' };
+        state.resumed = false;
+        state.openStep = record.step;
+        state.openAttempt = record.attempt;
+      } else if (state.openAttempt !== record.attempt) {
+        yield* supersede(state, record.attempt);
+      }
+      for (const part of record.parts) {
+        if (!filter.sendReasoning && part.type.startsWith('reasoning-')) continue;
+        if (!filter.sendSources && part.type === 'source') continue;
+        const chunk = toUIChunk(part, record.attempt);
+        if (!chunk) continue;
+        if (chunk.type === 'text-start' || chunk.type === 'reasoning-start') state.openParts.set(chunk.id, chunk.type === 'text-start' ? 'text' : 'reasoning');
+        if (chunk.type === 'text-end' || chunk.type === 'reasoning-end') state.openParts.delete(chunk.id);
+        yield chunk;
+      }
+      break;
+    }
+    case 'model-end':
+      // The stream outlives the call: the workflow may run more calls, so only its end (or closeDurableStream) ends the turn.
+      if (record.attempt === state.openAttempt) state.finishReason = record.aborted ? 'other' : record.finishReason?.unified;
+      break;
+    case 'tool':
+      // Local tool errors are masked like the AI SDK does; provider-executed ones (in model records) pass through verbatim.
+      yield record.errorText !== undefined
+        ? { type: 'tool-output-error', toolCallId: record.toolCallId, errorText: filter.onError(new Error(record.errorText)) }
+        : { type: 'tool-output-available', toolCallId: record.toolCallId, output: record.output };
+      break;
+    case 'ui':
+      yield* record.chunks;
+      break;
+    case 'end':
+      // The terminal chunk comes last, so the offset goes out first.
+      yield offsetChunk(state);
+      yield* closeStep(state);
+      yield { type: 'finish', finishReason: record.finishReason as UIFinishReason };
+      state.ended = true;
+      return;
+  }
+  yield offsetChunk(state);
+}
+
+type UIFinishReason = Extract<UIMessageChunk, { type: 'finish' }>['finishReason'];
+
+function parseInput(input: string): unknown {
+  try {
+    return JSON.parse(input);
+  } catch {
+    return input;
+  }
+}
+
+// Text and reasoning ids are only unique within one model call; the attempt id keeps calls, and re-executions, apart in one message.
+function toUIChunk(part: LanguageModelV4StreamPart, attempt: string): UIMessageChunk | undefined {
+  const id = 'id' in part ? `${attempt}:${part.id}` : '';
+  switch (part.type) {
+    case 'text-start':
+      return { type: 'text-start', id, providerMetadata: part.providerMetadata };
+    case 'text-delta':
+      return { type: 'text-delta', id, delta: part.delta, providerMetadata: part.providerMetadata };
+    case 'text-end':
+      return { type: 'text-end', id, providerMetadata: part.providerMetadata };
+    case 'reasoning-start':
+      return { type: 'reasoning-start', id, providerMetadata: part.providerMetadata };
+    case 'reasoning-delta':
+      return { type: 'reasoning-delta', id, delta: part.delta, providerMetadata: part.providerMetadata };
+    case 'reasoning-end':
+      return { type: 'reasoning-end', id, providerMetadata: part.providerMetadata };
+    case 'tool-input-start':
+      return {
+        type: 'tool-input-start',
+        toolCallId: part.id,
+        toolName: part.toolName,
+        providerExecuted: part.providerExecuted,
+        dynamic: part.dynamic,
+        title: part.title,
+        providerMetadata: part.providerMetadata,
+      };
+    case 'tool-input-delta':
+      return { type: 'tool-input-delta', toolCallId: part.id, inputTextDelta: part.delta };
+    case 'tool-call':
+      return {
+        type: 'tool-input-available',
+        toolCallId: part.toolCallId,
+        toolName: part.toolName,
+        input: parseInput(part.input),
+        providerExecuted: part.providerExecuted,
+        dynamic: part.dynamic,
+        providerMetadata: part.providerMetadata,
+      };
+    case 'tool-result':
+      return part.isError
+        ? { type: 'tool-output-error', toolCallId: part.toolCallId, errorText: JSON.stringify(part.result), providerExecuted: true, dynamic: part.dynamic }
+        : { type: 'tool-output-available', toolCallId: part.toolCallId, output: part.result, providerExecuted: true, dynamic: part.dynamic, preliminary: part.preliminary };
+    case 'source':
+      return part.sourceType === 'url'
+        ? { type: 'source-url', sourceId: part.id, url: part.url, title: part.title, providerMetadata: part.providerMetadata }
+        : { type: 'source-document', sourceId: part.id, mediaType: part.mediaType, title: part.title, filename: part.filename, providerMetadata: part.providerMetadata };
+    case 'file': {
+      const url = part.data.type === 'url' ? String(part.data.url) : part.data.type === 'data' ? `data:${part.mediaType};base64,${String(part.data.data)}` : undefined;
+      return url === undefined ? undefined : { type: 'file', url, mediaType: part.mediaType, providerMetadata: part.providerMetadata };
+    }
+    default:
+      return undefined;
+  }
+}

package/src/index.ts

@@ -1,2 +1,13 @@
-export { durableCalls, durableEmbeddingCalls, durableImageCalls } from './middleware';
+export { durableCalls, DurableCallsOptions, durableEmbeddingCalls, durableImageCalls } from './middleware';
 export { durableMCPTools, DurableMCPToolsOptions, MCPClientLike } from './mcp';
+export { durableTools, DurableToolsOptions } from './tools';
+export {
+  closeDurableStream,
+  DurableStreamOptions,
+  DurableStreamRecord,
+  DurableStreamSource,
+  ReadDurableStreamOptions,
+  readDurableStream,
+  writeDurableStream,
+} from './durable-stream';
+export { agentTool, AgentTool, AgentToolOptions } from './agent-tool';

package/src/internal.ts

@@ -11,6 +11,20 @@ export function assertNotInTransaction(operation: string): void {
   }
 }
 
+// Run fn as a durable step named `name` inside a workflow; elsewhere call it directly.
+export function runDurableStep<T>(name: string, fn: () => Promise<T>, config: StepConfig): Promise<T> {
+  assertNotInTransaction(name);
+  if (!isInWorkflowFunction()) return fn();
+  // Restore the AI SDK error identity a replay revival strips, so the SDK's retry/catch logic behaves the same.
+  return DBOS.runStep(fn, { ...config, name }).catch((error: unknown) => {
+    throw restoreAISDKErrorIdentity(error);
+  });
+}
+
+export function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
+  return typeof (value as AsyncIterable<unknown> | null | undefined)?.[Symbol.asyncIterator] === 'function';
+}
+
 // Aborts/timeouts are deliberate cancellations, never transient; retrying just re-runs an already-cancelled call.
 function isAbortError(error: unknown): boolean {
   const name = (error as { name?: unknown } | null)?.name;

package/src/mcp.ts

@@ -1,6 +1,7 @@
-import { DBOS, StepConfig } from '@dbos-inc/dbos-sdk';
+import { StepConfig } from '@dbos-inc/dbos-sdk';
 import type { ToolSet } from 'ai' with { 'resolution-mode': 'import' };
-import { assertNotInTransaction, isInWorkflowFunction, restoreAISDKErrorIdentity, withErrorClassification } from './internal';
+import { isAsyncIterable, runDurableStep, withErrorClassification } from './internal';
+import { writeToolRecord } from './durable-stream';
 
 // Structural type for an MCP client (e.g. from @ai-sdk/mcp) — deliberately loose: the AI SDK ecosystem
 // exact-pins @ai-sdk/provider-utils, so precise Tool types fail to match across skewed copies.
@@ -23,6 +24,8 @@ interface MCPToolLike {
 export interface DurableMCPToolsOptions extends StepConfig {
   /** Forwarded to client.tools() on listing and on each call (e.g. { schemas } for subsetting and output schemas). */
   toolOptions?: unknown;
+  /** Write each tool call's output (or error) to this durable stream from inside its step. */
+  durableStream?: string;
 }
 
 interface DurableToolDef {
@@ -36,10 +39,6 @@ interface DurableToolDef {
 
 type ToolModelOutput = Awaited<ReturnType<NonNullable<ToolSet[string]['toModelOutput']>>>;
 
-function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
-  return typeof (value as AsyncIterable<unknown> | null | undefined)?.[Symbol.asyncIterator] === 'function';
-}
-
 // Mirror of @ai-sdk/mcp's toModelOutput: MCP content becomes model content (text stays text, images become files).
 function mcpToolOutput(output: unknown): ToolModelOutput {
   const result = output as { content?: unknown };
@@ -66,17 +65,11 @@ function mcpToolOutput(output: unknown): ToolModelOutput {
  * checkpointed so a recovered workflow replays results instead of re-invoking the tool.
  */
 export async function durableMCPTools(client: MCPClientLike, options: DurableMCPToolsOptions = {}): Promise<ToolSet> {
-  const { toolOptions, ...stepOptions } = options;
+  const { toolOptions, durableStream, ...stepOptions } = options;
   const stepConfig = withErrorClassification(stepOptions);
   const { asSchema, dynamicTool, jsonSchema } = await import('ai');
-  const run = <T>(name: string, fn: () => Promise<T>, config: StepConfig = stepConfig): Promise<T> => {
-    assertNotInTransaction(name);
-    if (!isInWorkflowFunction()) return fn();
-    // Restore the AI SDK error identity a replay revival strips, so the SDK's retry/catch logic behaves the same.
-    return DBOS.runStep(fn, { ...config, name }).catch((error: unknown) => {
-      throw restoreAISDKErrorIdentity(error);
-    });
-  };
+  const run = <T>(name: string, fn: () => Promise<T>, config: StepConfig = stepConfig): Promise<T> =>
+    runDurableStep(name, fn, config);
 
   // Checkpoint the tool list as plain JSON schemas, so replay reconstructs tools without the live client.
   const listed = await run('mcp.listTools', async () => {
@@ -108,25 +101,35 @@ export async function durableMCPTools(client: MCPClientLike, options: DurableMCP
       toModelOutput: def.convertsOutput ? ({ output }) => mcpToolOutput(output) : undefined,
       // Re-fetch the live tool inside the step (its execute closure can't be checkpointed); replay returns the recorded result.
       execute: (input: unknown, execOptions) => {
-        const signal = (execOptions as { abortSignal?: AbortSignal } | undefined)?.abortSignal;
+        const { abortSignal: signal, toolCallId } = (execOptions ?? {}) as { abortSignal?: AbortSignal; toolCallId?: string };
         // An aborted consumer is done with this call, whatever the failure looks like; a retry would re-run a cancelled side effect.
         const callConfig: StepConfig = {
           ...stepConfig,
           shouldRetry: async (error: unknown) =>
             !signal?.aborted && (stepConfig.shouldRetry ? await stepConfig.shouldRetry(error) : true),
         };
+        // The tool call id comes from the checkpointed model result, so a reordered parallel step fails replay instead of swapping results.
         return run(
-          `mcp.tool.${name}`,
+          `mcp.tool.${name}.${toolCallId ?? 'call'}`,
           async () => {
-            const tool = (await client.tools(toolOptions))[name] as MCPToolLike | undefined;
-            if (typeof tool?.execute !== 'function') throw new Error(`MCP tool "${name}" is not executable.`);
-            const output = await tool.execute(input, execOptions);
-            // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
-            if (isAsyncIterable(output)) {
-              let last: unknown;
-              for await (last of output);
-              return last;
+            let output: unknown;
+            try {
+              const tool = (await client.tools(toolOptions))[name] as MCPToolLike | undefined;
+              if (typeof tool?.execute !== 'function') throw new Error(`MCP tool "${name}" is not executable.`);
+              output = await tool.execute(input, execOptions);
+              // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
+              if (isAsyncIterable(output)) {
+                let last: unknown;
+                for await (last of output);
+                output = last;
+              }
+            } catch (error) {
+              if (durableStream && toolCallId) {
+                await writeToolRecord(durableStream, toolCallId, { errorText: error instanceof Error ? error.message : String(error) });
+              }
+              throw error;
             }
+            if (durableStream && toolCallId) await writeToolRecord(durableStream, toolCallId, { output });
             return output;
           },
           callConfig,