fix(relay): preserve origin controls during drain

Co-authored-by: Orca <help@stably.ai>

src/main/runtime/relay/desktop-relay-service.ts

@@ -65,8 +65,7 @@ export class DesktopRelayService {
           mobileSocketWiring,
           isCurrent,
           refreshAccessToken,
-          onStatus: options.onStatus,
-          onResolveDirector: () => this.coordinator.restart()
+          onStatus: options.onStatus
         })
         void this.flushRevokeOutbox(broker)
         return broker

src/main/runtime/relay/relay-auth-coordinator.ts

@@ -63,11 +63,6 @@ export class RelayAuthCoordinator {
     this.options.onStatus('offline')
   }
 
-  restart(): void {
-    this.fenceAndCloseNow()
-    this.reconcile()
-  }
-
   getActiveBroker(): CoordinatedRelayBroker | null {
     return this.ownership?.valid ? this.ownership.broker : null
   }

src/main/runtime/relay/relay-control-client.ts

@@ -106,6 +106,10 @@ export class RelayControlClient {
     })
   }
 
+  get pendingRequestCount(): number {
+    return this.requests.size
+  }
+
   refreshAuthorization(relayJwt: string): void {
     this.sendActive({ type: 'auth-refresh', relayJwt })
   }

src/main/runtime/relay/relay-control-origin.ts

@@ -0,0 +1,243 @@
+import type WebSocket from 'ws'
+import type { E2EEKeypair } from '../e2ee-keypair'
+import { CloudRelayTransport } from '../rpc/relay-transport'
+import type { MobileSocketWiring } from '../rpc/mobile-socket-wiring'
+import { RelayControlClient } from './relay-control-client'
+import type {
+  RelayConnectionOpenMessage,
+  RelayDrainMessage,
+  RelayHostHelloAckMessage
+} from './relay-control-protocol'
+import type { RelayIdentity } from './relay-session-broker-contract'
+import type { RelayAssignment } from './relay-http-client'
+
+type RelayControlOriginOptions = {
+  assignment: RelayAssignment
+  relayJwt: string
+  relayHostId: string
+  identity: RelayIdentity
+  keypair: E2EEKeypair
+  appVersion: string
+  mobileSocketWiring: MobileSocketWiring
+  createControlSocket?: (url: string, relayJwt: string) => WebSocket
+  createDataSocket?: (url: string) => WebSocket
+  onConnectionOwned: (connectionId: string, origin: RelayControlOrigin) => void
+  onConnectionReleased: (connectionId: string, origin: RelayControlOrigin) => void
+  onDrain: (origin: RelayControlOrigin, message: RelayDrainMessage) => void
+  onClose: (origin: RelayControlOrigin, code: number) => void
+}
+
+export class RelayControlOrigin {
+  readonly assignment: RelayAssignment
+  readonly transport: CloudRelayTransport
+  private readonly options: RelayControlOriginOptions
+  private readonly controls = new Set<RelayControlClient>()
+  private readonly retiredControlTimers = new Map<
+    RelayControlClient,
+    ReturnType<typeof setTimeout>
+  >()
+  private activeControl: RelayControlClient | null = null
+  private generation = 0
+  private controlResumeSecret: string | null = null
+  private leaseExpiresAt = 0
+  private acceptingConnections = true
+  private closed = false
+
+  constructor(options: RelayControlOriginOptions) {
+    this.options = options
+    this.assignment = options.assignment
+    this.transport = new CloudRelayTransport({
+      cellUrl: options.assignment.cellUrl,
+      relayHostId: options.relayHostId,
+      generation: 0,
+      createSocket: options.createDataSocket,
+      onConnectionClosed: (connectionId) => options.onConnectionReleased(connectionId, this)
+    })
+    options.mobileSocketWiring.attachTransport(this.transport, (ws) =>
+      this.transport.metadataFor(ws)
+    )
+  }
+
+  get control(): RelayControlClient {
+    if (!this.activeControl) {
+      throw new Error('relay_control_not_active')
+    }
+    return this.activeControl
+  }
+
+  get availableControl(): RelayControlClient | null {
+    return this.activeControl
+  }
+
+  get cellUrl(): string {
+    return this.assignment.cellUrl
+  }
+
+  get assignmentEpoch(): number {
+    return this.assignment.assignmentEpoch
+  }
+
+  get controlLeaseExpiresAt(): number {
+    return this.leaseExpiresAt
+  }
+
+  get pendingRequestCount(): number {
+    let count = 0
+    for (const control of this.controls) {
+      count += control.pendingRequestCount
+    }
+    return count
+  }
+
+  async open(): Promise<void> {
+    await this.transport.start()
+    const { control, ack } = await this.openControl()
+    this.activate(control, ack)
+  }
+
+  async rebind(relayJwt: string, assignment: RelayAssignment): Promise<void> {
+    if (assignment.cellUrl !== this.cellUrl || !this.controlResumeSecret || this.generation <= 0) {
+      throw new Error('relay_control_rebind_origin_mismatch')
+    }
+    const previous = this.activeControl
+    const { control, ack } = await this.openControl({
+      relayJwt,
+      assignmentEpoch: assignment.assignmentEpoch,
+      previousGeneration: this.generation,
+      controlResumeSecret: this.controlResumeSecret
+    })
+    this.activate(control, ack)
+    this.acceptingConnections = true
+    // Why: the resumed control owns the same server generation and splices;
+    // the predecessor remains only long enough for any idempotent reply in flight.
+    if (previous && previous.pendingRequestCount === 0) {
+      this.closeRetiredControl(previous)
+    } else if (previous) {
+      // Why: basis-bound requests keep their original control through its
+      // bounded request deadline; afterward the resumed control is sole owner.
+      this.retiredControlTimers.set(
+        previous,
+        setTimeout(() => this.closeRetiredControl(previous), 10_100)
+      )
+    }
+  }
+
+  markDraining(): void {
+    // The relay changes the control's protocol state when it sends drain. This
+    // marker exists for the broker's ownership policy, not a second wire event.
+    this.acceptingConnections = false
+  }
+
+  refreshAuthorization(relayJwt: string): void {
+    for (const control of this.controls) {
+      try {
+        control.refreshAuthorization(relayJwt)
+      } catch {
+        // A closing drain-only origin cannot block refresh on the active target.
+      }
+    }
+  }
+
+  async close(): Promise<void> {
+    if (this.closed) {
+      return
+    }
+    this.closed = true
+    for (const timer of this.retiredControlTimers.values()) {
+      clearTimeout(timer)
+    }
+    this.retiredControlTimers.clear()
+    for (const control of this.controls) {
+      control.closeNow()
+    }
+    this.controls.clear()
+    this.activeControl = null
+    await this.transport.stop()
+  }
+
+  closeNow(): void {
+    void this.close()
+  }
+
+  private async openControl(overrides?: {
+    relayJwt: string
+    assignmentEpoch: number
+    previousGeneration: number
+    controlResumeSecret: string
+  }): Promise<{ control: RelayControlClient; ack: RelayHostHelloAckMessage }> {
+    let control!: RelayControlClient
+    control = new RelayControlClient({
+      cellUrl: this.cellUrl,
+      relayJwt: overrides?.relayJwt ?? this.options.relayJwt,
+      relayHostId: this.options.relayHostId,
+      assignmentEpoch: overrides?.assignmentEpoch ?? this.assignmentEpoch,
+      identity: this.options.identity,
+      keypair: this.options.keypair,
+      appVersion: this.options.appVersion,
+      ...(overrides
+        ? {
+            previousGeneration: overrides.previousGeneration,
+            controlResumeSecret: overrides.controlResumeSecret
+          }
+        : {}),
+      onConnectionOpen: (message) => this.openConnection(message),
+      onDrain: (message) => this.options.onDrain(this, message),
+      onClose: (code) => {
+        this.controls.delete(control)
+        const timer = this.retiredControlTimers.get(control)
+        if (timer) {
+          clearTimeout(timer)
+          this.retiredControlTimers.delete(control)
+        }
+        if (this.activeControl === control) {
+          this.activeControl = null
+          this.options.onClose(this, code)
+        }
+      },
+      createSocket: this.options.createControlSocket
+    })
+    this.controls.add(control)
+    try {
+      return { control, ack: await control.connect() }
+    } catch (error) {
+      this.controls.delete(control)
+      control.closeNow()
+      throw error
+    }
+  }
+
+  private closeRetiredControl(control: RelayControlClient): void {
+    const timer = this.retiredControlTimers.get(control)
+    if (timer) {
+      clearTimeout(timer)
+      this.retiredControlTimers.delete(control)
+    }
+    if (this.activeControl !== control && this.controls.delete(control)) {
+      control.closeNow()
+    }
+  }
+
+  private activate(control: RelayControlClient, ack: RelayHostHelloAckMessage): void {
+    if (ack.generation <= 0) {
+      throw new Error('invalid_relay_generation')
+    }
+    this.transport.setGeneration(ack.generation)
+    this.generation = ack.generation
+    this.controlResumeSecret = ack.controlResumeSecret
+    this.leaseExpiresAt = ack.leaseExpiresAt
+    this.activeControl = control
+    for (const connectionId of ack.activeConnIds) {
+      this.options.onConnectionOwned(connectionId, this)
+    }
+  }
+
+  private openConnection(message: RelayConnectionOpenMessage): void {
+    if (!this.acceptingConnections) {
+      return
+    }
+    this.options.onConnectionOwned(message.connId, this)
+    void this.transport.openConnection(message).catch(() => {
+      this.options.onConnectionReleased(message.connId, this)
+    })
+  }
+}

src/main/runtime/relay/relay-control-requests.ts

@@ -25,6 +25,10 @@ export type DeviceCredentialInstallAuthorization =
 export class RelayControlRequests {
   private readonly pending = new Map<string, PendingRequest>()
 
+  get size(): number {
+    return this.pending.size
+  }
+
   createInvite(
     reqId: string,
     relayDeviceId: string,

src/main/runtime/relay/relay-origin-pool.ts

@@ -0,0 +1,291 @@
+import type WebSocket from 'ws'
+import type { E2EEKeypair } from '../e2ee-keypair'
+import type { MobileSocketWiring } from '../rpc/mobile-socket-wiring'
+import { RelayControlOrigin } from './relay-control-origin'
+import type { RelayControlClient } from './relay-control-client'
+import type { RelayDrainMessage } from './relay-control-protocol'
+import { requestRelayAssignment, type RelayAssignment } from './relay-http-client'
+import type { RelayBrokerStatus, RelayIdentity } from './relay-session-broker-contract'
+
+type RelayOriginPoolOptions = {
+  directorUrl: string
+  relayHostId: string
+  identity: RelayIdentity
+  keypair: E2EEKeypair
+  appVersion: string
+  mobileSocketWiring: MobileSocketWiring
+  isCurrent: () => boolean
+  onStatus: (status: RelayBrokerStatus) => void
+  fetch?: typeof globalThis.fetch
+  createControlSocket?: (url: string, relayJwt: string) => WebSocket
+  createDataSocket?: (url: string) => WebSocket
+  random?: () => number
+  now?: () => number
+}
+
+export class RelayOriginPool {
+  private readonly options: RelayOriginPoolOptions
+  private activeOrigin: RelayControlOrigin | null = null
+  private readonly origins = new Set<RelayControlOrigin>()
+  private readonly drainingOrigins = new Set<RelayControlOrigin>()
+  private readonly basisOrigins = new Map<string, RelayControlOrigin>()
+  private readonly drainTimers = new Map<RelayControlOrigin, ReturnType<typeof setTimeout>>()
+  private assignment: RelayAssignment | null = null
+  private relayJwt: string | null = null
+  private rotationTimer: ReturnType<typeof setTimeout> | null = null
+  private rotationPromise: Promise<void> | null = null
+  private closed = false
+
+  constructor(options: RelayOriginPoolOptions) {
+    this.options = options
+  }
+
+  get activeAssignment(): RelayAssignment | null {
+    return this.assignment
+  }
+
+  get activeControl(): RelayControlClient | null {
+    return this.activeOrigin?.availableControl ?? null
+  }
+
+  controlForBasis(basisConnId: string): RelayControlClient | null {
+    return this.basisOrigins.get(basisConnId)?.availableControl ?? null
+  }
+
+  async openInitial(assignment: RelayAssignment, relayJwt: string): Promise<void> {
+    this.assignment = assignment
+    this.relayJwt = relayJwt
+    const origin = this.createOrigin(assignment, relayJwt)
+    this.origins.add(origin)
+    await origin.open()
+    this.assertCurrent()
+    this.activeOrigin = origin
+    this.scheduleControlRotation()
+  }
+
+  refreshAuthorization(relayJwt: string): void {
+    this.relayJwt = relayJwt
+    for (const origin of this.origins) {
+      origin.refreshAuthorization(relayJwt)
+    }
+  }
+
+  closeNow(): void {
+    if (this.closed) {
+      return
+    }
+    this.closed = true
+    if (this.rotationTimer) {
+      clearTimeout(this.rotationTimer)
+      this.rotationTimer = null
+    }
+    for (const timer of this.drainTimers.values()) {
+      clearTimeout(timer)
+    }
+    this.drainTimers.clear()
+    for (const origin of this.origins) {
+      origin.closeNow()
+    }
+    this.origins.clear()
+    this.drainingOrigins.clear()
+    this.basisOrigins.clear()
+    this.activeOrigin = null
+  }
+
+  private createOrigin(assignment: RelayAssignment, relayJwt: string): RelayControlOrigin {
+    return new RelayControlOrigin({
+      assignment,
+      relayJwt,
+      relayHostId: this.options.relayHostId,
+      identity: this.options.identity,
+      keypair: this.options.keypair,
+      appVersion: this.options.appVersion,
+      mobileSocketWiring: this.options.mobileSocketWiring,
+      createControlSocket: this.options.createControlSocket,
+      createDataSocket: this.options.createDataSocket,
+      onConnectionOwned: (connectionId, origin) => {
+        if (this.isCurrent() && this.origins.has(origin)) {
+          this.basisOrigins.set(connectionId, origin)
+        }
+      },
+      onConnectionReleased: (connectionId, origin) => {
+        if (this.basisOrigins.get(connectionId) === origin) {
+          this.basisOrigins.delete(connectionId)
+        }
+        this.maybeCloseDrainedOrigin(origin)
+      },
+      onDrain: (origin, message) => this.handleDrain(origin, message),
+      onClose: (origin) => {
+        if (origin === this.activeOrigin && this.isCurrent()) {
+          this.options.onStatus('offline')
+          this.handleDrain(origin, {
+            type: 'drain',
+            graceMs: 0,
+            recovery: 'resolve-director'
+          })
+        }
+      }
+    })
+  }
+
+  private handleDrain(origin: RelayControlOrigin, message: RelayDrainMessage): void {
+    if (!this.isCurrent() || origin !== this.activeOrigin) {
+      return
+    }
+    origin.markDraining()
+    this.drainingOrigins.add(origin)
+    this.options.onStatus('draining')
+    if (!this.rotationPromise) {
+      this.rotationPromise = this.resolveDrainTarget(origin, message).finally(() => {
+        this.rotationPromise = null
+      })
+    }
+  }
+
+  private async resolveDrainTarget(
+    origin: RelayControlOrigin,
+    message: RelayDrainMessage
+  ): Promise<void> {
+    try {
+      if (!this.relayJwt) {
+        throw new Error('relay_authorization_unavailable')
+      }
+      // Why: only the configured director can choose a migration target.
+      const assignment = await requestRelayAssignment({
+        directorUrl: this.options.directorUrl,
+        relayToken: this.relayJwt,
+        relayHostId: this.options.relayHostId,
+        fetch: this.options.fetch
+      })
+      this.assertCurrent()
+      if (assignment.cellUrl === origin.cellUrl) {
+        await origin.rebind(this.relayJwt, assignment)
+        this.assertCurrent()
+        this.activeOrigin = origin
+        this.assignment = assignment
+        this.drainingOrigins.delete(origin)
+      } else {
+        await this.activateTarget(origin, assignment, this.relayJwt, message.graceMs)
+      }
+      this.options.onStatus('registered')
+      this.scheduleControlRotation()
+    } catch {
+      if (this.isCurrent()) {
+        const random = this.options.random ?? Math.random
+        setTimeout(() => this.handleDrain(origin, message), 250 + Math.floor(random() * 751))
+      }
+    }
+  }
+
+  private async activateTarget(
+    origin: RelayControlOrigin,
+    assignment: RelayAssignment,
+    relayJwt: string,
+    graceMs: number
+  ): Promise<void> {
+    const target = this.createOrigin(assignment, relayJwt)
+    this.origins.add(target)
+    try {
+      await target.open()
+      this.assertCurrent()
+    } catch (error) {
+      this.origins.delete(target)
+      target.closeNow()
+      throw error
+    }
+    this.activeOrigin = target
+    this.assignment = assignment
+    this.scheduleDrainDeadline(origin, graceMs)
+    this.maybeCloseDrainedOrigin(origin)
+  }
+
+  private scheduleControlRotation(): void {
+    if (this.rotationTimer) {
+      clearTimeout(this.rotationTimer)
+    }
+    const origin = this.activeOrigin
+    if (!origin || this.closed) {
+      this.rotationTimer = null
+      return
+    }
+    const now = (this.options.now ?? Date.now)()
+    const random = this.options.random ?? Math.random
+    const earlyMs = 60_000 + Math.floor(random() * 60_001)
+    const delay = Math.max(0, origin.controlLeaseExpiresAt - earlyMs - now)
+    this.rotationTimer = setTimeout(() => void this.rebindActiveControl(origin), delay)
+  }
+
+  private async rebindActiveControl(origin: RelayControlOrigin): Promise<void> {
+    this.rotationTimer = null
+    if (!this.isCurrent() || origin !== this.activeOrigin || this.rotationPromise) {
+      return
+    }
+    if (!this.relayJwt || !this.assignment) {
+      return
+    }
+    try {
+      await origin.rebind(this.relayJwt, this.assignment)
+      this.assertCurrent()
+      this.scheduleControlRotation()
+    } catch {
+      if (this.isCurrent() && origin === this.activeOrigin) {
+        const random = this.options.random ?? Math.random
+        this.rotationTimer = setTimeout(
+          () => void this.rebindActiveControl(origin),
+          5_000 + Math.floor(random() * 10_001)
+        )
+      }
+    }
+  }
+
+  private scheduleDrainDeadline(origin: RelayControlOrigin, graceMs: number): void {
+    const existing = this.drainTimers.get(origin)
+    if (existing) {
+      clearTimeout(existing)
+    }
+    this.drainTimers.set(
+      origin,
+      setTimeout(() => this.closeOrigin(origin), graceMs)
+    )
+  }
+
+  private maybeCloseDrainedOrigin(origin: RelayControlOrigin): void {
+    if (
+      !this.drainingOrigins.has(origin) ||
+      origin.pendingRequestCount > 0 ||
+      [...this.basisOrigins.values()].includes(origin)
+    ) {
+      return
+    }
+    this.closeOrigin(origin)
+  }
+
+  private closeOrigin(origin: RelayControlOrigin): void {
+    if (origin === this.activeOrigin) {
+      return
+    }
+    const timer = this.drainTimers.get(origin)
+    if (timer) {
+      clearTimeout(timer)
+      this.drainTimers.delete(origin)
+    }
+    for (const [connectionId, owner] of this.basisOrigins) {
+      if (owner === origin) {
+        this.basisOrigins.delete(connectionId)
+      }
+    }
+    this.drainingOrigins.delete(origin)
+    this.origins.delete(origin)
+    origin.closeNow()
+  }
+
+  private assertCurrent(): void {
+    if (!this.isCurrent()) {
+      throw new Error('stale_relay_origin_pool')
+    }
+  }
+
+  private isCurrent(): boolean {
+    return !this.closed && this.options.isCurrent()
+  }
+}

src/main/runtime/relay/relay-session-broker-contract.ts

@@ -22,7 +22,6 @@ export type RelaySessionBrokerOptions = {
   isCurrent: () => boolean
   refreshAccessToken: () => Promise<string | null>
   onStatus: (status: RelayBrokerStatus) => void
-  onResolveDirector: () => void
   fetch?: typeof globalThis.fetch
   createControlSocket?: (url: string, relayJwt: string) => WebSocket
   createDataSocket?: (url: string) => WebSocket

src/main/runtime/relay/relay-session-broker.test.ts

@@ -6,14 +6,30 @@ import type * as RelayHttpClientModule from './relay-http-client'
 
 const fakes = vi.hoisted(() => ({
   controls: [] as {
+    options: {
+      onConnectionOpen(message: {
+        connId: string
+        connTicket: string
+        kind: 'invite' | 'resume'
+        relayDeviceId: string
+        attachDeadlineMs: number
+      }): void
+      onDrain(message: { type: 'drain'; graceMs: number; recovery: 'resolve-director' }): void
+      previousGeneration?: number
+      controlResumeSecret?: string
+    }
     connect: ReturnType<typeof vi.fn>
     closeNow: ReturnType<typeof vi.fn>
+    confirmResume: ReturnType<typeof vi.fn>
+    installCredential: ReturnType<typeof vi.fn>
+    pendingRequestCount: number
   }[],
   transports: [] as {
     start: ReturnType<typeof vi.fn>
     stop: ReturnType<typeof vi.fn>
     setGeneration: ReturnType<typeof vi.fn>
     metadataFor: ReturnType<typeof vi.fn>
+    openConnection: ReturnType<typeof vi.fn>
   }[],
   controlConnect: vi.fn(),
   exchange: vi.fn(),
@@ -30,8 +46,26 @@ vi.mock('./relay-control-client', () => ({
   RelayControlClient: class {
     connect = fakes.controlConnect
     closeNow = vi.fn()
+    confirmResume = vi.fn().mockResolvedValue({
+      type: 'device-resume-confirmed',
+      v: 1,
+      reqId: 'confirm-1',
+      currentVersion: 1,
+      acceptedAs: 'current',
+      renewed: true,
+      resumeExpiresAt: 100_000
+    })
+    installCredential = vi.fn().mockResolvedValue({
+      type: 'device-credential-installed',
+      v: 1,
+      reqId: 'install-1',
+      authorizationMode: 'relay-basis',
+      currentVersion: 1,
+      resumeExpiresAt: 100_000
+    })
+    pendingRequestCount = 0
 
-    constructor() {
+    constructor(readonly options: (typeof fakes.controls)[number]['options']) {
       fakes.controls.push(this)
     }
   }
@@ -43,6 +77,7 @@ vi.mock('../rpc/relay-transport', () => ({
     stop = vi.fn().mockResolvedValue(undefined)
     setGeneration = vi.fn()
     metadataFor = vi.fn()
+    openConnection = vi.fn().mockResolvedValue(undefined)
 
     constructor() {
       fakes.transports.push(this)
@@ -65,7 +100,7 @@ describe('RelaySessionBroker lifecycle ownership', () => {
     fakes.controls.length = 0
     fakes.transports.length = 0
     fakes.controlConnect.mockReset()
-    fakes.exchange.mockReset().mockResolvedValue({ relayToken: 'relay-jwt', expiresAt: 60_000 })
+    fakes.exchange.mockReset().mockResolvedValue({ relayToken: 'relay-jwt', expiresAt: 1_000_000 })
     fakes.assign.mockReset().mockResolvedValue({
       cellUrl: 'https://relay.example.test',
       assignmentEpoch: 1,
@@ -94,8 +129,7 @@ describe('RelaySessionBroker lifecycle ownership', () => {
       mobileSocketWiring: { attachTransport: vi.fn() } as never,
       isCurrent: () => current,
       refreshAccessToken: async () => null,
-      onStatus: (status) => statuses.push(status),
-      onResolveDirector: vi.fn()
+      onStatus: (status) => statuses.push(status)
     })
     await vi.waitFor(() => expect(fakes.controls).toHaveLength(1))
     current = false
@@ -104,7 +138,7 @@ describe('RelaySessionBroker lifecycle ownership', () => {
       v: 1,
       generation: 1,
       controlResumeSecret: 'A'.repeat(43),
-      leaseExpiresAt: 60_000,
+      leaseExpiresAt: 1_000_000,
       activeConnIds: [],
       pendingConns: []
     })
@@ -114,4 +148,134 @@ describe('RelaySessionBroker lifecycle ownership', () => {
     expect(fakes.transports[0]!.stop).toHaveBeenCalledOnce()
     expect(statuses).toEqual(['connecting'])
   })
+
+  it('activates a new origin while keeping basis-bound work on the drained origin', async () => {
+    const firstAck: RelayHostHelloAckMessage = {
+      type: 'host-hello-ack',
+      v: 1,
+      generation: 1,
+      controlResumeSecret: 'A'.repeat(43),
+      leaseExpiresAt: 1_000_000,
+      activeConnIds: [],
+      pendingConns: []
+    }
+    fakes.controlConnect.mockResolvedValueOnce(firstAck).mockResolvedValueOnce({
+      ...firstAck,
+      generation: 2,
+      controlResumeSecret: 'B'.repeat(43)
+    })
+    fakes.assign
+      .mockResolvedValueOnce({
+        cellUrl: 'https://relay-c1.example.test',
+        assignmentEpoch: 1,
+        leaseExpiresAt: 1_000_000
+      })
+      .mockResolvedValueOnce({
+        cellUrl: 'https://relay-c2.example.test',
+        assignmentEpoch: 2,
+        leaseExpiresAt: 2_000_000
+      })
+    const broker = await RelaySessionBroker.connect(brokerOptions({ onStatus: vi.fn() }))
+    fakes.controls[0]!.options.onConnectionOpen({
+      connId: 'old-basis',
+      connTicket: 'T'.repeat(43),
+      kind: 'resume',
+      relayDeviceId: 'device-1',
+      attachDeadlineMs: 1_000
+    })
+    expect(brokerBasisIds(broker)).toEqual(['old-basis'])
+    fakes.controls[0]!.options.onDrain({
+      type: 'drain',
+      graceMs: 30_000,
+      recovery: 'resolve-director'
+    })
+    await vi.waitFor(() => expect(fakes.controls).toHaveLength(2))
+
+    expect(broker.endpoint?.cellUrl).toBe('https://relay-c2.example.test')
+    expect(fakes.transports[0]!.openConnection).toHaveBeenCalledOnce()
+    expect(brokerBasisIds(broker)).toEqual(['old-basis'])
+    expect(fakes.controls[0]!.closeNow).not.toHaveBeenCalled()
+    expect(fakes.transports[0]!.stop).not.toHaveBeenCalled()
+    await broker.confirmResume('old-basis', 'confirm-1')
+    expect(fakes.controls[0]!.confirmResume).toHaveBeenCalledWith('old-basis', 'confirm-1')
+    await broker.installCredential(
+      'device-1',
+      { reqId: 'install-1', newResumeTokenHash: 'H'.repeat(43) },
+      { mode: 'relay-basis', basisConnId: 'old-basis' }
+    )
+    expect(fakes.controls[0]!.installCredential).toHaveBeenCalledOnce()
+  })
+
+  it('rebinds the same process generation with its control resume secret', async () => {
+    const ack: RelayHostHelloAckMessage = {
+      type: 'host-hello-ack',
+      v: 1,
+      generation: 7,
+      controlResumeSecret: 'R'.repeat(43),
+      leaseExpiresAt: 1_000_000,
+      activeConnIds: ['existing-basis'],
+      pendingConns: []
+    }
+    fakes.controlConnect.mockResolvedValueOnce(ack).mockResolvedValueOnce({
+      ...ack,
+      leaseExpiresAt: 2_000_000
+    })
+    fakes.assign
+      .mockResolvedValueOnce({
+        cellUrl: 'https://relay.example.test',
+        assignmentEpoch: 1,
+        leaseExpiresAt: 1_000_000
+      })
+      .mockResolvedValueOnce({
+        cellUrl: 'https://relay.example.test',
+        assignmentEpoch: 1,
+        leaseExpiresAt: 2_000_000
+      })
+    const broker = await RelaySessionBroker.connect(brokerOptions())
+    expect(brokerBasisIds(broker)).toEqual(['existing-basis'])
+    fakes.controls[0]!.options.onDrain({
+      type: 'drain',
+      graceMs: 5_000,
+      recovery: 'resolve-director'
+    })
+    await vi.waitFor(() => expect(fakes.controls).toHaveLength(2))
+
+    expect(fakes.transports).toHaveLength(1)
+    expect(fakes.controls[1]!.options.previousGeneration).toBe(7)
+    expect(fakes.controls[1]!.options.controlResumeSecret).toBe('R'.repeat(43))
+    await vi.waitFor(() => expect(brokerBasisIds(broker)).toEqual(['existing-basis']))
+    await broker.confirmResume('existing-basis', 'confirm-1')
+    expect(fakes.controls[1]!.confirmResume).toHaveBeenCalledOnce()
+  })
 })
+
+function brokerBasisIds(broker: RelaySessionBroker): string[] {
+  const pool = (broker as unknown as { originPool: unknown }).originPool
+  return [...(pool as { basisOrigins: Map<string, unknown> }).basisOrigins.keys()]
+}
+
+function brokerOptions(
+  overrides: Partial<Parameters<typeof RelaySessionBroker.connect>[0]> = {}
+): Parameters<typeof RelaySessionBroker.connect>[0] {
+  const keypair = nacl.box.keyPair()
+  return {
+    authConfig: {
+      relayTokenEndpoint: 'https://auth.example.test/v1/relay-token',
+      relayDirectorUrl: 'https://relay.example.test'
+    } as OrcaCloudAuthConfig,
+    accessToken: 'access-token',
+    identity: { userId: 'user-1', profileId: 'profile-1', organizationId: 'org-1' },
+    keypair: {
+      ...keypair,
+      publicKeyB64: Buffer.from(keypair.publicKey).toString('base64')
+    },
+    appVersion: '1.0.0',
+    mobileSocketWiring: { attachTransport: vi.fn() } as never,
+    isCurrent: () => true,
+    refreshAccessToken: async () => null,
+    onStatus: vi.fn(),
+    now: () => 0,
+    random: () => 0,
+    ...overrides
+  }
+}

src/main/runtime/relay/relay-session-broker.ts

@@ -6,8 +6,6 @@ import type {
   MobileRelayEndpoint,
   PairingProvisionRelayParams
 } from '../../../shared/mobile-relay-credential-contract'
-import { CloudRelayTransport } from '../rpc/relay-transport'
-import { RelayControlClient } from './relay-control-client'
 import type { DeviceCredentialInstallAuthorization } from './relay-control-requests'
 import {
   deriveRelayHostId,
@@ -16,6 +14,7 @@ import {
   type RelayAuthorization,
   type RelayAssignment
 } from './relay-http-client'
+import { RelayOriginPool } from './relay-origin-pool'
 import type { RelayBrokerStatus, RelaySessionBrokerOptions } from './relay-session-broker-contract'
 
 export type { RelayBrokerStatus } from './relay-session-broker-contract'
@@ -29,16 +28,29 @@ export class StaleRelayBrokerError extends Error {
 export class RelaySessionBroker {
   private readonly options: RelaySessionBrokerOptions
   private readonly relayHostId: string
-  private control: RelayControlClient | null = null
-  private transport: CloudRelayTransport | null = null
+  private readonly originPool: RelayOriginPool
   private authorization: RelayAuthorization | null = null
-  private assignment: RelayAssignment | null = null
   private refreshTimer: ReturnType<typeof setTimeout> | null = null
   private closed = false
 
   private constructor(options: RelaySessionBrokerOptions) {
     this.options = options
     this.relayHostId = deriveRelayHostId(options.keypair.publicKey)
+    this.originPool = new RelayOriginPool({
+      directorUrl: options.authConfig.relayDirectorUrl,
+      relayHostId: this.relayHostId,
+      identity: options.identity,
+      keypair: options.keypair,
+      appVersion: options.appVersion,
+      mobileSocketWiring: options.mobileSocketWiring,
+      isCurrent: () => this.isCurrent(),
+      onStatus: (status) => this.publishStatus(status),
+      fetch: options.fetch,
+      createControlSocket: options.createControlSocket,
+      createDataSocket: options.createDataSocket,
+      random: options.random,
+      now: options.now
+    })
   }
 
   static async connect(options: RelaySessionBrokerOptions): Promise<RelaySessionBroker> {
@@ -57,7 +69,7 @@ export class RelaySessionBroker {
   }
 
   get currentAssignment(): RelayAssignment | null {
-    return this.assignment
+    return this.originPool.activeAssignment
   }
 
   get ownerIdentityKey(): string {
@@ -66,7 +78,7 @@ export class RelaySessionBroker {
   }
 
   get endpoint(): MobileRelayEndpoint | null {
-    const assignment = this.assignment
+    const assignment = this.originPool.activeAssignment
     if (!assignment) {
       return null
     }
@@ -80,19 +92,21 @@ export class RelaySessionBroker {
     }
   }
 
-  createInvite(relayDeviceId: string): ReturnType<RelayControlClient['createInvite']> {
-    if (!this.control) {
+  createInvite(relayDeviceId: string) {
+    const control = this.originPool.activeControl
+    if (!control) {
       return Promise.reject(new Error('relay_control_not_active'))
     }
-    return this.control.createInvite(relayDeviceId)
+    return control.createInvite(relayDeviceId)
   }
 
   async createPairingRelay(relayDeviceId: string): Promise<PairingRelay> {
-    const assignment = this.assignment
-    if (!assignment || !this.control) {
+    const assignment = this.originPool.activeAssignment
+    const control = this.originPool.activeControl
+    if (!assignment || !control) {
       throw new Error('relay_control_not_active')
     }
-    const invite = await this.control.createInvite(relayDeviceId)
+    const invite = await control.createInvite(relayDeviceId)
     this.assertCurrent()
     return {
       v: 1,
@@ -107,10 +121,11 @@ export class RelaySessionBroker {
   }
 
   revokeDevice(relayDeviceId: string, reqId?: string): Promise<void> {
-    if (!this.control) {
+    const control = this.originPool.activeControl
+    if (!control) {
       return Promise.reject(new Error('relay_control_not_active'))
     }
-    return this.control.revokeDevice(relayDeviceId, reqId)
+    return control.revokeDevice(relayDeviceId, reqId)
   }
 
   async installCredential(
@@ -118,10 +133,14 @@ export class RelaySessionBroker {
     params: PairingProvisionRelayParams,
     authorization: DeviceCredentialInstallAuthorization
   ): Promise<DeviceCredentialInstalled> {
-    if (!this.control) {
+    const control =
+      authorization.mode === 'relay-basis'
+        ? this.originPool.controlForBasis(authorization.basisConnId)
+        : this.originPool.activeControl
+    if (!control) {
       throw new Error('relay_control_not_active')
     }
-    const message = await this.control.installCredential({
+    const message = await control.installCredential({
       relayDeviceId,
       authorization,
       ...params
@@ -135,20 +154,22 @@ export class RelaySessionBroker {
     relayDeviceId: string,
     reqId: string
   ): Promise<DeviceCredentialInstallStatusResult> {
-    if (!this.control) {
+    const control = this.originPool.activeControl
+    if (!control) {
       throw new Error('relay_control_not_active')
     }
-    const message = await this.control.credentialInstallStatus(relayDeviceId, reqId)
+    const message = await control.credentialInstallStatus(relayDeviceId, reqId)
     this.assertCurrent()
     const { type: _type, ...result } = message
     return result
   }
 
   async confirmResume(basisConnId: string, reqId: string): Promise<DeviceResumeConfirmed> {
-    if (!this.control) {
-      throw new Error('relay_control_not_active')
+    const control = this.originPool.controlForBasis(basisConnId)
+    if (!control) {
+      throw new Error('relay_basis_origin_not_found')
     }
-    const message = await this.control.confirmResume(basisConnId, reqId)
+    const message = await control.confirmResume(basisConnId, reqId)
     this.assertCurrent()
     const { type: _type, ...result } = message
     return result
@@ -164,10 +185,7 @@ export class RelaySessionBroker {
       clearTimeout(this.refreshTimer)
       this.refreshTimer = null
     }
-    this.control?.closeNow()
-    this.control = null
-    void this.transport?.stop()
-    this.transport = null
+    this.originPool.closeNow()
     if (publishOffline) {
       this.options.onStatus('offline')
     }
@@ -189,56 +207,16 @@ export class RelaySessionBroker {
       fetch: this.options.fetch
     })
     this.assertCurrent()
-    const transport = new CloudRelayTransport({
-      cellUrl: assignment.cellUrl,
-      relayHostId: this.relayHostId,
-      generation: 0,
-      createSocket: this.options.createDataSocket
-    })
-    // Why: stale-epoch cleanup must own partially opened resources before the
-    // next await, even though the broker is not externally active yet.
-    this.transport = transport
-    this.options.mobileSocketWiring.attachTransport(transport, (ws) => transport.metadataFor(ws))
-    await transport.start()
-    this.assertCurrent()
-    const control = new RelayControlClient({
-      cellUrl: assignment.cellUrl,
-      relayJwt: authorization.relayToken,
-      relayHostId: this.relayHostId,
-      assignmentEpoch: assignment.assignmentEpoch,
-      identity: this.options.identity,
-      keypair: this.options.keypair,
-      appVersion: this.options.appVersion,
-      onConnectionOpen: (message) => {
-        if (this.isCurrent()) {
-          void transport.openConnection(message).catch(() => {})
-        }
-      },
-      onDrain: () => {
-        if (!this.isCurrent()) {
-          return
-        }
-        this.publishStatus('draining')
-        // Why: the data plane never chooses recovery targets; every drain
-        // returns to the configured stable director.
-        this.options.onResolveDirector()
-      },
-      onClose: () => {
-        if (this.isCurrent()) {
-          this.publishStatus('offline')
-        }
-      },
-      createSocket: this.options.createControlSocket
-    })
-    this.control = control
-    const ack = await control.connect()
-    this.assertCurrent()
-    if (ack.generation <= 0) {
-      throw new Error('invalid_relay_generation')
+    try {
+      await this.originPool.openInitial(assignment, authorization.relayToken)
+    } catch (error) {
+      if (!this.isCurrent()) {
+        throw new StaleRelayBrokerError()
+      }
+      throw error
     }
-    transport.setGeneration(ack.generation)
+    this.assertCurrent()
     this.authorization = authorization
-    this.assignment = assignment
     this.publishStatus('registered')
     this.scheduleRefresh()
   }
@@ -271,7 +249,7 @@ export class RelaySessionBroker {
         fetch: this.options.fetch
       })
       this.assertCurrent()
-      this.control?.refreshAuthorization(authorization.relayToken)
+      this.originPool.refreshAuthorization(authorization.relayToken)
       this.authorization = authorization
       this.scheduleRefresh()
     } catch {

src/main/runtime/rpc/relay-transport.test.ts

@@ -39,6 +39,7 @@ describe('CloudRelayTransport', () => {
       server.once('connection', (socket, request) => resolve({ socket, path: request.url ?? '' }))
     })
     let clientSocket: WebSocketClient | null = null
+    const onConnectionClosed = vi.fn()
     const transport = new CloudRelayTransport({
       cellUrl: `http://127.0.0.1:${address.port}`,
       relayHostId: 'AbCdEf0123_-xyZ9',
@@ -46,7 +47,8 @@ describe('CloudRelayTransport', () => {
       createSocket: (url) => {
         clientSocket = new WebSocketClient(url, { perMessageDeflate: false })
         return clientSocket
-      }
+      },
+      onConnectionClosed
     })
     transports.push(transport)
     const received: (string | Uint8Array<ArrayBufferLike>)[] = []
@@ -87,6 +89,8 @@ describe('CloudRelayTransport', () => {
       basisConnId: 'conn/with spaces',
       credentialKind: 'resume'
     })
+    socket.close()
+    await vi.waitFor(() => expect(onConnectionClosed).toHaveBeenCalledWith('conn/with spaces'))
   })
 
   it('rejects non-origin cell URLs before opening a socket', () => {

src/main/runtime/rpc/relay-transport.ts

@@ -19,6 +19,7 @@ type CloudRelayTransportOptions = {
   relayHostId: string
   generation: number
   createSocket?: (url: string) => WebSocket
+  onConnectionClosed?: (connectionId: string) => void
 }
 
 function relayWebSocketOrigin(cellUrl: string): string {
@@ -41,6 +42,7 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport
   private readonly relayHostId: string
   private generation: number
   private readonly createSocket: (url: string) => WebSocket
+  private readonly onConnectionClosed: ((connectionId: string) => void) | undefined
   private readonly socketsByConnectionId = new Map<string, WebSocket>()
   private readonly metadataBySocket = new Map<WebSocket, MobileSocketTransportMetadata>()
   private readonly clientIds = new Map<WebSocket, string>()
@@ -52,6 +54,7 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport
     this.cellWebSocketOrigin = relayWebSocketOrigin(options.cellUrl)
     this.relayHostId = options.relayHostId
     this.generation = options.generation
+    this.onConnectionClosed = options.onConnectionClosed
     this.createSocket =
       options.createSocket ??
       ((url) =>
@@ -81,6 +84,9 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport
   }
 
   setGeneration(generation: number): void {
+    if (generation === this.generation) {
+      return
+    }
     if (
       this.socketsByConnectionId.size > 0 ||
       !Number.isSafeInteger(generation) ||
@@ -153,6 +159,7 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport
         this.metadataBySocket.delete(socket)
         const clientId = this.clientIds.get(socket) ?? null
         this.clientIds.delete(socket)
+        this.onConnectionClosed?.(connection.connId)
         const hasOtherConnections =
           clientId !== null && [...this.clientIds.values()].includes(clientId)
         this.closeHandler?.(clientId, socket, hasOtherConnections)