@replit/river 0.9.3 → 0.10.0

package/dist/router/server.d.ts

@@ -1,10 +1,10 @@
 import { Static } from '@sinclair/typebox';
 import { Connection, Transport } from '../transport/transport';
-import { AnyProcedure, AnyService, PayloadType } from './builder';
+import { AnyProcedure, PayloadType } from './builder';
 import type { Pushable } from 'it-pushable';
-import { TransportMessage } from '../transport/message';
 import { ServiceContext } from './context';
 import { Result, RiverError } from './result';
+import { ServiceDefs } from './defs';
 /**
  * Represents a server with a set of services. Use {@link createServer} to create it.
  * @template Services - The type of services provided by the server.
@@ -19,8 +19,8 @@ interface ProcStream {
     serviceName: string;
     procedureName: string;
     procedure: AnyProcedure;
-    incoming: Pushable<TransportMessage>;
-    outgoing: Pushable<TransportMessage<Result<Static<PayloadType>, Static<RiverError>>>>;
+    incoming: Pushable<PayloadType>;
+    outgoing: Pushable<Result<Static<PayloadType>, Static<RiverError>>>;
     promises: {
         outputHandler: Promise<unknown>;
         inputHandler: Promise<unknown>;
@@ -34,6 +34,6 @@ interface ProcStream {
  * @param extendedContext - An optional object containing additional context to be passed to all services.
  * @returns A promise that resolves to a server instance with the registered services.
  */
-export declare function createServer<Services extends Record<string, AnyService>>(transport: Transport<Connection>, services: Services, extendedContext?: Omit<ServiceContext, 'state'>): Promise<Server<Services>>;
+export declare function createServer<Services extends ServiceDefs>(transport: Transport<Connection>, services: Services, extendedContext?: Omit<ServiceContext, 'state'>): Server<Services>;
 export {};
 //# sourceMappingURL=server.d.ts.map

package/dist/router/server.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"server.d.ts","sourceRoot":"","sources":["../../router/server.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,MAAM,EAAE,MAAM,mBAAmB,CAAC;AAC3C,OAAO,EAAE,UAAU,EAAE,SAAS,EAAE,MAAM,wBAAwB,CAAC;AAC/D,OAAO,EAAE,YAAY,EAAE,UAAU,EAAE,WAAW,EAAE,MAAM,WAAW,CAAC;AAElE,OAAO,KAAK,EAAE,QAAQ,EAAE,MAAM,aAAa,CAAC;AAC5C,OAAO,EAGL,gBAAgB,EAKjB,MAAM,sBAAsB,CAAC;AAC9B,OAAO,EAAE,cAAc,EAA2B,MAAM,WAAW,CAAC;AAGpE,OAAO,EAEL,MAAM,EACN,UAAU,EAGX,MAAM,UAAU,CAAC;AAElB;;;GAGG;AACH,MAAM,WAAW,MAAM,CAAC,QAAQ;IAC9B,QAAQ,EAAE,QAAQ,CAAC;IACnB,OAAO,EAAE,GAAG,CAAC,MAAM,EAAE,UAAU,CAAC,CAAC;IACjC,KAAK,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;CACxB;AAED,UAAU,UAAU;IAClB,EAAE,EAAE,MAAM,CAAC;IACX,WAAW,EAAE,MAAM,CAAC;IACpB,aAAa,EAAE,MAAM,CAAC;IACtB,SAAS,EAAE,YAAY,CAAC;IACxB,QAAQ,EAAE,QAAQ,CAAC,gBAAgB,CAAC,CAAC;IACrC,QAAQ,EAAE,QAAQ,CAChB,gBAAgB,CAAC,MAAM,CAAC,MAAM,CAAC,WAAW,CAAC,EAAE,MAAM,CAAC,UAAU,CAAC,CAAC,CAAC,CAClE,CAAC;IACF,QAAQ,EAAE;QACR,aAAa,EAAE,OAAO,CAAC,OAAO,CAAC,CAAC;QAChC,YAAY,EAAE,OAAO,CAAC,OAAO,CAAC,CAAC;KAChC,CAAC;CACH;AAED;;;;;;;GAOG;AACH,wBAAsB,YAAY,CAAC,QAAQ,SAAS,MAAM,CAAC,MAAM,EAAE,UAAU,CAAC,EAC5E,SAAS,EAAE,SAAS,CAAC,UAAU,CAAC,EAChC,QAAQ,EAAE,QAAQ,EAClB,eAAe,CAAC,EAAE,IAAI,CAAC,cAAc,EAAE,OAAO,CAAC,GAC9C,OAAO,CAAC,MAAM,CAAC,QAAQ,CAAC,CAAC,CAqQ3B"}
+{"version":3,"file":"server.d.ts","sourceRoot":"","sources":["../../router/server.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,MAAM,EAAE,MAAM,mBAAmB,CAAC;AAC3C,OAAO,EAAE,UAAU,EAAE,SAAS,EAAE,MAAM,wBAAwB,CAAC;AAC/D,OAAO,EAAE,YAAY,EAAc,WAAW,EAAE,MAAM,WAAW,CAAC;AAElE,OAAO,KAAK,EAAE,QAAQ,EAAE,MAAM,aAAa,CAAC;AAU5C,OAAO,EAAE,cAAc,EAA2B,MAAM,WAAW,CAAC;AAGpE,OAAO,EAEL,MAAM,EACN,UAAU,EAGX,MAAM,UAAU,CAAC;AAElB,OAAO,EAAE,WAAW,EAAE,MAAM,QAAQ,CAAC;AAErC;;;GAGG;AACH,MAAM,WAAW,MAAM,CAAC,QAAQ;IAC9B,QAAQ,EAAE,QAAQ,CAAC;IACnB,OAAO,EAAE,GAAG,CAAC,MAAM,EAAE,UAAU,CAAC,CAAC;IACjC,KAAK,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;CACxB;AAED,UAAU,UAAU;IAClB,EAAE,EAAE,MAAM,CAAC;IACX,WAAW,EAAE,MAAM,CAAC;IACpB,aAAa,EAAE,MAAM,CAAC;IACtB,SAAS,EAAE,YAAY,CAAC;IACxB,QAAQ,EAAE,QAAQ,CAAC,WAAW,CAAC,CAAC;IAChC,QAAQ,EAAE,QAAQ,CAAC,MAAM,CAAC,MAAM,CAAC,WAAW,CAAC,EAAE,MAAM,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;IACpE,QAAQ,EAAE;QACR,aAAa,EAAE,OAAO,CAAC,OAAO,CAAC,CAAC;QAChC,YAAY,EAAE,OAAO,CAAC,OAAO,CAAC,CAAC;KAChC,CAAC;CACH;AAsUD;;;;;;;GAOG;AACH,wBAAgB,YAAY,CAAC,QAAQ,SAAS,WAAW,EACvD,SAAS,EAAE,SAAS,CAAC,UAAU,CAAC,EAChC,QAAQ,EAAE,QAAQ,EAClB,eAAe,CAAC,EAAE,IAAI,CAAC,cAAc,EAAE,OAAO,CAAC,GAC9C,MAAM,CAAC,QAAQ,CAAC,CAElB"}

package/dist/router/server.js

@@ -3,104 +3,109 @@ import { ControlMessagePayloadSchema, isStreamClose, isStreamOpen, reply, closeS
 import { log } from '../logging';
 import { Value } from '@sinclair/typebox/value';
 import { Err, UNCAUGHT_ERROR, } from './result';
-/**
- * Creates a server instance that listens for incoming messages from a transport and routes them to the appropriate service and procedure.
- * The server tracks the state of each service along with open streams and the extended context object.
- * @param transport - The transport to listen to.
- * @param services - An object containing all the services to be registered on the server.
- * @param extendedContext - An optional object containing additional context to be passed to all services.
- * @returns A promise that resolves to a server instance with the registered services.
- */
-export async function createServer(transport, services, extendedContext) {
-    const contextMap = new Map();
+class RiverServer {
+    transport;
+    services;
+    contextMap;
     // map of streamId to ProcStream
-    const streamMap = new Map();
-    async function cleanupStream(id) {
-        const stream = streamMap.get(id);
-        if (!stream) {
-            return;
+    streamMap;
+    // map of client to their open streams by streamId
+    clientStreams;
+    constructor(transport, services, extendedContext) {
+        this.transport = transport;
+        this.services = services;
+        this.contextMap = new Map();
+        for (const service of Object.values(services)) {
+            this.contextMap.set(service, {
+                ...extendedContext,
+                state: service.state,
+            });
         }
-        stream.incoming.end();
-        await stream.promises.inputHandler;
-        stream.outgoing.end();
-        await stream.promises.outputHandler;
-        streamMap.delete(id);
+        this.streamMap = new Map();
+        this.clientStreams = new Map();
+        this.transport.addEventListener('message', this.handler);
+        this.transport.addEventListener('connectionStatus', this.onDisconnect);
     }
-    function getContext(service) {
-        const context = contextMap.get(service);
-        if (!context) {
-            const err = `${transport.clientId} -- no context found for ${service.name}`;
-            log?.error(err);
-            throw new Error(err);
-        }
-        return context;
-    }
-    // populate the context map
-    for (const service of Object.values(services)) {
-        contextMap.set(service, { ...extendedContext, state: service.state });
+    get streams() {
+        return this.streamMap;
     }
-    const handler = async (message) => {
-        if (message.to !== transport.clientId) {
-            log?.info(`${transport.clientId} -- got msg with destination that isn't the server, ignoring`);
+    handler = async (message) => {
+        if (message.to !== this.transport.clientId) {
+            log?.info(`${this.transport.clientId} -- got msg with destination that isn't the server, ignoring`);
             return;
         }
-        const streamIdx = message.streamId;
-        const procStream = streamMap.get(streamIdx);
-        if (procStream) {
-            // If the stream is a continuation, we do not admit the init messages.
-            if (Value.Check(procStream.procedure.input, message.payload)) {
-                procStream.incoming.push(message);
-            }
-            else if (!Value.Check(ControlMessagePayloadSchema, message.payload)) {
-                log?.error(`${transport.clientId} -- procedure ${procStream.serviceName}.${procStream.procedureName} received invalid payload: ${JSON.stringify(message.payload)}`);
-            }
-            if (isStreamClose(message.controlFlags)) {
-                await cleanupStream(streamIdx);
-            }
+        let procStream = this.streamMap.get(message.streamId);
+        const isInitMessage = !procStream;
+        // create a proc stream if it doesnt exist
+        procStream ||= this.createNewProcStream(message);
+        if (!procStream) {
+            // if we fail to create a proc stream, just abort
             return;
         }
+        await this.pushToStream(procStream, message, isInitMessage);
+    };
+    // cleanup streams on unexpected disconnections
+    onDisconnect = async (evt) => {
+        if (evt.status !== 'disconnect') {
+            return;
+        }
+        const disconnectedClientId = evt.conn.connectedTo;
+        log?.info(`${this.transport.clientId} -- got unexpected disconnect from ${disconnectedClientId}, cleaning up streams`);
+        const streamsFromThisClient = this.clientStreams.get(disconnectedClientId);
+        if (!streamsFromThisClient) {
+            return;
+        }
+        await Promise.all(Array.from(streamsFromThisClient).map(this.cleanupStream));
+        this.clientStreams.delete(disconnectedClientId);
+    };
+    async close() {
+        this.transport.removeEventListener('message', this.handler);
+        this.transport.removeEventListener('connectionStatus', this.onDisconnect);
+        await Promise.all([...this.streamMap.keys()].map(this.cleanupStream));
+    }
+    createNewProcStream(message) {
         if (!isStreamOpen(message.controlFlags)) {
-            log?.warn(`${transport.clientId} -- couldn't find a matching procedure stream for ${message.serviceName}.${message.procedureName}:${message.streamId}`);
+            log?.warn(`${this.transport.clientId} -- couldn't find a matching procedure stream for ${message.serviceName}.${message.procedureName}:${message.streamId}`);
             return;
         }
-        if (!message.serviceName || !(message.serviceName in services)) {
-            log?.warn(`${transport.clientId} -- couldn't find service ${message.serviceName}`);
+        if (!message.serviceName || !(message.serviceName in this.services)) {
+            log?.warn(`${this.transport.clientId} -- couldn't find service ${message.serviceName}`);
             return;
         }
-        const service = services[message.serviceName];
-        const serviceContext = getContext(service);
+        const service = this.services[message.serviceName];
+        const serviceContext = this.getContext(service);
         if (!message.procedureName ||
             !(message.procedureName in service.procedures)) {
-            log?.warn(`${transport.clientId} -- couldn't find a matching procedure for ${message.serviceName}.${message.procedureName}`);
+            log?.warn(`${this.transport.clientId} -- couldn't find a matching procedure for ${message.serviceName}.${message.procedureName}`);
             return;
         }
         const procedure = service.procedures[message.procedureName];
-        const procHasInitMessage = 'init' in procedure;
         const incoming = pushable({ objectMode: true });
         const outgoing = pushable({ objectMode: true });
         const outputHandler = 
         // sending outgoing messages back to client
         (async () => {
             for await (const response of outgoing) {
-                transport.send(response);
+                this.transport.send(reply(message, response));
             }
             // we ended, send a close bit back to the client
             // only subscriptions and streams have streams the
             // handler can close
             if (procedure.type === 'subscription' || procedure.type === 'stream') {
-                transport.send(closeStream(transport.clientId, message.from, message.streamId));
+                this.transport.send(closeStream(this.transport.clientId, message.from, message.streamId));
             }
         })();
-        function errorHandler(err) {
+        const errorHandler = (err) => {
             const errorMsg = err instanceof Error ? err.message : `[coerced to error] ${err}`;
-            log?.error(`${transport.clientId} -- procedure ${message.serviceName}.${message.procedureName}:${message.streamId} threw an error: ${errorMsg}`);
-            outgoing.push(reply(message, Err({
+            log?.error(`${this.transport.clientId} -- procedure ${message.serviceName}.${message.procedureName}:${message.streamId} threw an error: ${errorMsg}`);
+            outgoing.push(Err({
                 code: UNCAUGHT_ERROR,
                 message: errorMsg,
-            })));
-        }
+            }));
+        };
         // pump incoming message stream -> handler -> outgoing message stream
         let inputHandler;
+        const procHasInitMessage = 'init' in procedure;
         if (procedure.type === 'stream') {
             if (procHasInitMessage) {
                 inputHandler = (async () => {
@@ -179,10 +184,10 @@ export async function createServer(transport, services, extendedContext) {
         else {
             // procedure is inferred to be never here as this is not a valid procedure type
             // we cast just to log
-            log?.warn(`${transport.clientId} -- got request for invalid procedure type ${procedure.type} at ${message.serviceName}.${message.procedureName}`);
+            log?.warn(`${this.transport.clientId} -- got request for invalid procedure type ${procedure.type} at ${message.serviceName}.${message.procedureName}`);
             return;
         }
-        streamMap.set(streamIdx, {
+        const procStream = {
             id: message.streamId,
             incoming,
             outgoing,
@@ -190,28 +195,66 @@ export async function createServer(transport, services, extendedContext) {
             procedureName: message.procedureName,
             procedure,
             promises: { inputHandler, outputHandler },
-        });
-        // This is the first message, so we parse is as the initialization message, if supplied.
-        if ((!procHasInitMessage && Value.Check(procedure.input, message.payload)) ||
-            (procHasInitMessage && Value.Check(procedure.init, message.payload))) {
-            incoming.push(message);
+        };
+        this.streamMap.set(message.streamId, procStream);
+        // add this stream to ones from that client so we can clean it up in the case of a disconnect without close
+        const streamsFromThisClient = this.clientStreams.get(message.from) ?? new Set();
+        streamsFromThisClient.add(message.streamId);
+        this.clientStreams.set(message.from, streamsFromThisClient);
+        return procStream;
+    }
+    async pushToStream(procStream, message, isInit) {
+        const procedure = procStream.procedure;
+        const procHasInitMessage = 'init' in procedure;
+        if ((isInit &&
+            procHasInitMessage &&
+            Value.Check(procedure.init, message.payload)) ||
+            Value.Check(procedure.input, message.payload)) {
+            procStream.incoming.push(message.payload);
         }
         else if (!Value.Check(ControlMessagePayloadSchema, message.payload)) {
-            log?.error(`${transport.clientId} -- procedure ${message.serviceName}.${message.procedureName} received invalid payload: ${JSON.stringify(message.payload)}`);
+            log?.error(`${this.transport.clientId} -- procedure ${procStream.serviceName}.${procStream.procedureName} received invalid payload: ${JSON.stringify(message.payload)}`);
         }
         if (isStreamClose(message.controlFlags)) {
-            await cleanupStream(streamIdx);
-        }
-    };
-    transport.addEventListener('message', handler);
-    return {
-        services,
-        streams: streamMap,
-        async close() {
-            transport.removeEventListener('message', handler);
-            for (const streamIdx of streamMap.keys()) {
-                await cleanupStream(streamIdx);
+            await this.cleanupStream(message.streamId);
+            const streamsFromThisClient = this.clientStreams.get(message.from);
+            if (streamsFromThisClient) {
+                streamsFromThisClient.delete(message.streamId);
+                if (streamsFromThisClient.size === 0) {
+                    this.clientStreams.delete(message.from);
+                }
             }
-        },
+        }
+    }
+    getContext(service) {
+        const context = this.contextMap.get(service);
+        if (!context) {
+            const err = `${this.transport.clientId} -- no context found for ${service.name}`;
+            log?.error(err);
+            throw new Error(err);
+        }
+        return context;
+    }
+    cleanupStream = async (id) => {
+        const stream = this.streamMap.get(id);
+        if (!stream) {
+            return;
+        }
+        stream.incoming.end();
+        await stream.promises.inputHandler;
+        stream.outgoing.end();
+        await stream.promises.outputHandler;
+        this.streamMap.delete(id);
     };
 }
+/**
+ * Creates a server instance that listens for incoming messages from a transport and routes them to the appropriate service and procedure.
+ * The server tracks the state of each service along with open streams and the extended context object.
+ * @param transport - The transport to listen to.
+ * @param services - An object containing all the services to be registered on the server.
+ * @param extendedContext - An optional object containing additional context to be passed to all services.
+ * @returns A promise that resolves to a server instance with the registered services.
+ */
+export function createServer(transport, services, extendedContext) {
+    return new RiverServer(transport, services, extendedContext);
+}

package/dist/transport/impls/stdio/stdio.test.js

@@ -1,8 +1,7 @@
 import { describe, test, expect } from 'vitest';
 import stream from 'node:stream';
 import { StdioTransport } from './stdio';
-import { waitForMessage } from '../..';
-import { payloadToTransportMessage } from '../../../util/testHelpers';
+import { payloadToTransportMessage, waitForMessage, } from '../../../util/testHelpers';
 import { ensureTransportIsClean } from '../../../__tests__/fixtures/cleanup';
 describe('sending and receiving across node streams works', () => {
     test('basic send/receive', async () => {

package/dist/transport/impls/ws/client.d.ts

@@ -27,6 +27,7 @@ export declare class WebSocketClientTransport extends Transport<WebSocketConnect
     options: Options;
     serverId: TransportClientId;
     reconnectPromises: Map<TransportClientId, Promise<WebSocketResult>>;
+    tryReconnecting: boolean;
     /**
      * Creates a new WebSocketTransport instance.
      * @param wsGetter A function that returns a Promise that resolves to a WebSocket instance.
@@ -36,10 +37,6 @@ export declare class WebSocketClientTransport extends Transport<WebSocketConnect
     constructor(wsGetter: () => Promise<WebSocket>, clientId: TransportClientId, serverId: TransportClientId, providedOptions?: Partial<Options>);
     setupConnectionStatusListeners(): void;
     createNewConnection(to: string, attempt?: number): Promise<void>;
-    /**
-     * Begins a new attempt to establish a WebSocket connection.
-     */
-    open(): Promise<void>;
 }
 export {};
 //# sourceMappingURL=client.d.ts.map

package/dist/transport/impls/ws/client.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"client.d.ts","sourceRoot":"","sources":["../../../../transport/impls/ws/client.ts"],"names":[],"mappings":";AAAA,OAAO,SAAS,MAAM,eAAe,CAAC;AACtC,OAAO,EAAE,SAAS,EAAE,MAAM,iBAAiB,CAAC;AAE5C,OAAO,EAAE,iBAAiB,EAAE,MAAM,eAAe,CAAC;AAElD,OAAO,EAAE,KAAK,KAAK,EAAE,MAAM,gBAAgB,CAAC;AAC5C,OAAO,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAEnD,UAAU,OAAO;IACf,eAAe,EAAE,MAAM,CAAC;IACxB,gBAAgB,EAAE,MAAM,CAAC;IACzB,KAAK,EAAE,KAAK,CAAC;CACd;AAQD,KAAK,eAAe,GAAG;IAAE,EAAE,EAAE,SAAS,CAAA;CAAE,GAAG;IAAE,GAAG,EAAE,MAAM,CAAA;CAAE,CAAC;AAE3D;;;;GAIG;AACH,qBAAa,wBAAyB,SAAQ,SAAS,CAAC,mBAAmB,CAAC;IAC1E;;OAEG;IACH,QAAQ,EAAE,CAAC,EAAE,EAAE,iBAAiB,KAAK,OAAO,CAAC,SAAS,CAAC,CAAC;IACxD,OAAO,EAAE,OAAO,CAAC;IACjB,QAAQ,EAAE,iBAAiB,CAAC;IAE5B,iBAAiB,EAAE,GAAG,CAAC,iBAAiB,EAAE,OAAO,CAAC,eAAe,CAAC,CAAC,CAAC;IAEpE;;;;;OAKG;gBAED,QAAQ,EAAE,MAAM,OAAO,CAAC,SAAS,CAAC,EAClC,QAAQ,EAAE,iBAAiB,EAC3B,QAAQ,EAAE,iBAAiB,EAC3B,eAAe,CAAC,EAAE,OAAO,CAAC,OAAO,CAAC;IAWpC,8BAA8B,IAAI,IAAI;IAIhC,mBAAmB,CAAC,EAAE,EAAE,MAAM,EAAE,OAAO,SAAI;IAuEjD;;OAEG;IACG,IAAI;CAGX"}
+{"version":3,"file":"client.d.ts","sourceRoot":"","sources":["../../../../transport/impls/ws/client.ts"],"names":[],"mappings":";AAAA,OAAO,SAAS,MAAM,eAAe,CAAC;AACtC,OAAO,EAAE,SAAS,EAAE,MAAM,iBAAiB,CAAC;AAE5C,OAAO,EAAE,iBAAiB,EAAE,MAAM,eAAe,CAAC;AAElD,OAAO,EAAE,KAAK,KAAK,EAAE,MAAM,gBAAgB,CAAC;AAC5C,OAAO,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAEnD,UAAU,OAAO;IACf,eAAe,EAAE,MAAM,CAAC;IACxB,gBAAgB,EAAE,MAAM,CAAC;IACzB,KAAK,EAAE,KAAK,CAAC;CACd;AAQD,KAAK,eAAe,GAAG;IAAE,EAAE,EAAE,SAAS,CAAA;CAAE,GAAG;IAAE,GAAG,EAAE,MAAM,CAAA;CAAE,CAAC;AAE3D;;;;GAIG;AACH,qBAAa,wBAAyB,SAAQ,SAAS,CAAC,mBAAmB,CAAC;IAC1E;;OAEG;IACH,QAAQ,EAAE,CAAC,EAAE,EAAE,iBAAiB,KAAK,OAAO,CAAC,SAAS,CAAC,CAAC;IACxD,OAAO,EAAE,OAAO,CAAC;IACjB,QAAQ,EAAE,iBAAiB,CAAC;IAC5B,iBAAiB,EAAE,GAAG,CAAC,iBAAiB,EAAE,OAAO,CAAC,eAAe,CAAC,CAAC,CAAC;IACpE,eAAe,EAAE,OAAO,CAAQ;IAEhC;;;;;OAKG;gBAED,QAAQ,EAAE,MAAM,OAAO,CAAC,SAAS,CAAC,EAClC,QAAQ,EAAE,iBAAiB,EAC3B,QAAQ,EAAE,iBAAiB,EAC3B,eAAe,CAAC,EAAE,OAAO,CAAC,OAAO,CAAC;IAWpC,8BAA8B,IAAI,IAAI;IAIhC,mBAAmB,CAAC,EAAE,EAAE,MAAM,EAAE,OAAO,SAAI;CA6ElD"}

package/dist/transport/impls/ws/client.js

@@ -20,6 +20,7 @@ export class WebSocketClientTransport extends Transport {
     options;
     serverId;
     reconnectPromises;
+    tryReconnecting = true;
     /**
      * Creates a new WebSocketTransport instance.
      * @param wsGetter A function that returns a Promise that resolves to a WebSocket instance.
@@ -44,6 +45,10 @@ export class WebSocketClientTransport extends Transport {
         }
         let reconnectPromise = this.reconnectPromises.get(to);
         if (!reconnectPromise) {
+            if (!this.tryReconnecting) {
+                log?.info(`${this.clientId} -- tryReconnecting is false, not attempting reconnect`);
+                return;
+            }
             reconnectPromise = new Promise(async (resolve) => {
                 log?.info(`${this.clientId} -- establishing a new websocket to ${to}`);
                 const ws = await this.wsGetter(to);
@@ -93,10 +98,4 @@ export class WebSocketClientTransport extends Transport {
             setTimeout(() => this.createNewConnection(to, attempt + 1), this.options.retryIntervalMs * attempt);
         }
     }
-    /**
-     * Begins a new attempt to establish a WebSocket connection.
-     */
-    async open() {
-        return this.createNewConnection(this.serverId);
-    }
 }

package/dist/transport/impls/ws/server.d.ts

@@ -2,6 +2,7 @@ import { Codec } from '../../../codec';
 import { TransportClientId } from '../../message';
 import { Transport } from '../../transport';
 import { Server } from 'ws';
+import { WebSocket } from 'isomorphic-ws';
 import { WebSocketConnection } from './connection';
 interface Options {
     codec: Codec;
@@ -10,8 +11,10 @@ export declare class WebSocketServerTransport extends Transport<WebSocketConnect
     wss: Server;
     clientId: TransportClientId;
     constructor(wss: Server, clientId: TransportClientId, providedOptions?: Partial<Options>);
+    connectionHandler: (ws: WebSocket) => void;
     setupConnectionStatusListeners(): void;
     createNewConnection(to: string): Promise<void>;
+    close(): Promise<void>;
 }
 export {};
 //# sourceMappingURL=server.d.ts.map

package/dist/transport/impls/ws/server.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"server.d.ts","sourceRoot":"","sources":["../../../../transport/impls/ws/server.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,KAAK,EAAkB,MAAM,gBAAgB,CAAC;AAEvD,OAAO,EAAE,iBAAiB,EAAE,MAAM,eAAe,CAAC;AAClD,OAAO,EAAE,SAAS,EAAE,MAAM,iBAAiB,CAAC;AAC5C,OAAO,EAAE,MAAM,EAAE,MAAM,IAAI,CAAC;AAC5B,OAAO,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAEnD,UAAU,OAAO;IACf,KAAK,EAAE,KAAK,CAAC;CACd;AAMD,qBAAa,wBAAyB,SAAQ,SAAS,CAAC,mBAAmB,CAAC;IAC1E,GAAG,EAAE,MAAM,CAAC;IACZ,QAAQ,EAAE,iBAAiB,CAAC;gBAG1B,GAAG,EAAE,MAAM,EACX,QAAQ,EAAE,iBAAiB,EAC3B,eAAe,CAAC,EAAE,OAAO,CAAC,OAAO,CAAC;IASpC,8BAA8B,IAAI,IAAI;IAiChC,mBAAmB,CAAC,EAAE,EAAE,MAAM;CAKrC"}
+{"version":3,"file":"server.d.ts","sourceRoot":"","sources":["../../../../transport/impls/ws/server.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,KAAK,EAAkB,MAAM,gBAAgB,CAAC;AAEvD,OAAO,EAAE,iBAAiB,EAAE,MAAM,eAAe,CAAC;AAClD,OAAO,EAAE,SAAS,EAAE,MAAM,iBAAiB,CAAC;AAC5C,OAAO,EAAE,MAAM,EAAE,MAAM,IAAI,CAAC;AAC5B,OAAO,EAAE,SAAS,EAAE,MAAM,eAAe,CAAC;AAC1C,OAAO,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAEnD,UAAU,OAAO;IACf,KAAK,EAAE,KAAK,CAAC;CACd;AAMD,qBAAa,wBAAyB,SAAQ,SAAS,CAAC,mBAAmB,CAAC;IAC1E,GAAG,EAAE,MAAM,CAAC;IACZ,QAAQ,EAAE,iBAAiB,CAAC;gBAG1B,GAAG,EAAE,MAAM,EACX,QAAQ,EAAE,iBAAiB,EAC3B,eAAe,CAAC,EAAE,OAAO,CAAC,OAAO,CAAC;IASpC,iBAAiB,OAAQ,SAAS,UA6BhC;IAEF,8BAA8B,IAAI,IAAI;IAIhC,mBAAmB,CAAC,EAAE,EAAE,MAAM;IAM9B,KAAK;CAIZ"}

package/dist/transport/impls/ws/server.js

@@ -15,34 +15,39 @@ export class WebSocketServerTransport extends Transport {
         this.clientId = clientId;
         this.setupConnectionStatusListeners();
     }
+    connectionHandler = (ws) => {
+        let conn = undefined;
+        ws.onmessage = (msg) => {
+            // when we establish WebSocketConnection, ws.onmessage
+            // gets overriden so this only runs on the first valid message
+            // the websocket receives
+            const parsedMsg = this.parseMsg(msg.data);
+            if (parsedMsg && !conn) {
+                conn = new WebSocketConnection(this, parsedMsg.from, ws);
+                this.onConnect(conn);
+                this.handleMsg(parsedMsg);
+            }
+        };
+        // close is always emitted, even on error, ok to do cleanup here
+        ws.onclose = () => {
+            if (conn) {
+                this.onDisconnect(conn);
+            }
+        };
+        ws.onerror = (msg) => {
+            log?.warn(`${this.clientId} -- ws error from client ${conn?.connectedTo ?? 'unknown'}: ${msg}`);
+        };
+    };
     setupConnectionStatusListeners() {
-        this.wss.on('connection', (ws) => {
-            let conn = undefined;
-            ws.onmessage = (msg) => {
-                // when we establish WebSocketConnection, ws.onmessage
-                // gets overriden so this only runs on the first valid message
-                // the websocket receives
-                const parsedMsg = this.parseMsg(msg.data);
-                if (parsedMsg && !conn) {
-                    conn = new WebSocketConnection(this, parsedMsg.from, ws);
-                    this.onConnect(conn);
-                    this.handleMsg(parsedMsg);
-                }
-            };
-            // close is always emitted, even on error, ok to do cleanup here
-            ws.onclose = () => {
-                if (conn) {
-                    this.onDisconnect(conn);
-                }
-            };
-            ws.onerror = (msg) => {
-                log?.warn(`${this.clientId} -- ws error from client ${conn?.connectedTo ?? 'unknown'}: ${msg}`);
-            };
-        });
+        this.wss.on('connection', this.connectionHandler);
     }
     async createNewConnection(to) {
         const err = `${this.clientId} -- failed to send msg to ${to}, client probably dropped`;
         log?.warn(err);
         return;
     }
+    async close() {
+        super.close();
+        this.wss.off('connection', this.connectionHandler);
+    }
 }

package/dist/transport/impls/ws/ws.test.js

@@ -1,10 +1,10 @@
 import http from 'http';
-import { describe, test, expect, afterAll } from 'vitest';
-import { createWebSocketServer, createWsTransports, createDummyTransportMessage, onServerReady, createLocalWebSocketClient, } from '../../../util/testHelpers';
-import { msg, waitForMessage } from '../..';
+import { describe, test, expect, afterAll, vi } from 'vitest';
+import { createWebSocketServer, createWsTransports, createDummyTransportMessage, onServerReady, createLocalWebSocketClient, waitForMessage, } from '../../../util/testHelpers';
+import { msg } from '../..';
 import { WebSocketServerTransport } from './server';
 import { WebSocketClientTransport } from './client';
-import { testFinishesCleanly } from '../../../__tests__/fixtures/cleanup';
+import { testFinishesCleanly, waitFor, } from '../../../__tests__/fixtures/cleanup';
 describe('sending and receiving across websockets works', async () => {
     const server = http.createServer();
     const port = await onServerReady(server);
@@ -16,8 +16,9 @@ describe('sending and receiving across websockets works', async () => {
     test('basic send/receive', async () => {
         const [clientTransport, serverTransport] = createWsTransports(port, wss);
         const msg = createDummyTransportMessage();
+        const msgPromise = waitForMessage(serverTransport, (recv) => recv.id === msg.id);
         clientTransport.send(msg);
-        await expect(waitForMessage(serverTransport, (recv) => recv.id === msg.id)).resolves.toStrictEqual(msg.payload);
+        await expect(msgPromise).resolves.toStrictEqual(msg.payload);
         await testFinishesCleanly({
             clientTransports: [clientTransport],
             serverTransport,
@@ -36,8 +37,9 @@ describe('sending and receiving across websockets works', async () => {
         const initClient = async (id) => {
             const client = new WebSocketClientTransport(() => createLocalWebSocketClient(port), id, 'SERVER');
             const initMsg = makeDummyMessage(id, serverId, 'hello server');
+            const initMsgPromise = waitForMessage(serverTransport, (recv) => recv.id === initMsg.id);
             client.send(initMsg);
-            await expect(waitForMessage(serverTransport, (recv) => recv.id === initMsg.id)).resolves.toStrictEqual(initMsg.payload);
+            await expect(initMsgPromise).resolves.toStrictEqual(initMsg.payload);
             return client;
         };
         const client1 = await initClient(clientId1);
@@ -70,43 +72,109 @@ describe('retry logic', async () => {
     // TODO: right now, we only test client-side disconnects, we probably
     // need to also write tests for server-side crashes (but this involves clearing/restoring state)
     // not going to worry about this rn but for future
-    test('ws transport is recreated after clean disconnect', async () => {
+    test('ws connection is recreated after clean disconnect', async () => {
         const [clientTransport, serverTransport] = createWsTransports(port, wss);
         const msg1 = createDummyTransportMessage();
         const msg2 = createDummyTransportMessage();
+        const msg1Promise = waitForMessage(serverTransport, (recv) => recv.id === msg1.id);
         clientTransport.send(msg1);
-        await expect(waitForMessage(serverTransport, (recv) => recv.id === msg1.id)).resolves.toStrictEqual(msg1.payload);
+        await expect(msg1Promise).resolves.toStrictEqual(msg1.payload);
+        // clean disconnect
         clientTransport.connections.forEach((conn) => conn.ws.close());
-        clientTransport.send(msg2);
+        const msg2Promise = waitForMessage(serverTransport, (recv) => recv.id === msg2.id);
         // by this point the client should have reconnected
-        await expect(waitForMessage(serverTransport, (recv) => recv.id === msg2.id)).resolves.toStrictEqual(msg2.payload);
+        clientTransport.send(msg2);
+        await expect(msg2Promise).resolves.toStrictEqual(msg2.payload);
         await testFinishesCleanly({
             clientTransports: [clientTransport],
             serverTransport,
         });
     });
-    test('ws transport is recreated after unclean disconnect', async () => {
+    test('ws connection is recreated after unclean disconnect', async () => {
         const [clientTransport, serverTransport] = createWsTransports(port, wss);
         const msg1 = createDummyTransportMessage();
         const msg2 = createDummyTransportMessage();
+        const msg1Promise = waitForMessage(serverTransport, (recv) => recv.id === msg1.id);
         clientTransport.send(msg1);
-        await expect(waitForMessage(serverTransport, (recv) => recv.id === msg1.id)).resolves.toStrictEqual(msg1.payload);
+        await expect(msg1Promise).resolves.toStrictEqual(msg1.payload);
+        // unclean disconnect
         clientTransport.connections.forEach((conn) => conn.ws.terminate());
+        const msg2Promise = waitForMessage(serverTransport, (recv) => recv.id === msg2.id);
+        // by this point the client should have reconnected
         clientTransport.send(msg2);
+        await expect(msg2Promise).resolves.toStrictEqual(msg2.payload);
+        await testFinishesCleanly({
+            clientTransports: [clientTransport],
+            serverTransport,
+        });
+    });
+    test('both client and server transport gets connection and disconnection notifs', async () => {
+        const [clientTransport, serverTransport] = createWsTransports(port, wss);
+        const msg1 = createDummyTransportMessage();
+        const msg2 = createDummyTransportMessage();
+        const onClientConnect = vi.fn();
+        const onClientDisconnect = vi.fn();
+        const clientHandler = (evt) => {
+            if (evt.conn.connectedTo !== serverTransport.clientId)
+                return;
+            if (evt.status === 'connect')
+                return onClientConnect();
+            if (evt.status === 'disconnect')
+                return onClientDisconnect();
+        };
+        const onServerConnect = vi.fn();
+        const onServerDisconnect = vi.fn();
+        const serverHandler = (evt) => {
+            if (evt.status === 'connect' &&
+                evt.conn.connectedTo === clientTransport.clientId)
+                return onServerConnect();
+            if (evt.status === 'disconnect' &&
+                evt.conn.connectedTo === clientTransport.clientId)
+                return onServerDisconnect();
+        };
+        clientTransport.addEventListener('connectionStatus', clientHandler);
+        serverTransport.addEventListener('connectionStatus', serverHandler);
+        expect(onClientConnect).toHaveBeenCalledTimes(0);
+        expect(onClientDisconnect).toHaveBeenCalledTimes(0);
+        expect(onServerConnect).toHaveBeenCalledTimes(0);
+        expect(onServerDisconnect).toHaveBeenCalledTimes(0);
+        const msg1Promise = waitForMessage(serverTransport, (recv) => recv.id === msg1.id);
+        clientTransport.send(msg1);
+        await expect(msg1Promise).resolves.toStrictEqual(msg1.payload);
+        expect(onClientConnect).toHaveBeenCalledTimes(1);
+        expect(onClientDisconnect).toHaveBeenCalledTimes(0);
+        expect(onServerConnect).toHaveBeenCalledTimes(1);
+        expect(onServerDisconnect).toHaveBeenCalledTimes(0);
+        // clean disconnect
+        clientTransport.connections.forEach((conn) => conn.ws.close());
+        // wait for connection status to propagate to server
+        await waitFor(() => expect(onClientConnect).toHaveBeenCalledTimes(1));
+        await waitFor(() => expect(onClientDisconnect).toHaveBeenCalledTimes(1));
+        await waitFor(() => expect(onServerConnect).toHaveBeenCalledTimes(1));
+        await waitFor(() => expect(onServerDisconnect).toHaveBeenCalledTimes(1));
+        const msg2Promise = waitForMessage(serverTransport, (recv) => recv.id === msg2.id);
         // by this point the client should have reconnected
-        await expect(waitForMessage(serverTransport, (recv) => recv.id === msg2.id)).resolves.toStrictEqual(msg2.payload);
-        // this is not expected to be clean because we destroyed the transport
+        clientTransport.send(msg2);
+        await expect(msg2Promise).resolves.toStrictEqual(msg2.payload);
+        expect(onClientConnect).toHaveBeenCalledTimes(2);
+        expect(onClientDisconnect).toHaveBeenCalledTimes(1);
+        expect(onServerConnect).toHaveBeenCalledTimes(2);
+        expect(onServerDisconnect).toHaveBeenCalledTimes(1);
+        // teardown
+        clientTransport.removeEventListener('connectionStatus', clientHandler);
+        serverTransport.removeEventListener('connectionStatus', serverHandler);
         await testFinishesCleanly({
             clientTransports: [clientTransport],
             serverTransport,
         });
     });
-    test('ws transport is not recreated after destroy', async () => {
+    test('ws connection is not recreated after destroy', async () => {
         const [clientTransport, serverTransport] = createWsTransports(port, wss);
         const msg1 = createDummyTransportMessage();
         const msg2 = createDummyTransportMessage();
+        const promise1 = waitForMessage(serverTransport, (recv) => recv.id === msg1.id);
         clientTransport.send(msg1);
-        await expect(waitForMessage(serverTransport, (recv) => recv.id === msg1.id)).resolves.toStrictEqual(msg1.payload);
+        await expect(promise1).resolves.toStrictEqual(msg1.payload);
         clientTransport.destroy();
         expect(() => clientTransport.send(msg2)).toThrow(new Error('transport is destroyed, cant send'));
         // this is not expected to be clean because we destroyed the transport

package/dist/transport/index.d.ts

@@ -1,13 +1,4 @@
-import { OpaqueTransportMessage } from './message';
-import { Transport, Connection } from './transport';
 export { Transport, Connection } from './transport';
 export { TransportMessageSchema, OpaqueTransportMessageSchema, msg, reply, } from './message';
 export type { TransportMessage, MessageId, OpaqueTransportMessage, TransportClientId, isStreamOpen, isStreamClose, } from './message';
-/**
- * Waits for a message from the transport.
- * @param {Transport} t - The transport to listen to.
- * @param filter - An optional filter function to apply to the received messages.
- * @returns A promise that resolves with the payload of the first message that passes the filter.
- */
-export declare function waitForMessage(t: Transport<Connection>, filter?: (msg: OpaqueTransportMessage) => boolean, rejectMismatch?: boolean): Promise<unknown>;
 //# sourceMappingURL=index.d.ts.map