fix (ai/core): do not construct object promise in streamObject result until requested (#2321)

Author: lgrammel

.changeset/long-crews-watch.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ai/core): do not construct object promise in streamObject result until requested

packages/core/core/generate-object/stream-object.test.ts

@@ -423,18 +423,6 @@ describe('result.object', () => {
     const result = await streamObject({
       model: new MockLanguageModelV1({
         doStream: async ({ prompt, mode }) => {
-          assert.deepStrictEqual(mode, { type: 'object-json' });
-          assert.deepStrictEqual(prompt, [
-            {
-              role: 'system',
-              content:
-                'JSON schema:\n' +
-                '{"type":"object","properties":{"content":{"type":"string"}},"required":["content"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}\n' +
-                'You MUST answer with a JSON object that matches the JSON schema above.',
-            },
-            { role: 'user', content: [{ type: 'text', text: 'prompt' }] },
-          ]);
-
           return {
             stream: convertArrayToReadableStream([
               { type: 'text-delta', textDelta: '{ ' },
@@ -470,18 +458,6 @@ describe('result.object', () => {
     const result = await streamObject({
       model: new MockLanguageModelV1({
         doStream: async ({ prompt, mode }) => {
-          assert.deepStrictEqual(mode, { type: 'object-json' });
-          assert.deepStrictEqual(prompt, [
-            {
-              role: 'system',
-              content:
-                'JSON schema:\n' +
-                '{"type":"object","properties":{"content":{"type":"string"}},"required":["content"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}\n' +
-                'You MUST answer with a JSON object that matches the JSON schema above.',
-            },
-            { role: 'user', content: [{ type: 'text', text: 'prompt' }] },
-          ]);
-
           return {
             stream: convertArrayToReadableStream([
               { type: 'text-delta', textDelta: '{ ' },
@@ -516,6 +492,39 @@ describe('result.object', () => {
         expect(TypeValidationError.isTypeValidationError(error)).toBeTruthy();
       });
   });
+
+  it('should not lead to unhandled promise rejections when the streamed object does not match the schema', async () => {
+    const result = await streamObject({
+      model: new MockLanguageModelV1({
+        doStream: async ({ prompt, mode }) => {
+          return {
+            stream: convertArrayToReadableStream([
+              { type: 'text-delta', textDelta: '{ ' },
+              { type: 'text-delta', textDelta: '"invalid": ' },
+              { type: 'text-delta', textDelta: `"Hello, ` },
+              { type: 'text-delta', textDelta: `world` },
+              { type: 'text-delta', textDelta: `!"` },
+              { type: 'text-delta', textDelta: ' }' },
+              {
+                type: 'finish',
+                finishReason: 'stop',
+                usage: { completionTokens: 10, promptTokens: 3 },
+              },
+            ]),
+            rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+          };
+        },
+      }),
+      schema: z.object({ content: z.string() }),
+      mode: 'json',
+      prompt: 'prompt',
+    });
+
+    // consume stream (runs in parallel)
+    convertAsyncIterableToArray(result.partialObjectStream);
+
+    // unhandled promise rejection should not be thrown (Vitest does this automatically)
+  });
 });
 
 describe('options.onFinish', () => {

packages/core/core/generate-object/stream-object.ts

@@ -28,6 +28,7 @@ import { convertZodToJSONSchema } from '../util/convert-zod-to-json-schema';
 import { prepareResponseHeaders } from '../util/prepare-response-headers';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { injectJsonSchemaIntoSystem } from './inject-json-schema-into-system';
+import { DelayedPromise } from '../util/delayed-promise';
 
 /**
 Generate a structured, typed object for a given prompt and schema using a language model.
@@ -316,17 +317,13 @@ The result of a `streamObject` call that contains the partial object stream and
  */
 export class StreamObjectResult<T> {
   private readonly originalStream: ReadableStream<ObjectStreamPart<T>>;
+  private readonly objectPromise: DelayedPromise<T>;
 
   /**
 Warnings from the model provider (e.g. unsupported settings)
    */
   readonly warnings: CallWarning[] | undefined;
 
-  /**
-The generated object (typed according to the schema). Resolved when the response is finished.
-   */
-  readonly object: Promise<T>;
-
   /**
 The token usage of the generated response. Resolved when the response is finished.
    */
@@ -363,12 +360,7 @@ Response headers.
     this.rawResponse = rawResponse;
 
     // initialize object promise
-    let resolveObject: (value: T | PromiseLike<T>) => void;
-    let rejectObject: (reason?: any) => void;
-    this.object = new Promise<T>((resolve, reject) => {
-      resolveObject = resolve;
-      rejectObject = reject;
-    });
+    this.objectPromise = new DelayedPromise<T>();
 
     // initialize usage promise
     let resolveUsage: (
@@ -388,6 +380,7 @@ Response headers.
     let delta = '';
     let latestObject: DeepPartial<T> | undefined = undefined;
 
+    const self = this;
     this.originalStream = stream.pipeThrough(
       new TransformStream<string | ObjectStreamInputPart, ObjectStreamPart<T>>({
         async transform(chunk, controller): Promise<void> {
@@ -445,10 +438,10 @@ Response headers.
 
               if (validationResult.success) {
                 object = validationResult.value;
-                resolveObject(object);
+                self.objectPromise.resolve(object);
               } else {
                 error = validationResult.error;
-                rejectObject(error);
+                self.objectPromise.reject(error);
               }
 
               break;
@@ -484,6 +477,13 @@ Response headers.
     );
   }
 
+  /**
+The generated object (typed according to the schema). Resolved when the response is finished.
+   */
+  get object(): Promise<T> {
+    return this.objectPromise.value;
+  }
+
   /**
 Stream of partial objects. It gets more complete as the stream progresses.
   

packages/core/core/util/delayed-promise.test.ts

@@ -0,0 +1,47 @@
+import { expect, it, describe } from 'vitest';
+import { DelayedPromise } from './delayed-promise';
+
+describe('DelayedPromise', () => {
+  it('should resolve when accessed after resolution', async () => {
+    const dp = new DelayedPromise<string>();
+    dp.resolve('success');
+    expect(await dp.value).toBe('success');
+  });
+
+  it('should reject when accessed after rejection', async () => {
+    const dp = new DelayedPromise<string>();
+    const error = new Error('failure');
+    dp.reject(error);
+    await expect(dp.value).rejects.toThrow('failure');
+  });
+
+  it('should resolve when accessed before resolution', async () => {
+    const dp = new DelayedPromise<string>();
+    const promise = dp.value;
+    dp.resolve('success');
+    expect(await promise).toBe('success');
+  });
+
+  it('should reject when accessed before rejection', async () => {
+    const dp = new DelayedPromise<string>();
+    const promise = dp.value;
+    const error = new Error('failure');
+    dp.reject(error);
+    await expect(promise).rejects.toThrow('failure');
+  });
+
+  it('should maintain the resolved state after multiple accesses', async () => {
+    const dp = new DelayedPromise<string>();
+    dp.resolve('success');
+    expect(await dp.value).toBe('success');
+    expect(await dp.value).toBe('success');
+  });
+
+  it('should maintain the rejected state after multiple accesses', async () => {
+    const dp = new DelayedPromise<string>();
+    const error = new Error('failure');
+    dp.reject(error);
+    await expect(dp.value).rejects.toThrow('failure');
+    await expect(dp.value).rejects.toThrow('failure');
+  });
+});

packages/core/core/util/delayed-promise.ts

@@ -0,0 +1,49 @@
+/**
+ * Delayed promise. It is only constructed once the value is accessed.
+ * This is useful to avoid unhandled promise rejections when the promise is created
+ * but not accessed.
+ */
+export class DelayedPromise<T> {
+  private status:
+    | { type: 'pending' }
+    | { type: 'resolved'; value: T }
+    | { type: 'rejected'; error: unknown } = { type: 'pending' };
+  private promise: Promise<T> | undefined;
+  private _resolve: undefined | ((value: T) => void) = undefined;
+  private _reject: undefined | ((error: unknown) => void) = undefined;
+
+  get value(): Promise<T> {
+    if (this.promise) {
+      return this.promise;
+    }
+
+    this.promise = new Promise<T>((resolve, reject) => {
+      if (this.status.type === 'resolved') {
+        resolve(this.status.value);
+      } else if (this.status.type === 'rejected') {
+        reject(this.status.error);
+      }
+
+      this._resolve = resolve;
+      this._reject = reject;
+    });
+
+    return this.promise;
+  }
+
+  resolve(value: T): void {
+    this.status = { type: 'resolved', value };
+
+    if (this.promise) {
+      this._resolve?.(value);
+    }
+  }
+
+  reject(error: unknown): void {
+    this.status = { type: 'rejected', error };
+
+    if (this.promise) {
+      this._reject?.(error);
+    }
+  }
+}