fix switching networks back and forth results in blocks not loading

src/chain.state.ts

@@ -12,7 +12,7 @@ import { createClient } from "polkadot-api"
 import { withLogsRecorder } from "polkadot-api/logs-provider"
 import { withPolkadotSdkCompat } from "polkadot-api/polkadot-sdk-compat"
 import { getSmProvider } from "polkadot-api/sm-provider"
-import { Chain, Client } from "polkadot-api/smoldot"
+import { Chain } from "polkadot-api/smoldot"
 import { startFromWorker } from "polkadot-api/smoldot/from-worker"
 import SmWorker from "polkadot-api/smoldot/worker?worker"
 import { getWsProvider } from "polkadot-api/ws-provider/web"
@@ -29,18 +29,39 @@ import {
   switchMap,
   tap,
 } from "rxjs"
-
 import ksmRawNetworks from "./networks/kusama.json"
-import polkadotRawNetworks from "./networks/polkadot.json"
 import paseoRawNetworks from "./networks/paseo.json"
+import polkadotRawNetworks from "./networks/polkadot.json"
 import westendRawNetworks from "./networks/westend.json"
 
-const [chainSpec, ksmChainSpec, paseoChainSpec, westendChainSpec] = [
+const [dotChainSpec, ksmChainSpec, paseoChainSpec, westendChainSpec] = [
   import("polkadot-api/chains/polkadot"),
   import("polkadot-api/chains/ksmcc3"),
   import("polkadot-api/chains/paseo"),
   import("polkadot-api/chains/westend2"),
 ].map((x) => x.then((y) => y.chainSpec))
+const smoldot = startFromWorker(new SmWorker(), {
+  logCallback: (level, target, message) => {
+    console.debug("smoldot[%s(%s)] %s", target, level, message)
+  },
+})
+const relayChains: Record<
+  string,
+  { chainSpec: Promise<string>; chain: Promise<Chain> | null }
+> = {
+  polkadot: { chainSpec: dotChainSpec, chain: null },
+  kusama: { chainSpec: ksmChainSpec, chain: null },
+  paseo: { chainSpec: paseoChainSpec, chain: null },
+  westend: { chainSpec: westendChainSpec, chain: null },
+}
+const getRelayChain = async (name: string) => {
+  if (!relayChains[name].chain) {
+    relayChains[name].chain = smoldot.addChain({
+      chainSpec: await relayChains[name].chainSpec,
+    })
+  }
+  return relayChains[name].chain
+}
 
 export type ChainSource = { id: string } & (
   | {
@@ -119,7 +140,6 @@ export const selectedChain$ = state<SelectedChain>(selectedChainChanged$, {
 })
 selectedChain$.subscribe()
 
-const relayChains = new Set(["polkadot", "kusama", "westend", "paseo"])
 const selectedSource$ = selectedChain$.pipe(
   switchMap(
     ({ endpoint, network }): Observable<ChainSource> | Promise<ChainSource> => {
@@ -130,12 +150,12 @@ const selectedSource$ = selectedChain$.pipe(
           type: "websocket",
           value: endpoint,
         })
-      if (relayChains.has(id)) {
-        return of({
+      if (id in relayChains) {
+        return relayChains[id].chainSpec.then((chainSpec) => ({
           id,
           type: "chainSpec",
-          value: { chainSpec: id },
-        })
+          value: { chainSpec },
+        }))
       }
       return import(`./chainspecs/${id}.ts`).then(({ chainSpec }) => {
         const parsed = JSON.parse(chainSpec)
@@ -240,46 +260,24 @@ export const dynamicBuilder$ = runtimeCtx$.pipeState(
   map((ctx) => ctx.dynamicBuilder),
 )
 
-let smoldot: {
-  client: Client
-  relayChains: Record<string, Promise<Chain>>
-} | null = null
 export function getProvider(source: ChainSource): JsonRpcProvider {
   if (source.type === "websocket") {
     return withPolkadotSdkCompat(getWsProvider(source.value))
   }
 
-  if (!smoldot) {
-    const client = startFromWorker(new SmWorker(), {
-      logCallback: (level, target, message) => {
-        console.debug("smoldot[%s(%s)] %s", target, level, message)
-      },
-    })
-    const createChain = async (input: Promise<string>) =>
-      client.addChain({
-        chainSpec: await input,
-      })
-    smoldot = {
-      client,
-      relayChains: {
-        polkadot: createChain(chainSpec),
-        kusama: createChain(ksmChainSpec),
-        westend: createChain(westendChainSpec),
-        paseo: createChain(paseoChainSpec),
-      },
-    }
-  }
   const chain = source.value.relayChain
-    ? smoldot.relayChains[source.value.relayChain].then((chain) => {
-        return smoldot!.client.addChain({
+    ? getRelayChain(source.value.relayChain).then((chain) =>
+        smoldot.addChain({
           chainSpec: source.value.chainSpec,
           potentialRelayChains: [chain],
-        })
-      })
-    : smoldot.relayChains[source.value.chainSpec] ||
-      smoldot.client.addChain({
+        }),
+      )
+    : smoldot.addChain({
         chainSpec: source.value.chainSpec,
       })
 
-  return withLogsRecorder(console.debug, getSmProvider(chain))
+  return withLogsRecorder(
+    (v) => v.includes("initialized") && console.debug(v),
+    getSmProvider(chain),
+  )
 }

src/pages/Network/Network.tsx

@@ -1,6 +1,19 @@
-import { useEffect, useRef, useState } from "react"
 import { Check, ChevronDown } from "lucide-react"
+import { FC, useEffect, useRef, useState } from "react"
 
+import {
+  Network,
+  networkCategories,
+  onChangeChain,
+  SelectedChain,
+  selectedChain$,
+} from "@/chain.state"
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from "@/components/ui/accordion"
 import { Button } from "@/components/ui/button"
 import {
   Command,
@@ -17,22 +30,10 @@ import {
   DialogTitle,
   DialogTrigger,
 } from "@/components/ui/dialog"
-import {
-  Accordion,
-  AccordionContent,
-  AccordionItem,
-  AccordionTrigger,
-} from "@/components/ui/accordion"
 import { Label } from "@/components/ui/label"
 import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"
 import { ScrollArea } from "@/components/ui/scroll-area"
 import { useStateObservable } from "@react-rxjs/core"
-import {
-  Network,
-  networkCategories,
-  onChangeChain,
-  selectedChain$,
-} from "@/chain.state"
 import { useCommandState } from "cmdk"
 
 const isValidUri = (input: string): boolean => {
@@ -99,6 +100,32 @@ const Empty: React.FC<{
 export function NetworkSwitcher() {
   const [open, setOpen] = useState(false)
   const selectedChain = useStateObservable(selectedChain$)
+
+  return (
+    <Dialog open={open} onOpenChange={setOpen}>
+      <DialogTrigger asChild>
+        <Button
+          variant="outline"
+          className="w-[200px] flex gap-0 justify-between text-base px-3 border border-border bg-input"
+        >
+          <span className="overflow-hidden text-ellipsis">
+            {selectedChain.network.display}
+          </span>
+          <ChevronDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+        </Button>
+      </DialogTrigger>
+      <NetworkSwitchDialogContent
+        selectedChain={selectedChain}
+        onClose={() => setOpen(false)}
+      />
+    </Dialog>
+  )
+}
+
+const NetworkSwitchDialogContent: FC<{
+  selectedChain: SelectedChain
+  onClose: () => void
+}> = ({ selectedChain, onClose }) => {
   const [selectedNetwork, setSelectedNetwork] = useState<Network>(
     selectedChain.network,
   )
@@ -122,119 +149,106 @@ export function NetworkSwitcher() {
   }
 
   const handleConfirm = () => {
-    setOpen(false)
     onChangeChain({
       network: selectedNetwork,
       endpoint: selectedRpc,
     })
+    onClose()
   }
 
   return (
-    <Dialog open={open} onOpenChange={setOpen}>
-      <DialogTrigger asChild>
-        <Button
-          variant="outline"
-          className="w-[200px] flex gap-0 justify-between text-base px-3 border border-border bg-input"
-        >
-          <span className="overflow-hidden text-ellipsis">
-            {selectedNetwork.display}
-          </span>
-          <ChevronDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
-        </Button>
-      </DialogTrigger>
-      <DialogContent className="sm:max-w-[425px]">
-        <DialogHeader>
-          <DialogTitle>Switch Network</DialogTitle>
-        </DialogHeader>
-        <Command className="rounded-lg border shadow-md">
-          <CommandInput
-            placeholder="Search or enter a custom URI"
-            value={enteredText}
-            onValueChange={setEnteredText}
+    <DialogContent className="sm:max-w-[425px]">
+      <DialogHeader>
+        <DialogTitle>Switch Network</DialogTitle>
+      </DialogHeader>
+      <Command className="rounded-lg border shadow-md">
+        <CommandInput
+          placeholder="Search or enter a custom URI"
+          value={enteredText}
+          onValueChange={setEnteredText}
+        />
+        <CommandList>
+          <EmptyOption
+            {...{
+              enteredText,
+              selectedNetwork,
+              selectedRpc,
+              setSelectedRpc,
+              setSelectedNetwork,
+            }}
           />
-          <CommandList>
-            <EmptyOption
-              {...{
-                enteredText,
-                selectedNetwork,
-                selectedRpc,
-                setSelectedRpc,
-                setSelectedNetwork,
-              }}
-            />
-            <ScrollArea className="h-[260px]">
-              {networkCategories.map((category) => (
-                <CommandGroup key={category.name} heading={category.name}>
-                  {category.networks.map((network) => (
-                    <CommandItem
-                      key={network.id}
-                      onSelect={() => handleNetworkSelect(network)}
-                    >
-                      <Check
-                        className={`mr-2 h-4 w-4 ${
-                          selectedNetwork.id === network.id
-                            ? "opacity-100"
-                            : "opacity-0"
-                        }`}
+          <ScrollArea className="h-[260px]">
+            {networkCategories.map((category) => (
+              <CommandGroup key={category.name} heading={category.name}>
+                {category.networks.map((network) => (
+                  <CommandItem
+                    key={network.id}
+                    onSelect={() => handleNetworkSelect(network)}
+                  >
+                    <Check
+                      className={`mr-2 h-4 w-4 ${
+                        selectedNetwork.id === network.id
+                          ? "opacity-100"
+                          : "opacity-0"
+                      }`}
+                    />
+                    {network.display}
+                  </CommandItem>
+                ))}
+              </CommandGroup>
+            ))}
+          </ScrollArea>
+        </CommandList>
+      </Command>
+      {selectedNetwork && selectedNetwork.id !== "custom-network" && (
+        <Accordion type="single" collapsible className="w-full -mt-3">
+          <AccordionItem value="connection-options">
+            <AccordionTrigger>Connection Options</AccordionTrigger>
+            <AccordionContent>
+              <ScrollArea className="h-[155px] rounded-lg border p-2">
+                <RadioGroup>
+                  {selectedNetwork.lightclient ? (
+                    <div className="flex items-center space-x-2">
+                      <RadioGroupItem
+                        value="light-client"
+                        id="light-client"
+                        checked={selectedRpc === "light-client"}
+                        onClick={() => setSelectedRpc("light-client")}
                       />
-                      {network.display}
-                    </CommandItem>
-                  ))}
-                </CommandGroup>
-              ))}
-            </ScrollArea>
-          </CommandList>
-        </Command>
-        {selectedNetwork && selectedNetwork.id !== "custom-network" && (
-          <Accordion type="single" collapsible className="w-full -mt-3">
-            <AccordionItem value="connection-options">
-              <AccordionTrigger>Connection Options</AccordionTrigger>
-              <AccordionContent>
-                <ScrollArea className="h-[155px] rounded-lg border p-2">
-                  <RadioGroup>
-                    {selectedNetwork.lightclient ? (
-                      <div className="flex items-center space-x-2">
+                      <Label htmlFor="light-client">
+                        Light Client (smoldot)
+                      </Label>
+                    </div>
+                  ) : null}
+                  {Object.entries(selectedNetwork.endpoints).map(
+                    ([rpcName, url]) => (
+                      <div key={url} className="flex items-center space-x-2">
                         <RadioGroupItem
-                          value="light-client"
-                          id="light-client"
-                          checked={selectedRpc === "light-client"}
-                          onClick={() => setSelectedRpc("light-client")}
+                          value={url}
+                          id={url}
+                          checked={selectedRpc === url}
+                          onClick={() => setSelectedRpc(url)}
                         />
-                        <Label htmlFor="light-client">
-                          Light Client (smoldot)
-                        </Label>
+                        <Label htmlFor={url}>{rpcName}</Label>
                       </div>
-                    ) : null}
-                    {Object.entries(selectedNetwork.endpoints).map(
-                      ([rpcName, url]) => (
-                        <div key={url} className="flex items-center space-x-2">
-                          <RadioGroupItem
-                            value={url}
-                            id={url}
-                            checked={selectedRpc === url}
-                            onClick={() => setSelectedRpc(url)}
-                          />
-                          <Label htmlFor={url}>{rpcName}</Label>
-                        </div>
-                      ),
-                    )}
-                  </RadioGroup>
-                </ScrollArea>
-              </AccordionContent>
-            </AccordionItem>
-          </Accordion>
-        )}
-        <Button
-          onClick={handleConfirm}
-          disabled={
-            !selectedNetwork ||
-            !hasChanged ||
-            (selectedNetwork.id === "custom-network" && !selectedRpc)
-          }
-        >
-          Confirm Selection
-        </Button>
-      </DialogContent>
-    </Dialog>
+                    ),
+                  )}
+                </RadioGroup>
+              </ScrollArea>
+            </AccordionContent>
+          </AccordionItem>
+        </Accordion>
+      )}
+      <Button
+        onClick={handleConfirm}
+        disabled={
+          !selectedNetwork ||
+          !hasChanged ||
+          (selectedNetwork.id === "custom-network" && !selectedRpc)
+        }
+      >
+        Confirm Selection
+      </Button>
+    </DialogContent>
   )
 }