feat: backport wallet-sdk heartbeat (#22948) to v4-next (#22956)

Backport of #22948
to `v4-next`.

## Original PR

> Implements a liveness check on wallets, avoiding promises hanging
forever. Due to proving taking a long time, the global backwards
compatible timer is set to 5min (since every request can refresh the
timer), but once we deprecate `v4` we can tighten it.

Closes:
-
https://linear.app/aztec-labs/issue/F-510/audit-143-silent-error-swallowing-in-wallet-extension-decryption
-
https://linear.app/aztec-labs/issue/F-511/audit-142-extensionwalletpostmessage-has-no-timeout-requests-hang

## Conflict resolution

The cherry-pick produced a single **modify/delete** conflict on
`docs/examples/webapp-tutorial/test-extension/src/background.ts` — that
tutorial example file exists only on `next`, not on `v4-next`. All 8
wallet-sdk source files applied cleanly. The conflict was resolved by
dropping the docs/example file (it's not part of v4-next); the
wallet-sdk content is identical to the original PR.

Three-commit history preserved per backport guidelines:
1. `feat: wallet-sdk heartbeat (#22948)` — original cherry-pick, with
the tutorial file added (as left in tree by git)
2. `fix: resolve cherry-pick conflicts in webapp-tutorial background.ts`
— remove the tutorial file
3. *(no build-fix commit needed — all foundation/aztec.js imports
resolve cleanly on v4-next)*

Detailed analysis:
https://gist.github.com/AztecBot/78418f5937bd4b60a797ea26f477479d


ClaudeBox log: https://claudebox.work/s/da6b542d03dd5623?run=1

Author: Thunkar

noir-projects/aztec-nr/aztec/src/contract_self/mod.nr

@@ -2,7 +2,7 @@ pub mod contract_self_private;
 pub mod contract_self_public;
 pub mod contract_self_utility;
 
-pub use contract_self_private::PrivateUtilityCalls;
 pub use contract_self_private::ContractSelfPrivate;
+pub use contract_self_private::PrivateUtilityCalls;
 pub use contract_self_public::ContractSelfPublic;
 pub use contract_self_utility::ContractSelfUtility;

yarn-project/wallet-sdk/src/extension/handlers/background_connection_handler.ts

@@ -17,6 +17,7 @@ import {
   type WalletMessage,
   WalletMessageType,
   type WalletResponse,
+  type WalletSdkLogger,
 } from '../../types.js';
 import {
   type BackgroundMessage,
@@ -131,6 +132,8 @@ export interface BackgroundConnectionConfig {
   walletVersion: string;
   /** Optional wallet icon URL. */
   walletIcon?: string;
+  /** Logger used for diagnostics. */
+  logger: WalletSdkLogger;
 }
 
 /**
@@ -149,6 +152,7 @@ export interface BackgroundConnectionConfig {
  *     walletId: 'my-wallet',
  *     walletName: 'My Wallet',
  *     walletVersion: '1.0.0',
+ *     logger: console,
  *   },
  *   {
  *     sendToTab: (tabId, message) => browser.tabs.sendMessage(tabId, message),
@@ -167,12 +171,15 @@ export interface BackgroundConnectionConfig {
 export class BackgroundConnectionHandler {
   private pendingDiscoveries = new Map<string, PendingDiscovery>();
   private activeSessions = new Map<string, ActiveSession>();
+  private log: WalletSdkLogger;
 
   constructor(
     private config: BackgroundConnectionConfig,
     private transport: BackgroundTransport,
     private callbacks: BackgroundConnectionCallbacks = {},
-  ) {}
+  ) {
+    this.log = config.logger;
+  }
 
   initialize(): void {
     this.transport.addContentListener(this.handleMessage);
@@ -198,8 +205,8 @@ export class BackgroundConnectionHandler {
         break;
       case InternalMessageType.KEY_EXCHANGE_REQUEST:
         if (sessionId) {
-          this.handleKeyExchangeRequest(sessionId, content as KeyExchangeRequest).catch(() => {
-            // Key exchange failed - session won't be established
+          this.handleKeyExchangeRequest(sessionId, content as KeyExchangeRequest).catch(err => {
+            this.log.warn('Key exchange failed — session will not be established', { sessionId, err });
           });
         }
         break;
@@ -213,9 +220,31 @@ export class BackgroundConnectionHandler {
           void this.handleEncryptedMessage(sessionId, content as EncryptedPayload);
         }
         break;
+      case InternalMessageType.PING:
+        if (sessionId) {
+          this.handlePing(sessionId);
+        }
+        break;
     }
   };
 
+  /**
+   * Reply to a dApp PING with a PONG. Used as a liveness probe so the dApp can
+   * tell the difference between a slow request and a dead extension.
+   * @param sessionId - The session that sent the PING.
+   */
+  private handlePing(sessionId: string): void {
+    const session = this.activeSessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+    this.transport.sendToTab(session.tabId, {
+      origin: MessageOrigin.BACKGROUND,
+      type: InternalMessageType.PONG,
+      sessionId,
+    });
+  }
+
   getWalletInfo(): WalletInfo {
     return {
       id: this.config.walletId,
@@ -315,8 +344,8 @@ export class BackgroundConnectionHandler {
       });
 
       this.callbacks.onSessionEstablished?.(session);
-    } catch {
-      // Key exchange failed silently - session won't be established
+    } catch (err) {
+      this.log.warn('Key exchange failed — session will not be established', { sessionId, err });
     }
   }
 
@@ -329,8 +358,8 @@ export class BackgroundConnectionHandler {
     try {
       const message = await decrypt<WalletMessage>(session.sharedKey, encrypted);
       this.callbacks.onWalletMessage?.(session, message);
-    } catch {
-      // Decryption failed - ignore malformed message
+    } catch (err) {
+      this.log.warn('Failed to decrypt incoming wallet message', { sessionId, err });
     }
   }
 
@@ -348,8 +377,12 @@ export class BackgroundConnectionHandler {
         sessionId,
         content: encrypted,
       });
-    } catch {
-      // Encryption failed - response won't be sent
+    } catch (err) {
+      this.log.error('Failed to encrypt wallet response — response will not be sent', {
+        sessionId,
+        messageId: response.messageId,
+        err,
+      });
     }
   }
 

yarn-project/wallet-sdk/src/extension/handlers/content_script_connection_handler.ts

@@ -139,9 +139,20 @@ export class ContentScriptConnectionHandler {
       case InternalMessageType.SESSION_DISCONNECTED:
         this.handleSessionDisconnected(sessionId);
         break;
+      case InternalMessageType.PONG:
+        this.handlePong(sessionId);
+        break;
     }
   };
 
+  private handlePong(sessionId: string): void {
+    const connection = this.ports.get(sessionId);
+    if (!connection) {
+      return;
+    }
+    connection.port.postMessage({ type: WalletMessageType.PONG });
+  }
+
   private handleDiscoveryRequest(request: DiscoveryRequest): void {
     this.transport.sendToBackground({
       origin: MessageOrigin.CONTENT_SCRIPT,
@@ -178,6 +189,13 @@ export class ContentScriptConnectionHandler {
             content: data,
           });
           break;
+        case WalletMessageType.PING:
+          this.transport.sendToBackground({
+            origin: MessageOrigin.CONTENT_SCRIPT,
+            type: InternalMessageType.PING,
+            sessionId,
+          });
+          break;
         default:
           this.transport.sendToBackground({
             origin: MessageOrigin.CONTENT_SCRIPT,

yarn-project/wallet-sdk/src/extension/handlers/internal_message_types.ts

@@ -9,11 +9,13 @@ export const InternalMessageType = {
   KEY_EXCHANGE_REQUEST: 'key-exchange-request',
   SECURE_MESSAGE: 'secure-message',
   DISCONNECT_REQUEST: 'disconnect-request',
+  PING: 'ping',
   // Background → Content script
   DISCOVERY_APPROVED: 'discovery-approved',
   KEY_EXCHANGE_RESPONSE: 'key-exchange-response',
   SECURE_RESPONSE: 'secure-response',
   SESSION_DISCONNECTED: 'session-disconnected',
+  PONG: 'pong',
 } as const;
 
 /**

yarn-project/wallet-sdk/src/extension/provider/extension_wallet.test.ts

@@ -0,0 +1,153 @@
+import type { ChainInfo } from '@aztec/aztec.js/account';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { jsonStringify } from '@aztec/foundation/json-rpc';
+import { promiseWithResolvers } from '@aztec/foundation/promise';
+import { sleep } from '@aztec/foundation/sleep';
+
+import {
+  type EncryptedPayload,
+  decrypt,
+  deriveSessionKeys,
+  encrypt,
+  exportPublicKey,
+  generateKeyPair,
+  importPublicKey,
+} from '../../crypto.js';
+import { type WalletMessage, WalletMessageType, type WalletResponse } from '../../types.js';
+import { ExtensionWallet } from './extension_wallet.js';
+
+async function makeSharedKeys(): Promise<{ appKey: CryptoKey; walletKey: CryptoKey }> {
+  const appPair = await generateKeyPair();
+  const walletPair = await generateKeyPair();
+
+  const appPub = await exportPublicKey(appPair.publicKey);
+  const walletPub = await exportPublicKey(walletPair.publicKey);
+
+  const appSession = await deriveSessionKeys(appPair, await importPublicKey(walletPub), true);
+  const walletSession = await deriveSessionKeys(walletPair, await importPublicKey(appPub), false);
+
+  return { appKey: appSession.encryptionKey, walletKey: walletSession.encryptionKey };
+}
+
+const chainInfo: ChainInfo = { chainId: new Fr(1), version: new Fr(1) };
+
+// Tight heartbeat tuning so tests run quickly. Real defaults (5s/300s) would take
+// a minute per test.
+const FAST_HEARTBEAT = { intervalMs: 50, deadAfterMs: 250 };
+
+describe('ExtensionWallet heartbeat', () => {
+  it('treats PONG as proof of liveness — no false disconnect on slow but alive peer', async () => {
+    const { appKey, walletKey } = await makeSharedKeys();
+    const channel = new MessageChannel();
+    const walletId = 'test-wallet';
+
+    const requestArrived = promiseWithResolvers<WalletMessage>();
+    let pongs = 0;
+
+    channel.port2.onmessage = async (event: MessageEvent) => {
+      const data = event.data;
+      if (data?.type === WalletMessageType.PING) {
+        pongs++;
+        channel.port2.postMessage({ type: WalletMessageType.PONG });
+        return;
+      }
+      const decoded = await decrypt<WalletMessage>(walletKey, data as EncryptedPayload);
+      requestArrived.resolve(decoded);
+    };
+    channel.port2.start();
+
+    const extWallet = ExtensionWallet.create(
+      walletId,
+      channel.port1,
+      appKey,
+      chainInfo,
+      'test-app',
+      undefined,
+      FAST_HEARTBEAT,
+    );
+    const wallet = extWallet.asWallet();
+
+    const pending = wallet.getChainInfo();
+    pending.catch(() => {});
+    const message = await requestArrived.promise;
+
+    // Let several heartbeat ticks fire while the peer responds with PONGs only —
+    // no real response yet. This is the slow-but-alive-wallet scenario. Wait
+    // ~3× the dead-after window to prove that PONGs alone keep the channel alive.
+    await sleep(FAST_HEARTBEAT.deadAfterMs * 3);
+
+    expect(pongs).toBeGreaterThan(0);
+    expect(extWallet.isDisconnected()).toBe(false);
+
+    // Now have the wallet respond. The dApp must settle the promise — the channel
+    // is alive. We send an error to avoid having to construct a schema-valid
+    // result object; what matters here is that the in-flight promise *settles*
+    // rather than hanging forever after the heartbeat-eligible window passes.
+    const response: WalletResponse = {
+      messageId: message.messageId,
+      error: 'simulated wallet response',
+      walletId,
+    };
+    const encrypted = await encrypt(walletKey, jsonStringify(response));
+    // Wallet side posts on port2 so the dApp receives on port1.
+    channel.port2.postMessage(encrypted);
+
+    await expect(pending).rejects.toThrow(/simulated/);
+
+    channel.port1.close();
+    channel.port2.close();
+  });
+
+  it('declares disconnect when the channel is fully silent past the dead window', async () => {
+    const { appKey } = await makeSharedKeys();
+    const channel = new MessageChannel();
+    // Deliberately do NOT wire up port2 — peer is fully unresponsive.
+    channel.port2.start();
+
+    const extWallet = ExtensionWallet.create(
+      'dead-wallet',
+      channel.port1,
+      appKey,
+      chainInfo,
+      'test-app',
+      undefined,
+      FAST_HEARTBEAT,
+    );
+    const wallet = extWallet.asWallet();
+
+    let disconnected = false;
+    extWallet.onDisconnect(() => {
+      disconnected = true;
+    });
+
+    const pending = wallet.getChainInfo();
+
+    await expect(pending).rejects.toThrow(/disconnected/i);
+    expect(disconnected).toBe(true);
+    expect(extWallet.isDisconnected()).toBe(true);
+
+    channel.port2.close();
+  });
+
+  it('does not start a heartbeat when there are no in-flight requests', async () => {
+    const { appKey } = await makeSharedKeys();
+    const channel = new MessageChannel();
+
+    let pings = 0;
+    channel.port2.onmessage = (event: MessageEvent) => {
+      if (event.data?.type === WalletMessageType.PING) {
+        pings++;
+      }
+    };
+    channel.port2.start();
+
+    ExtensionWallet.create('idle-wallet', channel.port1, appKey, chainInfo, 'test-app', undefined, FAST_HEARTBEAT);
+
+    await sleep(150);
+
+    expect(pings).toBe(0);
+
+    channel.port1.close();
+    channel.port2.close();
+  });
+});