@dbos-inc/vercel-ai 0.2.5 → 0.3.7

package/src/middleware.ts

@@ -21,24 +21,26 @@ import type {
   SharedV4Warning,
 } from '@ai-sdk/provider' with { 'resolution-mode': 'import' };
 import { assertNotInTransaction, isInWorkflowFunction, restoreAISDKErrorIdentity, 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;
+}
 
 // In-flight durable model calls per workflow; concurrent calls have a nondeterministic DBOS step order on replay, so we reject them.
 const inflightModelCalls = new Map<string, number>();
+// A stream call whose consumer detached (abort/cancel) but whose step is still settling; the next call in that workflow waits for it.
+const settlingModelCalls = new Map<string, Promise<void>>();
 
-// Consumer abort signals withheld from the provider inside a workflow, keyed by the transformed params wrapStream receives.
-const detachedAbortSignals = new WeakMap<object, AbortSignal>();
-
-function enterDurableModelCall(operation: 'generate' | 'stream' | 'embed'): string {
+async function enterDurableModelCall(): Promise<string> {
   const workflowID = DBOS.workflowID!;
+  // Sequence after a detached call's checkpoint, so steps are always recorded in the order they were started.
+  await settlingModelCalls.get(workflowID);
   const inflight = inflightModelCalls.get(workflowID) ?? 0;
   if (inflight > 0) {
-    // embedMany parallelizes its batches; maxParallelCalls: 1 serializes them deterministically. Other callers use child workflows.
-    const remedy =
-      operation === 'embed'
-        ? 'pass maxParallelCalls: 1 to embedMany, or run each call in its own child workflow with DBOS.startWorkflow'
-        : 'run each call in its own child workflow with DBOS.startWorkflow';
     throw new Error(
-      `Concurrent durable model calls in workflow "${workflowID}" are not supported because their step order is nondeterministic on replay; ${remedy}.`,
+      `Concurrent durable model calls in workflow "${workflowID}" are not supported because their step order is nondeterministic on replay; run each call in its own child workflow with DBOS.startWorkflow.`,
     );
   }
   inflightModelCalls.set(workflowID, inflight + 1);
@@ -55,35 +57,36 @@ 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: StepConfig = {}): LanguageModelMiddleware {
-  const stepConfig = withErrorClassification(options);
+export function durableCalls(options: DurableCallsOptions = {}): LanguageModelMiddleware {
+  const { durableStream, ...stepOptions } = options;
+  const stepConfig = withErrorClassification(stepOptions);
+  const streamConfig = resolveDurableStream(durableStream);
   return {
     specificationVersion: 'v4',
 
-    // A consumer abort (a stop button, an AI SDK timeout, a shutting-down worker) must not cut a durable stream short:
-    // the truncated output would be checkpointed as the step's permanent result and replayed on recovery as the whole
-    // answer. Withhold the signal from the provider so the call records a complete result; keep it here for the guard.
-    // Generate is left alone: an aborted doGenerate rejects (nothing partial to record), and the AI SDK's tool loop
-    // relies on that rejection to stop — stripping it there would let the loop keep calling the model after an abort.
-    transformParams: async ({ params, type }) => {
-      if (type !== 'stream' || !isInWorkflowFunction() || !params.abortSignal) {
-        return params;
-      }
-      const { abortSignal, ...rest } = params;
-      detachedAbortSignals.set(rest, abortSignal);
-      return rest;
-    },
-
-    wrapGenerate: async ({ doGenerate, model }) => {
+    wrapGenerate: async ({ doGenerate, params, model }) => {
       assertNotInTransaction('generate');
       if (!isInWorkflowFunction()) {
         return await doGenerate();
       }
-      const workflowID = enterDurableModelCall('generate');
+      const workflowID = await enterDurableModelCall();
       try {
-        return await DBOS.runStep(async () => ensureResponseMetadata(encodeBinaryContent(await doGenerate())), {
+        return await DBOS.runStep(
+          async () => {
+            const result = ensureResponseMetadata(encodeBinaryContent(await doGenerate()));
+            // 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);
+              for (const part of replayParts(result)) streamWriter.push(part);
+              await streamWriter.end({ finishReason: result.finishReason });
+            }
+            return result;
+          },
+          {
           ...stepConfig,
           name: stepConfig.name ?? stepName(model, 'generate'),
+          // An aborted call is never retried, whatever the provider's rejection looks like.
+          shouldRetry: async (error: unknown) => params.abortSignal?.aborted !== true && (await stepConfig.shouldRetry!(error)),
         });
       } catch (error) {
         // Restore the AI SDK error identity a replay revival strips, so the SDK's retry/catch logic behaves the same.
@@ -98,41 +101,36 @@ export function durableCalls(options: StepConfig = {}): LanguageModelMiddleware
       if (!isInWorkflowFunction()) {
         return await doStream();
       }
-      const workflowID = enterDurableModelCall('stream');
-      // The consumer's signal, detached from the provider call by transformParams (undefined outside a workflow).
-      const abortSignal = detachedAbortSignals.get(params);
-      // An aborted consumer is done with this call, like a cancelled one: post-abort failures must checkpoint as a (partial) success, or replay would fail where the live run ended gracefully.
+      const workflowID = await enterDurableModelCall();
+      // The caller's abortSignal merged with the AI SDK's timeouts; an abort stops the provider call and is recorded as the step's outcome.
+      const abortSignal = params.abortSignal;
       const aborted = () => abortSignal?.aborted === true;
 
-      // Free the concurrency guard as soon as the consumer detaches (abort/cancel), not only when the step
-      // settles: this call's step is already sequenced, so a sequential follow-up in the same workflow is
-      // deterministic on replay and must not be rejected as concurrent while this step drains in the background.
-      let guardReleased = false;
-      const releaseGuard = () => {
-        if (guardReleased) return;
-        guardReleased = true;
-        abortSignal?.removeEventListener('abort', releaseGuard);
-        exitDurableModelCall(workflowID);
-      };
-      abortSignal?.addEventListener('abort', releaseGuard, { once: true });
-
       let executed = false;
       let cancelled = false;
       let emittedLive = false;
+      let released = false;
       let controller!: ReadableStreamDefaultController<LanguageModelV4StreamPart>;
       let step!: Promise<LanguageModelV4GenerateResult>;
+      let settled!: Promise<void>;
+      // A detached consumer (abort/cancel) may issue a sequential follow-up while this step still settles: make that call wait for this checkpoint instead of refusing it as concurrent.
+      const detach = () => {
+        abortSignal?.removeEventListener('abort', detach);
+        if (!released) settlingModelCalls.set(workflowID, settled);
+      };
+      abortSignal?.addEventListener('abort', detach, { once: true });
       const stream = new ReadableStream<LanguageModelV4StreamPart>({
         start(c) {
           controller = c;
         },
-        // Await the step so an early cancel still blocks until the model result is checkpointed.
+        // Await the step so a direct cancel blocks until the checkpoint; through the AI SDK's pipeline a consumer's early exit returns sooner, hence the rule to drain or abort instead.
         async cancel() {
           cancelled = true;
-          releaseGuard();
+          detach();
           // Post-cancel failures normally checkpoint as a success; anything else (e.g. a failed checkpoint write) is only visible here.
-          await step.catch((error: unknown) =>
-            DBOS.logger.warn(`Durable model call step failed after consumer cancel: ${String(error)}`),
-          );
+          await step.catch((error: unknown) => {
+            if (!aborted()) DBOS.logger.warn(`Durable model call step failed after consumer cancel: ${String(error)}`);
+          });
         },
       });
       const emit = (part: LanguageModelV4StreamPart) => {
@@ -143,11 +141,11 @@ export function durableCalls(options: StepConfig = {}): LanguageModelMiddleware
         }
       };
 
-      // Once any output part has streamed live, a retry would re-stream from scratch and duplicate output, so stop retrying.
+      // Once any output part has streamed live, a retry would re-stream from scratch and duplicate output; an aborted call is never retried.
       const streamStepConfig: StepConfig = {
         ...stepConfig,
         shouldRetry: async (error: unknown) =>
-          !emittedLive && (stepConfig.shouldRetry ? await stepConfig.shouldRetry(error) : true),
+          !emittedLive && !aborted() && (stepConfig.shouldRetry ? await stepConfig.shouldRetry(error) : true),
       };
 
       try {
@@ -162,22 +160,26 @@ export function durableCalls(options: StepConfig = {}): LanguageModelMiddleware
             let sawFinish = false;
             const abandon = () => void reader?.cancel().catch(() => {});
             timeoutSignal?.addEventListener('abort', abandon, { once: true });
-            // A consumer abort deliberately does not abandon the drain: the step keeps reading so the checkpoint is complete.
+            // 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();
               for (;;) {
                 const { done, value: part } = await reader.read();
                 if (timeoutSignal?.aborted) throw (timeoutSignal.reason ?? new Error('step attempt timed out'));
-                if (done) break;
+                if (done || aborted()) break;
                 if (part.type === 'error') {
-                  // A cancelled or aborted 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 || aborted()) break;
+                  // 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;
                   throw toStepError(part.error);
                 }
                 // Stream deltas live but withhold 'finish' until the checkpoint is durable: the AI SDK runs tool calls (and their durable steps) on 'finish', which must not checkpoint before this model step.
                 if (part.type === 'finish') sawFinish = true;
                 else emit(part);
+                streamWriter?.push(part);
                 accumulator.add(part);
               }
               // No terminal part and no output: fail (retryably) like the AI SDK's NoOutputGeneratedError, instead of checkpointing a permanent empty success.
@@ -185,33 +187,55 @@ export function durableCalls(options: StepConfig = {}): LanguageModelMiddleware
                 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 or abort.
-              if (!cancelled && !aborted()) throw 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()) {
+                await streamWriter?.abandon();
+                throw error;
+              }
             } finally {
               timeoutSignal?.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(() => {});
             }
+            // 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 });
+              throw toAbortError(abortSignal!.reason);
+            }
             // Give the response a durable id/timestamp when the provider sent none, and emit it live (before the
             // withheld 'finish') so the SDK sees the same values live and on replay instead of a fresh fallback.
             // 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);
-            return encodeBinaryContent(accumulator.result(streamResult?.request, streamResult?.response));
+            const recorded = encodeBinaryContent(accumulator.result(streamResult?.request, streamResult?.response));
+            // 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;
           },
           { ...streamStepConfig, name: stepConfig.name ?? stepName(model, 'stream') },
         );
       } catch (error) {
         // runStep can throw synchronously (e.g. a shutdown race); don't leak the guard entry.
-        releaseGuard();
+        abortSignal?.removeEventListener('abort', detach);
+        exitDurableModelCall(workflowID);
         throw error;
       }
+      // Hold the guard until the step has settled, so the next call in this workflow always starts after this checkpoint.
+      settled = step.then(
+        () => undefined,
+        () => undefined,
+      ).then(() => {
+        released = true;
+        abortSignal?.removeEventListener('abort', detach);
+        exitDurableModelCall(workflowID);
+        if (settlingModelCalls.get(workflowID) === settled) settlingModelCalls.delete(workflowID);
+      });
 
       // Drive the returned stream from the settled step: a live run emits only the withheld 'finish' (deltas already
       // streamed); a recovered run synthesizes the whole stream from the checkpoint. Either way consumers finish
       // only after the result is durable.
-      void step
-        .then(
+      void step.then(
           (recorded) => {
             if (cancelled) return;
             if (executed) {
@@ -227,10 +251,11 @@ export function durableCalls(options: StepConfig = {}): LanguageModelMiddleware
             controller.close();
           },
           (error: unknown) => {
-            if (!cancelled) controller.error(restoreAISDKErrorIdentity(error));
+            if (cancelled) return;
+            // Live abort: surface the signal's own reason so the AI SDK takes its abort path; a replayed abort is an ordinary error.
+            controller.error(executed && aborted() ? (abortSignal?.reason ?? error) : restoreAISDKErrorIdentity(error));
           },
-        )
-        .finally(releaseGuard);
+        );
 
       return { stream };
     },
@@ -242,12 +267,17 @@ export function durableEmbeddingCalls(options: StepConfig = {}): EmbeddingModelM
   const stepConfig = withErrorClassification(options);
   return {
     specificationVersion: 'v4',
+    // embedMany awaits this per call: inside a workflow the batches run sequentially (deterministic step order on replay); elsewhere the model's own answer stands.
+    overrideSupportsParallelCalls: ({ model }) => ({
+      then: (onfulfilled, onrejected) =>
+        Promise.resolve(isInWorkflowFunction() ? false : model.supportsParallelCalls).then(onfulfilled, onrejected),
+    }),
     wrapEmbed: async ({ doEmbed, model }) => {
       assertNotInTransaction('embed');
       if (!isInWorkflowFunction()) {
         return await doEmbed();
       }
-      const workflowID = enterDurableModelCall('embed');
+      const workflowID = await enterDurableModelCall();
       try {
         return await DBOS.runStep(async () => doEmbed(), {
           ...stepConfig,
@@ -288,6 +318,12 @@ export function durableImageCalls(options: StepConfig = {}): ImageModelMiddlewar
   };
 }
 
+// The abort reason as a recordable Error (an AbortSignal's reason may be any value); the name is what the AI SDK's abort checks read.
+function toAbortError(reason: unknown): Error {
+  if (reason instanceof Error) return reason;
+  return Object.assign(new Error(String(reason)), { name: 'AbortError' });
+}
+
 /** Convert generated image bytes (Uint8Array) to base64 (spec-allowed) to keep checkpoints compact. */
 function encodeImageResult(result: ImageModelV4Result): ImageModelV4Result {
   const images = result.images as (string | Uint8Array)[];

package/src/tools.ts

@@ -0,0 +1,86 @@
+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 { AGENT_TOOL } from './agent-tool';
+import { 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;
+}
+
+// Loose view of a tool's execute; the AI SDK validates input and supplies the options.
+type ToolExecute = (input: unknown, options: { toolCallId: string; abortSignal?: AbortSignal }) => unknown;
+
+/**
+ * Wraps each tool's `execute` so that, inside a workflow, every tool call runs as a durable DBOS step named
+ * `<tool>.<toolCallId>` and replays from its checkpoint on recovery; outside a workflow tools run unchanged.
+ * 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 durable: ToolSet = {};
+  for (const [name, definition] of Object.entries(tools)) {
+    const override = perTool?.[name];
+    // An agent tool is a child workflow, not a step: leave it unwrapped, binding this durable stream to it.
+    const bindAgentTool = (definition as { [AGENT_TOOL]?: (key: string) => ToolSet[string] })[AGENT_TOOL];
+    if (bindAgentTool) {
+      durable[name] = durableStream ? bindAgentTool(durableStream) : definition;
+      continue;
+    }
+    if (typeof definition.execute !== 'function' || override === false) {
+      durable[name] = definition;
+      continue;
+    }
+    const execute = definition.execute as ToolExecute;
+    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 };
+    const prefix = stepConfig.name ?? name;
+    durable[name] = {
+      ...definition,
+      execute: (input: unknown, execOptions: Parameters<ToolExecute>[1]) => {
+        assertNotInTransaction(name);
+        if (!isInWorkflowFunction()) return 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 = {
+          ...stepConfig,
+          shouldRetry: async (error: unknown) => !signal?.aborted && (await stepConfig.shouldRetry!(error)),
+        };
+        // The tool call id comes from the checkpointed model result, so a reordered parallel step fails replay instead of swapping results.
+        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);
+            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)) {
+                let last: unknown;
+                for await (last of output);
+                output = last;
+              }
+            } catch (error) {
+              if (durableStream) await writeToolRecord(durableStream, execOptions.toolCallId, { errorText: errorMessage(error) });
+              throw error;
+            }
+            if (durableStream) await writeToolRecord(durableStream, execOptions.toolCallId, { output });
+            return output;
+          },
+          callConfig,
+        );
+      },
+    } as ToolSet[string];
+  }
+  return durable as TOOLS;
+}
+
+function errorMessage(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}