fix: finish all artifacts when browser exits (#6151)

src/dispatchers/browserDispatcher.ts

@@ -45,6 +45,10 @@ export class BrowserDispatcher extends Dispatcher<Browser, channels.BrowserIniti
     await this._object.close();
   }
 
+  async killForTests(): Promise<void> {
+    await this._object.killForTests();
+  }
+
   async newBrowserCDPSession(): Promise<channels.BrowserNewBrowserCDPSessionResult> {
     if (!this._object.options.isChromium)
       throw new Error(`CDP session is only available in Chromium`);

src/protocol/channels.ts

@@ -429,6 +429,7 @@ export type BrowserInitializer = {
 export interface BrowserChannel extends Channel {
   on(event: 'close', callback: (params: BrowserCloseEvent) => void): this;
   close(params?: BrowserCloseParams, metadata?: Metadata): Promise<BrowserCloseResult>;
+  killForTests(params?: BrowserKillForTestsParams, metadata?: Metadata): Promise<BrowserKillForTestsResult>;
   newContext(params: BrowserNewContextParams, metadata?: Metadata): Promise<BrowserNewContextResult>;
   newBrowserCDPSession(params?: BrowserNewBrowserCDPSessionParams, metadata?: Metadata): Promise<BrowserNewBrowserCDPSessionResult>;
   startTracing(params: BrowserStartTracingParams, metadata?: Metadata): Promise<BrowserStartTracingResult>;
@@ -438,6 +439,9 @@ export type BrowserCloseEvent = {};
 export type BrowserCloseParams = {};
 export type BrowserCloseOptions = {};
 export type BrowserCloseResult = void;
+export type BrowserKillForTestsParams = {};
+export type BrowserKillForTestsOptions = {};
+export type BrowserKillForTestsResult = void;
 export type BrowserNewContextParams = {
   sdkLanguage: string,
   noDefaultViewport?: boolean,

src/protocol/protocol.yml

@@ -440,6 +440,8 @@ Browser:
 
     close:
 
+    killForTests:
+
     newContext:
       parameters:
         $mixin: ContextOptions

src/protocol/validator.ts

@@ -254,6 +254,7 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     timeout: tOptional(tNumber),
   });
   scheme.BrowserCloseParams = tOptional(tObject({}));
+  scheme.BrowserKillForTestsParams = tOptional(tObject({}));
   scheme.BrowserNewContextParams = tObject({
     sdkLanguage: tString,
     noDefaultViewport: tOptional(tBoolean),

src/server/browser.ts

@@ -24,6 +24,7 @@ import { RecentLogsCollector } from '../utils/debugLogger';
 import * as registry from '../utils/registry';
 import { SdkObject } from './instrumentation';
 import { Artifact } from './artifact';
+import { kBrowserClosedError } from '../utils/errors';
 
 export interface BrowserProcess {
   onclose?: ((exitCode: number | null, signal: string | null) => void);
@@ -117,6 +118,9 @@ export abstract class Browser extends SdkObject {
       context._browserClosed();
     if (this._defaultContext)
       this._defaultContext._browserClosed();
+    for (const video of this._idToVideo.values())
+      video.artifact.reportFinished(kBrowserClosedError);
+    this._idToVideo.clear();
     this.emit(Browser.Events.Disconnected);
   }
 
@@ -128,4 +132,8 @@ export abstract class Browser extends SdkObject {
     if (this.isConnected())
       await new Promise(x => this.once(Browser.Events.Disconnected, x));
   }
+
+  async killForTests() {
+    await this.options.browserProcess.kill();
+  }
 }

src/server/browserContext.ts

@@ -92,6 +92,7 @@ export abstract class BrowserContext extends SdkObject {
       return;
     }
     this._closedStatus = 'closed';
+    this._deleteAllDownloads();
     this._downloads.clear();
     this._closePromiseFulfill!(new Error('Context closed'));
     this.emit(BrowserContext.Events.Close);
@@ -221,6 +222,10 @@ export abstract class BrowserContext extends SdkObject {
     return this._closedStatus !== 'open';
   }
 
+  private async _deleteAllDownloads(): Promise<void> {
+    await Promise.all(Array.from(this._downloads).map(download => download.artifact.deleteOnContextClose()));
+  }
+
   async close(metadata: CallMetadata) {
     if (this._closedStatus === 'open') {
       this.emit(BrowserContext.Events.BeforeClose);
@@ -244,11 +249,9 @@ export abstract class BrowserContext extends SdkObject {
         if (context === this)
           promises.push(artifact.finishedPromise());
       }
-      for (const download of this._downloads) {
-        // We delete downloads after context closure
-        // so that browser does not write to the download file anymore.
-        promises.push(download.artifact.deleteOnContextClose());
-      }
+      // We delete downloads after context closure
+      // so that browser does not write to the download file anymore.
+      promises.push(this._deleteAllDownloads());
       await Promise.all(promises);
 
       // Persistent context should also close the browser.

src/server/webkit/wkConnection.ts

@@ -23,6 +23,7 @@ import { rewriteErrorMessage } from '../../utils/stackTrace';
 import { debugLogger, RecentLogsCollector } from '../../utils/debugLogger';
 import { ProtocolLogger } from '../types';
 import { helper } from '../helper';
+import { kBrowserClosedError } from '../../utils/errors';
 
 // WKPlaywright uses this special id to issue Browser.close command which we
 // should ignore.
@@ -49,7 +50,7 @@ export class WKConnection {
     this._onDisconnect = onDisconnect;
     this._protocolLogger = protocolLogger;
     this._browserLogsCollector = browserLogsCollector;
-    this.browserSession = new WKSession(this, '', 'Browser has been closed.', (message: any) => {
+    this.browserSession = new WKSession(this, '', kBrowserClosedError, (message: any) => {
       this.rawSend(message);
     });
   }

tests/download.spec.ts

@@ -407,4 +407,30 @@ it.describe('download event', () => {
     expect(downloadPath).toBe(null);
     expect(saveError.message).toContain('File deleted upon browser context closure.');
   });
+
+  it('should throw if browser dies', async ({ server, browserType, browserName, browserOptions, platform}, testInfo) => {
+    it.skip(browserName === 'webkit' && platform === 'linux', 'WebKit on linux does not convert to the download immediately upon receiving headers');
+    server.setRoute('/downloadStall', (req, res) => {
+      res.setHeader('Content-Type', 'application/octet-stream');
+      res.setHeader('Content-Disposition', 'attachment; filename=file.txt');
+      res.writeHead(200);
+      res.flushHeaders();
+      res.write(`Hello world`);
+    });
+
+    const browser = await browserType.launch(browserOptions);
+    const page = await browser.newPage({ acceptDownloads: true });
+    await page.setContent(`<a href="${server.PREFIX}/downloadStall">click me</a>`);
+    const [download] = await Promise.all([
+      page.waitForEvent('download'),
+      page.click('a')
+    ]);
+    const [downloadPath, saveError] = await Promise.all([
+      download.path(),
+      download.saveAs(testInfo.outputPath('download.txt')).catch(e => e),
+      (browser as any)._channel.killForTests(),
+    ]);
+    expect(downloadPath).toBe(null);
+    expect(saveError.message).toContain('File deleted upon browser context closure.');
+  });
 });

tests/screencast.spec.ts

@@ -591,4 +591,26 @@ it.describe('screencast', () => {
     const saveResult = await page.video().saveAs(file).catch(e => e);
     expect(saveResult.message).toContain('browser has been closed');
   });
+
+  it('should throw if browser dies', async ({browserType, browserOptions, contextOptions}, testInfo) => {
+    const size = { width: 320, height: 240 };
+    const browser = await browserType.launch(browserOptions);
+
+    const context = await browser.newContext({
+      ...contextOptions,
+      recordVideo: {
+        dir: testInfo.outputPath(''),
+        size,
+      },
+      viewport: size,
+    });
+
+    const page = await context.newPage();
+    await new Promise(r => setTimeout(r, 1000));
+    await (browser as any)._channel.killForTests();
+
+    const file = testInfo.outputPath('saved-video-');
+    const saveResult = await page.video().saveAs(file).catch(e => e);
+    expect(saveResult.message).toContain('rowser has been closed');
+  });
 });