diff --git a/packages/blocks/azure/src/lib/actions/custom-azure-api-action.ts b/packages/blocks/azure/src/lib/actions/custom-azure-api-action.ts index c797669dc2..44640ec33b 100644 --- a/packages/blocks/azure/src/lib/actions/custom-azure-api-action.ts +++ b/packages/blocks/azure/src/lib/actions/custom-azure-api-action.ts @@ -1,15 +1,30 @@ -import { createCustomApiCallAction } from '@openops/blocks-common'; +import { + createCustomApiCallAction, + httpClient, + HttpError, + HttpHeaders, + HttpRequest, +} from '@openops/blocks-common'; import { Property } from '@openops/blocks-framework'; import { azureAuth, getUseHostSessionProperty } from '@openops/common'; import { getAzureAccessToken } from '../auth/get-azure-access-token'; import { getSubscriptionsDropdownForHostSession } from '../common-properties'; +const DEFAULT_RETRY_DELAY_MS = 60000; +const MAX_RETRY_ATTEMPTS = 4; +const COST_MANAGEMENT_RETRY_AFTER_HEADER_PATTERN = + /^x-ms-ratelimit-microsoft\.costmanagement.*-retry-after$/i; +const AZURE_RETRY_AFTER_HEADERS = new Set([ + 'x-ms-ratelimit-microsoft.consumption-retry-after', + 'x-ms-ratelimit-retailprices-retry-after', +]); + export const customAzureApiCallAction = createCustomApiCallAction({ auth: azureAuth, + baseUrl: () => 'https://management.azure.com/?api-version=2025-04-01', name: 'custom_azure_api_call', description: 'Make a custom REST API call to Azure.', displayName: 'Custom Azure API Call', - baseUrl: () => 'https://management.azure.com/?api-version=2025-04-01', additionalProps: { documentation: Property.MarkDown({ value: @@ -54,4 +69,100 @@ export const customAzureApiCallAction = createCustomApiCallAction({ 'Content-Type': 'application/json', }; }, + requestHandler: sendWithRetry, }); + +async function sendWithRetry( + request: HttpRequest>, +): Promise { + for (let attempt = 1; attempt <= MAX_RETRY_ATTEMPTS; attempt++) { + try { + return await httpClient.sendRequest(request); + } catch (error) { + if (!(error instanceof HttpError) || error.response.status !== 429) { + throw error; + } + + if (attempt === MAX_RETRY_ATTEMPTS) { + throw error; + } + + await sleep(getRetryDelayMs(error.response.headers, attempt)); + } + } + + throw new Error(`Failed to send Azure API request to ${request.url}`); +} + +export function getRetryDelayMs( + headers: HttpHeaders | undefined, + attempt: number, +): number { + const retryDelayMs = getHeaderRetryDelayMs(headers); + if (retryDelayMs === null) { + return Math.pow(2, attempt - 1) * DEFAULT_RETRY_DELAY_MS; + } + + return retryDelayMs; +} + +function getHeaderRetryDelayMs( + headers: HttpHeaders | undefined, +): number | null { + if (!headers) { + return null; + } + + const retryValues: string[] = []; + + for (const [headerName, headerValue] of Object.entries(headers)) { + if ( + AZURE_RETRY_AFTER_HEADERS.has(headerName.toLowerCase()) || + COST_MANAGEMENT_RETRY_AFTER_HEADER_PATTERN.test(headerName) + ) { + retryValues.push(...normalizeHeaderValues(headerValue)); + } + } + + if (retryValues.length === 0) { + return null; + } + + const retryDelaysMs = retryValues + .map((value) => parseRetryDelayMs(value)) + .filter((value): value is number => value !== null); + + if (retryDelaysMs.length === 0) { + return null; + } + + return Math.max(...retryDelaysMs); +} + +function normalizeHeaderValues( + headerValue: string | string[] | undefined, +): string[] { + if (!headerValue) { + return []; + } + + return Array.isArray(headerValue) ? headerValue : [headerValue]; +} + +function parseRetryDelayMs(headerValue: string): number | null { + const trimmedValue = headerValue.trim(); + if (trimmedValue.length === 0) { + return null; + } + + const retryAfterSeconds = Number(trimmedValue); + if (!Number.isFinite(retryAfterSeconds)) { + return null; + } + + return retryAfterSeconds * 1000; +} + +function sleep(delayMs: number): Promise { + return new Promise((resolve) => setTimeout(resolve, delayMs)); +} diff --git a/packages/blocks/azure/test/custom-azure-api-action.test.ts b/packages/blocks/azure/test/custom-azure-api-action.test.ts new file mode 100644 index 0000000000..15b732d619 --- /dev/null +++ b/packages/blocks/azure/test/custom-azure-api-action.test.ts @@ -0,0 +1,187 @@ +const authenticateUserWithAzureMock = jest.fn(); + +const openOpsMock = { + ...jest.requireActual('@openops/common'), + getUseHostSessionProperty: jest.fn().mockReturnValue({ + type: 'DYNAMIC', + required: true, + }), + authenticateUserWithAzure: authenticateUserWithAzureMock, +}; + +jest.mock('@openops/common', () => openOpsMock); + +import { HttpError } from '@openops/blocks-common'; +import axios from 'axios'; +import { + customAzureApiCallAction, + getRetryDelayMs, +} from '../src/lib/actions/custom-azure-api-action'; + +describe('customAzureApiCallAction', () => { + beforeEach(() => { + jest.clearAllMocks(); + jest.useFakeTimers(); + authenticateUserWithAzureMock.mockResolvedValue({ + access_token: 'mock-access-token', + }); + }); + + afterEach(() => { + jest.restoreAllMocks(); + jest.useRealTimers(); + }); + + test('should retry 429 responses using Azure Cost Management retry header', async () => { + jest + .spyOn(axios, 'request') + .mockRejectedValueOnce( + buildAxiosError({ + status: 429, + headers: { + 'x-ms-ratelimit-microsoft.costmanagement-entity-retry-after': '1', + }, + }), + ) + .mockResolvedValueOnce({ + status: 200, + headers: {}, + data: { ok: true }, + } as any); + + const runPromise = customAzureApiCallAction.run(createContext() as any); + + await jest.runAllTimersAsync(); + + await expect(runPromise).resolves.toEqual({ + status: 200, + headers: {}, + body: { ok: true }, + }); + }); + + test('should retry 429 responses using Azure Retail Prices retry header', async () => { + jest + .spyOn(axios, 'request') + .mockRejectedValueOnce( + buildAxiosError({ + status: 429, + headers: { + 'x-ms-ratelimit-retailprices-retry-after': '60', + }, + }), + ) + .mockResolvedValueOnce({ + status: 200, + headers: {}, + data: { ok: true }, + } as any); + + const runPromise = customAzureApiCallAction.run(createContext() as any); + + await jest.runAllTimersAsync(); + + await expect(runPromise).resolves.toEqual({ + status: 200, + headers: {}, + body: { ok: true }, + }); + }); + + test('should use Azure Cost Management retry header delay', () => { + expect( + getRetryDelayMs( + { + 'x-ms-ratelimit-microsoft.costmanagement-entity-retry-after': '1', + }, + 1, + ), + ).toBe(1000); + }); + + test('should use Azure Retail Prices retry header delay', () => { + expect( + getRetryDelayMs( + { + 'x-ms-ratelimit-retailprices-retry-after': '60', + }, + 1, + ), + ).toBe(60000); + }); + + test('should use the longest Azure Cost Management retry header value', () => { + expect( + getRetryDelayMs( + { + 'x-ms-ratelimit-microsoft.costmanagement-entity-retry-after': '3', + 'x-ms-ratelimit-microsoft.costmanagement-clienttype-retry-after': '1', + }, + 1, + ), + ).toBe(3000); + }); + + test('should retry 429 responses with fallback backoff when retry header is missing', () => { + expect(getRetryDelayMs({}, 1)).toBe(60000); + }); + + test('should not retry non-429 responses', async () => { + jest + .spyOn(axios, 'request') + .mockRejectedValueOnce(buildAxiosError({ status: 400 })); + + await expect( + customAzureApiCallAction.run(createContext() as any), + ).rejects.toBeInstanceOf(HttpError); + }); +}); + +function createContext(overrides?: Record) { + return { + auth: { + clientId: 'client-id', + clientSecret: 'client-secret', + tenantId: 'tenant-id', + }, + propsValue: { + useHostSession: { useHostSessionCheckbox: false }, + subscriptions: {}, + url: { + url: 'https://management.azure.com/subscriptions/sub-id/providers/Microsoft.CostManagement/query?api-version=2023-11-01', + }, + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + queryParams: {}, + body: { + type: 'ActualCost', + timeframe: 'LastMonth', + }, + ...overrides, + }, + }; +} + +function buildAxiosError({ + status, + headers = {}, +}: { + status: number; + headers?: Record; +}) { + return { + isAxiosError: true, + response: { + status, + headers, + data: { + error: { + code: String(status), + message: 'Request failed', + }, + }, + }, + }; +} diff --git a/packages/blocks/common/src/lib/helpers/index.ts b/packages/blocks/common/src/lib/helpers/index.ts index 2b1f3b3e55..cc51ce3cd9 100644 --- a/packages/blocks/common/src/lib/helpers/index.ts +++ b/packages/blocks/common/src/lib/helpers/index.ts @@ -26,6 +26,10 @@ export const getAccessTokenOrThrow = ( return accessToken; }; +type CustomApiCallRequestHandler = ( + request: HttpRequest>, +) => Promise; + export function createCustomApiCallAction({ auth, baseUrl, @@ -34,6 +38,7 @@ export function createCustomApiCallAction({ displayName, name, additionalProps, + requestHandler, }: { auth?: BlockAuthProperty; baseUrl: (auth?: unknown) => string; @@ -45,6 +50,7 @@ export function createCustomApiCallAction({ displayName?: string | null; name?: string | null; additionalProps?: Record; + requestHandler?: CustomApiCallRequestHandler; }) { return createAction({ name: name ? name : 'custom_api_call', @@ -139,6 +145,10 @@ export function createCustomApiCallAction({ } try { + if (requestHandler) { + return await requestHandler(request); + } + return await httpClient.sendRequest(request); } catch (error) { if (failsafe) { diff --git a/packages/blocks/common/src/lib/http/core/http-error.ts b/packages/blocks/common/src/lib/http/core/http-error.ts index d13fa90888..9604da28a2 100644 --- a/packages/blocks/common/src/lib/http/core/http-error.ts +++ b/packages/blocks/common/src/lib/http/core/http-error.ts @@ -1,4 +1,5 @@ import { AxiosError } from 'axios'; +import { HttpHeaders } from './http-headers'; export class HttpError extends Error { constructor( @@ -33,6 +34,7 @@ export class HttpError extends Error { get response() { return { status: this._err?.response?.status || 500, + headers: (this._err?.response?.headers as HttpHeaders | undefined) ?? {}, body: this._err?.response?.data, }; }