fix(ssh): keep recovery leases stable across reconnects

src/relay/relay-pty-consumer-owner-displacement.test.ts

@@ -6,6 +6,7 @@ import {
 } from './dispatcher'
 import { encodeJsonRpcFrame, MessageType } from './protocol'
 import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter'
+import { PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR } from '../shared/pty-consumer-session'
 
 const endpointIdentity: RelayClientSessionIdentity = {
   principal: 'endpoint-principal',
@@ -272,4 +273,78 @@ describe('relay PTY consumer owner displacement', () => {
     expect(adapter.deliveryMode(1)).toBe('subscriber')
     expect(detached).toContain(1)
   })
+
+  it('serializes overlapping reconnect grants without invalidating the original proof', async () => {
+    const incumbentWrites: Buffer[] = []
+    dispatcher = new RelayDispatcher(
+      (data, settle) => {
+        incumbentWrites.push(Buffer.from(data))
+        settle({ ok: true })
+        return true
+      },
+      { supportsWriteCallback: true },
+      endpointIdentity
+    )
+    const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a')
+
+    dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams()))
+    await flushRequests()
+    const incumbentGrant = response(incumbentWrites, 1)!.result as Record<string, unknown>
+    const resume = {
+      ownerGeneration: incumbentGrant.ownerGeneration,
+      ownerLease: incumbentGrant.ownerLease
+    }
+
+    let firstReconnectSettlement: ((result: SinkWriteSettlement) => void) | undefined
+    const firstReconnectWrites: Buffer[] = []
+    const firstReconnectClientId = dispatcher.attachClient(
+      (data, settle) => {
+        firstReconnectWrites.push(Buffer.from(data))
+        firstReconnectSettlement = settle
+        return true
+      },
+      { supportsWriteCallback: true },
+      endpointIdentity
+    )
+    dispatcher.feedClient(
+      firstReconnectClientId,
+      requestFrame(2, 'pty.openClient', ownerHelloParams(resume))
+    )
+    await flushRequests()
+
+    const retryWrites: Buffer[] = []
+    const retryClientId = dispatcher.attachClient(
+      (data, settle) => {
+        retryWrites.push(Buffer.from(data))
+        settle({ ok: true })
+        return true
+      },
+      { supportsWriteCallback: true },
+      endpointIdentity
+    )
+    dispatcher.feedClient(
+      retryClientId,
+      requestFrame(3, 'pty.openClient', ownerHelloParams(resume))
+    )
+    await flushRequests()
+
+    expect(response(retryWrites, 3)!.error).toMatchObject({
+      code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
+    })
+
+    firstReconnectSettlement!({ ok: true })
+    dispatcher.feedClient(
+      retryClientId,
+      requestFrame(4, 'pty.openClient', ownerHelloParams(resume))
+    )
+    await flushRequests()
+
+    expect(response(retryWrites, 4)!.result).toMatchObject({
+      role: 'session-owner',
+      ownerGeneration: 3,
+      ownerLease: incumbentGrant.ownerLease
+    })
+    expect(adapter.deliveryMode(firstReconnectClientId)).toBe('subscriber')
+    expect(adapter.deliveryMode(retryClientId)).toBe('source-owner')
+  })
 })

src/shared/pty-consumer-session.test.ts

@@ -122,7 +122,7 @@ describe('PtyConsumerSession', () => {
     ).toThrow('authentication required')
   })
 
-  it('rotates the lease and increments owner generation on valid recovery', () => {
+  it('keeps the lease stable and increments owner generation on valid recovery', () => {
     const session = createSession()
     const first = session.admit(ownerHello(), auth('connection-1'))
     first.commitPublication()
@@ -142,7 +142,7 @@ describe('PtyConsumerSession', () => {
     expect(recovered.grant).toMatchObject({
       role: 'session-owner',
       ownerGeneration: 2,
-      ownerLease: 'lease-2'
+      ownerLease: 'lease-1'
     })
   })
 
@@ -251,6 +251,56 @@ describe('PtyConsumerSession', () => {
     )
   })
 
+  it('retries overlapping recovery against the stable lease after publication commits', () => {
+    const session = createSession()
+    const first = session.admit(ownerHello(), auth('connection-1'))
+    first.commitPublication()
+    const resume = {
+      ownerGeneration: first.grant.ownerGeneration!,
+      ownerLease: first.grant.ownerLease!
+    }
+    const replacement = session.admit(ownerHello({ resume }), auth('connection-2'))
+
+    expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow(
+      expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR })
+    )
+
+    replacement.commitPublication()
+    const retry = session.admit(ownerHello({ resume }), auth('connection-3'))
+
+    expect(retry.grant).toMatchObject({
+      role: 'session-owner',
+      ownerGeneration: 3,
+      ownerLease: first.grant.ownerLease
+    })
+    expect(retry.displacedOwner?.connectionId).toBe('connection-2')
+  })
+
+  it('retries overlapping recovery against the incumbent after publication rolls back', () => {
+    const session = createSession()
+    const first = session.admit(ownerHello(), auth('connection-1'))
+    first.commitPublication()
+    const resume = {
+      ownerGeneration: first.grant.ownerGeneration!,
+      ownerLease: first.grant.ownerLease!
+    }
+    const replacement = session.admit(ownerHello({ resume }), auth('connection-2'))
+
+    expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow(
+      expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR })
+    )
+
+    replacement.rollbackPublication()
+    const retry = session.admit(ownerHello({ resume }), auth('connection-3'))
+
+    expect(retry.grant).toMatchObject({
+      role: 'session-owner',
+      ownerGeneration: 3,
+      ownerLease: first.grant.ownerLease
+    })
+    expect(retry.displacedOwner?.connectionId).toBe('connection-1')
+  })
+
   it('types mismatched recovery without disturbing principal or lease ownership', () => {
     const session = createSession()
     const first = session.admit(ownerHello(), auth('connection-1'))
@@ -289,7 +339,7 @@ describe('PtyConsumerSession', () => {
     expect(recovered.grant).toMatchObject({
       role: 'session-owner',
       ownerGeneration: 2,
-      ownerLease: 'lease-2'
+      ownerLease: 'lease-1'
     })
   })
 

src/shared/pty-consumer-session.ts

@@ -229,7 +229,6 @@ export class PtyConsumerSession {
       return null
     }
     const recoveryMatches =
-      hello.resume.ownerGeneration === current.generation &&
       hello.resume.ownerLease === current.lease &&
       hello.clientInstanceId === current.clientInstanceId &&
       authentication.principal === current.principal
@@ -284,7 +283,7 @@ export class PtyConsumerSession {
     authentication: PtyConsumerAuthentication,
     replaces: OwnerRecord | null
   ): OwnerRecord {
-    const lease = this.createLease()
+    const lease = replaces?.lease ?? this.createLease()
     assertNonEmptyString(lease, 'ownerLease')
     return {
       connectionId: authentication.connectionId,