fix sessionStatus type, remove no connection -> no connection transition on server, extra invariant checks

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.205.1",
+  "version": "0.205.2",
   "type": "module",
   "exports": {
     ".": {

testUtil/fixtures/mockTransport.ts

@@ -15,14 +15,15 @@ export class InMemoryConnection extends Connection {
   constructor(pipe: Duplex) {
     super();
     this.conn = pipe;
+    this.conn.allowHalfOpen = false;
 
     this.conn.on('data', (data: Uint8Array) => {
       for (const cb of this.dataListeners) {
         cb(data);
       }
     });
 
-    this.conn.on('end', () => {
+    this.conn.on('close', () => {
       for (const cb of this.closeListeners) {
         cb();
       }
@@ -46,6 +47,7 @@ export class InMemoryConnection extends Connection {
   close(): void {
     setImmediate(() => {
       this.conn.end();
+      this.conn.emit('close');
     });
   }
 }
@@ -153,6 +155,7 @@ export function createMockTransportNetwork(
     simulatePhantomDisconnect() {
       for (const conn of Object.values(connections.get())) {
         conn.serverToClient.pause();
+        conn.clientToServer.pause();
       }
     },
     async restartServer() {

transport/events.ts

@@ -2,6 +2,7 @@ import { type Static } from '@sinclair/typebox';
 import { Connection } from './connection';
 import { OpaqueTransportMessage, HandshakeErrorResponseCodes } from './message';
 import { Session, SessionState } from './sessionStateMachine';
+import { SessionId } from './sessionStateMachine/common';
 import { TransportStatus } from './transport';
 
 export const ProtocolError = {
@@ -26,11 +27,11 @@ export interface EventMap {
         session: Pick<Session<Connection>, 'id' | 'to'>;
       };
   sessionTransition:
-    | { state: SessionState.Connected }
-    | { state: SessionState.Handshaking }
-    | { state: SessionState.Connecting }
-    | { state: SessionState.BackingOff }
-    | { state: SessionState.NoConnection };
+    | { state: SessionState.Connected; id: SessionId }
+    | { state: SessionState.Handshaking; id: SessionId }
+    | { state: SessionState.Connecting; id: SessionId }
+    | { state: SessionState.BackingOff; id: SessionId }
+    | { state: SessionState.NoConnection; id: SessionId };
   protocolError:
     | {
         type: (typeof ProtocolError)['HandshakeFailed'];

transport/server.ts

@@ -395,9 +395,8 @@ export abstract class ServerTransport<
           );
 
         oldSession = noConnectionSession;
+        this.updateSession(oldSession);
       }
-
-      this.updateSession(oldSession);
     } else if (oldSession) {
       connectCase = 'hard reconnection';
 

transport/sessionStateMachine/SessionConnected.ts

@@ -4,6 +4,7 @@ import {
   ControlMessageAckSchema,
   OpaqueTransportMessage,
   PartialTransportMessage,
+  TransportMessage,
   isAck,
 } from '../message';
 import {
@@ -48,10 +49,25 @@ export class SessionConnected<
     this.heartbeatMisses = 0;
   }
 
+  private assertSendOrdering(constructedMsg: TransportMessage) {
+    if (constructedMsg.seq > this.seqSent + 1) {
+      const msg = `invariant violation: would have sent out of order msg (seq: ${constructedMsg.seq}, expected: ${this.seqSent} + 1)`;
+      this.log?.error(msg, {
+        ...this.loggingMetadata,
+        transportMessage: constructedMsg,
+        tags: ['invariant-violation'],
+      });
+
+      throw new Error(msg);
+    }
+  }
+
   send(msg: PartialTransportMessage): string {
     const constructedMsg = this.constructMsg(msg);
+    this.assertSendOrdering(constructedMsg);
     this.sendBuffer.push(constructedMsg);
     this.conn.send(this.options.codec.toBuffer(constructedMsg));
+    this.seqSent = constructedMsg.seq;
 
     return constructedMsg.id;
   }
@@ -75,7 +91,9 @@ export class SessionConnected<
       );
 
       for (const msg of this.sendBuffer) {
+        this.assertSendOrdering(msg);
         this.conn.send(this.options.codec.toBuffer(msg));
+        this.seqSent = msg.seq;
       }
     }
 
@@ -165,7 +183,7 @@ export class SessionConnected<
         );
       } else {
         const reason = `received out-of-order msg, closing connection (got seq: ${parsedMsg.seq}, wanted seq: ${this.ack})`;
-        this.log?.warn(reason, {
+        this.log?.error(reason, {
           ...this.loggingMetadata,
           transportMessage: parsedMsg,
           tags: ['invariant-violation'],

transport/sessionStateMachine/common.ts

@@ -208,6 +208,7 @@ export interface IdentifiedSessionProps extends CommonSessionProps {
   to: TransportClientId;
   seq: number;
   ack: number;
+  seqSent: number;
   sendBuffer: Array<OpaqueTransportMessage>;
   telemetry: TelemetryInfo;
   protocolVersion: ProtocolVersion;
@@ -224,15 +225,29 @@ export abstract class IdentifiedSession extends CommonSession {
    */
   seq: number;
 
+  /**
+   * Last seq we sent over the wire this session (excluding handshake) and retransmissions
+   */
+  seqSent: number;
+
   /**
    * Number of unique messages we've received this session (excluding handshake)
    */
   ack: number;
   sendBuffer: Array<OpaqueTransportMessage>;
 
   constructor(props: IdentifiedSessionProps) {
-    const { id, to, seq, ack, sendBuffer, telemetry, log, protocolVersion } =
-      props;
+    const {
+      id,
+      to,
+      seq,
+      ack,
+      sendBuffer,
+      telemetry,
+      log,
+      protocolVersion,
+      seqSent: messagesSent,
+    } = props;
     super(props);
     this.id = id;
     this.to = to;
@@ -242,6 +257,7 @@ export abstract class IdentifiedSession extends CommonSession {
     this.telemetry = telemetry;
     this.log = log;
     this.protocolVersion = protocolVersion;
+    this.seqSent = messagesSent;
   }
 
   get loggingMetadata(): MessageMetadata {

transport/sessionStateMachine/stateMachine.test.ts

@@ -1942,16 +1942,19 @@ describe('session state machine', () => {
       expect(conn.send).toHaveBeenCalledTimes(0);
 
       // send a heartbeat
-      conn.emitData(
-        session.options.codec.toBuffer(
-          session.constructMsg({
-            streamId: 'heartbeat',
-            controlFlags: ControlFlags.AckBit,
-            payload: {
-              type: 'ACK',
-            } satisfies Static<typeof ControlMessageAckSchema>,
-          }),
-        ),
+      conn.onData(
+        session.options.codec.toBuffer({
+          id: 'msgid',
+          to: 'SERVER',
+          from: 'client',
+          seq: 0,
+          ack: 0,
+          streamId: 'heartbeat',
+          controlFlags: ControlFlags.AckBit,
+          payload: {
+            type: 'ACK',
+          } satisfies Static<typeof ControlMessageAckSchema>,
+        }),
       );
 
       // make sure the session acks the heartbeat
@@ -1964,16 +1967,19 @@ describe('session state machine', () => {
       const conn = session.conn;
 
       // send a heartbeat
-      conn.emitData(
-        session.options.codec.toBuffer(
-          session.constructMsg({
-            streamId: 'heartbeat',
-            controlFlags: ControlFlags.AckBit,
-            payload: {
-              type: 'ACK',
-            } satisfies Static<typeof ControlMessageAckSchema>,
-          }),
-        ),
+      conn.onData(
+        session.options.codec.toBuffer({
+          id: 'msgid',
+          to: 'SERVER',
+          from: 'client',
+          seq: 0,
+          ack: 0,
+          streamId: 'heartbeat',
+          controlFlags: ControlFlags.AckBit,
+          payload: {
+            type: 'ACK',
+          } satisfies Static<typeof ControlMessageAckSchema>,
+        }),
       );
 
       expect(sessionHandle.onMessage).not.toHaveBeenCalled();

transport/sessionStateMachine/transitions.ts

@@ -50,6 +50,7 @@ function inheritSharedSession(
     to: session.to,
     seq: session.seq,
     ack: session.ack,
+    seqSent: session.seqSent,
     sendBuffer: session.sendBuffer,
     telemetry: session.telemetry,
     options: session.options,
@@ -90,6 +91,7 @@ export const SessionStateGraph = {
         to,
         seq: 0,
         ack: 0,
+        seqSent: 0,
         graceExpiryTime: Date.now() + options.sessionDisconnectGraceMs,
         sendBuffer,
         telemetry,
@@ -251,12 +253,13 @@ export const SessionStateGraph = {
         ? // old session exists, inherit state
           inheritSharedSession(oldSession)
         : // old session does not exist, create new state
-          {
+          ({
             id: sessionId,
             from,
             to,
             seq: 0,
             ack: 0,
+            seqSent: 0,
             sendBuffer: [],
             telemetry: createSessionTelemetryInfo(
               pendingSession.tracer,
@@ -269,7 +272,7 @@ export const SessionStateGraph = {
             tracer: pendingSession.tracer,
             log: pendingSession.log,
             protocolVersion,
-          };
+          } satisfies IdentifiedSessionProps);
 
       pendingSession._handleStateExit();
       oldSession?._handleStateExit();