fix(ws): mark and reroute on admission-time usage limit

packages/opencode/src/core/quota-manager.ts

@@ -53,23 +53,27 @@ export function quotaWindowResetIsPast(
 }
 
 /**
- * Resolve when a mid-stream-exhausted account's window actually resets, from
- * its own last-known cached quota snapshot (a response.failed frame carries no
- * reset time itself). The frame always names the exhausted window (primary or
- * secondary), so use only THAT window's cached reset; if it's unknown, fall
- * back to a conservative bounded default rather than borrowing the OTHER
- * window's reset — an unrelated window can reset far in the future and would
- * blackhole the account long past its actual rate-limit. The default is
- * self-correcting: if the account is still exhausted, the next response.failed
- * frame re-marks it. Pure — single source of truth shared by the fetch-override
- * reroute and its unit tests.
+ * Resolve when a WS-exhausted account becomes usable. Admission errors can
+ * carry the provider's authoritative reset; response.failed frames do not, so
+ * those fall back to the named window in the last-known quota snapshot. Never
+ * borrow the other window's reset — an unrelated window can blackhole the
+ * account long past its actual limit. The bounded default self-corrects when a
+ * later request reports exhaustion again.
  */
 export function resolveMidStreamRateLimitResetAt(
   quota: OAuthQuotaSnapshot | undefined,
   window: string,
   now: number,
   defaultMs: number,
+  explicitResetAt?: number,
 ): number {
+  if (
+    explicitResetAt !== undefined &&
+    Number.isFinite(explicitResetAt) &&
+    explicitResetAt > now
+  ) {
+    return explicitResetAt
+  }
   const named =
     window === PRIMARY
       ? quota?.primary

packages/opencode/src/index.ts

@@ -1277,13 +1277,12 @@ export async function CodexAuthPlugin(
           await writeMachineSidebarState(quotaManager, latestStorage)
         }
 
-        // Resolve when a mid-stream-exhausted account's window actually resets,
-        // from its own last-known cached quota. The resolution logic itself
-        // lives in quota-manager.ts (single source of truth, unit-tested
-        // there) — this wrapper only supplies the account's cached snapshot.
+        // The pure resolver prefers an admission error's explicit reset and
+        // otherwise uses this account's cached named-window quota.
         function midStreamRateLimitResetAt(
           accountKey: string,
           window: string,
+          explicitResetAt?: number,
         ): number {
           const quota =
             accountKey === 'main'
@@ -1294,6 +1293,7 @@ export async function CodexAuthPlugin(
             window,
             Date.now(),
             DEFAULT_MID_STREAM_RATE_LIMIT_RESET_MS,
+            explicitResetAt,
           )
         }
 
@@ -1314,13 +1314,10 @@ export async function CodexAuthPlugin(
                   true,
                 )
               },
-              // Mid-stream quota exhaustion (response.failed carrying
-              // rate_limit_reached_type) — the frame itself is the authority,
-              // so this marks the account rate-limited without any quota-API
-              // call. The fetch override then reroutes away from it: a
-              // rate-limited main is treated like a killswitch block, and a
-              // rate-limited fallback is dropped from candidate selection.
-              onRateLimitReached: (window, accountId) => {
+              // WS quota exhaustion is authoritative without a quota-API call.
+              // Admission errors supply their own reset; response.failed uses
+              // the cached named-window reset or the bounded default.
+              onRateLimitReached: (window, accountId, explicitResetAt) => {
                 if (!accountId) {
                   // The loader always sets the internal quota-account header,
                   // so this should never fire today — but a future call site
@@ -1332,7 +1329,11 @@ export async function CodexAuthPlugin(
                   )
                 }
                 const accountKey = accountId ?? 'main'
-                const resetAt = midStreamRateLimitResetAt(accountKey, window)
+                const resetAt = midStreamRateLimitResetAt(
+                  accountKey,
+                  window,
+                  explicitResetAt,
+                )
                 quotaManager.markRateLimited(accountKey, resetAt)
                 logQ.debug('mid-stream rate limit mark', {
                   pid: process.pid,

packages/opencode/src/raw-ws-bun.ts

@@ -15,6 +15,7 @@
 // send(string), close(), and a `url` field. Text frames only on receive.
 
 import { dumpDiagnostic } from './dump'
+import { rejectedUpgradeEvent, rejectedUpgradeStatus } from './raw-ws-upgrade'
 
 const WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
 
@@ -65,6 +66,7 @@ export class RawWebSocket {
   }
   private rxBuffer: Uint8Array<ArrayBufferLike> = new Uint8Array(0)
   private handshakeDone = false
+  private upgradeFailureEmitted = false
   private expectedAccept = ''
   // message reassembly across fragmented frames
   private fragOpcode = 0
@@ -166,12 +168,14 @@ export class RawWebSocket {
           data: (_s: unknown, d: Uint8Array) => this.onData(d),
           drain: () => this.flushWrites(),
           error: (_s: unknown, e: unknown) => {
+            if (this.emitRejectedUpgrade(true)) return
             this.readyState = 3
             const message = e instanceof Error ? e.message : String(e)
             void this.log('socket_error', { message })
             this.emit('error', { message })
           },
           close: () => {
+            if (this.emitRejectedUpgrade(true)) return
             this.readyState = 3
             void this.log('socket_close')
             this.emit('close', { code: 1006, reason: 'socket closed' })
@@ -195,15 +199,25 @@ export class RawWebSocket {
       const headerText = text.slice(0, idx)
       const statusLine = headerText.split('\r\n')[0] ?? ''
       const acceptMatch = headerText.match(/sec-websocket-accept:\s*(\S+)/i)
-      if (!statusLine.includes(' 101')) {
-        this.readyState = 3
+      const status = rejectedUpgradeStatus(statusLine)
+      if (status !== 101) {
+        if (status === undefined) {
+          this.readyState = 3
+          try {
+            this.socket?.end()
+          } catch {
+            /* ignore */
+          }
+          void this.log('upgrade_failed', { statusLine })
+          this.emit('error', { message: `WS upgrade failed: ${statusLine}` })
+          return
+        }
+        if (!this.emitRejectedUpgrade()) return
         try {
           this.socket?.end()
         } catch {
           /* ignore */
         }
-        void this.log('upgrade_failed', { statusLine })
-        this.emit('error', { message: `WS upgrade failed: ${statusLine}` })
         return
       }
       if (!acceptMatch || acceptMatch[1] !== this.expectedAccept) {
@@ -228,6 +242,27 @@ export class RawWebSocket {
     this.drainFrames()
   }
 
+  private emitRejectedUpgrade(finalizePartial = false): boolean {
+    if (this.upgradeFailureEmitted) return true
+    if (this.handshakeDone) return false
+    const headerEnd = Buffer.from(this.rxBuffer).indexOf('\r\n\r\n')
+    if (headerEnd === -1) return false
+    const errorEvent = rejectedUpgradeEvent(
+      this.rxBuffer,
+      headerEnd,
+      finalizePartial,
+    )
+    if (!errorEvent) return false
+    this.upgradeFailureEmitted = true
+    this.readyState = 3
+    void this.log('upgrade_failed', {
+      status: errorEvent.status,
+      partial: finalizePartial,
+    })
+    this.emit('error', errorEvent)
+    return true
+  }
+
   private drainFrames(): void {
     // RFC 6455 frame parsing for server->client (never masked).
     for (;;) {

packages/opencode/src/raw-ws-node.ts

@@ -18,6 +18,7 @@ import { connect as netConnect } from 'node:net'
 import { connect as tlsConnect } from 'node:tls'
 
 import { dumpDiagnostic } from './dump'
+import { rejectedUpgradeEvent, rejectedUpgradeStatus } from './raw-ws-upgrade'
 
 const WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
 
@@ -61,6 +62,7 @@ export class RawWebSocket {
   }
   private rxBuffer: Uint8Array<ArrayBufferLike> = new Uint8Array(0)
   private handshakeDone = false
+  private upgradeFailureEmitted = false
   private expectedAccept = ''
   // message reassembly across fragmented frames
   private fragOpcode = 0
@@ -159,12 +161,14 @@ export class RawWebSocket {
         )
       })
       socket.on('error', (e) => {
+        if (this.emitRejectedUpgrade(true)) return
         this.readyState = 3
         const message = e instanceof Error ? e.message : String(e)
         void this.log('socket_error', { message })
         this.emit('error', { message })
       })
       socket.on('close', (hadError) => {
+        if (this.emitRejectedUpgrade(true)) return
         this.readyState = 3
         void this.log('socket_close', { hadError })
         this.emit('close', { code: 1006, reason: 'socket closed' })
@@ -189,11 +193,17 @@ export class RawWebSocket {
       const headerText = text.slice(0, idx)
       const statusLine = headerText.split('\r\n')[0] ?? ''
       const acceptMatch = headerText.match(/sec-websocket-accept:\s*(\S+)/i)
-      if (!statusLine.includes(' 101')) {
-        this.readyState = 3
+      const status = rejectedUpgradeStatus(statusLine)
+      if (status !== 101) {
+        if (status === undefined) {
+          this.readyState = 3
+          this.socket?.destroy?.()
+          void this.log('upgrade_failed', { statusLine })
+          this.emit('error', { message: `WS upgrade failed: ${statusLine}` })
+          return
+        }
+        if (!this.emitRejectedUpgrade()) return
         this.socket?.destroy?.()
-        void this.log('upgrade_failed', { statusLine })
-        this.emit('error', { message: `WS upgrade failed: ${statusLine}` })
         return
       }
       if (!acceptMatch || acceptMatch[1] !== this.expectedAccept) {
@@ -214,6 +224,27 @@ export class RawWebSocket {
     this.drainFrames()
   }
 
+  private emitRejectedUpgrade(finalizePartial = false): boolean {
+    if (this.upgradeFailureEmitted) return true
+    if (this.handshakeDone) return false
+    const headerEnd = Buffer.from(this.rxBuffer).indexOf('\r\n\r\n')
+    if (headerEnd === -1) return false
+    const errorEvent = rejectedUpgradeEvent(
+      this.rxBuffer,
+      headerEnd,
+      finalizePartial,
+    )
+    if (!errorEvent) return false
+    this.upgradeFailureEmitted = true
+    this.readyState = 3
+    void this.log('upgrade_failed', {
+      status: errorEvent.status,
+      partial: finalizePartial,
+    })
+    this.emit('error', errorEvent)
+    return true
+  }
+
   private drainFrames(): void {
     // RFC 6455 frame parsing for server->client (never masked).
     for (;;) {

packages/opencode/src/raw-ws-upgrade.ts

@@ -0,0 +1,60 @@
+import { isRecord } from './util/record'
+
+export function rejectedUpgradeStatus(statusLine: string): number | undefined {
+  const status = Number(statusLine.match(/^HTTP\/\d(?:\.\d)?\s+(\d{3})/)?.[1])
+  return Number.isFinite(status) ? status : undefined
+}
+
+export function rejectedUpgradeEvent(
+  buffer: Uint8Array,
+  headerEnd: number,
+  finalizePartial = false,
+): Record<string, unknown> | undefined {
+  const headerText = Buffer.from(buffer.slice(0, headerEnd)).toString('latin1')
+  const lines = headerText.split('\r\n')
+  const statusLine = lines[0] ?? ''
+  const status = rejectedUpgradeStatus(statusLine)
+  if (status === undefined || status === 101) return undefined
+
+  const headers = Object.fromEntries(
+    lines.slice(1).flatMap((line) => {
+      const separator = line.indexOf(':')
+      return separator === -1
+        ? []
+        : [
+            [
+              line.slice(0, separator).trim().toLowerCase(),
+              line.slice(separator + 1).trim(),
+            ],
+          ]
+    }),
+  )
+  const bodyBytes = buffer.slice(headerEnd + 4)
+  const contentLengthHeader = headers['content-length']
+  const contentLength = Number(contentLengthHeader)
+  if (
+    !finalizePartial &&
+    (contentLengthHeader === undefined ||
+      (Number.isFinite(contentLength) && bodyBytes.length < contentLength))
+  ) {
+    return undefined
+  }
+
+  const body = Buffer.from(bodyBytes).toString('utf8')
+  const parsed = (() => {
+    try {
+      const value = JSON.parse(body)
+      return isRecord(value) ? value : undefined
+    } catch {
+      return undefined
+    }
+  })()
+  return {
+    message: `WS upgrade failed: ${statusLine}`,
+    status,
+    status_code: status,
+    headers,
+    body,
+    ...(isRecord(parsed?.error) ? { error: parsed.error } : {}),
+  }
+}