diff --git a/packages/bot-runner/src/work-mirror.ts b/packages/bot-runner/src/work-mirror.ts index ab92c4ec..dcb17850 100644 --- a/packages/bot-runner/src/work-mirror.ts +++ b/packages/bot-runner/src/work-mirror.ts @@ -13,6 +13,7 @@ import { import { basename, dirname, join, resolve } from "node:path"; import { z } from "zod"; import { + AGENT_FRAME_MAX_BYTES, EMPTY_WORK_HEAD, HubLaneSnapshot, advanceWorkHead, @@ -32,6 +33,9 @@ import { } from "@earendil-works/pi-coding-agent"; import type { WorkAppender, WorkAppendOutcome } from "./lane-fence.js"; +/** 留给 work_append 其余字段和可靠上行信封。 */ +const WORK_APPEND_LINES_MAX_BYTES = AGENT_FRAME_MAX_BYTES - 64 * 1024; + const MirrorStream = z.object({ sessionFile: z.string().nullable(), seq: z.number().int().nonnegative(), @@ -291,13 +295,26 @@ export class WorkMirror { this.#snapshot !== undefined && this.#snapshot.toolFingerprint !== toolListFingerprint(localOnlyToolNames) ? this.#filter(entries.slice(0, this.#uploaded).map((entry) => JSON.stringify(entry))).snapshot : this.#snapshot; - const retainer = createRetainer({ localOnlyToolNames }, start); + let retainer = createRetainer({ localOnlyToolNames }, start); const lines: string[] = []; + let bytes = 0; for (const entry of pending) { const retained = retainer.retain(JSON.stringify(entry)); if (retained.status === "rejected") return this.#stop({ kind: "invalid", reason: retained.reason }); + bytes += Buffer.byteLength(JSON.stringify(retained.line)) + 1; + if (bytes > WORK_APPEND_LINES_MAX_BYTES) { + if (lines.length === 0) { + this.#options.log(`工作记录第 ${this.#seq} 条超过上行帧上限`); + return this.#stop({ kind: "invalid" }); + } + // 超限的那条已经记进了 retainer 见过的 id，不重来的话下一帧再送它会被判成重复。 + retainer = createRetainer({ localOnlyToolNames }, start); + for (const kept of pending.slice(0, lines.length)) retainer.retain(JSON.stringify(kept)); + break; + } lines.push(retained.line); } + const batch = pending.slice(0, lines.length); const head = await advanceWorkHead(this.#head, lines); if (revision !== this.#revision || this.#control.signal.aborted) return { kind: "disconnected" }; let result: WorkAppendOutcome; @@ -379,7 +396,7 @@ export class WorkMirror { filtered.snapshot, ); const replacement: string[] = []; - for (const entry of pending) { + for (const entry of batch) { const retained = refreshed.retain(JSON.stringify(entry)); if (retained.status === "rejected") return this.#stop({ kind: "invalid", reason: retained.reason }); replacement.push(retained.line); diff --git a/packages/bot-runner/test/work-mirror.test.ts b/packages/bot-runner/test/work-mirror.test.ts index 91b69c06..36a37581 100644 --- a/packages/bot-runner/test/work-mirror.test.ts +++ b/packages/bot-runner/test/work-mirror.test.ts @@ -2,6 +2,7 @@ import { existsSync, mkdtempSync, mkdirSync, readFileSync, rmSync } from "node:f import { tmpdir } from "node:os"; import { join } from "node:path"; import { + AGENT_FRAME_MAX_BYTES, BotId, EMPTY_WORK_HEAD, GroupId, @@ -15,6 +16,7 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import type { WorkAppendOutcome } from "../src/lane-fence.js"; import { WorkMirror } from "../src/work-mirror.js"; +const MiB = 1024 * 1024; const roots: string[] = []; afterEach(() => { @@ -715,4 +717,44 @@ describe("WorkMirror", () => { expect(h.sent.at(-1)?.fromSeq).toBe(1); expect(h.logs.some((line) => line.includes("head_mismatch"))).toBe(true); }); + + it("一批条目加起来超过上行帧上限时按字节分批，每帧都在上限以内，条目一条不落", async () => { + const accepted: string[] = []; + const h = setup(async (frame) => { + accepted.push(...frame.lines); + return { kind: "ok" }; + }); + const manager = await h.mirror.prepare(1, { seq: 0, head: EMPTY_WORK_HEAD }, RunId.parse("run-big")); + for (let index = 0; index < 24; index += 1) + appendRawMessage(manager, assistant([{ type: "thinking", thinking: `${index}`.padEnd(MiB, "x") }], "stop")); + + await expect(h.mirror.flush("turn_end", RunId.parse("run-big"))).resolves.toEqual({ kind: "ok" }); + + expect(h.sent.length, "24 条装不进一帧，要分成几帧").toBeGreaterThan(1); + for (const frame of h.sent) + expect(Buffer.byteLength(JSON.stringify(frame)), "每帧都要能过 /agent 的帧上限").toBeLessThan( + AGENT_FRAME_MAX_BYTES, + ); + expect(accepted).toHaveLength(24); + expect(h.sent.map((frame) => frame.fromSeq)).toEqual( + h.sent.map((_, index) => h.sent.slice(0, index).reduce((sum, frame) => sum + frame.lines.length, 0)), + ); + }); + + it("单条就超过上行帧上限时停流，前面的照传，超限的那条不发出去", async () => { + const h = setup(async () => ({ kind: "ok" })); + const manager = await h.mirror.prepare(1, { seq: 0, head: EMPTY_WORK_HEAD }, RunId.parse("run-huge")); + appendUser(manager, "超限之前的一条"); + appendRawMessage( + manager, + assistant([{ type: "thinking", thinking: "x".repeat(AGENT_FRAME_MAX_BYTES + 1) }], "stop"), + ); + + await expect(h.mirror.flush("turn_end", RunId.parse("run-huge"))).resolves.toEqual({ kind: "invalid" }); + await expect(h.mirror.flush("turn_end", RunId.parse("run-huge"))).resolves.toEqual({ kind: "invalid" }); + + expect(h.sent).toHaveLength(1); + expect(h.sent[0]?.lines).toHaveLength(1); + expect(h.logs.some((line) => line.includes("超过上行帧上限"))).toBe(true); + }); }); diff --git a/packages/control-plane/src/server.ts b/packages/control-plane/src/server.ts index 445a1c21..e7516f9e 100644 --- a/packages/control-plane/src/server.ts +++ b/packages/control-plane/src/server.ts @@ -1,4 +1,5 @@ import { serve } from "@hono/node-server"; +import { AGENT_FRAME_MAX_BYTES } from "@baize/shared"; import { parse as parseCookie } from "hono/utils/cookie"; import type { IncomingMessage } from "node:http"; import type { Duplex } from "node:stream"; @@ -38,6 +39,9 @@ import { PokerSweeper, type PokerSweeperOptions } from "./poker/sweeper.js"; export const WS_PATH = "/api/ws"; +/** 网页在 /api/ws 上不发数据帧，来了超过 1 KiB 的就不是我们的网页。 */ +const PUSH_MAX_PAYLOAD = 1024; + /** 扫一遍放了一天还没有结论的审批和还没人确认的交接。提醒本身只发一次，扫得勤一点只是让它准时一点。 */ const REMIND_INTERVAL_MS = 60 * 60 * 1000; @@ -181,8 +185,8 @@ export async function startServer(options: StartOptions): Promise port: options.port, ...(options.hostname === undefined ? {} : { hostname: options.hostname }), }); - const clients = new WebSocketServer({ noServer: true }); - const agentSockets = new WebSocketServer({ noServer: true }); + const clients = new WebSocketServer({ noServer: true, maxPayload: PUSH_MAX_PAYLOAD }); + const agentSockets = new WebSocketServer({ noServer: true, maxPayload: AGENT_FRAME_MAX_BYTES }); // 网页往上只发输入、结束和看密码，最大的是 2000 字的整段上屏，64KB 绰绰有余。 const viewers = new WebSocketServer({ noServer: true, maxPayload: 64 * 1024 }); @@ -349,6 +353,8 @@ function acceptClient( return; } clients.handleUpgrade(request, socket, head, (ws) => { + // ws 遇到超过 maxPayload 之类的协议错误会自己关连接并 emit error，没人监听的 error 会变成未捕获异常。 + ws.on("error", () => ws.close()); hub.add(user.id, ws); for (const snapshot of groupTurns.snapshots(user.id)) ws.send(JSON.stringify(snapshot)); void notifications.connected(user.id).catch((error: unknown) => console.error("补发免打扰汇总出错", error)); diff --git a/packages/control-plane/test/ws-max-payload.test.ts b/packages/control-plane/test/ws-max-payload.test.ts new file mode 100644 index 00000000..39f791f3 --- /dev/null +++ b/packages/control-plane/test/ws-max-payload.test.ts @@ -0,0 +1,114 @@ +import { AgentCloseCode, CreateBotResponse, SendMessageResponse, UplinkAck, type ServerPush } from "@baize/shared"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import type { RawData } from "ws"; +import { helloAgent, openAgent, type AgentClient } from "./helpers/agent.js"; +import { startTestServer, type TestServer } from "./helpers/client.js"; +import { createTestDb, type TestDb } from "./helpers/db.js"; +import { PASSWORD, makeComputer, makeUser } from "./helpers/fixtures.js"; +import { connectPush, type PushClient } from "./helpers/ws.js"; + +const MiB = 1024 * 1024; +const MESSAGE_TOO_BIG = 1009; + +let database: TestDb; +let server: TestServer; + +beforeAll(async () => { + database = await createTestDb(); + server = await startTestServer(database.db); +}); + +afterAll(async () => { + await server.close(); + await database.close(); +}); + +async function ownerWithGroup() { + const user = await makeUser(database.db); + const client = server.client(); + expect((await client.login(user.email, PASSWORD)).status).toBe(200); + const created = CreateBotResponse.parse(await (await client.post("/api/bots", {})).json()); + return { user, client, groupId: created.group.id, cookie: client.cookie ?? "" }; +} + +function pushClosed(push: PushClient): Promise { + return new Promise((resolve) => push.socket.once("close", (code) => resolve(code))); +} + +async function nextMessagePush(push: PushClient): Promise> { + for (;;) { + const next = await push.next(); + if (next.type === "message") return next; + } +} + +function reliableAcked(agent: AgentClient, seq: number): Promise { + return new Promise((resolve, reject) => { + const received = agent.uplinkAcks().find((entry) => entry.upTo >= seq); + if (received !== undefined) return resolve(received); + const timer = setTimeout(() => reject(new Error(`等上行 ${seq} 确认超时`)), 10_000); + const onMessage = (raw: RawData) => { + const value: unknown = JSON.parse(raw.toString()); + if (typeof value !== "object" || value === null || !("type" in value) || value.type !== "uplink_ack") return; + const ack = UplinkAck.parse(value); + if (ack.upTo < seq) return; + clearTimeout(timer); + agent.socket.off("message", onMessage); + resolve(ack); + }; + agent.socket.on("message", onMessage); + }); +} + +describe("/api/ws 帧大小", () => { + it("网页发来超过 1 KiB 的帧，这条连接以 1009 关闭，控制面照常推送", async () => { + const { client, groupId, cookie } = await ownerWithGroup(); + const oversized = await connectPush(server.server.port, cookie); + const closed = pushClosed(oversized); + oversized.socket.send("x".repeat(2048)); + expect(await closed, "超限帧要以 1009 关闭").toBe(MESSAGE_TOO_BIG); + + const push = await connectPush(server.server.port, cookie); + try { + push.socket.send("x".repeat(512)); + const { message } = SendMessageResponse.parse( + await (await client.post(`/api/groups/${groupId}/messages`, { text: "还在吗" })).json(), + ); + const pushed = await nextMessagePush(push); + expect(pushed.message.id, "上限以内的帧不断连，推送照常到").toBe(message.id); + } finally { + push.close(); + } + }); +}); + +describe("/agent 帧大小与握手", () => { + it("没报 hello 就发来超过 16 MiB 的帧，以 1009 关闭", async () => { + const agent = await openAgent(server.server.port); + const closed = agent.closed(10_000); + agent.socket.send(Buffer.alloc(16 * MiB + 1, 0x20)); + expect(await closed).toBe(MESSAGE_TOO_BIG); + }); + + it("握手之后 4 MiB 的上行照收照确认，超过 16 MiB 的以 1009 关闭", async () => { + const user = await makeUser(database.db); + const computer = await makeComputer(database.db, user.id); + const { agent } = await helloAgent(server.server.port, computer.computerId, computer.token); + try { + const seq = agent.sendMalformedReliable({ type: "padding", text: "x".repeat(4 * MiB) }); + const ack = await reliableAcked(agent, seq); + expect(ack.upTo, "4 MiB 的帧到了控制面并被确认").toBeGreaterThanOrEqual(seq); + + const closed = agent.closed(10_000); + agent.sendMalformedReliable({ type: "padding", text: "x".repeat(16 * MiB) }); + expect(await closed).toBe(MESSAGE_TOO_BIG); + } finally { + agent.close(); + } + }); + + it("连上 10 秒不报 hello，控制面按坏 hello 关掉", { timeout: 20_000 }, async () => { + const agent = await openAgent(server.server.port); + expect(await agent.closed(15_000)).toBe(AgentCloseCode.badHello); + }); +}); diff --git a/packages/shared/src/protocol.ts b/packages/shared/src/protocol.ts index 4a0c7a16..54b27ac2 100644 --- a/packages/shared/src/protocol.ts +++ b/packages/shared/src/protocol.ts @@ -677,6 +677,9 @@ export const UplinkRejection = z.object({ }); export type UplinkRejection = z.infer; +/** 控制面 /agent 收的单帧上限。超限的可靠帧被关后重连照样重发，所以发送方要先按它切，不能等控制面拒。 */ +export const AGENT_FRAME_MAX_BYTES = 16 * 1024 * 1024; + export const UplinkFrame = z.object({ type: z.literal("uplink"), conn: z.number().int().positive(),