feat(agent): add credential refresh and unify learning tools

src/auth/credentialRefresh.ts

@@ -0,0 +1,232 @@
+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;
+};