refactor: normalize wait-for implementations (#20)

Author: bbatha

src/index.ts

@@ -22,12 +22,6 @@ export {
 } from './core/error';
 
 // Export knowledge base helpers
-export {
-  waitForDatabase,
-  type WaitForDatabaseOptions,
-  WaitForDatabaseTimeoutError,
-  WaitForDatabaseFailedError,
-} from './resources/knowledge-bases/wait-for-database';
 export {
   IndexingJobAbortedError,
   IndexingJobNotFoundError,

src/resources/agents/agents.ts

@@ -306,9 +306,11 @@ export class Agents extends APIResource {
    * });
    * ```
    */
-  async waitForReady(uuid: string, options: WaitForAgentOptions): Promise<AgentReadinessResponse> {
+  async waitForReady(uuid: string, options: WaitForAgentOptions = {}): Promise<AgentReadinessResponse> {
     const { interval = 3000, timeout = 180000, signal } = options;
     const start = Date.now();
+    let pollCount = 0;
+    let currentInterval = interval;
 
     while (true) {
       signal?.throwIfAborted();
@@ -331,7 +333,14 @@ export class Agents extends APIResource {
         throw new AgentDeploymentError(uuid, status);
       }
 
-      await sleep(interval);
+      await sleep(currentInterval);
+
+      // Apply exponential backoff: double the interval after each poll, up to maxInterval
+      pollCount++;
+      if (pollCount > 2) {
+        // Start exponential backoff after 2 polls
+        currentInterval = Math.min(currentInterval * 2, 30000);
+      }
     }
   }
 }

src/resources/knowledge-bases/indexing-jobs.ts

@@ -7,6 +7,14 @@ import { RequestOptions } from '../../internal/request-options';
 import { path } from '../../internal/utils/path';
 import { sleep } from '../../internal/utils/sleep';
 
+export interface IndexingJobWaitForCompletionOptions extends RequestOptions {
+  /**
+   * Initial polling interval in milliseconds (default: 5000ms).
+   * This interval will be used for the first few polls, then exponential backoff applies.
+   */
+  interval?: number;
+}
+
 /**
  * Error thrown when an indexing job polling operation is aborted
  */
@@ -239,39 +247,20 @@ export class IndexingJobs extends APIResource {
    */
   async waitForCompletion(
     uuid: string,
-    options: {
-      /**
-       * Initial polling interval in milliseconds (default: 5000ms).
-       * This interval will be used for the first few polls, then exponential backoff applies.
-       */
-      interval?: number;
-      /**
-       * Maximum time to wait in milliseconds (default: 600000ms = 10 minutes)
-       */
-      timeout?: number;
-      /**
-       * Maximum polling interval in milliseconds (default: 30000ms = 30 seconds).
-       * Exponential backoff will not exceed this value.
-       */
-      maxInterval?: number;
-      /**
-       * Request options to pass to each poll request
-       */
-      requestOptions?: RequestOptions;
-    } = {},
+    options: IndexingJobWaitForCompletionOptions = {},
   ): Promise<IndexingJobRetrieveResponse> {
-    const { interval = 5000, timeout = 600000, maxInterval = 30000, requestOptions } = options;
+    const { interval = 5000, timeout = 600000, signal } = options;
     const startTime = Date.now();
     let currentInterval = interval;
     let pollCount = 0;
 
     while (true) {
       // Check if operation was aborted
-      if (requestOptions?.signal?.aborted) {
+      if (signal?.aborted) {
         throw new IndexingJobAbortedError();
       }
 
-      const response = await this.retrieve(uuid, requestOptions);
+      const response = await this.retrieve(uuid, options);
       const job = response.job;
 
       if (!job) {
@@ -303,7 +292,7 @@ export class IndexingJobs extends APIResource {
       pollCount++;
       if (pollCount > 2) {
         // Start exponential backoff after 2 polls
-        currentInterval = Math.min(currentInterval * 2, maxInterval);
+        currentInterval = Math.min(currentInterval * 2, 30000);
       }
     }
   }

src/resources/knowledge-bases/knowledge-bases.ts

@@ -36,7 +36,112 @@ import {
 import { APIPromise } from '../../core/api-promise';
 import { RequestOptions } from '../../internal/request-options';
 import { path } from '../../internal/utils/path';
-import { waitForDatabase } from './wait-for-database';
+import { GradientError } from '../../core/error';
+import { sleep } from '../../internal/utils/sleep';
+
+// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+export interface WaitForDatabaseOptions extends RequestOptions {
+  /**
+   * The polling interval in milliseconds. Defaults to 5000 (5 seconds).
+   */
+  interval?: number;
+}
+
+/**
+ * Database status values that indicate a successful deployment.
+ */
+const ONLINE_STATUSES = ['ONLINE'] as const;
+
+/**
+ * Database status values that indicate a failed deployment.
+ */
+const FAILED_STATUSES = ['DECOMMISSIONED', 'UNHEALTHY'] as const;
+
+/**
+ * Database status values that indicate the deployment is still in progress.
+ */
+const PENDING_STATUSES = [
+  'CREATING',
+  'POWEROFF',
+  'REBUILDING',
+  'REBALANCING',
+  'FORKING',
+  'MIGRATING',
+  'RESIZING',
+  'RESTORING',
+  'POWERING_ON',
+] as const;
+
+export class WaitForDatabaseTimeoutError extends GradientError {
+  constructor(message: string, kbId?: string, timeout?: number) {
+    super(message);
+    this.name = 'WaitForDatabaseTimeoutError';
+    if (kbId) {
+      (this as any).knowledgeBaseId = kbId;
+      (this as any).timeout = timeout;
+    }
+  }
+}
+
+export class WaitForDatabaseFailedError extends GradientError {
+  constructor(message: string, kbId?: string, status?: string) {
+    super(message);
+    this.name = 'WaitForDatabaseFailedError';
+    if (kbId) {
+      (this as any).knowledgeBaseId = kbId;
+      (this as any).databaseStatus = status;
+    }
+  }
+}
+
+/**
+ * Polls for knowledge base database creation to complete.
+ *
+ * This helper function polls the knowledge base status until the database is ONLINE,
+ * handling various error states and providing configurable timeout and polling intervals.
+ *
+ * @example
+ * ```ts
+ * import Gradient from '@digitalocean/gradient';
+ *
+ * const client = new Gradient();
+ *
+ * // Basic usage
+ * try {
+ *   const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000');
+ *   console.log('Database is ready:', kb.database_status); // 'ONLINE'
+ * } catch (error) {
+ *   if (error instanceof WaitForDatabaseTimeoutError) {
+ *     console.log('Polling timed out');
+ *   } else if (error instanceof WaitForDatabaseFailedError) {
+ *     console.log('Database deployment failed');
+ *   }
+ * }
+ *
+ * // With AbortSignal
+ * const controller = new AbortController();
+ * setTimeout(() => controller.abort(), 30000); // Cancel after 30 seconds
+ *
+ * try {
+ *   const kb = await client.knowledgeBases.waitForDatabase('123e4567-e89b-12d3-a456-426614174000', {
+ *     signal: controller.signal
+ *   });
+ * } catch (error) {
+ *   if (error.message === 'Operation was aborted') {
+ *     console.log('Operation was cancelled');
+ *   }
+ * }
+ * ```
+ *
+ * @param client - The Gradient client instance
+ * @param uuid - The knowledge base UUID to poll for
+ * @param options - Configuration options for polling behavior
+ * @returns Promise<KnowledgeBaseRetrieveResponse> - The knowledge base with ONLINE database status
+ * @throws WaitForDatabaseTimeoutError - If polling times out
+ * @throws WaitForDatabaseFailedError - If the database enters a failed state
+ * @throws Error - If the operation is aborted via AbortSignal
+ */
 
 export class KnowledgeBases extends APIResource {
   dataSources: DataSourcesAPI.DataSources = new DataSourcesAPI.DataSources(this._client);
@@ -159,9 +264,75 @@ export class KnowledgeBases extends APIResource {
    */
   async waitForDatabase(
     uuid: string,
-    options?: import('./wait-for-database').WaitForDatabaseOptions,
+    options: WaitForDatabaseOptions = {},
   ): Promise<KnowledgeBaseRetrieveResponse> {
-    return waitForDatabase(this._client, uuid, options || {});
+    const { interval = 5000, timeout = 600000, signal } = options;
+
+    const startTime = Date.now();
+
+    while (true) {
+      // Check if operation was aborted
+      if (signal?.aborted) {
+        throw new Error('Operation was aborted');
+      }
+
+      const elapsed = Date.now() - startTime;
+
+      if (elapsed > timeout) {
+        throw new WaitForDatabaseTimeoutError(
+          `Knowledge base database ${uuid} did not become ONLINE within ${timeout}ms`,
+          uuid,
+          timeout,
+        );
+      }
+
+      try {
+        const response = await this.retrieve(uuid, options);
+        const status = response.database_status;
+
+        if (!status) {
+          // If database_status is not present, continue polling
+          await sleep(interval);
+          continue;
+        }
+
+        // Check for successful completion
+        if (ONLINE_STATUSES.includes(status as any)) {
+          return response;
+        }
+
+        // Check for failed states
+        if (FAILED_STATUSES.includes(status as any)) {
+          throw new WaitForDatabaseFailedError(
+            `Knowledge base database ${uuid} entered failed state: ${status}`,
+            uuid,
+            status,
+          );
+        }
+
+        // Check if still in progress
+        if (PENDING_STATUSES.includes(status as any)) {
+          await sleep(interval);
+          continue;
+        }
+
+        // Unknown status - treat as error for safety
+        throw new WaitForDatabaseFailedError(
+          `Knowledge base database ${uuid} entered unknown state: ${status}`,
+          uuid,
+          status,
+        );
+      } catch (error) {
+        // If it's our custom error, re-throw it
+        if (error instanceof WaitForDatabaseFailedError || error instanceof WaitForDatabaseTimeoutError) {
+          throw error;
+        }
+
+        // For other errors (network issues, etc.), try waiting a bit longer before retrying
+        await sleep(Math.min(interval * 2, 30000)); // Max 30 seconds between retries
+        continue;
+      }
+    }
   }
 }
 
@@ -450,13 +621,6 @@ export interface KnowledgeBaseListParams {
 KnowledgeBases.DataSources = DataSources;
 KnowledgeBases.IndexingJobs = IndexingJobs;
 
-export {
-  waitForDatabase,
-  WaitForDatabaseOptions,
-  WaitForDatabaseTimeoutError,
-  WaitForDatabaseFailedError,
-} from './wait-for-database';
-
 export declare namespace KnowledgeBases {
   export {
     type APIKnowledgeBase as APIKnowledgeBase,