/** * Z.AI Coding Plan Web Search client. * Uses the Coding Plan MCP endpoint, not the paid REST Web Search API. * Docs: https://docs.z.ai/devpack/mcp/search-mcp-server */ import type { ExtensionContext } from "@earendil-works/pi-coding-agent"; import type { SearchResult } from "./constants.js"; import { callMcpTool, McpError, type McpToolResult } from "./mcp.js"; import { type SettingsContext, type ZaiSearchFreshness, zaiContentSize, zaiFreshness, zaiLocation, zaiMcpUrl, } from "./settings.ts"; type ZaiRecencyFilter = | "oneDay" | "oneWeek" | "oneMonth" | "oneYear" | "noLimit"; interface ZaiSearchObjectResponse { title?: string; content?: string; link?: string; url?: string; media?: string; icon?: string; refer?: string; publish_date?: string; date?: string; } interface ZaiErrorResponse { code?: number | string; message?: string; error?: { code?: number | string; message?: string }; } export class ZaiApiError extends Error { readonly code?: string; readonly status?: number; readonly rawMessage: string; constructor( message: string, code?: string, status?: number, rawMessage = "unknown error", ) { super(message); this.name = "ZaiApiError"; this.code = code; this.status = status; this.rawMessage = rawMessage; } } const ZAI_PROVIDER_NAMES = ["zai", "z-ai", "z.ai", "bigmodel"] as const; const ZAI_API_KEY_ENV = ["ZAI_API_KEY", "Z_AI_API_KEY"] as const; const FRESHNESS_TO_RECENCY: Record = { day: "oneDay", week: "oneWeek", month: "oneMonth", year: "oneYear", }; export async function resolveZaiApiKey( ctx?: ExtensionContext, ): Promise { if (ctx) { for (const provider of ZAI_PROVIDER_NAMES) { const key = await ctx.modelRegistry.getApiKeyForProvider(provider); if (key) return key; } } for (const envName of ZAI_API_KEY_ENV) { const value = process.env[envName]?.trim(); if (value) return value; } return undefined; } export async function isZaiAvailable(ctx?: ExtensionContext): Promise { return !!(await resolveZaiApiKey(ctx)); } function normalizeMcpUrl(input: string): string { const value = input.trim().replace(/\/+$/, ""); try { const url = new URL(value); if (url.pathname.endsWith("/mcp")) return value; return `${url.origin}/api/mcp/web_search_prime/mcp`; } catch {} return value.endsWith("/mcp") ? value : `${value}/api/mcp/web_search_prime/mcp`; } function makeZaiApiError( status: number | undefined, code: unknown, message: unknown, ): ZaiApiError { const codeText = code === undefined || code === null ? undefined : String(code); const rawMessage = typeof message === "string" && message ? message : "unknown error"; const tag = codeText ? ` (${codeText})` : status !== undefined && status >= 400 ? ` (HTTP ${status})` : ""; return new ZaiApiError( `Z.AI API error${tag}: ${rawMessage}`, codeText, status, rawMessage, ); } async function parseZaiError(response: Response): Promise { const text = await response.text().catch(() => ""); if (!text) return makeZaiApiError( response.status, undefined, response.statusText || "request failed", ); try { const data = JSON.parse(text) as ZaiErrorResponse; return makeZaiApiError( response.status, data.error?.code ?? data.code, data.error?.message ?? data.message ?? response.statusText, ); } catch {} return makeZaiApiError( response.status, undefined, `${response.statusText}: ${text.slice(0, 500)}`, ); } function decodeSearchItems(text: string): ZaiSearchObjectResponse[] { let value: unknown = text.trim(); for (let i = 0; i < 3 && typeof value === "string"; i++) { try { value = JSON.parse(value); } catch { break; } } if (Array.isArray(value)) return value as ZaiSearchObjectResponse[]; if (value && typeof value === "object") { const record = value as { search_result?: ZaiSearchObjectResponse[]; results?: ZaiSearchObjectResponse[]; }; if (Array.isArray(record.search_result)) return record.search_result; if (Array.isArray(record.results)) return record.results; } return []; } function mcpResultToSearchResults(result: McpToolResult): SearchResult[] { const text = result.content?.find((item) => item.type === "text" && item.text)?.text ?? result.content?.find((item) => item.text)?.text; if (result.isError) { throw makeZaiApiError(200, undefined, text ?? "MCP tool returned an error"); } if (!text) return []; return decodeSearchItems(text) .map((r) => ({ title: r.title ?? r.link ?? r.url ?? "Untitled", url: r.link ?? r.url ?? "", snippet: r.content ?? r.media ?? "", date: r.publish_date ?? r.date, })) .filter((r) => r.url); } export async function zaiSearch( query: string, ctx: ExtensionContext | undefined, signal?: AbortSignal, options?: { count?: number; freshness?: ZaiSearchFreshness; timeoutMs?: number; }, ): Promise { const apiKey = await resolveZaiApiKey(ctx); if (!apiKey) { throw new Error( "Z.AI API key not configured. Set ZAI_API_KEY or configure a zai provider.", ); } // Without a session, project settings stay unread. const settings: SettingsContext = ctx ?? { cwd: process.cwd(), isProjectTrusted: () => false, }; const freshness = options?.freshness ?? zaiFreshness(settings); const toolArgs: Record = { search_query: query, content_size: zaiContentSize(settings), location: zaiLocation(settings), }; if (freshness) toolArgs.search_recency_filter = FRESHNESS_TO_RECENCY[freshness]; const timeoutSignal = AbortSignal.timeout(options?.timeoutMs ?? 30_000); const requestSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal; const result = await callMcpTool({ url: normalizeMcpUrl(zaiMcpUrl(settings)), protocolVersion: "2025-03-26", headers: { Authorization: `Bearer ${apiKey}` }, signal: requestSignal, name: "web_search_prime", arguments: toolArgs, }).catch(async (error: unknown) => { if (!(error instanceof McpError)) throw error; if (error.response) throw await parseZaiError(error.response); throw makeZaiApiError(error.status, error.code, error.message); }); const results = mcpResultToSearchResults(result); const count = Math.min( Math.max(Math.trunc(options?.count ?? results.length), 1), 50, ); return results.slice(0, count); }