repositories / pi-ext
pi-ext
bugabingas pi extensions
owned by admin
extensions/web/__tests__/zai.test.ts
Rawimport { afterEach, describe, expect, test } from "vitest";
import { ZaiApiError, zaiSearch } from "../zai.js";
const originalFetch = globalThis.fetch;
const originalZaiKey = process.env.ZAI_API_KEY;
const originalZaiBase = process.env.ZAI_BASE_URL;
const originalZaiSearchBase = process.env.ZAI_SEARCH_BASE_URL;
const originalZaiSearchMcp = process.env.ZAI_SEARCH_MCP_URL;
const originalZaiSearchLocation = process.env.ZAI_SEARCH_LOCATION;
const originalZaiSearchContentSize = process.env.ZAI_SEARCH_CONTENT_SIZE;
afterEach(() => {
globalThis.fetch = originalFetch;
if (originalZaiKey === undefined) delete process.env.ZAI_API_KEY;
else process.env.ZAI_API_KEY = originalZaiKey;
if (originalZaiBase === undefined) delete process.env.ZAI_BASE_URL;
else process.env.ZAI_BASE_URL = originalZaiBase;
if (originalZaiSearchBase === undefined)
delete process.env.ZAI_SEARCH_BASE_URL;
else process.env.ZAI_SEARCH_BASE_URL = originalZaiSearchBase;
if (originalZaiSearchMcp === undefined) delete process.env.ZAI_SEARCH_MCP_URL;
else process.env.ZAI_SEARCH_MCP_URL = originalZaiSearchMcp;
if (originalZaiSearchLocation === undefined)
delete process.env.ZAI_SEARCH_LOCATION;
else process.env.ZAI_SEARCH_LOCATION = originalZaiSearchLocation;
if (originalZaiSearchContentSize === undefined)
delete process.env.ZAI_SEARCH_CONTENT_SIZE;
else process.env.ZAI_SEARCH_CONTENT_SIZE = originalZaiSearchContentSize;
delete process.env.ZAI_SEARCH_FRESHNESS;
});
function sse(data: unknown): string {
return `id:1\nevent:message\ndata:${JSON.stringify(data)}\n\n`;
}
describe("zaiSearch", () => {
test("calls Coding Plan MCP search and maps results", async () => {
process.env.ZAI_API_KEY = "test-key";
const requestUrls: string[] = [];
let callBody: Record<string, unknown> = {};
globalThis.fetch = (async (input, init) => {
requestUrls.push(String(input));
const body = JSON.parse(String(init?.body));
expect(init?.headers).toMatchObject({ Authorization: "Bearer test-key" });
if (body.method === "initialize") {
return new Response(
sse({
jsonrpc: "2.0",
id: 1,
result: { protocolVersion: "2024-11-05" },
}),
{
status: 200,
headers: {
"Content-Type": "text/event-stream",
"mcp-session-id": "session-1",
},
},
);
}
if (body.method === "notifications/initialized") {
expect(init?.headers).toMatchObject({ "mcp-session-id": "session-1" });
return new Response("", { status: 200 });
}
callBody = body;
expect(init?.headers).toMatchObject({ "mcp-session-id": "session-1" });
const toolText = JSON.stringify(
JSON.stringify([
{
title: "Example",
content: "Summary",
link: "https://example.com",
publish_date: "2026-05-09",
},
]),
);
return new Response(
sse({
jsonrpc: "2.0",
id: 2,
result: { content: [{ type: "text", text: toolText }] },
}),
{ status: 200, headers: { "Content-Type": "text/event-stream" } },
);
}) as typeof fetch;
const results = await zaiSearch("example", undefined, undefined, {
count: 3,
freshness: "week",
});
expect(requestUrls).toEqual([
"https://api.z.ai/api/mcp/web_search_prime/mcp",
"https://api.z.ai/api/mcp/web_search_prime/mcp",
"https://api.z.ai/api/mcp/web_search_prime/mcp",
]);
expect(callBody).toMatchObject({
jsonrpc: "2.0",
method: "tools/call",
params: {
name: "web_search_prime",
arguments: {
search_query: "example",
content_size: "medium",
location: "us",
search_recency_filter: "oneWeek",
},
},
});
expect(results).toEqual([
{
title: "Example",
url: "https://example.com",
snippet: "Summary",
date: "2026-05-09",
},
]);
});
test("uses ZAI_SEARCH_FRESHNESS default", async () => {
process.env.ZAI_API_KEY = "test-key";
process.env.ZAI_SEARCH_FRESHNESS = "month";
let callArgs: Record<string, unknown> = {};
globalThis.fetch = (async (_input, init) => {
const body = JSON.parse(String(init?.body));
if (body.method === "initialize") {
return new Response(
sse({
jsonrpc: "2.0",
id: 1,
result: { protocolVersion: "2025-03-26" },
}),
{
status: 200,
headers: { "mcp-session-id": "s" },
},
);
}
if (body.method === "notifications/initialized")
return new Response("", { status: 200 });
callArgs = body.params.arguments;
return new Response(
sse({
jsonrpc: "2.0",
id: 2,
result: { content: [{ type: "text", text: "[]" }] },
}),
{ status: 200 },
);
}) as typeof fetch;
await zaiSearch("example", undefined, undefined, { count: 3 });
expect(callArgs).toMatchObject({ search_recency_filter: "oneMonth" });
});
test("normalizes base URLs to the MCP endpoint", async () => {
process.env.ZAI_API_KEY = "test-key";
process.env.ZAI_SEARCH_MCP_URL =
"https://api.z.ai/api/mcp/web_search_prime/mcp";
let requestUrl = "";
globalThis.fetch = (async (input, init) => {
requestUrl = String(input);
const body = JSON.parse(String(init?.body));
if (body.method === "initialize")
return new Response(
sse({
jsonrpc: "2.0",
id: 1,
result: { protocolVersion: "2025-03-26" },
}),
{
status: 200,
headers: { "mcp-session-id": "s" },
},
);
if (body.method === "notifications/initialized")
return new Response("", { status: 200 });
return new Response(
sse({
jsonrpc: "2.0",
id: 2,
result: { content: [{ type: "text", text: "[]" }] },
}),
{ status: 200 },
);
}) as typeof fetch;
await zaiSearch("example", undefined);
expect(requestUrl).toBe("https://api.z.ai/api/mcp/web_search_prime/mcp");
});
test("ZAI_SEARCH_BASE_URL overrides with MCP endpoint normalization", async () => {
process.env.ZAI_API_KEY = "test-key";
process.env.ZAI_SEARCH_BASE_URL = "https://api.z.ai/api/paas/v4";
let requestUrl = "";
globalThis.fetch = (async (input, init) => {
requestUrl = String(input);
const body = JSON.parse(String(init?.body));
if (body.method === "initialize")
return new Response(
sse({
jsonrpc: "2.0",
id: 1,
result: { protocolVersion: "2025-03-26" },
}),
{
status: 200,
headers: { "mcp-session-id": "s" },
},
);
if (body.method === "notifications/initialized")
return new Response("", { status: 200 });
return new Response(
sse({
jsonrpc: "2.0",
id: 2,
result: { content: [{ type: "text", text: "[]" }] },
}),
{ status: 200 },
);
}) as typeof fetch;
await zaiSearch("example", undefined);
expect(requestUrl).toBe("https://api.z.ai/api/mcp/web_search_prime/mcp");
});
test("preserves JSON API error code, status, and raw message", async () => {
process.env.ZAI_API_KEY = "test-key";
globalThis.fetch = (async () =>
new Response(
JSON.stringify({ error: { code: "1002", message: "bad key" } }),
{ status: 401, statusText: "Unauthorized" },
)) as typeof fetch;
try {
await zaiSearch("example", undefined);
throw new Error("expected failure");
} catch (err) {
expect(err).toBeInstanceOf(ZaiApiError);
expect("kind" in (err as object)).toBe(false);
expect((err as ZaiApiError).code).toBe("1002");
expect((err as ZaiApiError).status).toBe(401);
expect((err as ZaiApiError).rawMessage).toBe("bad key");
expect((err as Error).message).toBe("Z.AI API error (1002): bad key");
}
});
test("preserves HTTP plaintext errors", async () => {
process.env.ZAI_API_KEY = "test-key";
globalThis.fetch = (async () =>
new Response("Not Found", {
status: 404,
statusText: "Not Found",
})) as typeof fetch;
try {
await zaiSearch("example", undefined);
throw new Error("expected failure");
} catch (err) {
expect(err).toBeInstanceOf(ZaiApiError);
expect("kind" in (err as object)).toBe(false);
expect((err as ZaiApiError).code).toBeUndefined();
expect((err as ZaiApiError).status).toBe(404);
expect((err as ZaiApiError).rawMessage).toBe("Not Found: Not Found");
expect((err as Error).message).toBe(
"Z.AI API error (HTTP 404): Not Found: Not Found",
);
}
});
test("preserves MCP JSON-RPC errors", async () => {
process.env.ZAI_API_KEY = "test-key";
globalThis.fetch = (async (_input, init) => {
const body = JSON.parse(String(init?.body));
if (body.method === "initialize")
return new Response(
sse({
jsonrpc: "2.0",
id: 1,
result: { protocolVersion: "2025-03-26" },
}),
{
status: 200,
headers: { "mcp-session-id": "s" },
},
);
if (body.method === "notifications/initialized")
return new Response("", { status: 200 });
return new Response(
sse({
jsonrpc: "2.0",
id: 2,
error: { code: -32602, message: "bad params" },
}),
{ status: 200 },
);
}) as typeof fetch;
try {
await zaiSearch("example", undefined);
throw new Error("expected failure");
} catch (err) {
expect(err).toBeInstanceOf(ZaiApiError);
expect("kind" in (err as object)).toBe(false);
expect((err as ZaiApiError).code).toBe("-32602");
expect((err as ZaiApiError).status).toBe(200);
expect((err as ZaiApiError).rawMessage).toBe("bad params");
expect((err as Error).message).toBe(
"Z.AI API error (-32602): bad params",
);
}
});
test("preserves MCP tool errors without fake HTTP 200 label", async () => {
process.env.ZAI_API_KEY = "test-key";
globalThis.fetch = (async (_input, init) => {
const body = JSON.parse(String(init?.body));
if (body.method === "initialize")
return new Response(
sse({
jsonrpc: "2.0",
id: 1,
result: { protocolVersion: "2025-03-26" },
}),
{
status: 200,
headers: { "mcp-session-id": "s" },
},
);
if (body.method === "notifications/initialized")
return new Response("", { status: 200 });
return new Response(
sse({
jsonrpc: "2.0",
id: 2,
result: {
content: [{ type: "text", text: "MCP denied" }],
isError: true,
},
}),
{ status: 200 },
);
}) as typeof fetch;
try {
await zaiSearch("example", undefined);
throw new Error("expected failure");
} catch (err) {
expect(err).toBeInstanceOf(ZaiApiError);
expect("kind" in (err as object)).toBe(false);
expect((err as ZaiApiError).code).toBeUndefined();
expect((err as ZaiApiError).status).toBe(200);
expect((err as ZaiApiError).rawMessage).toBe("MCP denied");
expect((err as Error).message).toBe("Z.AI API error: MCP denied");
}
});
});