refactor: centralize relay state into zustand store, deduplicate playwright client maps

Introduce relay-state.ts with a single Zustand vanilla store holding all relay
server state (extensions, playwright clients, targets). All state transitions are
pure functions: data in, data out, no I/O inside setState().

Key changes:
- Move scattered mutable Maps (extensionConnections, extensionKeyIndex,
  playwrightClients) into one RelayState atom managed by zustand/vanilla
- Co-locate I/O handles (ws, pingInterval, pendingRequests) on ExtensionEntry
  instead of separate Maps, following the same pattern for PlaywrightClient
- Merge PlaywrightClientState + PlaywrightClientIO into single PlaywrightClient
  type with { id, extensionId, ws } — eliminates the duplicate playwrightClientIO
  map that had to stay in sync (and was already buggy: removeClientsForExtension
  only cleaned playwrightClients but leaked playwrightClientIO entries)
- Derive extensionKeyIndex on demand via findExtensionByStableKey() linear scan
  instead of maintaining a separate index Map
- Add relay-state.test.ts with 38 pure data-in/data-out tests covering all
  state transitions and derivation helpers

playwriter/package.json

@@ -58,7 +58,8 @@
     "string-dedent": "^3.0.2",
     "strip-ansi": "^7.1.2",
     "ws": "^8.18.3",
-    "zod": "^4.3.5"
+    "zod": "^4.3.5",
+    "zustand": "^5.0.8"
   },
   "optionalDependencies": {
     "sharp": "^0.34.5"

playwriter/src/relay-state.test.ts

@@ -0,0 +1,615 @@
+/**
+ * Unit tests for relay state transitions.
+ * Data-in / data-out transitions for the unified extension map.
+ */
+import { describe, test, expect } from 'vitest'
+import type { WSContext } from 'hono/ws'
+import type { Protocol } from './cdp-types.js'
+import * as relayState from './relay-state.js'
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function emptyState(): relayState.RelayState {
+  return {
+    extensions: new Map(),
+    playwrightClients: new Map(),
+  }
+}
+
+function fakeWs(): WSContext {
+  return {} as WSContext
+}
+
+function makeTargetInfo(overrides: Partial<Protocol.Target.TargetInfo> = {}): Protocol.Target.TargetInfo {
+  return {
+    targetId: 'target-1',
+    type: 'page',
+    title: 'Test Page',
+    url: 'https://example.com',
+    attached: true,
+    canAccessOpener: false,
+    ...overrides,
+  }
+}
+
+function stateWithExtension(
+  extensionId = 'ext-1',
+  info: relayState.ExtensionInfo = { browser: 'Chrome' },
+  stableKey = 'profile:chrome-1',
+): relayState.RelayState {
+  return relayState.addExtension(emptyState(), { id: extensionId, info, stableKey, ws: fakeWs() })
+}
+
+// ---------------------------------------------------------------------------
+// createRelayStore
+// ---------------------------------------------------------------------------
+
+describe('createRelayStore', () => {
+  test('creates store with empty maps', () => {
+    const store = relayState.createRelayStore()
+    const state = store.getState()
+    expect(state.extensions.size).toBe(0)
+    expect(state.playwrightClients.size).toBe(0)
+  })
+})
+
+// ---------------------------------------------------------------------------
+// addExtension / removeExtension
+// ---------------------------------------------------------------------------
+
+describe('addExtension', () => {
+  test('adds extension to empty state', () => {
+    const before = emptyState()
+    const after = relayState.addExtension(before, {
+      id: 'ext-1',
+      info: { browser: 'Chrome' },
+      stableKey: 'profile:chrome-1',
+      ws: fakeWs(),
+    })
+
+    expect(after.extensions.size).toBe(1)
+    const ext = after.extensions.get('ext-1')!
+    expect(ext.stableKey).toBe('profile:chrome-1')
+    expect(ext.connectedTargets.size).toBe(0)
+    expect(ext.ws).toBeTruthy()
+    expect(ext.messageId).toBe(0)
+    expect(ext.pendingRequests.size).toBe(0)
+    expect(ext.pingInterval).toBeNull()
+    // Original unchanged (immutable)
+    expect(before.extensions.size).toBe(0)
+  })
+
+  test('adding extension with same stableKey keeps old entry (removed on socket close)', () => {
+    const s1 = relayState.addExtension(emptyState(), {
+      id: 'ext-old',
+      info: { browser: 'Chrome' },
+      stableKey: 'profile:chrome-1',
+      ws: fakeWs(),
+    })
+    const s2 = relayState.addExtension(s1, {
+      id: 'ext-new',
+      info: { browser: 'Chrome', email: 'test@example.com' },
+      stableKey: 'profile:chrome-1',
+      ws: fakeWs(),
+    })
+
+    // Both coexist — old stays routable until its socket closes
+    expect(s2.extensions.size).toBe(2)
+    expect(s2.extensions.has('ext-old')).toBe(true)
+    expect(s2.extensions.has('ext-new')).toBe(true)
+    // findExtensionByStableKey returns newest
+    expect(relayState.findExtensionByStableKey(s2, 'profile:chrome-1')?.id).toBe('ext-new')
+    // Original unchanged
+    expect(s1.extensions.size).toBe(1)
+  })
+
+  test('allows multiple extensions with different stableKeys', () => {
+    let state = emptyState()
+    state = relayState.addExtension(state, { id: 'ext-1', info: { browser: 'Chrome' }, stableKey: 'profile:a', ws: fakeWs() })
+    state = relayState.addExtension(state, { id: 'ext-2', info: { browser: 'Firefox' }, stableKey: 'profile:b', ws: fakeWs() })
+
+    expect(state.extensions.size).toBe(2)
+  })
+})
+
+describe('removeExtension', () => {
+  test('removes existing extension', () => {
+    const before = stateWithExtension('ext-1')
+    const after = relayState.removeExtension(before, { extensionId: 'ext-1' })
+
+    expect(after.extensions.size).toBe(0)
+    // Original unchanged
+    expect(before.extensions.size).toBe(1)
+  })
+
+  test('no-op for nonexistent extension', () => {
+    const before = stateWithExtension('ext-1')
+    const after = relayState.removeExtension(before, { extensionId: 'ext-999' })
+
+    expect(after).toBe(before) // Same reference — no allocation
+  })
+})
+
+// ---------------------------------------------------------------------------
+// addPlaywrightClient / removePlaywrightClient
+// ---------------------------------------------------------------------------
+
+describe('addPlaywrightClient', () => {
+  test('adds client with ws handle', () => {
+    const before = emptyState()
+    const ws = fakeWs()
+    const after = relayState.addPlaywrightClient(before, { id: 'client-1', extensionId: 'ext-1', ws })
+
+    expect(after.playwrightClients.size).toBe(1)
+    const client = after.playwrightClients.get('client-1')!
+    expect(client.extensionId).toBe('ext-1')
+    expect(client.ws).toBe(ws)
+    expect(before.playwrightClients.size).toBe(0)
+  })
+})
+
+describe('removePlaywrightClient', () => {
+  test('removes client', () => {
+    const before = relayState.addPlaywrightClient(emptyState(), { id: 'c1', extensionId: null, ws: fakeWs() })
+    const after = relayState.removePlaywrightClient(before, { clientId: 'c1' })
+
+    expect(after.playwrightClients.size).toBe(0)
+    expect(before.playwrightClients.size).toBe(1)
+  })
+
+  test('no-op for nonexistent client', () => {
+    const before = emptyState()
+    const after = relayState.removePlaywrightClient(before, { clientId: 'nope' })
+
+    expect(after).toBe(before)
+  })
+})
+
+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')
+
+    const pending = {
+      resolve: (_result: unknown) => {
+        return
+      },
+      reject: (_error: Error) => {
+        return
+      },
+    }
+
+    state = relayState.addExtensionPendingRequest(state, {
+      extensionId: 'ext-1',
+      requestId: 7,
+      pendingRequest: pending,
+    })
+    expect(state.extensions.get('ext-1')?.pendingRequests.get(7)).toBe(pending)
+
+    state = relayState.removeExtensionPendingRequest(state, { extensionId: 'ext-1', requestId: 7 })
+    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()
+
+    const interval = setInterval(() => {}, 99999)
+    try {
+      state = relayState.updateExtensionIO(state, { extensionId: 'ext-1', pingInterval: interval })
+      expect(state.extensions.get('ext-1')?.pingInterval).toBe(interval)
+
+      state = relayState.updateExtensionIO(state, { extensionId: 'ext-1', pingInterval: null })
+      expect(state.extensions.get('ext-1')?.pingInterval).toBeNull()
+    } finally {
+      clearInterval(interval)
+    }
+  })
+
+  test('playwright client ws handle is co-located with state', () => {
+    let state = emptyState()
+    const ws = fakeWs()
+    state = relayState.addPlaywrightClient(state, { id: 'c1', extensionId: null, ws })
+    expect(state.playwrightClients.get('c1')?.ws).toBe(ws)
+
+    state = relayState.removePlaywrightClient(state, { clientId: 'c1' })
+    expect(state.playwrightClients.size).toBe(0)
+  })
+})
+
+describe('rebindClientsToExtension', () => {
+  test('rebinds all clients from old extension to successor extension', () => {
+    let state = stateWithExtension('ext-old')
+    state = relayState.addExtension(state, {
+      id: 'ext-new',
+      info: { browser: 'Chrome' },
+      stableKey: 'profile:chrome-1',
+      ws: fakeWs(),
+    })
+    state = relayState.addPlaywrightClient(state, { id: 'c1', extensionId: 'ext-old', ws: fakeWs() })
+    state = relayState.addPlaywrightClient(state, { id: 'c2', extensionId: 'ext-old', ws: fakeWs() })
+    state = relayState.addPlaywrightClient(state, { id: 'c3', extensionId: 'ext-new', ws: fakeWs() })
+
+    const after = relayState.rebindClientsToExtension(state, {
+      fromExtensionId: 'ext-old',
+      toExtensionId: 'ext-new',
+    })
+
+    expect(after.playwrightClients.get('c1')?.extensionId).toBe('ext-new')
+    expect(after.playwrightClients.get('c2')?.extensionId).toBe('ext-new')
+    expect(after.playwrightClients.get('c3')?.extensionId).toBe('ext-new')
+  })
+})
+
+// ---------------------------------------------------------------------------
+// addTarget / removeTarget / removeTargetByCrash
+// ---------------------------------------------------------------------------
+
+describe('addTarget', () => {
+  test('adds target to extension', () => {
+    const before = stateWithExtension('ext-1')
+    const targetInfo = makeTargetInfo()
+    const after = relayState.addTarget(before, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo,
+    })
+
+    const ext = after.extensions.get('ext-1')!
+    expect(ext.connectedTargets.size).toBe(1)
+    expect(ext.connectedTargets.get('pw-tab-1')?.targetId).toBe('target-1')
+    // Original unchanged
+    expect(before.extensions.get('ext-1')!.connectedTargets.size).toBe(0)
+  })
+
+  test('no-op if extension does not exist', () => {
+    const before = emptyState()
+    const after = relayState.addTarget(before, {
+      extensionId: 'ext-nope',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+
+    expect(after).toBe(before)
+  })
+
+  test('preserves existing frameIds on update', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+      existingFrameIds: new Set(['frame-A']),
+    })
+
+    // Update the same target with new targetInfo but no explicit frameIds
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo({ url: 'https://updated.com' }),
+    })
+
+    const target = state.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!
+    expect(target.targetInfo.url).toBe('https://updated.com')
+    expect(target.frameIds.has('frame-A')).toBe(true)
+  })
+})
+
+describe('removeTarget', () => {
+  test('removes target by sessionId', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+    const after = relayState.removeTarget(state, { extensionId: 'ext-1', sessionId: 'pw-tab-1' })
+
+    expect(after.extensions.get('ext-1')!.connectedTargets.size).toBe(0)
+  })
+
+  test('no-op if target does not exist', () => {
+    const before = stateWithExtension('ext-1')
+    const after = relayState.removeTarget(before, { extensionId: 'ext-1', sessionId: 'pw-tab-nope' })
+
+    expect(after).toBe(before)
+  })
+})
+
+describe('removeTargetByCrash', () => {
+  test('removes target by targetId', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+    const after = relayState.removeTargetByCrash(state, { extensionId: 'ext-1', targetId: 'target-1' })
+
+    expect(after.extensions.get('ext-1')!.connectedTargets.size).toBe(0)
+  })
+
+  test('no-op if targetId not found', () => {
+    const before = stateWithExtension('ext-1')
+    const after = relayState.removeTargetByCrash(before, { extensionId: 'ext-1', targetId: 'nope' })
+
+    expect(after).toBe(before)
+  })
+})
+
+// ---------------------------------------------------------------------------
+// updateTargetInfo
+// ---------------------------------------------------------------------------
+
+describe('updateTargetInfo', () => {
+  test('updates targetInfo matched by targetId', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo({ title: 'Old Title' }),
+    })
+
+    const newInfo = makeTargetInfo({ title: 'New Title' })
+    const after = relayState.updateTargetInfo(state, { extensionId: 'ext-1', targetInfo: newInfo })
+
+    expect(after.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!.targetInfo.title).toBe('New Title')
+    // Original unchanged
+    expect(state.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!.targetInfo.title).toBe('Old Title')
+  })
+})
+
+// ---------------------------------------------------------------------------
+// addFrameId / removeFrameId
+// ---------------------------------------------------------------------------
+
+describe('addFrameId', () => {
+  test('adds frameId to target', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+    const after = relayState.addFrameId(state, { extensionId: 'ext-1', sessionId: 'pw-tab-1', frameId: 'frame-1' })
+
+    expect(after.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!.frameIds.has('frame-1')).toBe(true)
+  })
+
+  test('no-op if frameId already present', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+    state = relayState.addFrameId(state, { extensionId: 'ext-1', sessionId: 'pw-tab-1', frameId: 'frame-1' })
+    const after = relayState.addFrameId(state, { extensionId: 'ext-1', sessionId: 'pw-tab-1', frameId: 'frame-1' })
+
+    expect(after).toBe(state) // Same reference
+  })
+})
+
+describe('removeFrameId', () => {
+  test('removes frameId from owning target', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+    state = relayState.addFrameId(state, { extensionId: 'ext-1', sessionId: 'pw-tab-1', frameId: 'frame-1' })
+    const after = relayState.removeFrameId(state, { extensionId: 'ext-1', frameId: 'frame-1' })
+
+    expect(after.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!.frameIds.has('frame-1')).toBe(false)
+  })
+
+  test('no-op if frameId not found', () => {
+    const before = stateWithExtension('ext-1')
+    const after = relayState.removeFrameId(before, { extensionId: 'ext-1', frameId: 'nope' })
+
+    expect(after).toBe(before)
+  })
+})
+
+// ---------------------------------------------------------------------------
+// updateTargetUrl
+// ---------------------------------------------------------------------------
+
+describe('updateTargetUrl', () => {
+  test('frameNavigated on top-level frame updates URL and title', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo({ url: 'https://old.com', title: 'Old' }),
+    })
+
+    const after = relayState.updateTargetUrl(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      url: 'https://new.com',
+      title: 'New Page',
+    })
+
+    const target = after.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!
+    expect(target.targetInfo.url).toBe('https://new.com')
+    expect(target.targetInfo.title).toBe('New Page')
+  })
+
+  test('navigatedWithinDocument updates URL only (no title)', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo({ url: 'https://example.com', title: 'Keep This' }),
+    })
+
+    const after = relayState.updateTargetUrl(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      url: 'https://example.com/new-path',
+    })
+
+    const target = after.extensions.get('ext-1')!.connectedTargets.get('pw-tab-1')!
+    expect(target.targetInfo.url).toBe('https://example.com/new-path')
+    expect(target.targetInfo.title).toBe('Keep This')
+  })
+
+  test('no-op if target does not exist', () => {
+    const before = stateWithExtension('ext-1')
+    const after = relayState.updateTargetUrl(before, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-nope',
+      url: 'https://example.com',
+    })
+
+    expect(after).toBe(before)
+  })
+})
+
+// ---------------------------------------------------------------------------
+// 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
+// ---------------------------------------------------------------------------
+
+describe('findExtensionByStableKey', () => {
+  test('finds extension by stableKey', () => {
+    const state = stateWithExtension('ext-1', { browser: 'Chrome' }, 'profile:chrome-1')
+
+    const found = relayState.findExtensionByStableKey(state, 'profile:chrome-1')
+    expect(found?.id).toBe('ext-1')
+  })
+
+  test('returns undefined if not found', () => {
+    const state = emptyState()
+    expect(relayState.findExtensionByStableKey(state, 'profile:nope')).toBeUndefined()
+  })
+})
+
+describe('findExtensionIdByCdpSession', () => {
+  test('finds extension owning a CDP sessionId', () => {
+    let state = stateWithExtension('ext-1')
+    state = relayState.addTarget(state, {
+      extensionId: 'ext-1',
+      sessionId: 'pw-tab-1',
+      targetId: 'target-1',
+      targetInfo: makeTargetInfo(),
+    })
+
+    expect(relayState.findExtensionIdByCdpSession(state, 'pw-tab-1')).toBe('ext-1')
+  })
+
+  test('returns null if session not found', () => {
+    const state = stateWithExtension('ext-1')
+    expect(relayState.findExtensionIdByCdpSession(state, 'pw-tab-nope')).toBeNull()
+  })
+})
+
+// ---------------------------------------------------------------------------
+// Zustand store integration
+// ---------------------------------------------------------------------------
+
+describe('store.setState with transitions', () => {
+  test('setState updates store atomically', () => {
+    const store = relayState.createRelayStore()
+
+    store.setState((s) => {
+      return relayState.addExtension(s, { id: 'ext-1', info: { browser: 'Chrome' }, stableKey: 'profile:1', ws: fakeWs() })
+    })
+
+    expect(store.getState().extensions.size).toBe(1)
+  })
+
+  test('chained transitions compose correctly', () => {
+    const store = relayState.createRelayStore()
+
+    store.setState((s) => {
+      let next = relayState.addExtension(s, { id: 'ext-1', info: {}, stableKey: 'k1', ws: fakeWs() })
+      next = relayState.addTarget(next, {
+        extensionId: 'ext-1',
+        sessionId: 'pw-tab-1',
+        targetId: 'target-1',
+        targetInfo: makeTargetInfo(),
+      })
+      next = relayState.addPlaywrightClient(next, { id: 'c1', extensionId: 'ext-1', ws: fakeWs() })
+      return next
+    })
+
+    const state = store.getState()
+    expect(state.extensions.get('ext-1')!.connectedTargets.size).toBe(1)
+    expect(state.playwrightClients.size).toBe(1)
+  })
+})

playwriter/src/relay-state.ts

@@ -0,0 +1,533 @@
+/**
+ * Centralized relay state: one immutable atom for domain state + runtime resources.
+ *
+ * Follows the zustand-centralized-state skill pattern:
+ * - Single Zustand vanilla store holds all relay state
+ * - setState() callbacks remain deterministic data transitions
+ * - Runtime resources (WebSocket/timers/pending callbacks) are co-located on
+ *   each extension entry in the same map
+ *
+ * See docs/plan-centralize-relay-state.md for the full refactor plan.
+ */
+import { createStore, type StoreApi } from 'zustand/vanilla'
+import type { WSContext } from 'hono/ws'
+import type { Protocol } from './cdp-types.js'
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+export type ConnectedTarget = {
+  sessionId: string
+  targetId: string
+  targetInfo: Protocol.Target.TargetInfo
+  frameIds: Set<string>
+}
+
+export type ExtensionInfo = {
+  browser?: string
+  email?: string
+  id?: string
+  /** playwriter package version the extension was built with (sent as ?v= query param) */
+  version?: string
+}
+
+export type ExtensionPendingRequest = {
+  resolve: (result: unknown) => void
+  reject: (error: Error) => void
+}
+
+/**
+ * Single aggregated extension object: domain state + runtime I/O.
+ */
+export type ExtensionEntry = {
+  id: string
+  info: ExtensionInfo
+  stableKey: string
+  connectedTargets: Map<string, ConnectedTarget>
+  // Runtime I/O fields
+  ws: WSContext | null
+  pendingRequests: Map<number, ExtensionPendingRequest>
+  messageId: number
+  pingInterval: ReturnType<typeof setInterval> | null
+}
+
+export type PlaywrightClient = {
+  id: string
+  extensionId: string | null
+  ws: WSContext
+}
+
+export type RelayState = {
+  extensions: Map<string, ExtensionEntry>
+  playwrightClients: Map<string, PlaywrightClient>
+}
+
+// ---------------------------------------------------------------------------
+// Store factory
+// ---------------------------------------------------------------------------
+
+export function createRelayStore(): StoreApi<RelayState> {
+  return createStore<RelayState>(() => ({
+    extensions: new Map(),
+    playwrightClients: new Map(),
+  }))
+}
+
+// ---------------------------------------------------------------------------
+// Derivation helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Linear scan over extensions to find one by stableKey. With <10 extensions this is free.
+ * Returns the LAST (newest) match because during reconnect both old and new connections
+ * coexist briefly. Map iteration order is insertion order, so last wins.
+ */
+export function findExtensionByStableKey(state: RelayState, stableKey: string): ExtensionEntry | undefined {
+  let match: ExtensionEntry | undefined
+  for (const ext of state.extensions.values()) {
+    if (ext.stableKey === stableKey) {
+      match = ext
+    }
+  }
+  return match
+}
+
+/** Find which extension owns a CDP tab sessionId (e.g. "pw-tab-1"). */
+export function findExtensionIdByCdpSession(state: RelayState, cdpSessionId: string): string | null {
+  for (const [connectionId, ext] of state.extensions.entries()) {
+    if (ext.connectedTargets.has(cdpSessionId)) {
+      return connectionId
+    }
+  }
+  return null
+}
+
+// ---------------------------------------------------------------------------
+// Pure state transition functions
+//
+// Each takes RelayState + event data and returns a new RelayState.
+// No I/O, no side effects. Testable with data in / data out.
+// ---------------------------------------------------------------------------
+
+/**
+ * Add a new extension connection.
+ * Does NOT remove an existing connection with the same stableKey — that old
+ * connection stays routable until its WebSocket onClose fires and calls
+ * removeExtension(). This preserves in-flight message routing during reconnect.
+ * findExtensionByStableKey() returns the newest match so stableKey lookups
+ * resolve to the new connection immediately.
+ */
+export function addExtension(
+  state: RelayState,
+  {
+    id,
+    info,
+    stableKey,
+    ws,
+  }: {
+    id: string
+    info: ExtensionInfo
+    stableKey: string
+    ws: WSContext | null
+  },
+): RelayState {
+  const newExtensions = new Map(state.extensions)
+
+  newExtensions.set(id, {
+    id,
+    info,
+    stableKey,
+    connectedTargets: new Map(),
+    ws,
+    pendingRequests: new Map(),
+    messageId: 0,
+    pingInterval: null,
+  })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Remove an extension and all its targets. */
+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 }
+}
+
+/** Add a playwright client (state + ws handle co-located). */
+export function addPlaywrightClient(
+  state: RelayState,
+  { id, extensionId, ws }: { id: string; extensionId: string | null; ws: WSContext },
+): RelayState {
+  const newClients = new Map(state.playwrightClients)
+  newClients.set(id, { id, extensionId, ws })
+  return { ...state, playwrightClients: newClients }
+}
+
+/** Remove a playwright client. */
+export function removePlaywrightClient(state: RelayState, { clientId }: { clientId: string }): RelayState {
+  if (!state.playwrightClients.has(clientId)) {
+    return state
+  }
+  const newClients = new Map(state.playwrightClients)
+  newClients.delete(clientId)
+  return { ...state, playwrightClients: newClients }
+}
+
+/** Rebind all clients from one extension id to another. */
+export function rebindClientsToExtension(
+  state: RelayState,
+  { fromExtensionId, toExtensionId }: { fromExtensionId: string; toExtensionId: string },
+): RelayState {
+  if (fromExtensionId === toExtensionId) {
+    return state
+  }
+
+  let updated = false
+  const newClients = new Map(state.playwrightClients)
+  for (const [clientId, client] of newClients) {
+    if (client.extensionId !== fromExtensionId) {
+      continue
+    }
+    newClients.set(clientId, { ...client, extensionId: toExtensionId })
+    updated = true
+  }
+
+  if (!updated) {
+    return state
+  }
+
+  return { ...state, playwrightClients: newClients }
+}
+
+/** Update an extension entry's I/O fields (ws, pingInterval). */
+export function updateExtensionIO(
+  state: RelayState,
+  {
+    extensionId,
+    ws,
+    pingInterval,
+  }: {
+    extensionId: string
+    ws?: WSContext | null
+    pingInterval?: ReturnType<typeof setInterval> | null
+  },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, {
+    ...ext,
+    ...(ws !== undefined ? { ws } : {}),
+    ...(pingInterval !== undefined ? { pingInterval } : {}),
+  })
+  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,
+  {
+    extensionId,
+    requestId,
+    pendingRequest,
+  }: {
+    extensionId: string
+    requestId: number
+    pendingRequest: ExtensionPendingRequest
+  },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+
+  const pendingRequests = new Map(ext.pendingRequests)
+  pendingRequests.set(requestId, pendingRequest)
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, pendingRequests })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Remove one pending extension request callback pair. */
+export function removeExtensionPendingRequest(
+  state: RelayState,
+  { extensionId, requestId }: { extensionId: string; requestId: number },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext || !ext.pendingRequests.has(requestId)) {
+    return state
+  }
+
+  const pendingRequests = new Map(ext.pendingRequests)
+  pendingRequests.delete(requestId)
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, pendingRequests })
+  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,
+  {
+    extensionId,
+    sessionId,
+    targetId,
+    targetInfo,
+    existingFrameIds,
+  }: {
+    extensionId: string
+    sessionId: string
+    targetId: string
+    targetInfo: Protocol.Target.TargetInfo
+    /** Preserve existing frameIds if target already existed (update scenario). */
+    existingFrameIds?: Set<string>
+  },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+
+  const existingTarget = ext.connectedTargets.get(sessionId)
+  const newTargets = new Map(ext.connectedTargets)
+  newTargets.set(sessionId, {
+    sessionId,
+    targetId,
+    targetInfo,
+    frameIds: existingFrameIds ?? existingTarget?.frameIds ?? new Set(),
+  })
+
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Remove a target by sessionId. No-op if extension or target doesn't exist. */
+export function removeTarget(
+  state: RelayState,
+  { extensionId, sessionId }: { extensionId: string; sessionId: string },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext || !ext.connectedTargets.has(sessionId)) {
+    return state
+  }
+
+  const newTargets = new Map(ext.connectedTargets)
+  newTargets.delete(sessionId)
+
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Remove a crashed target by targetId (not sessionId). */
+export function removeTargetByCrash(
+  state: RelayState,
+  { extensionId, targetId }: { extensionId: string; targetId: string },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+
+  let found = false
+  const newTargets = new Map(ext.connectedTargets)
+  for (const [sid, target] of newTargets) {
+    if (target.targetId === targetId) {
+      newTargets.delete(sid)
+      found = true
+      break
+    }
+  }
+
+  if (!found) {
+    return state
+  }
+
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Update targetInfo on a target matched by targetId. */
+export function updateTargetInfo(
+  state: RelayState,
+  { extensionId, targetInfo }: { extensionId: string; targetInfo: Protocol.Target.TargetInfo },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+
+  let updated = false
+  const newTargets = new Map(ext.connectedTargets)
+  for (const [sid, target] of newTargets) {
+    if (target.targetId === targetInfo.targetId) {
+      newTargets.set(sid, { ...target, targetInfo })
+      updated = true
+      break
+    }
+  }
+
+  if (!updated) {
+    return state
+  }
+
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Add a frameId to a target's frameIds set. */
+export function addFrameId(
+  state: RelayState,
+  { extensionId, sessionId, frameId }: { extensionId: string; sessionId: string; frameId: string },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+  const target = ext.connectedTargets.get(sessionId)
+  if (!target) {
+    return state
+  }
+
+  // Already present — no-op
+  if (target.frameIds.has(frameId)) {
+    return state
+  }
+
+  const newFrameIds = new Set(target.frameIds)
+  newFrameIds.add(frameId)
+
+  const newTargets = new Map(ext.connectedTargets)
+  newTargets.set(sessionId, { ...target, frameIds: newFrameIds })
+
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+  return { ...state, extensions: newExtensions }
+}
+
+/** Remove a frameId from the target that owns it (scans all targets in the extension). */
+export function removeFrameId(
+  state: RelayState,
+  { extensionId, frameId }: { extensionId: string; frameId: string },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+
+  for (const [sid, target] of ext.connectedTargets) {
+    if (target.frameIds.has(frameId)) {
+      const newFrameIds = new Set(target.frameIds)
+      newFrameIds.delete(frameId)
+
+      const newTargets = new Map(ext.connectedTargets)
+      newTargets.set(sid, { ...target, frameIds: newFrameIds })
+
+      const newExtensions = new Map(state.extensions)
+      newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+      return { ...state, extensions: newExtensions }
+    }
+  }
+
+  return state
+}
+
+/**
+ * Update URL (and optionally title) on a target.
+ * Used by Page.frameNavigated (top-level) and Page.navigatedWithinDocument.
+ */
+export function updateTargetUrl(
+  state: RelayState,
+  {
+    extensionId,
+    sessionId,
+    url,
+    title,
+  }: {
+    extensionId: string
+    sessionId: string
+    url: string
+    title?: string
+  },
+): RelayState {
+  const ext = state.extensions.get(extensionId)
+  if (!ext) {
+    return state
+  }
+  const target = ext.connectedTargets.get(sessionId)
+  if (!target) {
+    return state
+  }
+
+  const newTargetInfo = {
+    ...target.targetInfo,
+    url,
+    ...(title !== undefined ? { title } : {}),
+  }
+
+  const newTargets = new Map(ext.connectedTargets)
+  newTargets.set(sessionId, { ...target, targetInfo: newTargetInfo })
+
+  const newExtensions = new Map(state.extensions)
+  newExtensions.set(extensionId, { ...ext, connectedTargets: newTargets })
+  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 }
+}

pnpm-lock.yaml

@@ -362,6 +362,9 @@ importers:
       zod:
         specifier: ^4.3.5
         version: 4.3.5
+      zustand:
+        specifier: ^5.0.8
+        version: 5.0.8(@types/react@19.2.7)(react@19.2.3)
     devDependencies:
       '@browserbasehq/stagehand':
         specifier: ^3.0.5
@@ -9493,6 +9496,14 @@ snapshots:
     optionalDependencies:
       vite: 7.3.0(@types/node@25.0.3)(jiti@2.6.1)(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0)
 
+  '@vitest/mocker@4.0.8(vite@7.2.2(@types/node@24.10.1)(jiti@2.6.1)(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0))':
+    dependencies:
+      '@vitest/spy': 4.0.8
+      estree-walker: 3.0.3
+      magic-string: 0.30.21
+    optionalDependencies:
+      vite: 7.2.2(@types/node@24.10.1)(jiti@2.6.1)(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0)
+
   '@vitest/mocker@4.0.8(vite@7.2.2(@types/node@25.2.0)(jiti@2.6.1)(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0))':
     dependencies:
       '@vitest/spy': 4.0.8
@@ -9544,7 +9555,7 @@ snapshots:
       sirv: 3.0.2
       tinyglobby: 0.2.15
       tinyrainbow: 3.0.3
-      vitest: 4.0.8(@types/node@25.2.0)(@vitest/ui@4.0.8)(jiti@2.6.1)(jsdom@27.2.0(bufferutil@4.0.9))(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0)
+      vitest: 4.0.8(@types/node@24.10.1)(@vitest/ui@4.0.8)(jiti@2.6.1)(jsdom@27.2.0(bufferutil@4.0.9))(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0)
 
   '@vitest/utils@4.0.16':
     dependencies:
@@ -13580,7 +13591,7 @@ snapshots:
   vitest@4.0.8(@types/node@24.10.1)(@vitest/ui@4.0.8)(jiti@2.6.1)(jsdom@27.2.0(bufferutil@4.0.9))(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0):
     dependencies:
       '@vitest/expect': 4.0.8
-      '@vitest/mocker': 4.0.8(vite@7.2.2(@types/node@25.2.0)(jiti@2.6.1)(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0))
+      '@vitest/mocker': 4.0.8(vite@7.2.2(@types/node@24.10.1)(jiti@2.6.1)(lightningcss@1.31.1)(tsx@4.20.6)(yaml@2.6.0))
       '@vitest/pretty-format': 4.0.8
       '@vitest/runner': 4.0.8
       '@vitest/snapshot': 4.0.8