diff --git a/playwriter/src/cdp-relay.ts b/playwriter/src/cdp-relay.ts index b2ca68b..c3ba9f4 100644 --- a/playwriter/src/cdp-relay.ts +++ b/playwriter/src/cdp-relay.ts @@ -387,7 +387,9 @@ export async function startPlayWriterCDPRelayServer({ return s } id = ext.messageId + 1 - return relayState.incrementExtensionMessageId(s, { extensionId: resolvedExtensionId }) + const newExtensions = new Map(s.extensions) + newExtensions.set(resolvedExtensionId, { ...ext, messageId: id }) + return { ...s, extensions: newExtensions } }) if (!id) { @@ -1559,14 +1561,13 @@ export async function startPlayWriterCDPRelayServer({ } recordingRelays.delete(connectionId) - // Reject all pending I/O requests + // Reject all pending I/O requests (state cleanup happens in removeExtension below) const closingExt = store.getState().extensions.get(connectionId) if (closingExt) { stopExtensionPing(connectionId) for (const pending of closingExt.pendingRequests.values()) { pending.reject(new Error('Extension connection closed')) } - store.setState((s) => relayState.clearExtensionPendingRequests(s, { extensionId: connectionId })) } const currentRelayState = store.getState() @@ -1601,12 +1602,8 @@ export async function startPlayWriterCDPRelayServer({ } } - // State transitions: remove extension + bound clients atomically - store.setState((s) => { - let next = relayState.removeExtension(s, { extensionId: connectionId }) - next = relayState.removeClientsForExtension(next, { extensionId: connectionId }) - return next - }) + // State transition: remove extension + its bound clients atomically + store.setState((s) => relayState.removeExtension(s, { extensionId: connectionId })) }, onError(event) { diff --git a/playwriter/src/relay-state.test.ts b/playwriter/src/relay-state.test.ts index f5b68e7..41c5e73 100644 --- a/playwriter/src/relay-state.test.ts +++ b/playwriter/src/relay-state.test.ts @@ -130,6 +130,19 @@ describe('removeExtension', () => { expect(after).toBe(before) // Same reference — no allocation }) + + test('also removes playwright clients bound to the extension', () => { + let state = stateWithExtension('ext-1') + state = relayState.addPlaywrightClient(state, { id: 'c1', extensionId: 'ext-1', ws: fakeWs() }) + state = relayState.addPlaywrightClient(state, { id: 'c2', extensionId: 'ext-1', ws: fakeWs() }) + state = relayState.addPlaywrightClient(state, { id: 'c3', extensionId: 'ext-2', ws: fakeWs() }) + + const after = relayState.removeExtension(state, { extensionId: 'ext-1' }) + + expect(after.extensions.size).toBe(0) + expect(after.playwrightClients.size).toBe(1) + expect(after.playwrightClients.has('c3')).toBe(true) + }) }) // --------------------------------------------------------------------------- @@ -168,15 +181,6 @@ describe('removePlaywrightClient', () => { }) describe('extension I/O fields', () => { - test('increments extension message id without mutating old state', () => { - const state = stateWithExtension('ext-1') - expect(state.extensions.get('ext-1')?.messageId).toBe(0) - - const after = relayState.incrementExtensionMessageId(state, { extensionId: 'ext-1' }) - expect(after.extensions.get('ext-1')?.messageId).toBe(1) - expect(state.extensions.get('ext-1')?.messageId).toBe(0) - }) - test('adds and removes pending extension requests', () => { let state = stateWithExtension('ext-1') @@ -200,32 +204,6 @@ describe('extension I/O fields', () => { expect(state.extensions.get('ext-1')?.pendingRequests.has(7)).toBe(false) }) - test('clears all pending extension requests', () => { - let state = emptyState() - state = relayState.addExtension(state, { - id: 'ext-1', - info: { browser: 'Chrome' }, - stableKey: 'profile:chrome-1', - ws: fakeWs(), - }) - // Manually add a pending request - state = relayState.addExtensionPendingRequest(state, { - extensionId: 'ext-1', - requestId: 1, - pendingRequest: { - resolve: (_result: unknown) => { - return - }, - reject: (_error: Error) => { - return - }, - }, - }) - - state = relayState.clearExtensionPendingRequests(state, { extensionId: 'ext-1' }) - expect(state.extensions.get('ext-1')?.pendingRequests.size).toBe(0) - }) - test('updateExtensionIO updates ws and pingInterval', () => { let state = stateWithExtension('ext-1') expect(state.extensions.get('ext-1')?.pingInterval).toBeNull() @@ -516,30 +494,7 @@ describe('updateTargetUrl', () => { }) }) -// --------------------------------------------------------------------------- -// removeClientsForExtension -// --------------------------------------------------------------------------- -describe('removeClientsForExtension', () => { - test('removes all clients bound to extension', () => { - let state = stateWithExtension('ext-1') - state = relayState.addPlaywrightClient(state, { id: 'c1', extensionId: 'ext-1', ws: fakeWs() }) - state = relayState.addPlaywrightClient(state, { id: 'c2', extensionId: 'ext-1', ws: fakeWs() }) - state = relayState.addPlaywrightClient(state, { id: 'c3', extensionId: 'ext-2', ws: fakeWs() }) - - const after = relayState.removeClientsForExtension(state, { extensionId: 'ext-1' }) - - expect(after.playwrightClients.size).toBe(1) - expect(after.playwrightClients.has('c3')).toBe(true) - }) - - test('no-op if no clients bound', () => { - const before = stateWithExtension('ext-1') - const after = relayState.removeClientsForExtension(before, { extensionId: 'ext-1' }) - - expect(after).toBe(before) - }) -}) // --------------------------------------------------------------------------- // Derivation helpers diff --git a/playwriter/src/relay-state.ts b/playwriter/src/relay-state.ts index b9138b7..c11cb65 100644 --- a/playwriter/src/relay-state.ts +++ b/playwriter/src/relay-state.ts @@ -147,14 +147,26 @@ export function addExtension( return { ...state, extensions: newExtensions } } -/** Remove an extension and all its targets. */ +/** Remove an extension, its targets, and any playwright clients bound to it. */ export function removeExtension(state: RelayState, { extensionId }: { extensionId: string }): RelayState { if (!state.extensions.has(extensionId)) { return state } const newExtensions = new Map(state.extensions) newExtensions.delete(extensionId) - return { ...state, extensions: newExtensions } + + // Also remove playwright clients bound to this extension + const clientsToRemove = Array.from(state.playwrightClients.values()) + .filter((client) => client.extensionId === extensionId) + if (clientsToRemove.length === 0) { + return { ...state, extensions: newExtensions } + } + + const newClients = new Map(state.playwrightClients) + for (const client of clientsToRemove) { + newClients.delete(client.id) + } + return { ...state, extensions: newExtensions, playwrightClients: newClients } } /** Add a playwright client (state + ws handle co-located). */ @@ -229,18 +241,6 @@ export function updateExtensionIO( return { ...state, extensions: newExtensions } } -/** Increment an extension's message id counter and return updated state. */ -export function incrementExtensionMessageId(state: RelayState, { extensionId }: { extensionId: string }): RelayState { - const ext = state.extensions.get(extensionId) - if (!ext) { - return state - } - - const newExtensions = new Map(state.extensions) - newExtensions.set(extensionId, { ...ext, messageId: ext.messageId + 1 }) - return { ...state, extensions: newExtensions } -} - /** Add or replace one pending extension request callback pair. */ export function addExtensionPendingRequest( state: RelayState, @@ -283,21 +283,6 @@ export function removeExtensionPendingRequest( return { ...state, extensions: newExtensions } } -/** Remove all pending extension request callback pairs. */ -export function clearExtensionPendingRequests( - state: RelayState, - { extensionId }: { extensionId: string }, -): RelayState { - const ext = state.extensions.get(extensionId) - if (!ext || ext.pendingRequests.size === 0) { - return state - } - - const newExtensions = new Map(state.extensions) - newExtensions.set(extensionId, { ...ext, pendingRequests: new Map() }) - return { ...state, extensions: newExtensions } -} - /** Add a target to an extension's connectedTargets. No-op if extension doesn't exist. */ export function addTarget( state: RelayState, @@ -509,25 +494,4 @@ export function updateTargetUrl( return { ...state, extensions: newExtensions } } -/** Remove all playwright clients bound to a specific extension. */ -export function removeClientsForExtension( - state: RelayState, - { extensionId }: { extensionId: string }, -): RelayState { - const toRemove: string[] = [] - for (const [clientId, client] of state.playwrightClients) { - if (client.extensionId === extensionId) { - toRemove.push(clientId) - } - } - if (toRemove.length === 0) { - return state - } - - const newClients = new Map(state.playwrightClients) - for (const clientId of toRemove) { - newClients.delete(clientId) - } - return { ...state, playwrightClients: newClients } -}