flynn/src/auth/openai.ts

import { readFileSync, writeFileSync, mkdirSync, chmodSync } from 'fs';
import { resolve } from 'path';
import { homedir } from 'os';

const ISSUER = 'https://auth.openai.com';
const CLIENT_ID = 'app_EMoamEEZ73f0CkXaXp7hrann';
const DEVICE_URL = `${ISSUER}/codex/device`;
const DEVICE_CODE_URL = `${ISSUER}/api/accounts/deviceauth/usercode`;
const DEVICE_TOKEN_URL = `${ISSUER}/api/accounts/deviceauth/token`;
const TOKEN_URL = `${ISSUER}/oauth/token`;

const POLLING_SAFETY_MARGIN_MS = 3000;
const REFRESH_SAFETY_MARGIN_MS = 30_000;

const AUTH_DIR = resolve(homedir(), '.config/flynn');
const AUTH_FILE = resolve(AUTH_DIR, 'auth.json');

export interface OpenAIOAuthInfo {
  access_token: string;
  refresh_token: string;
  /** Epoch millis. */
  expires_at: number;
  /** Optional account/org id used for subscription routing. */
  account_id?: string;
  created_at: string;
}

export interface OpenAIApiKeyInfo {
  api_key: string;
  created_at: string;
}

interface OpenAIStoreEntry {
  oauth?: OpenAIOAuthInfo;
  api_key?: OpenAIApiKeyInfo;
}

interface AuthStore {
  // Leave github entry untyped here so this module does not depend on github.ts.
  github?: unknown;
  /** OpenAI credentials. Backward compatible with legacy OAuth-only entries. */
  openai?: OpenAIStoreEntry | OpenAIOAuthInfo;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null;
}

function isOpenAIOAuthInfo(value: unknown): value is OpenAIOAuthInfo {
  if (!isRecord(value)) {
    return false;
  }
  return typeof value.access_token === 'string'
    && typeof value.refresh_token === 'string'
    && typeof value.expires_at === 'number'
    && typeof value.created_at === 'string';
}

function isOpenAIApiKeyInfo(value: unknown): value is OpenAIApiKeyInfo {
  if (!isRecord(value)) {
    return false;
  }
  return typeof value.api_key === 'string'
    && typeof value.created_at === 'string';
}

function readOpenAIEntry(store: AuthStore): OpenAIStoreEntry | null {
  const raw = store.openai as unknown;
  if (!raw) {
    return null;
  }

  // Legacy format: auth.json.openai stored the OAuth info directly.
  if (isOpenAIOAuthInfo(raw)) {
    return { oauth: raw };
  }

  if (!isRecord(raw)) {
    return null;
  }

  const oauth = isOpenAIOAuthInfo(raw.oauth) ? raw.oauth : undefined;
  const apiKey = isOpenAIApiKeyInfo(raw.api_key) ? raw.api_key : undefined;
  return { oauth, api_key: apiKey };
}

function writeOpenAIEntry(store: AuthStore, entry: OpenAIStoreEntry | null): void {
  if (!entry || (!entry.oauth && !entry.api_key)) {
    delete store.openai;
    return;
  }
  store.openai = entry;
}

interface DeviceAuthResponse {
  device_auth_id: string;
  user_code: string;
  interval: string;
}

interface DeviceTokenResponse {
  authorization_code: string;
  code_verifier: string;
}

interface TokenResponse {
  id_token?: string;
  access_token: string;
  refresh_token: string;
  expires_in?: number;
}

export interface IdTokenClaims {
  chatgpt_account_id?: string;
  organizations?: Array<{ id: string }>;
  'https://api.openai.com/auth'?: {
    chatgpt_account_id?: string;
  };
}

function safeJsonParse<T>(raw: string): T | null {
  try {
    return JSON.parse(raw) as T;
  } catch {
    return null;
  }
}

function readAuthStore(): AuthStore {
  try {
    const raw = readFileSync(AUTH_FILE, 'utf-8');
    const parsed = safeJsonParse<AuthStore>(raw);
    return parsed ?? {};
  } catch {
    return {};
  }
}

function writeAuthStore(store: AuthStore): void {
  mkdirSync(AUTH_DIR, { recursive: true });
  writeFileSync(AUTH_FILE, JSON.stringify(store, null, 2) + '\n', 'utf-8');
  chmodSync(AUTH_FILE, 0o600);
}

export function loadStoredOpenAIAuth(): OpenAIOAuthInfo | null {
  const store = readAuthStore();
  const entry = readOpenAIEntry(store);
  return entry?.oauth ?? null;
}

export function storeOpenAIAuth(info: OpenAIOAuthInfo): void {
  const store = readAuthStore();
  const entry = readOpenAIEntry(store) ?? {};
  entry.oauth = info;
  writeOpenAIEntry(store, entry);
  writeAuthStore(store);
}

export function clearOpenAIAuth(): void {
  const store = readAuthStore();
  const entry = readOpenAIEntry(store);
  if (entry) {
    delete entry.oauth;
    writeOpenAIEntry(store, entry);
  } else {
    delete store.openai;
  }
  writeAuthStore(store);
}

export function loadStoredOpenAIApiKey(): string | null {
  const store = readAuthStore();
  const entry = readOpenAIEntry(store);
  return entry?.api_key?.api_key ?? null;
}

export function storeOpenAIApiKey(key: string): void {
  const trimmed = key.trim();
  if (!trimmed) {
    throw new Error('OpenAI API key is empty');
  }
  const store = readAuthStore();
  const entry = readOpenAIEntry(store) ?? {};
  entry.api_key = { api_key: trimmed, created_at: new Date().toISOString() };
  writeOpenAIEntry(store, entry);
  writeAuthStore(store);
}

export function clearOpenAIApiKey(): void {
  const store = readAuthStore();
  const entry = readOpenAIEntry(store);
  if (!entry) {
    return;
  }
  delete entry.api_key;
  writeOpenAIEntry(store, entry);
  writeAuthStore(store);
}

/**
 * Get an OpenAI API key from any available source.
 * Priority: OPENAI_API_KEY → stored auth.json.
 */
export function getOpenAIApiKey(): string | null {
  return process.env.OPENAI_API_KEY
    ?? loadStoredOpenAIApiKey()
    ?? null;
}

export function parseJwtClaims(token: string): IdTokenClaims | undefined {
  const parts = token.split('.');
  if (parts.length !== 3) {return undefined;}
  try {
    return JSON.parse(Buffer.from(parts[1], 'base64url').toString()) as IdTokenClaims;
  } catch {
    return undefined;
  }
}

function extractAccountIdFromClaims(claims: IdTokenClaims): string | undefined {
  return claims.chatgpt_account_id
    ?? claims['https://api.openai.com/auth']?.chatgpt_account_id
    ?? claims.organizations?.[0]?.id;
}

export function extractAccountId(tokens: TokenResponse): string | undefined {
  const idToken = tokens.id_token;
  if (idToken) {
    const claims = parseJwtClaims(idToken);
    const id = claims && extractAccountIdFromClaims(claims);
    if (id) {return id;}
  }
  const accessToken = tokens.access_token;
  if (accessToken) {
    const claims = parseJwtClaims(accessToken);
    return claims ? extractAccountIdFromClaims(claims) : undefined;
  }
  return undefined;
}

async function requestDeviceAuth(): Promise<DeviceAuthResponse> {
  const response = await fetch(DEVICE_CODE_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'User-Agent': 'flynn',
    },
    body: JSON.stringify({ client_id: CLIENT_ID }),
  });

  if (!response.ok) {
    const body = await response.text();
    throw new Error(`OpenAI device auth start failed (${response.status}): ${body}`);
  }

  return response.json() as Promise<DeviceAuthResponse>;
}

async function pollDeviceToken(deviceAuthId: string, userCode: string, intervalMs: number): Promise<DeviceTokenResponse> {
  while (true) {
    await new Promise(r => setTimeout(r, intervalMs + POLLING_SAFETY_MARGIN_MS));

    const response = await fetch(DEVICE_TOKEN_URL, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'User-Agent': 'flynn',
      },
      body: JSON.stringify({ device_auth_id: deviceAuthId, user_code: userCode }),
    });

    if (response.ok) {
      return response.json() as Promise<DeviceTokenResponse>;
    }

    // OpenCode treats 403/404 as "pending".
    if (response.status === 403 || response.status === 404) {
      continue;
    }

    const body = await response.text();
    throw new Error(`OpenAI device auth token failed (${response.status}): ${body}`);
  }
}

async function exchangeAuthorizationCode(authCode: string, codeVerifier: string): Promise<TokenResponse> {
  const response = await fetch(TOKEN_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/x-www-form-urlencoded',
      'User-Agent': 'flynn',
    },
    body: new URLSearchParams({
      grant_type: 'authorization_code',
      code: authCode,
      redirect_uri: `${ISSUER}/deviceauth/callback`,
      client_id: CLIENT_ID,
      code_verifier: codeVerifier,
    }).toString(),
  });

  if (!response.ok) {
    const body = await response.text();
    throw new Error(`OpenAI token exchange failed (${response.status}): ${body}`);
  }

  return response.json() as Promise<TokenResponse>;
}

export async function refreshOpenAIAuth(refreshToken: string): Promise<TokenResponse> {
  const response = await fetch(TOKEN_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/x-www-form-urlencoded',
      'User-Agent': 'flynn',
    },
    body: new URLSearchParams({
      grant_type: 'refresh_token',
      refresh_token: refreshToken,
      client_id: CLIENT_ID,
    }).toString(),
  });

  if (!response.ok) {
    const body = await response.text();
    throw new Error(`OpenAI token refresh failed (${response.status}): ${body}`);
  }

  return response.json() as Promise<TokenResponse>;
}

/**
 * Ensure we have a valid (non-expired) OpenAI OAuth access token.
 * Refreshes and persists the token if needed.
 */
export async function ensureValidOpenAIAuth(): Promise<OpenAIOAuthInfo> {
  const current = loadStoredOpenAIAuth();
  if (!current) {
    throw new Error('OpenAI OAuth is not configured. Run `flynn openai-auth` to authenticate.');
  }

  if (current.expires_at > Date.now() + REFRESH_SAFETY_MARGIN_MS) {
    return current;
  }

  const refreshed = await refreshOpenAIAuth(current.refresh_token);
  const expiresAt = Date.now() + (refreshed.expires_in ?? 3600) * 1000;
  const accountId = extractAccountId(refreshed) ?? current.account_id;

  const updated: OpenAIOAuthInfo = {
    access_token: refreshed.access_token,
    refresh_token: refreshed.refresh_token,
    expires_at: expiresAt,
    account_id: accountId,
    created_at: current.created_at,
  };

  storeOpenAIAuth(updated);
  return updated;
}

/**
 * Run the OpenAI Codex device flow interactively.
 * @param onPrompt Callback to display the user code and verification URL to the user.
 */
export async function loginOpenAI(
  onPrompt: (userCode: string, verificationUri: string) => void,
): Promise<OpenAIOAuthInfo> {
  const device = await requestDeviceAuth();
  const intervalMs = Math.max(parseInt(device.interval) || 5, 1) * 1000;

  onPrompt(device.user_code, DEVICE_URL);

  const deviceToken = await pollDeviceToken(device.device_auth_id, device.user_code, intervalMs);
  const tokens = await exchangeAuthorizationCode(deviceToken.authorization_code, deviceToken.code_verifier);

  const expiresAt = Date.now() + (tokens.expires_in ?? 3600) * 1000;
  const accountId = extractAccountId(tokens);

  const info: OpenAIOAuthInfo = {
    access_token: tokens.access_token,
    refresh_token: tokens.refresh_token,
    expires_at: expiresAt,
    ...(accountId ? { account_id: accountId } : {}),
    created_at: new Date().toISOString(),
  };

  storeOpenAIAuth(info);
  return info;
}