chore: move agent cache management into ArtifactsRecorder (#38518)

Author: pavelfeldman

packages/playwright-core/src/client/browser.ts

@@ -74,11 +74,9 @@ export class Browser extends ChannelOwner<channels.BrowserChannel> implements ap
     await this._channel.disconnectFromReusedContext({ reason });
   }
 
-  async _innerNewContext(options: BrowserContextOptions = {}, forReuse: boolean): Promise<BrowserContext> {
-    options = this._browserType._playwright.selectors._withSelectorOptions({
-      ...this._browserType._playwright._defaultContextOptions,
-      ...options,
-    });
+  async _innerNewContext(userOptions: BrowserContextOptions = {}, forReuse: boolean): Promise<BrowserContext> {
+    const options = this._browserType._playwright.selectors._withSelectorOptions(userOptions);
+    await this._instrumentation.runBeforeCreateBrowserContext(options);
     const contextOptions = await prepareBrowserContextParams(this._platform, options);
     const response = forReuse ? await this._channel.newContextForReuse(contextOptions) : await this._channel.newContext(contextOptions);
     const context = BrowserContext.from(response.context);

packages/playwright-core/src/client/browserType.ts

@@ -91,13 +91,14 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel> imple
   }
 
   async launchPersistentContext(userDataDir: string, options: LaunchPersistentContextOptions = {}): Promise<BrowserContext> {
-    const logger = options.logger || this._playwright._defaultLaunchOptions?.logger;
     assert(!(options as any).port, 'Cannot specify a port without launching as a server.');
     options = this._playwright.selectors._withSelectorOptions({
       ...this._playwright._defaultLaunchOptions,
-      ...this._playwright._defaultContextOptions,
       ...options,
     });
+    await this._instrumentation.runBeforeCreateBrowserContext(options);
+
+    const logger = options.logger || this._playwright._defaultLaunchOptions?.logger;
     const contextParams = await prepareBrowserContextParams(this._platform, options);
     const persistentParams: channels.BrowserTypeLaunchPersistentContextParams = {
       ...contextParams,

packages/playwright-core/src/client/clientInstrumentation.ts

@@ -15,9 +15,10 @@
  */
 
 import type { BrowserContext } from './browserContext';
-import type { APIRequestContext } from './fetch';
+import type { APIRequestContext, NewContextOptions } from './fetch';
 import type { StackFrame } from '@protocol/channels';
 import type { Page } from './page';
+import type { BrowserContextOptions } from './types';
 
 // Instrumentation can mutate the data, for example change apiName or stepId.
 export interface ApiCallData {
@@ -38,6 +39,8 @@ export interface ClientInstrumentation {
   onWillPause(options: { keepTestTimeout: boolean }): void;
   onPage(page: Page): void;
 
+  runBeforeCreateBrowserContext(options: BrowserContextOptions): Promise<void>;
+  runBeforeCreateRequestContext(options: NewContextOptions): Promise<void>;
   runAfterCreateBrowserContext(context: BrowserContext): Promise<void>;
   runAfterCreateRequestContext(context: APIRequestContext): Promise<void>;
   runBeforeCloseBrowserContext(context: BrowserContext): Promise<void>;
@@ -49,6 +52,8 @@ export interface ClientInstrumentationListener {
   onApiCallEnd?(apiCall: ApiCallData): void;
   onWillPause?(options: { keepTestTimeout: boolean }): void;
   onPage?(page: Page): void;
+  runBeforeCreateBrowserContext?(options: BrowserContextOptions): Promise<void>;
+  runBeforeCreateRequestContext?(options: NewContextOptions): Promise<void>;
   runAfterCreateBrowserContext?(context: BrowserContext): Promise<void>;
   runAfterCreateRequestContext?(context: APIRequestContext): Promise<void>;
   runBeforeCloseBrowserContext?(context: BrowserContext): Promise<void>;

packages/playwright-core/src/client/fetch.ts

@@ -48,7 +48,7 @@ export type FetchOptions = {
   maxRetries?: number,
 };
 
-type NewContextOptions = Omit<channels.PlaywrightNewRequestOptions, 'extraHTTPHeaders' | 'clientCertificates' | 'storageState' | 'tracesDir'> & {
+export type NewContextOptions = Omit<channels.PlaywrightNewRequestOptions, 'extraHTTPHeaders' | 'clientCertificates' | 'storageState' | 'tracesDir'> & {
   extraHTTPHeaders?: Headers,
   storageState?: string | SetStorageState,
   clientCertificates?: ClientCertificate[];
@@ -65,10 +65,8 @@ export class APIRequest implements api.APIRequest {
   }
 
   async newContext(options: NewContextOptions & TimeoutOptions = {}): Promise<APIRequestContext> {
-    options = {
-      ...this._playwright._defaultContextOptions,
-      ...options,
-    };
+    options = { ...options };
+    await this._playwright._instrumentation.runBeforeCreateRequestContext(options);
     const storageState = typeof options.storageState === 'string' ?
       JSON.parse(await this._playwright._platform.fs().promises.readFile(options.storageState, 'utf8')) :
       options.storageState;

packages/playwright-core/src/client/playwright.ts

@@ -24,7 +24,7 @@ import { APIRequest } from './fetch';
 import { Selectors } from './selectors';
 
 import type * as channels from '@protocol/channels';
-import type { BrowserContextOptions, LaunchOptions } from 'playwright-core';
+import type { LaunchOptions } from 'playwright-core';
 
 export class Playwright extends ChannelOwner<channels.PlaywrightChannel> {
   readonly _android: Android;
@@ -39,7 +39,6 @@ export class Playwright extends ChannelOwner<channels.PlaywrightChannel> {
 
   // Instrumentation.
   _defaultLaunchOptions?: LaunchOptions;
-  _defaultContextOptions?: BrowserContextOptions;
   _defaultContextTimeout?: number;
   _defaultContextNavigationTimeout?: number;
 

packages/playwright/src/index.ts

@@ -31,7 +31,7 @@ import type { ClientInstrumentationListener } from '../../playwright-core/src/cl
 import type { Playwright as PlaywrightImpl } from '../../playwright-core/src/client/playwright';
 import type { Browser as BrowserImpl } from '../../playwright-core/src/client/browser';
 import type { BrowserContext as BrowserContextImpl } from '../../playwright-core/src/client/browserContext';
-import type { APIRequestContext as APIRequestContextImpl } from '../../playwright-core/src/client/fetch';
+import type { APIRequestContext as APIRequestContextImpl, NewContextOptions as APIRequestContextOptions } from '../../playwright-core/src/client/fetch';
 import type { ChannelOwner } from '../../playwright-core/src/client/channelOwner';
 import type { Page as PageImpl } from '../../playwright-core/src/client/page';
 import type { BrowserContext, BrowserContextOptions, LaunchOptions, Page, Tracing } from 'playwright-core';
@@ -228,42 +228,33 @@ const playwrightFixtures: Fixtures<TestFixtures, WorkerFixtures> = ({
     if (serviceWorkers !== undefined)
       options.serviceWorkers = serviceWorkers;
 
-    const workerFile = await agentCacheWorkerFile(agent, testInfo as TestInfoImpl);
-    if (agent && workerFile)
-      options.agent = { ...agent, cacheFile: workerFile };
-
     await use({
       ...contextOptions,
       ...options,
     });
-
-    if (testInfo.status === 'passed' && workerFile)
-      await (testInfo as TestInfoImpl)._upstreamStorage(workerFile);
   }, { box: true }],
 
-  _setupContextOptions: [async ({ playwright, _combinedContextOptions, actionTimeout, navigationTimeout, testIdAttribute }, use, testInfo) => {
+  _setupContextOptions: [async ({ playwright, actionTimeout, navigationTimeout, testIdAttribute }, use, testInfo) => {
     if (testIdAttribute)
       playwrightLibrary.selectors.setTestIdAttribute(testIdAttribute);
     testInfo.snapshotSuffix = process.platform;
     if (debugMode() === 'inspector')
       (testInfo as TestInfoImpl)._setDebugMode();
 
-    playwright._defaultContextOptions = _combinedContextOptions;
     playwright._defaultContextTimeout = actionTimeout || 0;
     playwright._defaultContextNavigationTimeout = navigationTimeout || 0;
     await use();
-    playwright._defaultContextOptions = undefined;
     playwright._defaultContextTimeout = undefined;
     playwright._defaultContextNavigationTimeout = undefined;
   }, { auto: 'all-hooks-included',  title: 'context configuration', box: true } as any],
 
-  _setupArtifacts: [async ({ playwright, screenshot }, use, testInfo) => {
+  _setupArtifacts: [async ({ playwright, screenshot, _combinedContextOptions, agent }, use, testInfo) => {
     // This fixture has a separate zero-timeout slot to ensure that artifact collection
     // happens even after some fixtures or hooks time out.
     // Now that default test timeout is known, we can replace zero with an actual value.
     testInfo.setTimeout(testInfo.project.timeout);
 
-    const artifactsRecorder = new ArtifactsRecorder(playwright, tracing().artifactsDir(), screenshot);
+    const artifactsRecorder = new ArtifactsRecorder(playwright, tracing().artifactsDir(), screenshot, agent);
     await artifactsRecorder.willStartTest(testInfo as TestInfoImpl);
 
     const tracingGroupSteps: TestStepInternal[] = [];
@@ -317,20 +308,33 @@ const playwrightFixtures: Fixtures<TestFixtures, WorkerFixtures> = ({
         if (!keepTestTimeout)
           currentTestInfo()?._setDebugMode();
       },
+      runBeforeCreateBrowserContext: async (options: BrowserContextOptions) => {
+        for (const [key, value] of Object.entries(_combinedContextOptions)) {
+          if (!(key in options))
+            options[key as keyof BrowserContextOptions] = value;
+        }
+        await artifactsRecorder.willCreateBrowserContext(options);
+      },
+      runBeforeCreateRequestContext: async (options: APIRequestContextOptions) => {
+        for (const [key, value] of Object.entries(_combinedContextOptions)) {
+          if (!(key in options))
+            options[key as keyof APIRequestContextOptions] = value;
+        }
+      },
       runAfterCreateBrowserContext: async (context: BrowserContextImpl) => {
-        await artifactsRecorder?.didCreateBrowserContext(context);
+        await artifactsRecorder.didCreateBrowserContext(context);
         const testInfo = currentTestInfo();
         if (testInfo)
           attachConnectedHeaderIfNeeded(testInfo, context.browser());
       },
       runAfterCreateRequestContext: async (context: APIRequestContextImpl) => {
-        await artifactsRecorder?.didCreateRequestContext(context);
+        await artifactsRecorder.didCreateRequestContext(context);
       },
       runBeforeCloseBrowserContext: async (context: BrowserContextImpl) => {
-        await artifactsRecorder?.willCloseBrowserContext(context);
+        await artifactsRecorder.willCloseBrowserContext(context);
       },
       runBeforeCloseRequestContext: async (context: APIRequestContextImpl) => {
-        await artifactsRecorder?.willCloseRequestContext(context);
+        await artifactsRecorder.willCloseRequestContext(context);
       },
     };
 
@@ -643,10 +647,12 @@ class ArtifactsRecorder {
 
   private _screenshotRecorder: SnapshotRecorder;
   private _pageSnapshot: string | undefined;
+  private _agent: PlaywrightTestOptions['agent'];
 
-  constructor(playwright: PlaywrightImpl, artifactsDir: string, screenshot: ScreenshotOption) {
+  constructor(playwright: PlaywrightImpl, artifactsDir: string, screenshot: ScreenshotOption, agent: PlaywrightTestOptions['agent']) {
     this._playwright = playwright;
     this._artifactsDir = artifactsDir;
+    this._agent = agent;
     const screenshotOptions = typeof screenshot === 'string' ? undefined : screenshot;
     this._startedCollectingArtifacts = Symbol('startedCollectingArtifacts');
 
@@ -673,10 +679,15 @@ class ArtifactsRecorder {
     await this._startTraceChunkOnContextCreation(context, context.tracing);
   }
 
+  async willCreateBrowserContext(options: BrowserContextOptions) {
+    await this._cloneAgentCache(options);
+  }
+
   async willCloseBrowserContext(context: BrowserContextImpl) {
     await this._stopTracing(context, context.tracing);
     await this._screenshotRecorder.captureTemporary(context);
     await this._takePageSnapshot(context);
+    await this._upstreamAgentCache(context);
   }
 
   private async _takePageSnapshot(context: BrowserContextImpl) {
@@ -698,6 +709,24 @@ class ArtifactsRecorder {
     } catch {}
   }
 
+  private async _cloneAgentCache(options: BrowserContextOptions) {
+    if (!this._agent || this._agent.cacheMode === 'ignore')
+      return;
+    if (!this._agent.cacheFile && !this._agent.cachePathTemplate)
+      return;
+
+    const cacheFile = this._agent.cacheFile ?? this._testInfo._applyPathTemplate(this._agent.cachePathTemplate!, 'cache', '.json');
+    const workerFile = await this._testInfo._cloneStorage(cacheFile);
+    if (this._agent && workerFile)
+      options.agent = { ...this._agent, cacheFile: workerFile };
+  }
+
+  private async _upstreamAgentCache(context: BrowserContextImpl) {
+    const agent = context._options.agent;
+    if (this._testInfo.status === 'passed' && agent?.cacheFile)
+      await this._testInfo._upstreamStorage(agent.cacheFile);
+  }
+
   async didCreateRequestContext(context: APIRequestContextImpl) {
     await this._startTraceChunkOnContextCreation(context, context._tracing);
   }
@@ -792,16 +821,6 @@ function tracing() {
   return (test.info() as TestInfoImpl)._tracing;
 }
 
-async function agentCacheWorkerFile(agent: PlaywrightTestOptions['agent'], testInfo: TestInfoImpl): Promise<string | undefined> {
-  if (!agent || agent.cacheMode === 'ignore')
-    return undefined;
-  if (!agent.cacheFile && !agent.cachePathTemplate)
-    return undefined;
-
-  const cacheFile = agent.cacheFile ?? testInfo._applyPathTemplate(agent.cachePathTemplate!, 'cache', '.json');
-  return await testInfo._cloneStorage(cacheFile);
-}
-
 export const test = _baseTest.extend<TestFixtures, WorkerFixtures>(playwrightFixtures);
 
 export { defineConfig } from './common/configLoader';

tests/library/browsercontext-reuse.spec.ts

@@ -37,7 +37,8 @@ class LaunchScenario {
     const browser = await this.browser();
     if (this._context)
       await (browser as any)._disconnectFromReusedContext('reusedContext');
-    const defaultContextOptions = (this._browserType as any)._playwright._defaultContextOptions;
+    const defaultContextOptions = {};
+    await (this._browserType as any)._instrumentation.runBeforeCreateBrowserContext(defaultContextOptions);
     this._context = await (browser as any)._newContextForReuse({ ...defaultContextOptions, ...options });
     return this._context;
   }
@@ -67,7 +68,8 @@ class ConnectScenario {
     if (this._browser)
       await this._browser.close();
     this._browser = await this._browserType.connect(server.wsEndpoint());
-    const defaultContextOptions = (this._browserType as any)._playwright._defaultContextOptions;
+    const defaultContextOptions = {};
+    await (this._browserType as any)._instrumentation.runBeforeCreateBrowserContext(defaultContextOptions);
     return await (this._browser as any)._newContextForReuse({ ...defaultContextOptions, ...options });
   }
 

tests/library/global-fetch.spec.ts

@@ -255,7 +255,9 @@ it('should set playwright as user-agent', async ({ playwright, server, isWindows
 });
 
 it('should be able to construct with context options', async ({ playwright, browserType, server }) => {
-  const request = await playwright.request.newContext((browserType as any)._playwright._defaultContextOptions);
+  const defaultContextOptions = {};
+  await (playwright as any)._instrumentation.runBeforeCreateRequestContext(defaultContextOptions);
+  const request = await playwright.request.newContext(defaultContextOptions);
   const response = await request.get(server.EMPTY_PAGE);
   expect(response.ok()).toBeTruthy();
   await request.dispose();