client transport .hardDisconnect

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@replit/river",
   "description": "It's like tRPC but... with JSON Schema Support, duplex streaming and support for service multiplexing. Transport agnostic!",
-  "version": "0.202.0",
+  "version": "0.203.0",
   "type": "module",
   "exports": {
     ".": {

transport/client.ts

@@ -365,6 +365,17 @@ export abstract class ClientTransport<
     this.updateSession(backingOffSession);
   }
 
+  /**
+   * Manually kills all sessions to the server (including all pending state).
+   * This is useful for when you want to close all connections to a server
+   * and don't want to wait for the grace period to elapse.
+   */
+  hardDisconnect() {
+    for (const session of this.sessions.values()) {
+      this.deleteSession(session);
+    }
+  }
+
   protected onBackoffFinished(session: SessionBackingOff) {
     const connPromise = tracer.startActiveSpan('connect', async (span) => {
       try {

transport/sessionStateMachine/common.ts

@@ -43,6 +43,14 @@ abstract class StateMachineState {
   // by consumers, the proxy will call this when .close is closed
   abstract _handleClose(): void;
 
+  /**
+   * Cleanup this state machine state and mark it as consumed.
+   * After calling close, it is an error to access any properties on the state.
+   * You should never need to call this as a consumer.
+   *
+   * If you're looking to close the session from the client,
+   * use `.hardDisconnect` on the client transport.
+   */
   close(): void {
     this._handleClose();
   }

transport/transport.test.ts

@@ -818,6 +818,43 @@ describe.each(testMatrix())(
       });
     });
 
+    test('client transport calling .hardDisconnect() immediately kills the session and updates bookkeeping', async () => {
+      const clientTransport = testHelpers.getClientTransport('client');
+      const serverTransport = testHelpers.getServerTransport();
+      clientTransport.connect(serverTransport.clientId);
+
+      addPostTestCleanup(async () => {
+        await cleanupTransports([clientTransport, serverTransport]);
+      });
+
+      await waitFor(() => {
+        expect(numberOfConnections(clientTransport)).toBe(1);
+        expect(numberOfConnections(serverTransport)).toBe(1);
+      });
+
+      const oldClientSessionId = serverTransport.sessions.get('client')?.id;
+      const oldServerSessionId = clientTransport.sessions.get('SERVER')?.id;
+      expect(oldClientSessionId).not.toBeUndefined();
+      expect(oldServerSessionId).not.toBeUndefined();
+
+      clientTransport.hardDisconnect();
+
+      expect(numberOfConnections(clientTransport)).toBe(0);
+      expect(clientTransport.sessions.size).toBe(0);
+
+      await advanceFakeTimersByDisconnectGrace();
+      await advanceFakeTimersBySessionGrace();
+      await waitFor(() => {
+        expect(numberOfConnections(serverTransport)).toBe(0);
+        expect(serverTransport.sessions.size).toBe(0);
+      });
+
+      await testFinishesCleanly({
+        clientTransports: [clientTransport],
+        serverTransport,
+      });
+    });
+
     // make a custom auth thing that rejects all connections
     // session grace should elapse at some point despite retry loop
     test('session grace elapses during long reconnect loop', async () => {