From c7132025f7703f64ae0af759edb7bab531f2e4bf Mon Sep 17 00:00:00 2001
From: yuetsh <517252939@qq.com>
Date: Sun, 30 Aug 2026 08:29:36 -0600
Subject: [PATCH] =?UTF-8?q?fix(collab):=20=E4=BA=8C=E8=BF=9B=E5=88=B6?=
 =?UTF-8?q?=E5=B8=A7=E7=9A=84=E5=AE=BD=E6=9D=BE=E9=99=90=E6=B5=81=E6=A1=A3?=
 =?UTF-8?q?=E5=90=8D=E5=AD=98=E5=AE=9E=E4=BA=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

两档共用一个令牌桶：桶按严格档的 20 初始化，之后每条文本帧（含 30 秒
一次的心跳）都会 Math.min(RATE_BURST, ...) 把它压回 20，二进制帧那档
标的 200 突发根本拿不到。实测连打 150 帧在第 101 帧被 1008 踢下线——
正是设计里要避免的「协作编辑时打字把自己踢掉」。

改成两个独立的桶。文本帧仍是 20 / 每秒 2，二进制帧 200 / 每秒 100。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016DHxhKNxXfG89JnVzHbvgj
---
 apps/api/src/websocket.ts | 52 ++++++++++++++++++++++++++++-----------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/apps/api/src/websocket.ts b/apps/api/src/websocket.ts
index 048ade7..b889809 100644
--- a/apps/api/src/websocket.ts
+++ b/apps/api/src/websocket.ts
@@ -55,14 +55,27 @@ export function isAllowedWebSocketOrigin(origin: string | null, url: URL) {
   )
 }
 
+interface RateBucket {
+  tokens: number
+  updatedAt: number
+}
+
 export interface SubmissionSocketData {
   userId: number
   /** 同一个 Bun.serve 只能挂一个 websocket handler，用它区分通道 */
   kind: "submissions" | "config" | "collab"
   /** 握手时那张会话的 token，留着定期确认它还没被登出 / 过期，见 sweepSessions */
   token: string
-  /** 令牌桶，open 时初始化，见 allowMessage */
-  rate?: { tokens: number; updatedAt: number }
+  /** 文本控制帧的令牌桶，open 时初始化，见 allowMessage */
+  rate?: RateBucket
+  /**
+   * collab 二进制帧的令牌桶，和 rate 分开。
+   *
+   * 共用一个桶的话宽松档名存实亡：每条文本帧（含 30 秒一次的心跳）都会
+   * `Math.min(RATE_BURST, ...)` 把桶压回 20，二进制帧再怎么标 200 突发也拿不到。
+   * 实测连打 150 帧会在第 101 帧被 1008 踢下线。
+   */
+  binaryRate?: RateBucket
   /** 握手时从会话里读，三种 kind 都会填；collab 通道用它判断老师身份、拼 room_open 里的姓名 */
   username?: string
   adminType?: string
@@ -89,21 +102,31 @@ const RATE_REFILL_PER_SECOND = 2
 const COLLAB_BINARY_BURST = 200
 const COLLAB_BINARY_REFILL_PER_SECOND = 100
 
-function allowMessage(
-  ws: Bun.ServerWebSocket<SubmissionSocketData>,
-  burst = RATE_BURST,
-  refillPerSecond = RATE_REFILL_PER_SECOND,
-) {
+function consume(bucket: RateBucket, burst: number, refillPerSecond: number) {
   const now = Date.now()
-  const rate = (ws.data.rate ??= { tokens: burst, updatedAt: now })
-  const refill = ((now - rate.updatedAt) / 1000) * refillPerSecond
-  rate.tokens = Math.min(burst, rate.tokens + refill)
-  rate.updatedAt = now
-  if (rate.tokens < 1) return false
-  rate.tokens -= 1
+  const refill = ((now - bucket.updatedAt) / 1000) * refillPerSecond
+  bucket.tokens = Math.min(burst, bucket.tokens + refill)
+  bucket.updatedAt = now
+  if (bucket.tokens < 1) return false
+  bucket.tokens -= 1
   return true
 }
 
+/** 文本帧：严格档。会查库，走这一档的都按最坏情况算 */
+function allowMessage(ws: Bun.ServerWebSocket<SubmissionSocketData>) {
+  const bucket = (ws.data.rate ??= { tokens: RATE_BURST, updatedAt: Date.now() })
+  return consume(bucket, RATE_BURST, RATE_REFILL_PER_SECOND)
+}
+
+/** collab 二进制帧：宽松档，独立的桶 —— 见 binaryRate 的注释 */
+function allowCollabBinary(ws: Bun.ServerWebSocket<SubmissionSocketData>) {
+  const bucket = (ws.data.binaryRate ??= {
+    tokens: COLLAB_BINARY_BURST,
+    updatedAt: Date.now(),
+  })
+  return consume(bucket, COLLAB_BINARY_BURST, COLLAB_BINARY_REFILL_PER_SECOND)
+}
+
 /**
  * 当前挂着的连接。Bun 不提供遍历连接的接口，要定期巡检就得自己登记。
  * open 时加入、close 时移除，见 sweepSessions。
@@ -190,6 +213,7 @@ export function submissionWebSocketHandler(): Bun.WebSocketHandler<SubmissionSoc
       liveSockets.add(ws)
       ws.data.rate = { tokens: RATE_BURST, updatedAt: Date.now() }
       if (ws.data.kind === "collab") {
+        ws.data.binaryRate = { tokens: COLLAB_BINARY_BURST, updatedAt: Date.now() }
         handleCollabOpen(ws)
         return
       }
@@ -203,7 +227,7 @@ export function submissionWebSocketHandler(): Bun.WebSocketHandler<SubmissionSoc
     message(ws, message) {
       if (ws.data.kind === "collab") {
         if (typeof message !== "string") {
-          if (!allowMessage(ws, COLLAB_BINARY_BURST, COLLAB_BINARY_REFILL_PER_SECOND)) {
+          if (!allowCollabBinary(ws)) {
             ws.close(1008, "Too many messages")
             return
           }