WebSocket 帧大小上限：/api/ws 1 KiB、/agent 16 MiB，运行器按字节分批上传工作记录

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);

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);
+  });
 });

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<RunningServer>
     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));

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<number> {
+  return new Promise((resolve) => push.socket.once("close", (code) => resolve(code)));
+}
+
+async function nextMessagePush(push: PushClient): Promise<Extract<ServerPush, { type: "message" }>> {
+  for (;;) {
+    const next = await push.next();
+    if (next.type === "message") return next;
+  }
+}
+
+function reliableAcked(agent: AgentClient, seq: number): Promise<UplinkAck> {
+  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);
+  });
+});

packages/shared/src/protocol.ts

@@ -677,6 +677,9 @@ export const UplinkRejection = z.object({
 });
 export type UplinkRejection = z.infer<typeof UplinkRejection>;
 
+/** 控制面 /agent 收的单帧上限。超限的可靠帧被关后重连照样重发，所以发送方要先按它切，不能等控制面拒。 */
+export const AGENT_FRAME_MAX_BYTES = 16 * 1024 * 1024;
+
 export const UplinkFrame = z.object({
   type: z.literal("uplink"),
   conn: z.number().int().positive(),