pi-provider-kimi-code 0.2.2 → 0.3.1

package/README.md

@@ -4,13 +4,13 @@
 
 **Reuse your [Kimi Code Plan](https://www.kimi.com/code/docs/en/) inside [pi-coding-agent](https://pi.dev/)** — no separate API credits, no second billing dashboard. Every request draws from your Kimi Code membership quota (the 5-hour token bucket) instead of billing per-token on the Moonshot Open Platform.
 
-It's **"Claude Code for Kimi"** — log in with your Kimi account, with `KIMI_API_KEY` still supported as a fallback for CI. You get Kimi for Coding, **Kimi K2.6 Code Preview**, Kimi K2.5, Kimi K2 Thinking Turbo, and the umbrella `kimi-code` model with a 262K-token context window, automatic prompt caching, automatic large-image and video upload, and compatibility with both Anthropic-style and OpenAI-style clients.
+It's **"Claude Code for Kimi"** — log in with your Kimi account, with `KIMI_API_KEY` still supported as a fallback for CI. You get `kimi-for-coding` (the Kimi-k2.6 latest alias, backed by the same Kimi Code Plan that covers K2.6 and K2.5) with a 262K-token context window, automatic prompt caching, automatic large-image and video upload, and compatibility with both Anthropic-style and OpenAI-style clients.
 
 ## Who is this for?
 
 - You already pay for a **[Kimi Code Plan](https://www.kimi.com/code/docs/en/)** and want to use it inside `pi-coding-agent` instead of the official `kimi-cli` — see [MoonshotAI/kimi-cli#757](https://github.com/MoonshotAI/kimi-cli/issues/757) for the canonical feature request this extension answers.
 - You want **"Claude Code for Kimi"**: log in with your Kimi account instead of buying separate API credits. (`KIMI_API_KEY` is also supported as a fallback for CI.)
-- You're in the [pi / pi-mono](https://github.com/badlogic/pi-mono) ecosystem and want **Kimi K2.6 Code Preview**, **Kimi K2.5**, or **Kimi K2 Thinking Turbo** as a pi provider.
+- You're in the [pi / pi-mono](https://github.com/badlogic/pi-mono) ecosystem and want **Kimi K2.6** (and the K2.5 coverage that comes with the Code Plan) as a pi provider.
 
 Pay-per-token via `KIMI_API_KEY` also works if you just want to try Kimi in CI or without a subscription.
 
@@ -72,18 +72,16 @@ KIMI_API_KEY=sk-... pi
 
 ## Models
 
-| ID                         | Name                   | Reasoning | Input              | Context | Max Output |
-| -------------------------- | ---------------------- | --------- | ------------------ | ------- | ---------- |
-| `kimi-for-coding`          | Kimi for Coding        | yes       | text, image, video | 262 144 | 32 000     |
-| `kimi-k2.6-code-preview`   | Kimi K2.6 Code Preview | yes       | text, image, video | 262 144 | 32 000     |
-| `kimi-code`                | Kimi Code (kimi-k2.5)  | yes       | text, image        | 262 144 | 32 000     |
-| `kimi-k2.5`                | Kimi K2.5              | yes       | text, image        | 262 144 | 32 000     |
-| `kimi-k2-thinking-turbo`   | Kimi K2 Thinking Turbo | yes       | text               | 262 144 | 32 000     |
+| ID                | Name            | Reasoning | Input              | Context | Max Output |
+| ----------------- | --------------- | --------- | ------------------ | ------- | ---------- |
+| `kimi-for-coding` | Kimi for Coding | yes       | text, image, video | 262 144 | 32 000     |
 
-Select a model inside `pi`:
+`kimi-for-coding` is the latest alias on the Kimi Code Plan — today it points at Kimi-k2.6. The Plan itself still covers K2.5; the upstream routes older model IDs to the current alias, but this provider only publishes the canonical `kimi-for-coding` entry.
+
+Select it inside `pi`:
 
 ```
-/model kimi-coding/kimi-code
+/model kimi-coding/kimi-for-coding
 ```
 
 ## Environment variables

package/index.ts

@@ -29,7 +29,8 @@ import {
   streamSimpleOpenAICompletions,
   AssistantMessageEventStream,
 } from "@mariozechner/pi-ai";
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, OAuthCredential } from "@mariozechner/pi-coding-agent";
+import { AuthStorage } from "@mariozechner/pi-coding-agent";
 
 // =============================================================================
 // Constants
@@ -39,6 +40,9 @@ const CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098";
 const DEFAULT_OAUTH_HOST = "https://auth.kimi.com";
 const PROTOCOL =
   process.env.KIMI_CODE_PROTOCOL === "openai" ? "openai-completions" : "anthropic-messages";
+// Use a custom api identifier so this provider never conflicts with the
+// built-in "anthropic-messages" or "openai-completions" stream handlers.
+const KIMI_API_TYPE = PROTOCOL === "openai-completions" ? "kimi-openai-completions" : "kimi-anthropic-messages";
 const DEFAULT_BASE_URL =
   PROTOCOL === "openai-completions"
     ? "https://api.kimi.com/coding/v1"
@@ -111,11 +115,32 @@ function getDeviceModel(): string {
     return version && arch ? `macOS ${version} ${arch}` : `macOS ${arch}`;
   }
   if (platform === "win32") {
-    const release = os.release();
+    // Only show the major release (e.g. "Windows 10", "Windows 11") to match
+    // the upstream kimi-cli behavior in `def _device_model() -> str:`.
+    // See: https://github.com/MoonshotAI/kimi-cli/blob/main/src/kimi_cli/auth/oauth.py
+    const parts = os.release().split(".");
+    let release = parts[0];
+    if (release === "10" && parts.length >= 3) {
+      const build = parseInt(parts[2], 10);
+      if (!isNaN(build) && build >= 22000) {
+        release = "11";
+      }
+    }
     return release && arch ? `Windows ${release} ${arch}` : `Windows ${arch}`;
   }
+  // Normalize platform name to match Python `platform.system()` casing used by
+  // upstream kimi-cli (e.g. "Linux", "FreeBSD", "SunOS"), since Node's
+  // `process.platform` is always lower-case.
+  const SYSTEM_NAME: Record<string, string> = {
+    aix: "AIX",
+    freebsd: "FreeBSD",
+    linux: "Linux",
+    openbsd: "OpenBSD",
+    sunos: "SunOS",
+  };
+  const system = SYSTEM_NAME[platform] ?? platform;
   const release = os.release();
-  return release && arch ? `${platform} ${release} ${arch}` : `${platform} ${arch}`;
+  return release && arch ? `${system} ${release} ${arch}` : `${system} ${arch}`;
 }
 
 function asciiHeaderValue(value: string, fallback = "unknown"): string {
@@ -709,6 +734,54 @@ async function* filterEmptyResponseStream(
   }
 }
 
+// =============================================================================
+// Auth refresh: recover from server-side token invalidation
+// =============================================================================
+// pi-coding-agent only refreshes an OAuth token when the locally cached
+// `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. We detect that situation by inspecting the first event of
+// the upstream stream, force a refresh through AuthStorage (which persists
+// the new credentials under a file lock), and retry once.
+
+const PROVIDER_ID = "kimi-coding";
+
+async function refreshKimiAuthToken(currentKey: string): Promise<string | null> {
+  try {
+    const storage = AuthStorage.create();
+    const cred = storage.get(PROVIDER_ID);
+    if (!cred || cred.type !== "oauth") {
+      console.error(
+        `[kimi-coding] auth refresh skipped: no OAuth credentials for ${PROVIDER_ID} on disk`,
+      );
+      return null;
+    }
+
+    // If disk already has a different valid token (e.g., another process or a
+    // previous retry refreshed it while the caller's in-memory cache went
+    // stale), reuse it without hitting the OAuth endpoint.
+    if (cred.access !== currentKey && Date.now() < cred.expires) {
+      console.error("[kimi-coding] auth refresh: reusing newer on-disk token");
+      return cred.access;
+    }
+
+    console.error("[kimi-coding] auth refresh: requesting new access token");
+    const refreshed = await refreshAccessToken(cred.refresh);
+    const newCred: OAuthCredential = {
+      type: "oauth",
+      access: refreshed.access_token,
+      refresh: refreshed.refresh_token,
+      expires: Date.now() + refreshed.expires_in * 1000,
+    };
+    storage.set(PROVIDER_ID, newCred);
+    console.error("[kimi-coding] auth refresh: new token persisted");
+    return newCred.access;
+  } catch (err) {
+    console.error("[kimi-coding] auth refresh failed:", err);
+    return null;
+  }
+}
+
 // =============================================================================
 // Stream wrapper: orchestrates payload mutation + event filter
 // =============================================================================
@@ -723,7 +796,7 @@ function streamSimpleKimi(
   options?: SimpleStreamOptions,
 ): AssistantMessageEventStream {
   const filtered = new AssistantMessageEventStream();
-  const apiKey = options?.apiKey || process.env.KIMI_API_KEY || "";
+  const initialKey = options?.apiKey || process.env.KIMI_API_KEY || "";
 
   const cacheKeyOverride = (
     options as (SimpleStreamOptions & { prompt_cache_key?: unknown }) | undefined
@@ -731,61 +804,143 @@ function streamSimpleKimi(
   const cacheKey = (typeof cacheKeyOverride === "string" && cacheKeyOverride) || options?.sessionId;
   const cacheRetention = resolveCacheRetention(options?.cacheRetention);
   const envOverrides = readEnvOverrides();
-  const upload: Uploader | undefined = apiKey
-    ? (mimeType, data) => uploadKimiFile(apiKey, mimeType, data)
-    : undefined;
-
   const originalOnPayload = options?.onPayload;
-  const patchedOptions: SimpleStreamOptions = {
-    ...options,
-    onPayload: async (payload, modelData) => {
-      let nextPayload: unknown = payload;
-
-      if (isRecord(nextPayload)) {
-        await applyKimiPayloadMutations(nextPayload, {
-          api: model.api,
-          upload,
-          cacheKey,
-          cacheRetention,
-          reasoning: options?.reasoning,
-          envOverrides,
-        });
-      }
 
-      if (originalOnPayload) {
-        const res = await originalOnPayload(nextPayload, modelData);
-        if (res !== undefined) nextPayload = res;
-      }
+  const buildPatchedOptions = (apiKey: string): SimpleStreamOptions => {
+    const upload: Uploader | undefined = apiKey
+      ? (mimeType, data) => uploadKimiFile(apiKey, mimeType, data)
+      : undefined;
+    // Only forward apiKey if we actually have one — never override the
+    // caller's credential (e.g. Claude Code OAuth token) with an empty string.
+    const apiKeyOverride = apiKey ? { apiKey } : {};
+    return {
+      ...options,
+      ...apiKeyOverride,
+      onPayload: async (payload, modelData) => {
+        let nextPayload: unknown = payload;
+
+        if (isRecord(nextPayload)) {
+          await applyKimiPayloadMutations(nextPayload, {
+            api: model.api,
+            upload,
+            cacheKey,
+            cacheRetention,
+            reasoning: options?.reasoning,
+            envOverrides,
+          });
+        }
 
-      return nextPayload;
-    },
+        if (originalOnPayload) {
+          const res = await originalOnPayload(nextPayload, modelData);
+          if (res !== undefined) nextPayload = res;
+        }
+
+        return nextPayload;
+      },
+    };
   };
 
   void (async () => {
-    const upstream =
-      model.api === "openai-completions"
-        ? streamSimpleOpenAICompletions(
-            model as Model<"openai-completions">,
-            context,
-            patchedOptions,
-          )
-        : streamSimpleAnthropic(model as Model<"anthropic-messages">, context, patchedOptions);
-
-    try {
-      for await (const event of filterEmptyResponseStream(upstream)) {
-        filtered.push(event);
+    let attempt = 0;
+    let currentKey = initialKey;
+
+    while (true) {
+      const patchedOptions = buildPatchedOptions(currentKey);
+      // Route by the module-level PROTOCOL, not model.api, since we register
+      // with a custom api type (kimi-openai-completions / kimi-anthropic-messages)
+      // to avoid overriding the built-in Anthropic/OpenAI stream handlers.
+      const upstream =
+        PROTOCOL === "openai-completions"
+          ? streamSimpleOpenAICompletions(
+              model as Model<"openai-completions">,
+              context,
+              patchedOptions,
+            )
+          : streamSimpleAnthropic(model as Model<"anthropic-messages">, context, patchedOptions);
+
+      let shouldRetry = false;
+      let prefixBuffer: AssistantMessageEvent[] = [];
+
+      try {
+        for await (const event of filterEmptyResponseStream(upstream)) {
+          // streamAnthropic emits a synthetic "start" event synchronously,
+          // before the for-await loop begins iterating and therefore before
+          // the HTTP request is actually made.  If the request 401s, the loop
+          // throws and the catch block emits "error".  Without buffering, the
+          // "start" event (which carries an empty AssistantMessage) leaks into
+          // the session history and the TUI, leaving a phantom empty assistant
+          // bubble.  We buffer "start" events and only flush them once we see
+          // a non-error event that proves the stream is alive.
+          if (event.type === "start") {
+            prefixBuffer.push(event);
+            continue;
+          }
+
+          // Speculative OAuth refresh on the first error.  We retry once so
+          // that short-lived Kimi tokens invalidated by the server before the
+          // local expires timestamp lapses don't surface as raw 401s to the
+          // user.  Non-auth errors still trigger one wasted refresh, but the
+          // retried request fails the same way and we forward it.
+          if (attempt === 0 && event.type === "error") {
+            console.error(
+              `[kimi-coding] upstream error on first event, attempting refresh: ${event.error?.errorMessage?.slice(0, 200)}`,
+            );
+            const refreshed = await refreshKimiAuthToken(currentKey);
+            if (refreshed && refreshed !== currentKey) {
+              console.error("[kimi-coding] retrying stream with refreshed token");
+              currentKey = refreshed;
+              shouldRetry = true;
+              break; // discard prefixBuffer — don't leak the stale start
+            }
+            console.error(
+              "[kimi-coding] refresh did not yield a new token, forwarding original error",
+            );
+          }
+
+          // First non-start, non-retry event: flush buffered prefix, then
+          // stream normally.
+          for (const e of prefixBuffer) filtered.push(e);
+          prefixBuffer = [];
+          filtered.push(event);
+        }
+
+        // Stream ended normally: flush any remaining buffered starts.
+        for (const e of prefixBuffer) filtered.push(e);
+      } catch (err) {
+        // Upstream threw rather than emitting a stream `error` event. This can
+        // be the same stale-token 401 surfaced as an exception (depending on
+        // the SDK path / network layer), so mirror the in-stream refresh
+        // branch: on attempt 0, try one OAuth refresh + retry. Either way,
+        // discard `prefixBuffer` — we never confirmed the stream actually
+        // started, and flushing the buffered `start` would resurrect the
+        // phantom empty assistant message this PR set out to fix.
+        console.error("[kimi-coding] stream error:", err);
+        if (attempt === 0) {
+          const refreshed = await refreshKimiAuthToken(currentKey);
+          if (refreshed && refreshed !== currentKey) {
+            console.error("[kimi-coding] retrying stream after thrown error with refreshed token");
+            currentKey = refreshed;
+            shouldRetry = true;
+          }
+        }
+        if (!shouldRetry) {
+          filtered.push({
+            type: "error",
+            reason: "error",
+            error: {
+              content: [],
+              stopReason: "error",
+              usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 },
+            },
+          } as AssistantMessageEvent & { type: "error" });
+        }
       }
-    } catch (err) {
-      console.error("[kimi-coding] stream error:", err);
-      filtered.push({
-        type: "error",
-        reason: "error",
-        error: {
-          content: [],
-          stopReason: "error",
-          usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 },
-        },
-      } as AssistantMessageEvent & { type: "error" });
+
+      if (shouldRetry) {
+        attempt++;
+        continue;
+      }
+      break;
     }
   })();
 
@@ -800,7 +955,7 @@ export default function (pi: ExtensionAPI) {
   pi.registerProvider("kimi-coding", {
     baseUrl: getBaseUrl(),
     apiKey: "KIMI_API_KEY",
-    api: PROTOCOL,
+    api: KIMI_API_TYPE,
     streamSimple: streamSimpleKimi,
 
     headers: getCommonHeaders(),
@@ -815,42 +970,6 @@ export default function (pi: ExtensionAPI) {
         contextWindow: 262144,
         maxTokens: 32000,
       },
-      {
-        id: "kimi-k2.6-code-preview",
-        name: "Kimi K2.6 Code Preview",
-        reasoning: true,
-        input: ["text", "image", "video"],
-        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-        contextWindow: 262144,
-        maxTokens: 32000,
-      },
-      {
-        id: "kimi-code",
-        name: "Kimi Code (kimi-k2.5)",
-        reasoning: true,
-        input: ["text", "image"],
-        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-        contextWindow: 262144,
-        maxTokens: 32000,
-      },
-      {
-        id: "kimi-k2.5",
-        name: "Kimi K2.5",
-        reasoning: true,
-        input: ["text", "image"],
-        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-        contextWindow: 262144,
-        maxTokens: 32000,
-      },
-      {
-        id: "kimi-k2-thinking-turbo",
-        name: "Kimi K2 Thinking Turbo",
-        reasoning: true,
-        input: ["text"],
-        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-        contextWindow: 262144,
-        maxTokens: 32000,
-      },
     ],
 
     oauth: {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-provider-kimi-code",
-  "version": "0.2.2",
+  "version": "0.3.1",
   "description": "Pi extension for Kimi Code — reuse your Kimi Code Plan in pi-coding-agent (K2.6, K2.5)",
   "keywords": [
     "pi-package",
@@ -14,13 +14,11 @@
     "llm",
     "llm-provider",
     "kimi",
-    "kimi-code",
     "kimi-for-coding",
     "kimi-coder",
     "kimi-k2",
     "kimi-k2.6",
     "kimi-k2.5",
-    "kimi-k2-thinking",
     "moonshot",
     "moonshotai",
     "oauth",