From dcb457f7b7780f5057e596f80323c49536022dc3 Mon Sep 17 00:00:00 2001
From: Neil <4138956+nwparker@users.noreply.github.com>
Date: Sun, 31 May 2026 05:15:58 -0700
Subject: [PATCH] fix: close removed runtime subscriptions (#4126)

---
 src/main/ipc/runtime-environments.test.ts | 64 +++++++++++++++++++++++
 src/main/ipc/runtime-environments.ts      | 18 ++++++-
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/src/main/ipc/runtime-environments.test.ts b/src/main/ipc/runtime-environments.test.ts
index 350099e37..993ed146d 100644
--- a/src/main/ipc/runtime-environments.test.ts
+++ b/src/main/ipc/runtime-environments.test.ts
@@ -448,6 +448,70 @@ describe('registerRuntimeEnvironmentHandlers', () => {
     markUsedSpy.mockRestore()
   })
 
+  it('closes streaming subscriptions when their saved runtime is removed', async () => {
+    registerRuntimeEnvironmentHandlers()
+    const close = vi.fn()
+    const sendBinary = vi.fn()
+    subscribeRemoteRuntimeRequestMock.mockResolvedValue({
+      requestId: 'stream-remove',
+      close,
+      sendBinary
+    })
+
+    const add = handler<
+      { name: string; pairingCode: string },
+      { environment: { id: string; name: string } }
+    >('runtimeEnvironments:addFromPairingCode')
+    const added = await add(null, { name: 'desk', pairingCode: pairingCode() })
+
+    const destroyedListenerRemoved = vi.fn()
+    const subscribe = handler<
+      {
+        selector: string
+        method: string
+        params?: unknown
+        subscriptionId?: string
+      },
+      { subscriptionId: string; requestId: string }
+    >('runtimeEnvironments:subscribe')
+    const result = await subscribe(
+      {
+        sender: {
+          id: 1,
+          isDestroyed: () => false,
+          send: vi.fn(),
+          once: vi.fn(),
+          removeListener: destroyedListenerRemoved
+        }
+      },
+      {
+        selector: 'desk',
+        method: 'terminal.subscribe',
+        params: { terminal: 't1' },
+        subscriptionId: 'removed-env-sub'
+      }
+    )
+
+    const remove = handler<{ selector: string }, { removed: { id: string; name: string } }>(
+      'runtimeEnvironments:remove'
+    )
+    expect(remove(null, { selector: added.environment.id })).toMatchObject({
+      removed: { id: added.environment.id, name: 'desk' }
+    })
+
+    expect(close).toHaveBeenCalledTimes(1)
+    expect(destroyedListenerRemoved).toHaveBeenCalledWith('destroyed', expect.any(Function))
+
+    const unsubscribe = handler<{ subscriptionId: string }, { unsubscribed: boolean }>(
+      'runtimeEnvironments:unsubscribe'
+    )
+    expect(
+      await unsubscribe({ sender: { id: 1 } }, { subscriptionId: result.subscriptionId })
+    ).toEqual({
+      unsubscribed: false
+    })
+  })
+
   it('rejects cross-window streaming subscription control', async () => {
     registerRuntimeEnvironmentHandlers()
     const close = vi.fn()
diff --git a/src/main/ipc/runtime-environments.ts b/src/main/ipc/runtime-environments.ts
index ca7e54ea7..3d206a149 100644
--- a/src/main/ipc/runtime-environments.ts
+++ b/src/main/ipc/runtime-environments.ts
@@ -40,6 +40,7 @@ const RUNTIME_ENVIRONMENT_HANDLER_CHANNELS = [
 ] as const
 
 type RetainedRemoteRuntimeSubscription = RemoteRuntimeSubscription & {
+  environmentId: string
   ownerWebContentsId: number
   removeDestroyedListener: () => void
 }
@@ -53,6 +54,18 @@ function shouldUseCachedRequestConnection(method: string): boolean {
   return method === 'terminal.send' || method === 'terminal.updateViewport'
 }
 
+function closeSubscriptionsForEnvironment(environmentId: string): void {
+  // Why: removing a saved runtime invalidates its streaming WebSockets too;
+  // otherwise terminal/browser subscriptions stay alive until renderer teardown.
+  for (const [subscriptionId, subscription] of remoteRuntimeSubscriptions) {
+    if (subscription.environmentId !== environmentId) {
+      continue
+    }
+    remoteRuntimeSubscriptions.delete(subscriptionId)
+    subscription.close()
+  }
+}
+
 export function registerRuntimeEnvironmentHandlers(): void {
   // Why: keep direct re-registration safe even though register-core-handlers
   // normally guards this path; otherwise the binary send listener can stack.
@@ -86,6 +99,7 @@ export function registerRuntimeEnvironmentHandlers(): void {
       if (args.selector !== removed.id) {
         closeRemoteRuntimeRequestConnection(args.selector)
       }
+      closeSubscriptionsForEnvironment(removed.id)
       return { removed: redactRuntimeEnvironment(removed) }
     }
   )
@@ -136,6 +150,7 @@ export function registerRuntimeEnvironmentHandlers(): void {
       if (remoteRuntimeSubscriptions.has(subscriptionId)) {
         throw new Error('Runtime environment subscription id already exists')
       }
+      const environment = resolveEnvironment(getUserDataPath(), args.selector)
       const sender = event.sender
       const ownerWebContentsId = sender.id
       let senderDestroyed = sender.isDestroyed()
@@ -163,7 +178,7 @@ export function registerRuntimeEnvironmentHandlers(): void {
       destroyedListenerAttached = true
       try {
         subscription = await subscribeRuntimeEnvironment(
-          args.selector,
+          environment.id,
           args.method,
           args.params,
           args.timeoutMs,
@@ -194,6 +209,7 @@ export function registerRuntimeEnvironmentHandlers(): void {
       }
       remoteRuntimeSubscriptions.set(subscriptionId, {
         requestId: subscription.requestId,
+        environmentId: environment.id,
         ownerWebContentsId,
         removeDestroyedListener,
         sendBinary: (bytes) => subscription?.sendBinary(bytes) ?? false,