pi-provider-kimi-code 0.3.1 → 0.4.0

package/src/stream.ts

@@ -0,0 +1,299 @@
+// Stream wrapper + empty-response filter. Reads every side-effect source
+// (process.env, options, apiKey) at the top and hands a plain
+// KimiPayloadContext to applyKimiPayloadMutations. The actual logic lives in
+// the pure units in payload.ts; this file just wires SDK streaming + filter +
+// error fallback together.
+
+import type {
+  Api,
+  AssistantMessageEvent,
+  AssistantMessageEventStream,
+  Context,
+  Model,
+  SimpleStreamOptions,
+} from "@earendil-works/pi-ai";
+import {
+  createAssistantMessageEventStream,
+  streamSimpleAnthropic,
+  streamSimpleOpenAICompletions,
+} from "@earendil-works/pi-ai";
+
+import { IS_OPENAI_PROTOCOL, PROTOCOL } from "./constants.ts";
+import { isKimiAuthErrorMessage, refreshKimiAuthToken } from "./oauth.ts";
+import {
+  type Uploader,
+  applyKimiPayloadMutations,
+  isRecord,
+  readEnvOverrides,
+  resolveCacheRetention,
+  uploadKimiFile,
+} from "./payload.ts";
+
+// =============================================================================
+// Event stream filter: suppress Kimi "(Empty response: ...)" text blocks
+// =============================================================================
+// The Kimi API wraps thinking-only responses (no text content) into a text
+// block like: (Empty response: {'content': [{'type': 'thinking', ...}]}).
+// This leaks internal state to the user. We buffer text_start/text_delta
+// events and drop the whole block if text_end starts with the marker.
+// Pure async generator — no closure dependencies, testable with synthetic
+// event arrays.
+
+const EMPTY_RESPONSE_PREFIX = "(Empty response:";
+
+export async function* filterEmptyResponseStream(
+  upstream: AsyncIterable<AssistantMessageEvent>,
+): AsyncIterable<AssistantMessageEvent> {
+  const suppressedIndices = new Set<number>();
+  let textBuffer: AssistantMessageEvent[] = [];
+  let bufferingIndex: number | null = null;
+
+  for await (const event of upstream) {
+    // Start buffering when a new text block begins.
+    if (event.type === "text_start") {
+      bufferingIndex = event.contentIndex;
+      textBuffer = [event];
+      continue;
+    }
+
+    // Accumulate text deltas + detect the empty-response marker on text_end.
+    if (
+      bufferingIndex !== null &&
+      "contentIndex" in event &&
+      event.contentIndex === bufferingIndex
+    ) {
+      if (event.type === "text_delta") {
+        textBuffer.push(event);
+        continue;
+      }
+      if (event.type === "text_end") {
+        if (event.content.startsWith(EMPTY_RESPONSE_PREFIX)) {
+          // Suppress entire text block. Do NOT splice the message content
+          // array: it is a shared reference into session state, and mutating
+          // it would shift subsequent contentIndex values, corrupting the
+          // stream.
+          suppressedIndices.add(bufferingIndex);
+        } else {
+          // Legitimate text block — flush buffered events + end event.
+          for (const buffered of textBuffer) yield buffered;
+          yield event;
+        }
+        textBuffer = [];
+        bufferingIndex = null;
+        continue;
+      }
+    }
+
+    // Skip any event targeting an already-suppressed content index.
+    if ("contentIndex" in event && suppressedIndices.has(event.contentIndex)) {
+      continue;
+    }
+
+    // Clean suppressed blocks out of the final message.
+    if (event.type === "done" && suppressedIndices.size > 0) {
+      event.message.content = event.message.content.filter(
+        (block) =>
+          !(
+            block.type === "text" &&
+            typeof block.text === "string" &&
+            block.text.startsWith(EMPTY_RESPONSE_PREFIX)
+          ),
+      );
+    }
+
+    yield event;
+  }
+}
+
+// =============================================================================
+// Stream wrapper: orchestrates payload mutation + event filter
+// =============================================================================
+
+export function streamSimpleKimi(
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+): AssistantMessageEventStream {
+  const filtered = createAssistantMessageEventStream();
+  const initialKey = options?.apiKey || process.env.KIMI_API_KEY || "";
+
+  const cacheKeyOverride = (
+    options as (SimpleStreamOptions & { prompt_cache_key?: unknown }) | undefined
+  )?.prompt_cache_key;
+  const cacheKey = (typeof cacheKeyOverride === "string" && cacheKeyOverride) || options?.sessionId;
+  const cacheRetention = resolveCacheRetention(options?.cacheRetention);
+  const envOverrides = readEnvOverrides();
+  const thinkingKeep = process.env.KIMI_MODEL_THINKING_KEEP;
+  const originalOnPayload = options?.onPayload;
+  // The pi-side model id ("kimi-for-coding") is what users select via /model
+  // and what gets persisted into sessions. The wire model id discovered at
+  // OAuth login (e.g. a versioned alias the server exposes) gets carried on
+  // the model object via modifyModels and rewritten into the request payload
+  // here so /v1/chat/completions and /v1/messages see the real wire id.
+  const wireModelId = (model as Model<Api> & { wireModelId?: unknown }).wireModelId;
+
+  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: PROTOCOL,
+            upload,
+            cacheKey,
+            cacheRetention,
+            reasoning: options?.reasoning,
+            thinkingKeep,
+            envOverrides,
+          });
+          if (
+            typeof wireModelId === "string" &&
+            wireModelId &&
+            nextPayload.model === "kimi-for-coding"
+          ) {
+            nextPayload.model = wireModelId;
+          }
+        }
+
+        if (originalOnPayload) {
+          const res = await originalOnPayload(nextPayload, modelData);
+          if (res !== undefined) nextPayload = res;
+        }
+
+        return nextPayload;
+      },
+    };
+  };
+
+  void (async () => {
+    let attempt = 0;
+    let currentKey = initialKey;
+
+    while (true) {
+      const patchedOptions = buildPatchedOptions(currentKey);
+      // Route by the module-level protocol flag, 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 = IS_OPENAI_PROTOCOL
+        ? 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 auth error. We retry once
+          // so short-lived Kimi tokens invalidated before the local expires
+          // timestamp lapses don't surface as raw 401s to the user.
+          if (
+            attempt === 0 &&
+            event.type === "error" &&
+            isKimiAuthErrorMessage(event.error?.errorMessage)
+          ) {
+            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.
+        if (!shouldRetry) {
+          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 && isKimiAuthErrorMessage(err instanceof Error ? err.message : err)) {
+          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: {
+              role: "assistant",
+              api: model.api,
+              provider: model.provider,
+              model: model.id,
+              content: [],
+              stopReason: "error",
+              errorMessage: err instanceof Error ? err.message : String(err),
+              usage: {
+                input: 0,
+                output: 0,
+                cacheRead: 0,
+                cacheWrite: 0,
+                totalTokens: 0,
+                cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+              },
+              timestamp: Date.now(),
+            },
+          });
+        }
+      }
+
+      if (shouldRetry) {
+        attempt++;
+        continue;
+      }
+      break;
+    }
+  })();
+
+  return filtered;
+}

package/src/tools/moonshot.ts

@@ -0,0 +1,316 @@
+import { AuthStorage, defineTool, type AgentToolResult } from "@earendil-works/pi-coding-agent";
+import { Type, type Static } from "@earendil-works/pi-ai";
+
+import { PROVIDER_ID } from "../constants.ts";
+import { getCommonHeaders } from "../device.ts";
+
+const MOONSHOT_BASE_V1 = "https://api.kimi.com/coding/v1";
+const MOONSHOT_TIMEOUT_MS = 180_000;
+
+export const moonshotSearchSchema = Type.Object({
+  query: Type.String({ description: "The query text to search for." }),
+  limit: Type.Optional(
+    Type.Number({
+      description: "The number of results to return.",
+      default: 5,
+      minimum: 1,
+      maximum: 20,
+    }),
+  ),
+  include_content: Type.Optional(
+    Type.Boolean({
+      description: "Whether to include fetched page content in the search results.",
+      default: false,
+    }),
+  ),
+});
+
+export const moonshotFetchSchema = Type.Object({
+  url: Type.String({ description: "The URL to fetch content from." }),
+});
+
+export type MoonshotSearchParams = Static<typeof moonshotSearchSchema>;
+export type MoonshotFetchParams = Static<typeof moonshotFetchSchema>;
+
+export interface MoonshotSearchResult {
+  url: string;
+  title: string;
+  snippet: string;
+  content?: string;
+}
+
+export interface MoonshotFetchResult {
+  url: string;
+  content: string;
+}
+
+interface MoonshotToolDeps {
+  fetch: typeof fetch;
+  getAccessToken: () => string | null;
+}
+
+export interface BuildMoonshotToolOptions {
+  deps?: Partial<MoonshotToolDeps>;
+  defaultCollapsed?: boolean;
+}
+
+interface MoonshotSearchResponse {
+  search_results?: unknown;
+}
+
+function defaultGetAccessToken(): string | null {
+  const credential = AuthStorage.create().get(PROVIDER_ID);
+  if (!credential || credential.type !== "oauth") return null;
+  return credential.access;
+}
+
+function buildDeps(options: BuildMoonshotToolOptions = {}): MoonshotToolDeps {
+  return {
+    fetch: options.deps?.fetch ?? fetch,
+    getAccessToken: options.deps?.getAccessToken ?? defaultGetAccessToken,
+  };
+}
+
+function clampLimit(value: number | undefined): number {
+  if (typeof value !== "number" || !Number.isFinite(value)) return 5;
+  return Math.max(1, Math.min(20, Math.trunc(value)));
+}
+
+function errorResult<T>(message: string): AgentToolResult<T> {
+  return {
+    content: [{ type: "text", text: message }],
+    details: undefined as T,
+  };
+}
+
+function buildTimeoutSignal(signal: AbortSignal | undefined): {
+  signal: AbortSignal;
+  cleanup: () => void;
+} {
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), MOONSHOT_TIMEOUT_MS);
+  const abort = () => controller.abort();
+  if (signal?.aborted) {
+    controller.abort();
+  } else {
+    signal?.addEventListener("abort", abort, { once: true });
+  }
+
+  return {
+    signal: controller.signal,
+    cleanup: () => {
+      clearTimeout(timeout);
+      signal?.removeEventListener("abort", abort);
+    },
+  };
+}
+
+function buildHeaders(accessToken: string, toolCallId: string): Record<string, string> {
+  return {
+    ...getCommonHeaders(),
+    Authorization: `Bearer ${accessToken}`,
+    "Content-Type": "application/json",
+    "X-Msh-Tool-Call-Id": toolCallId,
+  };
+}
+
+function stringField(record: Record<string, unknown>, key: string): string {
+  const value = record[key];
+  return typeof value === "string" ? value : "";
+}
+
+function mapSearchResults(value: unknown, includeContent: boolean): MoonshotSearchResult[] {
+  if (!Array.isArray(value)) return [];
+  return value
+    .filter((item): item is Record<string, unknown> => typeof item === "object" && item !== null)
+    .map((item) => {
+      const result: MoonshotSearchResult = {
+        url: stringField(item, "url"),
+        title: stringField(item, "title"),
+        snippet: stringField(item, "snippet"),
+      };
+      if (includeContent) {
+        result.content = stringField(item, "content");
+      }
+      return result;
+    });
+}
+
+function searchResultsText(results: MoonshotSearchResult[]): string {
+  if (results.length === 0) return "[]";
+  return JSON.stringify(results, null, 2);
+}
+
+function textComponent(text: string) {
+  return {
+    render: () => text.split("\n"),
+    invalidate: () => {},
+  };
+}
+
+function firstText(result: AgentToolResult<unknown>): string {
+  const first = result.content[0];
+  return first?.type === "text" ? first.text : "";
+}
+
+function shouldCollapse(defaultCollapsed: boolean, expanded: boolean | undefined): boolean {
+  return defaultCollapsed && expanded !== true;
+}
+
+function renderSearchResult(
+  result: AgentToolResult<MoonshotSearchResult[]>,
+  expanded: boolean | undefined,
+  defaultCollapsed: boolean,
+) {
+  if (!Array.isArray(result.details)) {
+    return textComponent(firstText(result));
+  }
+  if (!shouldCollapse(defaultCollapsed, expanded)) {
+    return textComponent(searchResultsText(result.details));
+  }
+
+  const count = result.details.length;
+  const first = result.details[0];
+  const suffix = first ? `; first: ${first.title || first.url}` : "";
+  return textComponent(`moonshot_search returned ${count} result(s)${suffix}`);
+}
+
+function renderFetchResult(
+  result: AgentToolResult<MoonshotFetchResult>,
+  expanded: boolean | undefined,
+  defaultCollapsed: boolean,
+) {
+  if (!result.details) {
+    return textComponent(firstText(result));
+  }
+  if (!shouldCollapse(defaultCollapsed, expanded)) {
+    return textComponent(result.details.content);
+  }
+
+  return textComponent(
+    `moonshot_fetch fetched ${result.details.url} (${result.details.content.length} chars)`,
+  );
+}
+
+async function readErrorBody(response: Response): Promise<string> {
+  return response.text().catch(() => "");
+}
+
+export function buildMoonshotSearchTool(options: BuildMoonshotToolOptions = {}) {
+  const deps = buildDeps(options);
+  const defaultCollapsed = options.defaultCollapsed ?? true;
+
+  return defineTool({
+    name: "moonshot_search",
+    label: "Moonshot Search",
+    description: "Search the web through Kimi Coding's server-side Moonshot search service.",
+    promptSnippet: "Search the web with Kimi Coding's Moonshot search service",
+    parameters: moonshotSearchSchema,
+
+    async execute(toolCallId, params, signal) {
+      const accessToken = deps.getAccessToken();
+      if (!accessToken) {
+        return errorResult<MoonshotSearchResult[]>(
+          "Missing Kimi Code OAuth credentials. Run /login kimi-coding first.",
+        );
+      }
+
+      const limit = clampLimit(params.limit);
+      const includeContent = params.include_content === true;
+      const timeout = buildTimeoutSignal(signal);
+      try {
+        const response = await deps.fetch(`${MOONSHOT_BASE_V1}/search`, {
+          method: "POST",
+          headers: buildHeaders(accessToken, toolCallId),
+          body: JSON.stringify({
+            text_query: params.query,
+            limit,
+            enable_page_crawling: includeContent,
+            timeout_seconds: 30,
+          }),
+          signal: timeout.signal,
+        });
+
+        if (!response.ok) {
+          const body = await readErrorBody(response);
+          return errorResult<MoonshotSearchResult[]>(
+            `Moonshot search failed: ${response.status}${body ? ` ${body}` : ""}`,
+          );
+        }
+
+        const data = (await response.json()) as MoonshotSearchResponse;
+        const results = mapSearchResults(data.search_results, includeContent);
+        return {
+          content: [{ type: "text", text: searchResultsText(results) }],
+          details: results,
+        };
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        return errorResult<MoonshotSearchResult[]>(`Moonshot search failed: ${message}`);
+      } finally {
+        timeout.cleanup();
+      }
+    },
+
+    renderResult(result, renderOptions) {
+      return renderSearchResult(result, renderOptions.expanded, defaultCollapsed);
+    },
+  });
+}
+
+export function buildMoonshotFetchTool(options: BuildMoonshotToolOptions = {}) {
+  const deps = buildDeps(options);
+  const defaultCollapsed = options.defaultCollapsed ?? true;
+
+  return defineTool({
+    name: "moonshot_fetch",
+    label: "Moonshot Fetch",
+    description: "Fetch a web page through Kimi Coding's server-side Moonshot fetch service.",
+    promptSnippet: "Fetch a web page with Kimi Coding's Moonshot fetch service",
+    parameters: moonshotFetchSchema,
+
+    async execute(toolCallId, params, signal) {
+      const accessToken = deps.getAccessToken();
+      if (!accessToken) {
+        return errorResult<MoonshotFetchResult>(
+          "Missing Kimi Code OAuth credentials. Run /login kimi-coding first.",
+        );
+      }
+
+      const timeout = buildTimeoutSignal(signal);
+      try {
+        const response = await deps.fetch(`${MOONSHOT_BASE_V1}/fetch`, {
+          method: "POST",
+          headers: {
+            ...buildHeaders(accessToken, toolCallId),
+            Accept: "text/markdown",
+          },
+          body: JSON.stringify({ url: params.url }),
+          signal: timeout.signal,
+        });
+
+        if (!response.ok) {
+          const body = await readErrorBody(response);
+          return errorResult<MoonshotFetchResult>(
+            `Moonshot fetch failed: ${response.status}${body ? ` ${body}` : ""}`,
+          );
+        }
+
+        const content = await response.text();
+        return {
+          content: [{ type: "text", text: content }],
+          details: { url: params.url, content },
+        };
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        return errorResult<MoonshotFetchResult>(`Moonshot fetch failed: ${message}`);
+      } finally {
+        timeout.cleanup();
+      }
+    },
+
+    renderResult(result, renderOptions) {
+      return renderFetchResult(result, renderOptions.expanded, defaultCollapsed);
+    },
+  });
+}