fix(mobile): keep terminal input working when reopening worktrees (#22505)

src/main/runtime/rpc/core.ts

@@ -67,6 +67,8 @@ export type RpcContext = {
   signal?: AbortSignal
   // Why: per-WebSocket key so the server reaps a closing socket's subscriptions without touching sibling sockets sharing the deviceToken.
   connectionId?: string
+  // An unsubscribe cannot retire a registration created after its dispatch began.
+  subscriptionRegistrationVersion?: number
   // Why: shared-control multiplexes many logical streams over one socket; the frame id lets handlers register cleanup per logical stream.
   requestId?: string
   // Why: paired mobile device token; state-owning handlers use it to clean up when that device disconnects.

src/main/runtime/rpc/dispatcher.ts

@@ -96,6 +96,10 @@ export class RpcDispatcher {
           runtime: this.runtime,
           signal: options?.signal,
           connectionId: options?.connectionId,
+          subscriptionRegistrationVersion:
+            request.method === 'terminal.unsubscribe'
+              ? this.runtime.getSubscriptionRegistrationVersion()
+              : undefined,
           requestId: request.id,
           clientId: options?.clientId,
           clientKind: options?.clientKind,

src/main/runtime/rpc/methods/terminal/terminal-viewport-methods.ts

@@ -82,12 +82,12 @@ export const TERMINAL_VIEWPORT_METHODS_AFTER_STREAMS = [
   defineMethod({
     name: 'terminal.unsubscribe',
     params: TerminalUnsubscribe,
-    handler: async (params, { runtime, connectionId }) => {
-      // Why: only the connection that owns the subscription may retire it — a stale
-      // unsubscribe from the pre-reconnect socket names an id the replacement now owns.
+    handler: async (params, { runtime, connectionId, subscriptionRegistrationVersion }) => {
+      // Fence both socket replacement and a newer subscription on the same socket.
       let unsubscribed = runtime.cleanupSubscriptionIfOwnedByConnection(
         params.subscriptionId,
-        connectionId
+        connectionId,
+        subscriptionRegistrationVersion
       )
       // Why: older builds send a bare-handle subscriptionId, so also try the reconstructed `${terminal}:${clientId}` composite key.
       // Why AND over the calls that ran: a clientless stream registers under the bare id
@@ -97,7 +97,8 @@ export const TERMINAL_VIEWPORT_METHODS_AFTER_STREAMS = [
         unsubscribed =
           runtime.cleanupSubscriptionIfOwnedByConnection(
             `${params.subscriptionId}:${params.client.id}`,
-            connectionId
+            connectionId,
+            subscriptionRegistrationVersion
           ) && unsubscribed
       }
       return { unsubscribed }

src/main/runtime/rpc/rpc-streaming-dispatcher.ts

@@ -65,6 +65,11 @@ export class RpcStreamingDispatcher {
 
     if (!isStreamingMethod(method)) {
       try {
+        // Capture before middleware yields to a replacement subscribe on the same connection.
+        const subscriptionRegistrationVersion =
+          request.method === 'terminal.unsubscribe'
+            ? runtime.getSubscriptionRegistrationVersion()
+            : undefined
         const clientHostedBrowser = await routeDispatcherClientHostedBrowserRpc(
           runtime,
           request.method,
@@ -109,6 +114,7 @@ export class RpcStreamingDispatcher {
             signal: options?.signal,
             requestId: request.id,
             connectionId: options?.connectionId,
+            subscriptionRegistrationVersion,
             clientId: options?.clientId,
             pairedDeviceId: options?.pairedDeviceId,
             clientKind: options?.clientKind,

src/main/runtime/rpc/subscription-registry-test-double.ts

@@ -1,6 +1,7 @@
 import type { SubscriptionRegistration } from '../orca-runtime'
 
 type Cleanup = () => void | Promise<void>
+type Entry = { cleanup: Cleanup; version: number }
 
 export type SubscriptionRegistryDouble = {
   registerSubscriptionCleanup: (id: string, cleanup: Cleanup, connectionId?: string) => void
@@ -10,7 +11,12 @@ export type SubscriptionRegistryDouble = {
     connectionId?: string
   ) => SubscriptionRegistration
   cleanupSubscription: (id: string) => void
-  cleanupSubscriptionIfOwnedByConnection: (id: string, connectionId: string | undefined) => boolean
+  cleanupSubscriptionIfOwnedByConnection: (
+    id: string,
+    connectionId: string | undefined,
+    throughVersion?: number
+  ) => boolean
+  getSubscriptionRegistrationVersion: () => number
   cleanupSubscriptionsForConnection: (connectionId: string) => void
   /** Test-only inspection; the runtime deliberately exposes no such accessor. */
   peekCleanup: (id: string) => Cleanup | undefined
@@ -32,8 +38,9 @@ export type SubscriptionRegistryDouble = {
  * subscribe/teardown must use this instead of a bare Map.
  */
 export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
-  const cleanups = new Map<string, Cleanup>()
-  const inFlight = new Map<string, { cleanup: Cleanup; promise: Promise<void> }>()
+  const cleanups = new Map<string, Entry>()
+  const inFlight = new Map<string, { entry: Entry; promise: Promise<void> }>()
+  let registrationVersion = 0
   const byConnection = new Map<string, Set<string>>()
   const connectionByEntry = new Map<string, string>()
 
@@ -54,25 +61,25 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
   }
 
   const cleanupAndWait = (id: string): Promise<void> => {
-    const cleanup = cleanups.get(id)
-    if (!cleanup) {
+    const entry = cleanups.get(id)
+    if (!entry) {
       return Promise.resolve()
     }
     // Mirrors cleanupSubscriptionAndWait: join an in-flight attempt for this exact owner.
     const existing = inFlight.get(id)
-    if (existing?.cleanup === cleanup) {
+    if (existing?.entry === entry) {
       return existing.promise
     }
     let result: void | Promise<void>
     try {
-      result = cleanup()
+      result = entry.cleanup()
     } catch (error) {
       result = Promise.reject(error)
     }
     const promise = Promise.resolve(result)
       .then(() => {
         // Only the generation that registered this callback may retire it.
-        if (cleanups.get(id) !== cleanup) {
+        if (cleanups.get(id) !== entry) {
           return
         }
         cleanups.delete(id)
@@ -83,7 +90,7 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
           inFlight.delete(id)
         }
       })
-    inFlight.set(id, { cleanup, promise })
+    inFlight.set(id, { entry, promise })
     return promise
   }
 
@@ -92,7 +99,7 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
     void cleanupAndWait(id).catch(() => undefined)
   }
 
-  const cleanupOwned = (id: string, expected: Cleanup): void => {
+  const cleanupOwned = (id: string, expected: Entry): void => {
     if (cleanups.get(id) !== expected) {
       return
     }
@@ -103,15 +110,16 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
     id: string,
     cleanup: Cleanup,
     connectionId?: string
-  ): void => {
+  ): Entry => {
     const existing = cleanups.get(id)
     if (existing) {
       removeIndex(id)
       cleanupOwned(id, existing)
     }
-    cleanups.set(id, cleanup)
+    const entry = { cleanup, version: ++registrationVersion }
+    cleanups.set(id, entry)
     if (!connectionId) {
-      return
+      return entry
     }
     let set = byConnection.get(connectionId)
     if (!set) {
@@ -120,27 +128,29 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
     }
     set.add(id)
     connectionByEntry.set(id, connectionId)
+    return entry
   }
 
   return {
     registerSubscriptionCleanup,
     registerOwnedSubscriptionCleanup: (id, cleanup, connectionId) => {
-      registerSubscriptionCleanup(id, cleanup, connectionId)
+      const entry = registerSubscriptionCleanup(id, cleanup, connectionId)
       return {
-        releaseIfCurrent: () => cleanupOwned(id, cleanup)
+        releaseIfCurrent: () => cleanupOwned(id, entry)
       }
     },
     cleanupSubscription,
-    cleanupSubscriptionIfOwnedByConnection: (id, connectionId) => {
-      if (!connectionId) {
-        cleanupSubscription(id)
-        return true
-      }
+    getSubscriptionRegistrationVersion: () => registrationVersion,
+    cleanupSubscriptionIfOwnedByConnection: (id, connectionId, throughVersion) => {
+      const entry = cleanups.get(id)
       // Mirrors the production early-out: an unregistered id is already gone, not refused.
-      if (!cleanups.has(id)) {
+      if (!entry) {
         return true
       }
-      if (connectionByEntry.get(id) !== connectionId) {
+      if (throughVersion !== undefined && entry.version > throughVersion) {
+        return false
+      }
+      if (connectionId && connectionByEntry.get(id) !== connectionId) {
         return false
       }
       cleanupSubscription(id)
@@ -163,6 +173,6 @@ export function createSubscriptionRegistryDouble(): SubscriptionRegistryDouble {
         byConnection.delete(connectionId)
       }
     },
-    peekCleanup: (id) => cleanups.get(id)
+    peekCleanup: (id) => cleanups.get(id)?.cleanup
   }
 }

src/main/runtime/rpc/terminal-subscribe-ownership.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it, vi } from 'vitest'
 import type { OrcaRuntimeService } from '../orca-runtime'
+import { RuntimeSubscriptionRegistry } from '../runtime-subscription-registry'
 import type { RpcRequest } from './core'
 import { RpcDispatcher } from './dispatcher'
 import { TERMINAL_METHODS } from './methods/terminal'
@@ -14,6 +15,7 @@ function stubRuntime(
   waiters: Waiter[],
   overrides: Record<string, unknown> = {}
 ): OrcaRuntimeService {
+  // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: This partial runtime supplies the terminal RPC methods these tests invoke.
   return {
     getRuntimeId: () => 'test-runtime',
     registerRemoteTerminalViewSubscriber: () => () => {},
@@ -36,6 +38,7 @@ function stubRuntime(
     registerOwnedSubscriptionCleanup: vi.fn(registry.registerOwnedSubscriptionCleanup),
     cleanupSubscription: vi.fn(registry.cleanupSubscription),
     cleanupSubscriptionIfOwnedByConnection: vi.fn(registry.cleanupSubscriptionIfOwnedByConnection),
+    getSubscriptionRegistrationVersion: registry.getSubscriptionRegistrationVersion,
     subscribeToPtyExit: vi.fn((_ptyId: string, listener: () => void) => {
       waiters.push({ resolve: listener })
       return vi.fn()
@@ -291,6 +294,87 @@ describe('terminal.unsubscribe connection ownership', () => {
     expect(JSON.parse(replies[0]!).result).toEqual({ unsubscribed: false })
   })
 
+  it.each(
+    [binaryParams, leaseOnlyParams].flatMap((params) =>
+      ['streaming', 'unary'].flatMap((transport) =>
+        [SUBSCRIPTION_ID, 'terminal-1'].map((subscriptionId) => ({
+          params,
+          transport,
+          subscriptionId
+        }))
+      )
+    )
+  )(
+    'keeps the replacement when $transport unsubscribe of $subscriptionId yields',
+    async ({ params, transport, subscriptionId }) => {
+      const registry = new RuntimeSubscriptionRegistry()
+      const runtime = stubRuntime(createSubscriptionRegistryDouble(), [], {
+        registerOwnedSubscriptionCleanup: registry.registerOwned.bind(registry),
+        cleanupSubscriptionIfOwnedByConnection: registry.cleanupIfOwnedByConnection.bind(registry),
+        getSubscriptionRegistrationVersion: registry.getRegistrationVersion.bind(registry)
+      })
+      const dispatcher = new RpcDispatcher({ runtime, methods: TERMINAL_METHODS })
+      const connection = streamOptions('conn-a')
+      const original = dispatcher.dispatchStreaming(makeRequest(params), vi.fn(), connection)
+      await vi.waitFor(() => expect(runtime.handleMobileSubscribe).toHaveBeenCalledTimes(1))
+
+      const replies: string[] = []
+      const retiring =
+        transport === 'streaming'
+          ? dispatcher.dispatchStreaming(
+              unsubscribeRequest(subscriptionId),
+              (reply) => replies.push(reply),
+              { connectionId: 'conn-a' }
+            )
+          : dispatcher
+              .dispatch(unsubscribeRequest(subscriptionId), { connectionId: 'conn-a' })
+              .then((reply) => replies.push(JSON.stringify(reply)))
+      const replacementMessages: string[] = []
+      const replacement = dispatcher.dispatchStreaming(
+        { ...makeRequest(params), id: 'req-replacement' },
+        (reply) => replacementMessages.push(reply),
+        connection
+      )
+      await retiring
+
+      try {
+        expect(runtime.handleMobileSubscribe).toHaveBeenCalledTimes(2)
+        expect(runtime.handleMobileUnsubscribe).toHaveBeenCalledTimes(1)
+        expect(replacementMessages.some((reply) => JSON.parse(reply).result?.type === 'end')).toBe(
+          false
+        )
+        expect(JSON.parse(replies[0]!).result).toEqual({ unsubscribed: false })
+      } finally {
+        registry.cleanupForConnection('conn-a')
+        await Promise.all([original, replacement])
+      }
+    }
+  )
+
+  it('rejects malformed unsubscribe params before capturing registration state', async () => {
+    const registry = createSubscriptionRegistryDouble()
+    const runtime = stubRuntime(registry, [])
+    const captureVersion = vi.spyOn(runtime, 'getSubscriptionRegistrationVersion')
+    const cleanup = vi.fn()
+    registry.registerSubscriptionCleanup(SUBSCRIPTION_ID, cleanup, 'conn-a')
+    const dispatcher = new RpcDispatcher({ runtime, methods: TERMINAL_METHODS })
+    const replies: string[] = []
+
+    await dispatcher.dispatchStreaming(
+      { ...unsubscribeRequest(SUBSCRIPTION_ID), params: { subscriptionId: 42 } },
+      (reply) => replies.push(reply),
+      { connectionId: 'conn-a' }
+    )
+
+    expect(JSON.parse(replies[0]!)).toMatchObject({
+      ok: false,
+      error: { code: 'invalid_argument' }
+    })
+    expect(captureVersion).not.toHaveBeenCalled()
+    expect(cleanup).not.toHaveBeenCalled()
+    registry.cleanupSubscriptionsForConnection('conn-a')
+  })
+
   // T5 same-connection: preservation — the owning connection may still unsubscribe.
   it('honors an unsubscribe from the connection that owns the subscription', async () => {
     const registry = createSubscriptionRegistryDouble()

src/main/runtime/runtime-service-command-surface.ts

@@ -25,6 +25,7 @@ export type RuntimeServiceCommandSurface = {
   cleanupSubscriptionsByPrefix: RuntimeSubscriptionRegistry['cleanupByPrefix']
   cleanupSubscriptionsForConnection: RuntimeSubscriptionRegistry['cleanupForConnection']
   cleanupSubscriptionIfOwnedByConnection: RuntimeSubscriptionRegistry['cleanupIfOwnedByConnection']
+  getSubscriptionRegistrationVersion: RuntimeSubscriptionRegistry['getRegistrationVersion']
   onNotificationDispatched: RuntimeMobileNotificationController['onDispatched']
   getMobileNotificationListenerCount: RuntimeMobileNotificationController['getListenerCount']
   dispatchMobileNotification: RuntimeMobileNotificationController['dispatch']
@@ -114,6 +115,7 @@ export function installRuntimeServiceCommandSurface(
     cleanupSubscriptionsForConnection: subscriptions.cleanupForConnection.bind(subscriptions),
     cleanupSubscriptionIfOwnedByConnection:
       subscriptions.cleanupIfOwnedByConnection.bind(subscriptions),
+    getSubscriptionRegistrationVersion: subscriptions.getRegistrationVersion.bind(subscriptions),
     onNotificationDispatched: notifications.onDispatched.bind(notifications),
     getMobileNotificationListenerCount: notifications.getListenerCount.bind(notifications),
     dispatchMobileNotification: notifications.dispatch.bind(notifications),

src/main/runtime/runtime-subscription-registry.test.ts

@@ -0,0 +1,125 @@
+import { describe, expect, it, vi } from 'vitest'
+import { RuntimeSubscriptionRegistry } from './runtime-subscription-registry'
+
+describe('subscription registration versions', () => {
+  it.each(['conn-a', undefined])(
+    'fences a delayed unsubscribe across replacement (%s)',
+    async (connectionId) => {
+      const registry = new RuntimeSubscriptionRegistry()
+      const originalCleanup = vi.fn()
+      const replacementCleanup = vi.fn()
+      registry.register('terminal:generation', originalCleanup, 'conn-a')
+      const admittedVersion = registry.getRegistrationVersion()
+      registry.register('terminal:generation', replacementCleanup, 'conn-a')
+
+      expect(
+        registry.cleanupIfOwnedByConnection('terminal:generation', connectionId, admittedVersion)
+      ).toBe(false)
+      await Promise.resolve()
+      expect(originalCleanup).toHaveBeenCalledTimes(1)
+      expect(replacementCleanup).not.toHaveBeenCalled()
+
+      expect(
+        registry.cleanupIfOwnedByConnection(
+          'terminal:generation',
+          connectionId,
+          registry.getRegistrationVersion()
+        )
+      ).toBe(true)
+      await registry.cleanupAndWait('terminal:generation')
+      expect(replacementCleanup).toHaveBeenCalledTimes(1)
+    }
+  )
+
+  it('does not let a missing target at admission cancel a later registration', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    const admittedVersion = registry.getRegistrationVersion()
+    const cleanup = vi.fn()
+    registry.register('terminal:later', cleanup, 'conn-a')
+    expect(registry.cleanupIfOwnedByConnection('terminal:later', 'conn-a', admittedVersion)).toBe(
+      false
+    )
+    expect(cleanup).not.toHaveBeenCalled()
+    await registry.cleanupAndWait('terminal:later')
+  })
+
+  it('allows admitted cleanup after unrelated subscriptions register', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    const cleanup = vi.fn()
+    registry.register('terminal:first', cleanup, 'conn-a')
+    const admittedVersion = registry.getRegistrationVersion()
+    registry.register('terminal:other', vi.fn(), 'conn-a')
+    expect(registry.cleanupIfOwnedByConnection('terminal:first', 'conn-a', admittedVersion)).toBe(
+      true
+    )
+    await registry.cleanupAndWait('terminal:first')
+    expect(cleanup).toHaveBeenCalledTimes(1)
+    await registry.cleanupAndWait('terminal:other')
+  })
+
+  it('still refuses a different connection with a current registration version', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    const cleanup = vi.fn()
+    registry.register('terminal:owned', cleanup, 'conn-a')
+    expect(
+      registry.cleanupIfOwnedByConnection(
+        'terminal:owned',
+        'conn-b',
+        registry.getRegistrationVersion()
+      )
+    ).toBe(false)
+    expect(cleanup).not.toHaveBeenCalled()
+    await registry.cleanupAndWait('terminal:owned')
+  })
+
+  it('reports an already-retired subscription as gone', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    registry.register('terminal:retired', vi.fn(), 'conn-a')
+    const admittedVersion = registry.getRegistrationVersion()
+    await registry.cleanupAndWait('terminal:retired')
+    expect(registry.cleanupIfOwnedByConnection('terminal:retired', 'conn-a', admittedVersion)).toBe(
+      true
+    )
+  })
+
+  it('keeps registration ownership even when a cleanup callback is reused', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    const cleanup = vi.fn()
+    const original = registry.registerOwned('terminal:reused', cleanup, 'conn')
+    const replacement = registry.registerOwned('terminal:reused', cleanup, 'conn')
+    original.releaseIfCurrent()
+    await Promise.resolve()
+    expect(cleanup).toHaveBeenCalledTimes(1)
+    replacement.releaseIfCurrent()
+    await registry.cleanupAndWait('terminal:reused')
+    expect(cleanup).toHaveBeenCalledTimes(2)
+  })
+
+  it('does not merge different registrations of the same async cleanup callback', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    const gate = Promise.withResolvers<void>()
+    const cleanup = vi.fn(() => gate.promise)
+    registry.register('terminal:async', cleanup, 'conn')
+    const originalCleanup = registry.cleanupAndWait('terminal:async')
+    registry.register('terminal:async', cleanup, 'conn')
+    const replacementCleanup = registry.cleanupAndWait('terminal:async')
+    expect(cleanup).toHaveBeenCalledTimes(2)
+    gate.resolve()
+    await Promise.all([originalCleanup, replacementCleanup])
+  })
+
+  it('does not retry a replaced registration that reused the cleanup callback', async () => {
+    const registry = new RuntimeSubscriptionRegistry()
+    const gate = Promise.withResolvers<void>()
+    const cleanup = vi.fn()
+    registry.register('terminal:retry', cleanup, 'conn')
+    registry.retryAfter('terminal:retry', cleanup, gate.promise)
+    registry.register('terminal:retry', cleanup, 'conn')
+    gate.resolve()
+    await Promise.resolve()
+    await Promise.resolve()
+    expect(cleanup).toHaveBeenCalledTimes(1)
+    await registry.cleanupAndWait('terminal:retry')
+    expect(cleanup).toHaveBeenCalledTimes(2)
+  })
+})

src/main/runtime/runtime-subscription-registry.ts

@@ -1,17 +1,24 @@
 type SubscriptionCleanup = () => void | Promise<void>
 
+type SubscriptionEntry = { cleanup: SubscriptionCleanup; version: number }
+
 export type SubscriptionRegistration = {
   releaseIfCurrent(): void
 }
 
 export class RuntimeSubscriptionRegistry {
-  private readonly cleanups = new Map<string, SubscriptionCleanup>()
+  private readonly cleanups = new Map<string, SubscriptionEntry>()
   private readonly cleanupPromises = new Map<
     string,
-    { cleanup: SubscriptionCleanup; promise: Promise<void> }
+    { entry: SubscriptionEntry; promise: Promise<void> }
   >()
   private readonly subscriptionsByConnection = new Map<string, Set<string>>()
   private readonly connectionBySubscription = new Map<string, string>()
+  private registrationVersion = 0
+
+  getRegistrationVersion(): number {
+    return this.registrationVersion
+  }
 
   register(subscriptionId: string, cleanup: SubscriptionCleanup, connectionId?: string): void {
     const existing = this.cleanups.get(subscriptionId)
@@ -19,7 +26,7 @@ export class RuntimeSubscriptionRegistry {
       this.removeConnectionIndex(subscriptionId)
       this.cleanup(subscriptionId)
     }
-    this.cleanups.set(subscriptionId, cleanup)
+    this.cleanups.set(subscriptionId, { cleanup, version: ++this.registrationVersion })
     if (!connectionId) {
       return
     }
@@ -38,18 +45,23 @@ export class RuntimeSubscriptionRegistry {
     connectionId?: string
   ): SubscriptionRegistration {
     this.register(subscriptionId, cleanup, connectionId)
-    return { releaseIfCurrent: () => this.cleanupOwned(subscriptionId, cleanup) }
+    const version = this.registrationVersion
+    return { releaseIfCurrent: () => this.cleanupOwned(subscriptionId, version) }
   }
 
-  cleanupIfOwnedByConnection(subscriptionId: string, connectionId?: string): boolean {
-    if (!connectionId) {
-      this.cleanup(subscriptionId)
+  cleanupIfOwnedByConnection(
+    subscriptionId: string,
+    connectionId?: string,
+    throughVersion?: number
+  ): boolean {
+    const entry = this.cleanups.get(subscriptionId)
+    if (!entry) {
       return true
     }
-    if (!this.cleanups.has(subscriptionId)) {
-      return true
+    if (throughVersion !== undefined && entry.version > throughVersion) {
+      return false
     }
-    if (this.connectionBySubscription.get(subscriptionId) !== connectionId) {
+    if (connectionId && this.connectionBySubscription.get(subscriptionId) !== connectionId) {
       return false
     }
     this.cleanup(subscriptionId)
@@ -63,15 +75,19 @@ export class RuntimeSubscriptionRegistry {
   }
 
   retryAfter(subscriptionId: string, cleanupOwner: SubscriptionCleanup, gate: Promise<void>): void {
+    const entry = this.cleanups.get(subscriptionId)
     const failedGeneration = this.cleanupPromises.get(subscriptionId)
     void gate.then(
       async () => {
-        await (failedGeneration?.cleanup === cleanupOwner
+        if (entry?.cleanup !== cleanupOwner) {
+          return
+        }
+        await (failedGeneration?.entry === entry
           ? failedGeneration.promise.catch(() => undefined)
           : undefined)
-        while (this.cleanups.get(subscriptionId) === cleanupOwner) {
+        while (this.cleanups.get(subscriptionId) === entry) {
           const newerGeneration = this.cleanupPromises.get(subscriptionId)
-          if (newerGeneration?.cleanup === cleanupOwner) {
+          if (newerGeneration?.entry === entry) {
             await newerGeneration.promise.catch(() => undefined)
             continue
           }
@@ -84,23 +100,23 @@ export class RuntimeSubscriptionRegistry {
   }
 
   async cleanupAndWait(subscriptionId: string): Promise<void> {
-    const cleanup = this.cleanups.get(subscriptionId)
-    if (!cleanup) {
+    const entry = this.cleanups.get(subscriptionId)
+    if (!entry) {
       return
     }
     const inFlight = this.cleanupPromises.get(subscriptionId)
-    if (inFlight?.cleanup === cleanup) {
+    if (inFlight?.entry === entry) {
       return inFlight.promise
     }
     let cleanupResult: void | Promise<void>
     try {
-      cleanupResult = cleanup()
+      cleanupResult = entry.cleanup()
     } catch (error) {
       cleanupResult = Promise.reject(error)
     }
     const promise = Promise.resolve(cleanupResult)
       .then(() => {
-        if (this.cleanups.get(subscriptionId) !== cleanup) {
+        if (this.cleanups.get(subscriptionId) !== entry) {
           return
         }
         this.cleanups.delete(subscriptionId)
@@ -111,7 +127,7 @@ export class RuntimeSubscriptionRegistry {
           this.cleanupPromises.delete(subscriptionId)
         }
       })
-    this.cleanupPromises.set(subscriptionId, { cleanup, promise })
+    this.cleanupPromises.set(subscriptionId, { entry, promise })
     return promise
   }
 
@@ -139,8 +155,8 @@ export class RuntimeSubscriptionRegistry {
     }
   }
 
-  private cleanupOwned(subscriptionId: string, expectedCleanup: SubscriptionCleanup): void {
-    if (this.cleanups.get(subscriptionId) !== expectedCleanup) {
+  private cleanupOwned(subscriptionId: string, expectedVersion: number): void {
+    if (this.cleanups.get(subscriptionId)?.version !== expectedVersion) {
       return
     }
     this.cleanup(subscriptionId)