From dcdf38f1197359b5cb6c5c202470c46c22e8bdf8 Mon Sep 17 00:00:00 2001
From: Max Schmitt <max@schmitt.mx>
Date: Mon, 26 Jun 2023 22:21:44 +0200
Subject: [PATCH] chore: keep UI Mode running when used with browser mode
 (#23876)

This will keep UI Mode running in browser mode. When launched in normal
persistent context mode, we know when the persistent context closes, so
we can run the project teardown code.

Fixes https://github.com/microsoft/playwright/issues/23801
---
 packages/playwright-core/src/cli/program.ts   | 14 +++++++----
 packages/playwright-core/src/server/index.ts  |  2 +-
 .../src/server/trace/viewer/traceViewer.ts    | 22 ++++--------------
 packages/playwright-test/src/runner/uiMode.ts | 23 ++++++++++---------
 packages/trace-viewer/src/ui/uiModeView.tsx   |  3 ++-
 5 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/packages/playwright-core/src/cli/program.ts b/packages/playwright-core/src/cli/program.ts
index bc9afcb322..80ade0a375 100644
--- a/packages/playwright-core/src/cli/program.ts
+++ b/packages/playwright-core/src/cli/program.ts
@@ -22,7 +22,8 @@ import path from 'path';
 import type { Command } from '../utilsBundle';
 import { program } from '../utilsBundle';
 import { runDriver, runServer, printApiJson, launchBrowserServer } from './driver';
-import { showTraceViewer } from '../server/trace/viewer/traceViewer';
+import type { OpenTraceViewerOptions } from '../server/trace/viewer/traceViewer';
+import { openTraceInBrowser, openTraceViewerApp } from '../server/trace/viewer/traceViewer';
 import * as playwright from '../..';
 import type { BrowserContext } from '../client/browserContext';
 import type { Browser } from '../client/browser';
@@ -297,14 +298,19 @@ program
       if (options.browser === 'wk')
         options.browser = 'webkit';
 
-      const openOptions = {
+      const openOptions: OpenTraceViewerOptions = {
         headless: false,
         host: options.host,
         port: +options.port,
         isServer: !!options.stdin,
-        openInBrowser: options.port !== undefined || options.host !== undefined
       };
-      showTraceViewer(traces, options.browser, openOptions).catch(logErrorAndExit);
+      if (options.port !== undefined || options.host !== undefined) {
+        openTraceInBrowser(traces, openOptions).catch(logErrorAndExit);
+      } else {
+        openTraceViewerApp(traces, options.browser, openOptions).then(page => {
+          page.on('close', () => process.exit(0));
+        }).catch(logErrorAndExit);
+      }
     }).addHelpText('afterAll', `
 Examples:
 
diff --git a/packages/playwright-core/src/server/index.ts b/packages/playwright-core/src/server/index.ts
index 05b4c4e590..602656aeef 100644
--- a/packages/playwright-core/src/server/index.ts
+++ b/packages/playwright-core/src/server/index.ts
@@ -29,5 +29,5 @@ export { createPlaywright } from './playwright';
 
 export type { DispatcherScope } from './dispatchers/dispatcher';
 export type { Playwright } from './playwright';
-export { showTraceViewer, openTraceViewerApp } from './trace/viewer/traceViewer';
+export { openTraceInBrowser, openTraceViewerApp } from './trace/viewer/traceViewer';
 export { serverSideCallMetadata } from './instrumentation';
\ No newline at end of file
diff --git a/packages/playwright-core/src/server/trace/viewer/traceViewer.ts b/packages/playwright-core/src/server/trace/viewer/traceViewer.ts
index a678522199..07c53db2fc 100644
--- a/packages/playwright-core/src/server/trace/viewer/traceViewer.ts
+++ b/packages/playwright-core/src/server/trace/viewer/traceViewer.ts
@@ -33,25 +33,16 @@ export type Transport = {
   onclose: () => void;
 };
 
-type Options = {
+export type OpenTraceViewerOptions = {
   app?: string;
   headless?: boolean;
   host?: string;
   port?: number;
   isServer?: boolean;
-  openInBrowser?: boolean;
   transport?: Transport;
 };
 
-export async function showTraceViewer(traceUrls: string[], browserName: string, options?: Options): Promise<void> {
-  if (options?.openInBrowser) {
-    await openTraceInBrowser(traceUrls, options);
-    return;
-  }
-  await openTraceViewerApp(traceUrls, browserName, options);
-}
-
-async function startTraceViewerServer(traceUrls: string[], options?: Options): Promise<{ server: HttpServer, url: string }> {
+async function startTraceViewerServer(traceUrls: string[], options?: OpenTraceViewerOptions): Promise<{ server: HttpServer, url: string }> {
   for (const traceUrl of traceUrls) {
     let traceFile = traceUrl;
     // If .json is requested, we'll synthesize it.
@@ -134,7 +125,7 @@ async function startTraceViewerServer(traceUrls: string[], options?: Options): P
   return { server, url };
 }
 
-export async function openTraceViewerApp(traceUrls: string[], browserName: string, options?: Options): Promise<Page> {
+export async function openTraceViewerApp(traceUrls: string[], browserName: string, options?: OpenTraceViewerOptions): Promise<Page> {
   const { url } = await startTraceViewerServer(traceUrls, options);
   const traceViewerPlaywright = createPlaywright({ sdkLanguage: 'javascript', isInternalPlaywright: true });
   const traceViewerBrowser = isUnderTest() ? 'chromium' : browserName;
@@ -171,24 +162,21 @@ export async function openTraceViewerApp(traceUrls: string[], browserName: strin
 
   if (isUnderTest())
     page.on('close', () => context.close(serverSideCallMetadata()).catch(() => {}));
-  else
-    page.on('close', () => process.exit());
 
   await page.mainFrame().goto(serverSideCallMetadata(), url);
   return page;
 }
 
-async function openTraceInBrowser(traceUrls: string[], options?: Options) {
+export async function openTraceInBrowser(traceUrls: string[], options?: OpenTraceViewerOptions) {
   const { url } = await startTraceViewerServer(traceUrls, options);
   // eslint-disable-next-line no-console
   console.log('\nListening on ' + url);
-  await open(url, { wait: true }).catch(() => {});
+  await open(url).catch(() => {});
 }
 
 class StdinServer implements Transport {
   private _pollTimer: NodeJS.Timeout | undefined;
   private _traceUrl: string | undefined;
-  private _page: Page | undefined;
 
   constructor() {
     process.stdin.on('data', data => {
diff --git a/packages/playwright-test/src/runner/uiMode.ts b/packages/playwright-test/src/runner/uiMode.ts
index a07738edb9..88b5648651 100644
--- a/packages/playwright-test/src/runner/uiMode.ts
+++ b/packages/playwright-test/src/runner/uiMode.ts
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { showTraceViewer } from 'playwright-core/lib/server';
+import { openTraceViewerApp, openTraceInBrowser } from 'playwright-core/lib/server';
 import { isUnderTest, ManualPromise } from 'playwright-core/lib/utils';
 import type { FullResult } from '../../reporter';
 import { clearCompilationCache, collectAffectedTestFiles, dependenciesForTestFile } from '../transform/compilationCache';
@@ -27,7 +27,7 @@ import { chokidar } from '../utilsBundle';
 import type { FSWatcher } from 'chokidar';
 import { open } from 'playwright-core/lib/utilsBundle';
 import ListReporter from '../reporters/list';
-import type { Transport } from 'playwright-core/lib/server/trace/viewer/traceViewer';
+import type { OpenTraceViewerOptions, Transport } from 'playwright-core/lib/server/trace/viewer/traceViewer';
 
 class UIMode {
   private _config: FullConfigInternal;
@@ -82,7 +82,6 @@ class UIMode {
   }
 
   async showUI(options: { host?: string, port?: number }) {
-    const exitPromise = new ManualPromise();
     let queue = Promise.resolve();
 
     this._transport = {
@@ -90,10 +89,6 @@ class UIMode {
         if (method === 'ping')
           return;
 
-        if (method === 'exit') {
-          exitPromise.resolve();
-          return;
-        }
         if (method === 'watch') {
           this._watchFiles(params.fileNames);
           return;
@@ -117,16 +112,22 @@ class UIMode {
         await queue;
       },
 
-      onclose: () => exitPromise.resolve(),
+      onclose: () => { },
     };
-    await showTraceViewer([], 'chromium', {
+    const openOptions: OpenTraceViewerOptions = {
       app: 'uiMode.html',
       headless: isUnderTest() && process.env.PWTEST_HEADED_FOR_TEST !== '1',
       transport: this._transport,
       host: options.host,
       port: options.port,
-      openInBrowser: options.host !== undefined || options.port !== undefined,
-    });
+    };
+    const exitPromise = new ManualPromise<void>();
+    if (options.host !== undefined || options.port !== undefined) {
+      await openTraceInBrowser([], openOptions);
+    } else {
+      const page = await openTraceViewerApp([], 'chromium', openOptions);
+      page.on('close', () => exitPromise.resolve());
+    }
 
     if (!process.env.PWTEST_DEBUG) {
       process.stdout.write = (chunk: string | Buffer) => {
diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx
index 5008642e18..4fd760d687 100644
--- a/packages/trace-viewer/src/ui/uiModeView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeView.tsx
@@ -168,7 +168,8 @@ export const UIModeView: React.FC<{}> = ({
 
   return <div className='vbox ui-mode'>
     {isDisconnected && <div className='drop-target'>
-      <div className='title'>Process disconnected</div>
+      <div className='title'>UI Mode disconnected</div>
+      <div><a href='#' onClick={() => window.location.reload()}>Reload the page</a> to reconnect</div>
     </div>}
     <SplitView sidebarSize={250} orientation='horizontal' sidebarIsFirst={true}>
       <div className='vbox'>