Support Phantom account switching and add sign-out

Wires Phantom's own accountChanged/disconnect provider events
(inject.ts) so switching accounts or disconnecting directly in
Phantom's UI is detected, not just our own connect/sign calls --
relayed as unsolicited postMessage events (relay.ts's onWalletEvent)
since they aren't a response to any request we made.

Adds a "Sign out" button in the popup (nexa:sign-out) that clears the
stored session, force-closes the WS connection via a new
ws-client.ts disconnect() (distinct from reconnectNow() -- it also
suppresses auto-reconnect until a new wallet connects), and asks the
content script to call provider.disconnect(), which revokes
Phantom's trust for the origin so the next silent connect correctly
fails until the user reconnects.

Fixes a real bug this surfaced: the existing "skip re-auth if a
session token exists" check in background.ts only checked for *any*
token, so switching Phantom accounts would have silently kept
authenticating as the old wallet. Session storage now tracks which
wallet it belongs to (backend-client.ts's storeSession(token,
walletAddress)) so the handler can tell "already signed in" apart
from "signed in as a different wallet than the one that just
connected."

Co-Authored-By: Claude Sonnet 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01YXiHuScXrjxBh7yFGAPq3B

src/background/backend-client.ts

@@ -1,24 +1,36 @@
 import { browser } from 'wxt/browser';
 
-const SESSION_TOKEN_STORAGE_KEY = 'nexa:sessionToken';
+const SESSION_STORAGE_KEY = 'nexa:session';
 
-/**
- * Session token storage. The token itself is obtained by wallet-auth.ts
- * (nonce -> Phantom signature -> verify) once a wallet connects — this
- * module just persists/retrieves it, since ws-client.ts and popup code
- * shouldn't know how a token was obtained.
- */
-export async function getSessionToken(): Promise<string> {
-  const stored = await browser.storage.local.get(SESSION_TOKEN_STORAGE_KEY);
-  const token = stored[SESSION_TOKEN_STORAGE_KEY] as string | undefined;
-  if (!token) throw new Error('Not authenticated yet — connect a wallet on a supported trading site.');
-  return token;
+interface StoredSession {
+  token: string;
+  walletAddress: string;
 }
 
-export async function storeSessionToken(token: string): Promise<void> {
-  await browser.storage.local.set({ [SESSION_TOKEN_STORAGE_KEY]: token });
+/**
+ * Session storage. The token itself is obtained by wallet-auth.ts
+ * (nonce -> Phantom signature -> verify) once a wallet connects — this
+ * module just persists/retrieves it, since ws-client.ts and popup code
+ * shouldn't know how a token was obtained. Stored alongside the wallet
+ * address it belongs to so callers can tell "already signed in" apart from
+ * "signed in as a *different* wallet than the one now connected" (account
+ * switching in Phantom) — a bare token can't distinguish those.
+ */
+export async function getSessionToken(): Promise<string> {
+  const session = await getStoredSession();
+  if (!session) throw new Error('Not authenticated yet — connect a wallet on a supported trading site.');
+  return session.token;
+}
+
+export async function getStoredSession(): Promise<StoredSession | undefined> {
+  const stored = await browser.storage.local.get(SESSION_STORAGE_KEY);
+  return stored[SESSION_STORAGE_KEY] as StoredSession | undefined;
+}
+
+export async function storeSession(token: string, walletAddress: string): Promise<void> {
+  await browser.storage.local.set({ [SESSION_STORAGE_KEY]: { token, walletAddress } satisfies StoredSession });
 }
 
 export async function clearSessionToken(): Promise<void> {
-  await browser.storage.local.remove(SESSION_TOKEN_STORAGE_KEY);
+  await browser.storage.local.remove(SESSION_STORAGE_KEY);
 }

src/background/wallet-auth.ts

@@ -1,7 +1,7 @@
 import { browser } from 'wxt/browser';
 import { BACKEND_HTTP_URL } from '@/shared/config';
 import type { NexaMessage, WalletSignResult } from '@/shared/messaging';
-import { storeSessionToken } from './backend-client';
+import { storeSession } from './backend-client';
 
 /**
  * Runs the REST auth flow (backend/CLAUDE.md, "Current milestone:
@@ -40,20 +40,25 @@ export async function handleWalletConnected(tabId: number | undefined, walletAdd
   if (!verifyRes.ok) throw new Error(`verify request failed: ${verifyRes.status}`);
   const { session_token: sessionToken } = (await verifyRes.json()) as { session_token: string };
 
-  await storeSessionToken(sessionToken);
+  await storeSession(sessionToken, walletAddress);
   console.debug('[nexa/wallet-auth] session token stored');
 }
 
-/** Prompts any open supported-site tab to attempt a silent (onlyIfTrusted) reconnect — used after the session is invalidated server-side. */
-export async function requestWalletReconnect(): Promise<void> {
+async function messageAllSupportedTabs(message: NexaMessage): Promise<void> {
   const tabs = await browser.tabs.query({ url: 'https://axiom.trade/*' });
   await Promise.allSettled(
     tabs
       .filter((tab): tab is typeof tab & { id: number } => tab.id != null)
-      .map((tab) =>
-        browser.tabs
-          .sendMessage(tab.id, { type: 'nexa:request-wallet-connect' } satisfies NexaMessage)
-          .catch(() => undefined),
-      ),
+      .map((tab) => browser.tabs.sendMessage(tab.id, message).catch(() => undefined)),
   );
 }
+
+/** Prompts any open supported-site tab to attempt a silent (onlyIfTrusted) reconnect — used after the session is invalidated server-side. */
+export async function requestWalletReconnect(): Promise<void> {
+  await messageAllSupportedTabs({ type: 'nexa:request-wallet-connect' });
+}
+
+/** Sign-out: tells any open supported-site tab to disconnect from Phantom (revokes this origin's trust, so the next silent connect correctly fails until the user reconnects). */
+export async function requestWalletDisconnect(): Promise<void> {
+  await messageAllSupportedTabs({ type: 'nexa:wallet-disconnect-request' });
+}

src/background/ws-client.ts

@@ -62,12 +62,15 @@ export interface WsClientOptions {
 export interface WsClientController {
   /** Cancels any pending backoff wait and connects immediately — call after obtaining a fresh token. */
   reconnectNow(): void;
+  /** Closes the current connection (if any) and stops auto-reconnecting until `reconnectNow()` is called — e.g. on sign-out. */
+  disconnect(): void;
 }
 
 export function connectWsClient(options: WsClientOptions): WsClientController {
   let reconnectAttempt = 0;
   let reconnectTimer: ReturnType<typeof setTimeout> | undefined;
   let currentSocket: WebSocket | undefined;
+  let manuallyDisconnected = false;
 
   function scheduleReconnect(): void {
     const delay = RECONNECT_DELAYS_MS[Math.min(reconnectAttempt, RECONNECT_DELAYS_MS.length - 1)];
@@ -131,6 +134,11 @@ export function connectWsClient(options: WsClientOptions): WsClientController {
     socket.addEventListener('close', (event) => {
       if (currentSocket === socket) currentSocket = undefined;
 
+      if (manuallyDisconnected) {
+        options.onStatusChange('disconnected');
+        return; // sign-out — wait for reconnectNow() after a new wallet connects
+      }
+
       if (authFailure || event.code === 4001) {
         options.onStatusChange('auth-error');
         options.onAuthExpired();
@@ -148,13 +156,21 @@ export function connectWsClient(options: WsClientOptions): WsClientController {
   }
 
   function reconnectNow(): void {
+    manuallyDisconnected = false;
     if (reconnectTimer) clearTimeout(reconnectTimer);
     reconnectAttempt = 0;
     if (currentSocket && currentSocket.readyState === WebSocket.OPEN) return;
     void connect();
   }
 
+  function disconnect(): void {
+    manuallyDisconnected = true;
+    if (reconnectTimer) clearTimeout(reconnectTimer);
+    options.onStatusChange('disconnected');
+    currentSocket?.close(1000, 'client sign-out');
+  }
+
   void connect();
 
-  return { reconnectNow };
+  return { reconnectNow, disconnect };
 }

src/content-scripts/wallet-bridge/inject.ts

@@ -16,18 +16,21 @@ import bs58 from 'bs58';
 
 const CALL_CHANNEL = 'nexa:wallet-call';
 const RESULT_CHANNEL = 'nexa:wallet-result';
+const EVENT_CHANNEL = 'nexa:wallet-event';
 
 interface WalletCallMessage {
   channel: typeof CALL_CHANNEL;
   id: string;
-  action: 'connect' | 'signMessage';
-  payload: { onlyIfTrusted?: boolean } | { message: string };
+  action: 'connect' | 'signMessage' | 'disconnect';
+  payload: { onlyIfTrusted?: boolean } | { message: string } | Record<string, never>;
 }
 
 interface PhantomProvider {
   isPhantom?: boolean;
   connect(opts?: { onlyIfTrusted?: boolean }): Promise<{ publicKey: { toString(): string } }>;
+  disconnect?(): Promise<void>;
   signMessage(message: Uint8Array, display?: string): Promise<{ signature: Uint8Array }>;
+  on?(event: 'accountChanged' | 'disconnect', handler: (publicKey?: { toString(): string } | null) => void): void;
 }
 
 function getProvider(): PhantomProvider | undefined {
@@ -44,9 +47,30 @@ function respond(id: string, ok: boolean, dataOrError: unknown): void {
   );
 }
 
+/** Wires Phantom's own account-switch/disconnect events so the isolated world finds out even when nothing on our side initiated it. */
+function wireProviderEvents(provider: PhantomProvider): void {
+  provider.on?.('accountChanged', (publicKey) => {
+    const walletAddress = publicKey ? publicKey.toString() : null;
+    console.debug('[nexa/wallet-bridge] accountChanged', walletAddress);
+    window.postMessage({ channel: EVENT_CHANNEL, name: 'accountChanged', walletAddress }, window.location.origin);
+  });
+
+  provider.on?.('disconnect', () => {
+    console.debug('[nexa/wallet-bridge] disconnect event');
+    window.postMessage({ channel: EVENT_CHANNEL, name: 'disconnect', walletAddress: null }, window.location.origin);
+  });
+}
+
 export function initWalletBridgeInjected(): void {
   console.debug('[nexa/wallet-bridge] injected script active on', window.location.href);
 
+  // Phantom usually injects by document_start too, but ordering between two
+  // extensions' content scripts isn't guaranteed — retry once shortly after
+  // if it wasn't there yet, so accountChanged/disconnect still get wired up.
+  const provider = getProvider();
+  if (provider) wireProviderEvents(provider);
+  else setTimeout(() => getProvider() && wireProviderEvents(getProvider()!), 500);
+
   window.addEventListener('message', (event) => {
     if (event.source !== window) return; // ignore iframes/other windows
     const message = event.data as Partial<WalletCallMessage> | undefined;
@@ -57,18 +81,21 @@ export function initWalletBridgeInjected(): void {
 
     void (async () => {
       try {
-        const provider = getProvider();
-        if (!provider) throw new Error('No Solana wallet (Phantom) detected on this page.');
+        const activeProvider = getProvider();
+        if (!activeProvider) throw new Error('No Solana wallet (Phantom) detected on this page.');
 
         if (action === 'connect') {
           const { onlyIfTrusted } = payload as { onlyIfTrusted?: boolean };
-          const result = await provider.connect(onlyIfTrusted ? { onlyIfTrusted: true } : undefined);
+          const result = await activeProvider.connect(onlyIfTrusted ? { onlyIfTrusted: true } : undefined);
           console.debug('[nexa/wallet-bridge] connect ok', result.publicKey.toString());
           respond(id, true, { walletAddress: result.publicKey.toString() });
         } else if (action === 'signMessage') {
           const { message: nonce } = payload as { message: string };
-          const { signature } = await provider.signMessage(new TextEncoder().encode(nonce), 'utf8');
+          const { signature } = await activeProvider.signMessage(new TextEncoder().encode(nonce), 'utf8');
           respond(id, true, { signature: bs58.encode(signature) });
+        } else if (action === 'disconnect') {
+          await activeProvider.disconnect?.();
+          respond(id, true, {});
         } else {
           throw new Error(`Unknown wallet action: ${String(action)}`);
         }

src/content-scripts/wallet-bridge/relay.ts

@@ -10,6 +10,7 @@
 
 const CALL_CHANNEL = 'nexa:wallet-call';
 const RESULT_CHANNEL = 'nexa:wallet-result';
+const EVENT_CHANNEL = 'nexa:wallet-event';
 
 interface WalletResultMessage {
   channel: typeof RESULT_CHANNEL;
@@ -19,9 +20,15 @@ interface WalletResultMessage {
   error?: string;
 }
 
+interface WalletEventMessage {
+  channel: typeof EVENT_CHANNEL;
+  name: string;
+  walletAddress: string | null;
+}
+
 const WALLET_CALL_TIMEOUT_MS = 30_000;
 
-export function callWallet<T>(action: 'connect' | 'signMessage', payload: unknown): Promise<T> {
+export function callWallet<T>(action: 'connect' | 'signMessage' | 'disconnect', payload: unknown): Promise<T> {
   return new Promise((resolve, reject) => {
     const id = crypto.randomUUID();
 
@@ -49,3 +56,13 @@ export function callWallet<T>(action: 'connect' | 'signMessage', payload: unknow
     window.postMessage({ channel: CALL_CHANNEL, id, action, payload }, window.location.origin);
   });
 }
+
+/** Unsolicited events from Phantom itself (account switched in its UI, or disconnected there) — not a response to any `callWallet` request. */
+export function onWalletEvent(handler: (name: string, walletAddress: string | null) => void): void {
+  window.addEventListener('message', (event) => {
+    if (event.source !== window) return;
+    const message = event.data as Partial<WalletEventMessage> | undefined;
+    if (message?.channel !== EVENT_CHANNEL) return;
+    handler(message.name ?? '', message.walletAddress ?? null);
+  });
+}

src/content-scripts/wallet-connect.ts

@@ -1,6 +1,6 @@
 import { browser } from 'wxt/browser';
 import type { NexaMessage } from '@/shared/messaging';
-import { callWallet } from './wallet-bridge/relay';
+import { callWallet, onWalletEvent } from './wallet-bridge/relay';
 import { hideConnectBanner, setConnectBannerError, showConnectBanner } from './wallet-bridge/banner';
 
 interface ConnectResult {
@@ -19,6 +19,15 @@ async function reportConnected(walletAddress: string): Promise<void> {
     .catch((err) => console.debug('[nexa/wallet-connect] failed to notify background:', err));
 }
 
+/** Wallet disconnected or switched to no account — clears the backend session (tied to the old wallet) and re-shows the connect prompt. */
+async function reportDisconnected(): Promise<void> {
+  console.debug('[nexa/wallet-connect] disconnected');
+  await browser.runtime
+    .sendMessage({ type: 'nexa:wallet-disconnected' } satisfies NexaMessage)
+    .catch((err) => console.debug('[nexa/wallet-connect] failed to notify background:', err));
+  showConnectBanner(() => void connectWithGesture());
+}
+
 /** Real user gesture (banner button click) — required for Phantom to show its connect approval popup on a first-ever connect. */
 async function connectWithGesture(): Promise<void> {
   console.debug('[nexa/wallet-connect] banner clicked, calling connect()');
@@ -52,6 +61,18 @@ export function initWalletConnect(): void {
   console.debug('[nexa/wallet-connect] init on', window.location.href);
   void attemptSilentConnect();
 
+  // Phantom's own account-switch/disconnect events — not initiated by us, so
+  // this catches the user changing accounts (or disconnecting) directly in
+  // Phantom's UI, not just our own sign-out flow below.
+  onWalletEvent((name, walletAddress) => {
+    console.debug('[nexa/wallet-connect] wallet event', name, walletAddress);
+    if (walletAddress) {
+      void reportConnected(walletAddress); // switched to a different account — re-auth as it
+    } else {
+      void reportDisconnected(); // accountChanged(null) or a 'disconnect' event
+    }
+  });
+
   browser.runtime.onMessage.addListener((message: NexaMessage, _sender, sendResponse) => {
     if (message?.type === 'nexa:wallet-sign-request') {
       callWallet<SignResult>('signMessage', { message: message.nonce })
@@ -65,6 +86,16 @@ export function initWalletConnect(): void {
       return false;
     }
 
+    if (message?.type === 'nexa:wallet-disconnect-request') {
+      // Sign-out, initiated from the popup (see background/wallet-auth.ts).
+      // Phantom's disconnect() revokes this origin's trust, so the next
+      // onlyIfTrusted attempt correctly fails until the user connects again.
+      callWallet('disconnect', {})
+        .catch((err) => console.debug('[nexa/wallet-connect] disconnect() failed:', err))
+        .finally(() => void reportDisconnected());
+      return false;
+    }
+
     return undefined;
   });
 }

src/entrypoints/background.ts

@@ -1,7 +1,7 @@
-import { clearSessionToken, getSessionToken } from '@/background/backend-client';
+import { clearSessionToken, getSessionToken, getStoredSession } from '@/background/backend-client';
 import { getConnectionStatus, setConnectionStatus } from '@/background/connection-status';
 import { applyLockState, getLockState } from '@/background/lock-state';
-import { handleWalletConnected, requestWalletReconnect } from '@/background/wallet-auth';
+import { handleWalletConnected, requestWalletDisconnect, requestWalletReconnect } from '@/background/wallet-auth';
 import { connectWsClient } from '@/background/ws-client';
 import type { NexaMessage } from '@/shared/messaging';
 
@@ -15,6 +15,12 @@ export default defineBackground(() => {
     },
   });
 
+  async function signOut(): Promise<void> {
+    await clearSessionToken();
+    ws.disconnect();
+    await requestWalletDisconnect();
+  }
+
   browser.runtime.onMessage.addListener((message: NexaMessage, sender, sendResponse) => {
     switch (message?.type) {
       case 'nexa:get-lock-state':
@@ -33,17 +39,28 @@ export default defineBackground(() => {
 
       case 'nexa:wallet-connected':
         // A content script reports this on every page load (it always tries
-        // a silent onlyIfTrusted connect first) — but re-authenticating
-        // means asking Phantom to sign a fresh nonce, which shows its own
-        // approval popup every time, unlike a silent connect. Only pay that
-        // cost when we don't already have a usable session.
-        getSessionToken()
-          .then(() => undefined) // already authenticated — nothing to do
-          .catch(() =>
-            handleWalletConnected(sender.tab?.id, message.walletAddress)
-              .then(() => ws.reconnectNow())
-              .catch((err) => console.debug('[nexa/background] wallet auth failed:', err)), // ws-client's own retry loop keeps trying regardless
-          );
+        // a silent onlyIfTrusted connect first), AND whenever the user
+        // switches accounts in Phantom's own UI. Only re-authenticate (which
+        // means asking Phantom to sign a fresh nonce — a popup every time,
+        // unlike a silent connect) when this isn't the wallet we're already
+        // signed in as; a bare "do we have a token" check can't tell those
+        // apart from an account switch.
+        getStoredSession()
+          .then((session) => {
+            if (session?.walletAddress === message.walletAddress) return;
+            return handleWalletConnected(sender.tab?.id, message.walletAddress).then(() => ws.reconnectNow());
+          })
+          .catch((err) => console.debug('[nexa/background] wallet auth failed:', err)); // ws-client's own retry loop keeps trying regardless
+        return false;
+
+      case 'nexa:wallet-disconnected':
+        // Disconnected directly in Phantom's UI (not via our own sign-out
+        // flow, which already clears/disconnects itself) — treat the same way.
+        void clearSessionToken().then(() => ws.disconnect());
+        return false;
+
+      case 'nexa:sign-out':
+        void signOut();
         return false;
 
       default:

src/entrypoints/popup/App.tsx

@@ -34,6 +34,10 @@ export function App() {
     return () => browser.runtime.onMessage.removeListener(listener);
   }, []);
 
+  function signOut(): void {
+    void browser.runtime.sendMessage({ type: 'nexa:sign-out' } satisfies NexaMessage);
+  }
+
   return (
     <>
       <header>
@@ -59,6 +63,12 @@ export function App() {
         )}
       </section>
 
+      {connectionStatus === 'connected' && (
+        <button type="button" className="sign-out-button" onClick={signOut}>
+          Sign out
+        </button>
+      )}
+
       <a
         className="thresholds-link"
         href="#"

src/entrypoints/popup/style.css

@@ -79,6 +79,24 @@ header .subtitle {
   opacity: 0.65;
 }
 
+.sign-out-button {
+  display: block;
+  width: 100%;
+  margin-top: 4px;
+  padding: 6px 10px;
+  border: 1px solid rgba(128, 128, 128, 0.4);
+  border-radius: 6px;
+  background: transparent;
+  color: inherit;
+  font: inherit;
+  font-weight: 600;
+  cursor: pointer;
+}
+
+.sign-out-button:hover {
+  background: rgba(128, 128, 128, 0.12);
+}
+
 .thresholds-link {
   display: block;
   margin-top: 14px;

src/shared/messaging.ts

@@ -21,8 +21,12 @@ export type NexaMessage =
   // Wallet auth (content script <-> background). See
   // content-scripts/wallet-connect.ts and background/wallet-auth.ts.
   | { type: 'nexa:wallet-connected'; walletAddress: string }
+  | { type: 'nexa:wallet-disconnected' }
   | { type: 'nexa:wallet-sign-request'; nonce: string }
-  | { type: 'nexa:request-wallet-connect' };
+  | { type: 'nexa:request-wallet-connect' }
+  | { type: 'nexa:wallet-disconnect-request' }
+  // Sign-out (popup -> background). See entrypoints/popup/App.tsx.
+  | { type: 'nexa:sign-out' };
 
 /** Response shape for 'nexa:wallet-sign-request', returned via sendResponse (not a dispatched NexaMessage). */
 export type WalletSignResult = { signature: string } | { error: string };