use new extension ws socket

extension/src/background.ts

@@ -266,10 +266,21 @@ async function ensureConnection(): Promise<void> {
   debugLog('WebSocket connected successfully, creating RelayConnection instance')
   const newConnection = new RelayConnection({
     ws: socket,
-    onClose: () => {
-      debugLog('=== Relay connection onClose callback triggered ===')
+    onClose: (reason, code) => {
+      debugLog('=== Relay connection onClose callback triggered ===', { reason, code })
       const { connectedTabs } = useExtensionStore.getState()
       debugLog('Connected tabs before potential reconnection:', Array.from(connectedTabs.keys()))
+
+      if (reason === 'Extension Replaced' || code === 4001) {
+        debugLog('Connection replaced by another extension instance. Not reconnecting.')
+        useExtensionStore.setState({
+          connection: undefined,
+          connectionState: 'error',
+          errorText: 'Disconnected: Replaced by another extension',
+        })
+        return
+      }
+
       useExtensionStore.setState({ connection: undefined, connectionState: 'disconnected' })
 
       if (connectedTabs.size > 0) {

extension/src/relayConnection.ts

@@ -42,12 +42,12 @@ export class RelayConnection {
   private _nextSessionId: number = 1;
   private _ws: WebSocket;
   private _closed = false;
-  private _onCloseCallback?: () => void;
+  private _onCloseCallback?: (reason: string, code: number) => void;
   private _onTabDetachedCallback?: (tabId: number, reason: `${chrome.debugger.DetachReason}`) => void
 
   constructor({ ws, onClose, onTabDetached }: {
     ws: WebSocket;
-    onClose?: () => void;
+    onClose?: (reason: string, code: number) => void;
     onTabDetached?: (tabId: number, reason: `${chrome.debugger.DetachReason}`) => void;
   }) {
     this._ws = ws;
@@ -88,7 +88,7 @@ export class RelayConnection {
         reason: event.reason,
         wasClean: event.wasClean
       });
-      this._onClose();
+      this._onClose(event.reason, event.code);
     };
     this._ws.onerror = (event: Event) => {
       debugLog('WebSocket onerror event:', event);
@@ -191,10 +191,10 @@ export class RelayConnection {
   close(message: string): void {
     debugLog('Closing RelayConnection, reason:', message, 'current state:', this._ws.readyState);
     this._ws.close(1000, message);
-    this._onClose();
+    this._onClose(message, 1000);
   }
 
-  private _onClose() {
+  private _onClose(reason: string = 'Unknown', code: number = 1000) {
     if (this._closed) {
       debugLog('_onClose called but already closed');
       return;
@@ -225,7 +225,7 @@ export class RelayConnection {
     debugLog('All tabs cleared from map. Chrome automation bar should disappear in a few seconds.');
 
     debugLog('Connection closed, calling onClose callback');
-    this._onCloseCallback?.();
+    this._onCloseCallback?.(reason, code);
   }
 
   private _onDebuggerEvent = (source: chrome.debugger.DebuggerSession, method: string, params: any): void => {

playwriter/src/extension/cdp-relay.ts

@@ -335,33 +335,17 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
     return {
       onOpen(_event, ws) {
         if (extensionWs) {
-          logger.log('Rejecting second extension connection')
-          ws.close(1000, 'Another extension connection already established')
-          return
+          logger.log(chalk.yellow('Closing existing extension connection to replace with new one'))
+          extensionWs.close(4001, 'Extension Replaced')
+          
+          // Clear state from the old connection to prevent leaks
+          connectedTargets.clear()
+          for (const pending of extensionPendingRequests.values()) {
+            pending.reject(new Error('Extension connection replaced'))
+          }
+          extensionPendingRequests.clear()
         }
 
-        // logger.log('Extension connecting, resetting server state')
-
-        // // Clear stale targets from previous connections
-        // const staleTargetCount = connectedTargets.size
-        // if (staleTargetCount > 0) {
-        //   logger.log(`Clearing ${staleTargetCount} stale targets from previous connection`)
-        //   connectedTargets.clear()
-        // }
-
-        // // Reject any pending requests from previous connection
-        // const pendingRequestCount = extensionPendingRequests.size
-        // if (pendingRequestCount > 0) {
-        //   logger.log(`Rejecting ${pendingRequestCount} pending requests from previous connection`)
-        //   for (const pending of extensionPendingRequests.values()) {
-        //     pending.reject(new Error('Extension reconnected, request cancelled'))
-        //   }
-        //   extensionPendingRequests.clear()
-        // }
-
-        // // Reset message ID counter
-        // extensionMessageId = 0
-
         extensionWs = ws
         logger.log('Extension connected with clean state')
       },
@@ -471,9 +455,16 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
         }
       },
 
-      onClose() {
+      onClose(event, ws) {
         logger.log('Extension disconnected')
 
+        // If this is an old connection closing after we've already established a new one,
+        // don't clear the global state
+        if (extensionWs && extensionWs !== ws) {
+           logger.log('Old extension connection closed, keeping new one active')
+           return
+        }
+
         for (const pending of extensionPendingRequests.values()) {
           pending.reject(new Error('Extension connection closed'))
         }