makelore/electron/services/works-square-ai-gateway.ts

import { WORKS_SQUARE_CONFIG } from '../api/works-config';
import { proxyAwareFetch } from '../utils/proxy-fetch';
import { logger } from '../utils/logger';
import { getValidWorksSquareAccessToken } from './works-square-session';

const AI_TOKEN_REFRESH_SKEW_MS = 60_000;

export type WorksSquareAIGatewayCredentialInput = {
  accessToken: string;
  expiresIn?: unknown;
  expiresAt?: number | null;
  oneApiBaseUrl: string;
};

export type WorksSquareAIGatewayCredential = {
  accessToken: string;
  expiresAt: number | null;
  oneApiBaseUrl: string;
};

type GatewaySessionPayload = {
  access_token?: unknown;
  accessToken?: unknown;
  api_key?: unknown;
  apiKey?: unknown;
  expires_in?: unknown;
  expiresIn?: unknown;
  api_key_expires_in?: unknown;
  apiKeyExpiresIn?: unknown;
  one_api_base_url?: unknown;
  oneApiBaseUrl?: unknown;
  base_url?: unknown;
  baseUrl?: unknown;
};

let currentCredential: WorksSquareAIGatewayCredential | null = null;
let refreshPromise: Promise<WorksSquareAIGatewayCredential | null> | null = null;

function asString(value: unknown): string | null {
  return typeof value === 'string' && value.trim() ? value.trim() : null;
}

function asNumber(value: unknown): number | null {
  if (typeof value === 'number' && Number.isFinite(value)) return value;
  if (typeof value === 'string') {
    const parsed = Number(value);
    return Number.isFinite(parsed) ? parsed : null;
  }
  return null;
}

function normalizeBaseUrl(value: string): string {
  const baseUrl = value.trim().replace(/\/+$/, '');
  if (!/^https?:\/\//i.test(baseUrl)) {
    throw new Error('one-api base URL must start with http:// or https://');
  }
  return baseUrl;
}

function expiresAtFromInput(input: WorksSquareAIGatewayCredentialInput, nowMs = Date.now()): number | null {
  if (typeof input.expiresAt === 'number' && Number.isFinite(input.expiresAt)) {
    return input.expiresAt;
  }
  const seconds = asNumber(input.expiresIn);
  return seconds && seconds > 0 ? nowMs + seconds * 1000 : null;
}

function createWorksUrl(pathname: string): string {
  return `${WORKS_SQUARE_CONFIG.apiBaseUrl.replace(/\/+$/, '')}${pathname}`;
}

function parseGatewaySessionPayload(
  payload: GatewaySessionPayload,
  nowMs = Date.now(),
): WorksSquareAIGatewayCredential {
  const accessToken = asString(
    payload.access_token
    ?? payload.accessToken
    ?? payload.api_key
    ?? payload.apiKey,
  );
  const oneApiBaseUrl = asString(
    payload.one_api_base_url
    ?? payload.oneApiBaseUrl
    ?? payload.base_url
    ?? payload.baseUrl,
  );

  if (!accessToken) {
    throw new Error('AI gateway session response did not include access_token');
  }
  if (!oneApiBaseUrl) {
    throw new Error('AI gateway session response did not include one_api_base_url');
  }

  return {
    accessToken,
    expiresAt: expiresAtFromInput({
      accessToken,
      expiresIn: payload.expires_in
        ?? payload.expiresIn
        ?? payload.api_key_expires_in
        ?? payload.apiKeyExpiresIn,
      oneApiBaseUrl,
    }, nowMs),
    oneApiBaseUrl: normalizeBaseUrl(oneApiBaseUrl),
  };
}

async function readResponsePayload(response: Response): Promise<unknown> {
  const text = await response.text();
  if (!text.trim()) return null;
  try {
    return JSON.parse(text) as unknown;
  } catch {
    return text;
  }
}

function isCredentialFresh(
  credential: WorksSquareAIGatewayCredential | null,
  nowMs = Date.now(),
): credential is WorksSquareAIGatewayCredential {
  if (!credential) return false;
  if (credential.expiresAt == null) return true;
  return credential.expiresAt > nowMs + AI_TOKEN_REFRESH_SKEW_MS;
}

export function seedWorksSquareAIGatewayCredential(
  input: WorksSquareAIGatewayCredentialInput,
  nowMs = Date.now(),
): WorksSquareAIGatewayCredential {
  const accessToken = input.accessToken.trim();
  if (!accessToken) {
    throw new Error('AI gateway access token is required');
  }

  currentCredential = {
    accessToken,
    expiresAt: expiresAtFromInput(input, nowMs),
    oneApiBaseUrl: normalizeBaseUrl(input.oneApiBaseUrl),
  };
  return { ...currentCredential };
}

export function clearWorksSquareAIGatewayCredential(): void {
  currentCredential = null;
  refreshPromise = null;
}

export function getWorksSquareAIGatewaySnapshot(): WorksSquareAIGatewayCredential | null {
  return currentCredential ? { ...currentCredential } : null;
}

export function markWorksSquareAIGatewayCredentialExpired(): void {
  if (!currentCredential) return;
  currentCredential = {
    ...currentCredential,
    expiresAt: 0,
  };
}

async function refreshWorksSquareAIGatewayCredential(
  options: { fetchImpl?: typeof fetch; nowMs?: number } = {},
): Promise<WorksSquareAIGatewayCredential | null> {
  const fetchImpl = options.fetchImpl ?? proxyAwareFetch;
  const nowMs = options.nowMs ?? Date.now();
  const worksAccessToken = await getValidWorksSquareAccessToken({
    fetchImpl,
    nowMs,
  });
  if (!worksAccessToken) {
    clearWorksSquareAIGatewayCredential();
    return null;
  }

  const response = await fetchImpl(createWorksUrl('/api/ai-gateway/session'), {
    method: 'POST',
    headers: {
      Accept: 'application/json',
      Authorization: `Bearer ${worksAccessToken}`,
    },
  });
  const payload = await readResponsePayload(response);

  if (!response.ok) {
    logger.warn('[works-square-ai-gateway] Session refresh failed', {
      status: response.status,
      payload,
    });
    if (response.status === 401) {
      clearWorksSquareAIGatewayCredential();
    }
    return null;
  }

  if (!payload || typeof payload !== 'object' || Array.isArray(payload)) {
    throw new Error('AI gateway session response was invalid');
  }

  currentCredential = parseGatewaySessionPayload(payload as GatewaySessionPayload, nowMs);
  return { ...currentCredential };
}

export async function getFreshWorksSquareAIGatewayCredential(
  options: { fetchImpl?: typeof fetch; nowMs?: number } = {},
): Promise<WorksSquareAIGatewayCredential | null> {
  const nowMs = options.nowMs ?? Date.now();
  if (isCredentialFresh(currentCredential, nowMs)) {
    return { ...currentCredential };
  }

  if (!refreshPromise) {
    refreshPromise = refreshWorksSquareAIGatewayCredential(options).finally(() => {
      refreshPromise = null;
    });
  }
  return await refreshPromise;
}