🐛 Enforcing the waitForGeneration timeout on in-flight polls and retries

.changeset/wait-for-generation-deadline.md

@@ -0,0 +1,5 @@
+---
+"pdfmonkey": patch
+---
+
+Fixing `documents.waitForGeneration()` so its `timeout` is a true total budget: in-flight polls and their retries are now aborted when it expires, instead of resolving with a late success or reporting the timeout only after a slow response came back.

src/__tests__/generation.test.ts

@@ -1,4 +1,5 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
+import { PDFMonkey } from '../client.js';
 import { PDFMonkeyError } from '../error.js';
 import type { DocumentCard } from '../resources/document-cards.js';
 import type { Document } from '../resources/documents.js';
@@ -150,6 +151,68 @@ describe('waitForGeneration', () => {
     ).rejects.toThrow(/timed out/);
   });
 
+  describe('with a slow, abort-aware fetch', () => {
+    // Responds after `delayMs` unless the request signal aborts first.
+    function slowClient(status: Document['status'], delayMs: number) {
+      const fetch = vi.fn(
+        (_url: string, init: RequestInit) =>
+          new Promise<Response>((resolve, reject) => {
+            const timer = setTimeout(
+              () => resolve(Response.json({ document: { ...docFixture, status } })),
+              delayMs,
+            );
+            init.signal?.addEventListener('abort', () => {
+              clearTimeout(timer);
+              reject(init.signal?.reason);
+            });
+          }),
+      );
+      return { client: new PDFMonkey({ apiKey: 'sk_test', fetch, maxRetries: 0 }), fetch };
+    }
+
+    it('rejects instead of returning a success that arrives after the budget', async () => {
+      const { client } = slowClient('success', 250);
+      const startedAt = Date.now();
+
+      await expect(client.documents.waitForGeneration('doc_1', { timeout: 50 })).rejects.toThrow(
+        /timed out after 50ms/,
+      );
+      expect(Date.now() - startedAt).toBeLessThan(200);
+    });
+
+    it('times out within the budget when a poll is still in flight', async () => {
+      const { client } = slowClient('generating', 250);
+      const startedAt = Date.now();
+
+      await expect(client.documents.waitForGeneration('doc_1', { timeout: 50 })).rejects.toThrow(
+        /timed out after 50ms/,
+      );
+      expect(Date.now() - startedAt).toBeLessThan(200);
+    });
+
+    it('reports caller cancellation as an abort, not a timeout', async () => {
+      const { client } = slowClient('success', 250);
+      const controller = new AbortController();
+      setTimeout(() => controller.abort(), 20);
+
+      const promise = client.documents.waitForGeneration('doc_1', {
+        timeout: 10_000,
+        signal: controller.signal,
+      });
+      await expect(promise).rejects.toThrow(/aborted/);
+      await expect(promise).rejects.not.toThrow(/timed out/);
+    });
+
+    it('rejects a pre-aborted signal without polling', async () => {
+      const { client, fetch } = slowClient('success', 0);
+
+      await expect(
+        client.documents.waitForGeneration('doc_1', { signal: AbortSignal.abort() }),
+      ).rejects.toThrow('waitForGeneration aborted');
+      expect(fetch).not.toHaveBeenCalled();
+    });
+  });
+
   it('rejects maxInterval < interval', async () => {
     const { client } = createClient([]);
 

src/resources/documents.ts

@@ -271,34 +271,60 @@ export class Documents extends APIResource {
       );
     }
 
-    const start = Date.now();
-    let currentInterval = interval;
-
-    while (true) {
-      if (signal?.aborted) {
-        throw new PDFMonkeyError('waitForGeneration aborted');
-      }
-
-      const doc = await this.get(id, signal ? { signal } : undefined);
+    if (signal?.aborted) {
+      throw new PDFMonkeyError('waitForGeneration aborted');
+    }
 
-      if (doc.status === 'success') {
-        return doc;
-      }
+    // The budget covers in-flight polls (and their retries), not just the
+    // sleeps between them, so a slow response can't push us past `timeout`.
+    const deadline = new AbortController();
+    let timedOut = false;
+    const timeoutId = setTimeout(() => {
+      timedOut = true;
+      deadline.abort();
+    }, timeout);
+    const onCallerAbort = () => deadline.abort(signal?.reason);
+    signal?.addEventListener('abort', onCallerAbort, { once: true });
 
-      if (doc.status === 'failure' || doc.status === 'error') {
-        throw new PDFMonkeyError(
-          `Document generation failed: ${doc.failure_cause ?? 'Unknown error'}`,
-        );
+    const start = Date.now();
+    let currentInterval = interval;
+    let lastStatus: Document['status'] | undefined;
+
+    try {
+      while (true) {
+        const doc = await this.get(id, { signal: deadline.signal });
+        lastStatus = doc.status;
+
+        if (doc.status === 'success') {
+          return doc;
+        }
+
+        if (doc.status === 'failure' || doc.status === 'error') {
+          throw new PDFMonkeyError(
+            `Document generation failed: ${doc.failure_cause ?? 'Unknown error'}`,
+          );
+        }
+
+        if (Date.now() - start + currentInterval > timeout) {
+          throw new PDFMonkeyError(
+            `Document generation timed out after ${timeout}ms (status: ${doc.status})`,
+          );
+        }
+
+        await abortableSleep(currentInterval, deadline.signal);
+        currentInterval = Math.min(currentInterval * 2, maxInterval);
       }
-
-      if (Date.now() - start + currentInterval > timeout) {
+    } catch (error) {
+      if (timedOut && !signal?.aborted) {
         throw new PDFMonkeyError(
-          `Document generation timed out after ${timeout}ms (status: ${doc.status})`,
+          `Document generation timed out after ${timeout}ms (status: ${lastStatus ?? 'unknown'})`,
+          { cause: error },
         );
       }
-
-      await abortableSleep(currentInterval, signal);
-      currentInterval = Math.min(currentInterval * 2, maxInterval);
+      throw error;
+    } finally {
+      clearTimeout(timeoutId);
+      signal?.removeEventListener('abort', onCallerAbort);
     }
   }
 }