TJWaterAgent/src/auth/credentialRefresh.ts

import { randomUUID } from "node:crypto";

import { type RuntimeSessionContext } from "../runtime/sessionContext.js";

export type CredentialRefreshReason =
  | "access_token_expired"
  | "access_token_rejected";

export type CredentialRefreshEvent =
  | {
      type: "credential_refresh_required";
      requestId: string;
      reason: CredentialRefreshReason;
      timeoutMs: number;
    }
  | {
      type: "credential_refreshed";
      requestId: string;
    }
  | {
      type: "credential_refresh_failed";
      requestId: string;
      message: string;
    };

type PendingRefresh = {
  deadlineAt: number;
  promise: Promise<RuntimeSessionContext>;
  reason: CredentialRefreshReason;
  reject: (error: Error) => void;
  requestId: string;
  resolve: (context: RuntimeSessionContext) => void;
  timer: ReturnType<typeof setTimeout>;
};

type CredentialRefreshListener = (event: CredentialRefreshEvent) => void;

export class CredentialRefreshError extends Error {
  override readonly name = "CredentialRefreshError";

  constructor(
    message: string,
    readonly code: "cancelled" | "failed" | "timeout" | "unavailable" = "failed",
  ) {
    super(message);
  }
}

const AUTH_EXPIRY_SKEW_MS = 30_000;

export const isRuntimeCredentialExpired = (
  context: RuntimeSessionContext,
  now = Date.now(),
) => {
  if (!context.tokenExpiresAt) {
    return false;
  }
  const expiresAt = Date.parse(context.tokenExpiresAt);
  return Number.isFinite(expiresAt) && now >= expiresAt - AUTH_EXPIRY_SKEW_MS;
};

export class CredentialRefreshCoordinator {
  private readonly listeners = new Map<
    string,
    Set<CredentialRefreshListener>
  >();
  private readonly pending = new Map<string, PendingRefresh>();

  constructor(private readonly timeoutMs = 30_000) {}

  subscribe(sessionId: string, listener: CredentialRefreshListener) {
    const listeners =
      this.listeners.get(sessionId) ?? new Set<CredentialRefreshListener>();
    listeners.add(listener);
    this.listeners.set(sessionId, listeners);
    return () => {
      listeners.delete(listener);
      if (listeners.size === 0) {
        this.listeners.delete(sessionId);
      }
    };
  }

  request(sessionId: string, reason: CredentialRefreshReason) {
    const existing = this.pending.get(sessionId);
    if (existing) {
      return existing.promise;
    }
    if (!this.listeners.get(sessionId)?.size) {
      return Promise.reject(
        new CredentialRefreshError(
          "credential refresh channel is unavailable",
          "unavailable",
        ),
      );
    }

    const requestId = `credential-${randomUUID()}`;
    let resolvePromise!: (context: RuntimeSessionContext) => void;
    let rejectPromise!: (error: Error) => void;
    const promise = new Promise<RuntimeSessionContext>((resolve, reject) => {
      resolvePromise = resolve;
      rejectPromise = reject;
    });
    const timer = setTimeout(() => {
      this.fail(sessionId, requestId, "credential refresh timed out", "timeout");
    }, this.timeoutMs);
    this.pending.set(sessionId, {
      deadlineAt: Date.now() + this.timeoutMs,
      promise,
      reason,
      reject: rejectPromise,
      requestId,
      resolve: resolvePromise,
      timer,
    });
    this.emit(sessionId, {
      type: "credential_refresh_required",
      requestId,
      reason,
      timeoutMs: this.timeoutMs,
    });
    return promise;
  }

  resolve(
    sessionId: string,
    requestId: string,
    context: RuntimeSessionContext,
  ) {
    const pending = this.pending.get(sessionId);
    if (!pending || pending.requestId !== requestId) {
      return false;
    }
    clearTimeout(pending.timer);
    this.pending.delete(sessionId);
    pending.resolve(context);
    this.emit(sessionId, {
      type: "credential_refreshed",
      requestId,
    });
    return true;
  }

  fail(
    sessionId: string,
    requestId: string,
    message: string,
    code: CredentialRefreshError["code"] = "failed",
    emitFailureEvent = true,
  ) {
    const pending = this.pending.get(sessionId);
    if (!pending || pending.requestId !== requestId) {
      return false;
    }
    clearTimeout(pending.timer);
    this.pending.delete(sessionId);
    pending.reject(new CredentialRefreshError(message, code));
    if (emitFailureEvent) {
      this.emit(sessionId, {
        type: "credential_refresh_failed",
        requestId,
        message,
      });
    }
    return true;
  }

  cancelSession(sessionId: string, message = "credential refresh cancelled") {
    const pending = this.pending.get(sessionId);
    if (!pending) {
      return false;
    }
    return this.fail(
      sessionId,
      pending.requestId,
      message,
      "cancelled",
      false,
    );
  }

  getPendingRequestId(sessionId: string) {
    return this.pending.get(sessionId)?.requestId;
  }

  getPendingEvent(
    sessionId: string,
  ): Extract<CredentialRefreshEvent, { type: "credential_refresh_required" }> | null {
    const pending = this.pending.get(sessionId);
    if (!pending) return null;
    return {
      type: "credential_refresh_required",
      requestId: pending.requestId,
      reason: pending.reason,
      timeoutMs: Math.max(0, pending.deadlineAt - Date.now()),
    };
  }

  private emit(sessionId: string, event: CredentialRefreshEvent) {
    for (const listener of this.listeners.get(sessionId) ?? []) {
      listener(event);
    }
  }
}

export const runWithCredentialRefresh = async <T extends { status: number }>(
  coordinator: CredentialRefreshCoordinator,
  context: RuntimeSessionContext,
  execute: (context: RuntimeSessionContext) => Promise<T>,
) => {
  let activeContext = context;
  let refreshed = false;
  if (isRuntimeCredentialExpired(activeContext)) {
    activeContext = await coordinator.request(
      activeContext.sessionId,
      "access_token_expired",
    );
    refreshed = true;
  }

  let result = await execute(activeContext);
  if (result.status !== 401 || refreshed) {
    return result;
  }
  activeContext = await coordinator.request(
    activeContext.sessionId,
    "access_token_rejected",
  );
  result = await execute(activeContext);
  return result;
};