pi-provider-kimi-code 0.5.1 → 0.6.1

package/src/tools/datasource.ts

@@ -0,0 +1,239 @@
+import { defineTool } from "@earendil-works/pi-coding-agent";
+import { Type, type Static } from "@earendil-works/pi-ai";
+
+import {
+  type BuildKimiToolOptions,
+  buildHeaders,
+  getKimiDatasourceUrl,
+  buildKimiToolDeps,
+  buildTimeoutSignal,
+  errorResult,
+  fetchWithAuthRetry,
+  firstText,
+  readErrorBody,
+  shouldCollapse,
+  textComponent,
+} from "./common.ts";
+
+const DATASOURCE_DEFINITIONS = [
+  {
+    name: "stock_finance_data",
+    description: "A/HK/US stocks, prices, financial reports, indicators, shareholders.",
+  },
+  {
+    name: "yahoo_finance",
+    description: "Global finance: quotes, analyst ratings, options, indexes, historical prices.",
+  },
+  {
+    name: "world_bank_open_data",
+    description: "World Bank macro data: GDP, inflation, population, country comparisons.",
+  },
+  {
+    name: "tianyancha",
+    description: "Chinese enterprise registry: shareholders, legal representative, risks, patents.",
+  },
+  { name: "arxiv", description: "arXiv preprints: paper search, abstracts, arXiv IDs." },
+  {
+    name: "scholar",
+    description: "Google Scholar: academic papers, authors, citations, high-citation surveys.",
+  },
+  {
+    name: "yuandian_law",
+    description: "Chinese laws, regulations, judicial cases; labor contracts and legal clauses.",
+  },
+] as const;
+
+const DATASOURCE_NAMES = DATASOURCE_DEFINITIONS.map((source) => source.name);
+
+export const kimiDatasourceSchema = Type.Object({
+  data_source_name: Type.Union(
+    DATASOURCE_NAMES.map((name) => Type.Literal(name)),
+    {
+      description:
+        "Source: stock_finance_data stocks/reports; yahoo_finance quotes/options; world_bank_open_data GDP/population; tianyancha Chinese company registry/legal rep; arxiv preprints; scholar citations; yuandian_law Chinese laws/cases/contracts.",
+    },
+  ),
+  api_name: Type.Optional(
+    Type.String({
+      description: "API name from the datasource docs. Omit to get docs first.",
+    }),
+  ),
+  params: Type.Optional(
+    Type.Record(Type.String(), Type.Unknown(), {
+      description: "Parameters matching the selected datasource API docs.",
+    }),
+  ),
+});
+
+export type KimiDatasourceParams = Static<typeof kimiDatasourceSchema>;
+
+interface DatasourceToolResponse {
+  is_success?: boolean;
+  error?: unknown;
+  result?: {
+    assistant?: unknown;
+    user?: unknown;
+  };
+}
+
+function extractTextFromResult(result: unknown): string {
+  if (!Array.isArray(result)) return "";
+  const texts: string[] = [];
+  for (const item of result) {
+    if (typeof item === "string") {
+      texts.push(item);
+    } else if (typeof item === "object" && item !== null) {
+      const record = item as Record<string, unknown>;
+      if (typeof record.text === "string") {
+        texts.push(record.text);
+      }
+    }
+  }
+  return texts.join("\n");
+}
+
+function extractResponseText(data: DatasourceToolResponse): string {
+  if (data.result?.assistant) {
+    const text = extractTextFromResult(data.result.assistant);
+    if (text) return text;
+  }
+  if (data.result?.user) {
+    const text = extractTextFromResult(data.result.user);
+    if (text) return text;
+  }
+  return JSON.stringify(data, null, 2);
+}
+
+function extractResponseFiles(data: DatasourceToolResponse): string[] {
+  const raw = data as Record<string, unknown>;
+  if (!Array.isArray(raw.files)) return [];
+  const names: string[] = [];
+  for (const file of raw.files) {
+    if (
+      typeof file === "object" &&
+      file !== null &&
+      typeof (file as Record<string, unknown>).name === "string"
+    ) {
+      names.push((file as Record<string, unknown>).name as string);
+    }
+  }
+  return names;
+}
+
+const REQUEST_ID_HEADERS = [
+  "x-request-id",
+  "x-trace-id",
+  "x-msh-request-id",
+  "x-msh-trace-id",
+  "request-id",
+];
+
+function extractRequestId(headers: Headers): string | undefined {
+  for (const key of REQUEST_ID_HEADERS) {
+    const value = headers.get(key);
+    if (typeof value === "string" && value.trim()) return value.trim();
+  }
+  return undefined;
+}
+
+function buildTraceLine(toolCallId: string, requestId: string | undefined): string {
+  const parts: string[] = [];
+  if (requestId) parts.push(`request-id: ${requestId}`);
+  parts.push(`tool-call-id: ${toolCallId}`);
+  return `\n\n[kimi-datasource] ${parts.join(" · ")}`;
+}
+
+export function buildKimiDatasourceTool(options: BuildKimiToolOptions = {}) {
+  const deps = buildKimiToolDeps(options);
+  const defaultCollapsed = options.defaultCollapsed ?? true;
+
+  return defineTool({
+    name: "kimi_datasource",
+    label: "Kimi Datasource",
+    description:
+      "Call Kimi datasource APIs. Covers stocks/finance, macro data, Chinese company registry, papers, and Chinese law/cases/contracts. First omit api_name to get docs, then call api_name with documented params.",
+    promptSnippet: "Query Kimi datasource APIs; get docs first, then call documented api_name",
+    parameters: kimiDatasourceSchema,
+
+    async execute(toolCallId, params, signal) {
+      const accessToken = deps.getAccessToken();
+      if (!accessToken) {
+        return errorResult<string>(
+          "Missing Kimi Code OAuth credentials. Run /login kimi-coding first.",
+        );
+      }
+
+      const timeout = buildTimeoutSignal(signal);
+      try {
+        const apiName = params.api_name?.trim();
+        const body = apiName
+          ? {
+              method: "call_data_source_tool",
+              params: {
+                data_source_name: params.data_source_name,
+                api_name: apiName,
+                params: params.params ?? {},
+              },
+            }
+          : {
+              method: "get_data_source_desc",
+              params: { name: params.data_source_name },
+            };
+
+        const response = await fetchWithAuthRetry(deps, accessToken, (token) =>
+          deps.fetch(getKimiDatasourceUrl(), {
+            method: "POST",
+            headers: buildHeaders(token, toolCallId),
+            body: JSON.stringify(body),
+            signal: timeout.signal,
+          }),
+        );
+
+        const requestId = extractRequestId(response.headers);
+        const trace = buildTraceLine(toolCallId, requestId);
+
+        if (!response.ok) {
+          const bodyText = await readErrorBody(response);
+          return errorResult<string>(
+            `kimi_datasource failed: ${response.status}${bodyText ? ` ${bodyText}` : ""}${trace}`,
+          );
+        }
+
+        const data = (await response.json()) as DatasourceToolResponse;
+        if (data.is_success === false) {
+          const errorText =
+            typeof data.error === "string"
+              ? data.error
+              : JSON.stringify(data.error ?? data, null, 2);
+          return errorResult<string>(`kimi_datasource error: ${errorText}${trace}`);
+        }
+
+        let text = extractResponseText(data);
+        const files = extractResponseFiles(data);
+        if (files.length > 0) {
+          text += `\n\nNote: the backend returned ${files.length} file(s) that cannot be written in this context: ${files.join(", ")}`;
+        }
+        text += trace;
+
+        return {
+          content: [{ type: "text", text }],
+          details: text,
+        };
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        return errorResult<string>(`kimi_datasource failed: ${message}`);
+      } finally {
+        timeout.cleanup();
+      }
+    },
+
+    renderResult(result, renderOptions) {
+      const text = firstText(result);
+      if (!shouldCollapse(defaultCollapsed, renderOptions.expanded)) {
+        return textComponent(text);
+      }
+      const preview = text.length > 200 ? `${text.slice(0, 200)}…` : text;
+      return textComponent(`kimi_datasource: ${text.length} chars\n${preview}`);
+    },
+  });
+}

package/src/tools/moonshot.ts

@@ -1,12 +1,19 @@
-import { AuthStorage, defineTool, type AgentToolResult } from "@earendil-works/pi-coding-agent";
+import { 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";
-import { refreshKimiAuthToken } from "../oauth.ts";
-
-const MOONSHOT_BASE_V1 = "https://api.kimi.com/coding/v1";
-const MOONSHOT_TIMEOUT_MS = 180_000;
+import {
+  type BuildKimiToolOptions,
+  buildHeaders,
+  getKimiBaseV1,
+  buildKimiToolDeps,
+  buildTimeoutSignal,
+  errorResult,
+  fetchWithAuthRetry,
+  firstText,
+  readErrorBody,
+  shouldCollapse,
+  textComponent,
+} from "./common.ts";
 
 export const moonshotSearchSchema = Type.Object({
   query: Type.String({ description: "The query text to search for." }),
@@ -45,76 +52,13 @@ export interface MoonshotFetchResult {
   content: string;
 }
 
-interface MoonshotToolDeps {
-  fetch: typeof fetch;
-  getAccessToken: () => string | null;
-  refreshAccessToken: (currentToken: string) => Promise<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,
-    refreshAccessToken: options.deps?.refreshAccessToken ?? refreshKimiAuthToken,
-  };
-}
-
 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,
-  };
+interface MoonshotSearchResponse {
+  search_results?: unknown;
 }
 
 function stringField(record: Record<string, unknown>, key: string): string {
@@ -144,22 +88,6 @@ function searchResultsText(results: MoonshotSearchResult[]): string {
   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,
@@ -195,25 +123,8 @@ function renderFetchResult(
   );
 }
 
-async function readErrorBody(response: Response): Promise<string> {
-  return response.text().catch(() => "");
-}
-
-async function fetchWithAuthRetry(
-  deps: MoonshotToolDeps,
-  accessToken: string,
-  request: (token: string) => Promise<Response>,
-): Promise<Response> {
-  const response = await request(accessToken);
-  if (response.status !== 401) return response;
-
-  const refreshed = await deps.refreshAccessToken(accessToken);
-  if (!refreshed || refreshed === accessToken) return response;
-  return request(refreshed);
-}
-
-export function buildMoonshotSearchTool(options: BuildMoonshotToolOptions = {}) {
-  const deps = buildDeps(options);
+export function buildMoonshotSearchTool(options: BuildKimiToolOptions = {}) {
+  const deps = buildKimiToolDeps(options);
   const defaultCollapsed = options.defaultCollapsed ?? true;
 
   return defineTool({
@@ -236,7 +147,7 @@ export function buildMoonshotSearchTool(options: BuildMoonshotToolOptions = {})
       const timeout = buildTimeoutSignal(signal);
       try {
         const response = await fetchWithAuthRetry(deps, accessToken, (token) =>
-          deps.fetch(`${MOONSHOT_BASE_V1}/search`, {
+          deps.fetch(`${getKimiBaseV1()}/search`, {
             method: "POST",
             headers: buildHeaders(token, toolCallId),
             body: JSON.stringify({
@@ -276,8 +187,8 @@ export function buildMoonshotSearchTool(options: BuildMoonshotToolOptions = {})
   });
 }
 
-export function buildMoonshotFetchTool(options: BuildMoonshotToolOptions = {}) {
-  const deps = buildDeps(options);
+export function buildMoonshotFetchTool(options: BuildKimiToolOptions = {}) {
+  const deps = buildKimiToolDeps(options);
   const defaultCollapsed = options.defaultCollapsed ?? true;
 
   return defineTool({
@@ -298,7 +209,7 @@ export function buildMoonshotFetchTool(options: BuildMoonshotToolOptions = {}) {
       const timeout = buildTimeoutSignal(signal);
       try {
         const response = await fetchWithAuthRetry(deps, accessToken, (token) =>
-          deps.fetch(`${MOONSHOT_BASE_V1}/fetch`, {
+          deps.fetch(`${getKimiBaseV1()}/fetch`, {
             method: "POST",
             headers: {
               ...buildHeaders(token, toolCallId),