@dbos-inc/vercel-ai 0.3.7 → 0.5.4

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];
   }