Use one websocket per client webpage

src/js/src/client.ts

@@ -1,42 +1,58 @@
 import {
   BaseReactPyClient,
-  type GenericReactPyClientProps,
   type ReactPyClientInterface,
   type ReactPyModule,
-  type ReactPyUrls,
 } from "@reactpy/client";
-import { createReconnectingWebSocket } from "./websocket";
+import { PageClient } from "./pageClient";
+import type { ComponentConfig } from "./mount";
+
+export type ReactPyDjangoClientProps = {
+  rootId: string;
+  pageClient: PageClient;
+  mountElement: HTMLElement;
+  jsModulesPath: string;
+  componentConfig: ComponentConfig;
+};
 
 export class ReactPyDjangoClient
   extends BaseReactPyClient
   implements ReactPyClientInterface
 {
-  urls: ReactPyUrls;
-  socket: { current?: WebSocket };
-  mountElement: HTMLElement;
-  prerenderElement: HTMLElement | null = null;
-  offlineElement: HTMLElement | null = null;
-  private readonly messageQueue: any[] = [];
+  readonly rootId: string;
+  readonly pageClient: PageClient;
+  readonly mountElement: HTMLElement;
+  readonly prerenderElement: HTMLElement | null = null;
+  readonly offlineElement: HTMLElement | null = null;
+  readonly jsModulesPath: string;
+  private readonly componentConfig: ComponentConfig;
 
-  constructor(props: GenericReactPyClientProps) {
+  constructor(props: ReactPyDjangoClientProps) {
     super();
 
-    this.urls = props.urls;
+    this.rootId = props.rootId;
+    this.pageClient = props.pageClient;
     this.mountElement = props.mountElement;
+    this.jsModulesPath = props.jsModulesPath;
+    this.componentConfig = props.componentConfig;
+
     this.prerenderElement = document.getElementById(
       props.mountElement.id + "-prerender",
     );
     this.offlineElement = document.getElementById(
       props.mountElement.id + "-offline",
     );
 
-    this.socket = createReconnectingWebSocket({
-      url: this.urls.componentUrl,
-      readyPromise: this.ready,
-      ...props.reconnectOptions,
-      // onMessage: Use standard ReactPy message routing
-      onMessage: async ({ data }) => this.handleIncoming(JSON.parse(data)),
-      // onClose: If offlineElement exists, show it and hide the mountElement/prerenderElement
+    // Register with the shared page client for message routing
+    this.pageClient.registerComponent(this.rootId, {
+      handleIncoming: (message: any) => this.handleIncoming(message),
+      onOpen: () => {
+        // Re-mount the component when the WebSocket reconnects
+        this.sendMountMessage();
+        if (this.offlineElement && this.mountElement) {
+          this.offlineElement.hidden = true;
+          this.mountElement.hidden = false;
+        }
+      },
       onClose: () => {
         if (this.prerenderElement) {
           this.prerenderElement.remove();
@@ -47,28 +63,36 @@ export class ReactPyDjangoClient
           this.offlineElement.hidden = false;
         }
       },
-      // onOpen: If offlineElement exists, hide it and show the mountElement
-      onOpen: () => {
-        if (this.offlineElement && this.mountElement) {
-          this.offlineElement.hidden = true;
-          this.mountElement.hidden = false;
-        }
-      },
+    });
+
+    // If the shared socket is already open, send the mount message immediately.
+    // This handles dynamically-added components registered after the initial
+    // page render (e.g. lazy-loaded tabs or conditionally rendered components).
+    if (this.pageClient.socket.current?.readyState === WebSocket.OPEN) {
+      this.sendMountMessage();
+    }
+  }
+
+  /** Send a mount-component message to the server so it constructs this component. */
+  sendMountMessage(): void {
+    this.pageClient.sendMessage(this.rootId, {
+      type: "mount-component",
+      rootId: this.rootId,
+      dottedPath: this.componentConfig.dottedPath,
+      componentUuid: this.componentConfig.componentUuid,
+      hasArgs: Boolean(this.componentConfig.hasArgs),
     });
   }
 
   sendMessage(message: any): void {
-    if (
-      this.socket.current &&
-      this.socket.current.readyState === WebSocket.OPEN
-    ) {
-      this.socket.current.send(JSON.stringify(message));
-    } else {
-      this.messageQueue.push(message);
-    }
+    this.pageClient.sendMessage(this.rootId, message);
   }
 
   loadModule(moduleName: string): Promise<ReactPyModule> {
-    return import(`${this.urls.jsModulesPath}${moduleName}`);
+    return import(`${this.jsModulesPath}${moduleName}`);
+  }
+
+  destroy(): void {
+    this.pageClient.unregisterComponent(this.rootId);
   }
 }

src/js/src/mount.tsx

@@ -1,51 +1,66 @@
 import { ReactPyDjangoClient } from "./client";
+import { getPageClient } from "./pageClient";
 import { Layout, React } from "@reactpy/client";
 
+export type ComponentConfig = {
+  dottedPath: string;
+  componentUuid: string;
+  hasArgs: number;
+};
+
 export function mountComponent(
   mountElement: HTMLElement,
   host: string,
   urlPrefix: string,
-  componentPath: string,
+  componentConfig: ComponentConfig,
   resolvedJsModulesPath: string,
   reconnectInterval: number,
   reconnectMaxInterval: number,
   reconnectMaxRetries: number,
   reconnectBackoffMultiplier: number,
 ) {
-  // WebSocket route for component rendering
+  // Shared WebSocket route per page
   const wsProtocol = `ws${window.location.protocol === "https:" ? "s" : ""}:`;
   const wsOrigin = host
     ? `${wsProtocol}//${host}`
     : `${wsProtocol}//${window.location.host}`;
-  const componentUrl = new URL(`${wsOrigin}/${urlPrefix}/${componentPath}`);
+  const wsUrl = new URL(`${wsOrigin}/${urlPrefix}/`);
 
   // Embed the initial HTTP path into the WebSocket URL
-  componentUrl.searchParams.append("path", window.location.pathname);
+  wsUrl.searchParams.set("path", window.location.pathname);
   if (window.location.search) {
-    componentUrl.searchParams.append("qs", window.location.search);
+    wsUrl.searchParams.set("qs", window.location.search);
   }
 
   // HTTP route for JavaScript modules
   const httpProtocol = window.location.protocol;
   const httpOrigin: string = host
     ? `${httpProtocol}//${host}`
     : `${httpProtocol}//${window.location.host}`;
-  const jsModulesPath: string =
-    resolvedJsModulesPath || `${urlPrefix}/web_module/`;
+  const jsModulesPath: string = resolvedJsModulesPath
+    ? `${httpOrigin}${resolvedJsModulesPath.startsWith("/") ? "" : "/"}${resolvedJsModulesPath}`
+    : `${httpOrigin}/${urlPrefix}/web_module/`;
 
-  // Configure a new ReactPy client
-  const client = new ReactPyDjangoClient({
-    urls: {
-      componentUrl: componentUrl,
-      jsModulesPath: `${httpOrigin}${jsModulesPath.startsWith("/") ? "" : "/"}${jsModulesPath}`,
-    },
-    reconnectOptions: {
+  // Get or create the shared page client (one per unique WS URL)
+  const pageClient = getPageClient(
+    wsUrl,
+    {
       interval: reconnectInterval,
       maxInterval: reconnectMaxInterval,
       maxRetries: reconnectMaxRetries,
       backoffMultiplier: reconnectBackoffMultiplier,
     },
-    mountElement: mountElement,
+    jsModulesPath,
+  );
+
+  // Create a per-component client that shares the page WebSocket
+  const rootId = componentConfig.componentUuid;
+  const client = new ReactPyDjangoClient({
+    rootId,
+    pageClient,
+    mountElement,
+    jsModulesPath,
+    componentConfig,
   });
 
   // Replace the prerender element with the real element on the first layout update
@@ -58,6 +73,10 @@ export function mountComponent(
     });
   }
 
+  // The mount-component message will be sent automatically when the shared WebSocket
+  // connects (via the onOpen callback in the ReactPyDjangoClient constructor).
+  // On reconnect the same callback handles re-mounting.
+
   // Start rendering the component
   if (client.mountElement) {
     React.render(<Layout client={client} />, client.mountElement);

src/js/src/pageClient.ts

@@ -0,0 +1,148 @@
+import { type ReactPyModule } from "@reactpy/client";
+import { createReconnectingWebSocket } from "./websocket";
+
+type ComponentRecord = {
+  handleIncoming: (message: any) => void;
+  onOpen?: () => void;
+  onClose?: () => void;
+};
+
+const pageClients = new Map<string, PageClient>();
+
+export function getPageClient(
+  wsUrl: URL,
+  reconnectOptions: {
+    interval: number;
+    maxInterval: number;
+    maxRetries: number;
+    backoffMultiplier: number;
+  },
+  jsModulesPath: string,
+): PageClient {
+  const key = `${wsUrl.origin}${wsUrl.pathname}`;
+  if (!pageClients.has(key)) {
+    pageClients.set(
+      key,
+      new PageClient(key, wsUrl, reconnectOptions, jsModulesPath),
+    );
+  }
+  return pageClients.get(key)!;
+}
+
+/**
+ * Dispose a cached PageClient, close its WebSocket, and remove it from the
+ * global map. Useful during SPA page transitions where the old page's WS
+ * should no longer be reused.
+ */
+export function disposePageClient(key: string): void {
+  const client = pageClients.get(key);
+  if (client) {
+    client.close();
+    pageClients.delete(key);
+  }
+}
+
+export class PageClient {
+  private components: Map<string, ComponentRecord> = new Map();
+  socket: { current?: WebSocket } = {};
+  private readonly messageQueue: any[] = [];
+  private readonly jsModulesPath: string;
+
+  constructor(
+    private readonly key: string,
+    private wsUrl: URL,
+    reconnectOptions: {
+      interval: number;
+      maxInterval: number;
+      maxRetries: number;
+      backoffMultiplier: number;
+    },
+    jsModulesPath: string,
+  ) {
+    this.jsModulesPath = jsModulesPath;
+
+    // Immediately-resolved promise — the shared socket connects right away
+    const immediatePromise = Promise.resolve();
+
+    this.socket = createReconnectingWebSocket({
+      url: wsUrl,
+      readyPromise: immediatePromise,
+      ...reconnectOptions,
+      onMessage: async ({ data }) => this.handleIncoming(JSON.parse(data)),
+      onOpen: () => {
+        // Notify all registered components that the connection is back
+        for (const [, record] of this.components) {
+          if (record.onOpen) record.onOpen();
+        }
+        // Drain queued messages
+        while (this.messageQueue.length > 0) {
+          this.sendRaw(this.messageQueue.shift());
+        }
+      },
+      onClose: () => {
+        for (const [, record] of this.components) {
+          if (record.onClose) record.onClose();
+        }
+      },
+    });
+  }
+
+  registerComponent(rootId: string, record: ComponentRecord): void {
+    this.components.set(rootId, record);
+  }
+
+  unregisterComponent(rootId: string): void {
+    this.components.delete(rootId);
+  }
+
+  /** Tag an outgoing message with rootId and send it through the shared socket. */
+  sendMessage(rootId: string, message: any): void {
+    message.rootId = rootId;
+    this.sendRaw(message);
+  }
+
+  private sendRaw(message: any): void {
+    if (
+      this.socket.current &&
+      this.socket.current.readyState === WebSocket.OPEN
+    ) {
+      this.socket.current.send(JSON.stringify(message));
+    } else {
+      this.messageQueue.push(message);
+    }
+  }
+
+  /** Route an incoming message to the correct component by rootId. */
+  private handleIncoming(message: any): void {
+    const { rootId } = message;
+    if (rootId && this.components.has(rootId)) {
+      // Strip rootId before forwarding to the component's handler
+      const componentMessage = { ...message };
+      delete componentMessage.rootId;
+      this.components.get(rootId)!.handleIncoming(componentMessage);
+    }
+  }
+
+  /** Close the underlying WebSocket and clear all registered components. */
+  close(): void {
+    if (this.socket.current) {
+      this.socket.current.close();
+      this.socket.current = undefined;
+    }
+    this.components.clear();
+  }
+
+  /**
+   * Dispose this PageClient — closes the WebSocket, clears components, and
+   * removes the entry from the global cache. Call during SPA page transitions
+   * or when the page is fully torn down.
+   */
+  dispose(): void {
+    this.close();
+    disposePageClient(this.key);
+  }
+
+  loadModule(moduleName: string): Promise<ReactPyModule> {
+    return import(`${this.jsModulesPath}${moduleName}`);
+  }
+}

src/js/src/websocket.ts

@@ -9,13 +9,12 @@ export function createReconnectingWebSocket(
   let retries = 0;
   let currentInterval = interval;
   let everConnected = false;
-  const closed = false;
   const socket: { current?: WebSocket } = {};
 
   async function connect() {
-    if (closed) {
-      return;
-    }
+    // Refresh the URL's path/query search params on each connection so that
+    // the server always sees the current page location, even after client-side
+    // navigation in an SPA.
     props.url.searchParams.set("path", window.location.pathname);
     props.url.searchParams.set("qs", window.location.search);
     socket.current = new WebSocket(props.url);