From 1563d6ffeef51f1eba4e4d6d5f12ff00444c9158 Mon Sep 17 00:00:00 2001
From: Jason Nall <jsonnull@gmail.com>
Date: Wed, 5 Jul 2023 23:26:03 -0400
Subject: [PATCH] Clean up subscriptions when frames perform navigation.

---
 .changeset/stale-dogs-collect.md              |  5 +++
 .../src/main/createIPCHandler.ts              | 35 +++++++++++++++++--
 .../src/main/handleIPCMessage.ts              |  4 +++
 .../electron-trpc/src/renderer/ipcLink.ts     |  6 ++--
 4 files changed, 44 insertions(+), 6 deletions(-)
 create mode 100644 .changeset/stale-dogs-collect.md

diff --git a/.changeset/stale-dogs-collect.md b/.changeset/stale-dogs-collect.md
new file mode 100644
index 00000000..8f789f65
--- /dev/null
+++ b/.changeset/stale-dogs-collect.md
@@ -0,0 +1,5 @@
+---
+'electron-trpc': minor
+---
+
+Subscriptions will now be cleaned up when the frame owning that subscription performs a navigation.
diff --git a/packages/electron-trpc/src/main/createIPCHandler.ts b/packages/electron-trpc/src/main/createIPCHandler.ts
index 658b52e7..0ee086db 100644
--- a/packages/electron-trpc/src/main/createIPCHandler.ts
+++ b/packages/electron-trpc/src/main/createIPCHandler.ts
@@ -6,6 +6,9 @@ import { CreateContextOptions } from './types';
 import { ELECTRON_TRPC_CHANNEL } from '../constants';
 import { ETRPCRequest } from '../types';
 import { Unsubscribable } from '@trpc/server/observable';
+import debugFactory from 'debug';
+
+const debug = debugFactory('electron-trpc:main:IPCHandler');
 
 type Awaitable<T> = T | Promise<T>;
 
@@ -46,23 +49,49 @@ class IPCHandler<TRouter extends AnyRouter> {
       return;
     }
 
+    debug('Attaching window', win.id);
+
     this.#windows.push(win);
-    this.#attachSubscriptionCleanupHandler(win);
+    this.#attachSubscriptionCleanupHandlers(win);
   }
 
   detachWindow(win: BrowserWindow) {
+    debug('Detaching window', win.id);
+
     this.#windows = this.#windows.filter((w) => w !== win);
+    this.#cleanUpSubscriptions({ webContentsId: win.webContents.id });
+  }
 
+  #cleanUpSubscriptions({
+    webContentsId,
+    frameRoutingId,
+  }: {
+    webContentsId: number;
+    frameRoutingId?: number;
+  }) {
     for (const [key, sub] of this.#subscriptions.entries()) {
-      if (key.startsWith(`${win.webContents.id}-`)) {
+      if (key.startsWith(`${webContentsId}-${frameRoutingId ?? ''}`)) {
+        debug('Closing subscription', key);
         sub.unsubscribe();
         this.#subscriptions.delete(key);
       }
     }
   }
 
-  #attachSubscriptionCleanupHandler(win: BrowserWindow) {
+  #attachSubscriptionCleanupHandlers(win: BrowserWindow) {
+    win.webContents.on('did-start-navigation', ({ frame }) => {
+      debug(
+        'Handling webContents `did-start-navigation` event',
+        `webContentsId: ${win.webContents.id}`,
+        `frameRoutingId: ${frame.routingId}`
+      );
+      this.#cleanUpSubscriptions({
+        webContentsId: win.webContents.id,
+        frameRoutingId: frame.routingId,
+      });
+    });
     win.webContents.on('destroyed', () => {
+      debug('Handling webContents `destroyed` event');
       this.detachWindow(win);
     });
   }
diff --git a/packages/electron-trpc/src/main/handleIPCMessage.ts b/packages/electron-trpc/src/main/handleIPCMessage.ts
index 325ee9cc..bce69a57 100644
--- a/packages/electron-trpc/src/main/handleIPCMessage.ts
+++ b/packages/electron-trpc/src/main/handleIPCMessage.ts
@@ -8,6 +8,9 @@ import { getTRPCErrorFromUnknown } from './utils';
 import { CreateContextOptions } from './types';
 import { ELECTRON_TRPC_CHANNEL } from '../constants';
 import { ETRPCRequest } from '../types';
+import debugFactory from 'debug';
+
+const debug = debugFactory('electron-trpc:main:handleIPCMessage');
 
 export async function handleIPCMessage<TRouter extends AnyRouter>({
   router,
@@ -107,6 +110,7 @@ export async function handleIPCMessage<TRouter extends AnyRouter>({
       },
     });
 
+    debug('Creating subscription', internalId);
     subscriptions.set(internalId, subscription);
   } catch (cause) {
     const error: TRPCError = getTRPCErrorFromUnknown(cause);
diff --git a/packages/electron-trpc/src/renderer/ipcLink.ts b/packages/electron-trpc/src/renderer/ipcLink.ts
index e416ab70..ab03aa78 100644
--- a/packages/electron-trpc/src/renderer/ipcLink.ts
+++ b/packages/electron-trpc/src/renderer/ipcLink.ts
@@ -4,9 +4,9 @@ import type { TRPCResponseMessage } from '@trpc/server/rpc';
 import type { RendererGlobalElectronTRPC } from '../types';
 import { observable, Observer } from '@trpc/server/observable';
 import { transformResult } from './utils';
-import debug from 'debug';
+import debugFactory from 'debug';
 
-const log = debug('electron-trpc:renderer:ipcLink');
+const debug = debugFactory('electron-trpc:renderer:ipcLink');
 
 type IPCCallbackResult<TRouter extends AnyRouter = AnyRouter> = TRPCResponseMessage<
   unknown,
@@ -47,7 +47,7 @@ class IPCClient {
   }
 
   #handleResponse(response: TRPCResponseMessage) {
-    log('handling response', response);
+    debug('handling response', response);
     const request = response.id && this.#pendingRequests.get(response.id);
     if (!request) {
       return;