@dbos-inc/vercel-ai 0.2.5 → 0.3.7

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,388 @@
1
+ import { randomUUID } from 'node:crypto';
2
+ import { DBOS, Error as DBOSErrors, StatusString } from '@dbos-inc/dbos-sdk';
3
+ import type { UIMessageChunk } from 'ai' with { 'resolution-mode': 'import' };
4
+ import type { LanguageModelV4FinishReason, LanguageModelV4StreamPart } from '@ai-sdk/provider' with { 'resolution-mode': 'import' };
5
+
6
+ /** Durable stream config: the DBOS stream key, or the key plus batching limits for the model step's writes. */
7
+ export type DurableStreamOptions = string | { key: string; maxBatchParts?: number; maxBatchDelayMs?: number };
8
+
9
+ /** One DBOS stream value; the reader turns these into AI SDK UI message chunks. */
10
+ export type DurableStreamRecord =
11
+ | { kind: 'model'; step: number; attempt: string; parts: LanguageModelV4StreamPart[] }
12
+ | { kind: 'model-end'; step: number; attempt: string; finishReason?: LanguageModelV4FinishReason; aborted?: true }
13
+ | { kind: 'tool'; step: number; attempt: number; toolCallId: string; output?: unknown; errorText?: string }
14
+ | { kind: 'ui'; step?: number; attempt?: number; chunks: UIMessageChunk[] }
15
+ | { kind: 'end'; finishReason: string };
16
+
17
+ interface ResolvedDurableStream {
18
+ key: string;
19
+ maxBatchParts: number;
20
+ maxBatchDelayMs: number;
21
+ }
22
+
23
+ export function resolveDurableStream(options: DurableStreamOptions | undefined): ResolvedDurableStream | undefined {
24
+ if (options === undefined) return undefined;
25
+ const config = typeof options === 'string' ? { key: options } : options;
26
+ return { key: config.key, maxBatchParts: config.maxBatchParts ?? 20, maxBatchDelayMs: config.maxBatchDelayMs ?? 25 };
27
+ }
28
+
29
+ function stepInfo(): { step: number; attempt: number } {
30
+ return { step: DBOS.stepID ?? -1, attempt: DBOS.stepStatus?.currentAttempt ?? 1 };
31
+ }
32
+
33
+ // Parts the reader can render; framing, metadata and terminal parts are recorded elsewhere or not at all.
34
+ function isContentPart(part: LanguageModelV4StreamPart): boolean {
35
+ switch (part.type) {
36
+ case 'stream-start':
37
+ case 'response-metadata':
38
+ case 'finish':
39
+ case 'error':
40
+ case 'raw':
41
+ return false;
42
+ default:
43
+ return true;
44
+ }
45
+ }
46
+
47
+ // Bytes become base64 so a file part stays compact in Postgres.
48
+ function encodePart(part: LanguageModelV4StreamPart): LanguageModelV4StreamPart {
49
+ if (part.type === 'file' && part.data.type === 'data' && part.data.data instanceof Uint8Array) {
50
+ return { ...part, data: { type: 'data', data: Buffer.from(part.data.data).toString('base64') } };
51
+ }
52
+ return part;
53
+ }
54
+
55
+ // A transient write error (the SDK already retries offset conflicts) gets a few attempts before it fails the model call.
56
+ async function writeWithRetry(key: string, record: DurableStreamRecord): Promise<void> {
57
+ for (let attempt = 1; ; attempt++) {
58
+ try {
59
+ return await DBOS.writeStream(key, record);
60
+ } catch (error) {
61
+ if (attempt === 3) throw error;
62
+ await new Promise((resolve) => setTimeout(resolve, 100 * attempt));
63
+ }
64
+ }
65
+ }
66
+
67
+ /** Batches a live model step's parts into step-scope stream writes; nothing is written on replay because the step body does not run. */
68
+ export class ModelStreamWriter {
69
+ private pending: LanguageModelV4StreamPart[] = [];
70
+ private chain: Promise<void> = Promise.resolve();
71
+ private failure: unknown;
72
+ private timer: ReturnType<typeof setTimeout> | undefined;
73
+ private readonly step = DBOS.stepID ?? -1;
74
+ // Unique per execution of the step, so a recovered run's re-execution is distinguishable from the crashed one.
75
+ private readonly attempt = randomUUID();
76
+
77
+ constructor(private readonly config: ResolvedDurableStream) {}
78
+
79
+ push(part: LanguageModelV4StreamPart): void {
80
+ if (!isContentPart(part)) return;
81
+ this.pending.push(encodePart(part));
82
+ if (this.pending.length >= this.config.maxBatchParts) this.flush();
83
+ else this.timer ??= setTimeout(() => this.flush(), this.config.maxBatchDelayMs);
84
+ }
85
+
86
+ /** Flushes, records how the call ended, and resolves once every write is durable; a write that failed after retries fails the call here. */
87
+ async end(outcome: { finishReason: LanguageModelV4FinishReason } | { aborted: true }): Promise<void> {
88
+ this.flush();
89
+ this.write({ kind: 'model-end', step: this.step, attempt: this.attempt, ...outcome });
90
+ await this.chain;
91
+ if (this.failure !== undefined) throw this.failure;
92
+ }
93
+
94
+ /** 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. */
95
+ async abandon(): Promise<void> {
96
+ this.flush();
97
+ await this.chain;
98
+ }
99
+
100
+ private flush(): void {
101
+ if (this.timer !== undefined) clearTimeout(this.timer);
102
+ this.timer = undefined;
103
+ if (this.pending.length === 0) return;
104
+ const parts = this.pending;
105
+ this.pending = [];
106
+ this.write({ kind: 'model', step: this.step, attempt: this.attempt, parts });
107
+ }
108
+
109
+ // Every link has a handler, so a rejection can never sit unobserved; after one failure later writes are skipped.
110
+ private write(record: DurableStreamRecord): void {
111
+ this.chain = this.chain
112
+ .then(() => (this.failure === undefined ? writeWithRetry(this.config.key, record) : undefined))
113
+ .catch((error: unknown) => {
114
+ this.failure ??= error;
115
+ });
116
+ }
117
+ }
118
+
119
+ /** Records a tool call's outcome from inside its step. */
120
+ export function writeToolRecord(key: string, toolCallId: string, outcome: { output: unknown } | { errorText: string }): Promise<void> {
121
+ const record: DurableStreamRecord = { kind: 'tool', ...stepInfo(), toolCallId, ...outcome };
122
+ return DBOS.writeStream(key, record);
123
+ }
124
+
125
+ /**
126
+ * Appends UI message chunks to a durable stream. From a step the write is cheap and at-least-once, so give data parts
127
+ * stable ids; from workflow code it is a checkpointed step, so the number of calls must be deterministic.
128
+ */
129
+ export function writeDurableStream(key: string, chunks: UIMessageChunk[]): Promise<void> {
130
+ const status = DBOS.stepStatus;
131
+ const record: DurableStreamRecord = { kind: 'ui', step: status?.stepID, attempt: status?.currentAttempt, chunks };
132
+ return DBOS.writeStream(key, record);
133
+ }
134
+
135
+ /** 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. */
136
+ export async function closeDurableStream(key: string, finishReason = 'stop'): Promise<void> {
137
+ const record: DurableStreamRecord = { kind: 'end', finishReason };
138
+ await DBOS.writeStream(key, record);
139
+ await DBOS.closeStream(key);
140
+ }
141
+
142
+ /** What the reader needs from DBOS: the `DBOS` class in a launched process, or a `DBOSClient` anywhere else. */
143
+ export interface DurableStreamSource {
144
+ readStream<T>(workflowID: string, key: string, options?: { offset?: number }): AsyncGenerator<T, void, unknown>;
145
+ readStreamOffset<T>(workflowID: string, key: string, offset: number, options?: { timeoutSeconds?: number }): Promise<T>;
146
+ retrieveWorkflow(workflowID: string): { getStatus(): Promise<{ status: string; error?: unknown } | null> };
147
+ }
148
+
149
+ export interface ReadDurableStreamOptions {
150
+ workflowID: string;
151
+ key: string;
152
+ /** Id for the `start` chunk; omitted on a resume (`offset` > 0). */
153
+ messageId?: string;
154
+ /** Number of records already consumed, from the last `data-dbos-offset` chunk. */
155
+ offset?: number;
156
+ /** Defaults to `DBOS`; pass a `DBOSClient` to read from a process that has not launched DBOS. */
157
+ client?: DurableStreamSource;
158
+ /** Emit reasoning parts (default true, as in the AI SDK). */
159
+ sendReasoning?: boolean;
160
+ /** Emit source parts (default false, as in the AI SDK). */
161
+ sendSources?: boolean;
162
+ /** 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. */
163
+ onError?: (error: unknown) => string;
164
+ }
165
+
166
+ /** Reads a durable stream as AI SDK UI message chunks, live or after the fact, resuming from `offset`. */
167
+ export function readDurableStream(options: ReadDurableStreamOptions): ReadableStream<UIMessageChunk> {
168
+ const iterator = uiChunks(options);
169
+ return new ReadableStream<UIMessageChunk>({
170
+ async pull(controller) {
171
+ const { value, done } = await iterator.next();
172
+ if (done) controller.close();
173
+ else controller.enqueue(value);
174
+ },
175
+ async cancel() {
176
+ await iterator.return(undefined);
177
+ },
178
+ });
179
+ }
180
+
181
+ async function* uiChunks(options: ReadDurableStreamOptions): AsyncGenerator<UIMessageChunk> {
182
+ const { workflowID, key, sendReasoning = true, sendSources = false, onError = () => 'An error occurred.' } = options;
183
+ const client: DurableStreamSource = options.client ?? DBOS;
184
+ const state: ReaderState = {
185
+ offset: options.offset ?? 0,
186
+ resumed: (options.offset ?? 0) > 0,
187
+ openParts: new Map(),
188
+ ended: false,
189
+ };
190
+ if (state.offset === 0) yield { type: 'start', messageId: options.messageId };
191
+ const emit = (record: DurableStreamRecord) => emitRecord(state, record, { sendReasoning, sendSources, onError });
192
+
193
+ // Phase 1: everything already stored, one value per query until an offset is empty; a superseded attempt is skipped whole.
194
+ const history: DurableStreamRecord[] = [];
195
+ for (;;) {
196
+ try {
197
+ history.push(await client.readStreamOffset<DurableStreamRecord>(workflowID, key, state.offset + history.length, { timeoutSeconds: 0 }));
198
+ } catch (error) {
199
+ if (!DBOSErrors.isStreamTimeoutError(error)) throw error;
200
+ break;
201
+ }
202
+ }
203
+ const finalAttempt = new Map<number, string>();
204
+ for (const record of history) {
205
+ if (record.kind === 'model' || record.kind === 'model-end') finalAttempt.set(record.step, record.attempt);
206
+ }
207
+ for (const record of history) {
208
+ const stale = (record.kind === 'model' || record.kind === 'model-end') && finalAttempt.get(record.step) !== record.attempt;
209
+ yield* stale ? skipRecord(state) : emit(record);
210
+ if (state.ended) return;
211
+ }
212
+
213
+ // Phase 2: live; a re-executed step shows up as a new attempt and is handed off in place.
214
+ for await (const record of client.readStream<DurableStreamRecord>(workflowID, key, { offset: state.offset })) {
215
+ yield* emit(record);
216
+ if (state.ended) return;
217
+ }
218
+
219
+ // No end record: the workflow's status decides how the turn ended (a stream closed while it still runs counts as finished).
220
+ const status = (await client.retrieveWorkflow(workflowID).getStatus())?.status;
221
+ yield* closeStep(state);
222
+ if (status === StatusString.CANCELLED) {
223
+ yield { type: 'abort' };
224
+ } else if (status === undefined || status === StatusString.SUCCESS || status === StatusString.PENDING || status === StatusString.ENQUEUED) {
225
+ // The AI SDK's finish schema has no 'unknown'; a turn with no model call ends with the reason omitted.
226
+ yield state.finishReason === undefined ? { type: 'finish' } : { type: 'finish', finishReason: state.finishReason as UIFinishReason };
227
+ } else {
228
+ const error = (await client.retrieveWorkflow(workflowID).getStatus())?.error;
229
+ yield { type: 'error', errorText: onError(error ?? new Error('The workflow ended before the response completed.')) };
230
+ yield { type: 'finish', finishReason: 'error' };
231
+ }
232
+ }
233
+
234
+ interface ReaderState {
235
+ offset: number;
236
+ resumed: boolean;
237
+ openStep?: number;
238
+ openAttempt?: string;
239
+ // Text/reasoning parts of the open attempt that have started but not ended, by UI part id.
240
+ openParts: Map<string, 'text' | 'reasoning'>;
241
+ finishReason?: string;
242
+ ended: boolean;
243
+ }
244
+
245
+ function offsetChunk(state: ReaderState): UIMessageChunk {
246
+ return { type: 'data-dbos-offset', data: { offset: state.offset }, transient: true } as UIMessageChunk;
247
+ }
248
+
249
+ function* skipRecord(state: ReaderState): Generator<UIMessageChunk> {
250
+ state.offset += 1;
251
+ yield offsetChunk(state);
252
+ }
253
+
254
+ function* closeStep(state: ReaderState): Generator<UIMessageChunk> {
255
+ if (state.openStep !== undefined) yield { type: 'finish-step' };
256
+ state.openStep = undefined;
257
+ state.openAttempt = undefined;
258
+ state.openParts.clear();
259
+ }
260
+
261
+ // A live re-execution of the open step: end the stale attempt's parts and tell the client which ones to discard.
262
+ function* supersede(state: ReaderState, attempt: string): Generator<UIMessageChunk> {
263
+ for (const [id, kind] of state.openParts) yield { type: kind === 'text' ? 'text-end' : 'reasoning-end', id };
264
+ yield {
265
+ type: 'data-dbos-superseded',
266
+ data: { attempt: state.openAttempt, parts: [...state.openParts.keys()] },
267
+ transient: true,
268
+ } as UIMessageChunk;
269
+ state.openParts.clear();
270
+ state.openAttempt = attempt;
271
+ }
272
+
273
+ function* emitRecord(
274
+ state: ReaderState,
275
+ record: DurableStreamRecord,
276
+ filter: { sendReasoning: boolean; sendSources: boolean; onError: (error: unknown) => string },
277
+ ): Generator<UIMessageChunk> {
278
+ state.offset += 1;
279
+ switch (record.kind) {
280
+ case 'model': {
281
+ if (state.openStep !== record.step) {
282
+ yield* closeStep(state);
283
+ if (!state.resumed) yield { type: 'start-step' };
284
+ state.resumed = false;
285
+ state.openStep = record.step;
286
+ state.openAttempt = record.attempt;
287
+ } else if (state.openAttempt !== record.attempt) {
288
+ yield* supersede(state, record.attempt);
289
+ }
290
+ for (const part of record.parts) {
291
+ if (!filter.sendReasoning && part.type.startsWith('reasoning-')) continue;
292
+ if (!filter.sendSources && part.type === 'source') continue;
293
+ const chunk = toUIChunk(part, record.attempt);
294
+ if (!chunk) continue;
295
+ if (chunk.type === 'text-start' || chunk.type === 'reasoning-start') state.openParts.set(chunk.id, chunk.type === 'text-start' ? 'text' : 'reasoning');
296
+ if (chunk.type === 'text-end' || chunk.type === 'reasoning-end') state.openParts.delete(chunk.id);
297
+ yield chunk;
298
+ }
299
+ break;
300
+ }
301
+ case 'model-end':
302
+ // The stream outlives the call: the workflow may run more calls, so only its end (or closeDurableStream) ends the turn.
303
+ if (record.attempt === state.openAttempt) state.finishReason = record.aborted ? 'other' : record.finishReason?.unified;
304
+ break;
305
+ case 'tool':
306
+ // Local tool errors are masked like the AI SDK does; provider-executed ones (in model records) pass through verbatim.
307
+ yield record.errorText !== undefined
308
+ ? { type: 'tool-output-error', toolCallId: record.toolCallId, errorText: filter.onError(new Error(record.errorText)) }
309
+ : { type: 'tool-output-available', toolCallId: record.toolCallId, output: record.output };
310
+ break;
311
+ case 'ui':
312
+ yield* record.chunks;
313
+ break;
314
+ case 'end':
315
+ // The terminal chunk comes last, so the offset goes out first.
316
+ yield offsetChunk(state);
317
+ yield* closeStep(state);
318
+ yield { type: 'finish', finishReason: record.finishReason as UIFinishReason };
319
+ state.ended = true;
320
+ return;
321
+ }
322
+ yield offsetChunk(state);
323
+ }
324
+
325
+ type UIFinishReason = Extract<UIMessageChunk, { type: 'finish' }>['finishReason'];
326
+
327
+ function parseInput(input: string): unknown {
328
+ try {
329
+ return JSON.parse(input);
330
+ } catch {
331
+ return input;
332
+ }
333
+ }
334
+
335
+ // Text and reasoning ids are only unique within one model call; the attempt id keeps calls, and re-executions, apart in one message.
336
+ function toUIChunk(part: LanguageModelV4StreamPart, attempt: string): UIMessageChunk | undefined {
337
+ const id = 'id' in part ? `${attempt}:${part.id}` : '';
338
+ switch (part.type) {
339
+ case 'text-start':
340
+ return { type: 'text-start', id, providerMetadata: part.providerMetadata };
341
+ case 'text-delta':
342
+ return { type: 'text-delta', id, delta: part.delta, providerMetadata: part.providerMetadata };
343
+ case 'text-end':
344
+ return { type: 'text-end', id, providerMetadata: part.providerMetadata };
345
+ case 'reasoning-start':
346
+ return { type: 'reasoning-start', id, providerMetadata: part.providerMetadata };
347
+ case 'reasoning-delta':
348
+ return { type: 'reasoning-delta', id, delta: part.delta, providerMetadata: part.providerMetadata };
349
+ case 'reasoning-end':
350
+ return { type: 'reasoning-end', id, providerMetadata: part.providerMetadata };
351
+ case 'tool-input-start':
352
+ return {
353
+ type: 'tool-input-start',
354
+ toolCallId: part.id,
355
+ toolName: part.toolName,
356
+ providerExecuted: part.providerExecuted,
357
+ dynamic: part.dynamic,
358
+ title: part.title,
359
+ providerMetadata: part.providerMetadata,
360
+ };
361
+ case 'tool-input-delta':
362
+ return { type: 'tool-input-delta', toolCallId: part.id, inputTextDelta: part.delta };
363
+ case 'tool-call':
364
+ return {
365
+ type: 'tool-input-available',
366
+ toolCallId: part.toolCallId,
367
+ toolName: part.toolName,
368
+ input: parseInput(part.input),
369
+ providerExecuted: part.providerExecuted,
370
+ dynamic: part.dynamic,
371
+ providerMetadata: part.providerMetadata,
372
+ };
373
+ case 'tool-result':
374
+ return part.isError
375
+ ? { type: 'tool-output-error', toolCallId: part.toolCallId, errorText: JSON.stringify(part.result), providerExecuted: true, dynamic: part.dynamic }
376
+ : { type: 'tool-output-available', toolCallId: part.toolCallId, output: part.result, providerExecuted: true, dynamic: part.dynamic, preliminary: part.preliminary };
377
+ case 'source':
378
+ return part.sourceType === 'url'
379
+ ? { type: 'source-url', sourceId: part.id, url: part.url, title: part.title, providerMetadata: part.providerMetadata }
380
+ : { type: 'source-document', sourceId: part.id, mediaType: part.mediaType, title: part.title, filename: part.filename, providerMetadata: part.providerMetadata };
381
+ case 'file': {
382
+ const url = part.data.type === 'url' ? String(part.data.url) : part.data.type === 'data' ? `data:${part.mediaType};base64,${String(part.data.data)}` : undefined;
383
+ return url === undefined ? undefined : { type: 'file', url, mediaType: part.mediaType, providerMetadata: part.providerMetadata };
384
+ }
385
+ default:
386
+ return undefined;
387
+ }
388
+ }
package/src/index.ts CHANGED
@@ -1,2 +1,13 @@
1
- export { durableCalls, durableEmbeddingCalls, durableImageCalls } from './middleware';
1
+ export { durableCalls, DurableCallsOptions, durableEmbeddingCalls, durableImageCalls } from './middleware';
2
2
  export { durableMCPTools, DurableMCPToolsOptions, MCPClientLike } from './mcp';
3
+ export { durableTools, DurableToolsOptions } from './tools';
4
+ export {
5
+ closeDurableStream,
6
+ DurableStreamOptions,
7
+ DurableStreamRecord,
8
+ DurableStreamSource,
9
+ ReadDurableStreamOptions,
10
+ readDurableStream,
11
+ writeDurableStream,
12
+ } from './durable-stream';
13
+ export { agentTool, AgentTool, AgentToolOptions } from './agent-tool';
package/src/internal.ts CHANGED
@@ -11,6 +11,20 @@ export function assertNotInTransaction(operation: string): void {
11
11
  }
12
12
  }
13
13
 
14
+ // Run fn as a durable step named `name` inside a workflow; elsewhere call it directly.
15
+ export function runDurableStep<T>(name: string, fn: () => Promise<T>, config: StepConfig): Promise<T> {
16
+ assertNotInTransaction(name);
17
+ if (!isInWorkflowFunction()) return fn();
18
+ // Restore the AI SDK error identity a replay revival strips, so the SDK's retry/catch logic behaves the same.
19
+ return DBOS.runStep(fn, { ...config, name }).catch((error: unknown) => {
20
+ throw restoreAISDKErrorIdentity(error);
21
+ });
22
+ }
23
+
24
+ export function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
25
+ return typeof (value as AsyncIterable<unknown> | null | undefined)?.[Symbol.asyncIterator] === 'function';
26
+ }
27
+
14
28
  // Aborts/timeouts are deliberate cancellations, never transient; retrying just re-runs an already-cancelled call.
15
29
  function isAbortError(error: unknown): boolean {
16
30
  const name = (error as { name?: unknown } | null)?.name;
package/src/mcp.ts CHANGED
@@ -1,6 +1,7 @@
1
- import { DBOS, StepConfig } from '@dbos-inc/dbos-sdk';
1
+ import { StepConfig } from '@dbos-inc/dbos-sdk';
2
2
  import type { ToolSet } from 'ai' with { 'resolution-mode': 'import' };
3
- import { assertNotInTransaction, isInWorkflowFunction, restoreAISDKErrorIdentity, withErrorClassification } from './internal';
3
+ import { isAsyncIterable, runDurableStep, withErrorClassification } from './internal';
4
+ import { writeToolRecord } from './durable-stream';
4
5
 
5
6
  // Structural type for an MCP client (e.g. from @ai-sdk/mcp) — deliberately loose: the AI SDK ecosystem
6
7
  // exact-pins @ai-sdk/provider-utils, so precise Tool types fail to match across skewed copies.
@@ -23,6 +24,8 @@ interface MCPToolLike {
23
24
  export interface DurableMCPToolsOptions extends StepConfig {
24
25
  /** Forwarded to client.tools() on listing and on each call (e.g. { schemas } for subsetting and output schemas). */
25
26
  toolOptions?: unknown;
27
+ /** Write each tool call's output (or error) to this durable stream from inside its step. */
28
+ durableStream?: string;
26
29
  }
27
30
 
28
31
  interface DurableToolDef {
@@ -36,10 +39,6 @@ interface DurableToolDef {
36
39
 
37
40
  type ToolModelOutput = Awaited<ReturnType<NonNullable<ToolSet[string]['toModelOutput']>>>;
38
41
 
39
- function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
40
- return typeof (value as AsyncIterable<unknown> | null | undefined)?.[Symbol.asyncIterator] === 'function';
41
- }
42
-
43
42
  // Mirror of @ai-sdk/mcp's toModelOutput: MCP content becomes model content (text stays text, images become files).
44
43
  function mcpToolOutput(output: unknown): ToolModelOutput {
45
44
  const result = output as { content?: unknown };
@@ -66,17 +65,11 @@ function mcpToolOutput(output: unknown): ToolModelOutput {
66
65
  * checkpointed so a recovered workflow replays results instead of re-invoking the tool.
67
66
  */
68
67
  export async function durableMCPTools(client: MCPClientLike, options: DurableMCPToolsOptions = {}): Promise<ToolSet> {
69
- const { toolOptions, ...stepOptions } = options;
68
+ const { toolOptions, durableStream, ...stepOptions } = options;
70
69
  const stepConfig = withErrorClassification(stepOptions);
71
70
  const { asSchema, dynamicTool, jsonSchema } = await import('ai');
72
- const run = <T>(name: string, fn: () => Promise<T>, config: StepConfig = stepConfig): Promise<T> => {
73
- assertNotInTransaction(name);
74
- if (!isInWorkflowFunction()) return fn();
75
- // Restore the AI SDK error identity a replay revival strips, so the SDK's retry/catch logic behaves the same.
76
- return DBOS.runStep(fn, { ...config, name }).catch((error: unknown) => {
77
- throw restoreAISDKErrorIdentity(error);
78
- });
79
- };
71
+ const run = <T>(name: string, fn: () => Promise<T>, config: StepConfig = stepConfig): Promise<T> =>
72
+ runDurableStep(name, fn, config);
80
73
 
81
74
  // Checkpoint the tool list as plain JSON schemas, so replay reconstructs tools without the live client.
82
75
  const listed = await run('mcp.listTools', async () => {
@@ -108,25 +101,35 @@ export async function durableMCPTools(client: MCPClientLike, options: DurableMCP
108
101
  toModelOutput: def.convertsOutput ? ({ output }) => mcpToolOutput(output) : undefined,
109
102
  // Re-fetch the live tool inside the step (its execute closure can't be checkpointed); replay returns the recorded result.
110
103
  execute: (input: unknown, execOptions) => {
111
- const signal = (execOptions as { abortSignal?: AbortSignal } | undefined)?.abortSignal;
104
+ const { abortSignal: signal, toolCallId } = (execOptions ?? {}) as { abortSignal?: AbortSignal; toolCallId?: string };
112
105
  // An aborted consumer is done with this call, whatever the failure looks like; a retry would re-run a cancelled side effect.
113
106
  const callConfig: StepConfig = {
114
107
  ...stepConfig,
115
108
  shouldRetry: async (error: unknown) =>
116
109
  !signal?.aborted && (stepConfig.shouldRetry ? await stepConfig.shouldRetry(error) : true),
117
110
  };
111
+ // The tool call id comes from the checkpointed model result, so a reordered parallel step fails replay instead of swapping results.
118
112
  return run(
119
- `mcp.tool.${name}`,
113
+ `mcp.tool.${name}.${toolCallId ?? 'call'}`,
120
114
  async () => {
121
- const tool = (await client.tools(toolOptions))[name] as MCPToolLike | undefined;
122
- if (typeof tool?.execute !== 'function') throw new Error(`MCP tool "${name}" is not executable.`);
123
- const output = await tool.execute(input, execOptions);
124
- // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
125
- if (isAsyncIterable(output)) {
126
- let last: unknown;
127
- for await (last of output);
128
- return last;
115
+ let output: unknown;
116
+ try {
117
+ const tool = (await client.tools(toolOptions))[name] as MCPToolLike | undefined;
118
+ if (typeof tool?.execute !== 'function') throw new Error(`MCP tool "${name}" is not executable.`);
119
+ output = await tool.execute(input, execOptions);
120
+ // A streaming execute can't checkpoint mid-flight; drain it and record the final value (the last yield).
121
+ if (isAsyncIterable(output)) {
122
+ let last: unknown;
123
+ for await (last of output);
124
+ output = last;
125
+ }
126
+ } catch (error) {
127
+ if (durableStream && toolCallId) {
128
+ await writeToolRecord(durableStream, toolCallId, { errorText: error instanceof Error ? error.message : String(error) });
129
+ }
130
+ throw error;
129
131
  }
132
+ if (durableStream && toolCallId) await writeToolRecord(durableStream, toolCallId, { output });
130
133
  return output;
131
134
  },
132
135
  callConfig,