diff --git a/src/relay/relay-pty-consumer-owner-displacement.test.ts b/src/relay/relay-pty-consumer-owner-displacement.test.ts index e91d41d6852..f8db3a37c8f 100644 --- a/src/relay/relay-pty-consumer-owner-displacement.test.ts +++ b/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 + 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') + }) }) diff --git a/src/shared/pty-consumer-session.test.ts b/src/shared/pty-consumer-session.test.ts index bbfed831e1a..8f315672bd9 100644 --- a/src/shared/pty-consumer-session.test.ts +++ b/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' }) }) diff --git a/src/shared/pty-consumer-session.ts b/src/shared/pty-consumer-session.ts index 084c6638ec5..4afad1dd822 100644 --- a/src/shared/pty-consumer-session.ts +++ b/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,