iddaai-be/src/common/utils/ai-engine-client.ts

import axios, {
  AxiosError,
  AxiosInstance,
  AxiosRequestConfig,
  AxiosResponse,
} from "axios";
import { Logger } from "@nestjs/common";

export type AiCircuitState = "closed" | "open";

export interface AiEngineClientOptions {
  baseUrl: string;
  logger: Logger;
  serviceName: string;
  timeoutMs?: number;
  maxRetries?: number;
  retryDelayMs?: number;
  circuitBreakerThreshold?: number;
  circuitBreakerCooldownMs?: number;
}

interface AiEngineRequestConfig extends AxiosRequestConfig {
  retryCount?: number;
}

export interface AiEngineClientSnapshot {
  state: AiCircuitState;
  consecutiveFailures: number;
  openedAt: string | null;
}

export class AiEngineRequestError extends Error {
  status?: number;
  detail?: unknown;
  isCircuitOpen: boolean;

  constructor(
    message: string,
    options: {
      status?: number;
      detail?: unknown;
      isCircuitOpen?: boolean;
    } = {},
  ) {
    super(message);
    this.name = "AiEngineRequestError";
    this.status = options.status;
    this.detail = options.detail;
    this.isCircuitOpen = options.isCircuitOpen ?? false;
  }
}

export class AiEngineClient {
  private readonly axiosClient: AxiosInstance;
  private readonly logger: Logger;
  private readonly serviceName: string;
  private readonly defaultTimeoutMs: number;
  private readonly maxRetries: number;
  private readonly retryDelayMs: number;
  private readonly circuitBreakerThreshold: number;
  private readonly circuitBreakerCooldownMs: number;

  private consecutiveFailures = 0;
  private circuitOpenedAt: number | null = null;

  constructor(options: AiEngineClientOptions) {
    this.logger = options.logger;
    this.serviceName = options.serviceName;
    this.defaultTimeoutMs = options.timeoutMs ?? 30000;
    this.maxRetries = options.maxRetries ?? 2;
    this.retryDelayMs = options.retryDelayMs ?? 750;
    this.circuitBreakerThreshold = options.circuitBreakerThreshold ?? 3;
    this.circuitBreakerCooldownMs = options.circuitBreakerCooldownMs ?? 30000;

    this.axiosClient = axios.create({
      baseURL: options.baseUrl,
      timeout: this.defaultTimeoutMs,
    });
  }

  async get<T>(
    path: string,
    config?: AiEngineRequestConfig,
  ): Promise<AxiosResponse<T>> {
    return this.request<T>({
      method: "get",
      url: path,
      ...config,
    });
  }

  async post<T>(
    path: string,
    data?: unknown,
    config?: AiEngineRequestConfig,
  ): Promise<AxiosResponse<T>> {
    return this.request<T>({
      method: "post",
      url: path,
      data,
      ...config,
    });
  }

  getSnapshot(): AiEngineClientSnapshot {
    return {
      state: this.isCircuitOpen() ? "open" : "closed",
      consecutiveFailures: this.consecutiveFailures,
      openedAt: this.circuitOpenedAt
        ? new Date(this.circuitOpenedAt).toISOString()
        : null,
    };
  }

  private async request<T>(
    config: AiEngineRequestConfig,
  ): Promise<AxiosResponse<T>> {
    this.ensureCircuitAvailable();

    const retries = this.resolveRetryCount(config);
    let lastError: unknown;

    for (let attempt = 0; attempt <= retries; attempt += 1) {
      try {
        const response = await this.axiosClient.request<T>({
          timeout: this.defaultTimeoutMs,
          ...config,
        });

        this.resetFailures();
        return response;
      } catch (error) {
        lastError = error;
        const shouldRetry = attempt < retries && this.isRetriableError(error);

        if (!shouldRetry) {
          // Only register circuit breaker failure for server/network errors, not client errors (4xx)
          if (this.isServerError(error)) {
            this.registerFailure(error);
          } else {
            // It's a successful contact with the engine (e.g. 404, 422), so reset failures
            this.resetFailures();
          }
          throw this.toRequestError(error);
        }

        this.logger.warn(
          `[${this.serviceName}] AI request retry ${attempt + 1}/${retries} for ${config.method?.toUpperCase()} ${config.url}`,
        );
        await this.delay(this.retryDelayMs * (attempt + 1));
      }
    }

    this.registerFailure(lastError);
    throw this.toRequestError(lastError);
  }

  private resolveRetryCount(config: AiEngineRequestConfig): number {
    if (typeof config.retryCount === "number" && config.retryCount >= 0) {
      return config.retryCount;
    }

    return this.maxRetries;
  }

  private ensureCircuitAvailable() {
    if (!this.isCircuitOpen()) {
      return;
    }

    const remainingCooldown =
      this.circuitBreakerCooldownMs -
      (Date.now() - (this.circuitOpenedAt ?? 0));

    if (remainingCooldown > 0) {
      throw new AiEngineRequestError("AI engine circuit breaker is open", {
        status: 503,
        detail: {
          cooldownRemainingMs: remainingCooldown,
        },
        isCircuitOpen: true,
      });
    }

    this.logger.warn(
      `[${this.serviceName}] AI circuit breaker cooldown elapsed, allowing a recovery attempt`,
    );
    this.circuitOpenedAt = null;
  }

  private isCircuitOpen(): boolean {
    return this.circuitOpenedAt !== null;
  }

  private resetFailures() {
    this.consecutiveFailures = 0;
    this.circuitOpenedAt = null;
  }

  private registerFailure(error: unknown) {
    this.consecutiveFailures += 1;

    const normalizedError = this.toRequestError(error);
    this.logger.warn(
      `[${this.serviceName}] AI request failed (${this.consecutiveFailures}/${this.circuitBreakerThreshold}): ${normalizedError.message}`,
    );

    if (this.consecutiveFailures >= this.circuitBreakerThreshold) {
      this.circuitOpenedAt = Date.now();
      this.logger.error(
        `[${this.serviceName}] AI circuit breaker opened after ${this.consecutiveFailures} consecutive failures`,
      );
    }
  }

  private isRetriableError(error: unknown): boolean {
    if (!axios.isAxiosError(error)) {
      return false;
    }

    if (!error.response) {
      return true;
    }

    const status = error.response.status;
    return status >= 500 || status === 429 || error.code === "ECONNABORTED";
  }

  private isServerError(error: unknown): boolean {
    if (!axios.isAxiosError(error)) {
      return true; // Not an axios error, assume internal/network error
    }
    if (!error.response) {
      return true; // Network error, timeout, etc.
    }
    const status = error.response.status;
    return status >= 500 || status === 429;
  }

  private toRequestError(error: unknown): AiEngineRequestError {
    if (error instanceof AiEngineRequestError) {
      return error;
    }

    if (axios.isAxiosError(error)) {
      const detail = error.response?.data ?? error.message;
      const status = error.response?.status;
      const message = this.buildAxiosErrorMessage(error);

      return new AiEngineRequestError(message, {
        status,
        detail,
      });
    }

    if (error instanceof Error) {
      return new AiEngineRequestError(error.message);
    }

    return new AiEngineRequestError("Unknown AI engine error", {
      detail: error,
    });
  }

  private buildAxiosErrorMessage(error: AxiosError): string {
    if (error.code === "ECONNABORTED") {
      return "AI engine request timed out";
    }

    if (!error.response) {
      return "AI engine is unreachable";
    }

    const detail =
      (error.response.data as Record<string, unknown> | undefined)?.detail ??
      error.message;

    return typeof detail === "string"
      ? detail
      : `AI engine request failed with status ${error.response.status}`;
  }

  private async delay(ms: number) {
    await new Promise((resolve) => setTimeout(resolve, ms));
  }
}