源码
tests/desktop/model-retry.spec.ts
import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { chunk, done, fakeServer, send } from "../fake-server.ts";
import { runCleanup } from "../helpers/cleanup.ts";
import { launchDesktop } from "../helpers/desktop.ts";
import { expect, test } from "../helpers/test.ts";
test("shows Pi agent reconnect status, survives reload, recovers, stops waits and only then exposes terminal errors", async () => {
test.setTimeout(60000);
const dir = await mkdtemp(join(tmpdir(), "zpi-model-retry-"));
const project = join(dir, "project");
await mkdir(project);
await mkdir(join(project, ".zpi"));
await writeFile(join(project, ".zpi/settings.json"), JSON.stringify({ retry: { baseDelayMs: 12000 } }));
const counts = new Map<string, number>();
const finishRecoveredStream = Promise.withResolvers<void>();
const server = await fakeServer((body, response) => {
const messages = body.messages as unknown as { role: string; content: string }[];
const prompt = messages.filter((m) => m.role === "user").at(-1)?.content ?? "";
const count = (counts.get(prompt) ?? 0) + 1;
counts.set(prompt, count);
if (prompt === "partial") {
send(response, chunk({ content: count === 1 ? "discarded partial" : "recovered stream" }));
if (count === 1) response.end();
else done(response);
return;
}
if (prompt === "recover" && count > 3) {
send(response, chunk({ content: "recovered answer" }));
void finishRecoveredStream.promise.then(() => done(response));
return;
}
const delay = prompt === "recover" ? "0.4" : prompt === "stop" ? "120" : "0";
response.writeHead(429, { "content-type": "application/json", "retry-after": delay });
response.end(
JSON.stringify({
error: {
message: "Rate limit reached for requests",
...(prompt === "quota" ? { code: "insufficient_quota" } : {}),
},
}),
);
});
let app: Awaited<ReturnType<typeof launchDesktop>>;
try {
app = await launchDesktop({ dir, project, url: server.url });
const page = await app.firstWindow();
await page.getByRole("button", { name: "添加项目", exact: true }).first().click();
await expect(page.getByRole("button", { name: "选择项目", exact: true })).toHaveText("project");
const submit = async (text: string) => {
await page.getByLabel("消息", { exact: true }).fill(text);
await page.getByLabel("发送", { exact: true }).click();
await expect(page.getByTestId("run").last().locator(".user-message-text")).toHaveText(text);
};
await submit("recover");
const run = () => page.getByTestId("run").last();
await expect(run().locator(".chat-loading-spinner")).toBeVisible();
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
await expect(page.getByTestId("api-retry-status")).toHaveText("重新连接中… 1/10");
await expect(run()).toHaveAttribute("data-status", "running");
await expect(run().locator(".notice-error")).toHaveCount(0);
await expect(run().getByRole("button", { name: "立即重试", exact: true })).toBeVisible();
await expect(page.getByTestId("api-retry-status")).toHaveAttribute("title", /HTTP 429/);
await expect(page.getByTestId("api-retry-status")).toHaveClass("model-wait-message");
await expect(page.getByTestId("api-retry-status")).toHaveCSS("font-size", "12px");
await expect(page.getByTestId("api-retry-status")).toHaveCSS("font-weight", "400");
await expect(page.getByTestId("api-retry-status")).toHaveCSS("animation-name", "none");
await expect(run().locator(".chat-loading-spinner")).toBeVisible();
await page.screenshot({ path: "test-results/model-retry-light.png" });
await page.getByRole("button", { name: "设置", exact: true }).click();
const settings = page.getByRole("region", { name: "设置", exact: true });
await settings.getByRole("button", { name: "界面设置", exact: true }).click();
await settings.getByRole("combobox", { name: "界面主题", exact: true }).click();
await page.getByRole("option", { name: "深色", exact: true }).click();
await settings.getByLabel("关闭设置").click();
await expect(page.locator("html")).toHaveAttribute("data-theme", "dark");
await page.screenshot({ path: "test-results/model-retry-dark.png" });
await page.reload();
await expect(page.getByTestId("api-retry-status")).toHaveText("重新连接中… 1/10");
await run().getByRole("button", { name: "立即重试", exact: true }).click();
await expect(run().locator(".answer")).toHaveText("recovered answer", { timeout: 15_000 });
await expect(run()).toHaveAttribute("data-status", "running");
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
await page.reload();
await expect(run().locator(".answer")).toHaveText("recovered answer");
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
finishRecoveredStream.resolve();
await expect(run()).toHaveAttribute("data-status", "completed");
expect(counts.get("recover")).toBe(4);
await submit("partial");
await expect(run()).toHaveAttribute("data-status", "completed", { timeout: 15000 });
expect(counts.get("partial")).toBe(2);
await expect(run().locator(".answer")).toHaveText("recovered stream");
await expect(run()).not.toContainText("discarded partial");
await submit("stop");
await expect(page.getByTestId("api-retry-status")).toHaveText("重新连接中… 1/10");
await page.getByLabel("停止", { exact: true }).click();
await expect(run()).toHaveAttribute("data-status", "aborted");
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
expect(counts.get("stop")).toBe(1);
await submit("quota");
await expect(run()).toHaveAttribute("data-status", "error");
await expect(run().locator(".notice-turn")).toContainText("服务请求过于频繁或额度不足");
// insufficient_quota 属于不可恢复的额度错误，不应消耗 SDK 或会话重试次数。
expect(counts.get("quota")).toBe(1);
await writeFile(join(project, ".zpi/settings.json"), JSON.stringify({ retry: { baseDelayMs: 0 } }));
await page.getByLabel("新建任务", { exact: true }).first().click();
await expect(page.locator(".draft-composer-header")).toBeVisible();
await submit("exhaust");
await expect(run()).toHaveAttribute("data-status", "error");
expect(counts.get("exhaust")).toBe(33);
await expect(run().locator(".notice-turn")).toContainText("服务请求过于频繁或额度不足");
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
} finally {
finishRecoveredStream.resolve();
await runCleanup([
() => rm(dir, { recursive: true, force: true }),
() => server.close(),
() => app?.close(),
]);
}
});
test("recovers a heartbeat-only stream and stops a live request without waiting for the idle deadline", async () => {
const dir = await mkdtemp(join(tmpdir(), "zpi-stream-stall-"));
const project = join(dir, "project");
await mkdir(project);
await mkdir(join(project, ".zpi"));
await writeFile(
join(project, ".zpi/settings.json"),
JSON.stringify({ streamIdleTimeoutMs: 1000, retry: { baseDelayMs: 3000 } }),
);
const counts = new Map<string, number>();
const closed: string[] = [];
const server = await fakeServer((body, response) => {
const messages = body.messages as unknown as { role: string; content: string }[];
const prompt = messages.filter((m) => m.role === "user").at(-1)?.content ?? "";
const count = (counts.get(prompt) ?? 0) + 1;
counts.set(prompt, count);
if (prompt === "stall" && count > 1) {
send(response, chunk({ content: "recovered after idle timeout" }));
done(response);
return;
}
response.writeHead(200, { "content-type": "text/event-stream" });
response.write(": keepalive\n\n");
const heartbeat = setInterval(() => response.write(": keepalive\n\n"), 100);
response.on("close", () => {
clearInterval(heartbeat);
closed.push(prompt);
});
});
let app: Awaited<ReturnType<typeof launchDesktop>>;
try {
app = await launchDesktop({ dir, project, url: server.url });
const page = await app.firstWindow();
await page.getByRole("button", { name: "添加项目", exact: true }).first().click();
await expect(page.getByRole("button", { name: "选择项目", exact: true })).toHaveText("project");
const submit = async (text: string) => {
await page.getByLabel("消息", { exact: true }).fill(text);
await page.getByLabel("发送", { exact: true }).click();
};
await submit("stall");
const run = () => page.getByTestId("run").last();
await expect(page.getByTestId("api-retry-status")).toHaveText("重新连接中… 1/10");
await expect(run().locator(".notice-error")).toHaveCount(0);
await expect.poll(() => closed).toContain("stall");
await page.reload();
await expect(page.getByTestId("api-retry-status")).toHaveText("重新连接中… 1/10");
await expect(run()).toHaveAttribute("data-status", "completed", { timeout: 10000 });
await expect(run().locator(".answer")).toHaveText("recovered after idle timeout");
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
expect(counts.get("stall")).toBe(2);
await submit("stop live request");
await expect.poll(() => counts.get("stop live request")).toBe(1);
await page.getByLabel("停止", { exact: true }).click();
await expect(run()).toHaveAttribute("data-status", "aborted");
await expect.poll(() => closed).toContain("stop live request");
expect(counts.get("stop live request")).toBe(1);
await expect(page.getByTestId("api-retry-status")).toHaveCount(0);
} finally {
await runCleanup([
() => rm(dir, { recursive: true, force: true }),
() => server.close(),
() => app?.close(),
]);
}
});
test("shows a quiet first request and retries it immediately within the same turn", async () => {
test.setTimeout(45000);
const dir = await mkdtemp(join(tmpdir(), "zpi-model-wait-"));
const project = join(dir, "project");
await mkdir(project);
await mkdir(join(project, ".zpi"));
await writeFile(
join(project, ".zpi/settings.json"),
JSON.stringify({ streamFirstResponseTimeoutMs: 60000 }),
);
let requests = 0;
let closed = false;
const finishRetry = Promise.withResolvers<void>();
const server = await fakeServer((_body, response) => {
requests++;
if (requests > 1) {
void finishRetry.promise.then(() => {
send(response, chunk({ content: "recovered immediately" }));
done(response);
});
return;
}
response.writeHead(200, { "content-type": "text/event-stream" });
response.write(": keepalive\n\n");
const heartbeat = setInterval(() => response.write(": keepalive\n\n"), 100);
response.on("close", () => {
clearInterval(heartbeat);
closed = true;
});
});
let app: Awaited<ReturnType<typeof launchDesktop>>;
try {
app = await launchDesktop({ dir, project, url: server.url });
const page = await app.firstWindow();
await page.getByRole("button", { name: "添加项目", exact: true }).first().click();
await expect(page.getByRole("button", { name: "选择项目", exact: true })).toHaveText("project");
await page.getByLabel("消息", { exact: true }).fill("retry the waiting request");
await page.getByLabel("发送", { exact: true }).click();
const run = page.getByTestId("run");
await expect(run).toContainText("等待模型响应", { timeout: 20000 });
await expect(run.locator(".model-wait-message")).toHaveText(/等待模型响应 · \d+ 秒/);
await expect(run.locator(".model-wait-message")).toHaveCSS("font-size", "12px");
await expect(run.locator(".model-wait-message")).toHaveCSS("font-weight", "400");
const waiting = run.locator(".model-wait-row");
const retry = waiting.getByRole("button", { name: "立即重试", exact: true });
await expect(retry).toHaveText("重试");
await expect(retry).toHaveCSS("text-decoration-line", "none");
await expect(retry).toHaveCSS("height", "24px");
await expect(retry).toHaveCSS("border-top-width", "1px");
await expect(waiting).toHaveCSS("border-top-width", "0px");
await expect(run.locator(".notice-turn")).toHaveCount(0);
const spinnerBounds = await waiting.locator(".chat-loading-icon").boundingBox();
const retryBounds = await retry.boundingBox();
if (!spinnerBounds || !retryBounds) throw new Error("Missing model wait controls");
expect(
Math.abs(spinnerBounds.y + spinnerBounds.height / 2 - retryBounds.y - retryBounds.height / 2),
).toBeLessThan(1);
await run.screenshot({ path: "test-results/model-wait-light.png" });
await page.emulateMedia({ colorScheme: "dark" });
await expect(page.locator("html")).toHaveAttribute("data-theme", "dark");
await run.evaluate(async (element) => {
await Promise.all(
element
.getAnimations({ subtree: true })
.filter((animation) => animation.effect?.getTiming().iterations !== Infinity)
.map((animation) => animation.finished.catch(() => {})),
);
});
await run.screenshot({ path: "test-results/model-wait-dark.png" });
expect(requests).toBe(1);
const waitStyle = await run.locator(".model-wait-message").evaluate((element) => {
const style = getComputedStyle(element);
return { color: style.color, fontSize: style.fontSize, lineHeight: style.lineHeight };
});
await page.reload();
await expect(run.getByRole("button", { name: "立即重试", exact: true })).toBeVisible();
await run.getByRole("button", { name: "立即重试", exact: true }).click();
await expect(run.locator(".model-wait-message")).toHaveText("正在重新请求…");
await expect.poll(() => requests).toBe(2);
await page.reload();
await expect(run.locator(".model-wait-message")).toHaveText("正在重新请求…");
await expect(run).not.toContainText("1/10");
expect(
await run.locator(".model-wait-message").evaluate((element) => {
const style = getComputedStyle(element);
return { color: style.color, fontSize: style.fontSize, lineHeight: style.lineHeight };
}),
).toEqual(waitStyle);
await run.screenshot({ path: "test-results/model-manual-retry-dark.png" });
finishRetry.resolve();
await expect(run).toHaveAttribute("data-status", "completed", { timeout: 5000 });
await expect(run.locator(".answer")).toHaveText("recovered immediately");
await expect(run).toHaveCount(1);
await expect(run.locator(".notice-turn")).toHaveCount(0);
expect(requests).toBe(2);
await expect.poll(() => closed).toBe(true);
} finally {
finishRetry.resolve();
await runCleanup([
() => rm(dir, { recursive: true, force: true }),
() => server.close(),
() => app?.close(),
]);
}
});