pi-provider-kimi-code 0.2.1 → 0.3.0

package/README.md

@@ -2,22 +2,22 @@
 [![npm](https://img.shields.io/npm/v/pi-provider-kimi-code)](https://www.npmjs.com/package/pi-provider-kimi-code)
 [![license](https://img.shields.io/npm/l/pi-provider-kimi-code)](./LICENSE)
 
-**Reuse your [Kimi Code Membership](https://www.kimi.com/code/docs/en/) inside [pi-coding-agent](https://github.com/badlogic/pi-mono/tree/main/packages/coding-agent)** — 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.
+**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 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 Membership](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 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.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.
 
 ## Features
 
 - **Browser login with your Kimi account** — reuse your Kimi Code Membership without buying separate API credits. Credentials are stored locally and refreshed automatically. `KIMI_API_KEY` is also accepted for pay-per-token or CI use.
-- **262K-token context window** on all three registered models.
+- **262K-token context window** on all registered models.
 - **Automatic prompt caching** — binds Kimi's prompt cache to your pi session so repeated calls hit the cache cheaply (observed TTL ~5-10 minutes). Honors pi's `PI_CACHE_RETENTION=none` if you want to disable caching entirely.
 - **Automatic large-image and video upload** — images over 1 MB and all videos are uploaded to Kimi ahead of time, so you don't hit inline payload limits.
 - **Works with both Anthropic- and OpenAI-compatible modes** — use whichever one your pi setup expects.
@@ -47,7 +47,7 @@ If you prefer not to use npm, download the tarball from the [latest release](htt
 curl -L https://github.com/Leechael/pi-provider-kimi-code/releases/latest/download/pi-provider-kimi-code.tar.gz | tar -xz -C /tmp
 
 # Install from the extracted directory
-pi install /tmp/package
+pi install /tmp/pi-provider-kimi-code
 ```
 
 ## Authentication
@@ -72,16 +72,16 @@ KIMI_API_KEY=sk-... pi
 
 ## Models
 
-| ID                       | Name                             | Reasoning | Input       | Context | Max Output |
-| ------------------------ | -------------------------------- | --------- | ----------- | ------- | ---------- |
-| `kimi-code`              | Kimi Code (powered by 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
@@ -709,6 +710,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 +772,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 +780,105 @@ 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;
+    return {
+      ...options,
+      apiKey,
+      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);
+      const upstream =
+        model.api === "openai-completions"
+          ? streamSimpleOpenAICompletions(
+              model as Model<"openai-completions">,
+              context,
+              patchedOptions,
+            )
+          : streamSimpleAnthropic(model as Model<"anthropic-messages">, context, patchedOptions);
+
+      let pushedAny = false;
+      let shouldRetry = false;
+
+      try {
+        for await (const event of filterEmptyResponseStream(upstream)) {
+          // If the upstream terminates with an error before we've emitted
+          // anything downstream, speculatively try an OAuth refresh and
+          // retry once. This handles the common case of a server-side token
+          // invalidation before the local `expires` lapses, without having
+          // to pattern-match every possible error-message format.
+          //
+          // Non-auth errors (overflow, network, rate-limit, etc.) still
+          // trigger one wasted refresh, but the retried request fails the
+          // same way and we forward it — pi-coding-agent's own recovery
+          // paths (compaction, retry) take over from there.
+          if (!pushedAny && 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;
+            }
+            console.error(
+              "[kimi-coding] refresh did not yield a new token, forwarding original error",
+            );
+          }
+          filtered.push(event);
+          pushedAny = true;
+        }
+      } 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" });
       }
-    } 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;
     }
   })();
 
@@ -807,28 +900,10 @@ export default function (pi: ExtensionAPI) {
 
     models: [
       {
-        id: "kimi-code",
-        name: "Kimi Code (powered by 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",
+        id: "kimi-for-coding",
+        name: "Kimi for Coding",
         reasoning: true,
-        input: ["text"],
+        input: ["text", "image", "video"],
         cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
         contextWindow: 262144,
         maxTokens: 32000,

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "pi-provider-kimi-code",
-  "version": "0.2.1",
-  "description": "Pi extension for Kimi Code — reuse your kimi.com Membership in pi-coding-agent (K2.5, K2 Thinking Turbo)",
+  "version": "0.3.0",
+  "description": "Pi extension for Kimi Code — reuse your Kimi Code Plan in pi-coding-agent (K2.6, K2.5)",
   "keywords": [
     "pi-package",
     "pi-extension",
@@ -14,12 +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",