refactor: inline/remove single-use relay-state functions

- Inline incrementExtensionMessageId: trivially messageId + 1, no value
  as a separate function with one call site and one test
- Remove clearExtensionPendingRequests: redundant before removeExtension
  which deletes the entire extension entry (including pendingRequests)
- Fold removeClientsForExtension into removeExtension: always called
  together, callers can no longer forget to clean up bound clients

Net: -3 exported functions, -3 tests replaced by 1 new test verifying
removeExtension also removes bound playwright clients.

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
+          // State transition: remove extension + its bound clients atomically
-          store.setState((s) => {
+          store.setState((s) => relayState.removeExtension(s, { extensionId: connectionId }))
-            let next = relayState.removeExtension(s, { extensionId: connectionId })
-            next = relayState.removeClientsForExtension(next, { extensionId: connectionId })
-            return next
-          })
         },
 
         onError(event) {

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

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 }
-}