playwriter 0.0.63 → 0.0.89

package/src/cdp-relay.ts

@@ -28,13 +28,7 @@ import { EventEmitter } from 'node:events'
 import { VERSION, EXTENSION_IDS } from './utils.js'
 import { createCdpLogger, type CdpLogEntry, type CdpLogger } from './cdp-log.js'
 import { RecordingRelay } from './recording-relay.js'
-
-type ConnectedTarget = {
-  sessionId: string
-  targetId: string
-  targetInfo: Protocol.Target.TargetInfo
-  frameIds: Set<string>
-}
+import * as relayState from './relay-state.js'
 
 /**
  * Checks if a target should be filtered out (not exposed to Playwright).
@@ -68,32 +62,6 @@ function isRestrictedTarget(targetInfo: Protocol.Target.TargetInfo): boolean {
   return blockedPrefixes.some((prefix) => url.startsWith(prefix))
 }
 
-type PlaywrightClient = {
-  id: string
-  ws: WSContext
-  extensionId: string | null
-}
-
-type ExtensionInfo = {
-  browser?: string
-  email?: string
-  id?: string
-  /** playwriter package version the extension was built with (sent as ?v= query param) */
-  version?: string
-}
-
-type ExtensionConnection = {
-  id: string
-  ws: WSContext
-  info: ExtensionInfo
-  stableKey: string
-  connectedTargets: Map<string, ConnectedTarget>
-  pendingRequests: Map<number, { resolve: (result: any) => void; reject: (error: Error) => void }>
-  messageId: number
-  pingInterval: ReturnType<typeof setInterval> | null
-}
-
-
 export type RelayServer = {
   close(): void
   on<K extends keyof RelayServerEvents>(event: K, listener: RelayServerEvents[K]): void
@@ -114,31 +82,45 @@ export async function startPlayWriterCDPRelayServer({
   cdpLogger?: CdpLogger
 } = {}): Promise<RelayServer> {
   const emitter = new EventEmitter()
-  const extensionConnections = new Map<string, ExtensionConnection>()
-  const extensionKeyIndex = new Map<string, string>()
+  const store = relayState.createRelayStore()
+  const extensionDownloadBehavior = new Map<string, Protocol.Browser.SetDownloadBehaviorRequest>()
 
   const resolvedCdpLogger = cdpLogger || createCdpLogger()
   const logCdpJson = (entry: CdpLogEntry) => {
     resolvedCdpLogger.log(entry)
   }
-  const playwrightClients = new Map<string, PlaywrightClient>()
 
   const getDefaultExtensionId = (): string | null => {
-    return extensionConnections.keys().next().value || null
+    return store.getState().extensions.keys().next().value || null
   }
 
+  /**
+   * Resolve an extension by ID, stableKey, or fallback.
+   * Returns the unified ExtensionEntry which includes both state and I/O.
+   */
   const getExtensionConnection = (
     extensionId?: string | null,
-    options: { allowFallback?: boolean } = {}
-  ): ExtensionConnection | null => {
+    options: { allowFallback?: boolean } = {},
+  ): relayState.ExtensionEntry | null => {
+    const currentRelayState = store.getState()
+    const { extensions } = currentRelayState
+
     if (extensionId) {
-      const direct = extensionConnections.get(extensionId)
-      if (direct) {
+      const direct = extensions.get(extensionId)
+      if (direct?.ws) {
         return direct
       }
-      const mappedId = extensionKeyIndex.get(extensionId)
-      if (mappedId) {
-        return extensionConnections.get(mappedId) || null
+      // Try stableKey lookup.
+      const byKey = relayState.findExtensionByStableKey(currentRelayState, extensionId)
+      if (byKey) {
+        const candidates = Array.from(extensions.values())
+          .filter((ext) => ext.stableKey === byKey.stableKey)
+          .reverse()
+        for (const candidate of candidates) {
+          if (candidate.ws) {
+            return candidate
+          }
+        }
       }
       return null
     }
@@ -147,14 +129,33 @@ export async function startPlayWriterCDPRelayServer({
       return null
     }
 
-    const fallbackId = getDefaultExtensionId()
-    if (fallbackId) {
-      return extensionConnections.get(fallbackId) || null
+    // Single extension — use it directly
+    if (extensions.size === 1) {
+      const fallbackId = getDefaultExtensionId()
+      if (fallbackId) {
+        const ext = extensions.get(fallbackId)
+        if (ext?.ws) {
+          return ext
+        }
+      }
+    }
+
+    // Multiple extensions — auto-select if exactly one has active targets.
+    // This handles the common case of multiple Chrome profiles with the extension
+    // installed, where only one profile has playwriter-enabled tabs. (#52)
+    if (extensions.size > 1) {
+      const activeExtensions = Array.from(extensions.values()).filter((ext) => {
+        return ext.connectedTargets.size > 0
+      })
+      if (activeExtensions.length === 1 && activeExtensions[0].ws) {
+        return activeExtensions[0]
+      }
     }
+
     return null
   }
 
-  const buildStableExtensionKey = (info: ExtensionInfo, connectionId: string): string => {
+  const buildStableExtensionKey = (info: relayState.ExtensionInfo, connectionId: string): string => {
     if (info.id) {
       return `profile:${info.id}`
     }
@@ -176,37 +177,41 @@ export async function startPlayWriterCDPRelayServer({
   }
 
   const getPageTargetForFrameId = ({
-    connection,
-    frameId
+    extensionState,
+    frameId,
   }: {
-    connection: ExtensionConnection
+    extensionState: relayState.ExtensionEntry
     frameId: string
-  }): ConnectedTarget | undefined => {
-    return Array.from(connection.connectedTargets.values()).find((target) => {
+  }): relayState.ConnectedTarget | undefined => {
+    return Array.from(extensionState.connectedTargets.values()).find((target) => {
       return target.targetInfo.type === 'page' && target.frameIds.has(frameId)
     })
   }
 
   const startExtensionPing = (extensionId: string): void => {
-    const connection = extensionConnections.get(extensionId)
-    if (!connection) {
+    const ext = store.getState().extensions.get(extensionId)
+    if (!ext) {
       return
     }
-    if (connection.pingInterval) {
-      clearInterval(connection.pingInterval)
+    if (ext.pingInterval) {
+      clearInterval(ext.pingInterval)
     }
-    connection.pingInterval = setInterval(() => {
-      connection.ws.send(JSON.stringify({ method: 'ping' }))
+
+    const pingInterval = setInterval(() => {
+      const latestExt = store.getState().extensions.get(extensionId)
+      latestExt?.ws?.send(JSON.stringify({ method: 'ping' }))
     }, 5000)
+
+    store.setState((s) => relayState.updateExtensionIO(s, { extensionId, pingInterval }))
   }
 
   const stopExtensionPing = (extensionId: string): void => {
-    const connection = extensionConnections.get(extensionId)
-    if (!connection || !connection.pingInterval) {
+    const ext = store.getState().extensions.get(extensionId)
+    if (!ext || !ext.pingInterval) {
       return
     }
-    clearInterval(connection.pingInterval)
-    connection.pingInterval = null
+    clearInterval(ext.pingInterval)
+    store.setState((s) => relayState.updateExtensionIO(s, { extensionId, pingInterval: null }))
   }
 
   function logCdpMessage({
@@ -216,7 +221,7 @@ export async function startPlayWriterCDPRelayServer({
     sessionId,
     params,
     id,
-    source
+    source,
   }: {
     direction: 'to-playwright' | 'from-playwright' | 'from-extension'
     clientId?: string
@@ -232,7 +237,7 @@ export async function startPlayWriterCDPRelayServer({
       'Network.responseReceivedExtraInfo',
       'Network.dataReceived',
       'Network.requestWillBeSent',
-      'Network.loadingFinished'
+      'Network.loadingFinished',
     ]
 
     if (noisyEvents.includes(method)) {
@@ -287,9 +292,7 @@ export async function startPlayWriterCDPRelayServer({
     source?: 'extension' | 'server'
     extensionId?: string | null
   }) {
-    const messageToSend = source === 'server' && 'method' in message
-      ? { ...message, __serverGenerated: true }
-      : message
+    const messageToSend = source === 'server' && 'method' in message ? { ...message, __serverGenerated: true } : message
 
     logCdpJson({
       timestamp: new Date().toISOString(),
@@ -306,7 +309,7 @@ export async function startPlayWriterCDPRelayServer({
         method: message.method,
         sessionId: 'sessionId' in message ? message.sessionId : undefined,
         params: 'params' in message ? message.params : undefined,
-        source
+        source,
       })
     }
 
@@ -318,7 +321,7 @@ export async function startPlayWriterCDPRelayServer({
     // 2. We might still have messages in flight or try to send
     // This can cause "Assertion error" in Playwright's crConnection.js if a response
     // arrives after callbacks were cleared. We wrap in try-catch to handle this gracefully.
-    const safeSend = (client: PlaywrightClient) => {
+    const safeSend = (client: relayState.PlaywrightClient) => {
       try {
         client.ws.send(messageStr)
       } catch (e) {
@@ -328,13 +331,13 @@ export async function startPlayWriterCDPRelayServer({
     }
 
     if (clientId) {
-      const client = playwrightClients.get(clientId)
+      const client = store.getState().playwrightClients.get(clientId)
       if (client) {
         safeSend(client)
       }
     } else {
-      const clients = Array.from(playwrightClients.values())
-      for (const client of clients) {
+      const { playwrightClients } = store.getState()
+      for (const client of playwrightClients.values()) {
         if (extensionId && client.extensionId !== extensionId) {
           continue
         }
@@ -372,12 +375,28 @@ export async function startPlayWriterCDPRelayServer({
     params?: unknown
     timeout?: number
   }): Promise<unknown> {
-    const connection = getExtensionConnection(extensionId)
-    if (!connection) {
+    const conn = getExtensionConnection(extensionId)
+    if (!conn) {
+      throw new Error('Extension not connected')
+    }
+    const resolvedExtensionId = conn.id
+
+    let id = 0
+    store.setState((s) => {
+      const ext = s.extensions.get(resolvedExtensionId)
+      if (!ext) {
+        return s
+      }
+      id = ext.messageId + 1
+      const newExtensions = new Map(s.extensions)
+      newExtensions.set(resolvedExtensionId, { ...ext, messageId: id })
+      return { ...s, extensions: newExtensions }
+    })
+
+    if (!id) {
       throw new Error('Extension not connected')
     }
 
-    const id = ++connection.messageId
     const message = { id, method, params }
 
     const forwardCdpParams = method === 'forwardCDPCommand' ? getForwardCdpParams(params) : undefined
@@ -393,15 +412,18 @@ export async function startPlayWriterCDPRelayServer({
       })
     }
 
-    connection.ws.send(JSON.stringify(message))
-
     return new Promise((resolve, reject) => {
       const timeoutId = setTimeout(() => {
-        connection.pendingRequests.delete(id)
+        store.setState((s) =>
+          relayState.removeExtensionPendingRequest(s, {
+            extensionId: resolvedExtensionId,
+            requestId: id,
+          }),
+        )
         reject(new Error(`Extension request timeout after ${timeout}ms: ${method}`))
       }, timeout)
 
-      connection.pendingRequests.set(id, {
+      const pendingRequest = {
         resolve: (result) => {
           clearTimeout(timeoutId)
           resolve(result)
@@ -409,30 +431,90 @@ export async function startPlayWriterCDPRelayServer({
         reject: (error) => {
           clearTimeout(timeoutId)
           reject(error)
-        }
-      })
+        },
+      }
+
+      store.setState((s) =>
+        relayState.addExtensionPendingRequest(s, {
+          extensionId: resolvedExtensionId,
+          requestId: id,
+          pendingRequest,
+        }),
+      )
+
+      const latestExt = store.getState().extensions.get(resolvedExtensionId)
+      if (!latestExt?.ws) {
+        clearTimeout(timeoutId)
+        store.setState((s) =>
+          relayState.removeExtensionPendingRequest(s, {
+            extensionId: resolvedExtensionId,
+            requestId: id,
+          }),
+        )
+        reject(new Error('Extension not connected'))
+        return
+      }
+
+      try {
+        latestExt.ws.send(JSON.stringify(message))
+      } catch (error) {
+        clearTimeout(timeoutId)
+        store.setState((s) =>
+          relayState.removeExtensionPendingRequest(s, {
+            extensionId: resolvedExtensionId,
+            requestId: id,
+          }),
+        )
+        const sendError = error instanceof Error ? error : new Error(String(error))
+        reject(new Error(`Extension send failed: ${method}`, { cause: sendError }))
+      }
     })
   }
 
   const recordingRelays = new Map<string, RecordingRelay>()
 
+  // Find which extension connection owns a CDP tab session ID (pw-tab-*).
+  // Used by recording routes where sessionId identifies the target tab.
+  // Delegates to the pure derivation function from relay-state.ts.
+  const findExtensionIdByCdpSession = (cdpSessionId: string): string | null => {
+    return relayState.findExtensionIdByCdpSession(store.getState(), cdpSessionId)
+  }
+
+  // Resolve recording route session ID (CDP tab session) to extension connection.
+  const resolveRecordingRoute = async ({
+    sessionId,
+  }: {
+    sessionId: string | null
+  }): Promise<{
+    extensionId: string | null
+    sessionId: string | null
+  }> => {
+    if (!sessionId) {
+      return { extensionId: null, sessionId: null }
+    }
+
+    const extensionId = findExtensionIdByCdpSession(sessionId)
+    return { extensionId, sessionId }
+  }
+
   const getRecordingRelay = (extensionId?: string | null): RecordingRelay | null => {
-    const allowDefault = !extensionId && extensionConnections.size === 1
-    const connection = getExtensionConnection(extensionId, { allowFallback: allowDefault })
-    if (!connection) {
+    const allowDefault = !extensionId && store.getState().extensions.size === 1
+    const conn = getExtensionConnection(extensionId, { allowFallback: allowDefault })
+    if (!conn) {
       return null
     }
-    if (!recordingRelays.has(connection.id)) {
+    const connId = conn.id
+    if (!recordingRelays.has(connId)) {
       recordingRelays.set(
-        connection.id,
+        connId,
         new RecordingRelay(
-          (params) => sendToExtension({ extensionId: connection.id, ...params }),
-          () => extensionConnections.has(connection.id),
+          (params) => sendToExtension({ extensionId: connId, ...params }),
+          () => store.getState().extensions.has(connId),
           logger,
-        )
+        ),
       )
     }
-    return recordingRelays.get(connection.id) || null
+    return recordingRelays.get(connId) || null
   }
 
   // Auto-create initial tab when PLAYWRITER_AUTO_ENABLE is set and no targets exist.
@@ -441,36 +523,127 @@ export async function startPlayWriterCDPRelayServer({
     if (!process.env.PLAYWRITER_AUTO_ENABLE) {
       return
     }
-    const connection = getExtensionConnection(extensionId)
-    if (!connection) {
+    const conn = getExtensionConnection(extensionId)
+    if (!conn) {
       return
     }
-    if (connection.connectedTargets.size > 0) {
+    if (conn.connectedTargets.size > 0) {
       return
     }
 
     try {
       logger?.log(pc.blue('Auto-creating initial tab for Playwright client'))
-      const result = await sendToExtension({ extensionId, method: 'createInitialTab', timeout: 10000 }) as {
+      const result = (await sendToExtension({ extensionId, method: 'createInitialTab', timeout: 10000 })) as {
         success: boolean
         tabId: number
         sessionId: string
         targetInfo: Protocol.Target.TargetInfo
       }
       if (result.success && result.sessionId && result.targetInfo) {
-        connection.connectedTargets.set(result.sessionId, {
-          sessionId: result.sessionId,
-          targetId: result.targetInfo.targetId,
-          targetInfo: result.targetInfo,
-          frameIds: new Set()
-        })
-        logger?.log(pc.blue(`Auto-created tab, now have ${connection.connectedTargets.size} targets, url: ${result.targetInfo.url}`))
+        store.setState((s) =>
+          relayState.addTarget(s, {
+            extensionId,
+            sessionId: result.sessionId,
+            targetId: result.targetInfo.targetId,
+            targetInfo: result.targetInfo,
+          }),
+        )
+        const updatedTargets = store.getState().extensions.get(extensionId)?.connectedTargets.size || 0
+        logger?.log(
+          pc.blue(`Auto-created tab, now have ${updatedTargets} targets, url: ${result.targetInfo.url}`),
+        )
       }
     } catch (e) {
       logger?.error('Failed to auto-create initial tab:', e)
     }
   }
 
+  function getPageTargetSessionIds({ extensionId }: { extensionId: string }): string[] {
+    const extensionState = store.getState().extensions.get(extensionId)
+    if (!extensionState) {
+      return []
+    }
+    return Array.from(extensionState.connectedTargets.values())
+      .filter((target) => {
+        return target.targetInfo.type === 'page'
+      })
+      .map((target) => {
+        return target.sessionId
+      })
+  }
+
+  function maybeEmitBrowserDownloadCompatEvent({
+    method,
+    params,
+    extensionId,
+  }: {
+    method: string
+    params: unknown
+    extensionId: string
+  }): void {
+    const browserEventMethod =
+      method === 'Page.downloadWillBegin'
+        ? 'Browser.downloadWillBegin'
+        : method === 'Page.downloadProgress'
+          ? 'Browser.downloadProgress'
+          : null
+    if (!browserEventMethod) {
+      return
+    }
+    sendToPlaywright({
+      message: {
+        method: browserEventMethod,
+        params,
+      } as CDPEventBase,
+      source: 'server',
+      extensionId,
+    })
+  }
+
+  async function applyDownloadBehaviorToTargets({
+    extensionId,
+    behavior,
+    source,
+    targetSessionIds,
+  }: {
+    extensionId: string
+    behavior: Protocol.Browser.SetDownloadBehaviorRequest
+    source?: CDPCommand['source']
+    targetSessionIds?: string[]
+  }): Promise<void> {
+    const pageBehavior: Protocol.Page.SetDownloadBehaviorRequest['behavior'] =
+      behavior.behavior === 'allowAndName' ? 'allow' : behavior.behavior
+    const pageParams: Protocol.Page.SetDownloadBehaviorRequest = (() => {
+      if (pageBehavior === 'allow' && behavior.downloadPath) {
+        return { behavior: pageBehavior, downloadPath: behavior.downloadPath }
+      }
+      return { behavior: pageBehavior }
+    })()
+    const sessions = targetSessionIds || getPageTargetSessionIds({ extensionId })
+    if (sessions.length === 0) {
+      return
+    }
+    await Promise.all(
+      sessions.map(async (targetSessionId) => {
+        try {
+          await sendToExtension({
+            extensionId,
+            method: 'forwardCDPCommand',
+            params: {
+              sessionId: targetSessionId,
+              method: 'Page.setDownloadBehavior',
+              params: pageParams,
+              source,
+            },
+          })
+        } catch (error) {
+          const message = error instanceof Error ? error.message : String(error)
+          logger?.log(pc.yellow(`[Server] Failed to apply Page.setDownloadBehavior to ${targetSessionId}: ${message}`))
+        }
+      }),
+    )
+  }
+
   async function routeCdpCommand({
     extensionId,
     method,
@@ -479,13 +652,14 @@ export async function startPlayWriterCDPRelayServer({
     source,
   }: {
     extensionId: string | null
-    method: string
-    params: any
-    sessionId?: string
-    source?: 'playwriter'
+    method: CDPCommand['method'] | (string & {})
+    params: CDPCommand['params']
+    sessionId?: CDPCommand['sessionId']
+    source?: CDPCommand['source']
   }) {
-    const extension = getExtensionConnection(extensionId)
-    const connectedTargets = extension?.connectedTargets || new Map<string, ConnectedTarget>()
+    const conn = getExtensionConnection(extensionId)
+    const connectedTargets = conn?.connectedTargets || new Map<string, relayState.ConnectedTarget>()
+    const resolvedExtensionId = conn?.id || extensionId
     switch (method) {
       case 'Browser.getVersion': {
         return {
@@ -493,11 +667,23 @@ export async function startPlayWriterCDPRelayServer({
           product: 'Chrome/Extension-Bridge',
           revision: '1.0.0',
           userAgent: 'CDP-Bridge-Server/1.0.0',
-          jsVersion: 'V8'
+          jsVersion: 'V8',
         } satisfies Protocol.Browser.GetVersionResponse
       }
 
       case 'Browser.setDownloadBehavior': {
+        const downloadBehaviorParams = params as Protocol.Browser.SetDownloadBehaviorRequest | undefined
+        if (!downloadBehaviorParams?.behavior) {
+          throw new Error('behavior is required for Browser.setDownloadBehavior')
+        }
+        if (resolvedExtensionId) {
+          extensionDownloadBehavior.set(resolvedExtensionId, downloadBehaviorParams)
+          await applyDownloadBehaviorToTargets({
+            extensionId: resolvedExtensionId,
+            behavior: downloadBehaviorParams,
+            source,
+          })
+        }
         return {}
       }
 
@@ -508,15 +694,15 @@ export async function startPlayWriterCDPRelayServer({
         if (sessionId) {
           break
         }
-        if (extension) {
-          await maybeAutoCreateInitialTab(extension.id)
+        if (conn) {
+          await maybeAutoCreateInitialTab(conn.id)
         }
         // Forward auto-attach so Chrome emits iframe Target.attachedToTarget events.
         // Playwright relies on these (with parentFrameId) when reconnecting over CDP.
         await sendToExtension({
-          extensionId: extension?.id || extensionId,
+          extensionId: resolvedExtensionId,
           method: 'forwardCDPCommand',
-          params: { method, params, source }
+          params: { method, params, source },
         })
         return {}
       }
@@ -526,22 +712,23 @@ export async function startPlayWriterCDPRelayServer({
       }
 
       case 'Target.attachToTarget': {
-        const targetId = params?.targetId
-        if (!targetId) {
+        const attachParams = params as Protocol.Target.AttachToTargetRequest
+        if (!attachParams?.targetId) {
           throw new Error('targetId is required for Target.attachToTarget')
         }
 
         for (const target of connectedTargets.values()) {
-          if (target.targetId === targetId) {
+          if (target.targetId === attachParams.targetId) {
             return { sessionId: target.sessionId } satisfies Protocol.Target.AttachToTargetResponse
           }
         }
 
-        throw new Error(`Target ${targetId} not found in connected targets`)
+        throw new Error(`Target ${attachParams.targetId} not found in connected targets`)
       }
 
       case 'Target.getTargetInfo': {
-        const targetId = params?.targetId
+        const infoReqParams = params as Protocol.Target.GetTargetInfoRequest | undefined
+        const targetId = infoReqParams?.targetId
 
         if (targetId) {
           for (const target of connectedTargets.values()) {
@@ -568,33 +755,33 @@ export async function startPlayWriterCDPRelayServer({
             .filter((t) => !isRestrictedTarget(t.targetInfo))
             .map((t) => ({
               ...t.targetInfo,
-              attached: true
-            }))
+              attached: true,
+            })),
         }
       }
 
       case 'Target.createTarget': {
         return await sendToExtension({
-          extensionId: extension?.id || extensionId,
+          extensionId: resolvedExtensionId,
           method: 'forwardCDPCommand',
-          params: { method, params, source }
+          params: { method, params, source },
         })
       }
 
       case 'Target.closeTarget': {
         return await sendToExtension({
-          extensionId: extension?.id || extensionId,
+          extensionId: resolvedExtensionId,
           method: 'forwardCDPCommand',
-          params: { method, params, source }
+          params: { method, params, source },
         })
       }
 
       // Ghost Browser API - forward to extension for chrome.ghostPublicAPI/ghostProxies/projects
       case 'ghost-browser': {
         return await sendToExtension({
-          extensionId: extension?.id || extensionId,
+          extensionId: resolvedExtensionId,
           method: 'ghost-browser',
-          params
+          params,
         })
       }
 
@@ -616,16 +803,20 @@ export async function startPlayWriterCDPRelayServer({
           }
           const timeout = setTimeout(() => {
             emitter.off('cdp:event', handler)
-            logger?.log(pc.yellow(`IMPORTANT: Runtime.enable timed out waiting for main frame executionContextCreated (sessionId: ${sessionId}). This may cause pages to not be visible immediately.`))
+            logger?.log(
+              pc.yellow(
+                `IMPORTANT: Runtime.enable timed out waiting for main frame executionContextCreated (sessionId: ${sessionId}). This may cause pages to not be visible immediately.`,
+              ),
+            )
             resolve()
           }, 3000)
           emitter.on('cdp:event', handler)
         })
 
         const result = await sendToExtension({
-          extensionId: extension?.id || extensionId,
+          extensionId: resolvedExtensionId,
           method: 'forwardCDPCommand',
-          params: { sessionId, method, params, source }
+          params: { sessionId, method, params, source },
         })
 
         await contextCreatedPromise
@@ -635,9 +826,9 @@ export async function startPlayWriterCDPRelayServer({
     }
 
     return await sendToExtension({
-      extensionId: extension?.id || extensionId,
+      extensionId: resolvedExtensionId,
       method: 'forwardCDPCommand',
-      params: { sessionId, method, params, source }
+      params: { sessionId, method, params, source },
     })
   }
 
@@ -645,19 +836,22 @@ export async function startPlayWriterCDPRelayServer({
   // CORS middleware for HTTP endpoints - only allows our specific extension IDs.
   // This prevents other extensions from reading responses via fetch/XHR.
   // WebSocket connections have their own separate origin validation.
-  app.use('*', cors({
-    origin: (origin) => {
-      if (!origin.startsWith('chrome-extension://')) {
-        return null
-      }
-      const extensionId = origin.replace('chrome-extension://', '')
-      if (!EXTENSION_IDS.includes(extensionId)) {
-        return null
-      }
-      return origin
-    },
-    allowMethods: ['GET', 'POST', 'HEAD', 'OPTIONS'],
-  }))
+  app.use(
+    '*',
+    cors({
+      origin: (origin) => {
+        if (!origin.startsWith('chrome-extension://')) {
+          return null
+        }
+        const extensionId = origin.replace('chrome-extension://', '')
+        if (!EXTENSION_IDS.includes(extensionId)) {
+          return null
+        }
+        return origin
+      },
+      allowMethods: ['GET', 'POST', 'HEAD', 'OPTIONS'],
+    }),
+  )
   const { injectWebSocket, upgradeWebSocket } = createNodeWebSocket({ app })
 
   const getCdpWsUrl = (c: { req: { header: (name: string) => string | undefined } }) => {
@@ -675,7 +869,7 @@ export async function startPlayWriterCDPRelayServer({
 
   app.get('/extension/status', (c) => {
     const defaultExtension = getExtensionConnection(null, { allowFallback: true })
-    const connected = extensionConnections.size > 0
+    const connected = store.getState().extensions.size > 0
     const activeTargets = defaultExtension?.connectedTargets.size || 0
     const info = defaultExtension?.info
 
@@ -689,14 +883,14 @@ export async function startPlayWriterCDPRelayServer({
   })
 
   app.get('/extensions/status', (c) => {
-    const extensions = Array.from(extensionConnections.values()).map((extension) => {
+    const extensions = Array.from(store.getState().extensions.values()).map((ext) => {
       return {
-        extensionId: extension.id,
-        stableKey: extension.stableKey,
-        browser: extension.info.browser || null,
-        profile: extension.info ? { email: extension.info.email || '', id: extension.info.id || '' } : null,
-        activeTargets: extension.connectedTargets.size,
-        playwriterVersion: extension.info?.version || null,
+        extensionId: ext.id,
+        stableKey: ext.stableKey,
+        browser: ext.info.browser || null,
+        profile: ext.info ? { email: ext.info.email || '', id: ext.info.id || '' } : null,
+        activeTargets: ext.connectedTargets.size,
+        playwriterVersion: ext.info?.version || null,
       }
     })
     return c.json({ extensions })
@@ -709,76 +903,76 @@ export async function startPlayWriterCDPRelayServer({
   app
     .on(['GET', 'PUT'], '/json/version', (c) => {
       return c.json({
-        'Browser': `Playwriter/${VERSION}`,
+        Browser: `Playwriter/${VERSION}`,
         'Protocol-Version': '1.3',
-        'webSocketDebuggerUrl': getCdpWsUrl(c)
+        webSocketDebuggerUrl: getCdpWsUrl(c),
       })
     })
     .on(['GET', 'PUT'], '/json/version/', (c) => {
       return c.json({
-        'Browser': `Playwriter/${VERSION}`,
+        Browser: `Playwriter/${VERSION}`,
         'Protocol-Version': '1.3',
-        'webSocketDebuggerUrl': getCdpWsUrl(c)
+        webSocketDebuggerUrl: getCdpWsUrl(c),
       })
     })
     .on(['GET', 'PUT'], '/json/list', (c) => {
       const wsUrl = getCdpWsUrl(c)
       const defaultTargets = getExtensionConnection(null, { allowFallback: true })?.connectedTargets || new Map()
       return c.json(
-        Array.from(defaultTargets.values()).map(t => ({
+        Array.from(defaultTargets.values()).map((t) => ({
           id: t.targetId,
           type: t.targetInfo.type,
           title: t.targetInfo.title,
           description: t.targetInfo.title,
           url: t.targetInfo.url,
           webSocketDebuggerUrl: wsUrl,
-          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`
-        }))
+          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`,
+        })),
       )
     })
     .on(['GET', 'PUT'], '/json/list/', (c) => {
       const wsUrl = getCdpWsUrl(c)
       const defaultTargets = getExtensionConnection(null, { allowFallback: true })?.connectedTargets || new Map()
       return c.json(
-        Array.from(defaultTargets.values()).map(t => ({
+        Array.from(defaultTargets.values()).map((t) => ({
           id: t.targetId,
           type: t.targetInfo.type,
           title: t.targetInfo.title,
           description: t.targetInfo.title,
           url: t.targetInfo.url,
           webSocketDebuggerUrl: wsUrl,
-          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`
-        }))
+          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`,
+        })),
       )
     })
     .on(['GET', 'PUT'], '/json', (c) => {
       const wsUrl = getCdpWsUrl(c)
       const defaultTargets = getExtensionConnection(null, { allowFallback: true })?.connectedTargets || new Map()
       return c.json(
-        Array.from(defaultTargets.values()).map(t => ({
+        Array.from(defaultTargets.values()).map((t) => ({
           id: t.targetId,
           type: t.targetInfo.type,
           title: t.targetInfo.title,
           description: t.targetInfo.title,
           url: t.targetInfo.url,
           webSocketDebuggerUrl: wsUrl,
-          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`
-        }))
+          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`,
+        })),
       )
     })
     .on(['GET', 'PUT'], '/json/', (c) => {
       const wsUrl = getCdpWsUrl(c)
       const defaultTargets = getExtensionConnection(null, { allowFallback: true })?.connectedTargets || new Map()
       return c.json(
-        Array.from(defaultTargets.values()).map(t => ({
+        Array.from(defaultTargets.values()).map((t) => ({
           id: t.targetId,
           type: t.targetInfo.type,
           title: t.targetInfo.title,
           description: t.targetInfo.title,
           url: t.targetInfo.url,
           webSocketDebuggerUrl: wsUrl,
-          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`
-        }))
+          devtoolsFrontendUrl: `/devtools/inspector.html?ws=${wsUrl.replace('ws://', '')}`,
+        })),
       )
     })
 
@@ -798,218 +992,274 @@ export async function startPlayWriterCDPRelayServer({
   // Browsers always send Origin header for WebSocket connections, but Node.js clients don't.
   // We only allow our specific extension IDs to prevent malicious websites or extensions
   // from connecting to the local WebSocket server.
-  app.get('/cdp/:clientId?', (c, next) => {
-    const origin = c.req.header('origin')
-
-    // Validate Origin header if present (Node.js clients don't send it)
-    if (origin) {
-      if (origin.startsWith('chrome-extension://')) {
-        const extensionId = origin.replace('chrome-extension://', '')
-        if (!EXTENSION_IDS.includes(extensionId)) {
-          logger?.log(pc.red(`Rejecting /cdp WebSocket from unknown extension: ${extensionId}`))
+  app.get(
+    '/cdp/:clientId?',
+    (c, next) => {
+      const origin = c.req.header('origin')
+
+      // Validate Origin header if present (Node.js clients don't send it)
+      if (origin) {
+        if (origin.startsWith('chrome-extension://')) {
+          const extensionId = origin.replace('chrome-extension://', '')
+          if (!EXTENSION_IDS.includes(extensionId)) {
+            logger?.log(pc.red(`Rejecting /cdp WebSocket from unknown extension: ${extensionId}`))
+            return c.text('Forbidden', 403)
+          }
+        } else {
+          logger?.log(pc.red(`Rejecting /cdp WebSocket from origin: ${origin}`))
           return c.text('Forbidden', 403)
         }
-      } else {
-        logger?.log(pc.red(`Rejecting /cdp WebSocket from origin: ${origin}`))
-        return c.text('Forbidden', 403)
       }
-    }
 
-    if (token) {
+      if (token) {
+        const url = new URL(c.req.url, 'http://localhost')
+        const providedToken = url.searchParams.get('token')
+        if (providedToken !== token) {
+          return c.text('Unauthorized', 401)
+        }
+      }
+      return next()
+    },
+    upgradeWebSocket((c) => {
+      const clientId = c.req.param('clientId') || 'default'
       const url = new URL(c.req.url, 'http://localhost')
-      const providedToken = url.searchParams.get('token')
-      if (providedToken !== token) {
-        return c.text('Unauthorized', 401)
+      const requestedExtensionId = url.searchParams.get('extensionId')
+      // When extensionId is explicit, resolve directly. Otherwise use fallback which
+      // handles single-extension and uniquely-active-extension cases (#52).
+      const resolvedExtension = requestedExtensionId
+        ? getExtensionConnection(requestedExtensionId)
+        : getExtensionConnection(null, { allowFallback: true })
+      const clientExtensionId = resolvedExtension?.id || null
+
+      const getBoundExtensionIdForClient = (): string | null => {
+        const client = store.getState().playwrightClients.get(clientId)
+        return client?.extensionId || null
       }
-    }
-    return next()
-  }, upgradeWebSocket((c) => {
-    const clientId = c.req.param('clientId') || 'default'
-    const url = new URL(c.req.url, 'http://localhost')
-    const requestedExtensionId = url.searchParams.get('extensionId')
-    const resolvedExtension = getExtensionConnection(requestedExtensionId)
-    const allowDefault = !requestedExtensionId && extensionConnections.size === 1
-    const defaultExtension = allowDefault ? getExtensionConnection(null, { allowFallback: true }) : null
-    const clientExtensionId = resolvedExtension?.id || defaultExtension?.id || null
 
-    return {
-      async onOpen(_event, ws) {
-        if (playwrightClients.has(clientId)) {
-          logger?.log(pc.yellow(`Rejecting duplicate Playwright clientId: ${clientId}`))
-          ws.close(4004, 'Duplicate Playwright clientId')
-          return
-        }
+      return {
+        async onOpen(_event, ws) {
+          if (store.getState().playwrightClients.has(clientId)) {
+            logger?.log(pc.yellow(`Rejecting duplicate Playwright clientId: ${clientId}`))
+            ws.close(4004, 'Duplicate Playwright clientId')
+            return
+          }
 
-        if (!clientExtensionId) {
-          const reason = requestedExtensionId
-            ? `Unknown extensionId: ${requestedExtensionId}`
-            : 'Multiple extensions connected. Specify extensionId.'
-          logger?.log(pc.yellow(`Rejecting Playwright client ${clientId}: ${reason}`))
-          ws.close(4003, reason)
-          return
-        }
+          if (!clientExtensionId) {
+            const reason = requestedExtensionId
+              ? `Unknown extensionId: ${requestedExtensionId}`
+              : 'Multiple extensions connected. Specify extensionId.'
+            logger?.log(pc.yellow(`Rejecting Playwright client ${clientId}: ${reason}`))
+            ws.close(4003, reason)
+            return
+          }
 
-        // Add client first so it can receive Target.attachedToTarget events
-        playwrightClients.set(clientId, { id: clientId, ws, extensionId: clientExtensionId })
-        const extensionConnection = getExtensionConnection(clientExtensionId)
-        const targetCount = extensionConnection?.connectedTargets.size || 0
-        logger?.log(pc.green(`Playwright client connected: ${clientId} (${playwrightClients.size} total) (extension? ${!!extensionConnection}) (${targetCount} pages)`))
-      },
+          // Add client first so it can receive Target.attachedToTarget events
+          store.setState((s) => {
+            return relayState.addPlaywrightClient(s, { id: clientId, extensionId: clientExtensionId, ws })
+          })
+          const extensionConnection = getExtensionConnection(clientExtensionId)
+          const targetCount = extensionConnection?.connectedTargets.size || 0
+          logger?.log(
+            pc.green(
+              `Playwright client connected: ${clientId} (${store.getState().playwrightClients.size} total) (extension? ${!!extensionConnection}) (${targetCount} pages)`,
+            ),
+          )
+        },
 
-      async onMessage(event, ws) {
-        let message: CDPCommand
+        async onMessage(event, ws) {
+          let message: CDPCommand
 
-        try {
-          message = JSON.parse(event.data.toString())
-        } catch {
-          return
-        }
+          try {
+            message = JSON.parse(event.data.toString())
+          } catch {
+            return
+          }
 
-        const { id, sessionId, method, params, source } = message
+          const { id, sessionId, method, params, source } = message
 
-        logCdpJson({
-          timestamp: new Date().toISOString(),
-          direction: 'from-playwright',
-          clientId,
-          message,
-        })
+          logCdpJson({
+            timestamp: new Date().toISOString(),
+            direction: 'from-playwright',
+            clientId,
+            message,
+          })
 
-        logCdpMessage({
-          direction: 'from-playwright',
-          clientId,
-          method,
-          sessionId,
-          id
-        })
+          logCdpMessage({
+            direction: 'from-playwright',
+            clientId,
+            method,
+            sessionId,
+            id,
+          })
 
-        emitter.emit('cdp:command', { clientId, command: message })
+          emitter.emit('cdp:command', { clientId, command: message })
 
-        const extensionConnection = getExtensionConnection(clientExtensionId)
-        if (!extensionConnection) {
-          sendToPlaywright({
-            message: {
-              id,
-              sessionId,
-              error: { message: 'Extension not connected' }
-            },
-            clientId
-          })
-          return
-        }
+          const boundExtensionId = getBoundExtensionIdForClient()
+          const extensionConn = getExtensionConnection(boundExtensionId)
+          if (!extensionConn) {
+            sendToPlaywright({
+              message: {
+                id,
+                sessionId,
+                error: { message: 'Extension not connected' },
+              },
+              clientId,
+            })
+            return
+          }
 
-        try {
-          const result: any = await routeCdpCommand({ extensionId: extensionConnection.id, method, params, sessionId, source })
+          try {
+            const result = await routeCdpCommand({
+              extensionId: extensionConn.id,
+              method,
+              params,
+              sessionId,
+              source,
+            })
 
-          if (method === 'Target.setAutoAttach' && !sessionId) {
-            for (const target of extensionConnection.connectedTargets.values()) {
-              // Skip restricted targets (extensions, chrome:// URLs, non-page types)
-              if (isRestrictedTarget(target.targetInfo)) {
-                continue
-              }
-              const attachedPayload = {
-                method: 'Target.attachedToTarget',
-                params: {
-                  sessionId: target.sessionId,
-                  targetInfo: {
-                    ...target.targetInfo,
-                    attached: true
+            if (method === 'Target.setAutoAttach' && !sessionId) {
+              // Re-read state after async routeCdpCommand — targets may have changed
+              const freshExt = store.getState().extensions.get(extensionConn.id)
+              const freshTargets = freshExt?.connectedTargets || new Map()
+              for (const target of freshTargets.values()) {
+                // Skip restricted targets (extensions, chrome:// URLs, non-page types)
+                if (isRestrictedTarget(target.targetInfo)) {
+                  continue
+                }
+                const attachedPayload = {
+                  method: 'Target.attachedToTarget',
+                  params: {
+                    sessionId: target.sessionId,
+                    targetInfo: {
+                      ...target.targetInfo,
+                      attached: true,
+                    },
+                    waitingForDebugger: false,
                   },
-                  waitingForDebugger: false
+                } satisfies CDPEventFor<'Target.attachedToTarget'>
+                if (!target.targetInfo.url) {
+                  logger?.error(
+                    pc.red('[Server] WARNING: Target.attachedToTarget sent with empty URL!'),
+                    JSON.stringify(attachedPayload),
+                  )
                 }
-              } satisfies CDPEventFor<'Target.attachedToTarget'>
-              if (!target.targetInfo.url) {
-                logger?.error(pc.red('[Server] WARNING: Target.attachedToTarget sent with empty URL!'), JSON.stringify(attachedPayload))
+                logger?.log(
+                  pc.magenta('[Server] Target.attachedToTarget full payload:'),
+                  JSON.stringify(attachedPayload),
+                )
+                sendToPlaywright({
+                  message: attachedPayload,
+                  clientId,
+                  source: 'server',
+                })
               }
-              logger?.log(pc.magenta('[Server] Target.attachedToTarget full payload:'), JSON.stringify(attachedPayload))
-              sendToPlaywright({
-                message: attachedPayload,
-                clientId,
-                source: 'server'
-              })
             }
-          }
 
-          if (method === 'Target.setDiscoverTargets' && (params as any)?.discover) {
-            for (const target of extensionConnection.connectedTargets.values()) {
-              // Skip restricted targets (extensions, chrome:// URLs, non-page types)
-              if (isRestrictedTarget(target.targetInfo)) {
-                continue
-              }
-              const targetCreatedPayload = {
-                method: 'Target.targetCreated',
-                params: {
-                  targetInfo: {
-                    ...target.targetInfo,
-                    attached: true
-                  }
+            if (method === 'Target.setDiscoverTargets' && (params as Protocol.Target.SetDiscoverTargetsRequest)?.discover) {
+              const freshExt2 = store.getState().extensions.get(extensionConn.id)
+              const freshTargets2 = freshExt2?.connectedTargets || new Map()
+              for (const target of freshTargets2.values()) {
+                // Skip restricted targets (extensions, chrome:// URLs, non-page types)
+                if (isRestrictedTarget(target.targetInfo)) {
+                  continue
+                }
+                const targetCreatedPayload = {
+                  method: 'Target.targetCreated',
+                  params: {
+                    targetInfo: {
+                      ...target.targetInfo,
+                      attached: true,
+                    },
+                  },
+                } satisfies CDPEventFor<'Target.targetCreated'>
+                if (!target.targetInfo.url) {
+                  logger?.error(
+                    pc.red('[Server] WARNING: Target.targetCreated sent with empty URL!'),
+                    JSON.stringify(targetCreatedPayload),
+                  )
                 }
-              } satisfies CDPEventFor<'Target.targetCreated'>
-              if (!target.targetInfo.url) {
-                logger?.error(pc.red('[Server] WARNING: Target.targetCreated sent with empty URL!'), JSON.stringify(targetCreatedPayload))
+                logger?.log(
+                  pc.magenta('[Server] Target.targetCreated full payload:'),
+                  JSON.stringify(targetCreatedPayload),
+                )
+                sendToPlaywright({
+                  message: targetCreatedPayload,
+                  clientId,
+                  source: 'server',
+                })
               }
-              logger?.log(pc.magenta('[Server] Target.targetCreated full payload:'), JSON.stringify(targetCreatedPayload))
-              sendToPlaywright({
-                message: targetCreatedPayload,
-                clientId,
-                source: 'server'
-              })
             }
-          }
 
-          if (method === 'Target.attachToTarget' && result?.sessionId) {
-            const targetId = params?.targetId
-            const target = Array.from(extensionConnection.connectedTargets.values()).find(t => t.targetId === targetId)
-            if (target) {
-              const attachedPayload = {
-                method: 'Target.attachedToTarget',
-                params: {
-                  sessionId: result.sessionId,
-                  targetInfo: {
-                    ...target.targetInfo,
-                    attached: true
-                  },
-                  waitingForDebugger: false
+            if (method === 'Target.attachToTarget') {
+              const attachResponse = result as Protocol.Target.AttachToTargetResponse | undefined
+              const attachRequestParams = params as Protocol.Target.AttachToTargetRequest | undefined
+              if (attachResponse?.sessionId) {
+                const freshExt3 = store.getState().extensions.get(extensionConn.id)
+                const freshTargets3 = freshExt3?.connectedTargets || new Map()
+                const target = Array.from(freshTargets3.values()).find((t) => {
+                  return t.targetId === attachRequestParams?.targetId
+                })
+                if (target) {
+                  const attachedPayload = {
+                    method: 'Target.attachedToTarget',
+                    params: {
+                      sessionId: attachResponse.sessionId,
+                      targetInfo: {
+                        ...target.targetInfo,
+                        attached: true,
+                      },
+                      waitingForDebugger: false,
+                    },
+                  } satisfies CDPEventFor<'Target.attachedToTarget'>
+                  if (!target.targetInfo.url) {
+                    logger?.error(
+                      pc.red('[Server] WARNING: Target.attachedToTarget (from attachToTarget) sent with empty URL!'),
+                      JSON.stringify(attachedPayload),
+                    )
+                  }
+                  logger?.log(
+                    pc.magenta('[Server] Target.attachedToTarget (from attachToTarget) payload:'),
+                    JSON.stringify(attachedPayload),
+                  )
+                  sendToPlaywright({
+                    message: attachedPayload,
+                    clientId,
+                    source: 'server',
+                  })
                 }
-              } satisfies CDPEventFor<'Target.attachedToTarget'>
-              if (!target.targetInfo.url) {
-                logger?.error(pc.red('[Server] WARNING: Target.attachedToTarget (from attachToTarget) sent with empty URL!'), JSON.stringify(attachedPayload))
               }
-              logger?.log(pc.magenta('[Server] Target.attachedToTarget (from attachToTarget) payload:'), JSON.stringify(attachedPayload))
-              sendToPlaywright({
-                message: attachedPayload,
-                clientId,
-                source: 'server'
-              })
             }
-          }
 
-          const response: CDPResponseBase = { id, sessionId, result }
-          sendToPlaywright({ message: response, clientId })
-          emitter.emit('cdp:response', { clientId, response, command: message })
-        } catch (e) {
-          logger?.error('Error handling CDP command:', method, params, e)
-          const errorResponse: CDPResponseBase = {
-            id,
-            sessionId,
-            error: { message: (e as Error).message }
+            const response: CDPResponseBase = { id, sessionId, result }
+            sendToPlaywright({ message: response, clientId })
+            emitter.emit('cdp:response', { clientId, response, command: message })
+          } catch (e) {
+            logger?.error('Error handling CDP command:', method, params, e)
+            const errorResponse: CDPResponseBase = {
+              id,
+              sessionId,
+              error: { message: (e as Error).message },
+            }
+            sendToPlaywright({ message: errorResponse, clientId })
+            emitter.emit('cdp:response', { clientId, response: errorResponse, command: message })
           }
-          sendToPlaywright({ message: errorResponse, clientId })
-          emitter.emit('cdp:response', { clientId, response: errorResponse, command: message })
-        }
-      },
+        },
 
-      onClose() {
-        playwrightClients.delete(clientId)
-        logger?.log(pc.yellow(`Playwright client disconnected: ${clientId} (${playwrightClients.size} remaining)`))
-      },
+        onClose() {
+          store.setState((s) => relayState.removePlaywrightClient(s, { clientId }))
+          logger?.log(pc.yellow(`Playwright client disconnected: ${clientId} (${store.getState().playwrightClients.size} remaining)`))
+        },
 
-      onError(event) {
-        logger?.error(`Playwright WebSocket error [${clientId}]:`, event)
+        onError(event) {
+          logger?.error(`Playwright WebSocket error [${clientId}]:`, event)
+        },
       }
-    }
-  }))
+    }),
+  )
 
-  const getExtensionInfoFromRequest = (c: { req: { query: (name: string) => string | undefined } }): ExtensionInfo => {
+  const getExtensionInfoFromRequest = (c: {
+    req: { query: (name: string) => string | undefined }
+  }): relayState.ExtensionInfo => {
     const browser = c.req.query('browser')
     const email = c.req.query('email')
     const id = c.req.query('id')
@@ -1022,402 +1272,467 @@ export async function startPlayWriterCDPRelayServer({
     }
   }
 
-  app.get('/extension', (c, next) => {
-    // 1. Host Validation: The extension endpoint must ONLY be accessed from localhost.
-    // This prevents attackers on the network from hijacking the browser session
-    // even if the server is exposed via 0.0.0.0.
-    const info = getConnInfo(c)
-    const remoteAddress = info.remote.address
-    const isLocalhost = remoteAddress === '127.0.0.1' || remoteAddress === '::1'
-
-    if (!isLocalhost) {
-      logger?.log(pc.red(`Rejecting /extension WebSocket from remote IP: ${remoteAddress}`))
-      return c.text('Forbidden - Extension must be local', 403)
-    }
-
-    // 2. Origin Validation: Prevent browser-based attacks (CSRF).
-    // Browsers cannot spoof the Origin header, so this ensures the connection
-    // is coming from our specific Chrome Extension, not a malicious website.
-    const origin = c.req.header('origin')
-    if (!origin || !origin.startsWith('chrome-extension://')) {
-      logger?.log(pc.red(`Rejecting /extension WebSocket: origin must be chrome-extension://, got: ${origin || 'none'}`))
-      return c.text('Forbidden', 403)
-    }
-
-    const extensionId = origin.replace('chrome-extension://', '')
-    if (!EXTENSION_IDS.includes(extensionId)) {
-      logger?.log(pc.red(`Rejecting /extension WebSocket from unknown extension: ${extensionId}`))
-      return c.text('Forbidden', 403)
-    }
+  app.get(
+    '/extension',
+    (c, next) => {
+      // 1. Host Validation: The extension endpoint must ONLY be accessed from localhost.
+      // This prevents attackers on the network from hijacking the browser session
+      // even if the server is exposed via 0.0.0.0.
+      const info = getConnInfo(c)
+      const remoteAddress = info.remote.address
+      const isLocalhost = remoteAddress === '127.0.0.1' || remoteAddress === '::1'
+
+      if (!isLocalhost) {
+        logger?.log(pc.red(`Rejecting /extension WebSocket from remote IP: ${remoteAddress}`))
+        return c.text('Forbidden - Extension must be local', 403)
+      }
 
-    return next()
-  }, upgradeWebSocket((c) => {
-    const incomingExtensionInfo = getExtensionInfoFromRequest(c)
-    const connectionId = `${Date.now().toString(36)}_${Math.random().toString(36).slice(2, 8)}`
-    return {
-      onOpen(_event, ws) {
-        const stableKey = buildStableExtensionKey(incomingExtensionInfo, connectionId)
-        const existingId = extensionKeyIndex.get(stableKey)
-        if (existingId && existingId !== connectionId) {
-          logger?.log(pc.yellow(`Replacing extension connection for ${stableKey} (${existingId} -> ${connectionId})`))
-          const existingConnection = extensionConnections.get(existingId)
-          if (existingConnection) {
-            existingConnection.ws.close(4001, 'Extension Replaced')
-          }
-        }
+      // 2. Origin Validation: Prevent browser-based attacks (CSRF).
+      // Browsers cannot spoof the Origin header, so this ensures the connection
+      // is coming from our specific Chrome Extension, not a malicious website.
+      const origin = c.req.header('origin')
+      if (!origin || !origin.startsWith('chrome-extension://')) {
+        logger?.log(
+          pc.red(`Rejecting /extension WebSocket: origin must be chrome-extension://, got: ${origin || 'none'}`),
+        )
+        return c.text('Forbidden', 403)
+      }
 
-        const connection: ExtensionConnection = {
-          id: connectionId,
-          ws,
-          info: incomingExtensionInfo,
-          stableKey,
-          connectedTargets: new Map(),
-          pendingRequests: new Map(),
-          messageId: 0,
-          pingInterval: null,
-        }
-        extensionConnections.set(connectionId, connection)
-        extensionKeyIndex.set(stableKey, connectionId)
-        startExtensionPing(connectionId)
-        logger?.log(`Extension connected (${connectionId})`)
-      },
+      const extensionId = origin.replace('chrome-extension://', '')
+      if (!EXTENSION_IDS.includes(extensionId)) {
+        logger?.log(pc.red(`Rejecting /extension WebSocket from unknown extension: ${extensionId}`))
+        return c.text('Forbidden', 403)
+      }
 
-      async onMessage(event, ws) {
-        const connection = extensionConnections.get(connectionId)
-        if (!connection) {
-          ws.close(1000, 'Extension not registered')
-          return
-        }
-        // Handle binary data (recording chunks)
-        if (event.data instanceof ArrayBuffer || Buffer.isBuffer(event.data)) {
-          const buffer = Buffer.isBuffer(event.data) ? event.data : Buffer.from(event.data)
-          const relay = getRecordingRelay(connectionId)
-          if (relay) {
-            relay.handleBinaryData(buffer)
+      return next()
+    },
+    upgradeWebSocket((c) => {
+      const incomingExtensionInfo = getExtensionInfoFromRequest(c)
+      const connectionId = `${Date.now().toString(36)}_${Math.random().toString(36).slice(2, 8)}`
+      return {
+        onOpen(_event, ws) {
+          const stableKey = buildStableExtensionKey(incomingExtensionInfo, connectionId)
+
+          // Check for existing connection with same stableKey and close it
+          const existingExt = relayState.findExtensionByStableKey(store.getState(), stableKey)
+          if (existingExt && existingExt.id !== connectionId) {
+            logger?.log(pc.yellow(`Replacing extension connection for ${stableKey} (${existingExt.id} -> ${connectionId})`))
+            if (existingExt.ws) {
+              existingExt.ws.close(4001, 'Extension Replaced')
+            }
           }
-          return
-        }
 
-        let message: ExtensionMessage
+          // State transition: add extension with ws handle included.
+          // Existing same-stableKey entry stays until old socket onClose.
+          store.setState((s) => {
+            return relayState.addExtension(s, { id: connectionId, info: incomingExtensionInfo, stableKey, ws })
+          })
 
-        try {
-          message = JSON.parse(event.data.toString())
-        } catch {
-          ws.close(1000, 'Invalid JSON')
-          return
-        }
+          startExtensionPing(connectionId)
+          logger?.log(`Extension connected (${connectionId})`)
+        },
 
-        if (message.id !== undefined) {
-          const pending = connection.pendingRequests.get(message.id)
-          if (!pending) {
-            logger?.log('Unexpected response with id:', message.id)
+        async onMessage(event, ws) {
+          const ext = store.getState().extensions.get(connectionId)
+          if (!ext) {
+            ws.close(1000, 'Extension not registered')
             return
           }
-
-          connection.pendingRequests.delete(message.id)
-
-          if (message.error) {
-            pending.reject(new Error(message.error))
-          } else {
-            pending.resolve(message.result)
-          }
-        } else if (message.method === 'pong') {
-          // Keep-alive response, nothing to do
-        } else if (message.method === 'log') {
-          const { level, args } = message.params
-          const logFn = (logger as Record<string, unknown>)?.[level] as ((...args: unknown[]) => void) | undefined
-          const logFunc = logFn || logger?.log
-          const prefix = pc.yellow(`[Extension] [${level.toUpperCase()}]`)
-          logFunc?.(prefix, ...args)
-        } else if (message.method === 'recordingData') {
-          const relay = getRecordingRelay(connectionId)
-          if (relay) {
-            relay.handleRecordingData(message as RecordingDataMessage)
-          }
-        } else if (message.method === 'recordingCancelled') {
-          const relay = getRecordingRelay(connectionId)
-          if (relay) {
-            relay.handleRecordingCancelled(message as RecordingCancelledMessage)
+          // Handle binary data (recording chunks)
+          if (event.data instanceof ArrayBuffer || Buffer.isBuffer(event.data)) {
+            const buffer = Buffer.isBuffer(event.data) ? event.data : Buffer.from(event.data)
+            const relay = getRecordingRelay(connectionId)
+            if (relay) {
+              relay.handleBinaryData(buffer)
+            }
+            return
           }
-        } else {
-          const extensionEvent = message as ExtensionEventMessage
 
-          if (extensionEvent.method !== 'forwardCDPEvent') {
+          let message: ExtensionMessage
+
+          try {
+            message = JSON.parse(event.data.toString())
+          } catch {
+            ws.close(1000, 'Invalid JSON')
             return
           }
 
-          const { method, params, sessionId } = extensionEvent.params
+          if (message.id !== undefined) {
+            const pending = (() => {
+              let pendingRequest: relayState.ExtensionPendingRequest | null = null
 
-           logCdpJson({
-             timestamp: new Date().toISOString(),
-             direction: 'from-extension',
-             message: { method, params, sessionId },
-           })
+              store.setState((s) => {
+                const extensionEntry = s.extensions.get(connectionId)
+                if (!extensionEntry) {
+                  return s
+                }
 
-          logCdpMessage({
-            direction: 'from-extension',
-            method,
-            sessionId,
-            params
-          })
+                const nextPendingRequest = extensionEntry.pendingRequests.get(message.id)
+                if (!nextPendingRequest) {
+                  return s
+                }
 
-          const cdpEvent: CDPEventBase = { method, sessionId, params }
-          emitter.emit('cdp:event', { event: cdpEvent, sessionId })
-
-          if (method === 'Target.attachedToTarget') {
-            const targetParams = params as Protocol.Target.AttachedToTargetEvent
-            const incomingSessionId = sessionId
-            const iframeParentFrameId = targetParams.targetInfo.parentFrameId
-            const iframeOwnerSessionId = targetParams.targetInfo.type === 'iframe' && iframeParentFrameId
-              ? getPageTargetForFrameId({ connection, frameId: iframeParentFrameId })?.sessionId
-              : undefined
-
-            // Filter out restricted targets (unsupported types, extension pages, chrome:// URLs, etc.)
-            if (isRestrictedTarget(targetParams.targetInfo)) {
-              if (targetParams.waitingForDebugger && targetParams.sessionId) {
-                void sendToExtension({
+                pendingRequest = nextPendingRequest
+                return relayState.removeExtensionPendingRequest(s, {
                   extensionId: connectionId,
-                  method: 'forwardCDPCommand',
-                  params: {
-                    sessionId: targetParams.sessionId,
-                    method: 'Runtime.runIfWaitingForDebugger',
-                    params: {},
-                    source: 'server',
-                  },
-                }).catch((error) => {
-                  const message = error instanceof Error ? error.message : String(error)
-                  logger?.log(pc.yellow('[Server] Failed to resume restricted target:'), message)
+                  requestId: message.id,
                 })
-              }
-              logger?.log(pc.gray(`[Server] Ignoring restricted target: ${targetParams.targetInfo.type} (${targetParams.targetInfo.url})`))
+              })
+
+              return pendingRequest
+            })() as relayState.ExtensionPendingRequest | null
+
+            if (!pending) {
+              logger?.log('Unexpected response with id:', message.id)
               return
             }
 
-            if (!targetParams.targetInfo.url) {
-              logger?.error(pc.red('[Extension] WARNING: Target.attachedToTarget received with empty URL!'), JSON.stringify({ method, params: targetParams, sessionId }))
+            if (message.error) {
+              pending.reject(new Error(message.error))
+            } else {
+              pending.resolve(message.result)
+            }
+          } else if (message.method === 'pong') {
+            // Keep-alive response, nothing to do
+          } else if (message.method === 'log') {
+            const { level, args } = message.params
+            const logFn = (logger as Record<string, unknown>)?.[level] as ((...args: unknown[]) => void) | undefined
+            const logFunc = logFn || logger?.log
+            const prefix = pc.yellow(`[Extension] [${level.toUpperCase()}]`)
+            logFunc?.(prefix, ...args)
+          } else if (message.method === 'recordingData') {
+            const relay = getRecordingRelay(connectionId)
+            if (relay) {
+              relay.handleRecordingData(message as RecordingDataMessage)
+            }
+          } else if (message.method === 'recordingCancelled') {
+            const relay = getRecordingRelay(connectionId)
+            if (relay) {
+              relay.handleRecordingCancelled(message as RecordingCancelledMessage)
+            }
+          } else {
+            const extensionEvent = message as ExtensionEventMessage
+
+            if (extensionEvent.method !== 'forwardCDPEvent') {
+              return
             }
-            logger?.log(pc.yellow('[Extension] Target.attachedToTarget full payload:'), JSON.stringify({ method, params: targetParams, sessionId }))
-
-            // Check if we already sent this target to clients (e.g., from Target.setAutoAttach response)
-            const alreadyConnected = connection.connectedTargets.has(targetParams.sessionId)
-            const existingTarget = connection.connectedTargets.get(targetParams.sessionId)
-
-            // Always update our local state with latest target info
-            connection.connectedTargets.set(targetParams.sessionId, {
-              sessionId: targetParams.sessionId,
-              targetId: targetParams.targetInfo.targetId,
-              targetInfo: targetParams.targetInfo,
-              frameIds: existingTarget?.frameIds ?? new Set()
+
+            const { method, params, sessionId } = extensionEvent.params
+
+            logCdpJson({
+              timestamp: new Date().toISOString(),
+              direction: 'from-extension',
+              message: { method, params, sessionId },
+            })
+
+            logCdpMessage({
+              direction: 'from-extension',
+              method,
+              sessionId,
+              params,
             })
 
-            // Only forward to Playwright if this is a new target to avoid duplicates
-            if (!alreadyConnected) {
+            const cdpEvent: CDPEventBase = { method, sessionId, params }
+            emitter.emit('cdp:event', { event: cdpEvent, sessionId })
+
+            maybeEmitBrowserDownloadCompatEvent({ method, params, extensionId: connectionId })
+
+            if (method === 'Target.attachedToTarget') {
+              const targetParams = params as Protocol.Target.AttachedToTargetEvent
+              const incomingSessionId = sessionId
+              const iframeParentFrameId = targetParams.targetInfo.parentFrameId
+              // Read current extension state for iframe parent lookup
+              const currentExtState = store.getState().extensions.get(connectionId)
+              const iframeOwnerSessionId =
+                targetParams.targetInfo.type === 'iframe' && iframeParentFrameId && currentExtState
+                  ? getPageTargetForFrameId({ extensionState: currentExtState, frameId: iframeParentFrameId })?.sessionId
+                  : undefined
+
+              // Filter out restricted targets (unsupported types, extension pages, chrome:// URLs, etc.)
+              if (isRestrictedTarget(targetParams.targetInfo)) {
+                if (targetParams.waitingForDebugger && targetParams.sessionId) {
+                  void sendToExtension({
+                    extensionId: connectionId,
+                    method: 'forwardCDPCommand',
+                    params: {
+                      sessionId: targetParams.sessionId,
+                      method: 'Runtime.runIfWaitingForDebugger',
+                      params: {},
+                      source: 'server',
+                    },
+                  }).catch((error) => {
+                    const msg = error instanceof Error ? error.message : String(error)
+                    logger?.log(pc.yellow('[Server] Failed to resume restricted target:'), msg)
+                  })
+                }
+                logger?.log(
+                  pc.gray(
+                    `[Server] Ignoring restricted target: ${targetParams.targetInfo.type} (${targetParams.targetInfo.url})`,
+                  ),
+                )
+                return
+              }
+
+              if (!targetParams.targetInfo.url) {
+                logger?.error(
+                  pc.red('[Extension] WARNING: Target.attachedToTarget received with empty URL!'),
+                  JSON.stringify({ method, params: targetParams, sessionId }),
+                )
+              }
+              logger?.log(
+                pc.yellow('[Extension] Target.attachedToTarget full payload:'),
+                JSON.stringify({ method, params: targetParams, sessionId }),
+              )
+
+              // Check if we already sent this target to clients (e.g., from Target.setAutoAttach response)
+              const alreadyConnected = currentExtState?.connectedTargets.has(targetParams.sessionId) ?? false
+
+              // State transition: add/update target
+              store.setState((s) =>
+                relayState.addTarget(s, {
+                  extensionId: connectionId,
+                  sessionId: targetParams.sessionId,
+                  targetId: targetParams.targetInfo.targetId,
+                  targetInfo: targetParams.targetInfo,
+                }),
+              )
+
+              const cachedDownloadBehavior = extensionDownloadBehavior.get(connectionId)
+              if (cachedDownloadBehavior && targetParams.targetInfo.type === 'page') {
+                void applyDownloadBehaviorToTargets({
+                  extensionId: connectionId,
+                  behavior: cachedDownloadBehavior,
+                  targetSessionIds: [targetParams.sessionId],
+                })
+              }
+
+              // Only forward to Playwright if this is a new target to avoid duplicates
+              if (!alreadyConnected) {
+                sendToPlaywright({
+                  message: {
+                    // Iframe targets must be routed to the parent page sessionId so Playwright attaches them under the right page.
+                    // - iframeOwnerSessionId: derived parent session via parentFrameId -> page sessionId (frameId tracking).
+                    // - incomingSessionId: extension event sessionId for the parent tab.
+                    // The frameId mapping is racy: Target.attachedToTarget can arrive before Page.frameAttached/Page.frameNavigated populate frameIds.
+                    // When iframeOwnerSessionId is missing we must fall back to incomingSessionId, otherwise Playwright receives the attach on the root
+                    // session, detaches it, and the iframe stays paused (waitingForDebugger) which can hang navigations.
+                    sessionId: iframeOwnerSessionId ?? incomingSessionId,
+                    method: 'Target.attachedToTarget',
+                    params: targetParams,
+                  } as CDPEventBase,
+                  source: 'extension',
+                  extensionId: connectionId,
+                })
+              }
+            } else if (method === 'Target.detachedFromTarget') {
+              const detachParams = params as Protocol.Target.DetachedFromTargetEvent
+              store.setState((s) =>
+                relayState.removeTarget(s, { extensionId: connectionId, sessionId: detachParams.sessionId }),
+              )
+
               sendToPlaywright({
                 message: {
-                  // Iframe targets must be routed to the parent page sessionId so Playwright attaches them under the right page.
-                  // - iframeOwnerSessionId: derived parent session via parentFrameId -> page sessionId (frameId tracking).
-                  // - incomingSessionId: extension event sessionId for the parent tab.
-                  // The frameId mapping is racy: Target.attachedToTarget can arrive before Page.frameAttached/Page.frameNavigated populate frameIds.
-                  // When iframeOwnerSessionId is missing we must fall back to incomingSessionId, otherwise Playwright receives the attach on the root
-                  // session, detaches it, and the iframe stays paused (waitingForDebugger) which can hang navigations.
-                  sessionId: iframeOwnerSessionId ?? incomingSessionId,
-                  method: 'Target.attachedToTarget',
-                  params: targetParams
+                  method: 'Target.detachedFromTarget',
+                  params: detachParams,
                 } as CDPEventBase,
                 source: 'extension',
                 extensionId: connectionId,
               })
-            }
-          } else if (method === 'Target.detachedFromTarget') {
-            const detachParams = params as Protocol.Target.DetachedFromTargetEvent
-            connection.connectedTargets.delete(detachParams.sessionId)
-
-            sendToPlaywright({
-              message: {
-                method: 'Target.detachedFromTarget',
-                params: detachParams
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else if (method === 'Target.targetCrashed') {
-            const crashParams = params as Protocol.Target.TargetCrashedEvent
-            for (const [sid, target] of connection.connectedTargets.entries()) {
-              if (target.targetId === crashParams.targetId) {
-                connection.connectedTargets.delete(sid)
-                logger?.log(pc.red('[Server] Target crashed, removing:'), crashParams.targetId)
-                break
-              }
-            }
+            } else if (method === 'Target.targetCrashed') {
+              const crashParams = params as Protocol.Target.TargetCrashedEvent
+              store.setState((s) =>
+                relayState.removeTargetByCrash(s, { extensionId: connectionId, targetId: crashParams.targetId }),
+              )
+              logger?.log(pc.red('[Server] Target crashed, removing:'), crashParams.targetId)
 
-            sendToPlaywright({
-              message: {
-                method: 'Target.targetCrashed',
-                params: crashParams
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else if (method === 'Target.targetInfoChanged') {
-            const infoParams = params as Protocol.Target.TargetInfoChangedEvent
-            for (const target of connection.connectedTargets.values()) {
-              if (target.targetId === infoParams.targetInfo.targetId) {
-                target.targetInfo = infoParams.targetInfo
-                break
-              }
-            }
+              sendToPlaywright({
+                message: {
+                  method: 'Target.targetCrashed',
+                  params: crashParams,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            } else if (method === 'Target.targetInfoChanged') {
+              const infoParams = params as Protocol.Target.TargetInfoChangedEvent
+              store.setState((s) =>
+                relayState.updateTargetInfo(s, { extensionId: connectionId, targetInfo: infoParams.targetInfo }),
+              )
 
-            sendToPlaywright({
-              message: {
-                method: 'Target.targetInfoChanged',
-                params: infoParams
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else if (method === 'Page.frameAttached') {
-            const frameParams = params as Protocol.Page.FrameAttachedEvent
-            if (sessionId) {
-              const target = connection.connectedTargets.get(sessionId)
-              if (target) {
-                target.frameIds.add(frameParams.frameId)
+              sendToPlaywright({
+                message: {
+                  method: 'Target.targetInfoChanged',
+                  params: infoParams,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            } else if (method === 'Page.frameAttached') {
+              const frameParams = params as Protocol.Page.FrameAttachedEvent
+              if (sessionId) {
+                store.setState((s) =>
+                  relayState.addFrameId(s, { extensionId: connectionId, sessionId, frameId: frameParams.frameId }),
+                )
               }
-            }
 
-            sendToPlaywright({
-              message: {
-                sessionId,
-                method,
-                params
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else if (method === 'Page.frameDetached') {
-            const frameParams = params as Protocol.Page.FrameDetachedEvent
-            const ownerTarget = getPageTargetForFrameId({ connection, frameId: frameParams.frameId })
-            if (ownerTarget) {
-              ownerTarget.frameIds.delete(frameParams.frameId)
-            }
+              sendToPlaywright({
+                message: {
+                  sessionId,
+                  method,
+                  params,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            } else if (method === 'Page.frameDetached') {
+              const frameParams = params as Protocol.Page.FrameDetachedEvent
+              store.setState((s) =>
+                relayState.removeFrameId(s, { extensionId: connectionId, frameId: frameParams.frameId }),
+              )
 
-            sendToPlaywright({
-              message: {
-                sessionId,
-                method,
-                params
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else if (method === 'Page.frameNavigated') {
-            const frameParams = params as Protocol.Page.FrameNavigatedEvent
-            if (sessionId) {
-              const target = connection.connectedTargets.get(sessionId)
-              if (target) {
-                target.frameIds.add(frameParams.frame.id)
+              sendToPlaywright({
+                message: {
+                  sessionId,
+                  method,
+                  params,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            } else if (method === 'Page.frameNavigated') {
+              const frameParams = params as Protocol.Page.FrameNavigatedEvent
+              if (sessionId) {
+                store.setState((s) =>
+                  relayState.addFrameId(s, { extensionId: connectionId, sessionId, frameId: frameParams.frame.id }),
+                )
               }
-            }
-            if (!frameParams.frame.parentId && sessionId) {
-              const target = connection.connectedTargets.get(sessionId)
-              if (target) {
-                target.targetInfo = {
-                  ...target.targetInfo,
-                  url: frameParams.frame.url,
-                  title: frameParams.frame.name || target.targetInfo.title,
-                }
-                logger?.log(pc.magenta('[Server] Updated target URL from Page.frameNavigated:'), frameParams.frame.url)
+              if (!frameParams.frame.parentId && sessionId) {
+                store.setState((s) =>
+                  relayState.updateTargetUrl(s, {
+                    extensionId: connectionId,
+                    sessionId,
+                    url: frameParams.frame.url,
+                    title: frameParams.frame.name || undefined,
+                  }),
+                )
+                logger?.log(
+                  pc.magenta('[Server] Updated target URL from Page.frameNavigated:'),
+                  frameParams.frame.url,
+                )
               }
-            }
 
-            sendToPlaywright({
-              message: {
-                sessionId,
-                method,
-                params
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else if (method === 'Page.navigatedWithinDocument') {
-            const navParams = params as Protocol.Page.NavigatedWithinDocumentEvent
-            if (sessionId) {
-              const target = connection.connectedTargets.get(sessionId)
-              if (target) {
-                target.targetInfo = {
-                  ...target.targetInfo,
-                  url: navParams.url,
-                }
-                logger?.log(pc.magenta('[Server] Updated target URL from Page.navigatedWithinDocument:'), navParams.url)
+              sendToPlaywright({
+                message: {
+                  sessionId,
+                  method,
+                  params,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            } else if (method === 'Page.navigatedWithinDocument') {
+              const navParams = params as Protocol.Page.NavigatedWithinDocumentEvent
+              if (sessionId) {
+                store.setState((s) =>
+                  relayState.updateTargetUrl(s, { extensionId: connectionId, sessionId, url: navParams.url }),
+                )
+                logger?.log(
+                  pc.magenta('[Server] Updated target URL from Page.navigatedWithinDocument:'),
+                  navParams.url,
+                )
               }
-            }
 
-            sendToPlaywright({
-              message: {
-                sessionId,
-                method,
-                params
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
-          } else {
-            sendToPlaywright({
-              message: {
-                sessionId,
-                method,
-                params
-              } as CDPEventBase,
-              source: 'extension',
-              extensionId: connectionId,
-            })
+              sendToPlaywright({
+                message: {
+                  sessionId,
+                  method,
+                  params,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            } else {
+              sendToPlaywright({
+                message: {
+                  sessionId,
+                  method,
+                  params,
+                } as CDPEventBase,
+                source: 'extension',
+                extensionId: connectionId,
+              })
+            }
           }
-        }
-      },
-
-      onClose(event, ws) {
-        logger?.log(`Extension disconnected: code=${event.code} reason=${event.reason || 'none'} (${connectionId})`)
-        stopExtensionPing(connectionId)
+        },
 
-        // Cancel any active recordings BEFORE removing connection (cancelRecording checks isExtensionConnected)
-        const recordingRelay = recordingRelays.get(connectionId)
-        if (recordingRelay) {
-          recordingRelay.cancelRecording({}).catch(() => {
-            // Ignore errors during cleanup
-          })
-        }
-        recordingRelays.delete(connectionId)
+        onClose(event) {
+          logger?.log(`Extension disconnected: code=${event.code} reason=${event.reason || 'none'} (${connectionId})`)
 
-        const connection = extensionConnections.get(connectionId)
-        if (connection) {
-          for (const pending of connection.pendingRequests.values()) {
-            pending.reject(new Error('Extension connection closed'))
+          // Cancel recordings BEFORE removing extension state (cancelRecording checks isExtensionConnected)
+          const recordingRelay = recordingRelays.get(connectionId)
+          if (recordingRelay) {
+            recordingRelay.cancelRecording({}).catch(() => {
+              // Ignore errors during cleanup
+            })
+          }
+          recordingRelays.delete(connectionId)
+
+          // 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'))
+            }
           }
-          connection.pendingRequests.clear()
-          connection.connectedTargets.clear()
-        }
 
-        if (connection) {
-          const mappedId = extensionKeyIndex.get(connection.stableKey)
-          if (mappedId === connectionId) {
-            extensionKeyIndex.delete(connection.stableKey)
+          const currentRelayState = store.getState()
+          const closingExtension = currentRelayState.extensions.get(connectionId)
+          const successorCandidates = closingExtension
+            ? Array.from(currentRelayState.extensions.values())
+                .reverse()
+                .filter((ext) => {
+                  return ext.id !== connectionId && ext.stableKey === closingExtension.stableKey && Boolean(ext.ws)
+                })
+            : []
+          const successorExtension = closingExtension
+            ? successorCandidates[0]
+            : undefined
+
+          if (successorExtension) {
+            logger?.log(
+              pc.yellow(
+                `Rebinding clients from ${connectionId} to ${successorExtension.id} (stableKey: ${successorExtension.stableKey})`,
+              ),
+            )
+            store.setState((s) => {
+              return relayState.rebindClientsToExtension(s, {
+                fromExtensionId: connectionId,
+                toExtensionId: successorExtension.id,
+              })
+            })
           }
-        }
-        extensionConnections.delete(connectionId)
 
-        for (const [clientId, client] of playwrightClients.entries()) {
-          if (client.extensionId !== connectionId) {
-            continue
+          // Close playwright clients bound to this extension when no successor exists.
+          if (!successorExtension) {
+            const { playwrightClients } = store.getState()
+            for (const client of playwrightClients.values()) {
+              if (client.extensionId === connectionId) {
+                client.ws.close(1000, 'Extension disconnected')
+              }
+            }
           }
-          client.ws.close(1000, 'Extension disconnected')
-          playwrightClients.delete(clientId)
-        }
-      },
 
-      onError(event) {
-        logger?.error('Extension WebSocket error:', event)
+          // State transition: remove extension + its bound clients atomically
+          store.setState((s) => relayState.removeExtension(s, { extensionId: connectionId }))
+        },
+
+        onError(event) {
+          logger?.error('Extension WebSocket error:', event)
+        },
       }
-    }
-  }))
+    }),
+  )
 
   // ============================================================================
   // CLI Execute Endpoints - For stateful code execution via CLI
@@ -1457,7 +1772,10 @@ export async function startPlayWriterCDPRelayServer({
   //    preflight as a fallback, which our CORS policy already blocks.
   // 3. When token mode is enabled (remote access), require the token.
   // ============================================================================
-  const privilegedRouteMiddleware = async (c: Parameters<Parameters<typeof app.use>[1]>[0], next: () => Promise<void>) => {
+  const privilegedRouteMiddleware = async (
+    c: Parameters<Parameters<typeof app.use>[1]>[0],
+    next: () => Promise<void>,
+  ) => {
     // Block cross-origin browser requests via Sec-Fetch-Site header.
     // Browsers always set this forbidden header; it cannot be spoofed.
     // Non-browser clients (Node.js, curl, MCP) don't send it.
@@ -1497,7 +1815,7 @@ export async function startPlayWriterCDPRelayServer({
 
   app.post('/cli/execute', async (c) => {
     try {
-      const body = await c.req.json() as { sessionId: string | number; code: string; timeout?: number }
+      const body = (await c.req.json()) as { sessionId: string | number; code: string; timeout?: number }
       const sessionId = normalizeSessionId(body.sessionId)
       const { code, timeout = 10000 } = body
 
@@ -1508,7 +1826,10 @@ export async function startPlayWriterCDPRelayServer({
       const manager = await getExecutorManager()
       const existingExecutor = manager.getSession(sessionId)
       if (!existingExecutor) {
-        return c.json({ text: `Session ${sessionId} not found. Run 'playwriter session new' first.`, images: [], isError: true }, 404)
+        return c.json(
+          { text: `Session ${sessionId} not found. Run 'playwriter session new' first.`, images: [], isError: true },
+          404,
+        )
       }
       const result = await existingExecutor.execute(code, timeout)
 
@@ -1521,7 +1842,7 @@ export async function startPlayWriterCDPRelayServer({
 
   app.post('/cli/reset', async (c) => {
     try {
-      const body = await c.req.json() as { sessionId: string | number }
+      const body = (await c.req.json()) as { sessionId: string | number }
       const sessionId = normalizeSessionId(body.sessionId)
 
       if (!sessionId) {
@@ -1556,13 +1877,13 @@ export async function startPlayWriterCDPRelayServer({
   })
 
   app.post('/cli/session/new', async (c) => {
-    const body = await c.req.json().catch(() => ({})) as { extensionId?: string | null; cwd?: string }
+    const body = (await c.req.json().catch(() => ({}))) as { extensionId?: string | null; cwd?: string }
     const sessionId = String(nextSessionNumber++)
     const extensionId = body.extensionId || null
     const cwd = body.cwd
-    const allowDefault = !extensionId && extensionConnections.size === 1
-    const extension = getExtensionConnection(extensionId, { allowFallback: allowDefault })
-    if (!extension) {
+    const allowDefault = !extensionId && store.getState().extensions.size === 1
+    const conn = getExtensionConnection(extensionId, { allowFallback: allowDefault })
+    if (!conn) {
       const error = extensionId
         ? `Extension not connected: ${extensionId}`
         : 'Multiple extensions connected. Specify extensionId.'
@@ -1573,9 +1894,9 @@ export async function startPlayWriterCDPRelayServer({
       sessionId,
       cwd,
       sessionMetadata: {
-        extensionId: extension.stableKey,
-        browser: extension.info.browser || null,
-        profile: extension.info ? { email: extension.info.email || '', id: extension.info.id || '' } : null,
+        extensionId: conn.stableKey,
+        browser: conn.info.browser || null,
+        profile: conn.info ? { email: conn.info.email || '', id: conn.info.id || '' } : null,
       },
     })
     const metadata = executor.getSessionMetadata()
@@ -1605,7 +1926,7 @@ export async function startPlayWriterCDPRelayServer({
 
   app.post('/cli/session/delete', async (c) => {
     try {
-      const body = await c.req.json() as { sessionId: string | number }
+      const body = (await c.req.json()) as { sessionId: string | number }
       const sessionId = normalizeSessionId(body.sessionId)
 
       if (!sessionId) {
@@ -1630,73 +1951,64 @@ export async function startPlayWriterCDPRelayServer({
   // ============================================================================
 
   app.post('/recording/start', async (c) => {
-    const body = await c.req.json() as { outputPath?: string; sessionId?: string | number; frameRate?: number; audio?: boolean; videoBitsPerSecond?: number; audioBitsPerSecond?: number }
+    const body = (await c.req.json()) as {
+      outputPath?: string
+      sessionId?: string | number
+      frameRate?: number
+      audio?: boolean
+      videoBitsPerSecond?: number
+      audioBitsPerSecond?: number
+    }
     const sessionId = normalizeSessionId(body.sessionId)
     const { sessionId: _sessionId, ...recordingOptions } = body
-    const manager = await getExecutorManager()
-    const executor = sessionId ? manager.getSession(sessionId) : null
-    if (sessionId && !executor) {
-      return c.json({ success: false, error: `Session ${sessionId} not found` }, 404)
-    }
-    const extensionId = executor?.getSessionMetadata().extensionId || null
+    const { extensionId, sessionId: resolvedSessionId } = await resolveRecordingRoute({ sessionId })
     const relay = getRecordingRelay(extensionId)
     if (!relay) {
       return c.json({ success: false, error: 'Extension not connected' }, 500)
     }
-    const recordingParams = (sessionId ? { ...recordingOptions, sessionId } : recordingOptions) as StartRecordingBody
+    const recordingParams = (resolvedSessionId
+      ? { ...recordingOptions, sessionId: resolvedSessionId }
+      : recordingOptions) as StartRecordingBody
     const result = await relay.startRecording(recordingParams)
-    const status = result.success ? 200 : (result.error?.includes('required') ? 400 : 500)
+    const status = result.success ? 200 : result.error?.includes('required') ? 400 : 500
     return c.json(result, status)
   })
 
   app.post('/recording/stop', async (c) => {
-    const body = await c.req.json() as { sessionId?: string | number }
+    const body = (await c.req.json()) as { sessionId?: string | number }
     const sessionId = normalizeSessionId(body.sessionId)
-    const manager = await getExecutorManager()
-    const executor = sessionId ? manager.getSession(sessionId) : null
-    if (sessionId && !executor) {
-      return c.json({ success: false, error: `Session ${sessionId} not found` }, 404)
-    }
-    const extensionId = executor?.getSessionMetadata().extensionId || null
+    const { extensionId, sessionId: resolvedSessionId } = await resolveRecordingRoute({ sessionId })
     const relay = getRecordingRelay(extensionId)
     if (!relay) {
       return c.json({ success: false, error: 'Extension not connected' }, 500)
     }
-    const stopParams: StopRecordingParams = sessionId ? { sessionId } : {}
+    const stopParams: StopRecordingParams = resolvedSessionId ? { sessionId: resolvedSessionId } : {}
     const result = await relay.stopRecording(stopParams)
-    const status = result.success ? 200 : (result.error?.includes('not found') ? 404 : 500)
+    const status = result.success ? 200 : result.error?.includes('not found') ? 404 : 500
     return c.json(result, status)
   })
 
   app.get('/recording/status', async (c) => {
     const sessionId = normalizeSessionId(c.req.query('sessionId'))
-    const normalizedSessionId = sessionId || undefined
-    const manager = await getExecutorManager()
-    const executor = normalizedSessionId ? manager.getSession(normalizedSessionId) : null
-    const extensionId = executor?.getSessionMetadata().extensionId || null
+    const { extensionId, sessionId: resolvedSessionId } = await resolveRecordingRoute({ sessionId })
     const relay = getRecordingRelay(extensionId)
     if (!relay) {
       return c.json({ isRecording: false })
     }
-    const isRecordingParams: IsRecordingParams = normalizedSessionId ? { sessionId: normalizedSessionId } : {}
+    const isRecordingParams: IsRecordingParams = resolvedSessionId ? { sessionId: resolvedSessionId } : {}
     const result = await relay.isRecording(isRecordingParams)
     return c.json(result)
   })
 
   app.post('/recording/cancel', async (c) => {
-    const body = await c.req.json() as { sessionId?: string | number }
+    const body = (await c.req.json()) as { sessionId?: string | number }
     const sessionId = normalizeSessionId(body.sessionId)
-    const manager = await getExecutorManager()
-    const executor = sessionId ? manager.getSession(sessionId) : null
-    if (sessionId && !executor) {
-      return c.json({ success: false, error: `Session ${sessionId} not found` }, 404)
-    }
-    const extensionId = executor?.getSessionMetadata().extensionId || null
+    const { extensionId, sessionId: resolvedSessionId } = await resolveRecordingRoute({ sessionId })
     const relay = getRecordingRelay(extensionId)
     if (!relay) {
       return c.json({ success: false, error: 'Extension not connected' }, 500)
     }
-    const cancelParams: CancelRecordingParams = sessionId ? { sessionId } : {}
+    const cancelParams: CancelRecordingParams = resolvedSessionId ? { sessionId: resolvedSessionId } : {}
     const result = await relay.cancelRecording(cancelParams)
     return c.json(result)
   })
@@ -1716,14 +2028,24 @@ export async function startPlayWriterCDPRelayServer({
 
   return {
     close() {
+      const { extensions, playwrightClients } = store.getState()
+
       for (const client of playwrightClients.values()) {
         client.ws.close(1000, 'Server stopped')
       }
-      playwrightClients.clear()
-      for (const extension of extensionConnections.values()) {
-        extension.ws.close(1000, 'Server stopped')
+
+      for (const ext of extensions.values()) {
+        if (ext.pingInterval) {
+          clearInterval(ext.pingInterval)
+        }
+        ext.ws?.close(1000, 'Server stopped')
       }
-      extensionConnections.clear()
+
+      // Reset store state
+      store.setState({
+        extensions: new Map(),
+        playwrightClients: new Map(),
+      })
       server.close()
       emitter.removeAllListeners()
     },
@@ -1732,6 +2054,6 @@ export async function startPlayWriterCDPRelayServer({
     },
     off<K extends keyof RelayServerEvents>(event: K, listener: RelayServerEvents[K]) {
       emitter.off(event, listener as (...args: unknown[]) => void)
-    }
+    },
   }
 }