pi-provider-kimi-code 0.6.6 → 0.6.7

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-provider-kimi-code",
-  "version": "0.6.6",
+  "version": "0.6.7",
   "description": "Use Kimi Code in Pi without losing the Kimi-specific bits",
   "keywords": [
     "ai-provider",
@@ -49,12 +49,16 @@
     "format:check": "oxfmt . --check",
     "pre-commit": "prek run"
   },
+  "dependencies": {
+    "proper-lockfile": "^4.1.2"
+  },
   "devDependencies": {
-    "@earendil-works/pi-ai": "^0.79.10",
-    "@earendil-works/pi-coding-agent": "^0.79.10",
-    "@earendil-works/pi-tui": "^0.79.10",
+    "@earendil-works/pi-ai": "^0.80.9",
+    "@earendil-works/pi-coding-agent": "^0.80.9",
+    "@earendil-works/pi-tui": "^0.80.9",
     "@j178/prek": "^0.4.4",
     "@types/node": "^24.0.0",
+    "@types/proper-lockfile": "^4.1.4",
     "oxfmt": "^0.44.0",
     "oxlint": "^1.59.0",
     "typescript": "^5.9.0"

package/src/constants.ts

@@ -41,7 +41,7 @@ export function getDefaultBaseUrl(protocol: KimiWireProtocol): string {
 
 export const DEFAULT_BASE_URL = getDefaultBaseUrl(ENV_KIMI_CODE_PROTOCOL);
 
-export const PROVIDER_VERSION = "0.6.6";
+export const PROVIDER_VERSION = "0.6.7";
 
 // Upstream kimi-code CLI version — used in User-Agent and X-Msh-Version
 // headers to match the official client's identity. Update this when

package/src/oauth.ts

@@ -2,13 +2,94 @@
 // reuse, the login / refresh handlers wired into pi's OAuth interface, and
 // the stream-level auth refresh used by the streaming handler.
 
-import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import os from "node:os";
 import { dirname, join } from "node:path";
 
+import { lock as acquireFileLock } from "proper-lockfile";
+
 import type { OAuthCredentials, OAuthLoginCallbacks } from "@earendil-works/pi-ai";
-import type { OAuthCredential } from "@earendil-works/pi-coding-agent";
-import { AuthStorage } from "@earendil-works/pi-coding-agent";
+import * as piAgent from "@earendil-works/pi-coding-agent";
+
+// Structural mirrors of pi's stored credential shapes. The exported names
+// moved between pi versions (pi-coding-agent <=0.79 vs pi-ai >=0.80.8), so
+// these are defined locally to type-check against either surface.
+type StoredCredential = { type: string } & Record<string, unknown>;
+
+export interface StoredOAuthCredential {
+  type: "oauth";
+  access: string;
+  refresh: string;
+  expires: number;
+  [key: string]: unknown;
+}
+
+// pi renamed its auth storage surface in 0.80.8: the exported AuthStorage
+// class (sync get/set under a file lock) was replaced by readStoredCredential
+// for one-off reads, with no exported write path. Detect whichever the host
+// pi provides so the extension keeps loading on both sides of the rename.
+const piAgentRuntime = piAgent as unknown as {
+  getAgentDir(): string;
+  readStoredCredential?(providerId: string): StoredCredential | undefined;
+  AuthStorage?: {
+    create(): {
+      get(provider: string): StoredCredential | undefined;
+      set(provider: string, credential: StoredCredential): void;
+    };
+  };
+};
+
+export function readStoredOAuthCredential(providerId: string): StoredOAuthCredential | null {
+  const credential = piAgentRuntime.readStoredCredential
+    ? piAgentRuntime.readStoredCredential(providerId)
+    : piAgentRuntime.AuthStorage?.create().get(providerId);
+  return credential?.type === "oauth" ? (credential as StoredOAuthCredential) : null;
+}
+
+// Mutual exclusion with pi's own credential writes: use the same library
+// (proper-lockfile) on the same file pi's FileAuthStorageBackend locks, with
+// the same options. The shared implementation keeps the lock mtime fresh
+// 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,
+    stale: 30_000,
+    retries: { retries: 10, factor: 2, minTimeout: 100, maxTimeout: 10_000, randomize: true },
+  });
+}
+
+async function writeStoredCredential(
+  providerId: string,
+  credential: StoredOAuthCredential,
+): Promise<void> {
+  const AuthStorage = piAgentRuntime.AuthStorage;
+  if (AuthStorage) {
+    AuthStorage.create().set(providerId, credential);
+    return;
+  }
+  // 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 authPath = join(piAgentRuntime.getAgentDir(), "auth.json");
+  const release = await acquireAuthFileLock(authPath);
+  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.
+    }
+    data[providerId] = credential;
+    writeFileSync(authPath, JSON.stringify(data, null, 2), { encoding: "utf-8", mode: 0o600 });
+    chmodSync(authPath, 0o600);
+  } finally {
+    await release();
+  }
+}
 
 import { CLIENT_ID, PROVIDER_ID, RETRYABLE_REFRESH_STATUSES, getOAuthHost } from "./constants.ts";
 import { getCommonHeaders } from "./device.ts";
@@ -424,8 +505,8 @@ export async function refreshKimiCodeToken(
 // `expires` is in the past. If the server rotates/revokes the access token
 // before that (common with short-lived session tokens), every request keeps
 // returning 401. The streaming handler invokes refreshKimiAuthToken on the
-// first auth error to force a refresh through AuthStorage (which persists the
-// new credentials under a file lock), and retry once.
+// first auth error to force a refresh and persist the new credentials, and
+// retry once.
 // =============================================================================
 
 export function getKimiApiKey(credentials: OAuthCredentials): string {
@@ -453,9 +534,7 @@ export async function refreshKimiAuthToken(
 ): Promise<string | null> {
   try {
     const kimiCred = readKimiCliCredentials();
-    const storage = AuthStorage.create();
-    const piCred = storage.get(PROVIDER_ID);
-    const piOAuth = piCred?.type === "oauth" ? piCred : null;
+    const piOAuth = readStoredOAuthCredential(PROVIDER_ID);
 
     if (piOAuth) {
       if (piOAuth.access !== currentKey && Date.now() < piOAuth.expires) {
@@ -471,14 +550,14 @@ export async function refreshKimiAuthToken(
         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) {
-          const recovered: OAuthCredential = {
+          const recovered: StoredOAuthCredential = {
             ...piOAuth,
             type: "oauth",
             access: kimiCred.access_token!,
             refresh: kimiCred.refresh_token,
             expires: kimiExpiresMs,
           };
-          storage.set(PROVIDER_ID, recovered);
+          await writeStoredCredential(PROVIDER_ID, recovered);
           console.error("[kimi-coding] auth refresh: recovered newer kimi-code token");
           return recovered.access;
         }
@@ -487,14 +566,14 @@ export async function refreshKimiAuthToken(
       }
 
       const newExpiresMs = Date.now() + refreshed.expires_in * 1000;
-      const newCred: OAuthCredential = {
+      const newCred: StoredOAuthCredential = {
         ...piOAuth,
         type: "oauth",
         access: refreshed.access_token,
         refresh: refreshed.refresh_token,
         expires: newExpiresMs,
       };
-      storage.set(PROVIDER_ID, newCred);
+      await writeStoredCredential(PROVIDER_ID, newCred);
       writeKimiCodeCredentials(refreshed.access_token, refreshed.refresh_token, newExpiresMs);
       console.error("[kimi-coding] auth refresh: new token persisted");
       return newCred.access;

package/src/stream.ts

@@ -12,11 +12,46 @@ import type {
   Model,
   SimpleStreamOptions,
 } from "@earendil-works/pi-ai";
-import {
-  createAssistantMessageEventStream,
-  streamSimpleAnthropic,
-  streamSimpleOpenAICompletions,
-} from "@earendil-works/pi-ai";
+import * as piAi from "@earendil-works/pi-ai";
+import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
+
+type KimiStreamSimple = (
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+) => AssistantMessageEventStream;
+
+// pi 0.80.8 unexported the global streamSimple* functions. Depending on how
+// this module is loaded, the pi-ai root now offers one of three shapes:
+//   - pi extension loader (jiti/Bun): root aliases to the compat entrypoint,
+//     which exposes lazy per-API factories (anthropicMessagesApi et al.)
+//   - pi <=0.79: root only has the legacy streamSimple* globals
+//   - plain Node with pi-ai >=0.80.8 (tests, SDK use): root has neither, but
+//     lazyApi + a dynamic subpath import reaches the same implementation
+// Detect whichever shape is present so the extension loads in all three.
+const piAiRuntime = piAi as unknown as {
+  anthropicMessagesApi?: () => { streamSimple: KimiStreamSimple };
+  openAICompletionsApi?: () => { streamSimple: KimiStreamSimple };
+  streamSimpleAnthropic?: KimiStreamSimple;
+  streamSimpleOpenAICompletions?: KimiStreamSimple;
+  lazyApi?: (load: () => Promise<object>) => { streamSimple: KimiStreamSimple };
+};
+
+// The ./api/* subpath exports only exist on pi-ai >=0.80.8, so a literal
+// specifier would fail type resolution against older pi-ai. Computed
+// specifiers keep the imports invisible to TypeScript; they are only ever
+// evaluated in plain-Node environments where the subpaths do resolve.
+const anthropicMessagesModule = "@earendil-works/pi-ai/api/anthropic-messages";
+const openAICompletionsModule = "@earendil-works/pi-ai/api/openai-completions";
+
+const streamSimpleAnthropic: KimiStreamSimple =
+  piAiRuntime.anthropicMessagesApi?.().streamSimple ??
+  piAiRuntime.streamSimpleAnthropic ??
+  piAiRuntime.lazyApi?.(() => import(anthropicMessagesModule)).streamSimple!;
+const streamSimpleOpenAICompletions: KimiStreamSimple =
+  piAiRuntime.openAICompletionsApi?.().streamSimple ??
+  piAiRuntime.streamSimpleOpenAICompletions ??
+  piAiRuntime.lazyApi?.(() => import(openAICompletionsModule)).streamSimple!;
 import {
   DEFAULT_KIMI_CODE_CONFIG,
   type KimiCodeConfig,
@@ -162,7 +197,16 @@ export function resolveKimiApiKey(apiKey: string | undefined): string {
   return apiKey || process.env.KIMI_API_KEY || "";
 }
 
-export function mergeKimiRequestHeaders(headers?: Record<string, string>): Record<string, string> {
+// SimpleStreamOptions.headers is Record<string, string> on pi <=0.79 and
+// ProviderHeaders (Record<string, string | null>, null = suppress the header
+// when pi-ai assembles the request) on >=0.80.8. Accept and preserve both
+// shapes; defined locally because ProviderHeaders is not exported on old pi.
+type HeaderMap = Record<string, string | null>;
+
+export function mergeKimiRequestHeaders(headers?: HeaderMap): HeaderMap {
+  // Caller values override our Kimi defaults. Nulls are retained rather than
+  // deleted so the suppression semantics reach pi-ai, which strips null
+  // entries (including its own defaults) when building the request.
   return { ...getKimiProviderHeaders(), ...headers };
 }
 
@@ -231,7 +275,9 @@ export function streamSimpleKimi(
     return {
       ...options,
       ...apiKeyOverride,
-      headers: mergeKimiRequestHeaders(options?.headers),
+      // Cast to whatever headers shape the installed pi-ai declares:
+      // Record<string, string> on pi <=0.79, ProviderHeaders on >=0.80.8.
+      headers: mergeKimiRequestHeaders(options?.headers) as SimpleStreamOptions["headers"],
       onPayload: async (payload, modelData) => {
         let nextPayload: unknown = payload;
 

package/src/tools/common.ts

@@ -1,9 +1,13 @@
-import { AuthStorage, type AgentToolResult } from "@earendil-works/pi-coding-agent";
+import type { AgentToolResult } from "@earendil-works/pi-coding-agent";
 import { truncateToWidth } from "@earendil-works/pi-tui";
 
 import { PROVIDER_ID, getBaseUrl } from "../constants.ts";
 import { getKimiProviderHeaders } from "../device.ts";
-import { KIMI_LOGIN_REQUIRED_MESSAGE, refreshKimiAuthToken } from "../oauth.ts";
+import {
+  KIMI_LOGIN_REQUIRED_MESSAGE,
+  readStoredOAuthCredential,
+  refreshKimiAuthToken,
+} from "../oauth.ts";
 
 export const KIMI_TOOL_TIMEOUT_MS = 180_000;
 
@@ -19,9 +23,7 @@ export interface BuildKimiToolOptions {
 }
 
 function defaultGetAccessToken(): string | null {
-  const credential = AuthStorage.create().get(PROVIDER_ID);
-  if (!credential || credential.type !== "oauth") return null;
-  return credential.access;
+  return readStoredOAuthCredential(PROVIDER_ID)?.access ?? null;
 }
 
 export function buildKimiToolDeps(options: BuildKimiToolOptions = {}): KimiToolDeps {

package/src/usage.ts

@@ -1,8 +1,6 @@
-import { AuthStorage } from "@earendil-works/pi-coding-agent";
-
 import { PROVIDER_ID, getBaseUrl } from "./constants.ts";
 import { getKimiProviderHeaders } from "./device.ts";
-import { refreshKimiAuthToken } from "./oauth.ts";
+import { readStoredOAuthCredential, refreshKimiAuthToken } from "./oauth.ts";
 
 const MEMBERSHIP_LEVEL_NAMES: Record<string, string> = {
   LEVEL_FREE: "Free",
@@ -52,8 +50,8 @@ interface BoosterWalletInfo {
 const FIXED_POINT_CENTS = 1_000_000n;
 
 export function getKimiUsageToken(): string | null {
-  const credential = AuthStorage.create().get(PROVIDER_ID);
-  if (credential?.type === "oauth" && credential.access) return credential.access;
+  const credential = readStoredOAuthCredential(PROVIDER_ID);
+  if (credential?.access) return credential.access;
   const apiKey = process.env.KIMI_API_KEY?.trim();
   return apiKey || null;
 }