next-tjwater-agent/src/product/credential-refresh.ts

import { randomUUID } from "node:crypto";

import { type ProductCredential } from "./server-auth-client.js";
import { getRuntimeSessionContext } 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<ProductCredential>;
  reason: CredentialRefreshReason;
  reject: (error: Error) => void;
  requestId: string;
  resolve: (context: ProductCredential) => 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: ProductCredential,
  now = Date.now(),
) => {
  if (!context.expiresAt) {
    return false;
  }
  const expiresAt = Date.parse(context.expiresAt);
  return Number.isFinite(expiresAt) && now >= expiresAt - AUTH_EXPIRY_SKEW_MS;
};

export class CredentialRefreshCoordinator {
  private readonly listeners = new Map<
    string,
    Set<CredentialRefreshListener>
  >();
  private readonly background = new Set<string>();
  enableBackground(sessionId: string) { this.background.add(sessionId); }
  disableBackground(sessionId: string) { this.background.delete(sessionId); }
  notifyFailure(sessionId: string, message: string) {
    this.emit(sessionId, { type: "credential_refresh_failed", requestId: `credential-${randomUUID()}`, message });
  }
  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 && !this.background.has(sessionId)) {
      return Promise.reject(
        new CredentialRefreshError(
          "credential refresh channel is unavailable",
          "unavailable",
        ),
      );
    }

    const requestId = `credential-${randomUUID()}`;
    let resolvePromise!: (context: ProductCredential) => void;
    let rejectPromise!: (error: Error) => void;
    const promise = new Promise<ProductCredential>((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: ProductCredential,
  ) {
    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,
  sessionId: string,
  context: ProductCredential,
  execute: (context: ProductCredential) => Promise<T>,
  getCurrent?: () => ProductCredential | undefined,
) => {
  let activeContext = getCurrent?.() ?? context;
  let refreshed = false;
  if (isRuntimeCredentialExpired(activeContext)) {
    activeContext = await coordinator.request(
      sessionId,
      "access_token_expired",
    );
    refreshed = true;
  }

  let result = await execute(activeContext);
  if (result.status !== 401 || refreshed) {
    return result;
  }
  if (result instanceof Response) await result.body?.cancel();
  const latest = getCurrent?.();
  activeContext = latest && latest.authorization !== activeContext.authorization
    ? latest : await coordinator.request(sessionId, "access_token_rejected");
  result = await execute(activeContext);
  return result;
};

const credentialVault = new Map<string, ProductCredential>();
const credentialFailures = new Set<string>();
export const credentialRefreshCoordinator = new CredentialRefreshCoordinator(45_000);
export const bindSessionCredential = (sessionId: string, credential: ProductCredential | null) => {
  if (credential) { credentialVault.set(sessionId, { ...credential }); credentialFailures.delete(sessionId); }
};
const credentialReleaseListeners = new Set<(id: string) => void>();
export const onSessionCredentialReleased = (listener: (id: string) => void) => { credentialReleaseListeners.add(listener); return () => credentialReleaseListeners.delete(listener); };
export const releaseSessionCredential = (sessionId: string) => {
  for (const listener of credentialReleaseListeners) listener(sessionId);
  credentialRefreshCoordinator.cancelSession(sessionId);
  credentialRefreshCoordinator.disableBackground(sessionId);
  credentialVault.delete(sessionId);
  credentialFailures.delete(sessionId);
};
export const refreshSessionCredential = (sessionId: string, credential: ProductCredential, requestId?: string) => {
  const existing = credentialVault.get(sessionId);
  if (!existing || existing.projectId !== credential.projectId) return false;
  const pendingId = credentialRefreshCoordinator.getPendingRequestId(sessionId);
  if (requestId && requestId !== pendingId) return false;
  bindSessionCredential(sessionId, credential);
  if (pendingId) credentialRefreshCoordinator.resolve(sessionId, pendingId, credential);
  return true;
};
export const requestWithSessionCredential = async (
  sessionId: string, execute: (credential: ProductCredential) => Promise<Response>,
) => {
  const credential = credentialVault.get(sessionId);
  if (!credential && getRuntimeSessionContext(sessionId)?.projectKey.startsWith("personal-")) {
    throw Object.assign(new Error("个人会话未关联项目，请选择项目会话后访问业务数据"), { status: 403, code: "PROJECT_CONTEXT_REQUIRED" });
  }
  if (!credential) throw Object.assign(new Error("task access requires a current product credential"), { status: 401 });
  if (credentialFailures.has(sessionId)) throw Object.assign(new CredentialRefreshError("credential renewal failed"), { status: 401 });
  try {
    const response = await runWithCredentialRefresh(credentialRefreshCoordinator, sessionId, credential, execute, () => credentialVault.get(sessionId));
    if (response.status === 401) {
      credentialFailures.add(sessionId);
      credentialRefreshCoordinator.notifyFailure(sessionId, "renewed credential was rejected");
    }
    return response;
  } catch (error) {
    if (error instanceof CredentialRefreshError && error.code !== "cancelled") {
      credentialFailures.add(sessionId);
      if (error.code === "unavailable") credentialRefreshCoordinator.notifyFailure(sessionId, error.message);
    }
    throw error;
  }
};