feat: Reconnect (#39)

Author: chakra-guy

apps/integration-tests/package.json

@@ -11,6 +11,7 @@
 	"dependencies": {
 		"@types/ws": "^8.18.1",
 		"centrifuge": "^5.3.5",
+		"toxiproxy-node-client": "^4.0.0",
 		"vitest": "^3.2.4",
 		"ws": "^8.18.3"
 	}

apps/integration-tests/src/end-to-end.integration.test.ts

@@ -1,12 +1,16 @@
 /** biome-ignore-all lint/suspicious/noExplicitAny: test code */
+/** biome-ignore-all lint/suspicious/noShadowRestrictedNames: test code */
 import { type ConnectionMode, type IKeyManager, type IKVStore, type KeyPair, type SessionRequest, SessionStore, WebSocketTransport } from "@metamask/mobile-wallet-protocol-core";
 import { DappClient, type OtpRequiredPayload } from "@metamask/mobile-wallet-protocol-dapp-client";
 import { WalletClient } from "@metamask/mobile-wallet-protocol-wallet-client";
 import { decrypt, encrypt, PrivateKey } from "eciesjs";
+import { type Proxy, Toxiproxy } from "toxiproxy-node-client";
 import * as t from "vitest";
 import WebSocket from "ws";
 
 const RELAY_URL = "ws://localhost:8000/connection/websocket";
+const PROXY_RELAY_URL = "ws://localhost:8001/connection/websocket";
+const TOXIPROXY_URL = "http://localhost:8474";
 
 class InMemoryKVStore implements IKVStore {
 	private store = new Map<string, string>();
@@ -69,6 +73,12 @@ async function connectClients(dappClient: DappClient, walletClient: WalletClient
 	await Promise.all([dappConnectPromise, walletConnectPromise]);
 }
 
+// Helper to assert that a promise does NOT resolve within a given time
+async function assertPromiseNotResolve(promise: Promise<unknown>, timeout: number, message: string) {
+	const timeoutPromise = new Promise((_, reject) => setTimeout(() => reject(new Error(message)), timeout));
+	await t.expect(Promise.race([promise, timeoutPromise])).rejects.toThrow(message);
+}
+
 t.describe("E2E Integration Test", () => {
 	let dappClient: DappClient;
 	let walletClient: WalletClient;
@@ -199,4 +209,114 @@ t.describe("E2E Integration Test", () => {
 		await resumedDappClient.sendRequest(testPayload);
 		await t.expect(messagePromise).resolves.toEqual(testPayload);
 	});
+
+});
+
+t.describe("E2E Integration Test via Proxy", () => {
+	let dappClient: DappClient;
+	let walletClient: WalletClient;
+	let dappKvStore: InMemoryKVStore;
+	let walletKvStore: InMemoryKVStore;
+	let dappSessionStore: SessionStore;
+	let walletSessionStore: SessionStore;
+
+	// Toxiproxy setup
+	let toxiproxy: Toxiproxy;
+	let proxy: Proxy;
+	const proxyConfig = {
+		listen: "0.0.0.0:8001",
+		upstream: "centrifugo:8000",
+	};
+
+	t.beforeAll(async () => {
+		toxiproxy = new Toxiproxy(TOXIPROXY_URL);
+		try {
+			proxy = await toxiproxy.get("centrifugo_proxy");
+			await proxy.remove();
+		} catch {
+			// Proxy doesn't exist, which is fine
+		}
+		proxy = await toxiproxy.createProxy({
+			name: "centrifugo_proxy",
+			...proxyConfig,
+		});
+	});
+
+	t.beforeEach(async () => {
+		// Ensure the proxy is enabled before each test
+		await proxy.update({ ...proxyConfig, enabled: true });
+
+		dappKvStore = new InMemoryKVStore();
+		walletKvStore = new InMemoryKVStore();
+		dappSessionStore = new SessionStore(dappKvStore);
+		walletSessionStore = new SessionStore(walletKvStore);
+
+		// DApp connects directly, Wallet connects through proxy
+		const dappTransport = await WebSocketTransport.create({ url: RELAY_URL, kvstore: dappKvStore, websocket: WebSocket });
+		const walletTransport = await WebSocketTransport.create({ url: PROXY_RELAY_URL, kvstore: walletKvStore, websocket: WebSocket });
+		const keyManager = new KeyManager();
+
+		dappClient = new DappClient({ transport: dappTransport, sessionstore: dappSessionStore, keymanager: keyManager });
+		walletClient = new WalletClient({ transport: walletTransport, sessionstore: walletSessionStore, keymanager: keyManager });
+	});
+
+	t.afterEach(async () => {
+		// Reset proxy state after each test
+		if (proxy) {
+			await proxy.update({ ...proxyConfig, enabled: true });
+		}
+		try {
+			await dappClient?.disconnect();
+		} catch (e) {
+			console.error("Error disconnecting dappClient:", e);
+		}
+		try {
+			await walletClient?.disconnect();
+		} catch (e) {
+			console.error("Error disconnecting walletClient:", e);
+		}
+	});
+
+	t.test(
+		"should recover from a one-sided stale connection and receive queued messages after reconnect",
+		async () => {
+			// 1. Establish a normal connection and exchange a message to confirm it works
+			await connectClients(dappClient, walletClient, "trusted");
+			const initialMessage = { step: "initial_message" };
+			const initialMessagePromise = new Promise((resolve) => walletClient.once("message", resolve));
+			await dappClient.sendRequest(initialMessage);
+			await t.expect(initialMessagePromise).resolves.toEqual(initialMessage);
+			t.expect((walletClient as any).state).toBe("CONNECTED");
+
+			// 2. Create a ONE-SIDED network partition using toxiproxy.
+			// This cuts off the Wallet's connection but leaves the DApp's connection intact.
+			await proxy.update({ ...proxyConfig, enabled: false });
+
+			// 3. Dapp sends a message. Its transport is still live, so it successfully publishes to the relay.
+			// The Wallet, however, is now on a stale connection and should NOT receive it.
+			const missedMessage = { step: "missed_message" };
+			const missedMessagePromise = new Promise((resolve) => walletClient.once("message", resolve));
+			await dappClient.sendRequest(missedMessage);
+
+			// Assert that the message is NOT received by the wallet within 2 seconds.
+			await assertPromiseNotResolve(missedMessagePromise, 2000, "Wallet incorrectly received message during network partition.");
+
+			// 4. Restore the network path and trigger the wallet's reconnect logic.
+			await proxy.update({ ...proxyConfig, enabled: true });
+			await walletClient.reconnect();
+
+			// 5. The wallet should now re-establish its connection. The transport's recovery logic
+			// should fetch the missed message from the channel's history.
+			const receivedMissedMessage = await missedMessagePromise;
+			t.expect(receivedMissedMessage).toEqual(missedMessage);
+			t.expect((walletClient as any).state).toBe("CONNECTED");
+
+			// 6. Send a final message to confirm the live connection is fully restored and working.
+			const finalMessage = { step: "final_message" };
+			const finalMessagePromise = new Promise((resolve) => walletClient.once("message", resolve));
+			await dappClient.sendRequest(finalMessage);
+			await t.expect(finalMessagePromise).resolves.toEqual(finalMessage);
+		},
+		15000,
+	);
 });

apps/rn-demo/context/WalletContext.tsx

@@ -3,7 +3,7 @@ import "react-native-get-random-values";
 import { type Session, SessionStore } from "@metamask/mobile-wallet-protocol-core";
 import Constants from "expo-constants";
 import { createContext, type ReactNode, useCallback, useContext, useEffect, useState } from "react";
-import { Platform } from "react-native";
+import { AppState, type AppStateStatus, Platform } from "react-native";
 import { AsyncStorageKVStore } from "@/lib/AsyncStorageKVStore";
 import { type GlobalActivityLogEntry, SessionManager } from "@/lib/SessionManager";
 
@@ -127,6 +127,22 @@ export function WalletProvider({ children }: { children: ReactNode }) {
 		};
 	}, [addLog]);
 
+	useEffect(() => {
+		const handleAppStateChange = (nextAppState: AppStateStatus) => {
+			console.log(`App state changed to: ${nextAppState}`);
+			if (nextAppState === "active") {
+				console.log("App is in the foreground, reconnecting all clients...");
+				sessionManager?.reconnectAllClients();
+			}
+		};
+
+		const subscription = AppState.addEventListener("change", handleAppStateChange);
+
+		return () => {
+			subscription.remove();
+		};
+	}, [sessionManager]);
+
 	const value = { sessionManager, sessions, globalActivityLog, isInitializing, error, addLog, otpToDisplay, clearOtpDisplay };
 
 	return <WalletContext.Provider value={value}>{children}</WalletContext.Provider>;

apps/rn-demo/lib/SessionManager.ts

@@ -103,6 +103,12 @@ export class SessionManager extends EventEmitter {
 		await Promise.all(allSessions.map((id) => this.deleteClient(id)));
 	}
 
+	public async reconnectAllClients(): Promise<void> {
+		console.log("SessionManager: Reconnecting all clients...");
+		const reconnectPromises = Array.from(this.clients.values()).map((client) => client.reconnect());
+		await Promise.all(reconnectPromises);
+	}
+
 	public getClient(sessionId: string): WalletClient | undefined {
 		return this.clients.get(sessionId);
 	}

backend/docker-compose.yml

@@ -6,4 +6,13 @@ services:
       - 8000:8000
     volumes:
       - ./config.json:/centrifugo/config.json
-    command: centrifugo -c config.json 
+    command: centrifugo -c config.json
+
+  toxiproxy:
+    image: shopify/toxiproxy:latest
+    ports:
+      # The test client will connect through this proxy port
+      - 8001:8001
+      # The test runner will control the proxy via this HTTP API port
+      - 8474:8474
+    command: -host 0.0.0.0 

packages/core/src/base-client.ts

@@ -50,6 +50,25 @@ export abstract class BaseClient extends EventEmitter {
 		});
 	}
 
+	/**
+	 * Proactively refreshes the underlying transport connection.
+	 * This is the recommended method for mobile clients to call when the application
+	 * returns to the foreground to ensure the connection is not stale.
+	 */
+	public async reconnect(): Promise<void> {
+		if (this.state === ClientState.CONNECTING || !this.session || !this.transport.reconnect) return;
+
+		try {
+			this.state = ClientState.CONNECTING;
+			await this.transport.reconnect();
+			this.state = ClientState.CONNECTED;
+			this.emit("connected");
+		} catch {
+			this.state = ClientState.DISCONNECTED;
+			throw new TransportError(ErrorCode.TRANSPORT_RECONNECT_FAILED, "Failed to reconnect");
+		}
+	}
+
 	/**
 	 * Resumes an existing session by loading it from storage and connecting to the
 	 * transport on the session's secure channel.

packages/core/src/domain/errors.ts

@@ -11,6 +11,7 @@ export enum ErrorCode {
 	TRANSPORT_SUBSCRIBE_FAILED = "TRANSPORT_SUBSCRIBE_FAILED",
 	TRANSPORT_HISTORY_FAILED = "TRANSPORT_HISTORY_FAILED",
 	TRANSPORT_PARSE_FAILED = "TRANSPORT_PARSE_FAILED",
+	TRANSPORT_RECONNECT_FAILED = "TRANSPORT_RECONNECT_FAILED",
 
 	// Crypto errors
 	DECRYPTION_FAILED = "DECRYPTION_FAILED",

packages/core/src/domain/transport.ts

@@ -49,4 +49,14 @@ export interface ITransport {
 	 * @param channel The channel to clear.
 	 */
 	clear(channel: string): Promise<void>;
+
+	/**
+	 * Forcibly drops the current connection and attempts to re-establish it.
+	 * This is designed to recover from stale or zombie connections without losing
+	 * the client's subscription state. It should trigger the transport's built-in
+	 * recovery mechanisms upon a successful new connection.
+	 *
+	 * @returns A promise that resolves when the connection is re-established.
+	 */
+	reconnect?(): Promise<void>;
 }

packages/core/src/transport/websocket/index.ts

@@ -130,6 +130,25 @@ export class WebSocketTransport extends EventEmitter implements ITransport {
 		});
 	}
 
+	/**
+	 * Disconnects and immediately reconnects the underlying Centrifuge client.
+	 * This is a proactive way to force a fresh connection while preserving all
+	 * existing subscription objects in memory, allowing for automatic recovery.
+	 */
+	public reconnect(): Promise<void> {
+		if (this.state === "connecting") {
+			return new Promise((resolve) => this.centrifuge.once("connected", () => resolve()));
+		}
+
+		this.centrifuge.disconnect();
+
+		return new Promise((resolve, reject) => {
+			this.centrifuge.once("connected", () => resolve());
+			this.centrifuge.once("error", (ctx) => reject(new TransportError(ErrorCode.TRANSPORT_RECONNECT_FAILED, ctx.error.message)));
+			this.centrifuge.connect();
+		});
+	}
+
 	/**
 	 * Subscribes to a channel and fetches historical messages and sends any queued messages.
 	 */