From 30ad2a728bf880489b8b52569d4fa3a8b6854fbd Mon Sep 17 00:00:00 2001 From: Ben Batha Date: Tue, 28 Oct 2025 10:44:18 -0400 Subject: [PATCH] refactor: normalize wait-for implementations --- src/index.ts | 6 - src/resources/agents/agents.ts | 13 +- .../knowledge-bases/indexing-jobs.ts | 37 ++-- .../knowledge-bases/knowledge-bases.ts | 184 +++++++++++++++- .../knowledge-bases/wait-for-database.ts | 197 ------------------ .../knowledge-bases/wait-for-database.test.ts | 81 ------- 6 files changed, 198 insertions(+), 320 deletions(-) delete mode 100644 src/resources/knowledge-bases/wait-for-database.ts delete mode 100644 tests/resources/knowledge-bases/wait-for-database.test.ts diff --git a/src/index.ts b/src/index.ts index 51dbd70..bb26480 100644 --- a/src/index.ts +++ b/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, diff --git a/src/resources/agents/agents.ts b/src/resources/agents/agents.ts index 586ada7..ba6d606 100644 --- a/src/resources/agents/agents.ts +++ b/src/resources/agents/agents.ts @@ -306,9 +306,11 @@ export class Agents extends APIResource { * }); * ``` */ - async waitForReady(uuid: string, options: WaitForAgentOptions): Promise { + async waitForReady(uuid: string, options: WaitForAgentOptions = {}): Promise { 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); + } } } } diff --git a/src/resources/knowledge-bases/indexing-jobs.ts b/src/resources/knowledge-bases/indexing-jobs.ts index 25fdd8b..657f7d8 100644 --- a/src/resources/knowledge-bases/indexing-jobs.ts +++ b/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 { - 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); } } } diff --git a/src/resources/knowledge-bases/knowledge-bases.ts b/src/resources/knowledge-bases/knowledge-bases.ts index aa4fa4b..f5f1d5c 100644 --- a/src/resources/knowledge-bases/knowledge-bases.ts +++ b/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 - 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 { - 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, diff --git a/src/resources/knowledge-bases/wait-for-database.ts b/src/resources/knowledge-bases/wait-for-database.ts deleted file mode 100644 index 650c1a9..0000000 --- a/src/resources/knowledge-bases/wait-for-database.ts +++ /dev/null @@ -1,197 +0,0 @@ -// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. - -import { Gradient } from '../../client'; -import { GradientError } from '../../core/error'; -import { sleep } from '../../internal/utils/sleep'; -import { KnowledgeBaseRetrieveResponse } from './knowledge-bases'; -import { RequestOptions } from '../../internal/request-options'; - -export interface WaitForDatabaseOptions { - /** - * The polling interval in milliseconds. Defaults to 5000 (5 seconds). - */ - interval?: number; - - /** - * The maximum time to wait in milliseconds. Defaults to 600000 (10 minutes). - */ - timeout?: number; - - /** - * AbortSignal to cancel the polling operation. - */ - signal?: AbortSignal; - - /** - * Additional request options to pass through to the knowledge base retrieval request. - */ - requestOptions?: RequestOptions; -} - -/** - * 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 - 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 async function waitForDatabase( - client: Gradient, - uuid: string, - options: WaitForDatabaseOptions, -): Promise { - const { interval = 5000, timeout = 600000, signal, requestOptions } = 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 client.knowledgeBases.retrieve(uuid, requestOptions); - 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; - } - } -} diff --git a/tests/resources/knowledge-bases/wait-for-database.test.ts b/tests/resources/knowledge-bases/wait-for-database.test.ts deleted file mode 100644 index 523dff6..0000000 --- a/tests/resources/knowledge-bases/wait-for-database.test.ts +++ /dev/null @@ -1,81 +0,0 @@ -// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. - -import Gradient from '@digitalocean/gradient'; -import { - WaitForDatabaseFailedError, - WaitForDatabaseTimeoutError, - waitForDatabase, -} from '../../../src/resources/knowledge-bases/wait-for-database'; - -const client = new Gradient({ - accessToken: 'My Access Token', - modelAccessKey: 'My Model Access Key', - agentAccessKey: 'My Agent Access Key', - baseURL: process.env['TEST_API_BASE_URL'] ?? 'http://127.0.0.1:4010', -}); - -describe('waitForDatabase', () => { - const kbUuid = '123e4567-e89b-12d3-a456-426614174000'; - - beforeEach(() => { - jest.resetAllMocks(); - }); - - describe('error classes', () => { - it('should create WaitForDatabaseTimeoutError with correct properties', () => { - const error = new WaitForDatabaseTimeoutError('Test timeout', kbUuid, 1000); - - expect(error.name).toBe('WaitForDatabaseTimeoutError'); - expect(error.message).toBe('Test timeout'); - expect(error).toBeInstanceOf(Error); - }); - - it('should create WaitForDatabaseFailedError with correct properties', () => { - const error = new WaitForDatabaseFailedError('Test failure', kbUuid, 'DECOMMISSIONED'); - - expect(error.name).toBe('WaitForDatabaseFailedError'); - expect(error.message).toBe('Test failure'); - expect(error).toBeInstanceOf(Error); - }); - }); - - describe('function parameters', () => { - it('should accept correct parameters', () => { - expect(typeof waitForDatabase).toBe('function'); - expect(waitForDatabase.length).toBe(3); - }); - - it('should use default options when none provided', () => { - const options = {}; - expect(waitForDatabase).toBeDefined(); - // Function should exist and be callable (will fail at runtime due to mocking, but should compile) - }); - }); - - describe('status constants', () => { - it('should handle different database status values', () => { - // Test that status strings are handled correctly - const statuses = [ - 'ONLINE', - 'CREATING', - 'REBUILDING', - 'RESIZING', - 'POWERING_ON', - 'DECOMMISSIONED', - 'UNHEALTHY', - ]; - - statuses.forEach((status) => { - expect(typeof status).toBe('string'); - }); - }); - }); - - describe('integration test placeholder', () => { - it('should integrate with knowledge base retrieval', async () => { - // This is a placeholder test - actual integration tests would require a running mock server - expect(client.knowledgeBases).toBeDefined(); - expect(client.knowledgeBases.retrieve).toBeDefined(); - }); - }); -});