fix: switch balances to a new endpoint as soon as it is chosen

Choosing another endpoint in settings updated balances only eventually. Three
things were wrong:

- The popup read the endpoint only when it rendered. Changing the setting
  didn't re-render the popup, so the switch waited for some unrelated update.
- The old subscription was never ended. Each switch added a subscription
  beside the old one, and whichever endpoint the background connected to last
  won.
- The custom-node field saved every keystroke. Once the switch is immediate,
  that means connecting to `w`, `ws`, `ws:`, and so on.

Now:
- the balance hook follows the setting, including a change made in another
  extension page;
- it clears balances at once, rather than showing another chain's while the
  new one loads;
- it ends the old subscription through a new pri(balances.unsubscribe);
- the background shares one in-flight connection, and discards a connection or
  read that a later choice has superseded;
- a typed endpoint is used on Enter or when the field loses focus.

Verified in Firefox against the live nodes, with crystal_alice imported:
- mainnet -> Heisenberg shows 477.325 HEI after 1.5s;
- Heisenberg -> Planck shows 0 PLK after 0.4s;
- Planck -> mainnet shows 0 QTC after 1.0s;
- each switch clears the old balance within 0.1s;
- typing a node address leaves the saved endpoint unchanged until blur.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012uDUodEcRbBwNRi3UCmw8f

packages/extension-base/src/background/Balances.ts

@@ -47,6 +47,8 @@ export default class Balances {
   readonly subject = new BehaviorSubject<AccountBalances>({});
 
   #chain: ChainInfo | null = null;
+  // A connection being made, so readers arriving together share one.
+  #connecting: { endpoint: string, promise: Promise<ChainInfo> } | null = null;
   #endpoint: string | null = null;
   #provider: ProviderInterface | null = null;
   #subscribers = 0;
@@ -62,9 +64,29 @@ export default class Balances {
    * parsing it is the expensive part, so it is held for as long as the
    * connection is.
    */
-  async #connect (endpoint: string): Promise<ChainInfo> {
+  #connect (endpoint: string): Promise<ChainInfo> {
     if (this.#chain && this.#endpoint === endpoint && this.#provider?.isConnected) {
-      return this.#chain;
+      return Promise.resolve(this.#chain);
+    }
+
+    if (this.#connecting?.endpoint !== endpoint) {
+      const promise = this.#open(endpoint).finally(() => {
+        if (this.#connecting?.promise === promise) {
+          this.#connecting = null;
+        }
+      });
+
+      this.#connecting = { endpoint, promise };
+    }
+
+    return this.#connecting.promise;
+  }
+
+  async #open (endpoint: string): Promise<ChainInfo> {
+    // Balances from the previous endpoint are another chain's: a switch from
+    // Heisenberg to mainnet must not go on showing HEI while mainnet loads.
+    if (this.#endpoint !== endpoint) {
+      this.subject.next({});
     }
 
     await this.#close();
@@ -77,6 +99,15 @@ export default class Balances {
       provider.send<string>('state_getMetadata', []),
       provider.send<Record<string, unknown>>('system_properties', [])
     ]);
+
+    // The user moved on while this was loading. Whatever connects last must not
+    // win; whatever was asked for last must.
+    if (this.#connecting?.endpoint !== endpoint) {
+      await provider.disconnect().catch(console.error);
+
+      throw new Error(`${endpoint} was replaced before it connected`);
+    }
+
     const runtime = Runtime.fromMetadata(hexToU8a(metadata));
     // `system_properties` reports these as either a scalar or a one-element
     // array, depending on the chain's spec.
@@ -142,6 +173,11 @@ export default class Balances {
     const values = await provider.send<(string | null)[]>('state_queryStorageAt', [keys.map((k) => k.key)])
       .then((results) => (results as unknown as { changes: [string, string | null][] }[])[0].changes.map(([, value]) => value));
 
+    // a read that outlived a switch of endpoint belongs to the old chain
+    if (this.#provider !== provider) {
+      return;
+    }
+
     this.subject.next(Object.fromEntries(addresses.map((address, index) => [address, read(index, values[index])])));
 
     await this.#watch(provider, addresses, keys.map((k) => k.key), read);

packages/extension-base/src/background/handlers/Extension.ts

@@ -8,7 +8,7 @@ import type { MetadataDef } from '@polkadot/extension-inject/types';
 import type { KeyringPair, KeyringPair$Json, KeyringPair$Meta } from '@polkadot/keyring/types';
 import type { SubjectInfo } from '@polkadot/ui-keyring/observable/types';
 import type { KeypairType } from '@polkadot/util-crypto/types';
-import type { AccountJson, AllowedPath, AuthorizeRequest, MessageTypes, MetadataRequest, RequestAccountBatchExport, RequestAccountChangePassword, RequestAccountCreateExternal, RequestAccountCreateSuri, RequestAccountEdit, RequestAccountExport, RequestAccountForget, RequestAccountShow, RequestAccountTie, RequestAccountValidate, RequestActiveTabsUrlUpdate, RequestAuthorizeApprove, RequestBalancesSubscribe, RequestBatchRestore, RequestDeriveCreate, RequestDeriveValidate, RequestJsonRestore, RequestMetadataApprove, RequestMetadataReject, RequestSeedCreate, RequestSeedValidate, RequestSigningApprovePassword, RequestSigningApproveSignature, RequestSigningCancel, RequestSigningIsLocked, RequestTypes, RequestUpdateAuthorizedAccounts, ResponseAccountExport, ResponseAccountsExport, ResponseAuthorizeList, ResponseDeriveValidate, ResponseJsonGetAccountInfo, ResponseSeedCreate, ResponseSeedValidate, ResponseSigningIsLocked, ResponseType, SigningRequest } from '../types.js';
+import type { AccountJson, AllowedPath, AuthorizeRequest, MessageTypes, MetadataRequest, RequestAccountBatchExport, RequestAccountChangePassword, RequestAccountCreateExternal, RequestAccountCreateSuri, RequestAccountEdit, RequestAccountExport, RequestAccountForget, RequestAccountShow, RequestAccountTie, RequestAccountValidate, RequestActiveTabsUrlUpdate, RequestAuthorizeApprove, RequestBalancesSubscribe, RequestBalancesUnsubscribe, RequestBatchRestore, RequestDeriveCreate, RequestDeriveValidate, RequestJsonRestore, RequestMetadataApprove, RequestMetadataReject, RequestSeedCreate, RequestSeedValidate, RequestSigningApprovePassword, RequestSigningApproveSignature, RequestSigningCancel, RequestSigningIsLocked, RequestTypes, RequestUpdateAuthorizedAccounts, ResponseAccountExport, ResponseAccountsExport, ResponseAuthorizeList, ResponseDeriveValidate, ResponseJsonGetAccountInfo, ResponseSeedCreate, ResponseSeedValidate, ResponseSigningIsLocked, ResponseType, SigningRequest } from '../types.js';
 import type { AuthorizedAccountsDiff } from './State.js';
 import type State from './State.js';
 
@@ -38,6 +38,8 @@ function getSuri (seed: string, type?: KeypairType): string {
 
 export default class Extension {
   readonly #balances = new Balances();
+  // Ends each balance subscription, by id; each runs once however it is reached.
+  readonly #balanceSubs = new Map<string, VoidFunction>();
   readonly #cachedUnlocks: CachedUnlocks;
 
   readonly #state: State;
@@ -204,7 +206,7 @@ export default class Extension {
    * belong to one person, so holding a socket open for the life of the browser
    * would report far more than the feature needs.
    */
-  private balancesSubscribe ({ endpoint }: RequestBalancesSubscribe, id: string, port: chrome.runtime.Port): boolean {
+  private balancesSubscribe ({ endpoint }: RequestBalancesSubscribe, id: string, port: chrome.runtime.Port): string {
     const cb = createSubscription<'pri(balances.subscribe)'>(id, port);
 
     this.#balances.retain();
@@ -221,12 +223,28 @@ export default class Extension {
     const balances = this.#balances.subject.subscribe(cb);
     const accounts = accountsObservable.subject.subscribe(update);
 
-    port.onDisconnect.addListener((): void => {
-      unsubscribe(id);
-      balances.unsubscribe();
-      accounts.unsubscribe();
-      this.#balances.release();
-    });
+    const end = (): void => {
+      if (this.#balanceSubs.delete(id)) {
+        unsubscribe(id);
+        balances.unsubscribe();
+        accounts.unsubscribe();
+        this.#balances.release();
+      }
+    };
+
+    this.#balanceSubs.set(id, end);
+    port.onDisconnect.addListener(end);
+
+    return id;
+  }
+
+  /**
+   * End a balance subscription while the page that made it stays open, as it
+   * does when the user switches endpoint. Left to the port disconnecting, the
+   * old endpoint's subscription would live on beside the new one.
+   */
+  private balancesUnsubscribe ({ id }: RequestBalancesUnsubscribe): boolean {
+    this.#balanceSubs.get(id)?.();
 
     return true;
   }
@@ -649,6 +667,9 @@ export default class Extension {
       case 'pri(balances.subscribe)':
         return port && this.balancesSubscribe(request as RequestBalancesSubscribe, id, port);
 
+      case 'pri(balances.unsubscribe)':
+        return this.balancesUnsubscribe(request as RequestBalancesUnsubscribe);
+
       case 'pri(accounts.tie)':
         return this.accountsTie(request as RequestAccountTie);
 

packages/extension-base/src/background/types.ts

@@ -90,7 +90,8 @@ export interface RequestSignatures {
   'pri(accounts.show)': [RequestAccountShow, boolean];
   'pri(accounts.tie)': [RequestAccountTie, boolean];
   'pri(accounts.subscribe)': [RequestAccountSubscribe, boolean, AccountJson[]];
-  'pri(balances.subscribe)': [RequestBalancesSubscribe, boolean, AccountBalances];
+  'pri(balances.subscribe)': [RequestBalancesSubscribe, string, AccountBalances];
+  'pri(balances.unsubscribe)': [RequestBalancesUnsubscribe, boolean];
   'pri(accounts.validate)': [RequestAccountValidate, boolean];
   'pri(accounts.changePassword)': [RequestAccountChangePassword, boolean];
   'pri(authorize.approve)': [RequestAuthorizeApprove, boolean];
@@ -310,6 +311,11 @@ export interface RequestBalancesSubscribe {
   endpoint: string;
 }
 
+/** The id `pri(balances.subscribe)` answered with. */
+export interface RequestBalancesUnsubscribe {
+  id: string;
+}
+
 export interface AccountBalance {
   /** The chain's decimal places, for formatting. */
   decimals: number;

packages/extension-ui/src/hooks/useBalances.ts

@@ -5,8 +5,8 @@ import type { AccountBalances } from '@polkadot/extension-base/background/types'
 
 import { useEffect, useState } from 'react';
 
-import { subscribeBalances } from '../messaging.js';
-import { getBalanceEndpoint } from '../util/balanceEndpoint.js';
+import { subscribeBalances, unsubscribeBalances } from '../messaging.js';
+import { getBalanceEndpoint, onBalanceEndpointChange } from '../util/balanceEndpoint.js';
 
 /**
  * Balances for every account, from the endpoint in settings.
@@ -22,16 +22,30 @@ import { getBalanceEndpoint } from '../util/balanceEndpoint.js';
  */
 export default function useBalances (): AccountBalances {
   const [balances, setBalances] = useState<AccountBalances>({});
-  const endpoint = getBalanceEndpoint();
+  const [endpoint, setEndpoint] = useState(getBalanceEndpoint);
+
+  // Follow the setting as it changes. Read once at render, a new endpoint
+  // waited for something unrelated to re-render the popup.
+  useEffect(() => onBalanceEndpointChange(setEndpoint), []);
+
+  useEffect(() => {
+    // The old endpoint's balances are another chain's; show none rather than
+    // those while the new one loads.
+    setBalances({});
 
-  useEffect((): void => {
     if (!endpoint) {
-      setBalances({});
-
       return;
     }
 
-    subscribeBalances(endpoint, setBalances).catch(console.error);
+    let isCurrent = true;
+    const id = subscribeBalances(endpoint, (next) => isCurrent && setBalances(next));
+
+    id.catch(console.error);
+
+    return (): void => {
+      isCurrent = false;
+      id.then(unsubscribeBalances).catch(console.error);
+    };
   }, [endpoint]);
 
   return balances;

packages/extension-ui/src/messaging.ts

@@ -203,10 +203,20 @@ export async function subscribeAccounts (cb: (accounts: AccountJson[]) => void):
  * to `localStorage`, which the MV3 service worker holding the connection does
  * not have.
  */
-export async function subscribeBalances (endpoint: string, cb: (balances: AccountBalances) => void): Promise<boolean> {
+/** Resolves to the subscription's id, which `unsubscribeBalances` takes. */
+export async function subscribeBalances (endpoint: string, cb: (balances: AccountBalances) => void): Promise<string> {
   return sendMessage('pri(balances.subscribe)', { endpoint }, cb);
 }
 
+export async function unsubscribeBalances (id: string): Promise<boolean> {
+  // stop listening and stop replaying first: whatever the background says
+  // next about this id is no longer wanted
+  client.forget(id);
+  delete handlers[id];
+
+  return sendMessage('pri(balances.unsubscribe)', { id });
+}
+
 export async function subscribeAuthorizeRequests (cb: (accounts: AuthorizeRequest[]) => void): Promise<boolean> {
   return sendMessage('pri(authorize.requests)', null, cb);
 }

packages/extension-ui/src/partials/MenuSettings.tsx

@@ -77,6 +77,14 @@ function MenuSettings ({ className, reference }: Props): React.ReactElement<Prop
     }, []
   );
 
+  // Typing is not choosing: saving each keystroke would connect to `w`, `ws`,
+  // `ws:` and so on. The typed endpoint is used on Enter or on leaving the field.
+  const _onCommitEndpoint = useCallback(
+    (): void => {
+      endpoint !== getBalanceEndpoint() && setBalanceEndpoint(endpoint);
+    }, [endpoint]
+  );
+
   const _onChangeNotification = useCallback(
     (value: string): void => {
       setNotification(value).catch(console.error);
@@ -157,7 +165,9 @@ function MenuSettings ({ className, reference }: Props): React.ReactElement<Prop
         />
         <InputWithLabel
           label={t('or a node of your own')}
-          onChange={_onChangeEndpoint}
+          onBlur={_onCommitEndpoint}
+          onChange={setEndpoint}
+          onEnter={_onCommitEndpoint}
           placeholder='wss://…'
           value={endpoint}
         />

packages/extension-ui/src/util/balanceEndpoint.ts

@@ -37,10 +37,33 @@ export function getBalanceEndpoint (): string {
   }
 }
 
+const CHANGED = 'quantus:balanceEndpointChanged';
+
 export function setBalanceEndpoint (endpoint: string): void {
   try {
     localStorage.setItem(KEY, endpoint);
   } catch (error) {
     console.error(error);
   }
+
+  // `storage` events reach other extension pages (a popup and a tab open side
+  // by side) but never the page that made the change, so tell this one directly.
+  window.dispatchEvent(new Event(CHANGED));
+}
+
+/** Call `cb` whenever the endpoint changes, here or in another extension page. */
+export function onBalanceEndpointChange (cb: (endpoint: string) => void): () => void {
+  const onChanged = (): void => cb(getBalanceEndpoint());
+
+  const onStorage = ({ key }: StorageEvent): void => {
+    key === KEY && onChanged();
+  };
+
+  window.addEventListener(CHANGED, onChanged);
+  window.addEventListener('storage', onStorage);
+
+  return (): void => {
+    window.removeEventListener(CHANGED, onChanged);
+    window.removeEventListener('storage', onStorage);
+  };
 }