pi-provider-kimi-code 0.6.7 → 0.6.9

package/src/oauth.ts

@@ -52,43 +52,107 @@ export function readStoredOAuthCredential(providerId: string): StoredOAuthCreden
 // while held, recovers stale locks from crashed processes, and retries for
 // ~40s so a live holder running a slow network refresh inside the lock is
 // waited out.
-async function acquireAuthFileLock(authPath: string): Promise<() => Promise<void>> {
-  return acquireFileLock(authPath, {
-    realpath: false,
+async function acquireSharedFileLock(
+  path: string,
+  signal: AbortSignal | undefined,
+  options: { maxRetries: number; stale: number; minTimeout: number; maxTimeout: number },
+): Promise<() => Promise<void>> {
+  for (let attempt = 0; ; attempt++) {
+    signal?.throwIfAborted();
+    try {
+      return await acquireFileLock(path, {
+        realpath: false,
+        stale: options.stale,
+        // Retry via the outer abortable sleep below. proper-lockfile's own
+        // retry loop is not signal-aware and would otherwise ignore an abort
+        // for the full configured retry window.
+        retries: 0,
+      });
+    } catch (error) {
+      const code =
+        typeof error === "object" && error !== null && "code" in error
+          ? String((error as { code?: unknown }).code)
+          : undefined;
+      if (code !== "ELOCKED" || attempt >= options.maxRetries) throw error;
+      const baseDelay = Math.min(options.minTimeout * 2 ** attempt, options.maxTimeout);
+      await sleep(baseDelay * (1 + Math.random()), signal);
+    }
+  }
+}
+
+async function acquireAuthFileLock(
+  authPath: string,
+  signal?: AbortSignal,
+): Promise<() => Promise<void>> {
+  return acquireSharedFileLock(authPath, signal, {
+    maxRetries: 10,
     stale: 30_000,
-    retries: { retries: 10, factor: 2, minTimeout: 100, maxTimeout: 10_000, randomize: true },
+    minTimeout: 100,
+    maxTimeout: 10_000,
   });
 }
 
-async function writeStoredCredential(
-  providerId: string,
-  credential: StoredOAuthCredential,
-): Promise<void> {
-  const AuthStorage = piAgentRuntime.AuthStorage;
-  if (AuthStorage) {
-    AuthStorage.create().set(providerId, credential);
-    return;
+async function acquireKimiCodeRefreshLock(signal?: AbortSignal): Promise<() => Promise<void>> {
+  if (process.platform === "win32" || process.env.KIMI_DISABLE_OAUTH_LOCK === "1") {
+    return async () => {};
   }
-  // pi >=0.80.8: no exported write path, so do a locked read-modify-write of
-  // auth.json in pi's on-disk format (indent 2). The read happens inside the
-  // lock so concurrent updates from other processes are never clobbered, and
-  // the explicit chmod covers pre-existing files (writeFileSync's mode only
-  // applies on creation).
+  const target = join(
+    process.env.KIMI_CODE_HOME || join(os.homedir(), ".kimi-code"),
+    "oauth",
+    "kimi-code",
+  );
+  mkdirSync(dirname(target), { recursive: true, mode: 0o700 });
+  writeFileSync(target, "", { encoding: "utf-8", flag: "a" });
+  return acquireSharedFileLock(target, signal, {
+    maxRetries: 120,
+    stale: 5_000,
+    minTimeout: 500,
+    maxTimeout: 1_000,
+  });
+}
+
+interface LockedOAuthCredential {
+  credential: StoredOAuthCredential | null;
+  store(credential: StoredOAuthCredential): void;
+  release(): Promise<void>;
+}
+
+async function lockStoredOAuthCredential(
+  providerId: string,
+  signal?: AbortSignal,
+): Promise<LockedOAuthCredential> {
+  // Hold the same lock Pi uses across the complete read-refresh-write
+  // transaction. Kimi rotates refresh tokens, so locking only the write lets
+  // concurrent Pi processes exchange the same token twice.
   const authPath = join(piAgentRuntime.getAgentDir(), "auth.json");
-  const release = await acquireAuthFileLock(authPath);
+  mkdirSync(dirname(authPath), { recursive: true, mode: 0o700 });
+  if (!existsSync(authPath)) {
+    writeFileSync(authPath, "{}", { encoding: "utf-8", mode: 0o600 });
+    chmodSync(authPath, 0o600);
+  }
+  const release = await acquireAuthFileLock(authPath, signal);
+  let data: Record<string, unknown> = {};
   try {
-    let data: Record<string, unknown> = {};
-    try {
-      data = JSON.parse(readFileSync(authPath, "utf-8"));
-    } catch {
-      // Missing or unreadable auth.json — start from an empty store.
+    const parsed = JSON.parse(readFileSync(authPath, "utf-8")) as unknown;
+    if (typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)) {
+      data = parsed as Record<string, unknown>;
     }
-    data[providerId] = credential;
-    writeFileSync(authPath, JSON.stringify(data, null, 2), { encoding: "utf-8", mode: 0o600 });
-    chmodSync(authPath, 0o600);
-  } finally {
-    await release();
+  } catch {
+    // Missing or unreadable auth.json — treat it as an empty store.
   }
+  const stored = data[providerId] as StoredCredential | undefined;
+  return {
+    credential: stored?.type === "oauth" ? (stored as unknown as StoredOAuthCredential) : null,
+    store(credential) {
+      data[providerId] = credential;
+      writeFileSync(authPath, JSON.stringify(data, null, 2), {
+        encoding: "utf-8",
+        mode: 0o600,
+      });
+      chmodSync(authPath, 0o600);
+    },
+    release,
+  };
 }
 
 import { CLIENT_ID, PROVIDER_ID, RETRYABLE_REFRESH_STATUSES, getOAuthHost } from "./constants.ts";
@@ -339,14 +403,10 @@ function readKimiCliCredentials(): KimiCliCredentialsFile | null {
   return null;
 }
 
-function kimiCodeCredentialExists(): boolean {
-  return existsSync(getKimiCodeCredentialPath());
-}
-
 function writeKimiCodeCredentials(access: string, refresh: string, expiresMs: number): void {
   const path = getKimiCodeCredentialPath();
   const dir = dirname(path);
-  mkdirSync(dir, { recursive: true });
+  mkdirSync(dir, { recursive: true, mode: 0o700 });
   const data: KimiCliCredentialsFile = {
     access_token: access,
     refresh_token: refresh,
@@ -355,7 +415,24 @@ function writeKimiCodeCredentials(access: string, refresh: string, expiresMs: nu
     token_type: "Bearer",
     expires_in: Math.max(0, Math.floor((expiresMs - Date.now()) / 1000)),
   };
-  writeFileSync(path, `${JSON.stringify(data, null, 2)}\n`, "utf-8");
+  writeFileSync(path, `${JSON.stringify(data, null, 2)}\n`, { encoding: "utf-8", mode: 0o600 });
+  chmodSync(path, 0o600);
+}
+
+interface LockedKimiCodeCredentials {
+  credential: KimiCliCredentialsFile | null;
+  store(access: string, refresh: string, expiresMs: number): void;
+  release(): Promise<void>;
+}
+
+async function lockKimiCodeCredentials(signal?: AbortSignal): Promise<LockedKimiCodeCredentials> {
+  const release = await acquireKimiCodeRefreshLock(signal);
+  const credential = readKimiCliCredentials();
+  return {
+    credential,
+    store: writeKimiCodeCredentials,
+    release,
+  };
 }
 
 async function tryReuseKimiCliCredentials(
@@ -476,26 +553,83 @@ export async function loginKimiCode(callbacks: OAuthLoginCallbacks): Promise<Kim
 export async function refreshKimiCodeToken(
   credentials: OAuthCredentials,
 ): Promise<KimiOAuthCredentials> {
-  const kimiCred = kimiCodeCredentialExists() ? readKimiCliCredentials() : null;
-  let token: Awaited<ReturnType<typeof refreshAccessToken>>;
+  const kimiLocked = await lockKimiCodeCredentials();
   try {
-    token = await refreshAccessToken(credentials.refresh);
-  } catch (error) {
-    if (kimiCred?.refresh_token && kimiCred.refresh_token !== credentials.refresh) {
-      token = await refreshAccessToken(kimiCred.refresh_token);
-    } else {
+    const initialKimiCred = kimiLocked.credential;
+    const initialKimiExpiresMs = (initialKimiCred?.expires_at ?? 0) * 1000;
+    const peerHasNewerValidAccess =
+      initialKimiCred?.access_token !== undefined &&
+      initialKimiCred.access_token !== credentials.access &&
+      Date.now() < initialKimiExpiresMs;
+    const activeRefreshToken = peerHasNewerValidAccess
+      ? initialKimiCred.refresh_token!
+      : credentials.refresh;
+    const activeAccessToken = peerHasNewerValidAccess
+      ? initialKimiCred.access_token!
+      : credentials.access;
+    if (peerHasNewerValidAccess && initialKimiCred.refresh_token !== credentials.refresh) {
+      const extras = await discoverKimiModelMetadata(activeAccessToken);
+      return {
+        access: activeAccessToken,
+        refresh: activeRefreshToken,
+        expires: initialKimiExpiresMs,
+        ...extras,
+      };
+    }
+    let token: Awaited<ReturnType<typeof refreshAccessToken>>;
+    try {
+      token = await refreshAccessToken(activeRefreshToken);
+    } catch (error) {
+      // Lock-disabled platforms (Windows / KIMI_DISABLE_OAUTH_LOCK=1) use
+      // storage re-read as the only coordinator. A peer may have rotated the
+      // shared sidecar while this process was exchanging the old refresh
+      // token; recover from that newer valid credential instead of surfacing
+      // the stale-token invalid_grant as a login failure.
+      await sleep(100);
+      const peerCredential = readKimiCliCredentials();
+      const peerAccess = peerCredential?.access_token;
+      const peerRefresh = peerCredential?.refresh_token;
+      const peerExpiresMs = (peerCredential?.expires_at ?? 0) * 1000;
+      if (
+        peerAccess !== undefined &&
+        peerRefresh !== undefined &&
+        peerAccess !== credentials.access &&
+        Date.now() < peerExpiresMs
+      ) {
+        const extras = await discoverKimiModelMetadata(peerAccess);
+        return {
+          access: peerAccess,
+          refresh: peerRefresh,
+          expires: peerExpiresMs,
+          ...extras,
+        };
+      }
       throw error;
     }
+    const expiresMs = Date.now() + token.expires_in * 1000;
+    const kimiCredentialAfterRefresh = readKimiCliCredentials();
+    const peerAccess = kimiCredentialAfterRefresh?.access_token;
+    const peerRefresh = kimiCredentialAfterRefresh?.refresh_token;
+    const peerExpiresMs = (kimiCredentialAfterRefresh?.expires_at ?? 0) * 1000;
+    const accessToken =
+      peerAccess !== undefined && peerAccess !== activeAccessToken && Date.now() < peerExpiresMs
+        ? peerAccess
+        : token.access_token;
+    const refreshToken = accessToken === peerAccess ? peerRefresh! : token.refresh_token;
+    const effectiveExpiresMs = accessToken === peerAccess ? peerExpiresMs : expiresMs;
+    if (accessToken !== peerAccess) {
+      kimiLocked.store(accessToken, refreshToken, effectiveExpiresMs);
+    }
+    const extras = await discoverKimiModelMetadata(accessToken);
+    return {
+      access: accessToken,
+      refresh: refreshToken,
+      expires: effectiveExpiresMs,
+      ...extras,
+    };
+  } finally {
+    await kimiLocked.release().catch(() => {});
   }
-  const expiresMs = Date.now() + token.expires_in * 1000;
-  writeKimiCodeCredentials(token.access_token, token.refresh_token, expiresMs);
-  const extras = await discoverKimiModelMetadata(token.access_token);
-  return {
-    access: token.access_token,
-    refresh: token.refresh_token,
-    expires: expiresMs,
-    ...extras,
-  };
 }
 
 // =============================================================================
@@ -524,7 +658,13 @@ export function isKimiAuthErrorMessage(message: unknown): boolean {
     /\b401\b/.test(text) ||
     text.includes("unauthorized") ||
     text.includes("incorrect api key") ||
-    text.includes("invalid api key")
+    text.includes("invalid api key") ||
+    // Kimi's production 401 body: "The API Key appears to be invalid or may
+    // have expired..." with type "invalid_authentication_error". Neither
+    // contains the substrings above, so without these the digit fallback is
+    // the only match and vanishes when a transport strips the status prefix.
+    text.includes("api key appears to be invalid") ||
+    text.includes("invalid_authentication_error")
   );
 }
 
@@ -532,21 +672,23 @@ export async function refreshKimiAuthToken(
   currentKey: string,
   options: { signal?: AbortSignal } = {},
 ): Promise<string | null> {
+  let locked: LockedOAuthCredential | undefined;
+  let kimiLocked: LockedKimiCodeCredentials | undefined;
   try {
-    const kimiCred = readKimiCliCredentials();
-    const piOAuth = readStoredOAuthCredential(PROVIDER_ID);
+    locked = await lockStoredOAuthCredential(PROVIDER_ID, options.signal);
+    const piOAuth = locked.credential;
 
     if (piOAuth) {
       if (piOAuth.access !== currentKey && Date.now() < piOAuth.expires) {
-        console.error("[kimi-coding] auth refresh: reusing newer on-disk token");
         return piOAuth.access;
       }
 
-      console.error("[kimi-coding] auth refresh: requesting new access token");
       let refreshed: Awaited<ReturnType<typeof refreshAccessToken>>;
+      kimiLocked = await lockKimiCodeCredentials(options.signal);
       try {
         refreshed = await refreshAccessToken(piOAuth.refresh, { signal: options.signal });
       } catch (error) {
+        const kimiCred = readKimiCliCredentials();
         if (!kimiCred?.refresh_token || kimiCred.refresh_token === piOAuth.refresh) throw error;
         const kimiExpiresMs = (kimiCred.expires_at ?? 0) * 1000;
         if (kimiCred.access_token !== currentKey && Date.now() < kimiExpiresMs) {
@@ -557,15 +699,30 @@ export async function refreshKimiAuthToken(
             refresh: kimiCred.refresh_token,
             expires: kimiExpiresMs,
           };
-          await writeStoredCredential(PROVIDER_ID, recovered);
-          console.error("[kimi-coding] auth refresh: recovered newer kimi-code token");
+          locked.store(recovered);
           return recovered.access;
         }
-        console.error("[kimi-coding] auth refresh: pi token rejected, trying kimi-code token");
         refreshed = await refreshAccessToken(kimiCred.refresh_token, { signal: options.signal });
       }
 
       const newExpiresMs = Date.now() + refreshed.expires_in * 1000;
+      const kimiCredentialAfterRefresh = readKimiCliCredentials();
+      const peerAccess = kimiCredentialAfterRefresh?.access_token;
+      const peerExpiresMs = (kimiCredentialAfterRefresh?.expires_at ?? 0) * 1000;
+      const peerIsNewerValid =
+        peerAccess !== undefined && peerAccess !== currentKey && Date.now() < peerExpiresMs;
+      if (peerIsNewerValid) {
+        const peerCred: StoredOAuthCredential = {
+          ...piOAuth,
+          type: "oauth",
+          access: peerAccess,
+          refresh: kimiCredentialAfterRefresh!.refresh_token!,
+          expires: peerExpiresMs,
+        };
+        locked.store(peerCred);
+        return peerCred.access;
+      }
+      kimiLocked.store(refreshed.access_token, refreshed.refresh_token, newExpiresMs);
       const newCred: StoredOAuthCredential = {
         ...piOAuth,
         type: "oauth",
@@ -573,32 +730,32 @@ export async function refreshKimiAuthToken(
         refresh: refreshed.refresh_token,
         expires: newExpiresMs,
       };
-      await writeStoredCredential(PROVIDER_ID, newCred);
-      writeKimiCodeCredentials(refreshed.access_token, refreshed.refresh_token, newExpiresMs);
-      console.error("[kimi-coding] auth refresh: new token persisted");
+      locked.store(newCred);
       return newCred.access;
     }
 
+    kimiLocked = await lockKimiCodeCredentials(options.signal);
+    const kimiCred = kimiLocked.credential;
     if (!kimiCred) {
-      console.error(
-        `[kimi-coding] auth refresh skipped: no OAuth credentials for ${PROVIDER_ID} on disk`,
-      );
       return null;
     }
 
     const kimiExpiresMs = (kimiCred.expires_at ?? 0) * 1000;
     if (kimiCred.access_token !== currentKey && Date.now() < kimiExpiresMs) {
-      console.error("[kimi-coding] auth refresh: reusing newer kimi-code token");
       return kimiCred.access_token!;
     }
 
-    console.error("[kimi-coding] auth refresh: requesting new access token via kimi-code");
     const refreshed = await refreshAccessToken(kimiCred.refresh_token!, {
       signal: options.signal,
     });
     const newExpiresMs = Date.now() + refreshed.expires_in * 1000;
-    writeKimiCodeCredentials(refreshed.access_token, refreshed.refresh_token, newExpiresMs);
-    console.error("[kimi-coding] auth refresh: new token persisted to kimi-code");
+    const kimiCredentialAfterRefresh = readKimiCliCredentials();
+    const peerAccess = kimiCredentialAfterRefresh?.access_token;
+    const peerExpiresMs = (kimiCredentialAfterRefresh?.expires_at ?? 0) * 1000;
+    if (peerAccess !== undefined && peerAccess !== currentKey && Date.now() < peerExpiresMs) {
+      return peerAccess;
+    }
+    kimiLocked.store(refreshed.access_token, refreshed.refresh_token, newExpiresMs);
     return refreshed.access_token;
   } catch (err) {
     if (options.signal?.aborted) return null;
@@ -608,5 +765,12 @@ export async function refreshKimiAuthToken(
       console.error("[kimi-coding] auth refresh failed:", err);
     }
     return null;
+  } finally {
+    if (kimiLocked) {
+      await kimiLocked.release().catch(() => {});
+    }
+    if (locked) {
+      await locked.release().catch(() => {});
+    }
   }
 }