@dxos/edge-client 0.6.14-staging.e15392e → 0.7.0

package/dist/types/src/edge-ws-connection.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"edge-ws-connection.d.ts","sourceRoot":"","sources":["../../../src/edge-ws-connection.ts"],"names":[],"mappings":"AAOA,OAAO,EAAW,QAAQ,EAAE,MAAM,eAAe,CAAC;AAIlD,OAAO,EAAiB,KAAK,OAAO,EAAE,MAAM,4CAA4C,CAAC;AAGzF,OAAO,EAAE,KAAK,YAAY,EAAE,MAAM,iBAAiB,CAAC;AAKpD,MAAM,MAAM,yBAAyB,GAAG;IACtC,WAAW,EAAE,MAAM,IAAI,CAAC;IACxB,SAAS,EAAE,CAAC,OAAO,EAAE,OAAO,KAAK,IAAI,CAAC;IACtC,iBAAiB,EAAE,MAAM,IAAI,CAAC;CAC/B,CAAC;AAEF,qBAAa,gBAAiB,SAAQ,QAAQ;IAK1C,OAAO,CAAC,QAAQ,CAAC,SAAS;IAC1B,OAAO,CAAC,QAAQ,CAAC,eAAe;IAChC,OAAO,CAAC,QAAQ,CAAC,UAAU;IAN7B,OAAO,CAAC,qBAAqB,CAAsB;IACnD,OAAO,CAAC,GAAG,CAAwB;gBAGhB,SAAS,EAAE,YAAY,EACvB,eAAe,EAAE;QAAE,GAAG,EAAE,GAAG,CAAC;QAAC,cAAc,CAAC,EAAE,MAAM,CAAA;KAAE,EACtD,UAAU,EAAE,yBAAyB;IAKxD,IACW,IAAI;;;;MAMd;IAEM,IAAI,CAAC,OAAO,EAAE,OAAO;cAMH,KAAK;cAkDL,MAAM;IAc/B,OAAO,CAAC,mBAAmB;IAe3B,OAAO,CAAC,2BAA2B;CAgBpC"}

package/dist/types/src/persistent-lifecycle.d.ts

@@ -1,14 +1,14 @@
 import { Resource } from '@dxos/context';
-export type PersistentLifecycleParams = {
+export type PersistentLifecycleParams<T> = {
     /**
      * Create connection.
      * If promise resolves successfully, connection is considered established.
      */
-    start: () => Promise<void>;
+    start: () => Promise<T | undefined>;
     /**
      * Reset connection to initial state.
      */
-    stop: () => Promise<void>;
+    stop: (state: T) => Promise<void>;
     /**
      * Called after successful start.
      */
@@ -23,17 +23,19 @@ export type PersistentLifecycleParams = {
  * Handles restarts (e.g. persists connection).
  * Restarts are scheduled with exponential backoff.
  */
-export declare class PersistentLifecycle extends Resource {
+export declare class PersistentLifecycle<T> extends Resource {
     private readonly _start;
     private readonly _stop;
     private readonly _onRestart?;
     private readonly _maxRestartDelay;
+    private _currentContext;
     private _restartTask?;
     private _restartAfter;
-    constructor({ start, stop, onRestart, maxRestartDelay }: PersistentLifecycleParams);
+    constructor({ start, stop, onRestart, maxRestartDelay }: PersistentLifecycleParams<T>);
     protected _open(): Promise<void>;
     protected _close(): Promise<void>;
     private _restart;
+    private _stopCurrentContext;
     /**
      * Scheduling restart should be done from outside.
      */

package/dist/types/src/persistent-lifecycle.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"persistent-lifecycle.d.ts","sourceRoot":"","sources":["../../../src/persistent-lifecycle.ts"],"names":[],"mappings":"AAKA,OAAO,EAAqC,QAAQ,EAAE,MAAM,eAAe,CAAC;AAO5E,MAAM,MAAM,yBAAyB,GAAG;IACtC;;;OAGG;IACH,KAAK,EAAE,MAAM,OAAO,CAAC,IAAI,CAAC,CAAC;IAE3B;;OAEG;IACH,IAAI,EAAE,MAAM,OAAO,CAAC,IAAI,CAAC,CAAC;IAE1B;;OAEG;IACH,SAAS,CAAC,EAAE,MAAM,OAAO,CAAC,IAAI,CAAC,CAAC;IAEhC;;;OAGG;IACH,eAAe,CAAC,EAAE,MAAM,CAAC;CAC1B,CAAC;AAEF;;;GAGG;AACH,qBAAa,mBAAoB,SAAQ,QAAQ;IAC/C,OAAO,CAAC,QAAQ,CAAC,MAAM,CAAsB;IAC7C,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAsB;IAC5C,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAC,CAAsB;IAClD,OAAO,CAAC,QAAQ,CAAC,gBAAgB,CAAS;IAE1C,OAAO,CAAC,YAAY,CAAC,CAA2B;IAChD,OAAO,CAAC,aAAa,CAAK;gBAEd,EAAE,KAAK,EAAE,IAAI,EAAE,SAAS,EAAE,eAA2C,EAAE,EAAE,yBAAyB;cASrF,KAAK;cAeL,MAAM;YAMjB,QAAQ;IAgBtB;;OAEG;IAEH,eAAe;CAMhB"}
+{"version":3,"file":"persistent-lifecycle.d.ts","sourceRoot":"","sources":["../../../src/persistent-lifecycle.ts"],"names":[],"mappings":"AAKA,OAAO,EAAqC,QAAQ,EAAE,MAAM,eAAe,CAAC;AAO5E,MAAM,MAAM,yBAAyB,CAAC,CAAC,IAAI;IACzC;;;OAGG;IACH,KAAK,EAAE,MAAM,OAAO,CAAC,CAAC,GAAG,SAAS,CAAC,CAAC;IAEpC;;OAEG;IACH,IAAI,EAAE,CAAC,KAAK,EAAE,CAAC,KAAK,OAAO,CAAC,IAAI,CAAC,CAAC;IAElC;;OAEG;IACH,SAAS,CAAC,EAAE,MAAM,OAAO,CAAC,IAAI,CAAC,CAAC;IAEhC;;;OAGG;IACH,eAAe,CAAC,EAAE,MAAM,CAAC;CAC1B,CAAC;AAEF;;;GAGG;AACH,qBAAa,mBAAmB,CAAC,CAAC,CAAE,SAAQ,QAAQ;IAClD,OAAO,CAAC,QAAQ,CAAC,MAAM,CAA+B;IACtD,OAAO,CAAC,QAAQ,CAAC,KAAK,CAA8B;IACpD,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAC,CAAsB;IAClD,OAAO,CAAC,QAAQ,CAAC,gBAAgB,CAAS;IAE1C,OAAO,CAAC,eAAe,CAA4B;IACnD,OAAO,CAAC,YAAY,CAAC,CAA2B;IAChD,OAAO,CAAC,aAAa,CAAK;gBAEd,EAAE,KAAK,EAAE,IAAI,EAAE,SAAS,EAAE,eAA2C,EAAE,EAAE,yBAAyB,CAAC,CAAC,CAAC;cASxF,KAAK;cAgBL,MAAM;YAMjB,QAAQ;YAkBR,mBAAmB;IAWjC;;OAEG;IAEH,eAAe;CAMhB"}

package/dist/types/src/testing/test-utils.d.ts

@@ -10,6 +10,7 @@ type TestEdgeWsServerParams = {
 export declare const createTestEdgeWsServer: (port?: number, params?: TestEdgeWsServerParams) => Promise<{
     server: WebSocket.Server;
     messageSink: any[];
+    messageSourceLog: any[];
     endpoint: string;
     cleanup: () => void;
     currentConnection: () => WebSocket | undefined;

package/dist/types/src/testing/test-utils.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"test-utils.d.ts","sourceRoot":"","sources":["../../../../src/testing/test-utils.ts"],"names":[],"mappings":"AAIA,OAAO,SAAS,MAAM,eAAe,CAAC;AAEtC,OAAO,EAAE,OAAO,EAAE,MAAM,aAAa,CAAC;AAGtC,OAAO,EAAoC,KAAK,OAAO,EAAE,MAAM,4CAA4C,CAAC;AAK5G,eAAO,MAAM,YAAY,OAAO,CAAC;AAEjC,KAAK,sBAAsB,GAAG;IAC5B,eAAe,CAAC,EAAE,OAAO,CAAC;IAC1B,cAAc,CAAC,EAAE,CAAC,OAAO,EAAE,UAAU,KAAK,GAAG,CAAC;IAC9C,cAAc,CAAC,EAAE,CAAC,OAAO,EAAE,GAAG,KAAK,OAAO,CAAC,UAAU,GAAG,SAAS,CAAC,CAAC;CACpE,CAAC;AAEF,eAAO,MAAM,sBAAsB,2BAAwC,sBAAsB;;;;;;mCAmE9E,OAAO,mBAAmB,UAAU;uBA1BhC,OAAO;;EAS7B,CAAC"}
+{"version":3,"file":"test-utils.d.ts","sourceRoot":"","sources":["../../../../src/testing/test-utils.ts"],"names":[],"mappings":"AAIA,OAAO,SAAS,MAAM,eAAe,CAAC;AAEtC,OAAO,EAAE,OAAO,EAAE,MAAM,aAAa,CAAC;AAGtC,OAAO,EAAoC,KAAK,OAAO,EAAE,MAAM,4CAA4C,CAAC;AAK5G,eAAO,MAAM,YAAY,OAAO,CAAC;AAEjC,KAAK,sBAAsB,GAAG;IAC5B,eAAe,CAAC,EAAE,OAAO,CAAC;IAC1B,cAAc,CAAC,EAAE,CAAC,OAAO,EAAE,UAAU,KAAK,GAAG,CAAC;IAC9C,cAAc,CAAC,EAAE,CAAC,OAAO,EAAE,GAAG,KAAK,OAAO,CAAC,UAAU,GAAG,SAAS,CAAC,CAAC;CACpE,CAAC;AAEF,eAAO,MAAM,sBAAsB,2BAAwC,sBAAsB;;;;;;;mCAsE9E,OAAO,mBAAmB,UAAU;uBA1BhC,OAAO;;EAS7B,CAAC"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dxos/edge-client",
-  "version": "0.6.14-staging.e15392e",
+  "version": "0.7.0",
   "description": "EDGE Client",
   "homepage": "https://dxos.org",
   "bugs": "https://github.com/dxos/dxos/issues",
@@ -35,21 +35,21 @@
   "dependencies": {
     "isomorphic-ws": "^5.0.0",
     "ws": "^8.14.2",
-    "@dxos/async": "0.6.14-staging.e15392e",
-    "@dxos/credentials": "0.6.14-staging.e15392e",
-    "@dxos/context": "0.6.14-staging.e15392e",
-    "@dxos/crypto": "0.6.14-staging.e15392e",
-    "@dxos/debug": "0.6.14-staging.e15392e",
-    "@dxos/invariant": "0.6.14-staging.e15392e",
-    "@dxos/keyring": "0.6.14-staging.e15392e",
-    "@dxos/keys": "0.6.14-staging.e15392e",
-    "@dxos/log": "0.6.14-staging.e15392e",
-    "@dxos/node-std": "0.6.14-staging.e15392e",
-    "@dxos/protocols": "0.6.14-staging.e15392e",
-    "@dxos/util": "0.6.14-staging.e15392e"
+    "@dxos/async": "0.7.0",
+    "@dxos/context": "0.7.0",
+    "@dxos/credentials": "0.7.0",
+    "@dxos/debug": "0.7.0",
+    "@dxos/crypto": "0.7.0",
+    "@dxos/keyring": "0.7.0",
+    "@dxos/invariant": "0.7.0",
+    "@dxos/keys": "0.7.0",
+    "@dxos/node-std": "0.7.0",
+    "@dxos/protocols": "0.7.0",
+    "@dxos/util": "0.7.0",
+    "@dxos/log": "0.7.0"
   },
   "devDependencies": {
-    "@dxos/test-utils": "0.6.14-staging.e15392e"
+    "@dxos/test-utils": "0.7.0"
   },
   "publishConfig": {
     "access": "public"

package/src/auth.ts

@@ -4,6 +4,7 @@
 
 import { createCredential, signPresentation } from '@dxos/credentials';
 import { type Signer } from '@dxos/crypto';
+import { invariant } from '@dxos/invariant';
 import { Keyring } from '@dxos/keyring';
 import { PublicKey } from '@dxos/keys';
 import { type Chain, type Credential } from '@dxos/protocols/proto/dxos/halo/credentials';
@@ -47,7 +48,7 @@ export const createChainEdgeIdentity = async (
   signer: Signer,
   identityKey: PublicKey,
   peerKey: PublicKey,
-  chain: Chain,
+  chain: Chain | undefined,
   credentials: Credential[],
 ): Promise<EdgeIdentity> => {
   const credentialsToSign =
@@ -70,6 +71,8 @@ export const createChainEdgeIdentity = async (
     identityKey: identityKey.toHex(),
     peerKey: peerKey.toHex(),
     presentCredentials: async ({ challenge }) => {
+      // TODO: make chain required after device invitation flow update release
+      invariant(chain);
       return signPresentation({
         presentation: {
           credentials: credentialsToSign,

package/src/edge-client.test.ts

@@ -12,33 +12,33 @@ import { openAndClose } from '@dxos/test-utils';
 import { createEphemeralEdgeIdentity, createTestHaloEdgeIdentity } from './auth';
 import { protocol } from './defs';
 import { EdgeClient } from './edge-client';
+import { type EdgeIdentity } from './edge-identity';
+import { EdgeConnectionClosedError, EdgeIdentityChangedError } from './errors';
 import { createTestEdgeWsServer } from './testing';
 
 describe('EdgeClient', () => {
-  const textMessage = (message: string) => protocol.createMessage(TextMessageSchema, { payload: { message } });
+  let wsServerPort = 8001;
 
   test('reconnects on error', async () => {
-    const { closeConnection, endpoint, cleanup } = await createTestEdgeWsServer(8001);
+    const { closeConnection, endpoint, cleanup } = await createTestEdgeWsServer(wsServerPort++);
     onTestFinished(cleanup);
 
-    const client = new EdgeClient(await createEphemeralEdgeIdentity(), { socketEndpoint: endpoint });
-    await openAndClose(client);
+    const { client, reconnectTrigger } = await openNewClient(endpoint);
     await client.send(textMessage('Hello world 1'));
     expect(client.isOpen).is.true;
 
-    const reconnected = client.reconnect.waitForCount(1);
+    reconnectTrigger.reset();
     await closeConnection();
-    await reconnected;
+    await reconnectTrigger.wait();
     await expect(client.send(textMessage('Hello world 2'))).resolves.not.toThrow();
   });
 
   test('isConnected', async () => {
     const admitConnection = new Trigger();
-    const { closeConnection, endpoint, cleanup } = await createTestEdgeWsServer(8001, { admitConnection });
+    const { closeConnection, endpoint, cleanup } = await createTestEdgeWsServer(wsServerPort++, { admitConnection });
     onTestFinished(cleanup);
 
-    const client = new EdgeClient(await createEphemeralEdgeIdentity(), { socketEndpoint: endpoint });
-    await openAndClose(client);
+    const { client } = await openNewClient(endpoint);
 
     expect(client.isConnected).toBeFalsy();
     admitConnection.wake();
@@ -54,20 +54,90 @@ describe('EdgeClient', () => {
   });
 
   test('set identity reconnects', async () => {
-    const { endpoint, cleanup } = await createTestEdgeWsServer(8002);
+    const { endpoint, cleanup } = await createTestEdgeWsServer(wsServerPort++);
     onTestFinished(cleanup);
 
-    const client = new EdgeClient(await createEphemeralEdgeIdentity(), { socketEndpoint: endpoint });
-    await openAndClose(client);
+    const { client, reconnectTrigger } = await openNewClient(endpoint);
     await client.send(textMessage('Hello world 1'));
     expect(client.isOpen).is.true;
 
-    const reconnected = client.reconnect.waitForCount(1);
+    reconnectTrigger.reset();
     client.setIdentity(await createEphemeralEdgeIdentity());
-    await reconnected;
+    await reconnectTrigger.wait();
     await expect(client.send(textMessage('Hello world 2'))).resolves.not.toThrow();
   });
 
+  test('send blocks until connection becomes ready', async () => {
+    const admitConnection = new Trigger();
+    const { endpoint, messageSink, cleanup } = await createTestEdgeWsServer(wsServerPort++, { admitConnection });
+    onTestFinished(cleanup);
+
+    const { client } = await openNewClient(endpoint);
+    setTimeout(() => admitConnection.wake(), 20);
+    await client.send(textMessage('Hello world 1'));
+    await expect.poll(() => messageSink.length).toBe(1);
+  });
+
+  test('send fails if identity changes before connection becomes ready', async () => {
+    const admitConnection = new Trigger();
+    const { endpoint, cleanup, messageSink } = await createTestEdgeWsServer(wsServerPort++, { admitConnection });
+    onTestFinished(cleanup);
+
+    const { client } = await openNewClient(endpoint);
+    setTimeout(async () => client.setIdentity(await createEphemeralEdgeIdentity()));
+    await expect(client.send(textMessage('Hello world 1'))).rejects.toThrow(EdgeIdentityChangedError);
+
+    // Test recovers.
+    setTimeout(() => admitConnection.wake(), 20);
+    await client.send(textMessage('Hello world 1'));
+    await expect.poll(() => messageSink.length).toBe(1);
+  });
+
+  test('send fails if client is closed before connection becomes ready', async () => {
+    const admitConnection = new Trigger();
+    const { endpoint, cleanup } = await createTestEdgeWsServer(wsServerPort++, { admitConnection });
+    onTestFinished(cleanup);
+
+    const { client } = await openNewClient(endpoint);
+    setTimeout(() => client.close());
+    await expect(client.send(textMessage('Hello world 1'))).rejects.toThrow(EdgeConnectionClosedError);
+  });
+
+  test('onReconnect trigger', async () => {
+    const admitConnection = new Trigger();
+    const { endpoint, cleanup, closeConnection } = await createTestEdgeWsServer(wsServerPort++, { admitConnection });
+    onTestFinished(cleanup);
+
+    const { client } = await openNewClient(endpoint);
+    let callCount = 0;
+    client.onReconnected(() => callCount++);
+    admitConnection.wake();
+
+    await expect.poll(() => callCount).toEqual(1);
+    await closeConnection();
+    await expect.poll(() => callCount).toEqual(2);
+
+    const trigger = new Trigger();
+    client.onReconnected(() => trigger.wake());
+    await trigger.wait();
+    expect(callCount).toEqual(2);
+  });
+
+  test('send message right after identity change is delivered successfully with the new identity', async () => {
+    const { endpoint, cleanup, messageSourceLog } = await createTestEdgeWsServer(wsServerPort++);
+    onTestFinished(cleanup);
+
+    const { client, identity: oldIdentity } = await openNewClient(endpoint);
+    await client.send(textMessage('Hello world 1', oldIdentity));
+    expect(client.isOpen).is.true;
+
+    const newIdentity = await createEphemeralEdgeIdentity();
+    client.setIdentity(newIdentity);
+    await client.send(textMessage('Hello world 2', newIdentity));
+    await expect.poll(() => messageSourceLog.length).toBe(2);
+    expect(messageSourceLog.map((m) => m.peerKey)).toStrictEqual([oldIdentity.peerKey, newIdentity.peerKey]);
+  });
+
   test.skipIf(!process.env.EDGE_ENDPOINT)('connect to local edge server', async () => {
     // const identity = await createEphemeralEdgeIdentity();
 
@@ -79,4 +149,21 @@ describe('EdgeClient', () => {
     await client.send(textMessage('Hello world 1'));
     expect(client.isOpen).is.true;
   });
+
+  const textMessage = (message: string, source?: EdgeIdentity) =>
+    protocol.createMessage(TextMessageSchema, {
+      source: source && { peerKey: source.peerKey, identityKey: source.identityKey },
+      payload: { message },
+    });
+
+  const openNewClient = async (endpoint: string) => {
+    const identity = await createEphemeralEdgeIdentity();
+    const client = new EdgeClient(identity, { socketEndpoint: endpoint });
+    await openAndClose(client);
+    const reconnectTrigger = new Trigger();
+    client.onReconnected(() => {
+      reconnectTrigger.wake();
+    });
+    return { client, reconnectTrigger, identity };
+  };
 });

package/src/edge-client.ts

@@ -2,37 +2,33 @@
 // Copyright 2024 DXOS.org
 //
 
-import WebSocket from 'isomorphic-ws';
-
-import { Trigger, Event, scheduleTaskInterval, scheduleTask, TriggerState } from '@dxos/async';
-import { Context, LifecycleState, Resource, type Lifecycle } from '@dxos/context';
-import { log } from '@dxos/log';
-import { buf } from '@dxos/protocols/buf';
-import { type Message, MessageSchema } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
+import { Trigger, scheduleMicroTask, TriggerState } from '@dxos/async';
+import { Resource, type Lifecycle } from '@dxos/context';
+import { log, logInfo } from '@dxos/log';
+import { type Message } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
 
 import { protocol } from './defs';
 import { type EdgeIdentity, handleAuthChallenge } from './edge-identity';
+import { EdgeWsConnection } from './edge-ws-connection';
 import { EdgeConnectionClosedError, EdgeIdentityChangedError } from './errors';
 import { PersistentLifecycle } from './persistent-lifecycle';
-import { type Protocol, toUint8Array } from './protocol';
+import { type Protocol } from './protocol';
 import { getEdgeUrlWithProtocol } from './utils';
 
 const DEFAULT_TIMEOUT = 10_000;
-const SIGNAL_KEEPALIVE_INTERVAL = 5_000;
 
-export type MessageListener = (message: Message) => void | Promise<void>;
+export type MessageListener = (message: Message) => void;
+export type ReconnectListener = () => void;
 
 export interface EdgeConnection extends Required<Lifecycle> {
-  connected: Event;
-  reconnect: Event;
-
   get info(): any;
   get identityKey(): string;
   get peerKey(): string;
   get isOpen(): boolean;
   get isConnected(): boolean;
   setIdentity(identity: EdgeIdentity): void;
-  addListener(listener: MessageListener): () => void;
+  onMessage(listener: MessageListener): () => void;
+  onReconnected(listener: ReconnectListener): () => void;
   send(message: Message): Promise<void>;
 }
 
@@ -44,25 +40,25 @@ export type MessengerConfig = {
 };
 
 /**
- * Messenger client.
+ * Messenger client for EDGE:
+ *  - While open, uses PersistentLifecycle to keep an open EdgeWsConnection, reconnecting on failures.
+ *  - Manages identity and re-create EdgeWsConnection when identity changes.
+ *  - Dispatches connection state and message notifications.
  */
 export class EdgeClient extends Resource implements EdgeConnection {
-  public readonly reconnect = new Event();
-  public readonly connected = new Event();
-  private readonly _persistentLifecycle = new PersistentLifecycle({
-    start: async () => this._openWebSocket(),
-    stop: async () => this._closeWebSocket(),
-    onRestart: async () => this.reconnect.emit(),
+  private readonly _persistentLifecycle = new PersistentLifecycle<EdgeWsConnection>({
+    start: async () => this._connect(),
+    stop: async (state: EdgeWsConnection) => this._disconnect(state),
   });
 
-  private readonly _listeners = new Set<MessageListener>();
-  private _ready = new Trigger();
-  private _ws?: WebSocket = undefined;
-  private _keepaliveCtx?: Context = undefined;
-  private _heartBeatContext?: Context = undefined;
+  private readonly _messageListeners = new Set<MessageListener>();
+  private readonly _reconnectListeners = new Set<ReconnectListener>();
 
-  private _baseWsUrl: string;
-  private _baseHttpUrl: string;
+  private readonly _baseWsUrl: string;
+  private readonly _baseHttpUrl: string;
+
+  private _currentConnection?: EdgeWsConnection = undefined;
+  private _ready = new Trigger();
 
   constructor(
     private _identity: EdgeIdentity,
@@ -73,7 +69,7 @@ export class EdgeClient extends Resource implements EdgeConnection {
     this._baseHttpUrl = getEdgeUrlWithProtocol(_config.socketEndpoint, 'http');
   }
 
-  // TODO(burdon): Attach logging.
+  @logInfo
   public get info() {
     return {
       open: this.isOpen,
@@ -83,7 +79,7 @@ export class EdgeClient extends Resource implements EdgeConnection {
   }
 
   get isConnected() {
-    return Boolean(this._ws) && this._ready.state === TriggerState.RESOLVED;
+    return Boolean(this._currentConnection) && this._ready.state === TriggerState.RESOLVED;
   }
 
   get identityKey() {
@@ -98,13 +94,32 @@ export class EdgeClient extends Resource implements EdgeConnection {
     if (identity.identityKey !== this._identity.identityKey || identity.peerKey !== this._identity.peerKey) {
       log('Edge identity changed', { identity, oldIdentity: this._identity });
       this._identity = identity;
+      this._closeCurrentConnection(new EdgeIdentityChangedError());
       this._persistentLifecycle.scheduleRestart();
     }
   }
 
-  public addListener(listener: MessageListener): () => void {
-    this._listeners.add(listener);
-    return () => this._listeners.delete(listener);
+  public onMessage(listener: MessageListener): () => void {
+    this._messageListeners.add(listener);
+    return () => this._messageListeners.delete(listener);
+  }
+
+  public onReconnected(listener: () => void): () => void {
+    this._reconnectListeners.add(listener);
+    if (this._ready.state === TriggerState.RESOLVED) {
+      // Microtask so that listener is always called asynchronously, no matter the state of the ready trigger
+      // at the moment of registration.
+      scheduleMicroTask(this._ctx, () => {
+        if (this._reconnectListeners.has(listener)) {
+          try {
+            listener();
+          } catch (error) {
+            log.catch(error);
+          }
+        }
+      });
+    }
+    return () => this._reconnectListeners.delete(listener);
   }
 
   /**
@@ -122,97 +137,96 @@ export class EdgeClient extends Resource implements EdgeConnection {
    */
   protected override async _close() {
     log('closing...', { peerKey: this._identity.peerKey });
+    this._closeCurrentConnection();
     await this._persistentLifecycle.close();
   }
 
-  private async _openWebSocket() {
+  private async _connect(): Promise<EdgeWsConnection | undefined> {
     if (this._ctx.disposed) {
-      return;
+      return undefined;
     }
-    const path = `/ws/${this._identity.identityKey}/${this._identity.peerKey}`;
+
+    const identity = this._identity;
+    const path = `/ws/${identity.identityKey}/${identity.peerKey}`;
     const protocolHeader = this._config.disableAuth ? undefined : await this._createAuthHeader(path);
+    if (this._identity !== identity) {
+      log('identity changed during auth header request');
+      return undefined;
+    }
 
+    const restartRequired = new Trigger();
     const url = new URL(path, this._baseWsUrl);
     log('Opening websocket', { url: url.toString(), protocolHeader });
-    this._ws = new WebSocket(url, protocolHeader ? [protocolHeader] : []);
-
-    this._ws.onopen = () => {
-      log('opened', this.info);
-      this._ready.wake();
-      this.connected.emit();
-    };
-    this._ws.onclose = () => {
-      log('closed', this.info);
-      this._persistentLifecycle.scheduleRestart();
-    };
-    this._ws.onerror = (event) => {
-      log.warn('EdgeClient socket error', { error: event.error, info: event.message });
-      this._persistentLifecycle.scheduleRestart();
-    };
-    /**
-     * https://developer.mozilla.org/en-US/docs/Web/API/MessageEvent/data
-     */
-    this._ws.onmessage = async (event) => {
-      if (event.data === '__pong__') {
-        this._onHeartbeat();
-        return;
-      }
-      const data = await toUint8Array(event.data);
-      const message = buf.fromBinary(MessageSchema, data);
-      log('received', { peerKey: this._identity.peerKey, payload: protocol.getPayloadType(message) });
-      if (message) {
-        for (const listener of this._listeners) {
-          try {
-            await listener(message);
-          } catch (err) {
-            log.error('processing', { err, payload: protocol.getPayloadType(message) });
+    const connection = new EdgeWsConnection(
+      identity,
+      { url, protocolHeader },
+      {
+        onConnected: () => {
+          if (this._isActive(connection)) {
+            this._ready.wake();
+            this._notifyReconnected();
+          } else {
+            log.verbose('connected callback ignored, because connection is not active');
           }
-        }
-      }
-    };
-
-    // TODO(dmaretskyi): Potential race condition here since web socket errors don't resolve this trigger.
-    await this._ready.wait({ timeout: this._config.timeout ?? DEFAULT_TIMEOUT });
-    log('Websocket is ready', { identity: this._identity.identityKey, peer: this._identity.peerKey });
-
-    // TODO(dmaretskyi): Potential leak: context re-assigned without disposing the previous one.
-    this._keepaliveCtx = new Context();
-    scheduleTaskInterval(
-      this._keepaliveCtx,
-      async () => {
-        // TODO(mykola): use RFC6455 ping/pong once implemented in the browser?
-        // Cloudflare's worker responds to this `without interrupting hibernation`. https://developers.cloudflare.com/durable-objects/api/websockets/#setwebsocketautoresponse
-        this._ws?.send('__ping__');
+        },
+        onRestartRequired: () => {
+          if (this._isActive(connection)) {
+            this._closeCurrentConnection();
+            this._persistentLifecycle.scheduleRestart();
+          } else {
+            log.verbose('restart requested by inactive connection');
+          }
+          restartRequired.wake();
+        },
+        onMessage: (message) => {
+          if (this._isActive(connection)) {
+            this._notifyMessageReceived(message);
+          } else {
+            log.verbose('ignored a message on inactive connection', {
+              from: message.source,
+              type: message.payload?.typeUrl,
+            });
+          }
+        },
       },
-      SIGNAL_KEEPALIVE_INTERVAL,
     );
-    this._ws.send('__ping__');
-    this._onHeartbeat();
+    this._currentConnection = connection;
+
+    await connection.open();
+    // Race with restartRequired so that restart is not blocked by _connect execution.
+    // Wait on ready to attempt a reconnect if it times out.
+    await Promise.race([this._ready.wait({ timeout: this._config.timeout ?? DEFAULT_TIMEOUT }), restartRequired]);
+
+    return connection;
+  }
+
+  private async _disconnect(state: EdgeWsConnection) {
+    await state.close();
+  }
+
+  private _closeCurrentConnection(error: Error = new EdgeConnectionClosedError()) {
+    this._currentConnection = undefined;
+    this._ready.throw(error);
+    this._ready.reset();
   }
 
-  private async _closeWebSocket() {
-    if (!this._ws) {
-      return;
+  private _notifyReconnected() {
+    for (const listener of this._reconnectListeners) {
+      try {
+        listener();
+      } catch (err) {
+        log.error('ws reconnect listener failed', { err });
+      }
     }
-    try {
-      this._ready.throw(this.isOpen ? new EdgeIdentityChangedError() : new EdgeConnectionClosedError());
-      this._ready.reset();
-      void this._keepaliveCtx?.dispose();
-      this._keepaliveCtx = undefined;
-      void this._heartBeatContext?.dispose();
-      this._heartBeatContext = undefined;
-
-      // NOTE: Remove event handlers to avoid scheduling restart.
-      this._ws.onopen = () => {};
-      this._ws.onclose = () => {};
-      this._ws.onerror = () => {};
-      this._ws.close();
-      this._ws = undefined;
-    } catch (err) {
-      if (err instanceof Error && err.message.includes('WebSocket is closed before the connection is established.')) {
-        return;
+  }
+
+  private _notifyMessageReceived(message: Message) {
+    for (const listener of this._messageListeners) {
+      try {
+        listener(message);
+      } catch (err) {
+        log.error('ws incoming message processing failed', { err, payload: protocol.getPayloadType(message) });
       }
-      log.warn('Error closing websocket', { err });
     }
   }
 
@@ -225,9 +239,11 @@ export class EdgeClient extends Resource implements EdgeConnection {
       log('waiting for websocket to become ready');
       await this._ready.wait({ timeout: this._config.timeout ?? DEFAULT_TIMEOUT });
     }
-    if (!this._ws) {
+
+    if (!this._currentConnection) {
       throw new EdgeConnectionClosedError();
     }
+
     if (
       message.source &&
       (message.source.peerKey !== this._identity.peerKey || message.source.identityKey !== this.identityKey)
@@ -235,23 +251,7 @@ export class EdgeClient extends Resource implements EdgeConnection {
       throw new EdgeIdentityChangedError();
     }
 
-    log('sending...', { peerKey: this._identity.peerKey, payload: protocol.getPayloadType(message) });
-    this._ws.send(buf.toBinary(MessageSchema, message));
-  }
-
-  private _onHeartbeat() {
-    if (this._lifecycleState !== LifecycleState.OPEN) {
-      return;
-    }
-    void this._heartBeatContext?.dispose();
-    this._heartBeatContext = new Context();
-    scheduleTask(
-      this._heartBeatContext,
-      () => {
-        this._persistentLifecycle.scheduleRestart();
-      },
-      2 * SIGNAL_KEEPALIVE_INTERVAL,
-    );
+    this._currentConnection.send(message);
   }
 
   private async _createAuthHeader(path: string): Promise<string | undefined> {
@@ -265,6 +265,8 @@ export class EdgeClient extends Resource implements EdgeConnection {
       return undefined;
     }
   }
+
+  private _isActive = (connection: EdgeWsConnection) => connection === this._currentConnection;
 }
 
 const encodePresentationWsAuthHeader = (encodedPresentation: Uint8Array): string => {