@nmtjs/client 0.15.0-beta.43 → 0.15.0-beta.45

package/src/core.ts

@@ -7,7 +7,13 @@ import type {
   ProtocolVersionInterface,
   ServerMessageTypePayload,
 } from '@nmtjs/protocol/client'
-import { anyAbortSignal, createFuture, MAX_UINT32, noopFn } from '@nmtjs/common'
+import {
+  anyAbortSignal,
+  createFuture,
+  MAX_UINT32,
+  noopFn,
+  withTimeout,
+} from '@nmtjs/common'
 import {
   ClientMessageType,
   ConnectionType,
@@ -23,6 +29,12 @@ import {
   versions,
 } from '@nmtjs/protocol/client'
 
+import type {
+  ClientDisconnectReason,
+  ClientPlugin,
+  ClientPluginEvent,
+  ClientPluginInstance,
+} from './plugins/types.ts'
 import type { BaseClientTransformer } from './transformers.ts'
 import type { ClientCallResponse, ClientTransportFactory } from './transport.ts'
 import type {
@@ -48,8 +60,7 @@ export type ProtocolClientCall = Future<any> & {
   signal?: AbortSignal
 }
 
-const DEFAULT_RECONNECT_TIMEOUT = 1000
-const DEFAULT_MAX_RECONNECT_TIMEOUT = 60000
+const DEFAULT_RECONNECT_REASON = 'connect_error'
 
 export interface BaseClientOptions<
   RouterContract extends TAnyRouterContract = TAnyRouterContract,
@@ -60,7 +71,7 @@ export interface BaseClientOptions<
   format: BaseClientFormat
   application?: string
   timeout?: number
-  autoreconnect?: boolean
+  plugins?: ClientPlugin[]
   safe?: SafeCall
 }
 
@@ -79,7 +90,8 @@ export abstract class BaseClient<
   OutputTypeProvider extends TypeProvider = TypeProvider,
 > extends EventEmitter<{
   connected: []
-  disconnected: [reason: 'server' | 'client' | (string & {})]
+  disconnected: [reason: ClientDisconnectReason]
+  pong: [nonce: number]
 }> {
   _!: {
     routes: ResolveAPIRouterRoutes<
@@ -111,11 +123,20 @@ export abstract class BaseClient<
   protected callId = 0
   protected streamId = 0
   protected cab: AbortController | null = null
-  protected reconnectTimeout = DEFAULT_RECONNECT_TIMEOUT
   protected connecting: Promise<void> | null = null
-  protected state: 'connected' | 'disconnected' = 'disconnected'
 
-  #auth: any
+  protected _state: 'connected' | 'disconnected' = 'disconnected'
+  protected _lastDisconnectReason: ClientDisconnectReason = 'server'
+  protected _disposed = false
+
+  protected pingNonce = 0
+  protected pendingPings = new Map<number, Future<void>>()
+  protected plugins: ClientPluginInstance[] = []
+
+  private clientDisconnectAsReconnect = false
+  private clientDisconnectOverrideReason: string | null = null
+
+  private authValue: any
 
   constructor(
     readonly options: BaseClientOptions<RouterContract, SafeCall>,
@@ -138,48 +159,54 @@ export abstract class BaseClient<
       this.transportOptions,
     ) as any
 
-    if (
-      this.transport.type === ConnectionType.Bidirectional &&
-      this.options.autoreconnect
-    ) {
-      this.on('disconnected', async (reason) => {
-        while (this.state === 'disconnected') {
-          const timeout = new Promise((resolve) =>
-            setTimeout(resolve, this.reconnectTimeout),
-          )
-          this.reconnectTimeout = Math.min(
-            this.reconnectTimeout * 2,
-            DEFAULT_MAX_RECONNECT_TIMEOUT,
-          )
-          await timeout
-          await this.connect().catch(noopFn)
-        }
-      })
-
-      this.on('connected', () => {
-        this.reconnectTimeout = DEFAULT_RECONNECT_TIMEOUT
-      })
+    this.plugins = this.options.plugins?.map((plugin) => plugin(this)) ?? []
+    for (const plugin of this.plugins) {
+      plugin.onInit?.()
+    }
+  }
 
-      if (globalThis.window) {
-        globalThis.window.addEventListener('pageshow', () => {
-          if (this.state === 'disconnected') this.connect()
-        })
-      }
+  dispose() {
+    this._disposed = true
+    this.stopAllPendingPings('dispose')
+    for (let i = this.plugins.length - 1; i >= 0; i--) {
+      this.plugins[i].dispose?.()
     }
   }
 
+  get state() {
+    return this._state
+  }
+
+  get lastDisconnectReason() {
+    return this._lastDisconnectReason
+  }
+
+  get transportType() {
+    return this.transport.type
+  }
+
+  isDisposed() {
+    return this._disposed
+  }
+
+  requestReconnect(reason?: string) {
+    return this.disconnect({ reconnect: true, reason })
+  }
+
   get auth() {
-    return this.#auth
+    return this.authValue
   }
 
   set auth(value) {
-    this.#auth = value
+    this.authValue = value
   }
 
   connect() {
-    if (this.state === 'connected') return Promise.resolve()
+    if (this._state === 'connected') return Promise.resolve()
     if (this.connecting) return this.connecting
 
+    if (this._disposed) return Promise.reject(new Error('Client is disposed'))
+
     const _connect = async () => {
       if (this.transport.type === ConnectionType.Bidirectional) {
         this.cab = new AbortController()
@@ -187,7 +214,7 @@ export abstract class BaseClient<
         const serverStreams = this.serverStreams
         const transport = {
           send: (buffer) => {
-            this.#send(buffer).catch(noopFn)
+            this.send(buffer).catch(noopFn)
           },
         }
         this.messageContext = {
@@ -195,7 +222,7 @@ export abstract class BaseClient<
           encoder: this.options.format,
           decoder: this.options.format,
           addClientStream: (blob) => {
-            const streamId = this.#getStreamId()
+            const streamId = this.getStreamId()
             return this.clientStreams.add(blob.source, streamId, blob.metadata)
           },
           addServerStream(streamId, metadata) {
@@ -234,7 +261,7 @@ export abstract class BaseClient<
               return stream
             }
           },
-          streamId: this.#getStreamId.bind(this),
+          streamId: this.getStreamId.bind(this),
         }
         return this.transport.connect({
           auth: this.auth,
@@ -246,19 +273,46 @@ export abstract class BaseClient<
       }
     }
 
+    let emitDisconnectOnFailure: 'server' | 'client' | (string & {}) | null =
+      null
+
     this.connecting = _connect()
       .then(() => {
-        this.state = 'connected'
+        this._state = 'connected'
+      })
+      .catch((error) => {
+        if (this.transport.type === ConnectionType.Bidirectional) {
+          emitDisconnectOnFailure = DEFAULT_RECONNECT_REASON
+        }
+        throw error
       })
       .finally(() => {
         this.connecting = null
+
+        if (emitDisconnectOnFailure && !this._disposed) {
+          this._state = 'disconnected'
+          this._lastDisconnectReason = emitDisconnectOnFailure
+          void this.onDisconnect(emitDisconnectOnFailure)
+        }
       })
 
     return this.connecting
   }
 
-  async disconnect() {
+  async disconnect(options: { reconnect?: boolean; reason?: string } = {}) {
     if (this.transport.type === ConnectionType.Bidirectional) {
+      // Ensure connect() won't short-circuit while the transport is closing.
+      this._state = 'disconnected'
+      this._lastDisconnectReason = 'client'
+
+      if (options.reconnect) {
+        this.clientDisconnectAsReconnect = true
+        this.clientDisconnectOverrideReason = options.reason ?? 'server'
+      } else {
+        this.clientDisconnectAsReconnect = false
+        this.clientDisconnectOverrideReason = null
+      }
+
       this.cab!.abort()
       await this.transport.disconnect()
       this.messageContext = null
@@ -290,12 +344,19 @@ export abstract class BaseClient<
 
     const signal = signals.length ? anyAbortSignal(...signals) : undefined
 
-    const callId = this.#getCallId()
+    const callId = this.getCallId()
     const call = createFuture() as ProtocolClientCall
     call.procedure = procedure
     call.signal = signal
 
     this.calls.set(callId, call)
+    this.emitClientEvent({
+      kind: 'rpc_request',
+      timestamp: Date.now(),
+      callId,
+      procedure,
+      body: payload,
+    })
 
     // Check if signal is already aborted before proceeding
     if (signal?.aborted) {
@@ -322,7 +383,7 @@ export abstract class BaseClient<
                 ClientMessageType.RpcAbort,
                 { callId },
               )
-              this.#send(buffer).catch(noopFn)
+              this.send(buffer).catch(noopFn)
             }
           },
           { once: true },
@@ -337,7 +398,7 @@ export abstract class BaseClient<
             ClientMessageType.Rpc,
             { callId, procedure, payload: transformedPayload },
           )
-          await this.#send(buffer, signal)
+          await this.send(buffer, signal)
         } else {
           const response = await this.transport.call(
             {
@@ -348,9 +409,16 @@ export abstract class BaseClient<
             { callId, procedure, payload: transformedPayload },
             { signal, _stream_response: options._stream_response },
           )
-          this.#handleCallResponse(callId, response)
+          this.handleCallResponse(callId, response)
         }
       } catch (error) {
+        this.emitClientEvent({
+          kind: 'rpc_error',
+          timestamp: Date.now(),
+          callId,
+          procedure,
+          error,
+        })
         call.reject(error)
       }
     }
@@ -362,6 +430,11 @@ export abstract class BaseClient<
             controller.abort()
           })
         }
+
+        if (options._stream_response && typeof value === 'function') {
+          return value
+        }
+
         controller.abort()
         return value
       },
@@ -386,30 +459,158 @@ export abstract class BaseClient<
   }
 
   protected async onConnect() {
-    this.state = 'connected'
+    this._state = 'connected'
+    this._lastDisconnectReason = 'server'
+    this.emitClientEvent({
+      kind: 'connected',
+      timestamp: Date.now(),
+      transportType:
+        this.transport.type === ConnectionType.Bidirectional
+          ? 'bidirectional'
+          : 'unidirectional',
+    })
+    for (const plugin of this.plugins) {
+      await plugin.onConnect?.()
+    }
     this.emit('connected')
   }
 
-  protected async onDisconnect(reason: 'client' | 'server' | (string & {})) {
-    this.state = 'disconnected'
-    this.emit('disconnected', reason)
-    this.clientStreams.clear(reason)
-    this.serverStreams.clear(reason)
-    this.rpcStreams.clear(reason)
+  protected async onDisconnect(reason: ClientDisconnectReason) {
+    const effectiveReason =
+      reason === 'client' && this.clientDisconnectAsReconnect
+        ? (this.clientDisconnectOverrideReason ?? 'server')
+        : reason
+
+    this.clientDisconnectAsReconnect = false
+    this.clientDisconnectOverrideReason = null
+
+    this._state = 'disconnected'
+    this._lastDisconnectReason = effectiveReason
+    this.emitClientEvent({
+      kind: 'disconnected',
+      timestamp: Date.now(),
+      reason: effectiveReason,
+    })
+
+    // Connection is gone, never keep old message context around.
+    this.messageContext = null
+
+    this.stopAllPendingPings(effectiveReason)
+
+    // Fail-fast: do not keep pending calls around across disconnects.
+    if (this.calls.size) {
+      const error = new ProtocolError(
+        ErrorCode.ConnectionError,
+        'Disconnected',
+        { reason: effectiveReason },
+      )
+      for (const call of this.calls.values()) {
+        call.reject(error)
+      }
+      this.calls.clear()
+    }
+
+    if (this.cab) {
+      try {
+        this.cab.abort(reason)
+      } catch {
+        this.cab.abort()
+      }
+      this.cab = null
+    }
+
+    this.emit('disconnected', effectiveReason)
+
+    for (let i = this.plugins.length - 1; i >= 0; i--) {
+      await this.plugins[i].onDisconnect?.(effectiveReason)
+    }
+
+    void this.clientStreams.clear(effectiveReason)
+    void this.serverStreams.clear(effectiveReason)
+    void this.rpcStreams.clear(effectiveReason)
+  }
+
+  protected nextPingNonce() {
+    if (this.pingNonce >= MAX_UINT32) this.pingNonce = 0
+    return this.pingNonce++
+  }
+
+  ping(timeout: number, signal?: AbortSignal) {
+    if (
+      !this.messageContext ||
+      this.transport.type !== ConnectionType.Bidirectional
+    ) {
+      return Promise.reject(new Error('Client is not connected'))
+    }
+
+    const nonce = this.nextPingNonce()
+    const future = createFuture<void>()
+    this.pendingPings.set(nonce, future)
+
+    const buffer = this.protocol.encodeMessage(
+      this.messageContext,
+      ClientMessageType.Ping,
+      { nonce },
+    )
+
+    return this.send(buffer, signal)
+      .then(() =>
+        withTimeout(future.promise, timeout, new Error('Heartbeat timeout')),
+      )
+      .finally(() => {
+        this.pendingPings.delete(nonce)
+      })
+  }
+
+  protected stopAllPendingPings(reason?: any) {
+    if (!this.pendingPings.size) return
+    const error = new Error('Heartbeat stopped', { cause: reason })
+    for (const pending of this.pendingPings.values()) pending.reject(error)
+    this.pendingPings.clear()
   }
 
   protected async onMessage(buffer: ArrayBufferView) {
     if (!this.messageContext) return
 
     const message = this.protocol.decodeMessage(this.messageContext, buffer)
+    for (const plugin of this.plugins) {
+      plugin.onServerMessage?.(message, buffer)
+    }
+    this.emitClientEvent({
+      kind: 'server_message',
+      timestamp: Date.now(),
+      messageType: message.type,
+      rawByteLength: buffer.byteLength,
+      body: message,
+    })
 
     switch (message.type) {
       case ServerMessageType.RpcResponse:
-        this.#handleRPCResponseMessage(message)
+        this.handleRPCResponseMessage(message)
         break
       case ServerMessageType.RpcStreamResponse:
-        this.#handleRPCStreamResponseMessage(message)
+        this.handleRPCStreamResponseMessage(message)
         break
+      case ServerMessageType.Pong: {
+        const pending = this.pendingPings.get(message.nonce)
+        if (pending) {
+          this.pendingPings.delete(message.nonce)
+          pending.resolve()
+        }
+        this.emit('pong', message.nonce)
+        break
+      }
+      case ServerMessageType.Ping: {
+        if (this.messageContext) {
+          const buffer = this.protocol.encodeMessage(
+            this.messageContext,
+            ClientMessageType.Pong,
+            { nonce: message.nonce },
+          )
+          this.send(buffer).catch(noopFn)
+        }
+        break
+      }
       case ServerMessageType.RpcStreamChunk:
         this.rpcStreams.push(message.callId, message.chunk)
         break
@@ -439,14 +640,14 @@ export abstract class BaseClient<
                 ClientMessageType.ClientStreamPush,
                 { streamId: message.streamId, chunk },
               )
-              this.#send(buffer).catch(noopFn)
+              this.send(buffer).catch(noopFn)
             } else {
               const buffer = this.protocol.encodeMessage(
                 this.messageContext!,
                 ClientMessageType.ClientStreamEnd,
                 { streamId: message.streamId },
               )
-              this.#send(buffer).catch(noopFn)
+              this.send(buffer).catch(noopFn)
               this.clientStreams.end(message.streamId)
             }
           },
@@ -456,7 +657,7 @@ export abstract class BaseClient<
               ClientMessageType.ClientStreamAbort,
               { streamId: message.streamId },
             )
-            this.#send(buffer).catch(noopFn)
+            this.send(buffer).catch(noopFn)
             this.clientStreams.remove(message.streamId)
           },
         )
@@ -467,19 +668,40 @@ export abstract class BaseClient<
     }
   }
 
-  #handleRPCResponseMessage(
+  private handleRPCResponseMessage(
     message: ServerMessageTypePayload[ServerMessageType.RpcResponse],
   ) {
     const { callId, result, error } = message
     const call = this.calls.get(callId)
     if (!call) return
     if (error) {
+      this.emitClientEvent({
+        kind: 'rpc_error',
+        timestamp: Date.now(),
+        callId,
+        procedure: call.procedure,
+        error,
+      })
       call.reject(new ProtocolError(error.code, error.message, error.data))
     } else {
       try {
         const transformed = this.transformer.decode(call.procedure, result)
+        this.emitClientEvent({
+          kind: 'rpc_response',
+          timestamp: Date.now(),
+          callId,
+          procedure: call.procedure,
+          body: transformed,
+        })
         call.resolve(transformed)
       } catch (error) {
+        this.emitClientEvent({
+          kind: 'rpc_error',
+          timestamp: Date.now(),
+          callId,
+          procedure: call.procedure,
+          error,
+        })
         call.reject(
           new ProtocolError(
             ErrorCode.ClientRequestError,
@@ -491,12 +713,19 @@ export abstract class BaseClient<
     }
   }
 
-  #handleRPCStreamResponseMessage(
+  private handleRPCStreamResponseMessage(
     message: ServerMessageTypePayload[ServerMessageType.RpcStreamResponse],
   ) {
     const call = this.calls.get(message.callId)
     if (message.error) {
       if (!call) return
+      this.emitClientEvent({
+        kind: 'rpc_error',
+        timestamp: Date.now(),
+        callId: message.callId,
+        procedure: call.procedure,
+        error: message.error,
+      })
       call.reject(
         new ProtocolError(
           message.error.code,
@@ -507,6 +736,13 @@ export abstract class BaseClient<
     } else {
       if (call) {
         const { procedure, signal } = call
+        this.emitClientEvent({
+          kind: 'rpc_response',
+          timestamp: Date.now(),
+          callId: message.callId,
+          procedure,
+          stream: true,
+        })
         const stream = new ProtocolServerRPCStream({
           start: (controller) => {
             if (signal) {
@@ -525,7 +761,7 @@ export abstract class BaseClient<
                           ClientMessageType.RpcAbort,
                           { callId: message.callId, reason: signal.reason },
                         )
-                        this.#send(buffer).catch(noopFn)
+                        this.send(buffer).catch(noopFn)
                       }
                     }
                   },
@@ -545,7 +781,7 @@ export abstract class BaseClient<
               ClientMessageType.RpcPull,
               { callId: message.callId },
             )
-            this.#send(buffer).catch(noopFn)
+            this.send(buffer).catch(noopFn)
           },
           readableStrategy: { highWaterMark: 0 },
         })
@@ -561,17 +797,24 @@ export abstract class BaseClient<
             ClientMessageType.RpcAbort,
             { callId: message.callId },
           )
-          this.#send(buffer).catch(noopFn)
+          this.send(buffer).catch(noopFn)
         }
       }
     }
   }
 
-  #handleCallResponse(callId: number, response: ClientCallResponse) {
+  private handleCallResponse(callId: number, response: ClientCallResponse) {
     const call = this.calls.get(callId)
 
     if (response.type === 'rpc_stream') {
       if (call) {
+        this.emitClientEvent({
+          kind: 'rpc_response',
+          timestamp: Date.now(),
+          callId,
+          procedure: call.procedure,
+          stream: true,
+        })
         const stream = new ProtocolServerStream({
           transform: (chunk) => {
             return this.transformer.decode(
@@ -582,7 +825,36 @@ export abstract class BaseClient<
         })
         this.rpcStreams.add(callId, stream)
         call.resolve(({ signal }: { signal?: AbortSignal }) => {
-          response.stream.pipeTo(stream.writable, { signal }).catch(noopFn)
+          const reader = response.stream.getReader()
+
+          let onAbort: (() => void) | undefined
+          if (signal) {
+            onAbort = () => {
+              reader.cancel(signal.reason).catch(noopFn)
+              this.rpcStreams.abort(callId).catch(noopFn)
+            }
+            if (signal.aborted) onAbort()
+            else signal.addEventListener('abort', onAbort, { once: true })
+          }
+
+          void (async () => {
+            try {
+              while (true) {
+                const { done, value } = await reader.read()
+                if (done) break
+                await this.rpcStreams.push(callId, value)
+              }
+              await this.rpcStreams.end(callId)
+            } catch {
+              await this.rpcStreams.abort(callId).catch(noopFn)
+            } finally {
+              reader.releaseLock()
+              if (signal && onAbort) {
+                signal.removeEventListener('abort', onAbort)
+              }
+            }
+          })()
+
           return stream
         })
       } else {
@@ -593,9 +865,16 @@ export abstract class BaseClient<
       }
     } else if (response.type === 'blob') {
       if (call) {
+        this.emitClientEvent({
+          kind: 'rpc_response',
+          timestamp: Date.now(),
+          callId,
+          procedure: call.procedure,
+          stream: true,
+        })
         const { metadata, source } = response
         const stream = new ProtocolServerBlobStream(metadata)
-        this.serverStreams.add(this.#getStreamId(), stream)
+        this.serverStreams.add(this.getStreamId(), stream)
         call.resolve(({ signal }: { signal?: AbortSignal }) => {
           source.pipeTo(stream.writable, { signal }).catch(noopFn)
           return stream
@@ -613,8 +892,22 @@ export abstract class BaseClient<
           call.procedure,
           response.result,
         )
+        this.emitClientEvent({
+          kind: 'rpc_response',
+          timestamp: Date.now(),
+          callId,
+          procedure: call.procedure,
+          body: transformed,
+        })
         call.resolve(transformed)
       } catch (error) {
+        this.emitClientEvent({
+          kind: 'rpc_error',
+          timestamp: Date.now(),
+          callId,
+          procedure: call.procedure,
+          error,
+        })
         call.reject(
           new ProtocolError(
             ErrorCode.ClientRequestError,
@@ -626,23 +919,32 @@ export abstract class BaseClient<
     }
   }
 
-  #send(buffer: ArrayBufferView, signal?: AbortSignal) {
+  protected send(buffer: ArrayBufferView, signal?: AbortSignal) {
     if (this.transport.type === ConnectionType.Unidirectional)
       throw new Error('Invalid transport type for send')
     return this.transport.send(buffer, { signal })
   }
 
-  #getStreamId() {
+  protected getStreamId() {
     if (this.streamId >= MAX_UINT32) {
       this.streamId = 0
     }
     return this.streamId++
   }
 
-  #getCallId() {
+  protected getCallId() {
     if (this.callId >= MAX_UINT32) {
       this.callId = 0
     }
     return this.callId++
   }
+
+  protected emitClientEvent(event: ClientPluginEvent) {
+    for (const plugin of this.plugins) {
+      try {
+        const result = plugin.onClientEvent?.(event)
+        Promise.resolve(result).catch(noopFn)
+      } catch {}
+    }
+  }
 }