Modify Custom Azure API Call action to retry on 429 errors

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<Record<string, unknown>>,
+): Promise<unknown> {
+  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<void> {
+  return new Promise((resolve) => setTimeout(resolve, delayMs));
+}

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<string, unknown>) {
+  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<string, string>;
+}) {
+  return {
+    isAxiosError: true,
+    response: {
+      status,
+      headers,
+      data: {
+        error: {
+          code: String(status),
+          message: 'Request failed',
+        },
+      },
+    },
+  };
+}

packages/blocks/common/src/lib/helpers/index.ts

@@ -26,6 +26,10 @@ export const getAccessTokenOrThrow = (
   return accessToken;
 };
 
+type CustomApiCallRequestHandler = (
+  request: HttpRequest<Record<string, unknown>>,
+) => Promise<unknown>;
+
 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<string, any>;
+  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) {

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,
     };
   }