@dbos-inc/vercel-ai 0.3.7 → 0.5.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dbos-inc/vercel-ai",
-  "version": "0.3.7",
+  "version": "0.5.4",
   "description": "DBOS durable execution integration for the Vercel AI SDK",
   "license": "MIT",
   "type": "commonjs",
@@ -44,14 +44,14 @@
     "agent"
   ],
   "peerDependencies": {
-    "@dbos-inc/dbos-sdk": "^4.21.0 || ^5.0.0",
+    "@dbos-inc/dbos-sdk": "^4.27.0 || ^5.0.0",
     "ai": "^7.0.0"
   },
   "devDependencies": {
     "@ai-sdk/gateway": "^4.0.0",
     "@ai-sdk/provider": "^4.0.0",
     "@arethetypeswrong/cli": "^0.18.0",
-    "@dbos-inc/dbos-sdk": "^5.0.2",
+    "@dbos-inc/dbos-sdk": "^5.2.11",
     "@types/node": "^24.0.0",
     "@types/pg": "^8.20.0",
     "ai": "^7.0.0",

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 };
 
@@ -30,17 +30,24 @@ function stepInfo(): { step: number; attempt: number } {
   return { step: DBOS.stepID ?? -1, attempt: DBOS.stepStatus?.currentAttempt ?? 1 };
 }
 
-// Parts the reader can render; framing, metadata and terminal parts are recorded elsewhere or not at all.
+// Only the parts toUIChunk renders; anything else would be written and never read.
 function isContentPart(part: LanguageModelV4StreamPart): boolean {
   switch (part.type) {
-    case 'stream-start':
-    case 'response-metadata':
-    case 'finish':
-    case 'error':
-    case 'raw':
-      return false;
-    default:
+    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;
   }
 }
 
@@ -116,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);
 }
 
@@ -151,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;
@@ -183,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.
@@ -201,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;
   }
 
@@ -233,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;
 }
@@ -256,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;
 }
 
@@ -280,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) {
@@ -300,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,12 +20,14 @@ 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 {
   /** Write each streamed model call's parts to this durable stream from inside its step (see readDurableStream). */
   durableStream?: DurableStreamOptions;
+  /** Checkpoint a generate call's raw provider request/response bodies; set to match generateText's `include` (both default false). */
+  include?: { requestBody?: boolean; responseBody?: boolean };
 }
 
 // In-flight durable model calls per workflow; concurrent calls have a nondeterministic DBOS step order on replay, so we reject them.
@@ -58,7 +60,7 @@ function exitDurableModelCall(workflowID: string): void {
 
 /** AI SDK language-model middleware that runs each model call as a durable, checkpointed DBOS step (replayed on recovery); outside a workflow it calls the model directly. */
 export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMiddleware {
-  const { durableStream, ...stepOptions } = options;
+  const { durableStream, include, ...stepOptions } = options;
   const stepConfig = withErrorClassification(stepOptions);
   const streamConfig = resolveDurableStream(durableStream);
   return {
@@ -73,7 +75,10 @@ export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMi
       try {
         return await DBOS.runStep(
           async () => {
-            const result = ensureResponseMetadata(encodeBinaryContent(await doGenerate()));
+            // 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);
@@ -155,22 +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;
-            let streamResult: Awaited<ReturnType<typeof doStream>> | undefined;
+            // 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 {
-              streamResult = await doStream();
-              reader = streamResult.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 });
@@ -208,7 +222,7 @@ export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMi
             // Skip the live emit for a timed-out (abandoned) attempt so it can't interleave with its retry.
             const responseMetadataPart = accumulator.fillResponseMetadata(randomUUID(), new Date());
             if (responseMetadataPart && !timeoutSignal?.aborted) emit(responseMetadataPart);
-            const recorded = encodeBinaryContent(accumulator.result(streamResult?.request, streamResult?.response));
+            const recorded = encodeBinaryContent(accumulator.result());
             // Every stream write lands before the checkpoint, so a reader that sees the next step has seen all of this one.
             await streamWriter?.end({ finishReason: recorded.finishReason });
             return recorded;
@@ -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'),
         });
@@ -371,6 +391,21 @@ function ensureResponseMetadata(result: LanguageModelV4GenerateResult): Language
   };
 }
 
+// generateText discards the raw bodies unless `include` asks for them, and the request body repeats the whole prompt.
+function omitBodies(
+  result: LanguageModelV4GenerateResult,
+  include: DurableCallsOptions['include'],
+): LanguageModelV4GenerateResult {
+  const dropRequest = !include?.requestBody && result.request?.body !== undefined;
+  const dropResponse = !include?.responseBody && result.response?.body !== undefined;
+  if (!dropRequest && !dropResponse) return result;
+  return {
+    ...result,
+    request: dropRequest ? { ...result.request, body: undefined } : result.request,
+    response: dropResponse ? { ...result.response, body: undefined } : result.response,
+  };
+}
+
 /** Convert generated-file bytes (Uint8Array) to base64 (spec-allowed) to keep checkpoints compact. */
 function encodeBinaryContent(result: LanguageModelV4GenerateResult): LanguageModelV4GenerateResult {
   const content = result.content.map(encodeBinaryPart);
@@ -515,18 +550,15 @@ class StreamAccumulator {
     return block;
   }
 
-  result(
-    request?: { body?: unknown },
-    response?: { headers?: Record<string, string> },
-  ): LanguageModelV4GenerateResult {
+  // Omits the provider's request/response (wrapStream returns only the stream): the request body repeats the whole prompt.
+  result(): LanguageModelV4GenerateResult {
     return {
       content: this.content,
       finishReason: this.finishReason,
       usage: this.usage,
       warnings: this.warnings,
       providerMetadata: this.providerMetadata,
-      request,
-      response: this.responseMetadata ? { ...this.responseMetadata, ...response } : response,
+      response: this.responseMetadata,
     };
   }
 }