@dbos-inc/vercel-ai 0.4.4 → 0.5.4

package/src/durable-stream.ts

@@ -1,4 +1,4 @@
-import { randomUUID } from 'node:crypto';
+import { createHash, 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' };
@@ -10,7 +10,7 @@ export type DurableStreamOptions = string | { key: string; maxBatchParts?: numbe
 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: 'tool'; step: number; attempt: number; toolCallId: string; output?: unknown; errorText?: string; chunks?: UIMessageChunk[] }
   | { kind: 'ui'; step?: number; attempt?: number; chunks: UIMessageChunk[] }
   | { kind: 'end'; finishReason: string };
 
@@ -123,9 +123,14 @@ export class ModelStreamWriter {
   }
 }
 
-/** 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 };
+/** Records a tool call's outcome, with the message chunks it wrote, from inside its step. */
+export function writeToolRecord(
+  key: string,
+  toolCallId: string,
+  outcome: { output: unknown } | { errorText: string },
+  chunks: UIMessageChunk[] = [],
+): Promise<void> {
+  const record: DurableStreamRecord = { kind: 'tool', ...stepInfo(), toolCallId, ...outcome, ...(chunks.length > 0 && { chunks }) };
   return DBOS.writeStream(key, record);
 }
 
@@ -158,7 +163,10 @@ export interface ReadDurableStreamOptions {
   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. */
+  /**
+   * Number of records already consumed, from the last `data-dbos-offset` chunk. A resume from mid-text can't be applied
+   * to a partial message by the AI SDK (it has no record of the open part); to rebuild a message, read from 0 instead.
+   */
   offset?: number;
   /** Defaults to `DBOS`; pass a `DBOSClient` to read from a process that has not launched DBOS. */
   client?: DurableStreamSource;
@@ -190,11 +198,13 @@ async function* uiChunks(options: ReadDurableStreamOptions): AsyncGenerator<UIMe
   const client: DurableStreamSource = options.client ?? DBOS;
   const state: ReaderState = {
     offset: options.offset ?? 0,
-    resumed: (options.offset ?? 0) > 0,
     openParts: new Map(),
+    detachedParts: new Set(),
+    toolChunksSent: new Map(),
     ended: false,
   };
   if (state.offset === 0) yield { type: 'start', messageId: options.messageId };
+  else await restoreOpenCall(client, workflowID, key, state);
   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.
@@ -208,12 +218,16 @@ async function* uiChunks(options: ReadDurableStreamOptions): AsyncGenerator<UIMe
     }
   }
   const finalAttempt = new Map<number, string>();
-  for (const record of history) {
+  // A re-executed tool call writes another record; the last one matches its checkpoint, so only its chunks (if any) are shown.
+  const finalToolRecord = new Map<string, number>();
+  history.forEach((record, index) => {
     if (record.kind === 'model' || record.kind === 'model-end') finalAttempt.set(record.step, record.attempt);
-  }
-  for (const record of history) {
+    if (record.kind === 'tool') finalToolRecord.set(toolKey(record), index);
+  });
+  for (const [index, record] of history.entries()) {
     const stale = (record.kind === 'model' || record.kind === 'model-end') && finalAttempt.get(record.step) !== record.attempt;
-    yield* stale ? skipRecord(state) : emit(record);
+    const staleChunks = record.kind === 'tool' && record.chunks !== undefined && finalToolRecord.get(toolKey(record)) !== index;
+    yield* stale ? skipRecord(state) : emit(staleChunks ? { ...record, chunks: undefined } : record);
     if (state.ended) return;
   }
 
@@ -240,15 +254,23 @@ async function* uiChunks(options: ReadDurableStreamOptions): AsyncGenerator<UIMe
 
 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'>;
+  // Open parts from before a resume: the client's resumed stream no longer tracks them, so they are never ended.
+  detachedParts: Set<string>;
+  // Digest of the chunks emitted for each tool call (by toolKey).
+  toolChunksSent: Map<string, string>;
   finishReason?: string;
   ended: boolean;
 }
 
+// A step's function id survives re-execution, and with it tell apart calls that reuse a provider's tool call id.
+function toolKey(record: Extract<DurableStreamRecord, { kind: 'tool' }>): string {
+  return `${record.step}:${record.toolCallId}`;
+}
+
 function offsetChunk(state: ReaderState): UIMessageChunk {
   return { type: 'data-dbos-offset', data: { offset: state.offset }, transient: true } as UIMessageChunk;
 }
@@ -263,17 +285,90 @@ function* closeStep(state: ReaderState): Generator<UIMessageChunk> {
   state.openStep = undefined;
   state.openAttempt = undefined;
   state.openParts.clear();
+  state.detachedParts.clear();
+}
+
+// Records outside the open model call a resume reads past before giving up; a reconnect mid-call finds it at once.
+const RESTORE_SCAN_LIMIT = 100;
+
+/**
+ * On resume, rebuild the model call that was open at `offset` by reading back to its start, so the reader continues
+ * exactly as an uninterrupted one would: a re-executed attempt is superseded, and step boundaries fall in the same place.
+ * Past RESTORE_SCAN_LIMIT unrelated records it gives up, and the resumed stream only lacks that step's `finish-step`.
+ */
+async function restoreOpenCall(client: DurableStreamSource, workflowID: string, key: string, state: ReaderState): Promise<void> {
+  let step: number | undefined;
+  let attempt: string | undefined;
+  // Set once an earlier step's record is reached: the open call is fully read, and only a finish reason is still wanted.
+  let callRead = false;
+  let budget = RESTORE_SCAN_LIMIT;
+  const ended = new Set<string>();
+  for (let index = state.offset - 1; index >= 0; index--) {
+    let record: DurableStreamRecord;
+    try {
+      record = await client.readStreamOffset<DurableStreamRecord>(workflowID, key, index, { timeoutSeconds: 0 });
+    } catch (error) {
+      // A real offset never passes the stored records, so a gap means there is nothing to restore.
+      if (DBOSErrors.isStreamTimeoutError(error)) break;
+      throw error;
+    }
+    if (record.kind === 'end') break;
+    const model = record.kind === 'model' || record.kind === 'model-end' ? record : undefined;
+    if (model && step === undefined) {
+      step = model.step;
+      attempt = model.attempt;
+    }
+    if (model && !callRead && model.step === step) {
+      if (model.attempt === attempt) restorePart(state, ended, model);
+      continue;
+    }
+    if (model) callRead = true;
+    // As in an uninterrupted read, the finish reason carries over from the latest earlier call that ended.
+    if (callRead && (state.finishReason !== undefined || model?.kind === 'model-end')) {
+      if (model?.kind === 'model-end') state.finishReason ??= finishReasonOf(model);
+      break;
+    }
+    if (--budget === 0) break;
+  }
+  if (step === undefined) return;
+  state.openStep = step;
+  state.openAttempt = attempt;
+}
+
+function finishReasonOf(record: Extract<DurableStreamRecord, { kind: 'model-end' }>): string | undefined {
+  return record.aborted ? 'other' : record.finishReason?.unified;
+}
+
+// One record of the open call, read backward: a part is open if its start comes before (i.e. is read after) any end.
+function restorePart(state: ReaderState, ended: Set<string>, record: Extract<DurableStreamRecord, { kind: 'model' | 'model-end' }>): void {
+  if (record.kind === 'model-end') {
+    state.finishReason ??= finishReasonOf(record);
+    return;
+  }
+  for (const part of [...record.parts].reverse()) {
+    const kind = part.type === 'text-start' || part.type === 'text-end' ? 'text' : part.type === 'reasoning-start' || part.type === 'reasoning-end' ? 'reasoning' : undefined;
+    if (kind === undefined || !('id' in part)) continue;
+    const id = `${record.attempt}:${part.id}`;
+    if (part.type.endsWith('-end')) ended.add(id);
+    else if (!ended.has(id)) {
+      state.openParts.set(id, kind);
+      state.detachedParts.add(id);
+    }
+  }
 }
 
 // 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 };
+  for (const [id, kind] of state.openParts) {
+    if (!state.detachedParts.has(id)) 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.detachedParts.clear();
   state.openAttempt = attempt;
 }
 
@@ -287,8 +382,7 @@ function* emitRecord(
     case 'model': {
       if (state.openStep !== record.step) {
         yield* closeStep(state);
-        if (!state.resumed) yield { type: 'start-step' };
-        state.resumed = false;
+        yield { type: 'start-step' };
         state.openStep = record.step;
         state.openAttempt = record.attempt;
       } else if (state.openAttempt !== record.attempt) {
@@ -307,14 +401,27 @@ function* emitRecord(
     }
     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;
+      if (record.attempt === state.openAttempt) state.finishReason = finishReasonOf(record);
       break;
-    case 'tool':
+    case 'tool': {
+      const key = toolKey(record);
+      const chunks = record.chunks ?? [];
+      const sent = state.toolChunksSent.get(key);
+      // A live re-execution that wrote the same chunks changes nothing; different ones (or none) replace what the client has.
+      if (chunks.length > 0 || sent !== undefined) {
+        const digest = createHash('sha256').update(JSON.stringify(chunks)).digest('base64');
+        if (sent !== digest) {
+          if (sent !== undefined) yield { type: 'data-dbos-tool-superseded', data: { toolCallId: record.toolCallId }, transient: true } as UIMessageChunk;
+          state.toolChunksSent.set(key, digest);
+          yield* chunks;
+        }
+      }
       // 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;

package/src/index.ts

@@ -1,6 +1,6 @@
 export { durableCalls, DurableCallsOptions, durableEmbeddingCalls, durableImageCalls } from './middleware';
 export { durableMCPTools, DurableMCPToolsOptions, MCPClientLike } from './mcp';
-export { durableTools, DurableToolsOptions } from './tools';
+export { durableTools, DurableToolsOptions, toolWriter } from './tools';
 export {
   closeDurableStream,
   DurableStreamOptions,

package/src/internal.ts

@@ -21,6 +21,31 @@ export function runDurableStep<T>(name: string, fn: () => Promise<T>, config: St
   });
 }
 
+// DBOS 5.1+ fires stepStatus.cancelSignal when the step's workflow is cancelled; older versions have none.
+export function stepCancelSignal(): AbortSignal | undefined {
+  return (DBOS.stepStatus as { cancelSignal?: AbortSignal } | undefined)?.cancelSignal;
+}
+
+/**
+ * DBOS discards a timed-out attempt's result, so a stream record of it would contradict the checkpoint. Returns the
+ * outcome to record instead: the timeout when the step ends with it, `null` when a retry follows, `undefined` if not timed out.
+ */
+export async function timedOutOutcome(shouldRetry: StepConfig['shouldRetry']): Promise<{ errorText: string } | null | undefined> {
+  const status = DBOS.stepStatus;
+  if (status?.timeoutSignal?.aborted !== true) return undefined;
+  const reason: unknown = status.timeoutSignal.reason;
+  const lastAttempt = status.currentAttempt === undefined || status.currentAttempt >= (status.maxAttempts ?? 1);
+  // DBOS asks the step's shouldRetry about this same error (the signal's reason) before retrying; a throw ends the step.
+  const retried = !lastAttempt && (await Promise.resolve(shouldRetry ? shouldRetry(reason) : true).catch(() => false));
+  return retried ? null : { errorText: reason instanceof Error ? reason.message : 'The step timed out.' };
+}
+
+// Fires when any given signal does; a lone signal is returned as-is.
+export function anySignal(...signals: (AbortSignal | undefined)[]): AbortSignal | undefined {
+  const defined = signals.filter((s): s is AbortSignal => s !== undefined);
+  return defined.length <= 1 ? defined[0] : AbortSignal.any(defined);
+}
+
 export function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
   return typeof (value as AsyncIterable<unknown> | null | undefined)?.[Symbol.asyncIterator] === 'function';
 }

package/src/mcp.ts

@@ -1,6 +1,6 @@
-import { StepConfig } from '@dbos-inc/dbos-sdk';
+import { DBOS, StepConfig } from '@dbos-inc/dbos-sdk';
 import type { ToolSet } from 'ai' with { 'resolution-mode': 'import' };
-import { isAsyncIterable, runDurableStep, withErrorClassification } from './internal';
+import { anySignal, isAsyncIterable, runDurableStep, stepCancelSignal, timedOutOutcome, 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
@@ -112,11 +112,20 @@ export async function durableMCPTools(client: MCPClientLike, options: DurableMCP
         return run(
           `mcp.tool.${name}.${toolCallId ?? 'call'}`,
           async () => {
+            // Stop the call when the attempt times out or the workflow is cancelled, as well as on the caller's abort.
+            const cancelSignal = stepCancelSignal();
+            const record = async (outcome: { output: unknown } | { errorText: string }) => {
+              if (!durableStream || !toolCallId) return;
+              const timeout = await timedOutOutcome(callConfig.shouldRetry);
+              if (timeout === null) return;
+              await writeToolRecord(durableStream, toolCallId, timeout ?? outcome);
+            };
+            const abortSignal = anySignal(signal, DBOS.stepStatus?.timeoutSignal, cancelSignal);
             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);
+              output = await tool.execute(input, abortSignal === signal ? execOptions : { ...execOptions, abortSignal });
               // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
               if (isAsyncIterable(output)) {
                 let last: unknown;
@@ -124,12 +133,10 @@ export async function durableMCPTools(client: MCPClientLike, options: DurableMCP
                 output = last;
               }
             } catch (error) {
-              if (durableStream && toolCallId) {
-                await writeToolRecord(durableStream, toolCallId, { errorText: error instanceof Error ? error.message : String(error) });
-              }
+              if (!cancelSignal?.aborted) await record({ errorText: error instanceof Error ? error.message : String(error) });
               throw error;
             }
-            if (durableStream && toolCallId) await writeToolRecord(durableStream, toolCallId, { output });
+            await record({ output });
             return output;
           },
           callConfig,

package/src/middleware.ts

@@ -20,7 +20,7 @@ import type {
   SharedV4ProviderMetadata,
   SharedV4Warning,
 } from '@ai-sdk/provider' with { 'resolution-mode': 'import' };
-import { assertNotInTransaction, isInWorkflowFunction, restoreAISDKErrorIdentity, withErrorClassification } from './internal';
+import { anySignal, assertNotInTransaction, isInWorkflowFunction, restoreAISDKErrorIdentity, stepCancelSignal, withErrorClassification } from './internal';
 import { type DurableStreamOptions, ModelStreamWriter, resolveDurableStream } from './durable-stream';
 
 export interface DurableCallsOptions extends StepConfig {
@@ -75,7 +75,10 @@ export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMi
       try {
         return await DBOS.runStep(
           async () => {
-            const result = ensureResponseMetadata(encodeBinaryContent(omitBodies(await doGenerate(), include)));
+            // The provider call also stops when the attempt times out or the workflow is cancelled.
+            const signal = anySignal(params.abortSignal, DBOS.stepStatus?.timeoutSignal, stepCancelSignal());
+            const raw = signal === params.abortSignal ? await doGenerate() : await model.doGenerate({ ...params, abortSignal: signal });
+            const result = ensureResponseMetadata(encodeBinaryContent(omitBodies(raw, include)));
             // A non-streaming call writes its whole output at once, so the stream holds every call the loop makes, not only streamed ones.
             if (streamConfig) {
               const streamWriter = new ModelStreamWriter(streamConfig);
@@ -157,20 +160,26 @@ export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMi
             const accumulator = new StreamAccumulator();
             // A timed-out attempt is abandoned by DBOS (its outcome is discarded) but keeps running; stop it so it can't emit alongside a retry.
             const timeoutSignal = DBOS.stepStatus?.timeoutSignal;
+            // A workflow cancel is not a consumer abort: it must fail the attempt so nothing is checkpointed and a resume re-runs the call.
+            const cancelSignal = stepCancelSignal();
+            const cancelledWorkflow = () => cancelSignal?.aborted === true;
+            const providerSignal = anySignal(abortSignal, timeoutSignal, cancelSignal);
             let reader: ReadableStreamDefaultReader<LanguageModelV4StreamPart> | undefined;
             let sawFinish = false;
             const abandon = () => void reader?.cancel().catch(() => {});
             timeoutSignal?.addEventListener('abort', abandon, { once: true });
+            cancelSignal?.addEventListener('abort', abandon, { once: true });
             // A consumer abort tears the provider call down too; the attempt is then recorded as aborted below.
             abortSignal?.addEventListener('abort', abandon, { once: true });
             // Step-scope stream writes: cheap, replay-safe, and never duplicated since a retry is refused once content has streamed.
             const streamWriter = streamConfig ? new ModelStreamWriter(streamConfig) : undefined;
             try {
-              reader = (await doStream()).stream.getReader();
+              const result = providerSignal === abortSignal ? await doStream() : await model.doStream({ ...params, abortSignal: providerSignal });
+              reader = result.stream.getReader();
               for (;;) {
                 const { done, value: part } = await reader.read();
                 if (timeoutSignal?.aborted) throw (timeoutSignal.reason ?? new Error('step attempt timed out'));
-                if (done || aborted()) break;
+                if (done || aborted() || cancelledWorkflow()) break;
                 if (part.type === 'error') {
                   // A cancelled consumer abandoned this call; don't let a late failure become the step outcome, or replay would fail where the live run succeeded.
                   if (cancelled) break;
@@ -183,21 +192,26 @@ export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMi
                 accumulator.add(part);
               }
               // No terminal part and no output: fail (retryably) like the AI SDK's NoOutputGeneratedError, instead of checkpointing a permanent empty success.
-              if (!sawFinish && !accumulator.hasContent && !cancelled && !aborted()) {
+              if (!sawFinish && !accumulator.hasContent && !cancelled && !aborted() && !cancelledWorkflow()) {
                 throw new Error('Model stream ended without a finish part or any output.');
               }
             } catch (error) {
-              // Same rule for stream-level failures (doStream or a read rejecting) after a cancel; after an abort, the abort is the outcome.
-              if (!cancelled && !aborted()) {
+              // Same rule for stream-level failures (doStream or a read rejecting) after a cancel; after an abort or a workflow cancel, that is the outcome.
+              if (!cancelled && !aborted() && !cancelledWorkflow()) {
                 await streamWriter?.abandon();
                 throw error;
               }
             } finally {
               timeoutSignal?.removeEventListener('abort', abandon);
+              cancelSignal?.removeEventListener('abort', abandon);
               abortSignal?.removeEventListener('abort', abandon);
               // Tear down the provider stream on early exits (error part, post-cancel break); a no-op after a clean drain.
               void reader?.cancel().catch(() => {});
             }
+            if (cancelledWorkflow()) {
+              await streamWriter?.abandon();
+              throw cancelSignal!.reason;
+            }
             // Record the abort as the step's failure; replay rethrows it, so the workflow must catch aborts it means to survive.
             if (aborted()) {
               await streamWriter?.end({ aborted: true });
@@ -272,14 +286,17 @@ export function durableEmbeddingCalls(options: StepConfig = {}): EmbeddingModelM
       then: (onfulfilled, onrejected) =>
         Promise.resolve(isInWorkflowFunction() ? false : model.supportsParallelCalls).then(onfulfilled, onrejected),
     }),
-    wrapEmbed: async ({ doEmbed, model }) => {
+    wrapEmbed: async ({ doEmbed, params, model }) => {
       assertNotInTransaction('embed');
       if (!isInWorkflowFunction()) {
         return await doEmbed();
       }
       const workflowID = await enterDurableModelCall();
       try {
-        return await DBOS.runStep(async () => doEmbed(), {
+        return await DBOS.runStep(async () => {
+          const signal = anySignal(params.abortSignal, DBOS.stepStatus?.timeoutSignal, stepCancelSignal());
+          return signal === params.abortSignal ? doEmbed() : model.doEmbed({ ...params, abortSignal: signal });
+        }, {
           ...stepConfig,
           name: stepConfig.name ?? stepName(model, 'embed'),
         });
@@ -301,13 +318,16 @@ export function durableImageCalls(options: StepConfig = {}): ImageModelMiddlewar
   const stepConfig = withErrorClassification(options);
   return {
     specificationVersion: 'v4',
-    wrapGenerate: async ({ doGenerate, model }) => {
+    wrapGenerate: async ({ doGenerate, params, model }) => {
       assertNotInTransaction('generateImage');
       if (!isInWorkflowFunction()) {
         return await doGenerate();
       }
       try {
-        return await DBOS.runStep(async () => encodeImageResult(await doGenerate()), {
+        return await DBOS.runStep(async () => {
+          const signal = anySignal(params.abortSignal, DBOS.stepStatus?.timeoutSignal, stepCancelSignal());
+          return encodeImageResult(signal === params.abortSignal ? await doGenerate() : await model.doGenerate({ ...params, abortSignal: signal }));
+        }, {
           ...stepConfig,
           name: stepConfig.name ?? stepName(model, 'image'),
         });

package/src/tools.ts

@@ -1,14 +1,108 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
 import { DBOS, StepConfig } from '@dbos-inc/dbos-sdk';
-import type { ToolSet } from 'ai' with { 'resolution-mode': 'import' };
-import { assertNotInTransaction, isAsyncIterable, isInWorkflowFunction, runDurableStep, withErrorClassification } from './internal';
+import type { ToolSet, UIMessageChunk, UIMessageStreamWriter } from 'ai' with { 'resolution-mode': 'import' };
+import { anySignal, assertNotInTransaction, isAsyncIterable, isInWorkflowFunction, runDurableStep, stepCancelSignal, timedOutOutcome, withErrorClassification } from './internal';
 import { AGENT_TOOL } from './agent-tool';
-import { writeToolRecord } from './durable-stream';
+import { writeDurableStream, writeToolRecord } from './durable-stream';
 
 export interface DurableToolsOptions extends StepConfig {
   /** Per-tool step config overriding the defaults; `false` leaves that tool non-durable. */
   tools?: Record<string, StepConfig | false>;
   /** Write each tool call's output (or error) to this durable stream from inside its step. */
   durableStream?: string;
+  /** Receives chunks tools write via `toolWriter()`; those that are part of the message are checkpointed and re-emitted on replay, before the tool's output. */
+  writer?: UIMessageStreamWriter;
+}
+
+const currentWriter = new AsyncLocalStorage<UIMessageStreamWriter>();
+
+/** Returns a UI message stream writer bound to the current `durableTools` tool call. */
+export function toolWriter(): UIMessageStreamWriter {
+  const writer = currentWriter.getStore();
+  if (!writer) throw new Error('toolWriter() can only be called from a tool wrapped by durableTools.');
+  return writer;
+}
+
+// Tool output plus the message chunks its call wrote, as checkpointed; a bare output is a call that wrote none.
+interface ToolEnvelope {
+  __dbosToolChunks: 1;
+  output: unknown;
+  chunks: UIMessageChunk[];
+}
+
+function isToolEnvelope(value: unknown): value is ToolEnvelope {
+  return typeof value === 'object' && value !== null && (value as Partial<ToolEnvelope>).__dbosToolChunks === 1;
+}
+
+function isTransient(chunk: UIMessageChunk): boolean {
+  return (chunk as { transient?: boolean }).transient === true;
+}
+
+// An async generator's body runs as it is iterated, outside run(); bind each step of the iteration to the writer.
+function runWithWriter(writer: UIMessageStreamWriter, fn: () => unknown): unknown {
+  const result = currentWriter.run(writer, fn);
+  if (!isAsyncIterable(result)) return result;
+  const iterator = currentWriter.run(writer, () => result[Symbol.asyncIterator]());
+  return {
+    [Symbol.asyncIterator]() {
+      return this;
+    },
+    next: (...args: [] | [unknown]) => currentWriter.run(writer, () => iterator.next(...args)),
+    return: (value?: unknown) => currentWriter.run(writer, () => iterator.return?.(value) ?? Promise.resolve({ done: true as const, value })),
+    throw: (error?: unknown) => currentWriter.run(writer, () => iterator.throw?.(error) ?? Promise.reject(error)),
+  };
+}
+
+// Outside a step: every chunk goes straight to the writer, or nowhere.
+function liveWriter(writer: UIMessageStreamWriter | undefined): UIMessageStreamWriter {
+  return writer ?? { write() {}, merge() {}, onError: undefined };
+}
+
+/** One attempt's writer: transient chunks go out live, the rest wait for the step to succeed. */
+class StepWriter implements UIMessageStreamWriter {
+  readonly chunks: UIMessageChunk[] = [];
+  private pending: Promise<unknown>[] = [];
+
+  constructor(
+    private readonly writer: UIMessageStreamWriter | undefined,
+    private readonly durableStream: string | undefined,
+  ) {}
+
+  get onError() {
+    return this.writer?.onError;
+  }
+
+  write(chunk: UIMessageChunk): void {
+    if (!isTransient(chunk)) {
+      this.chunks.push(chunk);
+      return;
+    }
+    this.writer?.write(chunk);
+    if (this.durableStream) this.track(writeDurableStream(this.durableStream, [chunk]));
+  }
+
+  merge(stream: ReadableStream<UIMessageChunk>): void {
+    this.track(
+      (async () => {
+        for await (const chunk of stream) this.write(chunk);
+      })(),
+    );
+  }
+
+  // Observed now so a failure after the tool throws is not an unhandled rejection; settle still sees it.
+  private track(promise: Promise<unknown>): void {
+    promise.catch(() => {});
+    this.pending.push(promise);
+  }
+
+  /** Waits for merges and live writes, including any started while waiting; a failed one fails the call. */
+  async settle(): Promise<void> {
+    while (this.pending.length > 0) {
+      const batch = this.pending;
+      this.pending = [];
+      await Promise.all(batch);
+    }
+  }
 }
 
 // Loose view of a tool's execute; the AI SDK validates input and supplies the options.
@@ -20,7 +114,8 @@ type ToolExecute = (input: unknown, options: { toolCallId: string; abortSignal?:
  * Retries are off by default (the AI SDK never retries tools); opt in per tool with `retriesAllowed`.
  */
 export function durableTools<TOOLS extends ToolSet>(tools: TOOLS, options: DurableToolsOptions = {}): TOOLS {
-  const { tools: perTool, durableStream, ...defaults } = options;
+  const { tools: perTool, durableStream, writer, ...defaults } = options;
+  const live = liveWriter(writer);
   const durable: ToolSet = {};
   for (const [name, definition] of Object.entries(tools)) {
     const override = perTool?.[name];
@@ -30,11 +125,18 @@ export function durableTools<TOOLS extends ToolSet>(tools: TOOLS, options: Durab
       durable[name] = durableStream ? bindAgentTool(durableStream) : definition;
       continue;
     }
-    if (typeof definition.execute !== 'function' || override === false) {
+    if (typeof definition.execute !== 'function') {
       durable[name] = definition;
       continue;
     }
     const execute = definition.execute as ToolExecute;
+    if (override === false) {
+      durable[name] = {
+        ...definition,
+        execute: (input: unknown, execOptions: Parameters<ToolExecute>[1]) => runWithWriter(live, () => execute(input, execOptions)),
+      } as ToolSet[string];
+      continue;
+    }
     const merged: StepConfig = { ...defaults, ...override };
     // Default classification (aborts and provider-declared non-retryable errors are terminal), but retries stay opt-in.
     const stepConfig: StepConfig = { ...withErrorClassification(merged), retriesAllowed: merged.retriesAllowed ?? false };
@@ -43,7 +145,7 @@ export function durableTools<TOOLS extends ToolSet>(tools: TOOLS, options: Durab
       ...definition,
       execute: (input: unknown, execOptions: Parameters<ToolExecute>[1]) => {
         assertNotInTransaction(name);
-        if (!isInWorkflowFunction()) return execute(input, execOptions);
+        if (!isInWorkflowFunction()) return runWithWriter(live, () => execute(input, execOptions));
         const signal = execOptions.abortSignal;
         // An aborted call is done whatever the failure looks like; a retry would re-run a cancelled side effect.
         const callConfig: StepConfig = {
@@ -54,27 +156,43 @@ export function durableTools<TOOLS extends ToolSet>(tools: TOOLS, options: Durab
         return runDurableStep(
           `${prefix}.${execOptions.toolCallId}`,
           async () => {
-            // A timed-out attempt is abandoned by DBOS but keeps running; forward its signal so the tool stops too.
-            const timeoutSignal = DBOS.stepStatus?.timeoutSignal;
-            const abortSignal = timeoutSignal && signal ? AbortSignal.any([signal, timeoutSignal]) : (timeoutSignal ?? signal);
+            // A timed-out attempt is abandoned by DBOS but keeps running, and a cancelled workflow's tool should stop: forward both signals.
+            const cancelSignal = stepCancelSignal();
+            const abortSignal = anySignal(signal, DBOS.stepStatus?.timeoutSignal, cancelSignal);
+            const stepWriter = new StepWriter(writer, durableStream);
+            const record = async (outcome: { output: unknown } | { errorText: string }, chunks?: UIMessageChunk[]) => {
+              if (!durableStream) return;
+              const timeout = await timedOutOutcome(callConfig.shouldRetry);
+              if (timeout === null) return;
+              await writeToolRecord(durableStream, execOptions.toolCallId, timeout ?? outcome, timeout ? [] : chunks);
+            };
             let output: unknown;
             try {
-              output = await execute(input, abortSignal === signal ? execOptions : { ...execOptions, abortSignal });
-              // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
-              if (isAsyncIterable(output)) {
+              output = await currentWriter.run(stepWriter, async () => {
+                const value = await execute(input, abortSignal === signal ? execOptions : { ...execOptions, abortSignal });
+                // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
+                if (!isAsyncIterable(value)) return value;
                 let last: unknown;
-                for await (last of output);
-                output = last;
-              }
+                for await (last of value);
+                return last;
+              });
+              await stepWriter.settle();
             } catch (error) {
-              if (durableStream) await writeToolRecord(durableStream, execOptions.toolCallId, { errorText: errorMessage(error) });
+              // A cancelled workflow's reader ends the turn with an abort; this call has no outcome of its own.
+              if (!cancelSignal?.aborted) await record({ errorText: errorMessage(error) });
               throw error;
             }
-            if (durableStream) await writeToolRecord(durableStream, execOptions.toolCallId, { output });
-            return output;
+            const chunks = stepWriter.chunks;
+            await record({ output }, chunks);
+            return chunks.length > 0 ? ({ __dbosToolChunks: 1, output, chunks } satisfies ToolEnvelope) : output;
           },
           callConfig,
-        );
+        ).then((result) => {
+          if (!isToolEnvelope(result)) return result;
+          // Runs on first execution and on replay alike, so the workflow's message gets the same parts either way.
+          for (const chunk of result.chunks) writer?.write(chunk);
+          return result.output;
+        });
       },
     } as ToolSet[string];
   }