feat(chat-ws): authenticate v2 after connect (#2309)

Signed-off-by: RainbowBird <git@luoling.moe>

server/apps/api/src/app.ts

@@ -50,6 +50,7 @@ import { createCharacterRoutes } from './routes/characters'
 import { createChatWsRuntime } from './routes/chat-ws/runtime'
 import { createChatWsV1Handlers } from './routes/chat-ws/v1'
 import { createChatWsV2Handlers } from './routes/chat-ws/v2'
+import { createChatWsPayloadLimit } from './routes/chat-ws/v2/payload-limit'
 import { createChatRoutes } from './routes/chats'
 import { createFluxRoutes } from './routes/flux'
 import { createInternalAuthRoutes } from './routes/internal-auth'
@@ -102,6 +103,8 @@ interface AppDeps {
   providerCatalogService: ProviderCatalogService
 }
 
+const MAX_UNAUTHENTICATED_CHAT_WS_FRAME_BYTES = 8192
+
 export async function buildApp(deps: AppDeps) {
   const logger = useLogger('app').useGlobalConfig()
 
@@ -144,19 +147,46 @@ export async function buildApp(deps: AppDeps) {
   }
 
   // WebSocket setup — must be registered BEFORE bodyLimit middleware
-  const { injectWebSocket, upgradeWebSocket } = createNodeWebSocket({ app })
+  const { injectWebSocket, upgradeWebSocket, wss } = createNodeWebSocket({ app })
+  const chatWsPayloadLimit = createChatWsPayloadLimit(MAX_UNAUTHENTICATED_CHAT_WS_FRAME_BYTES)
+  wss.on('connection', (socket, request) => {
+    if (new URL(request.url ?? '/', 'http://localhost').pathname !== '/ws/v2/chat')
+      return
+
+    // NOTICE:
+    // @hono/node-ws creates one ws server with a 100 MiB default frame limit.
+    // The library has no per-route maxPayload option, so use ws's receiver limit.
+    // Source: @hono/node-ws@1.3.0 dist/index.js; ws@8.20.0 Receiver._maxPayload.
+    // Removal condition: @hono/node-ws supports maxPayload per upgrade route.
+    chatWsPayloadLimit.restrict(socket)
+  })
   // Per-process stable id used by the chat-ws sub callback to skip echoes of
   // its own publishes. Falls back to a random nanoid when ops do not provide
   // SERVER_INSTANCE_ID, which is fine because we only need uniqueness across
   // simultaneously-running api instances, not across restarts.
   const instanceId = process.env.SERVER_INSTANCE_ID || nanoid()
   const chatWsRuntime = createChatWsRuntime(deps.redis, instanceId, deps.otel?.engagement ?? null)
-  const chatWsV2Setup = createChatWsV2Handlers(deps.chatService, deps.redis, instanceId, deps.otel?.engagement ?? null, chatWsRuntime)
+  const chatWsV2Setup = createChatWsV2Handlers(
+    deps.chatService,
+    deps.redis,
+    instanceId,
+    async (token) => {
+      const session = await resolveRequestAuth(
+        deps.db,
+        deps.env,
+        new Headers({ Authorization: `Bearer ${token}` }),
+      )
+      return session?.user?.id ?? null
+    },
+    deps.otel?.engagement ?? null,
+    chatWsRuntime,
+    chatWsPayloadLimit.restore,
+  )
   const chatWsV1Setup = createChatWsV1Handlers(deps.chatService, deps.redis, instanceId, deps.otel?.engagement ?? null, chatWsRuntime)
 
   // `/ws/chat` keeps query-token authentication for deployed clients. The
   // Eventa beta.15 adapter accepts their beta.13 envelopes. `/ws/v2/chat`
-  // keeps the versioned endpoint for its updated authentication flow.
+  // authenticates after the connection opens.
   app.get('/ws/chat', upgradeWebSocket(async (c) => {
     const token = c.req.query('token')
     if (!token)
@@ -173,21 +203,7 @@ export async function buildApp(deps: AppDeps) {
     return chatWsV1Setup(session.user.id)
   }))
 
-  app.get('/ws/v2/chat', upgradeWebSocket(async (c) => {
-    const token = c.req.query('token')
-    if (!token)
-      return createUnauthorizedWsEvents()
-
-    const session = await resolveRequestAuth(
-      deps.db,
-      deps.env,
-      new Headers({ Authorization: `Bearer ${token}` }),
-    )
-    if (!session?.user)
-      return createUnauthorizedWsEvents()
-
-    return chatWsV2Setup(session.user.id)
-  }))
+  app.get('/ws/v2/chat', upgradeWebSocket(() => chatWsV2Setup()))
 
   // Bidirectional streaming TTS proxy. The handler factory builds one ws-to-ws
   // bridge per connection: client ↔ server/apps/api ↔ unspeech ↔ upstream

server/apps/api/src/libs/request-auth.ts

@@ -7,7 +7,7 @@ import { timingSafeEqual } from 'node:crypto'
 
 import { isUserBannedNow } from '@proj-airi/auth-shared'
 import { eq } from 'drizzle-orm'
-import { createRemoteJWKSet, jwtVerify } from 'jose'
+import { createRemoteJWKSet, errors, jwtVerify } from 'jose'
 
 import * as authSchema from '@proj-airi/auth-shared'
 
@@ -98,44 +98,50 @@ async function resolveJWTAccessToken(
   env: TokenIssuerEnv,
   accessToken: string,
 ): Promise<AuthSession | null> {
+  const jwks = getJWKS(env)
+  let payload: Awaited<ReturnType<typeof jwtVerify>>['payload']
   try {
-    const jwks = getJWKS(env)
     // NOTICE: better-auth's jwt() plugin sets issuer to the full baseURL
     // including the path prefix (e.g. "http://localhost:3000/api/auth"),
     // not just the server origin.
-    const { payload } = await jwtVerify(accessToken, jwks, {
+    const verified = await jwtVerify(accessToken, jwks, {
       issuer: `${env.AUTH_SERVER_URL}/api/auth`,
       audience: env.AUTH_SERVER_URL,
     })
-
-    if (!payload.sub)
-      return null
-
-    // The resource server deliberately reads only its authorization projection.
-    // It does not instantiate Better Auth or depend on its internal adapter.
-    const user = await db.query.user.findFirst({
-      where: eq(authSchema.user.id, payload.sub),
-    })
-    if (!user)
-      return null
-
-    return {
-      user,
-      session: {
-        id: payload.jti ?? payload.sub,
-        token: accessToken,
-        userId: payload.sub,
-        createdAt: payload.iat ? new Date(payload.iat * 1000) : new Date(),
-        updatedAt: payload.iat ? new Date(payload.iat * 1000) : new Date(),
-        expiresAt: payload.exp ? new Date(payload.exp * 1000) : new Date(),
-        ipAddress: null,
-        userAgent: null,
-      },
-    }
+    payload = verified.payload
   }
-  catch {
+  catch (error) {
+    // A fetch failure while resolving JWKS is temporary. Let WebSocket auth
+    // return its retryable close code instead of treating a valid token as bad.
+    if (error instanceof TypeError || error instanceof errors.JWKSTimeout)
+      throw error
     return null
   }
+
+  if (!payload.sub)
+    return null
+
+  // The resource server deliberately reads only its authorization projection.
+  // It does not instantiate Better Auth or depend on its internal adapter.
+  const user = await db.query.user.findFirst({
+    where: eq(authSchema.user.id, payload.sub),
+  })
+  if (!user)
+    return null
+
+  return {
+    user,
+    session: {
+      id: payload.jti ?? payload.sub,
+      token: accessToken,
+      userId: payload.sub,
+      createdAt: payload.iat ? new Date(payload.iat * 1000) : new Date(),
+      updatedAt: payload.iat ? new Date(payload.iat * 1000) : new Date(),
+      expiresAt: payload.exp ? new Date(payload.exp * 1000) : new Date(),
+      ipAddress: null,
+      userAgent: null,
+    },
+  }
 }
 
 export async function resolveRequestAuth(

server/apps/api/src/libs/tests/request-auth.test.ts

@@ -7,10 +7,13 @@ import { resolveRequestAuth } from '../request-auth'
 
 vi.mock('jose', () => ({
   createRemoteJWKSet: vi.fn(() => 'mock-jwks'),
+  errors: {
+    JWKSTimeout: class JWKSTimeout extends Error {},
+  },
   jwtVerify: vi.fn(),
 }))
 
-const { createRemoteJWKSet, jwtVerify } = await import('jose')
+const { createRemoteJWKSet, errors, jwtVerify } = await import('jose')
 const mockedCreateRemoteJWKSet = vi.mocked(createRemoteJWKSet)
 const mockedJwtVerify = vi.mocked(jwtVerify)
 
@@ -39,11 +42,15 @@ function createUser(overrides: Partial<RequestAuthSession['user']> = {}): Reques
   }
 }
 
-function createDb(user: RequestAuthSession['user'] | null): Database {
+function createDb(user: RequestAuthSession['user'] | null, failure?: Error): Database {
   return {
     query: {
       user: {
-        findFirst: vi.fn(async () => user),
+        findFirst: vi.fn(async () => {
+          if (failure)
+            throw failure
+          return user
+        }),
       },
     },
   } as unknown as Database
@@ -175,4 +182,34 @@ describe('resolveRequestAuth', () => {
       new Headers({ Authorization: 'Bearer subjectless' }),
     )).toBeNull()
   })
+
+  it('propagates temporary JWKS failures', async () => {
+    mockedJwtVerify.mockRejectedValueOnce(new TypeError('fetch failed'))
+
+    await expect(resolveRequestAuth(
+      createDb(null),
+      mockEnv,
+      new Headers({ Authorization: 'Bearer valid-token' }),
+    )).rejects.toThrow('fetch failed')
+  })
+
+  // https://github.com/moeru-ai/airi/pull/2309#discussion_r3818708556
+  it('propagates JWKS timeouts so WebSocket clients can retry', async () => {
+    mockedJwtVerify.mockRejectedValueOnce(new errors.JWKSTimeout())
+
+    await expect(resolveRequestAuth(
+      createDb(null),
+      mockEnv,
+      new Headers({ Authorization: 'Bearer valid-token' }),
+    )).rejects.toBeInstanceOf(errors.JWKSTimeout)
+  })
+
+  it('propagates database failures after JWT verification', async () => {
+    mockValidJwt()
+    await expect(resolveRequestAuth(
+      createDb(null, new Error('database unavailable')),
+      mockEnv,
+      new Headers({ Authorization: 'Bearer valid-token' }),
+    )).rejects.toThrow('database unavailable')
+  })
 })

server/apps/api/src/libs/tests/ws-auth.test.ts

@@ -15,7 +15,8 @@ describe('createUnauthorizedWsEvents', () => {
   //
   // Accepting the upgrade first and closing with 4001 in `onOpen` gives
   // the client a structured signal so its `onDisconnected` handler can
-  // stop the reconnect loop until the URL (token) actually changes.
+  // stop the reconnect loop until the token actually changes and the client
+  // starts a fresh connection.
   it('closes the socket with WS_CLOSE_UNAUTHORIZED immediately on open', () => {
     const close = vi.fn<(code?: number, reason?: string) => void>()
     const events = createUnauthorizedWsEvents()

server/apps/api/src/libs/ws-auth.ts

@@ -5,18 +5,20 @@ import type { WSEvents } from 'hono/ws'
  * signal "auth failed, do not reconnect with this token".
  *
  * Use when:
- * - Rejecting a WebSocket connection because the bearer token in the
- *   `?token=` query parameter is missing, invalid, expired, or revoked.
+ * - Rejecting a WebSocket connection because its bearer token is missing,
+ *   invalid, expired, or revoked. Legacy chat clients send it as `?token=`;
+ *   version-two chat clients send it after the upgrade.
  *
  * Expects:
  * - Server: accept the upgrade first, then close with this code inside
- *   `onOpen`. Throwing inside `upgradeWebSocket` produces an HTTP 401
+ *   `onOpen` or an authenticated protocol handler. Throwing inside
+ *   `upgradeWebSocket` produces an HTTP 401
  *   that browsers swallow before the connection enters the WS state
  *   machine, leaving clients with only `code=1006` — indistinguishable
  *   from a transient network drop.
  * - Client: in `onDisconnected`, treat `ev.code === 4001` as a terminal
  *   auth failure for the current token and stop the autoReconnect loop
- *   until a fresh token rotates the URL.
+ *   until a fresh token rotates the connection.
  *
  * NOTICE:
  * - This constant is duplicated in
@@ -27,6 +29,16 @@ import type { WSEvents } from 'hono/ws'
  */
 export const WS_CLOSE_UNAUTHORIZED = 4001
 
+/**
+ * Standard WebSocket close code for a temporary condition. Clients retry it.
+ */
+export const WS_CLOSE_TRY_AGAIN_LATER = 1013
+
+/**
+ * Standard WebSocket close code for a server-side error. Clients retry it.
+ */
+export const WS_CLOSE_INTERNAL_ERROR = 1011
+
 /**
  * Build a `WSEvents` shape that immediately closes the socket with
  * `WS_CLOSE_UNAUTHORIZED` after the upgrade completes.

server/apps/api/src/routes/chat-ws/v1/rpc.contract.test.ts

@@ -1,4 +1,5 @@
 import { parsePullMessagesRequest, parseSendMessagesRequest } from '@proj-airi/server-sdk-shared'
+import { parseAuthenticateRequest, parseAuthenticateResponse } from '@proj-airi/server-sdk-shared/v2'
 import { describe, expect, it } from 'vitest'
 
 describe('v1 chat WebSocket request contracts', () => {
@@ -20,4 +21,14 @@ describe('v1 chat WebSocket request contracts', () => {
     expect(() => parsePullMessagesRequest({ chatId: 'chat-1', afterSeq: -1 }))
       .toThrow()
   })
+
+  // https://github.com/moeru-ai/airi/pull/2309#discussion_r3796726614
+  it('rejects malformed authenticate requests from the shared contract', () => {
+    expect(() => parseAuthenticateRequest({ token: '' })).toThrow()
+    expect(() => parseAuthenticateRequest({ token: 'a'.repeat(4097) })).toThrow()
+  })
+
+  it('rejects malformed authenticate responses from the shared contract', () => {
+    expect(() => parseAuthenticateResponse({ userId: '' })).toThrow()
+  })
 })

server/apps/api/src/routes/chat-ws/v2/auth.test.ts

@@ -0,0 +1,101 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { WS_CLOSE_INTERNAL_ERROR, WS_CLOSE_TRY_AGAIN_LATER, WS_CLOSE_UNAUTHORIZED } from '../../../libs/ws-auth'
+import { createChatWsV2Authentication } from './auth'
+
+interface Deferred<T> {
+  promise: Promise<T>
+  reject: (error: Error) => void
+  resolve: (value: T) => void
+}
+
+function createDeferred<T>(): Deferred<T> {
+  let reject: (error: Error) => void = () => {}
+  let resolve: (value: T) => void = () => {}
+  const promise = new Promise<T>((resolvePromise, rejectPromise) => {
+    resolve = resolvePromise
+    reject = rejectPromise
+  })
+  return { promise, reject, resolve }
+}
+
+function createAuthentication(resolveUserId: (token: string) => Promise<string | null>) {
+  const close = vi.fn<(code?: number, reason?: string) => void>()
+  const onAuthenticated = vi.fn()
+  const authentication = createChatWsV2Authentication({
+    socket: { close },
+    resolveUserId,
+    onAuthenticated,
+  })
+  return { authentication, close, onAuthenticated }
+}
+
+describe('v2 chat WebSocket authentication', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  // https://github.com/moeru-ai/airi/pull/2309#discussion_r3796626514
+  // https://github.com/moeru-ai/airi/pull/2309#discussion_r3796626537
+  // ROOT CAUSE:
+  //
+  // An authentication resolver can finish after the socket closes or after a
+  // timeout. The old handler registered that disconnected context and allowed
+  // concurrent requests to repeat the resolver work.
+  //
+  // The authentication session now accepts one attempt, marks the socket
+  // inactive before it closes, and refuses a late resolver result.
+  it('does not authenticate after disconnecting during authentication', async () => {
+    const deferred = createDeferred<string | null>()
+    const resolveUserId = vi.fn(() => deferred.promise)
+    const { authentication, onAuthenticated } = createAuthentication(resolveUserId)
+
+    const request = authentication.authenticate({ token: 'valid-token' })
+    authentication.disconnect()
+    deferred.resolve('user-1')
+
+    await expect(request).rejects.toThrow('WebSocket closed during authentication')
+    expect(onAuthenticated).not.toHaveBeenCalled()
+  })
+
+  it('limits each socket to one authentication attempt', async () => {
+    const deferred = createDeferred<string | null>()
+    const resolveUserId = vi.fn(() => deferred.promise)
+    const { authentication } = createAuthentication(resolveUserId)
+
+    const first = authentication.authenticate({ token: 'valid-token' })
+    await expect(authentication.authenticate({ token: 'valid-token' })).rejects.toThrow('WebSocket authentication already attempted')
+    expect(resolveUserId).toHaveBeenCalledTimes(1)
+
+    deferred.resolve('user-1')
+    await expect(first).resolves.toEqual({ userId: 'user-1' })
+  })
+
+  it('uses a retryable code when authentication times out', async () => {
+    const deferred = createDeferred<string | null>()
+    const { authentication, close, onAuthenticated } = createAuthentication(() => deferred.promise)
+
+    const request = authentication.authenticate({ token: 'valid-token' })
+    await vi.advanceTimersByTimeAsync(15_000)
+    deferred.resolve('user-1')
+
+    await expect(request).rejects.toThrow('WebSocket closed during authentication')
+    expect(close).toHaveBeenCalledWith(WS_CLOSE_TRY_AGAIN_LATER, 'authentication timeout')
+    expect(close).not.toHaveBeenCalledWith(WS_CLOSE_UNAUTHORIZED, 'unauthorized')
+    expect(onAuthenticated).not.toHaveBeenCalled()
+  })
+
+  it('uses a retryable code when the resolver fails temporarily', async () => {
+    const resolveUserId = vi.fn(async () => {
+      throw new Error('database unavailable')
+    })
+    const { authentication, close } = createAuthentication(resolveUserId)
+
+    await expect(authentication.authenticate({ token: 'valid-token' })).rejects.toThrow('database unavailable')
+    expect(close).toHaveBeenCalledWith(WS_CLOSE_INTERNAL_ERROR, 'authentication unavailable')
+  })
+})

server/apps/api/src/routes/chat-ws/v2/auth.ts

@@ -0,0 +1,94 @@
+import type { WSContext } from 'hono/ws'
+
+import { parseAuthenticateRequest } from '@proj-airi/server-sdk-shared/v2'
+
+import { WS_CLOSE_INTERNAL_ERROR, WS_CLOSE_TRY_AGAIN_LATER, WS_CLOSE_UNAUTHORIZED } from '../../../libs/ws-auth'
+
+const CHAT_AUTH_TIMEOUT_MS = 15_000
+
+export interface ChatWsAuthResolver {
+  (token: string): Promise<string | null>
+}
+
+interface CreateChatWsV2AuthenticationOptions {
+  socket?: Pick<WSContext, 'close'>
+  resolveUserId: ChatWsAuthResolver
+  onAuthenticated: (userId: string) => void
+}
+
+export interface ChatWsV2Authentication {
+  /** Handles the only authentication request accepted for this socket. */
+  authenticate: (request: unknown) => Promise<{ userId: string }>
+  /** Stops authentication when the websocket disconnects. */
+  disconnect: () => void
+}
+
+/**
+ * Owns authentication lifetime for one version-two websocket connection.
+ *
+ * The session timer begins when the socket opens. A close or timeout marks the
+ * session inactive before the token resolver completes, preventing a late
+ * resolver result from registering a disconnected peer.
+ */
+export function createChatWsV2Authentication(options: CreateChatWsV2AuthenticationOptions): ChatWsV2Authentication {
+  let authenticationStarted = false
+  let connectionActive = true
+  let authenticationComplete = false
+  let authTimer: ReturnType<typeof setTimeout>
+
+  function stopAuthentication(code: number, reason: string): void {
+    connectionActive = false
+    clearTimeout(authTimer)
+    options.socket?.close(code, reason)
+  }
+
+  authTimer = setTimeout(() => {
+    if (!authenticationComplete)
+      stopAuthentication(WS_CLOSE_TRY_AGAIN_LATER, 'authentication timeout')
+  }, CHAT_AUTH_TIMEOUT_MS)
+
+  return {
+    async authenticate(request) {
+      if (authenticationStarted)
+        throw new Error('WebSocket authentication already attempted')
+
+      authenticationStarted = true
+
+      let parsedRequest
+      try {
+        parsedRequest = parseAuthenticateRequest(request)
+      }
+      catch {
+        stopAuthentication(WS_CLOSE_UNAUTHORIZED, 'unauthorized')
+        throw new Error('WebSocket authentication failed')
+      }
+
+      let userId: string | null
+      try {
+        userId = await options.resolveUserId(parsedRequest.token)
+      }
+      catch (error) {
+        if (connectionActive)
+          stopAuthentication(WS_CLOSE_INTERNAL_ERROR, 'authentication unavailable')
+        throw error
+      }
+
+      if (!connectionActive)
+        throw new Error('WebSocket closed during authentication')
+
+      if (!userId) {
+        stopAuthentication(WS_CLOSE_UNAUTHORIZED, 'unauthorized')
+        throw new Error('WebSocket authentication failed')
+      }
+
+      authenticationComplete = true
+      clearTimeout(authTimer)
+      options.onAuthenticated(userId)
+      return { userId }
+    },
+    disconnect() {
+      connectionActive = false
+      clearTimeout(authTimer)
+    },
+  }
+}

server/apps/api/src/routes/chat-ws/v2/index.ts

@@ -1,42 +1,141 @@
+import type { WSContext, WSEvents } from 'hono/ws'
 import type Redis from 'ioredis'
 
 import type { EngagementMetrics } from '../../../otel'
 import type { ChatService } from '../../../services/domain/chats'
 import type { ChatWsRuntime } from '../runtime'
+import type { ChatWsAuthResolver } from './auth'
 
-import { createPeerHooks } from '@moeru/eventa/adapters/websocket/hono'
+import { defineInvokeHandler } from '@moeru/eventa'
+import { createPeerHooks, wsDisconnectedEvent } from '@moeru/eventa/adapters/websocket/hono'
+import { authenticate } from '@proj-airi/server-sdk-shared/v2'
 
+import { WS_CLOSE_TRY_AGAIN_LATER, WS_CLOSE_UNAUTHORIZED } from '../../../libs/ws-auth'
 import { registerChatWsPeer } from '../peer'
 import { createChatWsRuntime } from '../runtime'
+import { createChatWsV2Authentication } from './auth'
+import { createChatWsUnauthenticatedPeerLimit } from './unauthenticated-peers'
+
+const MAX_UNAUTHENTICATED_CHAT_WS_CONNECTIONS = 100
+const MAX_UNAUTHENTICATED_CHAT_WS_FRAME_BYTES = 8192
+
+function isTerminableSocket(raw: unknown): raw is { terminate: () => void } {
+  return typeof raw === 'object'
+    && raw !== null
+    && 'terminate' in raw
+    && typeof raw.terminate === 'function'
+}
+
+function isAuthenticateFrame(data: unknown): boolean {
+  if (typeof data !== 'string' || data.length > MAX_UNAUTHENTICATED_CHAT_WS_FRAME_BYTES)
+    return false
+
+  try {
+    const frame = JSON.parse(data) as { eventa?: { id?: unknown }, payload?: { id?: unknown } }
+    return frame.eventa?.id === 'chat:authenticate-send'
+      || frame.payload?.id === 'chat:authenticate-send'
+  }
+  catch {
+    return false
+  }
+}
 
 /**
- * Creates websocket handlers for chat sync RPC and message fanout.
+ * Creates version-two WebSocket handlers for chat sync and message fanout.
  *
- * Use when:
- * - Mounting an already authenticated chat peer.
- *
- * Expects:
- * - `instanceId` is stable for this process so Redis echo suppression works.
- * - Redis Pub/Sub is used only for best-effort cross-instance notification.
- *
- * Returns:
- * - A per-user Hono websocket setup function.
+ * `/ws/v2/chat` accepts an anonymous WebSocket upgrade. The client must invoke
+ * `chat:authenticate` before it joins the shared authenticated peer runtime.
  */
 export function createChatWsV2Handlers(
   chatService: ChatService,
   redis: Redis,
   instanceId: string,
+  resolveUserId: ChatWsAuthResolver,
   metrics?: EngagementMetrics | null,
   runtime?: ChatWsRuntime,
+  restoreAuthenticatedPayloadLimit?: (socket: unknown) => void,
 ) {
   const chatRuntime = runtime ?? createChatWsRuntime(redis, instanceId, metrics)
+  // The v2 upgrade is anonymous by design. Keep a bounded number of peers in
+  // the authentication window so a burst cannot retain unbounded timers and
+  // Eventa contexts in this process.
+  const unauthenticatedPeers = createChatWsUnauthenticatedPeerLimit(MAX_UNAUTHENTICATED_CHAT_WS_CONNECTIONS)
+
+  return function setupPeer() {
+    let socket: WSContext | undefined
+    let ownsUnauthenticatedSlot = false
+    let authenticated = false
+
+    function releaseUnauthenticatedSlot(): void {
+      if (!ownsUnauthenticatedSlot)
+        return
+
+      ownsUnauthenticatedSlot = false
+      unauthenticatedPeers.release()
+    }
 
-  return function setupPeer(userId: string) {
     const { hooks } = createPeerHooks({
       onContext: (ctx) => {
-        registerChatWsPeer({ ctx, userId, chatService, runtime: chatRuntime, metrics })
+        const authentication = createChatWsV2Authentication({
+          socket,
+          resolveUserId,
+          onAuthenticated(userId) {
+            if (socket)
+              restoreAuthenticatedPayloadLimit?.(socket.raw)
+            authenticated = true
+            releaseUnauthenticatedSlot()
+            registerChatWsPeer({ ctx, userId, chatService, runtime: chatRuntime, metrics })
+          },
+        })
+        const unregisterAuthenticate = defineInvokeHandler(ctx, authenticate, authentication.authenticate)
+
+        ctx.on(wsDisconnectedEvent, () => {
+          authentication.disconnect()
+          unregisterAuthenticate()
+        })
       },
     })
-    return hooks
+
+    const originalOnOpen = hooks.onOpen
+    const originalOnClose = hooks.onClose
+    const originalOnError = hooks.onError
+    const originalOnMessage = hooks.onMessage
+    const v2Hooks: WSEvents = {
+      ...hooks,
+      onOpen(event, ws) {
+        if (!unauthenticatedPeers.tryAcquire()) {
+          // Do not wait for a hostile peer to answer a close frame. The slot is
+          // full, so terminating releases this connection immediately.
+          if (isTerminableSocket(ws.raw))
+            ws.raw.terminate()
+          else
+            ws.close(WS_CLOSE_TRY_AGAIN_LATER, 'too many unauthenticated connections')
+          return
+        }
+
+        ownsUnauthenticatedSlot = true
+        socket = ws
+        originalOnOpen?.(event, ws)
+      },
+      onClose(event, ws) {
+        releaseUnauthenticatedSlot()
+        originalOnClose?.(event, ws)
+      },
+      onError(event, ws) {
+        releaseUnauthenticatedSlot()
+        originalOnError?.(event, ws)
+      },
+      onMessage(event, ws) {
+        // Before authentication, accept one bounded Eventa invoke only. This
+        // prevents arbitrary frames from reaching the adapter parser and log.
+        if (!authenticated && !isAuthenticateFrame(event.data)) {
+          ws.close(WS_CLOSE_UNAUTHORIZED, 'unauthorized')
+          return
+        }
+
+        originalOnMessage?.(event, ws)
+      },
+    }
+    return v2Hooks
   }
 }

server/apps/api/src/routes/chat-ws/v2/payload-limit.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it } from 'vitest'
+
+import { createChatWsPayloadLimit } from './payload-limit'
+
+describe('v2 chat WebSocket payload limit', () => {
+  // https://github.com/moeru-ai/airi/pull/2309#discussion_r3818708552
+  // ROOT CAUSE:
+  //
+  // The pre-authentication frame limit stayed on a socket after successful
+  // authentication. Normal chat batches larger than that limit then failed.
+  //
+  // The limiter records each receiver's original limit and restores it when
+  // the authentication handler completes.
+  it('restores the original frame limit after authentication', () => {
+    const socket = { _receiver: { _maxPayload: 100 * 1024 * 1024 } }
+    const payloadLimit = createChatWsPayloadLimit(8192)
+
+    payloadLimit.restrict(socket)
+    expect(socket._receiver._maxPayload).toBe(8192)
+
+    payloadLimit.restore(socket)
+    expect(socket._receiver._maxPayload).toBe(100 * 1024 * 1024)
+  })
+
+  it('does not change sockets without a ws receiver', () => {
+    const socket = {}
+    const payloadLimit = createChatWsPayloadLimit(8192)
+
+    payloadLimit.restrict(socket)
+    payloadLimit.restore(socket)
+
+    expect(socket).toEqual({})
+  })
+})

server/apps/api/src/routes/chat-ws/v2/payload-limit.ts

@@ -0,0 +1,52 @@
+interface PayloadReceiverSocket {
+  _receiver: { _maxPayload: number }
+}
+
+function hasReceiverPayloadLimit(socket: unknown): socket is PayloadReceiverSocket {
+  return typeof socket === 'object'
+    && socket !== null
+    && '_receiver' in socket
+    && typeof socket._receiver === 'object'
+    && socket._receiver !== null
+    && '_maxPayload' in socket._receiver
+    && typeof socket._receiver._maxPayload === 'number'
+}
+
+export interface ChatWsPayloadLimit {
+  /** Limits frames while the peer has not authenticated. */
+  restrict: (socket: unknown) => void
+  /** Restores the transport limit after the peer authenticates. */
+  restore: (socket: unknown) => void
+}
+
+/**
+ * Owns the temporary frame limit for post-connect WebSocket authentication.
+ *
+ * The `ws` receiver is created before Hono runs the route hooks. Record its
+ * configured limit per socket so authenticated chat traffic keeps the normal
+ * capacity after the small authentication frame has been accepted.
+ */
+export function createChatWsPayloadLimit(unauthenticatedMaximum: number): ChatWsPayloadLimit {
+  const originalLimits = new WeakMap<object, number>()
+
+  return {
+    restrict(socket) {
+      if (!hasReceiverPayloadLimit(socket))
+        return
+
+      originalLimits.set(socket, socket._receiver._maxPayload)
+      socket._receiver._maxPayload = unauthenticatedMaximum
+    },
+    restore(socket) {
+      if (!hasReceiverPayloadLimit(socket))
+        return
+
+      const originalLimit = originalLimits.get(socket)
+      if (originalLimit === undefined)
+        return
+
+      socket._receiver._maxPayload = originalLimit
+      originalLimits.delete(socket)
+    },
+  }
+}

server/apps/api/src/routes/chat-ws/v2/unauthenticated-peers.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, it } from 'vitest'
+
+import { createChatWsUnauthenticatedPeerLimit } from './unauthenticated-peers'
+
+describe('v2 chat unauthenticated peer limit', () => {
+  // https://github.com/moeru-ai/airi/pull/2309#discussion_r3811138375
+  // ROOT CAUSE:
+  //
+  // v2 upgrades before the client proves identity. Without a cap, unauthenticated
+  // sockets each retained an Eventa context and an authentication timer for up
+  // to fifteen seconds.
+  //
+  // The per-process limiter rejects peers once full and returns capacity only
+  // when a reserved slot is released.
+  it('does not admit more unauthenticated peers than its configured capacity', () => {
+    const limit = createChatWsUnauthenticatedPeerLimit(2)
+
+    expect(limit.tryAcquire()).toBe(true)
+    expect(limit.tryAcquire()).toBe(true)
+    expect(limit.tryAcquire()).toBe(false)
+
+    limit.release()
+
+    expect(limit.tryAcquire()).toBe(true)
+  })
+
+  it('does not create capacity when a peer is released more than once', () => {
+    const limit = createChatWsUnauthenticatedPeerLimit(1)
+
+    expect(limit.tryAcquire()).toBe(true)
+    limit.release()
+    limit.release()
+
+    expect(limit.tryAcquire()).toBe(true)
+    expect(limit.tryAcquire()).toBe(false)
+  })
+})

server/apps/api/src/routes/chat-ws/v2/unauthenticated-peers.ts

@@ -0,0 +1,31 @@
+export interface ChatWsUnauthenticatedPeerLimit {
+  /** Reserves an unauthenticated connection slot when capacity remains. */
+  tryAcquire: () => boolean
+  /** Releases a previously reserved unauthenticated connection slot. */
+  release: () => void
+}
+
+/**
+ * Limits concurrent peers that have upgraded but have not authenticated yet.
+ *
+ * One instance owns only its local sockets, so the count protects the memory,
+ * timer, and Eventa-context capacity of that process. Releasing more than once
+ * is safe because close and error events can both arrive for one peer.
+ */
+export function createChatWsUnauthenticatedPeerLimit(maximumConnections: number): ChatWsUnauthenticatedPeerLimit {
+  let activeConnections = 0
+
+  return {
+    tryAcquire() {
+      if (activeConnections >= maximumConnections)
+        return false
+
+      activeConnections += 1
+      return true
+    },
+    release() {
+      if (activeConnections > 0)
+        activeConnections -= 1
+    },
+  }
+}

server/packages/server-sdk-shared/README.md

@@ -19,7 +19,7 @@ The package uses Eventa `1.0.0-beta.15`. Its WebSocket adapter accepts beta.13
 `id/type/payload` envelopes and sends these fields with current envelopes.
 
 `/ws/chat` keeps query-token authentication for deployed clients. `/ws/v2/chat`
-authenticates after the WebSocket opens.
+authenticates after the WebSocket opens with `chat:authenticate`.
 
 ## License
 

server/packages/server-sdk-shared/src/v2.ts

@@ -2,6 +2,8 @@ import type { NewMessagesPayload, PullMessagesRequest, PullMessagesResponse, Sen
 
 import { defineInvokeEventa, defineOutboundEventa } from '@moeru/eventa'
 
+import * as v from 'valibot'
+
 export type {
   MessageRole,
   NewMessagesPayload,
@@ -18,6 +20,29 @@ export {
   SendMessagesRequestSchema,
 } from './chat'
 
+export const AuthenticateRequestSchema = v.object({
+  token: v.pipe(v.string(), v.minLength(1), v.maxLength(4096)),
+})
+
+export type AuthenticateRequest = v.InferOutput<typeof AuthenticateRequestSchema>
+
+export const AuthenticateResponseSchema = v.object({
+  userId: v.pipe(v.string(), v.minLength(1)),
+})
+
+export type AuthenticateResponse = v.InferOutput<typeof AuthenticateResponseSchema>
+
+/** Parses a `chat:authenticate` payload at the WebSocket boundary. */
+export function parseAuthenticateRequest(request: unknown): AuthenticateRequest {
+  return v.parse(AuthenticateRequestSchema, request)
+}
+
+/** Parses a `chat:authenticate` response at the WebSocket boundary. */
+export function parseAuthenticateResponse(response: unknown): AuthenticateResponse {
+  return v.parse(AuthenticateResponseSchema, response)
+}
+
+export const authenticate = defineInvokeEventa<AuthenticateResponse, AuthenticateRequest>('chat:authenticate')
 export const sendMessages = defineInvokeEventa<SendMessagesResponse, SendMessagesRequest>('chat:send-messages')
 export const pullMessages = defineInvokeEventa<PullMessagesResponse, PullMessagesRequest>('chat:pull-messages')
 export const newMessages = defineOutboundEventa<NewMessagesPayload>('chat:new-messages')