From b43a52727529832c50f2149e557e4df97fb4cf7e Mon Sep 17 00:00:00 2001 From: LemonNeko Date: Fri, 3 Apr 2026 16:11:45 +0800 Subject: [PATCH] refactor(server-sdk): use a external injected websocket --- packages/server-sdk/src/client.ts | 37 ++++++++++-------- packages/server-sdk/src/index.ts | 1 + packages/server-sdk/src/websocket-like.ts | 30 +++++++++++++++ packages/server-sdk/test/client.test.ts | 46 +++++++++++++++++++++++ 4 files changed, 99 insertions(+), 15 deletions(-) create mode 100644 packages/server-sdk/src/websocket-like.ts diff --git a/packages/server-sdk/src/client.ts b/packages/server-sdk/src/client.ts index c57e866ac..ec5a96244 100644 --- a/packages/server-sdk/src/client.ts +++ b/packages/server-sdk/src/client.ts @@ -8,7 +8,9 @@ import type { WebSocketEvents, } from '@proj-airi/server-shared/types' -import WebSocket from 'crossws/websocket' +import type { WebSocketLike, WebSocketLikeConstructor, WebSocketMessageEventLike } from './websocket-like' + +import NativeWebSocket from 'crossws/websocket' import superjson from 'superjson' import { errorMessageFrom, sleep } from '@moeru/std' @@ -46,6 +48,7 @@ export interface ClientOptions { url?: string name: string token?: string + websocketConstructor?: WebSocketLikeConstructor possibleEvents?: Array> identity?: MetadataEventSource @@ -72,7 +75,7 @@ interface ConnectionAttempt { promise: Promise reject: (error: Error) => void resolve: () => void - socket: WebSocket + socket: WebSocketLike } function createInstanceId() { @@ -107,7 +110,7 @@ function normalizeHeartbeatOptions(heartbeat?: ClientHeartbeatOptions): Required } export class Client { - private websocket?: WebSocket + private websocket?: WebSocketLike private shouldClose = false private connectTask?: Promise private heartbeatTimer?: ReturnType @@ -120,8 +123,9 @@ export class Client { private status: ClientStatus = 'idle' private readonly identity: MetadataEventSource private readonly heartbeat: Required + private readonly websocketConstructor: WebSocketLikeConstructor - private readonly opts: Required, 'token' | 'heartbeat'>> & Pick, 'token'> & { + private readonly opts: Required, 'token' | 'heartbeat' | 'websocketConstructor'>> & Pick, 'token'> & { heartbeat: Required } @@ -133,6 +137,7 @@ export class Client { private readonly stateListeners = new Set<(context: ClientStateChangeContext) => void>() constructor(options: ClientOptions) { + const { websocketConstructor, ...clientOptions } = options const identity = options.identity ?? { kind: 'plugin', plugin: { id: options.name }, @@ -155,13 +160,14 @@ export class Client { autoConnect: true, autoReconnect: true, maxReconnectAttempts: -1, - ...options, + ...clientOptions, heartbeat, identity, } this.identity = identity this.heartbeat = heartbeat + this.websocketConstructor = websocketConstructor ?? (NativeWebSocket as unknown as WebSocketLikeConstructor) if (this.opts.autoConnect) { void this.connect() @@ -177,7 +183,7 @@ export class Client { } get isSocketOpen() { - return this.websocket?.readyState === WebSocket.OPEN + return this.websocket?.readyState === this.websocketConstructor.OPEN } get lastError() { @@ -294,7 +300,7 @@ export class Client { const websocket = this.websocket this.websocket = undefined - if (websocket && websocket.readyState !== WebSocket.CLOSED && websocket.readyState !== WebSocket.CLOSING) { + if (websocket && websocket.readyState !== this.websocketConstructor.CLOSED && websocket.readyState !== this.websocketConstructor.CLOSING) { websocket.close() } @@ -347,7 +353,8 @@ export class Client { } private connectOnce(): Promise { - const ws = new WebSocket(this.opts.url) + const WebSocketConstructor = this.websocketConstructor + const ws = new WebSocketConstructor(this.opts.url) this.websocket = ws this.lastReadAt = Date.now() this.lastPingAt = 0 @@ -366,7 +373,7 @@ export class Client { const isCurrentSocket = () => this.websocket === ws - ws.onmessage = (event: MessageEvent) => { + ws.onmessage = (event: WebSocketMessageEventLike) => { if (!isCurrentSocket()) { return } @@ -433,7 +440,7 @@ export class Client { return attempt.promise } - private handleSocketFailure(error: Error, socket?: WebSocket) { + private handleSocketFailure(error: Error, socket?: WebSocketLike) { if (socket && this.websocket !== socket) { return } @@ -441,14 +448,14 @@ export class Client { const currentSocket = socket ?? this.websocket this.cleanupSocket(socket) - if (currentSocket && currentSocket.readyState !== WebSocket.CLOSED && currentSocket.readyState !== WebSocket.CLOSING) { + if (currentSocket && currentSocket.readyState !== this.websocketConstructor.CLOSED && currentSocket.readyState !== this.websocketConstructor.CLOSING) { currentSocket.close() } this.rejectAttempt(error) } - private cleanupSocket(socket?: WebSocket) { + private cleanupSocket(socket?: WebSocketLike) { if (socket && this.websocket !== socket) { return } @@ -576,7 +583,7 @@ export class Client { }) } - private async handleMessage(event: MessageEvent) { + private async handleMessage(event: WebSocketMessageEventLike) { this.lastReadAt = Date.now() try { @@ -755,7 +762,7 @@ export class Client { } private sendNativeHeartbeat(kind: 'ping' | 'pong') { - const websocket = this.websocket as WebSocket & { + const websocket = this.websocket as WebSocketLike & { ping?: () => void pong?: () => void } @@ -809,7 +816,7 @@ export class Client { this.cleanupSocket(websocket) this.rejectAttempt(error) - if (websocket && websocket.readyState !== WebSocket.CLOSED && websocket.readyState !== WebSocket.CLOSING) { + if (websocket && websocket.readyState !== this.websocketConstructor.CLOSED && websocket.readyState !== this.websocketConstructor.CLOSING) { websocket.close() } diff --git a/packages/server-sdk/src/index.ts b/packages/server-sdk/src/index.ts index 78508edf6..7d00e73ea 100644 --- a/packages/server-sdk/src/index.ts +++ b/packages/server-sdk/src/index.ts @@ -1,3 +1,4 @@ export * from './client' +export type * from './websocket-like' export type * from '@proj-airi/server-shared/types' export { ContextUpdateStrategy, WebSocketEventSource } from '@proj-airi/server-shared/types' diff --git a/packages/server-sdk/src/websocket-like.ts b/packages/server-sdk/src/websocket-like.ts new file mode 100644 index 000000000..95e4867bc --- /dev/null +++ b/packages/server-sdk/src/websocket-like.ts @@ -0,0 +1,30 @@ +export interface WebSocketMessageEventLike { + data: T +} + +export interface WebSocketErrorEventLike { + error?: Error +} + +export interface WebSocketLike { + readonly readyState: number + + onopen?: (event?: unknown) => void + onmessage?: (event: WebSocketMessageEventLike) => void + onerror?: (event: WebSocketErrorEventLike | unknown) => void + onclose?: (event?: unknown) => void + + send: (data: string | ArrayBufferLike | ArrayBufferView) => void + close: (code?: number, reason?: string) => void + + ping?: () => void + pong?: () => void +} + +export interface WebSocketLikeConstructor { + readonly OPEN: number + readonly CLOSING: number + readonly CLOSED: number + + new (url: string): WebSocketLike +} diff --git a/packages/server-sdk/test/client.test.ts b/packages/server-sdk/test/client.test.ts index 653b67373..10a3b2b03 100644 --- a/packages/server-sdk/test/client.test.ts +++ b/packages/server-sdk/test/client.test.ts @@ -36,6 +36,15 @@ class MockWebSocket { pong() {} } +class InjectedMockWebSocket extends MockWebSocket { + static instances: InjectedMockWebSocket[] = [] + + constructor(url: string) { + super(url) + InjectedMockWebSocket.instances.push(this) + } +} + vi.mock('crossws/websocket', () => ({ default: MockWebSocket, })) @@ -73,6 +82,7 @@ function emitMessage(socket: MockWebSocket, event: WebSocketEvent) { afterEach(() => { MockWebSocket.instances.length = 0 + InjectedMockWebSocket.instances.length = 0 vi.useRealTimers() }) @@ -166,6 +176,42 @@ describe('client', () => { expect(() => client.offEvent('input:text', listener)).not.toThrow() }) + it('uses an injected websocket constructor when provided', async () => { + const client = new Client({ + autoConnect: false, + autoReconnect: false, + name: 'test-plugin', + websocketConstructor: InjectedMockWebSocket, + }) + + const connected = client.connect() + const socket = InjectedMockWebSocket.instances.at(-1) + + expect(socket).toBeDefined() + expect(MockWebSocket.instances).toHaveLength(1) + + if (!socket) { + throw new Error('No custom mock websocket instance created') + } + + emitOpen(socket) + const announceEvent = parseSent(socket) + + emitMessage(socket, { + type: 'module:announced', + data: { + name: 'test-plugin', + identity: announceEvent.data.identity, + }, + metadata: { + source: { kind: 'plugin', plugin: { id: 'server' }, id: 'server-1' }, + event: { id: 'announce-1' }, + }, + }) + + await expect(connected).resolves.toBeUndefined() + }) + it('supports timeout-aware ensureConnected without cancelling the shared connect task', async () => { vi.useFakeTimers()