From 1d744e7995f2c9d2d88b115d1af990514bdd5c01 Mon Sep 17 00:00:00 2001
From: "Tommy D. Rossi" <beats.by.morse@gmail.com>
Date: Sat, 24 Jan 2026 14:23:50 +0100
Subject: [PATCH] fix CLI relay server startup and extension reconnection

- detect source vs compiled via __filename.endsWith('.ts') instead of env var
- add waitForExtension() to poll for extension connection after server start
- CLI waits up to 10s for extension to reconnect before executing commands
---
 playwriter/src/cli.ts          |  7 ++++-
 playwriter/src/relay-client.ts | 56 +++++++++++++++++++++++++++++++---
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/playwriter/src/cli.ts b/playwriter/src/cli.ts
index acf9fcd..cef9ba8 100644
--- a/playwriter/src/cli.ts
+++ b/playwriter/src/cli.ts
@@ -2,7 +2,7 @@
 
 import { cac } from '@xmorse/cac'
 import { VERSION, LOG_FILE_PATH } from './utils.js'
-import { ensureRelayServer, RELAY_PORT } from './relay-client.js'
+import { ensureRelayServer, RELAY_PORT, waitForExtension } from './relay-client.js'
 
 const cliRelayEnv = { PLAYWRITER_AUTO_ENABLE: '1' }
 
@@ -57,6 +57,11 @@ async function executeCode(options: {
   // Ensure relay server is running (only for local)
   if (!host && !process.env.PLAYWRITER_HOST) {
     await ensureRelayServer({ logger: console, env: cliRelayEnv })
+    // Wait for extension to reconnect after server (re)start
+    const connected = await waitForExtension({ logger: console, timeoutMs: 10000 })
+    if (!connected) {
+      console.error('Warning: Extension not connected. Commands may fail.')
+    }
   }
 
   // Session is required
diff --git a/playwriter/src/relay-client.ts b/playwriter/src/relay-client.ts
index 4df310f..077cc59 100644
--- a/playwriter/src/relay-client.ts
+++ b/playwriter/src/relay-client.ts
@@ -29,6 +29,47 @@ export async function getRelayServerVersion(port: number = RELAY_PORT): Promise<
   }
 }
 
+export async function getExtensionStatus(port: number = RELAY_PORT): Promise<{ connected: boolean; activeTargets: number } | null> {
+  try {
+    const response = await fetch(`http://127.0.0.1:${port}/extension/status`, {
+      signal: AbortSignal.timeout(500),
+    })
+    if (!response.ok) {
+      return null
+    }
+    return await response.json() as { connected: boolean; activeTargets: number }
+  } catch {
+    return null
+  }
+}
+
+/**
+ * Wait for the extension to connect to the relay server.
+ * Returns true if connected within timeout, false otherwise.
+ */
+export async function waitForExtension(options: {
+  port?: number
+  timeoutMs?: number
+  logger?: { log: (...args: any[]) => void }
+} = {}): Promise<boolean> {
+  const { port = RELAY_PORT, timeoutMs = 5000, logger } = options
+  const startTime = Date.now()
+
+  logger?.log('Waiting for extension to connect...')
+
+  while (Date.now() - startTime < timeoutMs) {
+    const status = await getExtensionStatus(port)
+    if (status?.connected) {
+      logger?.log('Extension connected')
+      return true
+    }
+    await sleep(200)
+  }
+
+  logger?.log('Extension did not connect within timeout')
+  return false
+}
+
 async function killRelayServer(port: number): Promise<void> {
   try {
     await killPortProcess(port)
@@ -95,12 +136,14 @@ export async function ensureRelayServer(options: EnsureRelayServerOptions = {}):
     logger?.log('CDP relay server not running, starting it...')
   }
 
-  const dev = process.env.PLAYWRITER_NODE_ENV === 'development'
-  const scriptPath = dev
-    ? path.resolve(__dirname, '../src/start-relay-server.ts')
+  // Detect if we're running from source (.ts) or compiled (.js)
+  // This handles: tsx, vite-node, ts-node, or direct node on compiled output
+  const isRunningFromSource = __filename.endsWith('.ts')
+  const scriptPath = isRunningFromSource
+    ? path.resolve(__dirname, './start-relay-server.ts')
     : path.resolve(__dirname, './start-relay-server.js')
 
-  const serverProcess = spawn(dev ? 'tsx' : process.execPath, [scriptPath], {
+  const serverProcess = spawn(isRunningFromSource ? 'tsx' : process.execPath, [scriptPath], {
     detached: true,
     stdio: 'ignore',
     env: { ...process.env, ...additionalEnv },
@@ -108,14 +151,19 @@ export async function ensureRelayServer(options: EnsureRelayServerOptions = {}):
 
   serverProcess.unref()
 
+
   for (let i = 0; i < 10; i++) {
     await sleep(500)
     const newVersion = await getRelayServerVersion(RELAY_PORT)
     if (newVersion) {
       logger?.log('CDP relay server started successfully')
+      await sleep(1000)
       return
     }
   }
 
+
+
+
   throw new Error(`Failed to start CDP relay server after 5 seconds. Check logs at: ${LOG_FILE_PATH}`)
 }