From e99f444be08940ff34f6fe17d0080addb14ea59c Mon Sep 17 00:00:00 2001
From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com>
Date: Wed, 22 Jul 2026 16:28:17 -0700
Subject: [PATCH] fix(relay): detach retired mobile transports

---
 src/main/runtime/relay/relay-control-origin.ts      | 13 ++++++++++---
 src/main/runtime/relay/relay-session-broker.test.ts | 10 ++++++++--
 src/main/runtime/rpc/mobile-socket-wiring.test.ts   |  9 ++++++++-
 src/main/runtime/rpc/mobile-socket-wiring.ts        |  5 ++++-
 4 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/main/runtime/relay/relay-control-origin.ts b/src/main/runtime/relay/relay-control-origin.ts
index ef607c9862b..8035e32c1b6 100644
--- a/src/main/runtime/relay/relay-control-origin.ts
+++ b/src/main/runtime/relay/relay-control-origin.ts
@@ -42,6 +42,7 @@ export class RelayControlOrigin {
   private leaseExpiresAt = 0
   private acceptingConnections = true
   private closed = false
+  private readonly detachMobileSocketTransport: () => void
 
   constructor(options: RelayControlOriginOptions) {
     this.options = options
@@ -53,8 +54,9 @@ export class RelayControlOrigin {
       createSocket: options.createDataSocket,
       onConnectionClosed: (connectionId) => options.onConnectionReleased(connectionId, this)
     })
-    options.mobileSocketWiring.attachTransport(this.transport, (ws) =>
-      this.transport.metadataFor(ws)
+    this.detachMobileSocketTransport = options.mobileSocketWiring.attachTransport(
+      this.transport,
+      (ws) => this.transport.metadataFor(ws)
     )
   }
 
@@ -152,7 +154,12 @@ export class RelayControlOrigin {
     }
     this.controls.clear()
     this.activeControl = null
-    await this.transport.stop()
+    try {
+      await this.transport.stop()
+    } finally {
+      // Why: detaching earlier would skip socket-close cleanup in MobileSocketWiring.
+      this.detachMobileSocketTransport()
+    }
   }
 
   closeNow(): void {
diff --git a/src/main/runtime/relay/relay-session-broker.test.ts b/src/main/runtime/relay/relay-session-broker.test.ts
index d2e0ff0021d..c25092e6784 100644
--- a/src/main/runtime/relay/relay-session-broker.test.ts
+++ b/src/main/runtime/relay/relay-session-broker.test.ts
@@ -115,6 +115,7 @@ describe('RelaySessionBroker lifecycle ownership', () => {
     let current = true
     const statuses: string[] = []
     const keypair = nacl.box.keyPair()
+    const detachTransport = vi.fn()
     const connecting = RelaySessionBroker.connect({
       authConfig: {
         relayTokenEndpoint: 'https://auth.example.test/v1/relay-token',
@@ -127,12 +128,14 @@ describe('RelaySessionBroker lifecycle ownership', () => {
         publicKeyB64: Buffer.from(keypair.publicKey).toString('base64')
       },
       appVersion: '1.0.0',
-      mobileSocketWiring: { attachTransport: vi.fn() } as never,
+      mobileSocketWiring: { attachTransport: vi.fn(() => detachTransport) } as never,
       isCurrent: () => current,
       refreshAccessToken: async () => null,
       onStatus: (status) => statuses.push(status)
     })
     await vi.waitFor(() => expect(fakes.controls).toHaveLength(1))
+    const transportStopped = deferred<void>()
+    fakes.transports[0]!.stop.mockReturnValue(transportStopped.promise)
     current = false
     controlAck.resolve({
       type: 'host-hello-ack',
@@ -147,6 +150,9 @@ describe('RelaySessionBroker lifecycle ownership', () => {
     await expect(connecting).rejects.toBeInstanceOf(StaleRelayBrokerError)
     expect(fakes.controls[0]!.closeNow).toHaveBeenCalledOnce()
     expect(fakes.transports[0]!.stop).toHaveBeenCalledOnce()
+    expect(detachTransport).not.toHaveBeenCalled()
+    transportStopped.resolve(undefined)
+    await vi.waitFor(() => expect(detachTransport).toHaveBeenCalledOnce())
     expect(statuses).toEqual(['connecting'])
   })
 
@@ -316,7 +322,7 @@ function brokerOptions(
       publicKeyB64: Buffer.from(keypair.publicKey).toString('base64')
     },
     appVersion: '1.0.0',
-    mobileSocketWiring: { attachTransport: vi.fn() } as never,
+    mobileSocketWiring: { attachTransport: vi.fn(() => () => {}) } as never,
     isCurrent: () => true,
     refreshAccessToken: async () => null,
     onStatus: vi.fn(),
diff --git a/src/main/runtime/rpc/mobile-socket-wiring.test.ts b/src/main/runtime/rpc/mobile-socket-wiring.test.ts
index b53f20e572e..ae6566cf3d8 100644
--- a/src/main/runtime/rpc/mobile-socket-wiring.test.ts
+++ b/src/main/runtime/rpc/mobile-socket-wiring.test.ts
@@ -84,12 +84,19 @@ describe('MobileSocketWiring', () => {
       onBinary: vi.fn(),
       onClose: vi.fn()
     })
-    wiring.attachTransport(direct)
+    const detachDirect = wiring.attachTransport(direct)
     wiring.attachTransport(relay)
 
     expect(wiring.terminateDeviceConnections('valid-token')).toBe(3)
     expect(direct.terminateClientConnections).toHaveBeenCalledWith('valid-token')
     expect(relay.terminateClientConnections).toHaveBeenCalledWith('valid-token')
+
+    detachDirect()
+    direct.terminateClientConnections.mockClear()
+    relay.terminateClientConnections.mockClear()
+    expect(wiring.terminateDeviceConnections('valid-token')).toBe(2)
+    expect(direct.terminateClientConnections).not.toHaveBeenCalled()
+    expect(relay.terminateClientConnections).toHaveBeenCalledWith('valid-token')
   })
 
   it('preserves the legacy direct handshake, identity, and close cleanup', () => {
diff --git a/src/main/runtime/rpc/mobile-socket-wiring.ts b/src/main/runtime/rpc/mobile-socket-wiring.ts
index 637b5185866..8fc7417b694 100644
--- a/src/main/runtime/rpc/mobile-socket-wiring.ts
+++ b/src/main/runtime/rpc/mobile-socket-wiring.ts
@@ -90,12 +90,15 @@ export class MobileSocketWiring {
     getMetadata: (ws: WebSocket) => MobileSocketTransportMetadata = () => ({
       transport: 'direct'
     })
-  ): void {
+  ): () => void {
     this.transports.add(transport)
     transport.onMessage((message, _reply, ws) => {
       this.handleRawMessage(transport, ws, message, getMetadata(ws))
     })
     transport.onConnectionClose((_clientId, ws) => this.handleClose(ws))
+    return () => {
+      this.transports.delete(transport)
+    }
   }
 
   getConnectionId(ws: WebSocket): string | undefined {