fix(SorokitProvider): handle disconnect errors, re-init client on network switch

src/context/SorokitContext.ts

@@ -40,6 +40,18 @@ export interface SorokitState {
 export interface SorokitProviderProps {
   client: SorokitClient;
   onError?: (error: string, source: string) => void;
+  /**
+   * Called after every successful network switch. Use it for side effects the
+   * provider can't know about — clearing caches, re-subscribing to feeds —
+   * instead of watching the `network` value and depending on its shape.
+   */
+  onNetworkChange?: (network: NetworkInfo) => void;
+  /**
+   * Builds a client configured for a given network. When provided, a successful
+   * `switchNetwork` re-initialises the `getClient()` singleton with the result,
+   * so calls made after a switch reach the new network's endpoints.
+   */
+  createClientForNetwork?: (network: NetworkInfo) => SorokitClient;
   children: React.ReactNode;
 }
 

src/context/SorokitProvider.test.tsx

@@ -1,4 +1,5 @@
-import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";import { StrictMode, useRef, useState } from "react";
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { StrictMode, useRef, useState } from "react";
 import { beforeEach,describe, expect, it, vi } from "vitest";
 
 import { renderWithProvider } from "@/__tests__/utils";
@@ -542,4 +543,146 @@ describe("SorokitProvider", () => {
       expect(screen.getByTestId("errorHistoryCount")).toHaveTextContent("0");
     });
   });
+
+  describe("disconnect errors, client re-init, and onNetworkChange", () => {
+    it("surfaces a disconnect failure and still clears the session", async () => {
+      mockClient.wallet.disconnect = vi
+        .fn()
+        .mockRejectedValue(new Error("Wallet extension unavailable"));
+
+      renderWithProvider(<TestComponent />, { client: mockClient });
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Connect"));
+      });
+      await waitFor(() => {
+        expect(screen.getByTestId("address")).toHaveTextContent("GABC");
+      });
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Disconnect"));
+      });
+
+      // The error is reported...
+      expect(screen.getByTestId("error")).toHaveTextContent(
+        "Wallet extension unavailable",
+      );
+      // ...and the session is still torn down.
+      expect(screen.getByTestId("address")).toHaveTextContent("none");
+      expect(screen.getByTestId("account")).toHaveTextContent("none");
+      expect(screen.getByTestId("balances")).toHaveTextContent("0");
+    });
+
+    it("reports a non-Error rejection with a fallback message", async () => {
+      mockClient.wallet.disconnect = vi.fn().mockRejectedValue("boom");
+
+      renderWithProvider(<TestComponent />, { client: mockClient });
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Disconnect"));
+      });
+
+      expect(screen.getByTestId("error")).toHaveTextContent(
+        "Failed to disconnect wallet.",
+      );
+    });
+
+    it("points getClient() at the provider's client on mount", async () => {
+      await act(async () => {
+        renderWithProvider(<TestComponent />, { client: mockClient });
+      });
+
+      expect(getClient()).toBe(mockClient);
+    });
+
+    it("re-initialises getClient() with the new network's client after a switch", async () => {
+      const switchedClient = {
+        ...mockClient,
+      } as unknown as ReturnType<typeof getClient>;
+      const createClientForNetwork = vi.fn().mockReturnValue(switchedClient);
+
+      await act(async () => {
+        render(
+          <SorokitProvider
+            client={mockClient}
+            createClientForNetwork={createClientForNetwork}
+          >
+            <TestComponent />
+          </SorokitProvider>,
+        );
+      });
+
+      expect(getClient()).toBe(mockClient);
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Switch"));
+      });
+
+      expect(createClientForNetwork).toHaveBeenCalledWith({ name: "testnet" });
+      expect(getClient()).toBe(switchedClient);
+    });
+
+    it("keeps the existing client when no factory is provided", async () => {
+      await act(async () => {
+        renderWithProvider(<TestComponent />, { client: mockClient });
+      });
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Switch"));
+      });
+
+      expect(getClient()).toBe(mockClient);
+    });
+
+    it("fires onNetworkChange with the new network after a successful switch", async () => {
+      const onNetworkChange = vi.fn();
+
+      await act(async () => {
+        render(
+          <SorokitProvider
+            client={mockClient}
+            onNetworkChange={onNetworkChange}
+          >
+            <TestComponent />
+          </SorokitProvider>,
+        );
+      });
+      onNetworkChange.mockClear();
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Switch"));
+      });
+
+      expect(onNetworkChange).toHaveBeenCalledTimes(1);
+      expect(onNetworkChange).toHaveBeenCalledWith({ name: "testnet" });
+    });
+
+    it("does not fire onNetworkChange when the switch fails", async () => {
+      mockClient.network.switchNetwork = vi
+        .fn()
+        .mockResolvedValue({ data: null, error: "Invalid network: nope" });
+      const onNetworkChange = vi.fn();
+
+      await act(async () => {
+        render(
+          <SorokitProvider
+            client={mockClient}
+            onNetworkChange={onNetworkChange}
+          >
+            <TestComponent />
+          </SorokitProvider>,
+        );
+      });
+      onNetworkChange.mockClear();
+
+      await act(async () => {
+        fireEvent.click(screen.getByText("Switch"));
+      });
+
+      expect(onNetworkChange).not.toHaveBeenCalled();
+      expect(screen.getByTestId("error")).toHaveTextContent(
+        "Invalid network: nope",
+      );
+    });
+  });
 });

src/context/SorokitProvider.tsx

@@ -6,13 +6,20 @@ import type {
   NetworkInfo,
   NetworkName,
 } from "@/lib/client";
+import { initClient } from "@/lib/client";
 
 import { SorokitContext, type SorokitProviderProps } from "./SorokitContext";
 
 const STORAGE_KEY_NETWORK = "sorokit_network";
 const STORAGE_KEY_CUSTOM_NETWORKS = "sorokit_custom_networks";
 
-export function SorokitProvider({ client, onError, children }: SorokitProviderProps) {
+export function SorokitProvider({
+  client,
+  onError,
+  onNetworkChange,
+  createClientForNetwork,
+  children,
+}: SorokitProviderProps) {
   const [address, setAddress] = useState<string | null>(null);
   const [walletName, setWalletName] = useState<string | null>(null);
   const [isConnecting, setIsConnecting] = useState(false);
@@ -46,8 +53,21 @@ export function SorokitProvider({ client, onError, children }: SorokitProviderPr
   const clientRef = useRef(client);
   useEffect(() => {
     clientRef.current = client;
+    // Keep the module singleton pointed at the client the provider is using,
+    // so components that reach for `getClient()` share this one.
+    initClient(client);
   }, [client]);
 
+  const onNetworkChangeRef = useRef(onNetworkChange);
+  useEffect(() => {
+    onNetworkChangeRef.current = onNetworkChange;
+  }, [onNetworkChange]);
+
+  const createClientForNetworkRef = useRef(createClientForNetwork);
+  useEffect(() => {
+    createClientForNetworkRef.current = createClientForNetwork;
+  }, [createClientForNetwork]);
+
   // #353 — guards refreshAccount against overlapping calls (e.g. a user
   // clicking "Refresh" again before the previous request settles), so a
   // slower earlier response can't land after a newer one and show stale data.
@@ -186,18 +206,31 @@ export function SorokitProvider({ client, onError, children }: SorokitProviderPr
   const disconnectWallet = useCallback(async () => {
     setIsDisconnecting(true);
     try {
-      await clientRef.current.wallet.disconnect();
+      // A wallet adapter that throws (e.g. the extension went away
+      // mid-session) used to surface as an unhandled rejection. Capture it and
+      // still tear the session down — the user asked to disconnect.
+      let disconnectError: string | null = null;
+      try {
+        await clientRef.current.wallet.disconnect();
+      } catch (e) {
+        disconnectError =
+          e instanceof Error ? e.message : "Failed to disconnect wallet.";
+      }
+
       setAddress(null);
       setWalletName(null);
       setAccount(null);
       setBalances([]);
       // #353 — a fresh session shouldn't carry over error history from
       // whatever the previous wallet connection ran into.
       setErrorHistory([]);
+
+      // Reported after the reset so the failure isn't cleared along with it.
+      if (disconnectError) reportError(disconnectError, "wallet", "error");
     } finally {
       setIsDisconnecting(false);
     }
-  }, []);
+  }, [reportError]);
 
   const switchNetwork = useCallback(
     async (param: NetworkName | NetworkInfo) => {
@@ -209,6 +242,18 @@ export function SorokitProvider({ client, onError, children }: SorokitProviderPr
       if (data) {
         setError(null);
         setNetwork(data);
+
+        // Re-point the `getClient()` singleton at the new network. Without
+        // this, callers like `getClient().transaction.submit()` keep hitting
+        // the previous network's endpoints after a switch.
+        const nextClient = createClientForNetworkRef.current?.(data);
+        if (nextClient) {
+          clientRef.current = nextClient;
+          initClient(nextClient);
+        } else {
+          initClient(clientRef.current);
+        }
+
         try {
           window.localStorage.setItem(
             STORAGE_KEY_NETWORK,
@@ -220,6 +265,8 @@ export function SorokitProvider({ client, onError, children }: SorokitProviderPr
         setAddress(null);
         setAccount(null);
         setBalances([]);
+
+        onNetworkChangeRef.current?.(data);
       }
     },
     [reportError],