From 84c053b44ff5196c7731ca010beefb101f4eb3d4 Mon Sep 17 00:00:00 2001
From: Neko Ayaka <neko@ayaka.moe>
Date: Wed, 17 Dec 2025 19:57:18 +0800
Subject: [PATCH] fix(server-sdk,airi-plugin-vscode): improve connect, ensure
 connected before sending

---
 packages/server-sdk/src/client.ts      | 74 ++++++++++++++++----------
 plugins/airi-plugin-vscode/src/airi.ts |  4 +-
 2 files changed, 50 insertions(+), 28 deletions(-)

diff --git a/packages/server-sdk/src/client.ts b/packages/server-sdk/src/client.ts
index 26ce841c1..068b348b2 100644
--- a/packages/server-sdk/src/client.ts
+++ b/packages/server-sdk/src/client.ts
@@ -4,13 +4,6 @@ import WebSocket from 'crossws/websocket'
 
 import { sleep } from '@moeru/std'
 
-class ReconnectingError extends Error {
-  constructor(message: string) {
-    super(message)
-    this.name = 'ReconnectingError'
-  }
-}
-
 export interface ClientOptions<C = undefined> {
   url?: string
   name: string
@@ -28,6 +21,8 @@ export class Client<C = undefined> {
   private connecting = false
   private websocket?: WebSocket
   private shouldClose = false
+  private connectAttempt?: Promise<void>
+  private connectTask?: Promise<void>
 
   private readonly opts: Required<Omit<ClientOptions<C>, 'token'>> & Pick<ClientOptions<C>, 'token'>
   private readonly eventListeners = new Map<
@@ -78,10 +73,6 @@ export class Client<C = undefined> {
         return
       }
       catch (err) {
-        if (err instanceof ReconnectingError) {
-          return
-        }
-
         this.opts.onError?.(err)
         const delay = Math.min(2 ** attempts * 1000, 30_000) // capped exponential backoff
         await sleep(delay)
@@ -92,7 +83,7 @@ export class Client<C = undefined> {
 
   private async tryReconnectWithExponentialBackoff() {
     if (this.shouldClose) {
-      return
+      throw new Error('Client is closed')
     }
 
     await this.retryWithExponentialBackoff(() => this._connect())
@@ -103,25 +94,42 @@ export class Client<C = undefined> {
       return Promise.resolve()
     }
     if (this.connecting) {
-      return Promise.reject(new ReconnectingError('Already connecting'))
+      return this.connectAttempt ?? Promise.resolve()
     }
 
-    return new Promise((resolve, reject) => {
+    this.connectAttempt = new Promise((resolve, reject) => {
       this.connecting = true
+      let settled = false
+
+      const settle = (fn: () => void) => {
+        if (settled)
+          return
+
+        settled = true
+        this.connecting = false
+        this.connectAttempt = undefined
+        fn()
+      }
 
       const ws = new WebSocket(this.opts.url)
       this.websocket = ws
 
       ws.onmessage = this.handleMessageBound
       ws.onerror = (event: any) => {
-        this.connecting = false
-        this.connected = false
+        settle(() => {
+          this.connected = false
 
-        this.opts.onError?.(event)
-        reject(event?.error ?? new Error('WebSocket error'))
+          this.opts.onError?.(event)
+          reject(event?.error ?? new Error('WebSocket error'))
+        })
       }
       ws.onclose = () => {
-        this.connecting = false
+        if (!settled && !this.connected) {
+          settle(() => {
+            reject(new Error('WebSocket closed before open'))
+          })
+          return
+        }
 
         if (this.connected) {
           this.connected = false
@@ -132,21 +140,33 @@ export class Client<C = undefined> {
         }
       }
       ws.onopen = () => {
-        this.connecting = false
-        this.connected = true
+        settle(() => {
+          this.connected = true
 
-        if (this.opts.token)
-          this.tryAuthenticate()
-        else
-          this.tryAnnounce()
+          if (this.opts.token)
+            this.tryAuthenticate()
+          else
+            this.tryAnnounce()
 
-        resolve()
+          resolve()
+        })
       }
     })
+
+    return this.connectAttempt
   }
 
   async connect() {
-    await this.tryReconnectWithExponentialBackoff()
+    if (this.connected) {
+      return
+    }
+    if (this.connectTask) {
+      return this.connectTask
+    }
+
+    this.connectTask = this.tryReconnectWithExponentialBackoff().finally(() => (this.connectTask = undefined))
+
+    return this.connectTask
   }
 
   private tryAnnounce() {
diff --git a/plugins/airi-plugin-vscode/src/airi.ts b/plugins/airi-plugin-vscode/src/airi.ts
index 2c7ebb46a..95402dfe0 100644
--- a/plugins/airi-plugin-vscode/src/airi.ts
+++ b/plugins/airi-plugin-vscode/src/airi.ts
@@ -9,6 +9,7 @@ export class Client {
   async connect(): Promise<boolean> {
     try {
       this.client = new ServerClient<Events>({ name: 'proj-airi:plugin-vscode' })
+      await this.client.connect()
       useLogger().log('AIRI connected to Server Channel')
       return true
     }
@@ -26,13 +27,14 @@ export class Client {
     }
   }
 
-  sendEvent(event: Events): void {
+  async sendEvent(event: Events): Promise<void> {
     if (!this.client) {
       useLogger().warn('Cannot send event: not connected to AIRI Server Channel')
       return
     }
 
     try {
+      await this.client.connect()
       this.client.send({ type: 'vscode:context', data: event })
 
       useLogger().log(`Sent event to AIRI: ${event.type}`, event)