import { AnyRouter , ProcedureType } from '@trpc/server' ; import { TRPCClientIncomingMessage , TRPCClientIncomingRequest , TRPCRequest , TRPCResponse , TRPCResult , } from '@trpc/server/rpc' ; import { TRPCErrorResponse } from '@trpc/server/rpc' ; import { ObservableCallbacks , UnsubscribeFn } from '../internals/observable' ; import { retryDelay } from '../internals/retryDelay' ; import { TRPCClientError } from '../TRPCClientError' ; import { TRPCLink } from './core' ; type Operation = { id : number | string ; type : ProcedureType ; input : unknown ; path : string ; } ; export interface WebSocketClientOptions { url : string ; WebSocket ? : WebSocket ; retryDelayMs ? : typeof retryDelay ; } export function createWSClient ( opts ) { const { url , WebSocket : WebSocketImpl = WebSocket , retryDelayMs : retryDelayFn = retryDelay , } = opts ; if ( ! WebSocketImpl ) { throw new Error ( "No WebSocket implementation found - you probably don't want to use this on the server, but if you do you need to pass a `WebSocket`-ponyfill" ) ; } /**
     * outgoing messages buffer whilst not open
     */ let outgoing = [ ] ; /**
     * pending outgoing requests that are awaiting callback
     */ type TCallbacks = ObservableCallbacks < TRPCResult , TRPCClientError < AnyRouter > | TRPCErrorResponse > ; type TRequest = { /**
         * Reference to the WebSocket instance this request was made to
         */ ws : WebSocket ; type : ProcedureType ; callbacks : TCallbacks ; op : Operation ; } ; const pendingRequests = Object . create ( null ) ; let connectAttempt = 0 ; let dispatchTimer = null ; let connectTimer = null ; let activeConnection = createWS ( ) ; let state = 'connecting' ; /**
     * tries to send the list of messages
     */ function dispatch ( ) { if ( state !== 'open' || dispatchTimer ) { return ; } dispatchTimer = setTimeout ( ( ) => { dispatchTimer = null ; if ( outgoing . length === 1 ) { activeConnection . send ( JSON . stringify ( outgoing . pop ( ) ) ) ; } else { activeConnection . send ( JSON . stringify ( outgoing ) ) ; } outgoing = [ ] ; } ) ; } function tryReconnect ( ) { if ( connectTimer || state === 'closed' ) { return ; } const timeout = retryDelayFn ( connectAttempt ++ ) ; reconnectInMs ( timeout ) ; } function reconnect ( ) { state = 'connecting' ; const oldConnection = activeConnection ; activeConnection = createWS ( ) ; closeIfNoPending ( oldConnection ) ; } function reconnectInMs ( ms ) { if ( connectTimer ) { return ; } state = 'connecting' ; connectTimer = setTimeout ( reconnect , ms ) ; } function closeIfNoPending ( conn ) { const hasPendingRequests = Object . values ( pendingRequests ) . some ( ( p ) => p . ws === conn ) ; if ( ! hasPendingRequests ) { conn . close ( ) ; } } function resumeSubscriptionOnReconnect ( req ) { if ( outgoing . some ( ( r ) => r . id === req . op . id ) ) { return ; } request ( req . op , req . callbacks ) ; } function createWS ( ) { const conn = new WebSocket ( url ) ; clearTimeout ( connectTimer as any ) ; connectTimer = null ; conn . addEventListener ( 'open' , ( ) => { if ( conn !== activeConnection ) { return ; } connectAttempt = 0 ; state = 'open' ; dispatch ( ) ; } ) ; conn . addEventListener ( 'error' , ( ) => { if ( conn === activeConnection ) { tryReconnect ( ) ; } } ) ; const handleIncomingRequest = ( req ) => { if ( req . method === 'reconnect' && conn === activeConnection ) { reconnect ( ) ; for ( const pendingReq of Object . values ( pendingRequests ) ) { if ( pendingReq . type === 'subscription' ) { resumeSubscriptionOnReconnect ( pendingReq ) ; } } } } ; const handleIncomingResponse = ( res ) => { const req = res . id !== null && pendingRequests [ res . id ] ; if ( ! req ) { return ; } if ( 'error' in res ) { req . callbacks . onError ?. ( res ) ; return ; } req . callbacks . onNext ?. ( res . result ) ; if ( req . ws !== activeConnection && conn === activeConnection ) { const oldWs = req . ws ; req . ws = activeConnection ; closeIfNoPending ( oldWs ) ; } if ( res . result . type === 'stopped' && conn === activeConnection ) { req . callbacks . onDone ?. ( ) ; } } ; conn . addEventListener ( 'message' , ( { data } ) => { const msg = JSON . parse ( data ) as TRPCClientIncomingMessage ; if ( 'method' in msg ) { handleIncomingRequest ( msg ) ; } else { handleIncomingResponse ( msg ) ; } if ( conn !== activeConnection || state === 'closed' ) { closeIfNoPending ( conn ) ; } } ) ; conn . addEventListener ( 'close' , ( ) => { if ( activeConnection === conn ) { tryReconnect ( ) ; } for ( const key in pendingRequests ) { const req = pendingRequests [ key ] ; if ( req . ws !== conn ) { continue ; } req . callbacks . onError ?. ( TRPCClientError . from ( new TRPCWebSocketClosedError ( 'WebSocket closed prematurely' ) ) ) ; if ( req . type !== 'subscription' ) { delete pendingRequests [ key ] ; req . callbacks . onDone ?. ( ) ; } else if ( state !== 'closed' ) { resumeSubscriptionOnReconnect ( req ) ; } } } ) ; return conn ; } function request ( op , callbacks ) { const { type , input , path , id } = op ; const envelope = { id , jsonrpc : '2.0' , method : type , params : { input , path } } ; pendingRequests [ id ] = { ws : activeConnection , type , callbacks , op } ; outgoing . push ( envelope ) ; dispatch ( ) ; return ( ) => { const callbacks = pendingRequests [ id ] ?. callbacks ; delete pendingRequests [ id ] ; outgoing = outgoing . filter ( ( msg ) => msg . id !== id ) ; callbacks ?. onDone ?. ( ) ; if ( op . type === 'subscription' ) { outgoing . push ( { id , method : 'subscription.stop' , params : undefined } ) ; dispatch ( ) ; } } ; } return { close : ( ) => { state = 'closed' ; closeIfNoPending ( activeConnection ) ; } , request , getConnection ( ) { return activeConnection ; } } ; } export type TRPCWebSocketClient = ReturnType < typeof createWSClient > ; export interface WebSocketLinkOptions { client : TRPCWebSocketClient ; } class TRPCWebSocketClosedError extends Error { constructor ( message ) { super ( message ) ; this . name = 'TRPCWebSocketClosedError' ; Object . setPrototypeOf ( this , TRPCWebSocketClosedError . prototype ) ; } } class TRPCSubscriptionEndedError extends Error { constructor ( message ) { super ( message ) ; this . name = 'TRPCSubscriptionEndedError' ; Object . setPrototypeOf ( this , TRPCSubscriptionEndedError . prototype ) ; } } export function wsLink < TRouter extends AnyRouter > ( opts ) { return ( rt ) => { const { client } = opts ; return ( { op , prev , onDestroy } ) => { const { type , input : rawInput , path , id } = op ; const input = rt . transformer . serialize ( rawInput ) ; let isDone = false ; const unsub = client . request ( { type , path , input , id } , { onNext ( result ) { if ( isDone ) { return ; } if ( 'data' in result ) { const data = rt . transformer . deserialize ( result . data ) ; prev ( { type : 'data' , data } ) ; } else { prev ( result ) ; } if ( op . type !== 'subscription' ) { isDone = true ; unsub ( ) ; } } , onError ( err ) { if ( isDone ) { return ; } prev ( err instanceof Error ? err : TRPCClientError . from ( { ... err , error : rt . transformer . deserialize ( err . error ) } ) ) ; } , onDone ( ) { if ( isDone ) { return ; } const result = new TRPCSubscriptionEndedError ( 'Operation ended prematurely' ) ; prev ( TRPCClientError . from ( result , { isDone : true } ) ) ; isDone = true ; } } ) ; onDestroy ( ( ) => { isDone = true ; unsub ( ) ; } ) ; } ; } ; }