pi-provider-kimi-code 0.2.2 → 0.3.0

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
@@ -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;
     }
   })();
 
@@ -815,42 +908,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.0",
   "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",