@bitkyc08/opencodex 2.7.42 → 2.7.43

package/src/oauth/kiro.ts

@@ -1,29 +1,44 @@
 /**
  * Kiro (AWS CodeWhisperer) OAuth — import-first.
  *
- * Unlike browser/PKCE providers, kiro reuses the locally installed kiro-cli login:
- * it reads the kiro-cli SQLite token store, falls back to KIRO_ACCESS_TOKEN env, then to a
- * manual access-token paste (CLI only). Refresh hits the Kiro desktop refresh endpoint.
+ * Normal login imports the locally installed kiro-cli session. Account-add login deliberately
+ * asks kiro-cli to switch identities in its supported browser flow, then imports that fresh session.
  *
  * Ported from jawcode packages/ai/src/providers/kiro.ts (readKiroCliSqlite, refreshKiroDesktopToken).
- * profileArn/region are NOT stored in the credential — the kiro ADAPTER resolves them at request
- * time (SQLite profile_arn / KIRO_PROFILE_ARN, KIRO_REGION) since getValidAccessToken surfaces
- * only the access token.
+ * profileArn/region/client registration are persisted per OCX account so switching the account pool
+ * never combines one account's access token with another account's local Kiro profile metadata.
  */
-import type { OAuthController, OAuthCredentials } from "./types";
+import type { KiroOAuthMetadata, OAuthController, OAuthCredentials } from "./types";
 import {
+  discardKiroCliSessionRecovery,
   inferRegionFromProfileArn,
   inspectKiroCliSqliteSources,
+  inspectKiroCliSessionSnapshot,
   normalizeKiroRegion,
+  persistKiroCliSessionRecovery,
   readImportedKiroCredential,
   readKiroCliSqliteCredential,
+  restoreKiroCliSession,
+  restoreStaleKiroCliSessionRecovery,
   requireKiroRegion,
+  type ImportedKiroCredential,
+  type KiroCliSessionSnapshot,
   type KiroImportDiagnostic,
 } from "./kiro-credentials";
+import { getAccountSet, saveAccountCredential } from "./store";
 
 const DEFAULT_REGION = "us-east-1";
 const REFRESH_URL = "https://prod.{region}.auth.desktop.kiro.dev/refreshToken";
 const OIDC_URL = "https://oidc.{region}.amazonaws.com/token";
+const KIRO_TERMINAL_REFRESH_ERRORS = new Set([
+  "invalid_grant",
+  "refresh_token_reused",
+  "revoked",
+  "revoked_token",
+  "refresh_token_revoked",
+  "access_denied",
+  "expired_token",
+]);
 
 interface ImportedKiroToken {
   access: string;
@@ -31,6 +46,205 @@ interface ImportedKiroToken {
   expires: number;
 }
 
+export interface KiroCliCommandResult {
+  exitCode: number;
+  stdout: string;
+}
+
+export class KiroTokenRefreshError extends Error {
+  constructor(
+    readonly httpStatus: number,
+    readonly oauthError?: string,
+  ) {
+    super(`Kiro token refresh failed: ${httpStatus}${oauthError ? ` (${oauthError})` : ""}`);
+    this.name = "KiroTokenRefreshError";
+  }
+}
+
+export type KiroCliRunner = (args: string[], signal?: AbortSignal) => Promise<KiroCliCommandResult>;
+
+export interface KiroLoginOptions {
+  forceLogin?: boolean;
+  cliRunner?: KiroCliRunner;
+}
+
+const pendingKiroLoginTransactions = new WeakMap<OAuthCredentials, KiroCliSessionSnapshot>();
+/** Forced logins that started with no native CLI DB must logout on persistence failure. */
+const pendingKiroEmptyPriorSessions = new WeakSet<OAuthCredentials>();
+
+function logoutKiroCliBestEffort(): void {
+  try {
+    Bun.spawnSync(["kiro-cli", "logout"], {
+      stdin: "ignore",
+      stdout: "ignore",
+      stderr: "ignore",
+      windowsHide: true,
+    });
+  } catch {
+    // Best-effort rollback when the prior CLI state was empty.
+  }
+}
+
+/** Settle the external CLI side of a forced login after OCX credential persistence resolves. */
+export function settleKiroLoginTransaction(credential: OAuthCredentials, persisted: boolean): void {
+  const snapshot = pendingKiroLoginTransactions.get(credential);
+  const emptyPrior = pendingKiroEmptyPriorSessions.has(credential);
+  pendingKiroLoginTransactions.delete(credential);
+  pendingKiroEmptyPriorSessions.delete(credential);
+  if (!persisted) {
+    if (snapshot) {
+      restoreKiroCliSession(snapshot);
+      discardKiroCliSessionRecovery(snapshot);
+      return;
+    }
+    if (emptyPrior) logoutKiroCliBestEffort();
+    return;
+  }
+  if (snapshot) discardKiroCliSessionRecovery(snapshot);
+}
+
+function restoreKiroLoginOrThrow(snapshot: KiroCliSessionSnapshot | null, emptyPrior: boolean, cause: unknown): never {
+  if (snapshot) {
+    try {
+      restoreKiroCliSession(snapshot);
+      discardKiroCliSessionRecovery(snapshot);
+    } catch (restoreError) {
+      throw new AggregateError(
+        [cause, restoreError],
+        "Kiro login failed and the previous Kiro CLI session could not be restored.",
+      );
+    }
+  } else if (emptyPrior) {
+    logoutKiroCliBestEffort();
+  }
+  throw cause;
+}
+
+function throwIfKiroLoginCancelled(signal?: AbortSignal): void {
+  if (signal?.aborted) throw new Error("Kiro login cancelled.");
+}
+
+async function defaultKiroCliRunner(args: string[], signal?: AbortSignal): Promise<KiroCliCommandResult> {
+  throwIfKiroLoginCancelled(signal);
+  let child: ReturnType<typeof Bun.spawn>;
+  try {
+    child = Bun.spawn(["kiro-cli", ...args], {
+      stdin: "ignore",
+      stdout: "pipe",
+      stderr: "ignore",
+      windowsHide: true,
+    });
+  } catch {
+    throw new Error("Kiro CLI is not installed or could not be started.");
+  }
+  const abort = () => child.kill();
+  signal?.addEventListener("abort", abort, { once: true });
+  // AbortSignal does not replay an abort that lands between the pre-check and listener registration.
+  if (signal?.aborted) abort();
+  try {
+    const [exitCode, stdout] = await Promise.all([
+      child.exited,
+      child.stdout instanceof ReadableStream ? new Response(child.stdout).text() : Promise.resolve(""),
+    ]);
+    throwIfKiroLoginCancelled(signal);
+    return { exitCode, stdout };
+  } finally {
+    signal?.removeEventListener("abort", abort);
+  }
+}
+
+async function readKiroCliIdentity(runner: KiroCliRunner, signal?: AbortSignal): Promise<{ email?: string }> {
+  try {
+    const result = await runner(["whoami", "--format", "json"], signal);
+    if (result.exitCode !== 0) return {};
+    const parsed = JSON.parse(result.stdout) as { email?: unknown };
+    const email = typeof parsed.email === "string" ? parsed.email.trim().toLowerCase() : "";
+    return email && email.length <= 320 ? { email } : {};
+  } catch {
+    return {};
+  }
+}
+
+function metadataFromImported(imported: ImportedKiroCredential): KiroOAuthMetadata | undefined {
+  const metadata: KiroOAuthMetadata = {
+    ...(imported.profileArn ? { profileArn: imported.profileArn } : {}),
+    ...(imported.ssoRegion ? { ssoRegion: imported.ssoRegion } : {}),
+    ...(imported.apiRegion ? { apiRegion: imported.apiRegion } : {}),
+    ...(imported.clientId ? { clientId: imported.clientId } : {}),
+    ...(imported.clientSecret ? { clientSecret: imported.clientSecret } : {}),
+  };
+  return Object.keys(metadata).length > 0 ? metadata : undefined;
+}
+
+/** Persistable routing subset from explicit KIRO_* environment overrides (never local CLI state). */
+export function environmentKiroRoutingMetadata(): Pick<KiroOAuthMetadata, "profileArn" | "apiRegion" | "ssoRegion"> | undefined {
+  const profileArn = process.env.KIRO_PROFILE_ARN?.trim() || undefined;
+  const apiRegion = process.env.KIRO_API_REGION !== undefined
+    ? requireKiroRegion(process.env.KIRO_API_REGION)
+    : undefined;
+  const ssoRegion = process.env.KIRO_REGION !== undefined
+    ? requireKiroRegion(process.env.KIRO_REGION)
+    : undefined;
+  if (!profileArn && !apiRegion && !ssoRegion) return undefined;
+  return {
+    ...(profileArn ? { profileArn } : {}),
+    ...(apiRegion ? { apiRegion } : {}),
+    ...(ssoRegion ? { ssoRegion } : {}),
+  };
+}
+
+/**
+ * Before Add-account switches the external CLI identity, bind any matching legacy identity-less
+ * OCX row to the current CLI session. Unmatched identity-less rows are left alone here so a
+ * cancelled browser login cannot destroy the only stored credential; after a successful add they
+ * remain selectable but cannot borrow the new CLI identity (and reauth refuses identity-less slots).
+ */
+async function bindMatchingLegacyIdentitylessKiroAccounts(
+  runner: KiroCliRunner,
+  signal?: AbortSignal,
+): Promise<void> {
+  const set = getAccountSet("kiro");
+  if (!set) return;
+  const legacyAccounts = set.accounts.filter(
+    account => account.credential.accountId === undefined && account.credential.email === undefined,
+  );
+  if (legacyAccounts.length === 0) return;
+
+  let imported: ImportedKiroCredential | null = null;
+  try {
+    imported = readKiroCliSqliteCredential();
+  } catch {
+    imported = null;
+  }
+  if (!imported?.refresh) return;
+
+  for (const account of legacyAccounts) {
+    const refresh = account.credential.refresh;
+    if (!refresh || imported.refresh !== refresh) continue;
+    const bound = await oauthCredentialFromImported(imported, runner, signal);
+    if (!bound.accountId && !bound.email) continue;
+    await saveAccountCredential("kiro", account.id, bound);
+  }
+}
+
+async function oauthCredentialFromImported(
+  imported: ImportedKiroCredential,
+  runner: KiroCliRunner,
+  signal?: AbortSignal,
+): Promise<OAuthCredentials> {
+  const identity = imported.source === "sqlite" ? await readKiroCliIdentity(runner, signal) : {};
+  const metadata = metadataFromImported(imported);
+  return {
+    access: imported.access,
+    refresh: imported.refresh,
+    expires: imported.expires,
+    source: imported.source === "json" ? "credential-file" : "local-cli",
+    ...(imported.profileArn ? { accountId: imported.profileArn } : {}),
+    ...(identity.email ? { email: identity.email } : {}),
+    ...(metadata ? { kiro: metadata } : {}),
+  };
+}
+
 export type KiroCliImportDiagnosticStatus = KiroImportDiagnostic["status"];
 export type KiroCliImportDiagnostic = KiroImportDiagnostic;
 
@@ -54,22 +268,73 @@ export function readKiroCliSqlite(): ImportedKiroToken | null {
  * so the GUI renders the paste-input field, then blocks on onManualCodeInput for the token.
  * If neither onAuth nor onManualCodeInput is available, throws a clear error.
  */
-export async function loginKiro(ctrl: OAuthController): Promise<OAuthCredentials> {
+export async function loginKiro(ctrl: OAuthController, options: KiroLoginOptions = {}): Promise<OAuthCredentials> {
+  const runner = options.cliRunner ?? defaultKiroCliRunner;
+  // A prior process may have exited after switching the external CLI account but before
+  // settlement. Recover that durable transaction before either importing or switching again.
+  restoreStaleKiroCliSessionRecovery();
+  if (options.forceLogin) {
+    throwIfKiroLoginCancelled(ctrl.signal);
+    ctrl.onAuth?.({
+      url: "",
+      instructions: "Kiro CLI is opening a fresh browser login. This also switches the account used by kiro-cli.",
+    });
+    ctrl.onProgress?.("Opening a fresh Kiro CLI browser login.");
+    // Never log out of a session we could not capture: the recovery contract promises an exact
+    // restore, so an uncapturable store (unreadable / schema-mismatched / ambiguous token) must
+    // abort here rather than let a later failure destroy it permanently.
+    const inspected = inspectKiroCliSessionSnapshot();
+    if (inspected.blocked) {
+      throw new Error(
+        "Kiro CLI session could not be backed up, so OCX will not sign it out. " +
+          "Repair or remove the unreadable kiro-cli credential database " +
+          "(usually `~/.local/share/kiro-cli/data.sqlite3` or " +
+          "`~/Library/Application Support/kiro-cli/data.sqlite3`), " +
+          "unset KIROCLI_DB_PATH / KIRO_CLI_DB_FILE if set for import-only overrides, then retry.",
+      );
+    }
+    const previousSession = inspected.snapshot;
+    await bindMatchingLegacyIdentitylessKiroAccounts(runner, ctrl.signal);
+    if (previousSession) persistKiroCliSessionRecovery(previousSession);
+    try {
+      const logout = await runner(["logout"], ctrl.signal);
+      throwIfKiroLoginCancelled(ctrl.signal);
+      if (logout.exitCode !== 0) throw new Error("Kiro CLI could not prepare a fresh login.");
+      const login = await runner(["login"], ctrl.signal);
+      throwIfKiroLoginCancelled(ctrl.signal);
+      if (login.exitCode !== 0) throw new Error("Kiro CLI login did not complete successfully.");
+      const fresh = readKiroCliSqliteCredential();
+      if (!fresh) throw new Error("Kiro CLI login completed but no credential could be imported.");
+      const credential = await oauthCredentialFromImported(fresh, runner, ctrl.signal);
+      throwIfKiroLoginCancelled(ctrl.signal);
+      if (!credential.accountId && !credential.email) {
+        throw new Error("Kiro login completed but OCX could not determine a stable account identity.");
+      }
+      if (previousSession) pendingKiroLoginTransactions.set(credential, previousSession);
+      else pendingKiroEmptyPriorSessions.add(credential);
+      return credential;
+    } catch (error) {
+      restoreKiroLoginOrThrow(previousSession, !previousSession, error);
+    }
+  }
+
   const imported = readImportedKiroCredential();
   if (imported) {
     ctrl.onProgress?.(imported.source === "json" ? "Imported token from Kiro credentials file." : "Imported token from installed kiro-cli login.");
-    return {
-      access: imported.access,
-      refresh: imported.refresh,
-      expires: imported.expires,
-      source: imported.source === "json" ? "credential-file" : "local-cli",
-    };
+    return oauthCredentialFromImported(imported, runner, ctrl.signal);
   }
 
   const envToken = process.env.KIRO_ACCESS_TOKEN;
   if (envToken) {
     ctrl.onProgress?.("Using KIRO_ACCESS_TOKEN from environment.");
-    return { access: envToken, refresh: process.env.KIRO_REFRESH_TOKEN ?? "", expires: Date.now() + 3600_000, source: "environment" };
+    const routing = environmentKiroRoutingMetadata();
+    return {
+      access: envToken,
+      refresh: process.env.KIRO_REFRESH_TOKEN ?? "",
+      expires: Date.now() + 3600_000,
+      source: "environment",
+      ...(routing ? { kiro: routing } : {}),
+    };
   }
 
   if (ctrl.onManualCodeInput) {
@@ -85,7 +350,16 @@ export async function loginKiro(ctrl: OAuthController): Promise<OAuthCredentials
     });
     ctrl.onProgress?.("No kiro-cli token found. Paste a Kiro access token (starts with 'aoa'), or install the Kiro CLI and run `kiro-cli login` first.");
     const raw = (await ctrl.onManualCodeInput()).trim();
-    if (raw) return { access: raw, refresh: "", expires: Date.now() + 3600_000, source: "manual" };
+    if (raw) {
+      const routing = environmentKiroRoutingMetadata();
+      return {
+        access: raw,
+        refresh: "",
+        expires: Date.now() + 3600_000,
+        source: "manual",
+        ...(routing ? { kiro: routing } : {}),
+      };
+    }
   }
 
   throw new Error(
@@ -95,16 +369,25 @@ export async function loginKiro(ctrl: OAuthController): Promise<OAuthCredentials
   );
 }
 
-/** Auth/SSO region precedence: KIRO_REGION → imported SSO region → default us-east-1. */
-export function resolveKiroRegion(): string {
+/** Account metadata is authoritative; legacy accountless calls use KIRO_REGION → local import → default. */
+export function resolveKiroRegion(account?: KiroOAuthMetadata): string {
+  if (account !== undefined) return normalizeKiroRegion(account.ssoRegion) || DEFAULT_REGION;
   if (process.env.KIRO_REGION !== undefined) return requireKiroRegion(process.env.KIRO_REGION);
   return normalizeKiroRegion(readImportedKiroCredential()?.ssoRegion) || DEFAULT_REGION;
 }
 
-/** Runtime API region precedence: KIRO_API_REGION → imported API/profile region → auth region. */
-export function resolveKiroApiRegion(): string {
-  const imported = readImportedKiroCredential();
+/** Account metadata is authoritative; legacy accountless calls may use KIRO_API_REGION/local import. */
+export function resolveKiroApiRegion(account?: Pick<KiroOAuthMetadata, "apiRegion" | "profileArn" | "ssoRegion">): string {
+  if (account !== undefined) {
+    return (
+      normalizeKiroRegion(account.apiRegion) ||
+      inferRegionFromProfileArn(account.profileArn) ||
+      normalizeKiroRegion(account.ssoRegion) ||
+      DEFAULT_REGION
+    );
+  }
   if (process.env.KIRO_API_REGION !== undefined) return requireKiroRegion(process.env.KIRO_API_REGION);
+  const imported = readImportedKiroCredential();
   return (
     normalizeKiroRegion(imported?.apiRegion) ||
     inferRegionFromProfileArn(imported?.profileArn) ||
@@ -116,15 +399,29 @@ export function resolveKiroApiRegion(): string {
 
 /**
  * Resolve the CodeWhisperer profileArn for request-time use by the adapter.
- * KIRO_PROFILE_ARN env → kiro-cli SQLite `profile_arn`. Returns undefined if absent
- * (the adapter decides whether that is fatal).
+ * Account metadata is authoritative. Legacy accountless calls use KIRO_PROFILE_ARN → local import.
+ * Returns undefined if absent (the adapter decides whether that is fatal).
  */
-export function resolveKiroProfileArn(): string | undefined {
+export function resolveKiroProfileArn(account?: Pick<KiroOAuthMetadata, "profileArn">): string | undefined {
+  if (account !== undefined) return account.profileArn;
   const env = process.env.KIRO_PROFILE_ARN;
   if (env) return env;
   return readImportedKiroCredential()?.profileArn;
 }
 
+async function kiroTokenRefreshError(response: Response): Promise<KiroTokenRefreshError> {
+  let oauthError: string | undefined;
+  try {
+    const payload = await response.json() as { error?: unknown };
+    if (typeof payload.error === "string" && KIRO_TERMINAL_REFRESH_ERRORS.has(payload.error)) {
+      oauthError = payload.error;
+    }
+  } catch {
+    // Error bodies are untrusted and intentionally excluded from the surfaced message.
+  }
+  return new KiroTokenRefreshError(response.status, oauthError);
+}
+
 async function readTokenResponse(res: Response, oldRefresh: string): Promise<OAuthCredentials> {
   const data = (await res.json()) as { accessToken?: string; refreshToken?: string; expiresIn?: number };
   if (!data.accessToken) throw new Error("Kiro refresh returned no accessToken");
@@ -135,43 +432,126 @@ async function readTokenResponse(res: Response, oldRefresh: string): Promise<OAu
   };
 }
 
-async function refreshKiroDesktopToken(refresh: string, signal?: AbortSignal): Promise<OAuthCredentials> {
-  const region = resolveKiroRegion();
+function kiroRefreshSignal(signal?: AbortSignal): AbortSignal {
+  const timeout = AbortSignal.timeout(30_000);
+  return signal ? AbortSignal.any([signal, timeout]) : timeout;
+}
+
+async function refreshKiroDesktopToken(refresh: string, signal?: AbortSignal, metadata?: KiroOAuthMetadata): Promise<OAuthCredentials> {
+  const region = resolveKiroRegion(metadata);
   const res = await fetch(REFRESH_URL.replace("{region}", region), {
     method: "POST",
     headers: { "Content-Type": "application/json" },
     body: JSON.stringify({ refreshToken: refresh }),
-    signal: signal ?? AbortSignal.timeout(30_000),
+    signal: kiroRefreshSignal(signal),
   });
-  if (!res.ok) throw new Error(`Kiro token refresh failed: ${res.status}`);
+  if (!res.ok) throw await kiroTokenRefreshError(res);
   return readTokenResponse(res, refresh);
 }
 
-async function refreshAwsSsoOidcToken(refresh: string, signal?: AbortSignal): Promise<OAuthCredentials> {
-  const imported = readImportedKiroCredential();
-  if (!imported?.clientId || !imported.clientSecret) return refreshKiroDesktopToken(refresh, signal);
-  const region = resolveKiroRegion();
+async function refreshAwsSsoOidcToken(
+  refresh: string,
+  signal?: AbortSignal,
+  credential?: OAuthCredentials,
+): Promise<OAuthCredentials> {
+  let metadata = credential?.kiro;
+  if (!metadata) {
+    const local = readImportedKiroCredential();
+    // Only a local store holding this exact refresh token describes this account. Anything else
+    // belongs to whichever account kiro-cli is currently signed into.
+    if (local?.refresh === refresh) metadata = metadataFromImported(local);
+  }
+  // A stored OCX account with no usable `kiro` metadata must still refresh account-scoped: falling
+  // through to `resolveKiroRegion(undefined)` would read KIRO_REGION or the local CLI import and
+  // borrow an unrelated account's region after a switch. Environment/manual credentials may still
+  // honor explicit KIRO_* routing; other stored accounts pin an empty marker (default region).
+  // Only a truly accountless refresh (no stored credential) keeps the legacy env/local fallback.
+  if (!metadata && credential) {
+    metadata = credential.source === "environment" || credential.source === "manual"
+      ? environmentKiroRoutingMetadata() ?? {}
+      : {};
+  }
+  const clientId = metadata?.clientId;
+  const clientSecret = metadata?.clientSecret;
+  if (!clientId || !clientSecret) return refreshKiroDesktopToken(refresh, signal, metadata);
+  const region = resolveKiroRegion(metadata);
   const run = async (refreshToken: string): Promise<Response> => fetch(OIDC_URL.replace("{region}", region), {
     method: "POST",
     headers: { "Content-Type": "application/json" },
     body: JSON.stringify({
       grantType: "refresh_token",
-      clientId: imported.clientId,
-      clientSecret: imported.clientSecret,
+      clientId,
+      clientSecret,
       refreshToken,
     }),
-    signal: signal ?? AbortSignal.timeout(30_000),
+    signal: kiroRefreshSignal(signal),
   });
-  let res = await run(refresh);
-  if (!res.ok && res.status === 400 && imported.source === "sqlite") {
-    const reloaded = readImportedKiroCredential();
-    if (reloaded?.refresh && reloaded.refresh !== refresh) res = await run(reloaded.refresh);
-  }
-  if (!res.ok) throw new Error(`Kiro AWS SSO OIDC refresh failed: ${res.status}`);
+  const res = await run(refresh);
+  if (!res.ok) throw await kiroTokenRefreshError(res);
   return readTokenResponse(res, refresh);
 }
 
-export async function refreshKiroToken(refresh: string, signal?: AbortSignal): Promise<OAuthCredentials> {
+function matchingRotatedKiroCliCredential(
+  refresh: string,
+  credential?: OAuthCredentials,
+): ImportedKiroCredential | undefined {
+  const storedIdentities = [credential?.kiro?.profileArn, credential?.accountId]
+    .filter((value): value is string => Boolean(value));
+  if (storedIdentities.length === 0 || new Set(storedIdentities).size !== 1) return undefined;
+  try {
+    const local = readKiroCliSqliteCredential();
+    if (!local?.profileArn || storedIdentities.some(identity => identity !== local.profileArn)) return undefined;
+    if (!local.refresh || local.refresh === refresh) return undefined;
+    return local;
+  } catch {
+    // An unrelated or malformed local store must not block the stored OCX account.
+    return undefined;
+  }
+}
+
+function metadataForRotatedKiroCliCredential(
+  credential: OAuthCredentials,
+  local: ImportedKiroCredential,
+): KiroOAuthMetadata {
+  const {
+    clientId: _storedClientId,
+    clientSecret: _storedClientSecret,
+    ...storedRouting
+  } = credential.kiro ?? {};
+  const localMetadata = metadataFromImported(local) ?? {};
+  const {
+    clientId: _localClientId,
+    clientSecret: _localClientSecret,
+    ...localRouting
+  } = localMetadata;
+  return {
+    ...storedRouting,
+    ...localRouting,
+    ...(local.authType === "aws_sso_oidc" && local.clientId && local.clientSecret
+      ? { clientId: local.clientId, clientSecret: local.clientSecret }
+      : {}),
+  };
+}
+
+export async function refreshKiroToken(
+  refresh: string,
+  signal?: AbortSignal,
+  credential?: OAuthCredentials,
+): Promise<OAuthCredentials> {
   if (!refresh) throw new Error("Kiro: no refresh token available (re-run `kiro-cli login`).");
-  return refreshAwsSsoOidcToken(refresh, signal);
+  try {
+    return await refreshAwsSsoOidcToken(refresh, signal, credential);
+  } catch (error) {
+    if (!(error instanceof KiroTokenRefreshError) || error.httpStatus !== 400) throw error;
+    if (!credential) throw error;
+    const local = matchingRotatedKiroCliCredential(refresh, credential);
+    if (!local) throw error;
+    const retryMetadata = metadataForRotatedKiroCliCredential(credential, local);
+    const fresh = await refreshAwsSsoOidcToken(local.refresh, signal, {
+      ...credential,
+      refresh: local.refresh,
+      kiro: retryMetadata,
+    });
+    return { ...fresh, kiro: retryMetadata };
+  }
 }

package/src/oauth/login-cli.ts

@@ -2,19 +2,21 @@ import * as readline from "node:readline";
 import { openUrl } from "../lib/open-url";
 import { loadConfig, saveConfig } from "../config";
 import { findLiveProxy, probeHostname } from "../server/proxy-liveness";
-import { isPublicOAuthProvider, listOAuthProviders, OAUTH_PROVIDERS, runLogin } from "./index";
+import { isPublicOAuthProvider, listOAuthProviders, runLogin } from "./index";
 import { KEY_LOGIN_PROVIDERS, isKeyLoginProvider, validateApiKey, type KeyLoginProvider } from "./key-providers";
 import type { OcxProviderConfig } from "../types";
+import { configuredAdminToken } from "../lib/admin-secrets";
+import { codexAccountNamespaceProviderCollisionError } from "../codex/account-namespace-match";
 
 export function runningProxyUpdateHeaders(): Headers {
   const headers = new Headers({ "Content-Type": "application/json" });
-  const apiToken = process.env.OPENCODEX_API_AUTH_TOKEN?.trim();
-  if (apiToken) headers.set("X-OpenCodex-API-Key", apiToken);
+  const adminToken = configuredAdminToken();
+  if (adminToken) headers.set("X-OpenCodex-API-Key", adminToken);
   return headers;
 }
 
 /** Push the new provider into a running proxy's live config so it routes without a restart. */
-async function notifyRunningProxy(name: string, provider: unknown): Promise<void> {
+export async function notifyRunningProxy(name: string, provider: unknown): Promise<void> {
   // Identity-checked runtime-port lookup: reaches a fallback-port proxy and avoids
   // posting credentials-adjacent config to whatever else answers on config.port.
   const live = await findLiveProxy();
@@ -30,6 +32,19 @@ async function notifyRunningProxy(name: string, provider: unknown): Promise<void
   }
 }
 
+/**
+ * After `runLogin()` has persisted the merged provider (including preserved apiKey /
+ * apiKeyPool / authMode), push that on-disk entry into a running proxy.
+ *
+ * Must not send `OAUTH_PROVIDERS[name].providerConfig`: POST /api/providers replaces the
+ * live entry and saves it, which would drop the preserved key billing state.
+ */
+export async function notifyRunningProxyAfterOAuthLogin(name: string): Promise<void> {
+  const provider = loadConfig().providers[name];
+  if (!provider) return;
+  await notifyRunningProxy(name, provider);
+}
+
 export async function handleLogin(provider?: string): Promise<void> {
   const name = (provider ?? "").trim().toLowerCase();
   if (isPublicOAuthProvider(name)) return handleOAuthLogin(name);
@@ -58,7 +73,7 @@ async function handleOAuthLogin(name: string): Promise<void> {
   } finally {
     rl.close();
   }
-  await notifyRunningProxy(name, OAUTH_PROVIDERS[name].providerConfig);
+  await notifyRunningProxyAfterOAuthLogin(name);
   console.log(`\n✅ Logged in to ${name}. Try: ocx sync`);
 }
 
@@ -67,6 +82,7 @@ export function providerConfigFromKeyLoginProvider(def: KeyLoginProvider, key: s
     adapter: def.adapter,
     baseUrl: baseUrlOverride ?? def.baseUrl,
     apiKey: key,
+    ...(def.apiKeyTransport !== undefined ? { apiKeyTransport: def.apiKeyTransport } : {}),
     ...(def.defaultModel ? { defaultModel: def.defaultModel } : {}),
     ...(def.models ? { models: [...def.models] } : {}),
     ...(def.contextWindow !== undefined ? { contextWindow: def.contextWindow } : {}),
@@ -92,6 +108,12 @@ export function providerConfigFromKeyLoginProvider(def: KeyLoginProvider, key: s
 
 async function handleKeyLogin(name: string): Promise<void> {
   const def = KEY_LOGIN_PROVIDERS[name];
+  const preflightConfig = loadConfig();
+  const namespaceCollision = codexAccountNamespaceProviderCollisionError(preflightConfig.codexAccountNamespaces, name);
+  if (namespaceCollision) {
+    console.error(`Error: ${namespaceCollision}.`);
+    process.exit(1);
+  }
   console.log(`\n🔑 ${def.label} — opening ${def.dashboardUrl} so you can create/copy an API key...`);
   openUrl(def.dashboardUrl);
   const rl = readline.createInterface({ input: process.stdin, output: process.stdout });
@@ -121,6 +143,11 @@ async function handleKeyLogin(name: string): Promise<void> {
   }
   const provider = providerConfigFromKeyLoginProvider(def, key, baseUrl);
   const config = loadConfig();
+  const commitCollision = codexAccountNamespaceProviderCollisionError(config.codexAccountNamespaces, name);
+  if (commitCollision) {
+    console.error(`Error: ${commitCollision}.`);
+    process.exit(1);
+  }
   config.providers[name] = provider;
   saveConfig(config);
   await notifyRunningProxy(name, provider);